Skip to content

fix(use-output-allocator): report graph-output rank for reshaped-view outputs - #557

Draft
anilmartha wants to merge 3 commits into
mainfrom
fix/output-allocator-graph-output-rank
Draft

fix(use-output-allocator): report graph-output rank for reshaped-view outputs#557
anilmartha wants to merge 3 commits into
mainfrom
fix/output-allocator-graph-output-rank

Conversation

@anilmartha

Copy link
Copy Markdown
Contributor

Summary

UseOutputAllocator rewrote a graph-output memref.alloc into hip.alloc_output using the producer alloc's type. When the returned value is a reshaping view of that alloc (e.g. a rank-2 Gemm result expand_shape'd to rank-3 with no trailing op — DETR's class_labels_classifier logits head), the EP registered the alloc's rank with ORT instead of the declared graph-output rank, so ORT rejected the output:

Invalid rank for output: logits Got: 2 Expected: 3

Outputs with a trailing op (e.g. pred_boxes' sigmoid, whose alloc is already at the output rank) were unaffected — which is why only logits failed.

Fix

  • Allocate hip.alloc_output at the func.return operand type (the graph-output shape).
  • When it differs from the producer alloc's type, rebuild the inverse of the view chain (cast / expand_shape / collapse_shape, static output) so the producer keeps writing into the same EP-owned buffer viewed as its original type — no extra copy.
  • Fast path (return type == alloc type) is unchanged; subview or dynamic-shaped outputs fall back to the legacy behavior, so no existing case regresses.

Applies to both the LLVM_IR and NATIVE artifact paths.

Test plan

  • facebook/detr-resnet-50 (fp16): logits now returns {1,100,92}; onnxruntime_perf_test runs clean.
  • Re-run existing lit/numeric suites for regressions (esp. conv/reshape output models).
  • Spot-check a model with a rank-4 conv output collapsed before return.

… outputs (fixes DETR Invalid rank for output: logits)
@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.

Co-authored-by: Cursor <cursoragent@cursor.com>
@amd-mingw

Copy link
Copy Markdown
Collaborator

#533 is fixing on it. cc @amd-genmingz

Co-authored-by: Cursor <cursoragent@cursor.com>
@github-actions

Copy link
Copy Markdown

L2 Accuracy Results (EP vs CPU)

Model Combined L2 Total Elems Skipped NaN/Inf
conv_test_hybrid 4.8668E-07 64 0
GroupQueryAttention_seq256 25.2366 2621440 0
MatMulNBits_o_seq128 259.906 368640 0
QMoE_seq128 34.957 368640 0

Threshold: 0.01 | Run: 3271 - Commit: c56b74c

@github-actions

Copy link
Copy Markdown

MorphiZen EP Performance Results

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 7.53 6.39 360 3 1244
GroupQueryAttention_seq128 4377.53 1.65253 12 6 310
matmul_down_seq128 524.48 2.34 75 3 352

EPContext Export Performance

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 7.58 44.35 369 3 15589

EPContext Import Performance

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 7.56 9.07 359 3 15759

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 164.7 80.2 1.33 13.57
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml 1 5 128 128 282.2 39.6 1.22 6.43

OGA Wheel Smoke (Python benchmark_e2e.py)

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

Run: 3271 - Commit: c56b74c

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.

2 participants