Add rewrite broadcasts for fuse_reduce - #5146
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a reusable “rewrite broadcasts onto pointwise inputs” transformation and wires it into the fuse_reduce pipeline (and associated tests) to enable additional fusion opportunities when broadcasts sit between pointwise ops and reductions.
Changes:
- Added a shared
rewrite_broadcasts(module_pass_manager&, op_name)utility (moved out offuse_pointwise) and integrated it into bothfuse_pointwiseandfuse_reduce. - Extended
fuse_reduceto optionally rewrite/hoist broadcasts to unlock fusions, and enabled this path infuse_pointwise_reduce. - Added new unit tests covering fusion cases enabled by broadcast rewriting/hoisting.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
src/rewrite_broadcasts.cpp |
New shared implementation of broadcast-rewrite matcher and cleanup passes. |
src/include/migraphx/rewrite_broadcasts.hpp |
Public declaration for the new rewrite helper. |
src/fuse_pointwise.cpp |
Replaces the local broadcast rewrite logic with the shared helper. |
src/include/migraphx/fuse_reduce.hpp |
Adds enable_rewrite_broadcasts flag to control the new behavior. |
src/fuse_reduce.cpp |
Adds optional broadcast rewriting/hoisting logic for reduce fusion, plus broader broadcast matching. |
src/fuse_pointwise_reduce.cpp |
Runs an additional fuse_reduce pass with broadcast rewriting enabled. |
test/fuse_pointwise.cpp |
Adds a test validating pointwise fusion enabled by broadcast rewrite. |
test/fuse_reduce.cpp |
Adds multiple tests validating reduce fusion enabled by broadcast rewrite/hoist. |
src/CMakeLists.txt |
Adds the new rewrite_broadcasts.cpp to the library build. |
Suppressed comments (1)
src/fuse_reduce.cpp:520
- In
reduce_reshape, convertingbroadcasttomultibroadcastcan change semantics for the 2-inputbroadcastform (dynamic broadcasting), whereout_lensis unused and may be unset.multibroadcastwith 2+ inputs ignores theout_lensattribute entirely, so this rewrite can miscompile graphs that contain 2-input broadcasts inside the fused submodule.
if(contains({"multibroadcast", "broadcast"}, sop.name()))
return make_op("multibroadcast", {{"out_lens", dims}});
assert(sop.name() == "pointwise");
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if(contains({"multibroadcast", "broadcast"}, op.name())) | ||
| { | ||
| auto out_lens = | ||
| expand_dims(op.to_value().at("out_lens").to_vector<std::size_t>(), | ||
| unreduced_axes, | ||
| blens); | ||
| return rmm.insert_instruction( | ||
| pos, make_op("multibroadcast", {{"out_lens", out_lens}}), inputs); | ||
| } |
There was a problem hiding this comment.
n the reduce-module rebuild callback, treating broadcast and multibroadcast uniformly and unconditionally reading op.to_value().at("out_lens") is unsafe for 2-input broadcast instructions (used for dynamic broadcasting), which typically don't set out_lens at all (e.g. src/onnx/onnx_parser.cpp:155). This can lead to an empty lens vector and out-of-range access in expand_dims, and converting a 2-input broadcast into a 2-input multibroadcast changes semantics (multibroadcast ignores out_lens when it has 2+ inputs).
I think this is valid, multi-input broadcasts need to be handled differently or just dont apply this matcher for that case for now
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #5146 +/- ##
===========================================
- Coverage 93.13% 93.12% -0.01%
===========================================
Files 625 626 +1
Lines 33252 33309 +57
===========================================
+ Hits 30967 31017 +50
- Misses 2285 2292 +7
🚀 New features to boost your workflow:
|
bdevorem
left a comment
bdevorem
left a comment
There was a problem hiding this comment.
lgtm. Copilot's comment about 2-input broadcast seems meaningful though
Motivation
This PR introduces a reusable “rewrite broadcasts onto pointwise inputs” transformation and wires it into the
fuse_reducepipeline (and associated tests) to enable additional fusion opportunities when broadcasts sit between pointwise ops and reductions.Technical Details
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot ApplicableFollow the LLVM AI Tool Use Policy for contributions using AI.