llvm / llvm/llvm-project

[AMDGPU][gfx950] CoExec does not consider persistent 32x32 MFMA chains for AGPR allocation

Open
#222,182 5 comments 0 reactions 0 assignees View on GitHub
backend:AMDGPU
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

## Summary

On upstream `llvm-project/main` at `66462afa7e4a33f0d720877f2228589d5b43de81`, CoExec does not consider any of the 32 `32x32x16` MFMA instructions containing the persistent accumulator live ranges for its pre-RA VGPR-to-AGPR rewrite. A later post-RA pass converts 10 of the 32 instructions, but the final intrinsic lowering uses 204 B of scratch versus 68 B for a contextual mixed AGPR/VGPR constrained lowering.

This is a static code-generation report. It does not claim that the constrained result is globally optimal, that bank selection is the only source of the difference, or that a particular LLVM implementation is required.

Reproducer bundle: [upstream-main revision `6321467dee2f8740df78eb1777bc91f9513bb62c`](https://gist.github.com/bangtianliu/136ea4c92624d19cc8e19b5a61c56b39/6321467dee2f8740df78eb1777bc91f9513bb62c)

## Input

- Target: `gfx950`
- Kernel: BF16 backward attention, `B=1, Hq=8, Hkv=1, N=16384, D=128`, non-causal
- Launch: `BLOCK_M=16, BLOCK_N=256`, four warps
- Primary input: `persistent-mfma.ll`, containing 96 `16x16x32` and 32 `32x32x16` LLVM MFMA intrinsic calls
- Contextual baseline: `constrained-baseline.ll`, in which 24 persistent `32x32x16` updates use inline assembly with 16 VGPR and eight AGPR destination constraints

Both inputs retain the same 26 empty inline-assembly handoff nodes elsewhere in the kernel: 20 side-effecting nodes and six pure nodes. The changed sites are the 24 persistent MFMA updates.

## Inline assembly versus direct SSA

The constrained LLVM IR makes the destination bank and fixed-delay padding explicit:

```llvm
; VGPR form
%out = call <16 x float> asm sideeffect
"s_nop 3\0Av_mfma_f32_32x32x16_bf16 $0, $1, $2, $0",
"=&v,v,v,0"(<8 x bfloat> %a, <8 x bfloat> %b, <16 x float> %acc)
%bits = bitcast <16 x float> %out to <16 x i32>
%padded = call <16 x i32> asm sideeffect
"s_nop 15\0As_nop 3", "=v,0"(<16 x i32> %bits)

; AGPR form uses "=a,v,v,0" and matching "=a,0" post-MFMA padding.
```

The corresponding direct-SSA LLVM IR keeps the MFMA visible to LLVM but carries no accumulator-bank constraint:

```llvm
%out = call <16 x float>
@llvm.amdgcn.mfma.f32.32x32x16.bf16(
<8 x bfloat> %a, <8 x bfloat> %b, <16 x float> %acc,
i32 0, i32 0, i32 0)
```

The examples are matched by debug location (`!110` for VGPR and `!168` for AGPR); SSA numbers differ because the files were captured independently. The direct form preserves the numerical SSA dependency and exposes the real MFMA opcode to scheduling and hazard handling. The constrained form communicates `VGPR` or `AGPR`, but hides the MFMA in a side-effecting `INLINEASM` node. It also adds early-clobber at the VGPR sites and a tied post-MFMA use/def that keeps the selected bank live through fixed-delay padding. Therefore this is not a pure register-class A/B. The bundle contains the matched full inputs and a [detailed paired IR comparison](https://gist.github.com/bangtianliu/136ea4c92624d19cc8e19b5a61c56b39/6321467dee2f8740df78eb1777bc91f9513bb62c#file-ir-comparison-md). The `s_nop` sequences are static delays, not scoreboard waits or synchronization operations.

## Final code generation

Both CoExec compilations pass `-verify-machineinstrs` and retain occupancy 1. All 96 `16x16x32` MFMAs have VGPR destinations in both outputs.

| Input | `32x32x16` AGPR/VGPR destinations | AGPR reads/writes | Scratch loads/stores | `ScratchSize` |
| --- | ---: | ---: | ---: | ---: |
| Direct SSA MFMA intrinsics | 10 / 22 | 138 / 128 | 26 / 26 | 204 B |
| Mixed constrained lowering | 15 / 17 | 152 / 156 | 12 / 12 | 68 B |

## Where the upstream automatic path stops

Immediately before machine scheduling, instruction selection has selected all 128 MFMAs in VGPR-destination form. Auditing the unmodified MIR against `RewriteMFMAFormStage::isRewriteCandidate` gives:

| MFMA form | Pre-RA instructions | Eligible candidates | Candidate components |
| --- | ---: | ---: | --- |
| `16x16x32` | 96 | 76 | `4 x 7` and `16 x 3` |
| `32x32x16` | 32 | 0 | none |

All 32 wide MFMAs are rejected directly by the current destination-use check:

```cpp
for (const MachineInstr &UseMI : DAG.MRI.use_nodbg_instructions(DstReg)) {
if (!TII->isMAI(UseMI) && !UseMI.isCopy())
return false;
}
```

Of the 32 wide `32x32x16` MFMAs, 16 are first rejected by `V_CVT_PK_BF16_F32_e64` and the other 16 by `V_PK_MUL_F32`. Thus the 24 persistent calls under study never enter the pre-RA bank-rewrite candidate set.

The 76 eligible candidates are unrelated `16x16x32` path prefixes. Diagnostic-only logging of upstream's unchanged cost calculation reports `-4096` spill benefit, `95040` `CopyForUse` cost, and total cost `90944`, so that complete small-MFMA rewrite is rejected. This cost is not a decision about the excluded wide MFMAs.

Immediately before `AMDGPURewriteAGPRCopyMFMA`, all 32 wide MFMAs still use VGPR form. That post-RA pass rewrites 10 to AGPR form; 22 remain VGPR. The final direct-SSA allocation therefore comes from opportunistic post-RA copy folding rather than a pre-RA mixed-bank decision for the persistent accumulator families.

## Reproduction

```bash
BUNDLE=/path/to/reproducer-bundle
LLC=/path/to/llvm-main-build/bin/llc
COMMON=(
-mtriple=amdgpu9.50-amd-amdhsa
-mcpu=gfx950
-verify-machineinstrs
-amdgpu-sched-strategy=coexec
)

"$LLC" "${COMMON[@]}" "$BUNDLE/persistent-mfma.ll" -o direct-ssa.s
"$LLC" "${COMMON[@]}" "$BUNDLE/constrained-baseline.ll" -o constrained.s
python3 "$BUNDLE/summarize.py" direct-ssa.s constrained.s
```

The bundle also includes an optional `LLVM_DEBUG`-only patch, its filtered output, exact hashes, and commands. The diagnostic build produces byte-for-byte identical assembly to the clean upstream build.

## Questions

1. Is the exclusion of these loop-carried `32x32x16` accumulators expected when their virtual registers also feed ordinary vector consumers?
2. Which current upstream pass is intended to choose a mixed AGPR/VGPR allocation for such persistent MFMA accumulator families?
3. Is `AMDGPURewriteAGPRCopyMFMA` expected to recover this case after allocation, or is the absence of a pre-RA candidate the relevant limitation?

## Scope

No runtime-performance or numerical-correctness claim is made. The constrained lowering is only evidence that a lower-scratch mixed-bank result exists in this kernel context; its opaque MFMA assembly, fixed-delay padding, extra tied use/defs, early-clobber constraints, and scheduling effects are not proposed as the desired compiler representation.

Contributor guide

Open the contributing guide

Research direction

Start by running llc with the supplied gfx950 commands on persistent-mfma.ll and constrained-baseline.ll, then compare the summarized output. Read RewriteMFMAFormStage::isRewriteCandidate and AMDGPURewriteAGPRCopyMFMA alongside the candidate counts in the report. Done means determining which pass should handle the persistent 32x32x16 accumulators and whether their exclusion is an upstream limitation.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.