llvm / llvm/llvm-project

[AMDGPU][gfx950] Tied 8-VGPR boundary changes HSTU dR codegen versus ordinary SSA

Open
#222,423 11 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

On gfx950, adding a numerical-identity tied `<8 x i32> "=v,0"` boundary around an FP32 `dR` value changes LLVM O3 vectorization and the final code for a production HSTU backward kernel. A freshly collected balanced-order pass measures the resident form about 4% faster, although both variants use 256 VGPRs, occupancy 1, and no scratch.

This issue now asks whether ordinary SSA at this `dR` site can recover the same result, and which part of the boundary matters: vector shape, the tied VGPR constraint, inline-assembly opacity, or changed use-def/liveness structure.

## Unresolved HSTU dR case

The only source-level toggle is:

```python
dr = tlx.amd_register_resident(
dr,
register_class="vgpr",
registers_per_group=8,
)
```

The source is [here](https://github.com/facebookexperimental/triton/blob/c54ae5685d9f3d8a881e3a5cbe46bb9587d3b880/third_party/tlx/tutorials/hstu_self_attn/tlx_gfx950_ragged_hstu_attention.py#L2789-L2797). For every thread-local group of eight FP32 values, the lowering builds an `<8 x i32>` value, passes it through `asm "", "=v,0"`, and unpacks it without changing the bits.

The operation changes all of those properties together. The current evidence does not isolate tuple grouping or coalescing as the cause.

“Ordinary SSA” is local to this `dR` site. The direct kernel still has 319 other inline-assembly calls, including 24 other tied `<8 x i32> "=v,0"` calls.

## Where the variants diverge

Before LLVM O3, the resident module has 12 inlined `dR` wrappers. Each wrapper adds eight inserts, two bitcasts, one tied inline-assembly call, and eight extracts. Removing those wrappers and rewiring their 96 outputs makes the pre-O3 modules identical under `llvm-diff`.

The pre-O3 arithmetic is identical: 576 `fadd`, 1,416 `fmul`, and 96 `fsub` operations in each variant. In a standalone O3 replay with the bundled LLVM:

- `SLPVectorizerPass` first chooses different scalar/vector arithmetic forms.
- `VectorCombinePass` first changes total arithmetic counts and forms four eight-wide FP32 `dR` groups in the resident variant.

The captured post-O3 IR already differs before standalone `llc`:

| LLVM IR opcode | Direct | Resident |
| --- | ---: | ---: |
| `fadd` | 480 | 468 |
| `fmul` | 948 | 924 |
| `fsub` | 48 | 36 |
| `shufflevector` | 876 | 984 |

This is therefore a whole-pipeline observation, not an isolated AMDGPU scheduler or register-allocation reproducer.

## Runtime and final assembly

Workload: BF16, `B=512`, `H=4`, `max_seq_len=2048`, `D=128`, 995,460 tokens, lower-triangular HSTU mask.

The balanced ABBA pass was collected after the fixed-AB pass in the same process.

| Method | Direct median | Resident median | Direct / resident |
| --- | ---: | ---: | ---: |
| Fixed AB, 31 pairs | 16.5826 ms | 15.9530 ms | 1.0395x |
| Balanced ABBA, 32 pairs | 16.5907 ms | 15.9428 ms | 1.0406x |

These are cross-variant comparisons, not checks against an independent reference. At the full shape, dK and dV are storage-bit identical. dQ is finite with 60.99 dB SNR (direct output as the signal) and maximum absolute difference `1.0914e-11`. A smaller `(B,H,N,D)=(8,4,1024,128)` control produces finite, storage-bit-identical dQ, dK, and dV.

| Metric | Direct | Resident |
| --- | ---: | ---: |
| VGPR | 256 | 256 |
| AGPR | 126 | 128 |
| Total vector registers | 382 | 384 |
| Scratch bytes | 0 | 0 |
| Compiler-reported occupancy | 1 | 1 |
| `v_perm_b32` | 72 | 40 |
| `v_pk_mul_f32` | 374 | 410 |
| Static MFMA sites | 392 | 392 |

Other instructions also change, so these counts are correlation rather than a claim that any one code-generation difference causes the speedup.

Exact pre/post-O3 IR, bundled assembly, benchmark source, and raw timing/correctness results: https://gist.github.com/bangtianliu/46792c16bc698fe387a82d031255c318/cf17e6c8382dea0f6e6b4bc6a5adad3364eda0fa

Provenance:

```text
Triton source: c54ae5685d9f3d8a881e3a5cbe46bb9587d3b880
Bundled LLVM: 850a2b1b975c061ae0fc982ba68064d305485cb2
Compile-only upstream LLVM: 2de38b4861ed4a7be7ddd42fc803e91d1fe81b14
GPU: AMD Instinct MI350X, gfx950
Driver / ROCm: 6.16.6 / 7.1.1
PyTorch / HIP: 2.13.0+rocm7.1 / 7.1.52802
```

Both post-O3 files compile with the pinned upstream LLVM revision. This report makes no upstream-runtime claim.

## Questions

1. Is ordinary SSA expected to reach the useful vectorization/codegen seen with this tied boundary?
2. Because the first divergence appears in SLP/vector combining, should this be treated primarily as a generic vectorization problem, an AMDGPU lowering problem, or an interaction between both?
3. What additional ablation or reduced IR would best separate vector grouping, register constraints, opacity, and liveness effects?

This report does not propose a new LLVM API or prescribe an implementation. Its goal is to identify what information is missing and which compiler stage should own it.

Resolved original BF16/AGPR source case

The issue originally asked whether a reused BF16 MFMA A/B source should be forced into AGPRs. The replies [here](https://github.com/llvm/llvm-project/issues/222423#issuecomment-5607170511) and [here](https://github.com/llvm/llvm-project/issues/222423#issuecomment-5607201133) explain that AMDGPU normally selects VGPR forms and rewrites to AGPRs later when register pressure makes that useful.

Removing only the hard BF16 `=a` preserve constraints from [Triton PR #2709](https://github.com/facebookexperimental/triton/pull/2709) retained correctness, occupancy 1, and zero scratch/spills. It measured 30.027 ms with normal allocation versus 30.182 ms with hard `=a`. The integer bitcast was also unnecessary for LLVM type legality.

That original case therefore does not motivate explicit AGPR residency. Its immutable reproducer remains at: https://gist.github.com/bangtianliu/a6719f44caefbe119be4f719e384c296/34be3cfc3038c24e363d5ee59280ff18ff39e3e6

Contributor guide

Open the contributing guide

Research direction

Start with third_party/tlx/tutorials/hstu_self_attn/tlx_gfx950_ragged_hstu_attention.py at the cited dR site, then compare the linked pre- and post-O3 IR and the standalone LLVM replay. Construct a reduced ablation that separates vector grouping, tied register constraints, opacity, and liveness, and identify whether SLP/vector combining or AMDGPU lowering owns the behavior.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.