Skip to content

fix(compiler): refine incomplete shapes, unused Split outputs, and unit-slice copies - #1306

Merged
asparkhi merged 2 commits into
mainfrom
fix/refine-incomplete-onnx-shapes
Oct 23, 2026
Merged

asparkhi merged 2 commits into
mainfrom
fix/refine-incomplete-onnx-shapes

Conversation

@pragyam32

@pragyam32 pragyam32 commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

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.

inferConcatResult returned a null type when its operands were unranked, and the caller passed that type to dyn_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

  • Return RankedTensorType from Concat inference and decline a null result instead of casting it.
  • Re-run shape rules for ranked results that are not fully static.
  • Preserve the input shape through LeakyRelu.
  • Infer ConvTranspose from the input rank, weight channels, group, and output_shape.
  • Fold constant Upsample scales, including scales produced through Div, Concat, and Cast, and propagate the refined shape through later shape-preserving ops.
  • Classify every Split result before creating ops. A none result becomes onnx.NoValue and still advances the slice offset. A static equal split uses constant lengths.
  • Treat a memref copy as one flat device transfer when every dimension wider than one element is contiguous. A stride on a dimension of extent 1 cannot introduce a hole.

Test plan

  • pre-commit run --files on 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-shapes with FileCheck on test/lit/Dialect/hip-infer-loop-body-shapes.mlir: passed. Cases cover unranked Concat, ConvTranspose output_shape, and Upsample dynamic-dimension refinement.
  • MorphizenMLIRLitTests: passed, including test_split.mlir for an unused none Split output and test_memref_copy.mlir for unit-extent strided copies.
  • Rebuilt the execution provider and completed an end-to-end session with CPU fallback disabled for a graph that previously failed conversion and then faulted during inference.

Notes for reviewers

  • ConvTranspose inference requires a present output_shape whose length matches the spatial rank. Other forms are left unchanged.
  • A none Split 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.
  • Copies whose only non-dense stride sits on an extent-1 dimension now use the flat memcpy path. The pitched 2D Concat tests use a batch of 2 so they still exercise a real row pitch.
  • No user-facing documentation change. The behavior stays inside the existing shape pass, ONNX-to-HIP conversion, and memref copy lowering.
  • Cursor assisted with the diagnosis, implementation, tests, and validation. I reviewed the result.

Checklist

  • The change is focused, or links a design/series explaining its scope.
  • Relevant tests were added or updated and the results are documented.
  • User-facing or design documentation was updated when needed.
  • Substantial AI assistance is disclosed, and I reviewed and understand the result.

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
@github-actions

Copy link
Copy Markdown

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

  1. The change is focused. Substantial work links the relevant issue
    or design discussion.
  2. The PR documents relevant test results and updates affected
    documentation.
  3. If AI tools provided substantial assistance, the description
    explains what was assisted and how it was validated, and commit
    trailers identify the tool. The contributor has reviewed and
    understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

@github-actions

github-actions Bot commented Oct 23, 2026 •

Copy link
Copy Markdown

L2 Accuracy Results (EP vs CPU)

Model Combined L2 Total Elems Skipped NaN/Inf
conv_test_hybrid 1.12256E-07 64 0
GroupQueryAttention_seq256 59.2561 2621440 0
MatMulNBits_o_seq128 259.905 368640 0
QMoE_seq128 24.2583 368640 0

Threshold: 0.01 | Run: 7837 - Commit: abcb14f

@github-actions

github-actions Bot commented Oct 23, 2026 •

Copy link
Copy Markdown

Binary sizes

Baseline: run 39935981292 - ae4e929 (dll-sizes)

File This (MB) Main (MB) Delta (MB)
hipgpu.dll 59.26 59.26 0.00
custom_kernels_gfx1170.dll 5.51 5.51 0.00
custom_kernels_gfx11-generic.dll 5.80 5.80 0.00

MorphiZen EP Performance Results

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 8.78 7.31 245 3 1157
GroupQueryAttention_seq128 11415.30 2.42397 1 5 230
matmul_down_seq128 510.67 3.09 126 3 276

EPContext Export Performance

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 8.82 49.40 247 3 15504

EPContext Import Performance

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 8.83 8.92 256 3 15638

OGA Benchmark Results

Model Warmup Reps Prompt Len Gen Tokens TTFT (ms) TPS Peak Mem (GB) GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32 1 5 128 128 95.1 75.0 1.25 12.24
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml 1 5 128 128 124.1 47.8 1.15 5.23

OGA Wheel Smoke (Python benchmark_e2e.py)

Model TTFT (ms) TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml 123 48.0

VLM Benchmark Results

Model TTFT (ms) TPS Output
Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx 14384.6 14.0 The user wants a detailed description of the image. Let me carefully observe all elements before writing. Subject: A white do

Run: 7837 - Commit: abcb14f

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
@pragyam32 pragyam32 changed the title fix(onnx): refine incomplete tensor shapes before conversion fix(compiler): refine incomplete shapes, unused Split outputs, and unit-slice copies Oct 23, 2026
@asparkhi
asparkhi merged commit d5979f7 into main Oct 23, 2026
11 checks passed
@asparkhi
asparkhi deleted the fix/refine-incomplete-onnx-shapes branch October 23, 2026 17:17
cliu1003 added a commit that referenced this pull request Oct 25, 2026
@donghanbj

donghanbj commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

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 main at d5979f7, Gemma-4 E4B (AWQ uint4, split fp16 export, OGA, gfx1151 / Windows) fails to compile 8 subgraphs (vision encoder and decoder), which then run on the CPU EP. The output text is still correct, but it is much slower:

