[ORT] Extend reshape fusion

[Honry]

Enable reshape fusion for models with symbolic (dynamic) dimensions, covering patterns found in MobileNetV2, vision encoders, and SLMs.

Changes:

• Match_Shape: allow dimension matching by symbolic dim_param, not just static values. Needed for SLMs (e.g. Qwen) where position_ids takes Shape from input_ids — same dim_param names, no static values.
• Match_Shape: allow matching when only the gathered dimension agrees between Shape source and reshape root. Needed for vision encoders (e.g. moondream2) where Shape is taken from pre-projection tensor but Reshape operates on post-projection output.
• Match_Shape: add GlobalAveragePool passthrough — accept Shape from pre-pool input when gathering batch dim. Needed for MobileNetV2.
• ReshapeHelper: handle zero-dim shapes (allow_zero=true) by computing unknown dim from non-zero dimension product only. Needed for MobileNetV2 reshape fusion output.
• Add corresponding unit tests.

[Honry]

@guschmue, PTAL, thanks!

[Honry]

@guschmue, friendly ping, or could you point me to the proper owner?

[fdwr]

@yuslepukhin @hariharans29 - In Guenther S's absence, do you know how might be appropriate to review this change? TY.

[hariharans29]

@yuslepukhin @hariharans29 - In Guenther S's absence, do you know how might be appropriate to review this change? TY.

