Fix moving shape transforms through a shared reduction and fold unsqueeze into multibroadcast

[pfultz2]

Motivation

simplify_reshapes moves a shape-transform chain between a reduction and a pointwise op into a common space so the two can fuse. When the chain broadcasts a new axis (for example reduce_mean -> unsqueeze -> multibroadcast -> mul), the pass rebuilt the old reduction for its other consumers with to_src_from_common(), which produced an invalid reshape because the broadcasted axis has no source axis to map back to. The result was a MULTIBROADCAST: input shape {1, 1, 1, 1, 6} cannot be broadcasted error at compile time (#5360).

While fixing this, the unsqueeze -> multibroadcast chain itself turned out to be redundant: multibroadcast aligns from the right, so an unsqueezed axis that only follows 1s is added by the broadcast. Folding it exposes the reduction to the pass as a bare broadcast, which previously was always skipped, so the descriptor also needs to tell a broadcast of a new axis apart from a re-broadcast of a reduced axis.

Technical Details

shape_transform_descriptor:

• to_src_from_common() hides the axis of a squeezed-away dimension that came from a broadcast with no source axis ([1:3] becomes [1:$3]). rebase() now refuses to give such a hidden dimension of 1 a different length, so the inverse cannot silently merge the broadcasted axis into a neighbor.
• generate() drops axis-less dimensions that only follow input dimensions of 1 before emitting a multibroadcast, so unsqueeze{3} + multibroadcast{1,1,1,2,3} on {1,1,1,3} becomes a single multibroadcast. optimize_shape_transforms and find_nested_shape_transforms pick this up.
• rebase(dims, broadcast=true) detaches a hidden broadcast whose source axis is rebased to a size other than 1 or its own length. The axes shift down into a preceding axis-less dimension of 1 and the broadcast becomes its own axis-less dimension, which is the same form the unsqueeze chain produced. Broadcasting a reduced axis back to its size is unaffected.
• New has_axisless_broadcast() reports a broadcasted dimension with no source axis.

simplify_reshapes (find_op_shape_transform_op):

• The old reduction is only restored when to_src_from_common() rebases onto the common dims. Otherwise the rewrite proceeds only if the reduction is consumed solely through the chain (is_private_chain), so it becomes dead code, and it bails when the reduction is shared.
• A chain of only broadcasts is now viable when it feeds from a reduction and the rebased descriptor broadcasts a new axis. Plain re-broadcasts of reduced axes are still skipped to avoid a no-op rewrite loop.
• The chain-collection loop is shared as chain_ops.

Tests: new simplify_reshapes cases for the expanding chain, the shared element-preserving chain and the shared expanding chain, plus descriptor cases for the fold, the hidden inverse and the rebase detachment. Existing expectations that contained an unsqueeze absorbed by the broadcast were updated to the shorter equivalent chains (conv_add_layernorm_conv, dequantizelinear_entry_shape_transform_zero_point, pointwise_reshapes_reduce_shadowed_broadcast, dot_skinny_rewrite, dot_softmax_dot_projection_rewrite, optimize_multibroadcast_transpose, generate_shape_transforms_for), and slice_reshape_multibroadcast_rebase_axis compares output lens since only a size-1 stride changed.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

• • Added: New functionality.
• • Changed: Changes to existing functionality.
• • Removed: Functionality or support that has been removed. (Compared to a previous release)
• • Optimized: Component performance that has been optimized or improved.
• • Resolved Issues: Known issues from a previous version that have been resolved.
• • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

[Copilot]

🟡 Changes recommended

A shared expanding reduction can shadow a later rewritable input and prevent the intended optimization.

1 open finding

• Avoid aborting before checking later eligible reshape inputs · New

What changed in this PR

Fixes invalid shape-transform rewrites around shared reductions and simplifies redundant unsqueeze/broadcast chains.

Changes:

• Strengthens descriptor rebasing and inverse-transform handling.
• Prevents unsafe shared-reduction rewrites and adds regression coverage.
• Updates affected expected graphs and changelog.

Review performed as a single pass without subagent fan-out. Relevant tests still need CI validation.

File	Description
src/​simplify_reshapes.cpp	Guards reduction rewrites and supports axis-expanding broadcasts.
src/​shape_transform_descriptor.cpp	Implements broadcast detachment, folding, and inverse validation.
src/​include/​migraphx/​shape_transform_descriptor.hpp	Exposes axisless-broadcast detection.
test/​simplify_reshapes_test.cpp	Adds shared and expanding reduction regressions.
test/​shape_transform_descriptor.cpp	Tests folding, rebasing, and inverse behavior.
test/​rewrite_reduce.cpp	Updates optimized graph expectations.
test/​fuse_reduce.cpp	Updates reshape expectation after folding.
CHANGELOG.md	Records the resolved issue.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[urpetkov-amd]

This fixes all five failing models from #PR5360 . Thanks

[codecov]

Codecov Report

❌ Patch coverage is 97.08738% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/simplify_reshapes.cpp	92.00%	2 Missing ⚠️
src/shape_transform_descriptor.cpp	98.61%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #5365      +/-   ##
===========================================
+ Coverage    92.68%   92.72%   +0.04%     
===========================================
  Files          633      635       +2     
  Lines        36087    36265     +178     
===========================================
+ Hits         33446    33625     +179     
+ Misses        2641     2640       -1     

Files with missing lines	Coverage Δ
...rc/include/migraphx/shape_transform_descriptor.hpp	100.00% <ø> (ø)
src/simplify_algebra.cpp	97.62% <100.00%> (+<0.01%)	⬆️
src/shape_transform_descriptor.cpp	94.23% <98.61%> (+0.21%)	⬆️
src/simplify_reshapes.cpp	97.74% <92.00%> (-0.14%)	⬇️

... and 9 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.