Repository navigation
feat(tensilelite): reorder wmma - #12767
chengchingwen wants to merge 1 commit into
Conversation
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
84eb697 to
728668b
Compare
|
🎉 All checks passed! This PR is ready for review. |
c255630 to
32d0f98
Compare
KKyang
left a comment
KKyang
left a comment
There was a problem hiding this comment.
Thanks for this. It's a big change, so here is what I'd want before it goes further. I read the changes to existing StinkyTofu and TensileLite code closely and only skimmed the new LoopModel/, Lowering/ and GIR pass bodies.
1. Unit tests
- The only tests added are
GirFrameAnalysisTest.cpp, one filecheck (gir_frame_waitcnt.stir) and the sweep yamlulm_gfx1250.yaml. The roughly 15k lines of new Python (Tensile/LoopModel,Tensile/Lowering,Components/LoopModel) have no unit tests, and neither doGirFencePlacementPassorGirFrameCounterFlow. - The SCC rule rewrite has no test either (see the inline comment).
- Please also show that existing kernels are unchanged with
UseLoopModel=False. The wave-separated TDM descriptor code was replaced for all kernels (initTDMDescriptorWaveSeparatedand friends became the...Grouphelpers driven by the newtdmFusedGroupsandtdmGroupWaveRanges). By my reading the predicate and the default grouping match the old behaviour, but a before/after assembly diff on the existing gfx1250 TDM yamls (tdm_gfx1250,decouple_pgr_tdm_fuse_gfx1250,tdmpluslds_tdm_gfx1250, the streamk TDM ones) would settle it.
2. Description
The "Test Result" section is empty, and the Technical Details are incomplete. Please add:
- The new StinkyTofu pieces:
OrderTokenDataandNoWaitCntData, the module optionsEnableGirFramePipeline,EnableGirFrameWaitCntInsertion,DisableWaitCntRemovalandDisableTensorcntInsertion, and the dag featuresuseMemoryTokenOrderingandpreserveTensorLoadOrder. - The change to
StinkyBuildImplicitDependencyPassandRegionDAG(order-token walls), theWaitDataflowgating, andPB.applyExtensionPoint(EndOfPipeline, ...)in the backend. - The SCC rule rewrite and the TDM group refactor of existing code.
- The rocisa
instruction.hpp/.cppchanges. - What was run and the results. Also, the StinkyTofu paragraph starts with a stray "—" and "Two new folder" needs fixing.
3. Please split this up if you can
At 137 files and +25k lines in one commit it's hard to review, and a regression on the default path wouldn't be attributable. The StinkyTofu behaviour changes in particular should land separately and first:
- SCC rule rewrite (
StinkyDAGSchedulerPass.cpp). - Order-token ordering (
OrderTokenData,RegionDAG,StinkyBuildImplicitDependencyPass) andpreserveTensorLoadOrder. - The
WaitDataflow/WaitAwareScheduleRepairPass/ backend wiring.
After that, the TDM grouping refactor in TensileLite, and then the loop model itself (LoopModel/, Lowering/, GIR passes, docs).
4. Aborts
See the inline comments. report_fatal_error (and STINKY_UNREACHABLE) end in std::abort(), so input-dependent failures should return errors that Tensile can turn into a solution rejection.
What I checked and found fine: the new StinkyTofu code is inert for existing kernels. OrderTokenData is new and only the adaptor sets it, GirFrameAnalysis returns early when no contract is loaded, and the legacy TDMSplit emitter that was removed was already rejected on the parent (Solution.py). I'd just like that to be demonstrated by tests and a before/after diff rather than my reading.
Generated by Claude Code
| // because its reader is waiting on that barrier. So pin to a fixpoint, then lock. | ||
| std::vector<std::vector<char>> reach; | ||
| std::vector<std::vector<char>> pinned(live.size(), std::vector<char>(barriers.size(), 0)); | ||
| for (bool changed = true; changed;) { |
There was a problem hiding this comment.
This rewrites the existing cluster-barrier SCC rule (applyClusterBarrierSccRule is split into pinChainAfterBarriers plus a pin-to-fixpoint loop, and lock decisions now use the post-pin reach). It isn't related to the loop model, and the PR description doesn't mention it.
I traced the single-chain case and it reduces to the old logic. What changes is the multi-chain interplay, which is exactly the situation the new comment describes (a lock and a barrier blocking each other). It changes scheduling for existing kernels whenever ClusterBarrier is on, not only under UseLoopModel.
Two asks:
- Please move this to its own PR so a scheduling change on the default path is attributable to one commit.
- Please add a unit test. The PR changes no scheduler test, so nothing exercises the fixpoint or the lock/barrier deadlock it fixes. A
DAGSchedulerPassTestcase with two chains where pinning one makes the other's reader depend on a barrier would do, and the existingClusterBarrierSccRule_*tests should be run against it.
Generated by Claude Code
|
OrderTokenData and NoWaitCntData related will be removed. |
|
Summarizing the concerns raised across the inline comments, since they add up to one overall point about this PR's design:
Net: 2-6 are the same pattern repeated -- new machinery standing in for something that already exists in this codebase, rather than extending it. I don't think this should merge as one 25k-line commit while that's unresolved; see the earlier review for the split-PR ask. Whether a GIR layer is the right direction at all for StinkyTofu (1) is worth resolving with @nakajee first, since that answer changes what 2-5 should even look like. Generated by Claude Code |
The FencePlacementPass add workgroup barrier, not cluster barrier. The pass inserts barrier after instruction scheduling.
GirWaitCntInsertionPass use exist StinkyWaitCntInsertionPass for dscnt/vgpr use-def, only the tensorcnt and tdm WAR dscnt is inserted by this gir pass based frame graph/hazard
The loop axis reordering would cooperate with the PLR/CLR/TDMSplitAB/TDMFuse behavior that is not expressible by what Subtile path currently has, and enable the full ST scheduling. |
|
|
07303b0 to
aeeedfb
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #12767 +/- ##
===========================================
- Coverage 77.77% 77.72% -0.05%
===========================================
Files 2905 2974 +69
Lines 491723 503457 +11734
Branches 74578 76933 +2355
===========================================
+ Hits 382425 391292 +8867
- Misses 84295 86525 +2230
- Partials 25003 25640 +637
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
2feee35 to
30fa5a0
Compare
3aab103 to
726cfd4
Compare
777670a to
b5a4383
Compare
| /// | ||
| /// Both halves are pinned. `owner.anchor` names the fence by its signal, and the token edges all | ||
| /// land on the wait, so pinning the anchor alone gives the halves different windows to drift in. | ||
| static void pinGirFenceWindows(RegionDAG& dag, const GirFrameHazardAnalysis::Result& hazards, |
| passFeatureConfig.dagFeatures.girFramePipeline = | ||
| moduleOptions.EnableGirFramePipeline; | ||
| passFeatureConfig.dagFeatures.clusterBarrierTensorLoadDelayCycles = | ||
| moduleOptions.EnableGirFramePipeline && moduleOptions.ClusterBarrier ? 200 : 0; |
There was a problem hiding this comment.
Why use this magic number to seperate the tensor_load and barrier? Could you find a rule when the tensor load can be issued?
| return true; | ||
| }; | ||
|
|
||
| // Merge a nearby later fence into the first fence when at most two configured WMMA queues |
There was a problem hiding this comment.
We may discuss how to merge the token for all kinds of kernels. This looks only for the large MT type. cc @aazz44ss
429db9a to
ce07952
Compare
fe3bfbf to
0d27126
Compare
0d27126 to
fa1fc22
Compare
Motivation
JIRA ID: AIHPBLAS-4229
WMMA GEMM inner loop is scheduled by hand-written logic, which fixes one traversal of the K/M/N axes and one placement for every local read, LDS copy and wait. This adds an opt-in path
UseLoopModelthat allows to change the order of loop axes.UseLoopModel=Falseremains the default and the normal path.Technical Details
New opt-in parameter.
UseLoopModel(defaultFalse) routes the inner-loop body throughthe new path.
New kernel parameters, all defaulting to the existing behaviour:
WmmaInnerOrder1..6WmmaOuterOrder-1..3-1derives it from the live splitsTDMSplitA/TDMSplitB0..2Two new folder
Tensile/LoopModelholds the Solution->LoopIR plus its checks.Tensile/Loweringlowers LoopIR to GIR and renders it.StinkyTofu barrier and wait placement. —
GirFrameContractImportPass,GirFencePlacementPass,GirWaitCntInsertionPassand theGirFrameCounterFlowanalysis — which read a frame contract emitted beside the GIR and decide fences ands_wait_*counts from the program itself.Test Plan
Tensile/Tests/common/gemm/gfx12/ulm_gfx1250.yaml
Test Result
Submission Checklist