[tensilelite] Fix UseSubtileImpl NonEdge store correctness and relax edge dispatch - #7064
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project status has failed because the head coverage (77.83%) 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 #7064 +/- ##
========================================
Coverage 65.26% 65.26%
========================================
Files 2083 2083
Lines 323506 323506
Branches 42438 42438
========================================
Hits 211111 211111
Misses 94821 94821
Partials 17574 17574
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
… guard, and alignSize - Fix vmcnt hazard in interleaved B1 NonEdge store path: when M-guard branches skip stores at runtime, the actual vmcnt counter has fewer outstanding operations than codegen assumed, making s_waitcnt values too permissive for C-loads. Set vscnt=0 for UseSubtileImpl GLS=0. - Fix SrdD incToNextRow skipped when MGuard branch skips the second paired store in an N-group. Defer the increment and emit it after the N-group end label. - Flush pending N-group end label before beta*C fmacs for the next N-group, so M-guard branches don't skip the fmacs. - Add scalar fallback store when only the lower M-block in a pair is valid (MGuard check: paired store needs both blocks valid). - Reduce checkIsEdgeSubtile M alignment from mBlockSize (32) to MatrixInstM (16) so more workgroups take the optimized NonEdge path. The paired store path already handles partial pairs via MGuard + scalar fallback. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
b-shi
force-pushed
the
users/brianshi/storeD_fix
branch
from
May 6, 2026 13:46
befa581 to
e028463
Compare
Collaborator
|
I am seeing significant speedup for some cases. If all tests pass, we want to merge it. |
msujon-AMD
approved these changes
May 6, 2026
aledudek
pushed a commit
that referenced
this pull request
May 20, 2026
…edge dispatch (#7064) ## Motivation Fix correctness failures in the `UseSubtileImpl` NonEdge store path for gfx950 BF16 and MXFP4 subtile kernels. These failures were caused by several interrelated bugs in the interleaved (GLS=0) store codegen that manifested when M-guard branches skipped stores at runtime, and by an overly strict edge/NonEdge dispatch check that forced subtile-aligned workgroups into the unoptimized edge path. ## Technical Details **1. vmcnt hazard in interleaved B1 NonEdge store path** (`GlobalWriteBatch.py`) On gfx950 (`SeparateVscnt=False`), loads and stores share the same vmcnt counter. The interleaved store path computed `vmcnt = vlcnt + vscnt` where `vscnt = self.storesIssued`. When M-guard branches skip stores at runtime, the actual vmcnt counter has fewer outstanding operations than codegen assumed, making `s_waitcnt vmcnt(N)` too permissive — the hardware doesn't wait long enough for C-loads to complete before fmacs consume them. Fixed by setting `vscnt = 0` for `UseSubtileImpl` with `GroupLoadStore=False`. This is more conservative but correct regardless of how many stores were skipped at runtime. **2. SrdD increment skipped by M-guard branch** (`GlobalWriteBatch.py`) The SrdD `incToNextRow` was emitted inside the M-guard-skippable region of the store loop. When the M-guard branch skipped the last store in an N-group, the SrdD increment was also skipped, causing subsequent N-group stores to write to the wrong row address. Fixed by deferring the increment and emitting it after the N-group end label. **3. N-group end label placed after next N-group's fmacs** (`GlobalWriteBatch.py`) When `GroupLoadStore=False`, fmacs and stores are interleaved in the same module. The fmacs for N-group K+1 were emitted before the N-group K end label was placed. The M-guard branch (targeting the end label) would skip both the last store of N-group K *and* the fmacs for N-group K+1, leaving the K+1 accumulators at zero. Fixed by flushing the pending N-group end label and deferred SrdD increment before emitting the next N-group's fmacs. **4. Paired store blockIdxM guard mismatch** (`GlobalWriteBatch.py`) `SubtileMGuard` counts valid M-blocks in `MatrixInstM` (16-row) units, but `blockIdxM` for paired and orphan stores was computed in 32-row units (`(tt0-1)//2` or `(tt0*16)//32`). This mismatch caused incorrect OOB guard decisions. Fixed by using `blockIdxM = tt0` (16-row index) to match MGuard units. **5. Scalar fallback for partial paired stores** (`GlobalWriteBatch.py`) Added a scalar `dwordx2` fallback when only the lower M-block in a pair is valid (`MGuard > tt0-1` but not `MGuard > tt0`). Previously, the paired `dwordx4` store would execute for both blocks even when the upper block was OOB. This case arises when the tile remainder has an odd number of `MatrixInstM`-sized blocks (e.g., remainder=48 = 3 blocks of 16 rows). **6. sba=1 orphan store for large macro tiles** (`GlobalWriteBatch.py`) When `MIWaveTile[0]` is large enough that batch boundaries split an (sba=0, sba=1) pair, the sba=1 element had no partner and was silently dropped. Added scalar store handling for this orphan case. **7. Relaxed edge/NonEdge dispatch alignment** (`KernelWriterAssembly.py`) The `checkIsEdgeSubtile` M-dimension alignment was `waveGroupM` (e.g., 48 for MIWT3), requiring the tile remainder to be a multiple of the full wave group height. Reduced to `MatrixInstM` (16), so any remainder that is a multiple of 16 rows takes the optimized NonEdge path. The NonEdge path's MGuard + scalar fallback (fix #5) handles partial wave groups correctly. All changes are guarded by `UseSubtileImpl` — no impact on non-subtile kernels. ## Test Plan - Run `subtile_bf16.yaml` (BF16 BBS/BSS, multiple MIWT configs, SK3, PGR0/PGR2) with tile-aligned sizes (M % 32 == 0, N % 16 == 0) - Run `subtile_mxfp4.yaml` (MXFP4 F4BS/F4HS/F4SS, bias, activations, ScaleAlphaVec, PGR0/PGR2) with tile-aligned sizes (M % 32 == 0, N % 32 == 0) - Verified edge stores are not exercised for tile-aligned sizes by temporarily disabling edge path stores and confirming all tests still pass - Verified with rocgdb breakpoints that MGuard/NGuard values, C-load data, and accumulator values are correct at N-group boundaries ## Test Result - `subtile_bf16.yaml` (tile-aligned sizes): all tests PASSED - `subtile_mxfp4.yaml` (tile-aligned sizes): all tests PASSED - Edge-stores-disabled verification: all tests PASSED (confirms tile-aligned sizes use NonEdge path exclusively) ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Claude Opus 4 <noreply@anthropic.com>
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
Fix correctness failures in the
UseSubtileImplNonEdge store path for gfx950 BF16 and MXFP4 subtile kernels. These failures were caused by several interrelated bugs in the interleaved (GLS=0) store codegen that manifested when M-guard branches skipped stores at runtime, and by an overly strict edge/NonEdge dispatch check that forced subtile-aligned workgroups into the unoptimized edge path.Technical Details
1. vmcnt hazard in interleaved B1 NonEdge store path (
GlobalWriteBatch.py)On gfx950 (
SeparateVscnt=False), loads and stores share the same vmcnt counter. The interleaved store path computedvmcnt = vlcnt + vscntwherevscnt = self.storesIssued. When M-guard branches skip stores at runtime, the actual vmcnt counter has fewer outstanding operations than codegen assumed, makings_waitcnt vmcnt(N)too permissive — the hardware doesn't wait long enough for C-loads to complete before fmacs consume them. Fixed by settingvscnt = 0forUseSubtileImplwithGroupLoadStore=False. This is more conservative but correct regardless of how many stores were skipped at runtime.2. SrdD increment skipped by M-guard branch (
GlobalWriteBatch.py)The SrdD
incToNextRowwas emitted inside the M-guard-skippable region of the store loop. When the M-guard branch skipped the last store in an N-group, the SrdD increment was also skipped, causing subsequent N-group stores to write to the wrong row address. Fixed by deferring the increment and emitting it after the N-group end label.3. N-group end label placed after next N-group's fmacs (
GlobalWriteBatch.py)When
GroupLoadStore=False, fmacs and stores are interleaved in the same module. The fmacs for N-group K+1 were emitted before the N-group K end label was placed. The M-guard branch (targeting the end label) would skip both the last store of N-group K and the fmacs for N-group K+1, leaving the K+1 accumulators at zero. Fixed by flushing the pending N-group end label and deferred SrdD increment before emitting the next N-group's fmacs.4. Paired store blockIdxM guard mismatch (
GlobalWriteBatch.py)SubtileMGuardcounts valid M-blocks inMatrixInstM(16-row) units, butblockIdxMfor paired and orphan stores was computed in 32-row units ((tt0-1)//2or(tt0*16)//32). This mismatch caused incorrect OOB guard decisions. Fixed by usingblockIdxM = tt0(16-row index) to match MGuard units.5. Scalar fallback for partial paired stores (
GlobalWriteBatch.py)Added a scalar
dwordx2fallback when only the lower M-block in a pair is valid (MGuard > tt0-1but notMGuard > tt0). Previously, the paireddwordx4store would execute for both blocks even when the upper block was OOB. This case arises when the tile remainder has an odd number ofMatrixInstM-sized blocks (e.g., remainder=48 = 3 blocks of 16 rows).6. sba=1 orphan store for large macro tiles (
GlobalWriteBatch.py)When
MIWaveTile[0]is large enough that batch boundaries split an (sba=0, sba=1) pair, the sba=1 element had no partner and was silently dropped. Added scalar store handling for this orphan case.7. Relaxed edge/NonEdge dispatch alignment (
KernelWriterAssembly.py)The
checkIsEdgeSubtileM-dimension alignment waswaveGroupM(e.g., 48 for MIWT3), requiring the tile remainder to be a multiple of the full wave group height. Reduced toMatrixInstM(16), so any remainder that is a multiple of 16 rows takes the optimized NonEdge path. The NonEdge path's MGuard + scalar fallback (fix #5) handles partial wave groups correctly.All changes are guarded by
UseSubtileImpl— no impact on non-subtile kernels.Test Plan
subtile_bf16.yaml(BF16 BBS/BSS, multiple MIWT configs, SK3, PGR0/PGR2) with tile-aligned sizes (M % 32 == 0, N % 16 == 0)subtile_mxfp4.yaml(MXFP4 F4BS/F4HS/F4SS, bias, activations, ScaleAlphaVec, PGR0/PGR2) with tile-aligned sizes (M % 32 == 0, N % 32 == 0)Test Result
subtile_bf16.yaml(tile-aligned sizes): all tests PASSEDsubtile_mxfp4.yaml(tile-aligned sizes): all tests PASSEDSubmission Checklist