feat(stinkytofu): wmma reorder pass with cross-iteration ds_load prefetch migration - #11478
Draft
KKyang wants to merge 5 commits into
Draft
feat(stinkytofu): wmma reorder pass with cross-iteration ds_load prefetch migration#11478KKyang wants to merge 5 commits into
KKyang wants to merge 5 commits into
Conversation
…tion The analysis lived only inside StinkyWmmaVgprReorderPass and published its results through a file-static map, so any consumer had to depend on that pass having run earlier in the same pipeline. Lift the body into analyzeWmmaVgprReorder(bb[, liveness, algorithm]). The pass becomes a thin wrapper over it, keeping getWmmaReorderResult() and its unit tests unchanged, while callers that only need one block's order can compute it themselves. No behaviour change.
Rewrites an unrolled-loop body to a new wmma issue order and re-places the ds_loads that feed it, including across the loop back-edge. The order comes from a swappable mode (IWmmaOrderProvider) so each reordering method is one implementation over shared rewrite machinery: VgprAnalysisOrder (minimize VGPR pressure), ExplicitOrder (a permutation, for tuning sweeps) and ReverseOrder (exercises the machinery without a tuner). Two invariants keep the rewrite honest. The loop body is only permuted, so the kernel still issues the same work per iteration. Everything added or removed happens in the preheader, where the software pipeline's iteration-0 prefetches live: a ds_load whose new slot falls before the top of the body gains a preheader clone and has its body copy retuned to the next iteration, and one that no longer needs the head start loses its preheader copy and is de-rotated. The per-iteration LDS stride is never invented. It is read off the preheader/body ds_load pairs the kernel already has; with no such pair, cross-iteration migration is disabled and the body is permuted in place. A wrapped load must issue after every wmma that reads its registers this iteration, not just the first, or it clobbers them. The requested prefetch distance yields to that rule, and a loop where an already-prefetched load has no legal slot is left untouched rather than silently corrupted. Also adds mod.wmma_pool serialization, without which WmmaPoolData had no .stir round-trip and the vgpr-analysis mode could not be tested, and a key=value argument helper for stinkytofu-opt.
…mmaReorder Adds the ModuleOptions flag, default false, and runs the pass in addGfx1250RegionPasses immediately before StinkyBuildImplicitDependencyPass. That slot is required, not incidental. WMMA is the scheduling anchor: scheduleInDAG builds wmmaIndex by walking the block in program order and the dsReadPriority pre-scan gives every ds_read the index of its earliest wmma consumer, so moving a wmma after the DAG has run would leave those priorities describing an order that no longer exists. Running here also puts the pass after StinkyRemoveWaitCntPass, so no leftover s_waitcnt chops the window the reorder is allowed to move in. The flag stays off, so no kernel changes. With it on, today's kernels are still unaffected: the default mode declines any block whose wmma carry no WmmaPoolData, which TensileLite does not yet emit.
github-actions
Bot
added
project: hipblaslt
project: hipsparselt
shared: stinkytofu
labels
❌ PR Check — Action Required
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🚫 Please fix the failed policies before requesting reviews. The following policy checks failed:
The |
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project check has failed because the head coverage (38.33%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #11478 +/- ##
========================================
Coverage 69.33% 69.33%
========================================
Files 2810 2810
Lines 465293 465293
Branches 68653 68653
========================================
Hits 322565 322565
Misses 119189 119189
Partials 23539 23539
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
getTerminator() returns null once every existing prefetch has been de-rotated out of a preheader, and dyn_cast doesn't tolerate a null argument -- classof dereferences it directly. Reordering a loop whose wmma pattern fully swaps which loads are prefetched (e.g. reverse on a kernel where the preheader holds only prefetches) hits this and segfaults intermittently, since which garbage classof reads back depends on heap layout. Found while stress-testing StinkyWmmaReorderPass under AddressSanitizer.
A wmma order is free to place any two ds_loads in any relative order, but it can't invent extra physical registers: when a segment double-buffers one RegGroup across more than one producer (ordinary software pipelining, not just the back-edge case), the new order can make two producers' live windows overlap on the same physical register. Detect this generally -- not just the existing wrapped-load-vs-its-own- twin check -- by tracking each ds_load's exact consumer set and original body position, scoping consumer attribution per producer so a shared register's reuses aren't all attributed to every producer indiscriminately. When two producers on one RegGroup collide, the later one's def, matched consumers, and preheader twin are retargeted together to a placeholder register (StinkyRegister::Virtual) instead of silently clobbering the earlier value or bailing the whole loop. Final physical assignment is left to register allocation.
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
1 participant
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.
Summary
A new mechanism for changing the wmma issue pattern of an unrolled GEMM loop, and repairing the LDS software pipeline that the new pattern invalidates.
The same wmma are issued either way — only the order changes — but that reorder moves which ds_loads need a head start. Some loads must now be issued an iteration early and some no longer need to be, so prefetches appear in and disappear from the loop preheader.
Default off. No kernel changes without
EnableWmmaReorder.Design
The order comes from a swappable mode.
IWmmaOrderProvideris the extension point; each reordering method is one implementation over shared rewrite machinery. Three ship:VgprAnalysisOrderProvider(default)analyzeWmmaVgprReorder— minimize VGPR pressureExplicitOrderProviderReverseOrderProviderTwo invariants keep the rewrite honest.
The per-iteration LDS stride is never invented. It is read off the preheader/body ds_load pairs the kernel already has (same dest VGPR group, same address operands, offset difference). With no such pair the pass cannot know how the kernel rotates buffers, so cross-iteration migration is disabled and the body is permuted in place.
One hard safety rule. A wrapped load issues for the next iteration, so it must sit after every wmma that reads its registers this iteration — not just the first. The requested prefetch distance yields to that. A loop where an already-prefetched load has no legal slot is left untouched rather than silently corrupted.
Worked example
Four wmma over A0..A3, arriving pipelined for order
(A0,A1,A2,A3)at distance 2, so A0/A1 are prefetched. Reversing to(A3,A2,A1,A0):prefetch +4/-4, with the 512-byte stride discovered from the input. This istests/filecheck/wmma_reorder_cross_iter.stir.Why this slot in the pipeline
WMMA is the scheduling anchor, so the pattern has to be final before the DAG runs.
scheduleInDAGbuildswmmaIndexby walking the block in program order, and thedsReadPrioritypre-scan gives every ds_read the index of its earliest wmma consumer. Moving a wmma after the DAG has run would leave those priorities describing an order that no longer exists.Running immediately before
StinkyBuildImplicitDependencyPassalso puts the pass afterStinkyRemoveWaitCntPass, so no leftovers_waitcntchops the window the reorder is allowed to move in.The in-body ds_load placement this pass produces is therefore a seed order — the DAG re-places local reads against its own model. What survives to the kernel is the wmma order plus the preheader add/remove, which the DAG cannot do at all since it never crosses the back-edge. The seed is safe under the DAG because
buildRegisterDependencyDAGemits WAR edges, so a wrapped load stays after the wmma reading its registers.Commits
refactor— lift the wmma vgpr reorder analysis out of its pass intoanalyzeWmmaVgprReorder(), so a mode can compute its own input instead of depending on a pipeline neighbour and a file-static map. No behaviour change.feat— the pass, its modes and its tests. Also addsmod.wmma_poolserialization (WmmaPoolDatahad no.stirround-trip, so the vgpr-analysis mode was untestable) and akey=valueargument helper forstinkytofu-opt.feat— theEnableWmmaReorderflag and the gfx1250 wiring.Testing
ctest1138/1138 on gfx1250. New coverage:tests/filecheck/wmma_reorder_cross_iter.stir— the worked example above, checking both the preheader swap and the rewritten body.tests/filecheck/wmma_reorder_no_pipeline.stir— a loop with no existing prefetch: no stride is available, so nothing crosses the back-edge and the body is permuted in place.tests/unit/asm/StinkyWmmaReorderPassTest.cpp— 7 tests over the mode ABI, cross-iteration migration, the body-is-a-permutation invariant, and the refusals.Not in this PR
prefetchDistanceis derived from ds_load latency over wmma issue rate. Worth flagging: it is the value that decides which loads wrap, so it is the one worth sweeping. One moreX(...)row when wanted.WmmaPoolData, which TensileLite does not yet emit. The wiring is ready for when it does, or for a sweep drivingExplicitOrderProvider.🤖 Generated with Claude Code