Hoist and horizontal dot#5014
Open
TedThemistokleous wants to merge 31 commits into
Open
Conversation
Hoist a pointwise (or SiLU diamond) chain that is replicated across several sibling slices of the same instruction above the slices: the chain is computed once on the bounding slice and each consumer reads its sub-range back out. This exposes a single wide pointwise op that downstream passes (find_splits, horizontal fusion) can recombine instead of N identical narrow ops. trace_pw_chain handles both linear pointwise chains and the SiLU diamond (x -> sigmoid -> mul(x, sigmoid)). Slices are only hoisted when they exactly tile the bounding range so no redundant elements are computed.
Cover the hoist transform in isolation (simplify_reshapes + DCE only): - relu hoisted above two sibling unit slices - SiLU diamond (sigmoid/mul) hoisted above sibling slices - no hoist when slices leave a gap in the bounding range - no hoist when sibling slices feed different pointwise chains
Batch structurally-identical dot operations (same activation/weight lens and output type, constant weights) into a single batched GEMM by stacking activations and weights along a new leading axis, then slice and squeeze the per-dot results back out. Per review on #4723, only the dots are fused here; downstream elementwise work (bias add, SiLU, ...) is left in place so the existing fusions (find_splits in simplify_algebra and the pointwise/MLIR fusions) can recombine it on top of the batched dot rather than duplicating that logic.
Cover the batched-dot fusion (fuse_horizontal + DCE): - two shape-identical dots with constant weights batch into one GEMM - dots with non-constant weights are not fused - dots with mismatched activation/weight shapes do not group
Exercise simplify_reshapes (hoist_pointwise_above_slices) and fuse_horizontal (dot_horizontal_fusion) in a single pipeline on a module containing both opportunities: a SiLU replicated across sibling slices and a set of parallel constant-weight dots. Verify the dots collapse to a single batched GEMM and the per-slice SiLUs collapse to one sigmoid/mul.
TedThemistokleous
added
the
high priority
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #5014 +/- ##
===========================================
+ Coverage 92.89% 92.92% +0.04%
===========================================
Files 603 603
Lines 32448 32672 +224
===========================================
+ Hits 30140 30360 +220
- Misses 2308 2312 +4
🚀 New features to boost your workflow:
|
pfultz2
requested changes
Jun 26, 2026
| } | ||
|
|
||
| EXPECT(m1.sort() == m2.sort()); | ||
| } |
Collaborator
There was a problem hiding this comment.
This should already be handled by find_splits.
Collaborator
Author
There was a problem hiding this comment.
Understood. Modified the changes so we can reuse as much of find_splits as possible, got rid of the netire tracing section so its just a case where we handle the smaller cases. Diamon pattern of ops to lift before slice should be handled like we talked about after fuse_pointwise.
Let me know what you think I'm just running some additional tests on this.
pfultz2
reviewed
Jun 26, 2026
Check flagged results 🔆 |
|
Replace the bespoke pointwise-chain hoisting engine (chain tracing, equivalence, replay) in find_splits with a unary-only hoist_bounded_splits that reuses the existing get_split_groups/is_fusable/split_groups_are_dependent helpers. When sibling slices only partially tile their tensor, a single-input op replicated across them is computed once on the bounding slice and each consumer re-slices its sub-range. Multi-op linear chains are peeled one op per fixpoint iteration, and SiLU/diamond patterns are handled once fuse_pointwise collapses them into a single pointwise op, so no diamond-specific code is needed here. Co-authored-by: Cursor <cursoragent@cursor.com>
The slim find_splits hoist only widens single-input ops, so a raw sigmoid/mul SiLU diamond is no longer hoisted directly. Update the two SiLU tests to run fuse_pointwise first (collapsing each per-slice SiLU into one pointwise op) and assert on the resulting fused pointwise: hoist_silu_above_slices now compares against a hoisted pointwise program, and the end-to-end test asserts a single pointwise plus a single batched dot. Co-authored-by: Cursor <cursoragent@cursor.com>
ensures we're not grabbing everything ith bias which was generated errnornous add_kernels by fusing dot early.
pfultz2
requested changes
Jun 30, 2026
- move tests - Adjust changes to find_splits with partial flag.
Collaborator
|
Can you revert the refactor into helper functions and just use a |
Collaborator
Author
Sure can do that |
CharlieL7
reviewed
Jul 10, 2026
CharlieL7
requested changes
Jul 10, 2026
CharlieL7
left a comment
CharlieL7
left a comment
Collaborator
There was a problem hiding this comment.
Looks good overall, recommending some test changes
A horizontally-fused wide dot gets sliced per group, and each slice feeds an optional bias add then a SiLU gate (mul(x, sigmoid(x))). find_splits cannot hoist this because its binary path only folds a constant second operand, while the gate's second operand is derived from the split, so the epilogue is stranded after the slice and cannot fuse into the GEMM. find_split_gate rewrites the (optional bias) + gate onto the un-sliced tensor (concat the per-slice biases, add, then g/mul on the whole result, slicing last) so MLIR can fuse dot+add+g+mul into the wide GEMM. Co-authored-by: Cursor <cursoragent@cursor.com>
The K>1535 cutoff in is_mlir_dot routed the wide horizontally-fused GEMMs (K=2112 and K=5536) to rocBLAS (gpu::hip_gemm), which has no epilogue fusion, so their bias+SiLU fell out as standalone kernels. Raise the threshold to 8096 so these GEMMs stay on the MLIR path and can fuse their epilogue. Co-authored-by: Cursor <cursoragent@cursor.com>
…tial groups that are skipped over within the find_splits matcher.
TedThemistokleous
force-pushed
the
hoist_and_horizontal_dot
branch
from
July 17, 2026 20:08
c4863b8 to
e10b146
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| // Fuse across distinct tables first, then same-table groups. Running same-table | ||
| // fusion first can shrink cross-table groups below their size threshold and miss it. | ||
| fuse_horizontal_ops(m, gather_horizontal_fusion{}, same_table_gather_horizontal_fusion{}); | ||
| fuse_horizontal_ops(m, gather_horizontal_fusion{}, same_table_gather_horizontal_fusion{}, dot_horizontal_fusion{}); |
Contributor
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
Suggested change
| fuse_horizontal_ops(m, gather_horizontal_fusion{}, same_table_gather_horizontal_fusion{}, dot_horizontal_fusion{}); | |
| fuse_horizontal_ops(m, | |
| gather_horizontal_fusion{}, | |
| same_table_gather_horizontal_fusion{}, | |
| dot_horizontal_fusion{}); |
Comment on lines
+1715
to
+1716
| auto lo = any_cast<op::slice>(splits[i]->get_operator()).starts.front(); | ||
| auto hi = any_cast<op::slice>(splits[j - 1]->get_operator()).ends.front(); |
Contributor
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
Suggested change
| auto lo = any_cast<op::slice>(splits[i]->get_operator()).starts.front(); | |
| auto hi = any_cast<op::slice>(splits[j - 1]->get_operator()).ends.front(); | |
| auto lo = any_cast<op::slice>(splits[i]->get_operator()).starts.front(); | |
| auto hi = any_cast<op::slice>(splits[j - 1]->get_operator()).ends.front(); |
Comment on lines
+1728
to
+1729
| c = m.insert_instruction( | ||
| std::next(base), info.consumer->get_operator(), {base}, info.consumer->module_inputs()); |
Contributor
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
Suggested change
| c = m.insert_instruction( | |
| std::next(base), info.consumer->get_operator(), {base}, info.consumer->module_inputs()); | |
| c = m.insert_instruction(std::next(base), | |
| info.consumer->get_operator(), | |
| {base}, | |
| info.consumer->module_inputs()); |
Comment on lines
+1738
to
+1740
| info.consumer->get_operator(), | ||
| args, | ||
| info.consumer->module_inputs()); |
Contributor
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
Suggested change
| info.consumer->get_operator(), | |
| args, | |
| info.consumer->module_inputs()); | |
| info.consumer->get_operator(), | |
| args, | |
| info.consumer->module_inputs()); |
Comment on lines
+2387
to
+2388
| auto r2 = m1.add_instruction(migraphx::make_op("relu"), s2); | ||
| auto s3 = m1.add_instruction( |
Contributor
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
Suggested change
| auto r2 = m1.add_instruction(migraphx::make_op("relu"), s2); | |
| auto s3 = m1.add_instruction( | |
| auto r2 = m1.add_instruction(migraphx::make_op("relu"), s2); | |
| auto s3 = m1.add_instruction( |
pfultz2
reviewed
Jul 17, 2026
pfultz2
reviewed
Jul 17, 2026
| // mixed full-cover group (e.g. SiLU heads next to a reshape head) still | ||
| // fuses its heads. (The `partial` path already covers partial tiling.) | ||
| if(not partial) | ||
| hoist_partial_groups(m, ins, splits); |
Collaborator
There was a problem hiding this comment.
What is this doing? It looks like it is reimplementing find_splits again. The comment doesnt explain it good. From the comments, it sounds like it is trying to handle the case where we have slice -> op, slice -> reshape -> op but that seems like we should have another matcher to rewrite the slice -> op, slice -> reshape -> op to slice -> op, slice -> op -> reshape, which I thought you already had a matcher to handle that so I am not sure what the issue is.
Collaborator
Author
There was a problem hiding this comment.
We were missing some branches. Let me clean this up today - Its to catch partial splits which were originally getting missed causing things to never get matched/changed then fused correctly
Collaborator
There was a problem hiding this comment.
From the tests, is this trying to handle the case where a binary takes both inputs as splits from different tensors?
I do know the current code doesnt handle that as it requires the second argument to be constant, but we could add a check that the second argument can be constant or another split.
…PS=dot,fused_dot instead
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Final piece split from #4723
Adds in a hoist pass into simplify reshapes to move point wise operations above slices that is repeated across several slices. This sets us up to perofrm a horizontal dot fusion across each of the branches.
Technical Details
Updating this after refining things further. Analysis after doing a compare with develop and some changes
Summary
Fuse the MLP prediction-tower bias + SiLU epilogue into the horizontally-fused GEMM path on the feed-gen model, eliminating per-tower loose pointwise kernels. Two root causes were fixed:
find_splitsonly hoisted a pointwise/reduction above sibling slices when every slice fed the same op. A wide horizontally-fused dot whose output is sliced per tower mixes SiLU heads with a differently-consumed head (e.g. a reshape), so the group was rejected and each head'sadd + SiLUfell out as a separate kernel. Added a contiguous-subset partial hoist tofind_splitsthat lifts the op of a maximal contiguous matching run over its bounding slice (supports unary ops and binary ops whose other operand is an identically-ranged slice of another wide tensor, i.e. a broadcast bias).is_mlir_dotskipped GEMMs withK > 1535, sending the wide concatenated-K tower GEMMs togpu::hip_gemm(rocBLAS), which cannot take a fused epilogue. Raised the threshold to8096so they stay on MLIR and fuse.Performance (feed-gen model, batch 1,
driver perf -n 10000)The rate/latency gain with flat total-instruction-time reflects fewer kernel launches (less dispatch overhead), not less compute.
Resultant kernels
add_sigmoid_mul_kernel(loose SiLU)mlir_dot_add_sigmoid_mulgpu::hip_gemm(rocBLAS, no epilogue fusion)mlir_dot_add_unsqueeze→mlir_slice_dot_add_unsqueezeTests
test/simplify_algebra_test.cpp:simplify_split_partial_pointwise_subset(unary) andsimplify_split_partial_binary_subset(binary bias+SiLU shape) — cover the mixed-group partial hoist.test/optimize_module_test.cpp: end-to-end MLP-tower SiLU pipeline guard.Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot Applicable