Skip to content

Fix stablehlo.transpose and stablehlo.dot_general MLIR emission#520

Merged
michalharakal merged 1 commit intodevelopfrom
feature/hlo-linalg-mlir-emission
Apr 18, 2026
Merged

Fix stablehlo.transpose and stablehlo.dot_general MLIR emission#520
michalharakal merged 1 commit intodevelopfrom
feature/hlo-linalg-mlir-emission

Conversation

@michalharakal
Copy link
Copy Markdown
Contributor

iree-compile rejected MLIR emitted by LinalgOperationsConverter because the converter used syntax that does not match the StableHLO spec:

  • transpose emitted : outputType instead of : (inputType) -> outputType, which iree-compile flagged as "invalid kind of type specified".
  • dot_general emitted contracting_dims = [[N], [M]] (ArrayAttr-style) instead of contracting_dims = [N] x [M], omitted batching_dims entirely for batched matmuls, and dropped the (lhsType, rhsType) -> input-type clause — all three are required by the StableHLO parser.

LinalgOperationsConverter now threads input TensorSpecs through to the emission, derives batching_dims from lhs rank for bmm (supporting 2D, 3D, and 4D+ inputs instead of the previous 3D-only assumption), and writes the (lhs, rhs) -> result type tuple. Tests updated accordingly; full :skainet-compile-hlo:jvmTest passes.

These are the first two of four items from a skainet-whisper IREE Vulkan bring-up on branch feature/iree-vulkan-gpu (SKaiNET 0.18.0). The remaining two are tracked separately.

Related: #518, #519

iree-compile rejected MLIR emitted by LinalgOperationsConverter because
the converter used syntax that does not match the StableHLO spec:

- transpose emitted `: outputType` instead of `: (inputType) -> outputType`,
  which iree-compile flagged as "invalid kind of type specified".
- dot_general emitted `contracting_dims = [[N], [M]]` (ArrayAttr-style)
  instead of `contracting_dims = [N] x [M]`, omitted batching_dims entirely
  for batched matmuls, and dropped the `(lhsType, rhsType) ->` input-type
  clause — all three are required by the StableHLO parser.

LinalgOperationsConverter now threads input TensorSpecs through to the
emission, derives batching_dims from lhs rank for bmm (supporting 2D,
3D, and 4D+ inputs instead of the previous 3D-only assumption), and
writes the `(lhs, rhs) -> result` type tuple. Tests updated accordingly;
full :skainet-compile-hlo:jvmTest passes.

These are the first two of four items from a skainet-whisper IREE Vulkan
bring-up on branch `feature/iree-vulkan-gpu` (SKaiNET 0.18.0). The
remaining two are tracked separately.

Related: #518, #519

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@michalharakal michalharakal merged commit b5c11e0 into develop Apr 18, 2026
4 checks passed
@michalharakal michalharakal deleted the feature/hlo-linalg-mlir-emission branch April 18, 2026 07:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant