Fix stablehlo.reshape source type for %arg0 consumers (closes #518)#521
Merged
michalharakal merged 1 commit intodevelopfrom Apr 18, 2026
Merged
Fix stablehlo.reshape source type for %arg0 consumers (closes #518)#521michalharakal merged 1 commit intodevelopfrom
michalharakal merged 1 commit intodevelopfrom
Conversation
ShapeOperationsConverter emitted `: ($outputType) -> $outputType` on every reshape / flatten / squeeze / unsqueeze, reusing the output type on the source side of the cast. When the operand was a function argument (`%arg0: tensor<80x3000xf32>`), the result was syntactically valid but wrong — e.g. `stablehlo.reshape %arg0 : (tensor<1x80x3000xf32>) -> tensor<1x80x3000xf32>` — and iree-compile rejected it with a type mismatch on every native SKaiNET DSL export that reached the IREE path. Fix threads an SSA-value-name -> MLIR-type map through ConversionContext: - StableHloConverter seeds `%argN -> mapTensorType(inputSpec)` when it writes the function signature, and records `resultValue -> outputType` after each successful op conversion. - ShapeOperationsConverter looks up the operand's real type via a new resolveOperandType helper (context.getValueType first, then node.inputs[0], then a dynamic fallback) and uses that on the source side of every reshape cast. The type map is consumed opt-in by converters that need it, so other converters (concat, slice, dot_general, transpose) are unaffected. Regression test `testReshapeOnArgUsesDeclaredArgType` builds an input -> unsqueeze graph and asserts the emitted MLIR contains `(tensor<80x3000xf32>) -> tensor<1x80x3000xf32>`. Full :skainet-compile-hlo:jvmTest passes. Closes #518 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ConversionContext; seed it with%argNdeclared types at function-signature emission and populate it after each successful op conversion.ShapeOperationsConverterto look up the operand's real type via the map (withnode.inputs[0]fallback) instead of reusingoutputTypeon the source side of everystablehlo.reshapecast.Background
iree-compilerejected SKaiNET-emitted MLIR on every reshape / unsqueeze consuming a function argument. The converter emitted: (outputType) -> outputType, producing e.g.stablehlo.reshape %arg0 : (tensor<1x80x3000xf32>) -> tensor<1x80x3000xf32>when%arg0was actuallytensor<80x3000xf32>. Third of four items from the skainet-whisper IREE Vulkan bring-up; first two landed in #520.Test plan
:skainet-compile:skainet-compile-hlo:jvmTestpasses, including new regressionShapeOperationsConverterTest.testReshapeOnArgUsesDeclaredArgType.*WhisperHloExportTest*+iree-compileonfeature/iree-vulkan-gpuafter merge to confirm the type-mismatch error is gone and only HLO: inline dense constants for large tensors — scale / compile-time blocker #519 (inline constants) remains blocking.Closes #518
🤖 Generated with Claude Code