Repository navigation
fix(compiler): refine incomplete shapes, unused Split outputs, and unit-slice copies - #1306
Conversation
Unranked Concat inference crashed on a null type, and ranked results with dynamic dimensions stayed unresolved. Rank LeakyRelu and ConvTranspose and fold static Upsample scales so conversion can bufferize those ops. Co-Authored-By: Cursor Grok 4.7 <cursoragent@cursor.com> Made-with: Cursor
|
Thanks for opening a PR! This project follows LLVM's incremental-development and AI-tool-use Before requesting review, please check that:
Reviewers are assigned through |
L2 Accuracy Results (EP vs CPU)
Threshold: 0.01 | Run: 7837 - Commit: |
Binary sizesBaseline: run 39935981292 -
MorphiZen EP Performance Results
EPContext Export Performance
EPContext Import Performance
OGA Benchmark Results
OGA Wheel Smoke (Python benchmark_e2e.py)
VLM Benchmark Results
Run: 7837 - Commit: |
An unused Split result typed none made the rewrite bail after creating ops, so conversion never converged and failed with no diagnostic. The one-element slices that Split bufferizes to were then copied with host loads of device memory. Co-Authored-By: Cursor Grok 4.7 <cursoragent@cursor.com> Made-with: Cursor
Regression: Gemma-4 E4B falls back to CPU after this PR (decode 48 → 8 TPS)Hi @pragyam32 @asparkhi (cc @cliu1003 @amd-mingw @qianglin-amd , since #1272 now carries a revert of this PR) Symptom. With
The log shows repeated: I bisected it to this commit: 940d3f0 is clean and d5979f7 fails. Root cause. The PR's "Re-run shape rules for ranked results that are not fully static" change makes Minimal lit repro ( func.func @slice_result_widened(%x: tensor<1x?x12x64xf16>,
%starts: tensor<1xi64>, %ends: tensor<1xi64>, %axes: tensor<1xi64>)
-> tensor<1x?x12x32xf16> {
%y = "onnx.Slice"(%x, %starts, %ends, %axes)
: (tensor<1x?x12x64xf16>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>)
-> tensor<1x?x12x32xf16>
return %y : tensor<1x?x12x32xf16>
}A second, smaller issue in the same PR (Gemma-4 does not hit it). A static equal Split now uses constant chunk lengths. If the declared result type leaves the split axis dynamic, the slice is still built with that declared type and fails the same way: // hip-mlir-opt --hip-add-context-arg --convert-onnx-to-hip
%a, %b = "onnx.Split"(%x) {axis = 3 : si64}
: (tensor<1x?x12x64xf16>) -> (tensor<1x?x12x?xf16>, tensor<1x?x12x?xf16>)
// error: expected type to be 'tensor<1x?x12x32xf16>' or a rank-reduced version. (size mismatch)Happy to share the full logs or IR dumps if useful. |
Hi @donghanbj and @nvunnam57128, I may be missing some context here. Could you share some background on why this loop-body pass was extended to do broader shape inference? I also briefly looked at #1216 and noticed that some shape-inference logic overlaps with existing helpers. I have been working on cleaning up this duplication. Could we discuss how to reuse the common rules where possible, while preserving the shape information we already have? For context, #639 includes details of the refactoring series I have been working on: |
Hi @fhanuman, thanks for looking into this. I hit the regression while working on Gemma-4 E4B decode performance and bisected it to this PR. From the history, as far as I can tell:
So @nvunnam57128 and @pragyam32 are better placed to explain the motivation for the broader inference. From the Gemma-4 side, preserving the shape information the importer already provides is exactly what matters. I can share the full logs or IR dumps (before/after IR dumps (before/after |
|
Mistral-Small (and others) fail locally and in CI: Could this be related to this change? (I didn't bisect yet) |
…, and unit-slice copies (#1306)" Shape rules again run only for unranked results, so a stored dimension is not replaced by a less specific rule. Unused Split outputs and unit-extent copies return to the previous lowering. Co-Authored-By: Cursor Grok 4.7 <cursoragent@cursor.com> Made-with: Cursor
## Summary - Merge a shape rule with the type already stored on the result: keep a known size, and fill `?` only when the rule has a number. An unranked `tensor<*xT>` still takes the rule type. - Read Slice, Pad, Unsqueeze, reduction, and Split bounds from a constant input when the attribute is absent. A runtime tensor returns no type, so the stored shape stays. - Rank AveragePool (`1x64x512x512`, kernel 2, stride 2 becomes `1x64x256x256`) and opset 7-8 Upsample from its `scales` attribute. An equal Split of a static axis writes the chunk size into `extract_slice` when that axis is still `?`. ## Related issue or design Follow-up to the Gemma-4 shape regression reported on #1306. None filed separately. ## Why `#1306` runs the shape rules on any ranked result that still contains a `?`, then replaces the stored type. Slice and Pad, when starts/ends/pads are inputs rather than attributes, answered with a same-rank all-`?` type, so a stored size such as `tensor<1x?x12x32xf16>` became `tensor<?x?x?x?xf16>`. The Topaz UNet also left `AveragePool` and attribute-form `Upsample` unranked, so conversion never saw a rank. ## What - `mergeInferredType` in `forwardInferUnrankedResults` combines the stored type with the rule result. - AveragePool uses the ONNX spatial formula. Upsample reads the `scales` attribute when there is no scales input. - Slice, Pad, Unsqueeze, reductions, and Split accept a constant input vector and return no type for a runtime vector. - `SplitToStdTensor` puts a constant chunk size on the `extract_slice` result and casts back to the declared type. ## Test plan - [x] `pre-commit run --files` on the four changed files (passed after lintrunner formatting) - [x] `hip-mlir-opt --hip-infer-loop-body-shapes test/lit/Dialect/hip-infer-loop-body-shapes.mlir | FileCheck` - [x] `hip-mlir-opt test/lit/Conversion/onnx-to-hip/test_split.mlir --hip-add-context-arg --convert-onnx-to-hip | FileCheck` - [ ] Gemma-4 and the Topaz UNet on a full session ## Notes for reviewers Drafted with Cursor assistance. I reviewed the merge rule, the constant-input fallback, and the lit checks. The Gemma Slice case is locked as `tensor<1x?x12x32xf16>` when starts, ends, and axes are runtime tensors. The pool and attribute Upsample cases are in `hip-infer-loop-body-shapes.mlir`. ## Checklist - [x] The change is focused, or links a design/series explaining its scope. - [x] Relevant tests were added or updated and the results are documented. - [x] User-facing or design documentation was updated when needed. - [x] Substantial AI assistance is disclosed, and I reviewed and understand the result. --------- Co-authored-by: Cursor <cursoragent@cursor.com>
Summary
Shape inference no longer crashes on an unranked Concat, and it fills static extents for ranked results that still contain dynamic dimensions. Conversion also finishes a Split whose unused output is typed
none, and a one-element strided slice of that Split is copied on the device instead of with host loads.Related issue or design
None.
Why
Three failures stopped a graph before a correct device run.
inferConcatResultreturned a null type when its operands were unranked, and the caller passed that type todyn_cast. In this MLIR build that cast is unsafe and aborted the compiler. The same pass refined only unranked results, so a ranked result that still used dynamic dimensions, including Upsample scales that fold to constants, was left unchanged. Conversion then rejected the dynamic spatial result and bufferization stopped.An exporter names an output it does not use with the empty string, and the importer types that result
none. The Split rewrite created slices and then gave up on that result. The leftover ops made the greedy driver re-run the pattern on every iteration, so ONNX-to-HIP failed with no operator diagnostic.Splitting a row into single elements bufferizes to a rank-1 window whose stride is the parent pitch. That copy was not recognized as contiguous, the pitched 2D path does not cover rank 1, and the default MemRef lowering emitted host loads and stores of device pointers.
What
RankedTensorTypefrom Concat inference and decline a null result instead of casting it.output_shape.noneresult becomesonnx.NoValueand still advances the slice offset. A static equal split uses constant lengths.Test plan
pre-commit run --fileson the changed files: passed. Lintrunner reformatted on the first run; the second run passed with no further edits.hip-mlir-opt --hip-infer-loop-body-shapeswith FileCheck ontest/lit/Dialect/hip-infer-loop-body-shapes.mlir: passed. Cases cover unranked Concat, ConvTransposeoutput_shape, and Upsample dynamic-dimension refinement.MorphizenMLIRLitTests: passed, includingtest_split.mlirfor an unusednoneSplit output andtest_memref_copy.mlirfor unit-extent strided copies.Notes for reviewers
output_shapewhose length matches the spatial rank. Other forms are left unchanged.noneSplit result is not sliced. The chunk it stands for still shifts later offsets, using the equal-split length because that result has no type of its own to measure.Checklist