I will take a look. (Sadly, going forward, it will be in @yuslepukhin 's absence too).

[Copilot]

Pull request overview

This PR extends the ReshapeFusion graph transformer to fuse additional reshape-shape construction patterns when tensor shapes involve symbolic (dim_param) dimensions or when the Shape source tensor differs from the reshape root (including a narrow GlobalAveragePool passthrough). It also updates ReshapeHelper’s handling of 0 dimensions when inferring an unknown (-1) dimension, and adds optimizer unit tests for the new fusion matches.

Changes:

• Extend ReshapeFusion::Match_Shape to allow matching by dim_param and (optionally) by only the gathered dimension, plus a GlobalAveragePool passthrough for batch-dim extraction.
• Update ReshapeHelper inference when zeros are present alongside an unknown dimension.
• Add graph transformation unit tests covering symbolic-dim and cross-tensor Shape sources, and the GlobalAveragePool passthrough case.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/core/optimizer/reshape_fusion.cc	Expands shape-source matching logic (symbolic dims, gathered-dim-only match, GAP passthrough) to enable additional reshape fusions.
onnxruntime/core/providers/cpu/tensor/reshape_helper.h	Adjusts unknown-dimension inference when zero dimensions are involved.
onnxruntime/test/optimizer/graph_transform_test.cc	Adds unit tests validating the new reshape fusion patterns.

[Honry]

Thank you, @fdwr, @hariharans29.

Comments addressed, PTAL again, thanks!

[hariharans29]

Review: PR #29027 — [ORT] Extend reshape fusion

Approve with minor observations. Real-world model-driven improvement — three additional Match_Shape paths let ReshapeFusion fold the Shape → Gather → Unsqueeze → Concat → Reshape construction found in Qwen (position_ids reshaped using Shape(input_ids)), moondream2 (Shape from pre-projection tensor, Reshape on post-projection tensor), and MobileNetV2 (GlobalAveragePool passthrough). Coupled with a companion fix in ReshapeHelper for the allowzero=1 + zero-dim + unknown-dim corner case that MobileNetV2's fused output hits. Tests cover the three new fusion patterns and the three ReshapeHelper cases. Existing static-value path is preserved as the first check, so no legitimate fusion regresses.

What looks correct

1. Ordering of the three new match paths. In Match_Shape the checks are layered from most-restrictive to most-permissive: (1) GlobalAveragePool producer-based passthrough (structural), (2) same-rank dim-by-dim static-value-or-dim_param match, (3) gathered-index-only static-value-or-dim_param match, and (4) the pre-existing multi-hop p_node_before_shape recursion still runs at the end. The static-value optimizer_utils::CompareShape remains the first attempt inside the shape-info branch, so nothing that previously fused stops fusing.
2. Path 1 (GlobalAveragePool passthrough) safety gates are correct. OpType() == "GlobalAveragePool" AND producer.InputDefs()[0] == shape_input AND expected_gather_index == 0. GlobalAveragePool preserves dim 0 (batch) and dim 1 (channel) and only reduces spatial dims to 1, so restricting the accept to index 0 is strictly conservative. root_input_producer->InputDefs().size() > 0 && ... != nullptr guards are defensive-and-cheap.
3. Path 2 (same-rank dim_param) is exhaustive. Every dim must match by dim_value == dim_value OR dim_param == dim_param; the mixed-form case (one has value, one has param) and both-unset case correctly fall through to all_dims_match = false. This closes the pre-fix hole where two tensors with identical [batch_size, sequence_length] symbolic shapes would fail CompareShape (which only handles dim_value).
4. Path 3 (gathered-index-only) semantics. The expected_gather_index threaded through from Match_One_Element_Output_Subgraph_1 is exactly the concat position, which by the surrounding IsInitializerWithExpectedValue(..., int64_t(shape_value.size()), false) check equals the Gather indices value. So the invariant being enforced ("the Shape source and reshape root agree on the gathered dim") is the right one — each Shape → Gather → Unsqueeze chain contributes exactly one dim to the concat, and only that dim needs to agree for the substitution to be sound. Bounds checks expected_gather_index < shape_input_shape->dim_size() and < root_input_shape->dim_size() are present.
5. ReshapeHelper zero-dim inference is internally consistent. Walked all three new tests:
   • [2,0,3] shape [0,-1] → size_for_inference = 1, input_shape_non_zero_size = 6, unknown = 6/1 = 6 → [0,6] ✓
   • [4,0] shape [0,2,-1] → size_for_inference = 2, input_shape_non_zero_size = 4, unknown = 2 → [0,2,2] ✓
   • [3,0,5] shape [0,-1] → size_for_inference = 1, input_shape_non_zero_size = 15, unknown = 15 → [0,15] ✓
     The has_zero_dim = allow_zero assignment is correct because upstream !allow_zero && requested_shape[i] == 0 has already been substituted to input_shape[i] before this line, so a residual dim == 0 here can only originate from the allow_zero == true path. The else (no zero) branch is algebraically identical to the pre-PR code (size == requested_shape_size when no zeros). Overflow guards are added for both size_for_inference *= dim and input_shape_non_zero_size *= input_shape[i].
6. Documentation update to the fusion diagram correctly reflects that Shape may now feed from Ancestor of Root or Root itself.

Non-blocking observations

1. Path 3 relies on the ONNX "same dim_param name = same runtime value" convention, which is a widely-followed convention but not a strictly enforced part of the spec. A graph produced by a sloppy or malicious exporter could conceivably have coincidentally-same dim_param names for unrelated axes, and this fusion would then substitute a symbolically-equal-but-runtime-different value. This is the standard risk taken by every optimizer that reasons about symbolic dims in ORT, so accepting it is consistent, but a one-line comment above the path 3 block noting "relies on ONNX convention: identical dim_param names denote identical runtime values" would document the trust boundary for future readers.
2. Missing negative tests. All three new tests exercise cases where fusion should happen. Cheap additions worth considering:
   • Same-rank shapes where the gathered dim's dim_param names differ ([batch, seq] vs [other_batch, seq], gather 0) — must NOT fuse via path 2 or 3.
   • GlobalAveragePool passthrough but with expected_gather_index == 1 (channel dim) — must NOT fuse via path 1 as currently gated (fusion via path 3 might still succeed if dim 1 matches — that's fine and correct, but explicitly asserting path-1's index == 0 gate would prevent an accidental future widening from silently going through).
   • Different rank between shape_input and root_input — must NOT fuse via path 2 (rank check enforces this), and must not fuse via path 3 unless the bounded gathered index happens to match — which is fine and correct.
3. Overlap between path 1 and path 3. For the GlobalAveragePool test, both path 1 and path 3 would allow fusion (dim 0 is static 1 in both input and pool_out). Path 1's real value is when shape inference hasn't populated NodeArg::Shape() on the pool output (so path 3's shape_input_shape != nullptr && root_input_shape != nullptr guard would fail). A short comment explaining that path 1 is the shape-info-missing fallback would help.
4. Path 1 could safely extend to expected_gather_index == 1 (channel dim). GlobalAveragePool preserves dims 0 AND 1 by ONNX spec. Restricting to == 0 is strictly conservative and fine — just calling out that a future extension is safe if a real model needs it.
5. Two logically separable changes in one PR. The reshape_fusion.cc extension and the reshape_helper.h zero-dim fix are motivated by the same MobileNetV2 model but exercise entirely independent code paths and could be reviewed separately. Not a blocker; noting for future PRs. If either one turns out to need a revert, splitting would have made that cleaner.
6. ReshapeHelper semantics for allowzero=1 + zero-dim + -1 is spec-ambiguous. ONNX's Reshape spec doesn't cleanly define behavior when the requested shape mixes explicit 0 (allowzero-preserved) with -1 (infer). The chosen convention here — "unknown dim = non-zero product of input / non-zero product of requested" — is defensible and internally consistent, and it's applied uniformly via reshape_helper.h so all EPs that consume the header stay in sync. Worth confirming this matches what the exporter that produced the affected MobileNetV2 graph expects at runtime, and worth a comment in the has_zero_dim branch explaining the chosen convention (it's non-obvious without walking the algebra).
7. Checks 1 on the PR is unusually low for an optimizer + EP-independent header change. Worth confirming that the full CI matrix has actually run (not just a doc/style gate) before merging — normally optimizer + reshape kernel changes fan out across a dozen+ EPs' test legs.
8. The added reshape_helper.h code path also runs for every allowzero=1 Reshape on every EP, not just the newly fused MobileNetV2 case. Nothing wrong with that (the semantics are strictly more permissive than the pre-PR ORT_ENFORCE(requested_shape_size != 0 && ...) which threw), but it's a good idea to be aware that this behavior change reaches beyond the fusion motivation.
9. Comment quality in Match_Shape. The added block comments are clear and self-describing. Nice.

Summary

Approve pending full CI green. Correctness of the three new fusion paths and the ReshapeHelper zero-dim handling checks out. The main soft concern is the reliance on ONNX's "same dim_param = same runtime value" convention in path 3, which is defensible and consistent with other optimizers, but worth a one-line acknowledgment in a code comment. Adding a couple of negative graph-transform tests and a comment on the chosen has_zero_dim inference convention would harden the change; both are non-blocking follow-ups.

[Honry]

@hariharans29, there's one CI failure and I've fixed it in latest commit, PTAL again, thanks!

[hariharans29]

Hm m - not sure why the CI checks didn't complete. Could you please rebase with main ? I ll kick off another round of CI with that.

[Honry]

Hm m - not sure why the CI checks didn't complete. Could you please rebase with main ? I ll kick off another round of CI with that.

@hariharans29, done, please help rerun the CI, thanks!

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[hariharans29]

Seems like there are some issues with our CI pipelines that we are working to resolve. I will take the next step on this PR once those are resolved. Thanks.