Repository navigation
Fix invalid multibroadcast shapes for shared reductions - #5360
urpetkov-amd wants to merge 6 commits into
Conversation
| // An element-expanding chain can be moved through a reduction only when | ||
| // the reduction is not shared. The rewrite replaces x_ins globally, so | ||
| // changing its shape would invalidate its other consumers. | ||
| return is_reduce(x_ins) and x_ins->outputs().size() == 1; |
There was a problem hiding this comment.
We already insert a reshape back to the source for other consumers.
We reshape back to the src here:
AMDMIGraphX/src/simplify_reshapes.cpp
Line 527 in 862a8fe
And then replace x_ins with this reshaped input:
AMDMIGraphX/src/simplify_reshapes.cpp
Line 542 in 862a8fe
So there shouldn't be incompatible shapes.
There was a problem hiding this comment.
You’re right that the pass intends to reshape back before replacing x_ins. The failure occurs before replace_instruction is reached.
For the failing expanding chain, to_src_from_common() generates a multibroadcast from {1,1,1,1,6} to {1,1,1,2,3}, which is invalid. I kept the shared-reduction guard because shared consumers require restoring the original output. For a single-consumer reduction, that restoration is unnecessary.
There was a problem hiding this comment.
But this will still bail even when to_src_common does work.
There was a problem hiding this comment.
Copilot review overview
🟢 Approval recommended
The guard addresses the global replacement risk and has focused regression coverage.
Review effort: Balanced
Findings: None
What changed in this PR
Prevents simplify_reshapes from invalidating other consumers of shared reductions.
Changes:
- Gates element-expanding rewrites on single-consumer reductions.
- Adds regression coverage and a changelog entry.
| File | Description |
|---|---|
src/simplify_reshapes.cpp |
Adds the consumer-count guard. |
test/simplify_reshapes_test.cpp |
Tests a shared expanding reduction chain. |
CHANGELOG.md |
Documents the resolved issue. |
Review was single-pass without agent fan-out; runtime tests were not executed.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #5360 +/- ##
========================================
Coverage 92.71% 92.71%
========================================
Files 635 635
Lines 36236 36239 +3
========================================
+ Hits 33595 33598 +3
Misses 2641 2641
🚀 New features to boost your workflow:
|
| auto image = m1.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}}); | ||
| auto scale = m1.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}}); | ||
| auto mean = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image); | ||
| auto unsqueeze = m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), mean); |
There was a problem hiding this comment.
The unsqueeze should be remove here as well.
|
So I opened #5365 to address these issues as this requires changes to the |
|
@pfultz2
or you have in mind some better or simpler approach, now I see your PR and changes you made and they are fixing the problem I assume, I'll test with them to confirm it. |
|
This has been fixed by #PR5365 |
…eeze into multibroadcast (#5365) `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.
Motivation
Fix invalid
MULTIBROADCASTshapes when compiling Topaz APF, APO, AIOB, and AION models.The
simplify_reshapespass could move an element-expanding shape-transform chain through a shared reduction. Because the rewrite replaces the reduction globally, its other consumers received an incompatible shape.Technical Details
op_shape_transform_shared_reduce_expanding_chainto cover a sharedreduce_meanwith an expanding shape-transform chain.This regression became reachable after the matcher changes in #5233.
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot ApplicableFollow the LLVM AI Tool Use Policy for contributions using AI.