940d3f0 (parent) d5979f7
subgraphs failing to compile 0 8
decode TPS (388-token prompt) 48.3 7.8
TTFT 332 ms 4.4 s

The log shows repeated:

error: expected type to be 'tensor<1x?x12x32xf16>' or a rank-reduced version. (size mismatch)
MlirCompiler.cpp:128] Compilation failed: MLIR pass pipeline failed

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 forwardInferUnrankedResults replace the result type with whatever the rule returns. A rule can know less than the importer did. For example, inferSliceResult cannot read opset-10 Slice bounds, because they are operands rather than attributes, so it returns an all-dynamic type of the same rank. On Gemma-4 E4B, hip-infer-loop-body-shapes changed 736 result types, every one of them to a less refined type and none to a more refined one. For example, the vision RoPE Slice goes from 1x?x12x64 -> 1x?x12x32 to 1x?x12x64 -> ?x?x?x?. Slice conversion then builds tensor.extract_slice with static sizes (…, 32) and the widened declared type, which fails verification in the next pass.

Minimal lit repro (--hip-infer-loop-body-shapes). The Slice result is widened to ?x?x?x?:

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.

@fhanuman

fhanuman commented Oct 25, 2026 •

Copy link
Copy Markdown
Collaborator

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 main at d5979f7, Gemma-4 E4B (AWQ uint4, split fp16 export, OGA, gfx1151 / Windows) fails to compile 8 subgraphs (vision encoder and decoder), which then run on the CPU EP. The output text is still correct, but it is much slower:

940d3f0 (parent) d5979f7
subgraphs failing to compile 0 8
decode TPS (388-token prompt) 48.3 7.8
TTFT 332 ms 4.4 s
The log shows repeated:

error: expected type to be 'tensor<1x?x12x32xf16>' or a rank-reduced version. (size mismatch)
MlirCompiler.cpp:128] Compilation failed: MLIR pass pipeline failed

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 forwardInferUnrankedResults replace the result type with whatever the rule returns. A rule can know less than the importer did. For example, inferSliceResult cannot read opset-10 Slice bounds, because they are operands rather than attributes, so it returns an all-dynamic type of the same rank. On Gemma-4 E4B, hip-infer-loop-body-shapes changed 736 result types, every one of them to a less refined type and none to a more refined one. For example, the vision RoPE Slice goes from 1x?x12x64 -> 1x?x12x32 to 1x?x12x64 -> ?x?x?x?. Slice conversion then builds tensor.extract_slice with static sizes (…, 32) and the widened declared type, which fails verification in the next pass.

Minimal lit repro (--hip-infer-loop-body-shapes). The Slice result is widened to ?x?x?x?:

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:
#639

@donghanbj

Copy link
Copy Markdown
Contributor

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 main at d5979f7, Gemma-4 E4B (AWQ uint4, split fp16 export, OGA, gfx1151 / Windows) fails to compile 8 subgraphs (vision encoder and decoder), which then run on the CPU EP. The output text is still correct, but it is much slower:
940d3f0 (parent) d5979f7
subgraphs failing to compile 0 8
decode TPS (388-token prompt) 48.3 7.8
TTFT 332 ms 4.4 s
The log shows repeated:

error: expected type to be 'tensor<1x?x12x32xf16>' or a rank-reduced version. (size mismatch)
MlirCompiler.cpp:128] Compilation failed: MLIR pass pipeline failed

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 forwardInferUnrankedResults replace the result type with whatever the rule returns. A rule can know less than the importer did. For example, inferSliceResult cannot read opset-10 Slice bounds, because they are operands rather than attributes, so it returns an all-dynamic type of the same rank. On Gemma-4 E4B, hip-infer-loop-body-shapes changed 736 result types, every one of them to a less refined type and none to a more refined one. For example, the vision RoPE Slice goes from 1x?x12x64 -> 1x?x12x32 to 1x?x12x64 -> ?x?x?x?. Slice conversion then builds tensor.extract_slice with static sizes (…, 32) and the widened declared type, which fails verification in the next pass.
Minimal lit repro (--hip-infer-loop-body-shapes). The Slice result is widened to ?x?x?x?:

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: #639

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 hip-infer-loop-body-shapes for the failing subgraphs), and I'm happy to validate any refactor on Gemma-4 E4B.

IR dumps (before/after hip-infer-loop-body-shapes for one vision and one decoder subgraph), the full list of widened results for all 10 compiled subgraphs, and the two repros are attached: gemma4-e4b-pr1306-ir-evidence.zip
gemma4-e4b-pr1306-ir-evidence.zip

@pschlan-amd

pschlan-amd commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Mistral-Small (and others) fail locally and in CI:
[2026-10-24T17:14:20.450Z] error: failed to legalize operation 'hip.rope' that was explicitly marked illegal: "hip.rope"(<<UNKNOWN SSA VALUE>>, %29873, %28385, %28432, %29892) <{interleaved = 1 : i64, num_heads = 32 : i64, rotary_embedding_dim = 64 : i64}> : (!hip.context, memref<?x?x?xf16>, memref<?x?x32xf16>, memref<?x?x32xf16>, memref<?x?x2048xf16>) -> ()
e.g. in http://xsjvitisaijenkins:8080/blue/organizations/jenkins/daily_hip_ep_genai/detail/daily_hip_ep_genai/2651/pipeline

Could this be related to this change? (I didn't bisect yet)

cliu1003 added a commit that referenced this pull request Oct 26, 2026
pragyam32 pushed a commit that referenced this pull request Oct 27, 2026
…, 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
asparkhi pushed a commit that referenced this pull request Oct 27, 2026
## 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>
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.

6 participants