Repository navigation
[stinkytofu] Add WmmaVgprReorder analysis pass with explicit pool tagging - #8697
Merged
Merged
Conversation
…ging Add StinkyWmmaVgprReorderPass: an analysis-only pass that determines the optimal wmma reordering to reduce VGPR usage in double/N-buffered GEMM kernels by aliasing A-register pools across buffer stages. Key design decisions: - Pool membership is explicit via WmmaPoolData modifier (new TypedModifier in StinkyModifiers.hpp) set by TensileLite at generation time, not inferred from register patterns. This avoids fragile heuristics that break for non-contiguous or interleaved pool layouts. - A vs B operand detection is automatic: whichever src index has disjoint register groups across pools is A (pool-varying); the other is B (shared). - Three-layer ABI: IRegLivenessAnalysis, IWmmaReorderAlgorithm, and WmmaReorderAnalysisResult are each independently swappable. - Pass bails out if any wmma is missing WmmaPoolData — partial tagging would produce incorrect alias pairs silently. - Output is analysis only: desiredWmmaOrder + RegReplacement list for downstream reorder and rename passes. [stinkytofu] Fix stale wmmaSeq reference in AMajorOuterAlgorithm::solve [stinkytofu] Gate WmmaVgprReorderPass diagnostics behind PASS_DEBUG [stinkytofu] Rename aMajorOuter parameter half -> pool [stinkytofu] Rename aMajorOuter -> reorderToMinimizeALiveness [stinkytofu] Generalize reorder algorithm to detect and work on A or B operand [stinkytofu] Fix solve() inflation detection; add B-variant pool-varying tests [stinkytofu] Rename AMajorOuterAlgorithm -> PoolVaryingReorderAlgorithm [stinkytofu] Update file-level comment to describe generic pool-varying reorder [stinkytofu] Replace stale planning doc with ABI reference doc
KKyang
force-pushed
the
users/kkyang/stinkytofu-wmma-reorder
branch
from
June 23, 2026 04:35
44f273c to
ec1e4b9
Compare
…iases on constrained liveness
C is both src2 and dst in v_wmma (accumulator pattern): C_out = A*B + C_in. Within a pool each wmma has a unique C tile, and matrix accumulation is commutative, so wmma-to-wmma reordering through C is always mathematically safe. Two corrections: - Remove WAR check for C (non-wmma reading intermediate C must not prevent reordering; the final accumulated value is order-independent). - Add cGroup to the RAW check (non-wmma writing C creates a real data dep on wmma's src2 accumulator input).
Tests the case: wmma(A0,B0,C0) ds_load → C0 ← RAW: minRank[w00]=1 wmma(A1,B0,C1) wmma(A0,B1,C2) wmma(A1,B1,C3) EDF scheduler must slide w00 to pool-position 1 and put w01 first. A0 group still ends up contiguous at positions 0-1 (intervals [0,1] vs [2,3]), so the saving is preserved despite the dependency.
A non-wmma instruction reading any C tile in the pool is observing an intermediate accumulation state. We cannot know what it does with that value, so nothing should cross it in either direction — same treatment as s_waitcnt or other hard barriers. Implements minRank/maxRank clamping for all pool wmma relative to the read point, mirroring the existing isBarrier() path.
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project status has failed because the head coverage (77.89%) 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 #8697 +/- ##
========================================
Coverage 71.43% 71.43%
========================================
Files 2612 2612
Lines 406986 406986
Branches 60770 60770
========================================
Hits 290724 290724
Misses 94958 94958
Partials 21304 21304
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
hcman2
approved these changes
Jun 24, 2026
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
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
StinkyWmmaVgprReorderPass, a read-only analysis pass that detects VGPR savings in N-buffered GEMM kernels by simulating A-major outer wmma orderingWmmaPoolDatamodifier (TypedModifier<WmmaPoolData>) so TensileLite can explicitly tag each wmma instruction with its buffer pool index at generation time — replacing the fragile register-heuristic pool boundary detectionIRegLivenessAnalysis/IWmmaReorderAlgorithm/WmmaReorderAnalysisResult, all independently swappableDesign
The core insight: TensileLite owns the full generation pipeline (same advantage as Triton), so pool structure can be passed down as explicit metadata rather than recovered from register patterns. Any heuristic approach (
findPoolBoundaries) breaks for interleaved or non-contiguous pool layouts.groupWmmaByPoolreadsWmmaPoolDatatags directly; bails out entirely if any wmma is untagged (partial tagging = misconfigured pipeline).detectABIndicesthen auto-identifies the A (pool-varying) vs B (pool-shared) src operand by checking cross-pool register intersection — no hardcoded operand order.AMajorOuterAlgorithm::solvechecks liveness to confirm B-major outer is in effect (interval width > nB), then produces the A-major permutation and emits alias candidates (pools 1..N-1 aliased onto pool 0).Test plan
BMajorOuter_DesiredOrderIsAMajor— B-major input → reorder fires, aliases emittedAMajorOuter_NotApplicable— A-major input → liveness detects no saving, pass skipsUntaggedWmma_PassSkipped— any untagged wmma →applicable=falseTripleBuffer_TwoAliasPairs— 3 pools → 2× alias pairsWmmaIntervalLivenessunit testsAMajorOuterAlgorithmunit tests via injected poolsOpen item
TensileLite-side coordination needed:
ToStinkyTofuUtils.cppattachesWmmaPoolDataduring rocisa → StinkyTofu conversion, but the rocisaModulenaming convention for wmma pools is TBD.