Skip to content

[COMGR][HotSwap] Fix M32 scaled WMMA lowering - #1

Draft
paulzzy wants to merge 12 commits into
users/pauzhang/rocm28594-amd-staging-basefrom
users/pauzhang/rocm28594-m32-scaled-wmma-on-3677-v1
Draft

[COMGR][HotSwap] Fix M32 scaled WMMA lowering#1
paulzzy wants to merge 12 commits into
users/pauzhang/rocm28594-amd-staging-basefrom
users/pauzhang/rocm28594-m32-scaled-wmma-on-3677-v1

Conversation

@paulzzy

@paulzzy paulzzy commented Jul 29, 2026

Copy link
Copy Markdown
Owner

Jira: ROCM-28594

Summary

Fix regular-scale and Scale16 M32 WMMA lowering for gfx1250 B0-to-A0
HotSwap.

The change:

  • splits regular-scale M32 operations into two exact M16 operations and uses
    MATRIX_SCALE_ROW1 for the upper M half;
  • replaces in-place Scale16 A/scale mutation with generated bank-zero scratch;
  • executes M0 low/high followed by M1 low/high;
  • keeps masked A contiguous while allowing four generated scales and the
    gather temporary to use independent dead low-bank VGPRs;
  • avoids unnecessary VGPR growth when 13 dead registers are available but are
    not one contiguous interval;
  • consumes legacy gfx1250 tensor instructions through canonical LLVM MC
    decoding instead of COMGR raw-byte classifiers;
  • exposes stable AMDGPU named operand roles to MC clients and builds the WMMA
    replacement instructions from structured MCOperands rather than parsing
    MCInstPrinter text or mirroring positional operand layouts;
  • reuses a locally dead, already-declared SGPR for saturated required tensor
    wrappers when current-text CFG liveness proves it safe;
  • rechecks only required tensor trampolines after final pool layout and
    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-staging at
7f7448fc5af9.

Please review these 11 signed commits in order:

  1. 63c2b1e1136e - decode the observed legacy gfx1250 tensor DMA aliases as
    canonical LLVM MC instructions;
  2. c72aecf655ac - exact M32 Scale16 lowering;
  3. 83608054a2b2 - make the M32 Scale16 lowering bank-safe;
  4. 892dafb85d8d - drain XCNT across M32 VGPR bank changes;
  5. bd3c63223442 - fix regular/Scale16 M32 splitting, preserve original
    operands, and support fragmented scratch;
  6. 452bd2650f3e - remove temporary raw tensor classifiers and consume LLVM
    MC semantics;
  7. 4f76e8635f84 - add legacy tensor integration coverage;
  8. d1dec3d5f3ff - route required tensor wrappers in saturated kernels;
  9. 09e071d5e59e - expose a public AMDGPU named-operand adapter;
  10. c2932c4e05d8 - add shared COMGR named-MC helpers;
  11. f456866e0d33 - remove printer parsing and build WMMA splits from
    structured 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_scaled failures. The failing kernels use the B0-only
32x16x128 FP4 WMMA shape:

  • 208 cases use E8M0 scale encoding;
  • 208 cases use E4M3 scale encoding;
  • both groups exercise MXFP4 matrix operands.

A representative reproducer is a non-transposed M=32, N=32, K=128
FP4-by-FP4 kernel with both E8M0 scales enabled. It emits
v_wmma_scale_f32_32x16x128_f4 as intended on B0, but the HotSwapped result
is 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 original MCOperand values,
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:

  • legacy and canonical gfx1250 tensor MC tests pass;
  • the exact 12-byte legacy record decodes as canonical
    TENSOR_LOAD_TO_LDS_d4_gfx1250.

Host Release:

  • HotswapMCTests: 194/194;
  • focused WMMA LIT: 41/41;
  • full supported COMGR LIT: 166/166;
  • CTest: 32/32.

Clang-19 ASan:

  • focused HotswapMCTests: 194/194;
  • focused WMMA LIT: 41/41;
  • full supported COMGR LIT: 166/166.

A0 numerical validation:

  • one-sided Scale16 regression: 1/1 passed with candidate-load proof;
  • complete M32 wmma_shape1 family: 208 passed, 3,536 skipped, zero failed;
  • pre/post device-holder audits were clean.

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:

  • parent: 20,382 pass / 58 fail / 2 established timeout;
  • candidate: 20,392 pass / 48 fail / 2 established timeout;
  • status changes: ten previously failing inputs -> pass, zero pass -> fail;
  • common outputs: 20,381 unchanged, one expected Scale16 output changed;
  • candidate-only outputs: ten;
  • all 20,392 emitted outputs passed ELF validation and byte-identical second
    rewrite.

The full host and A0 validation above was completed on the merged-ROCm#3677
amd-staging tip used during development. The final refresh onto live
amd-staging (7f7448fc5af9) was conflict-free; all 11 stable patch IDs and
the complete git range-diff are identical, and git diff --check passes.

@paulzzy
paulzzy force-pushed the users/pauzhang/rocm28594-m32-scaled-wmma-pr3677-legacy-tensor-base branch from e46b3ff to 1385ffe Compare July 29, 2026 06:15
@paulzzy
paulzzy force-pushed the users/pauzhang/rocm28594-m32-scaled-wmma-on-3677-v1 branch from 5c4ec4e to ac61ffe Compare July 29, 2026 06:16
paulzzy and others added 11 commits July 29, 2026 11:30
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)
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
paulzzy force-pushed the users/pauzhang/rocm28594-m32-scaled-wmma-on-3677-v1 branch from ac61ffe to f456866 Compare July 29, 2026 18:40
@paulzzy
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants