[COMGR][HotSwap] Fix M32 scaled WMMA lowering - #1
Draft
paulzzy wants to merge 12 commits into
Draft
Conversation
paulzzy
force-pushed
the
users/pauzhang/rocm28594-m32-scaled-wmma-pr3677-legacy-tensor-base
branch
from
July 29, 2026 06:15
e46b3ff to
1385ffe
Compare
paulzzy
force-pushed
the
users/pauzhang/rocm28594-m32-scaled-wmma-on-3677-v1
branch
from
July 29, 2026 06:16
5c4ec4e to
ac61ffe
Compare
Add dedicated gfx1250 real records for the legacy d031 load and store forms. Model their zero dim, dmask, and unused vaddr4 fields while keeping canonical assembly on the existing d071 records. Normalize legacy decoder results to the canonical tensor opcodes before operand post-processing so MC clients observe the existing 12-byte instruction model. Cover d2 and d4 load and store forms in both B0 and A0 feature states, and retain the canonical encoder bytes. (cherry picked from commit 2774376fc88c615b723ccf573c91fc2f4d450a5e)
(cherry picked from commit fef5323)
(cherry picked from commit dc80021)
(cherry picked from commit a655e97)
Lower regular-scale M32 WMMA into two exact M16 operations and select MATRIX_SCALE_ROW1 for the upper M half. Replace the M32 Scale16 in-place A and scale mutation with a non-mutating scratch schedule. Keep masked A in one contiguous bank-zero interval, gather the four generated scales into dedicated bank-zero registers, and process M0 low/high before M1 low/high. Do not require all 13 scratch registers to be contiguous. When a full dead interval is unavailable, allocate the eight masked-A registers as one block and the four scales plus temporary from independent dead low-bank slots. This avoids unnecessary VGPR growth and the associated metadata-size rewrite failure. Add focused regular-scale, bank-zero, non-mutating, and split-scratch regressions. Co-authored-by: Harsh Menon <harsh.menon@amd.com>
Remove the M32 raw-byte tensor/VOP3 classifiers now that the AMDGPU MC dependency decodes legacy gfx1250 tensor DMA aliases as canonical 12-byte instructions. Let control-flow, VGPR-MSB, and physical-liveness analyses consume the canonical MCInsts through their ordinary paths. Drop the classifier-only unit test; the LLVM MC and COMGR integration tests cover the real encodings.
Assemble legacy gfx1250 tensor load and store d2 and d4 forms through the new LLVM MC records. Verify canonical 12-byte decode boundaries and reuse of the existing tensor-load mask wrappers. Check that stores keep their canonical operands without receiving the load-only wrapper. Also cover ELF validity and byte-identical A0 rewrite. (cherry picked from commit 1dc3ae15754bddb45b591db7533c8044f476db3a)
Reuse a kernel-local SGPR only when current-text CFG liveness proves its incoming value is unused by the tensor wrapper and continuation. Keep the fallback within the declared numbered-SGPR footprint so it needs no metadata growth. Mark required tensor trampolines for a final-layout recheck after every earlier long-trampoline island dword is known. Promote only marked entries whose short edge shifted out of range, keeping unrelated optional rewrite outputs unchanged. Add saturated legacy-tensor, genuine no-scratch fail-closed, and exact negative branch-boundary regressions.
Expose a stable subset of AMDGPU's TableGen named-operand table through an installed header. This lets MC clients such as COMGR query operand roles without depending on backend-private AMDGPUBaseInfo headers or mirroring positional layouts.
Link the AMDGPU utility component explicitly and add shared helpers for parsing, encoding, copying, and transforming named MC operands. Cover the installed AMDGPU operand-role adapter with a gfx1250 scaled-WMMA unit test.
Replace MCInstPrinter parsing and mirrored positional WMMA layouts with AMDGPU named operands. Parse structural replacement MCInst templates, copy inline constants and modifiers as MCOperands, and encode the transformed instructions through MCCodeEmitter. Derive the Scale16 mask plan from the decoded matrix tuple width.
paulzzy
force-pushed
the
users/pauzhang/rocm28594-m32-scaled-wmma-on-3677-v1
branch
from
July 29, 2026 18:40
ac61ffe to
f456866
Compare
paulzzy
changed the base branch from
users/pauzhang/rocm28594-m32-scaled-wmma-pr3677-legacy-tensor-base
to
users/pauzhang/rocm28594-amd-staging-base
July 29, 2026 18:41
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.
Jira: ROCM-28594
Summary
Fix regular-scale and Scale16 M32 WMMA lowering for gfx1250 B0-to-A0
HotSwap.
The change:
MATRIX_SCALE_ROW1for the upper M half;gather temporary to use independent dead low-bank VGPRs;
not one contiguous interval;
decoding instead of COMGR raw-byte classifiers;
replacement instructions from structured
MCOperands rather than parsingMCInstPrintertext or mirroring positional operand layouts;wrappers when current-text CFG liveness proves it safe;
promotes a shifted out-of-range short return into the transactional
branch-island path merged in PR [comgr][hotswap] Make branch-island allocation transactional ROCm/llvm-project#3677.
Review scope
The PR base branch mirrors
ROCm/llvm-project:amd-stagingat7f7448fc5af9.Please review these 11 signed commits in order:
63c2b1e1136e- decode the observed legacy gfx1250 tensor DMA aliases ascanonical LLVM MC instructions;
c72aecf655ac- exact M32 Scale16 lowering;83608054a2b2- make the M32 Scale16 lowering bank-safe;892dafb85d8d- drain XCNT across M32 VGPR bank changes;bd3c63223442- fix regular/Scale16 M32 splitting, preserve originaloperands, and support fragmented scratch;
452bd2650f3e- remove temporary raw tensor classifiers and consume LLVMMC semantics;
4f76e8635f84- add legacy tensor integration coverage;d1dec3d5f3ff- route required tensor wrappers in saturated kernels;09e071d5e59e- expose a public AMDGPU named-operand adapter;c2932c4e05d8- add shared COMGR named-MC helpers;f456866e0d33- remove printer parsing and build WMMA splits fromstructured operands.
Commits 2-4 preserve @harsh-amd's original implementation history. The primary
corrective review is commits 5-11.
Root cause
ROCM-28594 tracks 416 parameterized Triton/Gluon
test_amd_wmma_scaledfailures. The failing kernels use the B0-only32x16x128FP4 WMMA shape:A representative reproducer is a non-transposed
M=32, N=32, K=128FP4-by-FP4 kernel with both E8M0 scales enabled. It emits
v_wmma_scale_f32_32x16x128_f4as intended on B0, but the HotSwapped resultis numerically incorrect on A0.
A0 does not provide the M32 instruction directly. Regular scaled WMMA must be
split into two M16 operations. Scale16 M32 additionally needs a K-subblock
split, producing four exact M16 block-32 WMMAs.
The earlier Scale16 lowering modified matrix A and both scale pairs in place.
That corrupts kernels where only one matrix operand is scaled because later
generated passes still need the original A and scale values. Fully scaled
variants can fail for the same reason when the generated M/K passes reuse
mutated inputs.
Larger parameter combinations also expose a scratch-allocation failure. A
representative 128-VGPR kernel has an eight-register dead low-bank interval
plus five independently dead low-bank registers, but no contiguous
13-register interval. Requiring all scratch to be contiguous grows the kernel
by 13 VGPRs. That changes the metadata note size, prevents a size-preserving
rewrite, and leaves the original B0-only instruction for A0 execution.
The final allocation keeps only the eight masked-A registers contiguous and
allocates the four generated scales and temporary independently. Original A
and scale operands are never modified.
The prior implementation also reconstructed WMMA replacements from printed
assembly text and duplicated backend operand positions. That was fragile for
inline floating-point immediates, modifiers, tied operands, and aliased
source/destination operands. The final implementation queries named operand
roles, transforms
MCInsts directly, copies the originalMCOperandvalues,and emits the result through
MCCodeEmitter.Canonical decoding also exposes required legacy tensor wrappers in kernels
whose declared SGPR footprint is saturated. Those kernels can still have a
register that is dead across the wrapper and its continuation. The fallback
uses current-text CFG liveness to reuse that register without metadata growth.
If final trampoline layout shifts a required short return out of range, the
site is promoted into PR ROCm#3677's transactional registerless island path.
Testing
LLVM MC:
TENSOR_LOAD_TO_LDS_d4_gfx1250.Host Release:
HotswapMCTests: 194/194;Clang-19 ASan:
HotswapMCTests: 194/194;A0 numerical validation:
wmma_shape1family: 208 passed, 3,536 skipped, zero failed;The numerical runs preserved 169 unique Triton input HSACOs. Rewriting all
169 with the prior numerically validated candidate and this candidate produced
169/169 byte-identical, idempotent outputs with zero status or hash changes.
20,442-object corpus:
rewrite.
The full host and A0 validation above was completed on the merged-ROCm#3677
amd-stagingtip used during development. The final refresh onto liveamd-staging(7f7448fc5af9) was conflict-free; all 11 stable patch IDs andthe complete
git range-diffare identical, andgit diff --checkpasses.