| Issue |
204254
|
| Summary |
[AMDGPU] Register allocator places all three v_wmma_f32_16x16x16_f16 operands on the same VGPR bank (gfx11/gfx12)
|
| Labels |
new issue
|
| Assignees |
|
| Reporter |
mgehre-amd
|
## Summary
On RDNA3 / RDNA3.5 (gfx11, e.g. `gfx1151`) the back-to-back issue latency of
`v_wmma_f32_16x16x16_f16` depends on the VGPR **bank** of its three operands
(bank = first VGPR number `% 4`). When the A (`src0`), B (`src1`) and
C (`src2`/`dst`) operands each start in a *distinct* bank, a dependent WMMA issues
every **32 cycles**. When any two operands share a bank, the instruction pays a
**+2 cycle** penalty (**34 cycles**).
The AMDGPU register allocator currently lays the operands out consecutively, which
on this instruction *always* produces a bank collision, so every WMMA emitted from
the intrinsic runs at 34 cyc instead of the achievable 32 cyc — a ~6% throughput
loss on WMMA-bound kernels that is invisible at the source level and cannot be
fixed without dropping to inline assembly.
## Environment
- Target: `amdgcn-amd-amdhsa`, `--offload-arch=gfx1151` (RDNA3.5; also applies to
other gfx11 and the gfx12 WMMA encoding).
- Compiler: `clang version 23.0.0git` (ROCm/llvm-project @ 43215c73116c).
- Device measured on: Radeon 8060S (gfx1151).
## Root cause
Each f16 `16x16x16` WMMA operand occupies **8 VGPRs**, and `8 % 4 == 0`. When the
allocator assigns operands to consecutive register ranges (A at `vN`, B at `vN+8`,
C at `vN+16`, …) every operand's starting register is congruent mod 4, so all three
land on the **same** bank. The allocator never inserts the 1-register gap that would
move the operands onto separate banks.
## Reproducer
No rocWMMA required — the builtin alone reproduces it. Compile **device-only to
assembly** and look at the operand registers:
```c++
// wmma_bank.hip
#include <hip/hip_runtime.h>
typedef __fp16 half16 __attribute__((ext_vector_type(16))); // A / B : 8 VGPRs each
typedef float float8 __attribute__((ext_vector_type(8))); // C : 8 VGPRs
__global__ void k(float* out, const half16* ga, const half16* gb, int n) {
half16 a = ga[threadIdx.x]; // loaded from distinct buffers so CSE cannot
half16 b = gb[threadIdx.x]; // merge A and B into one register
float8 c = {0,0,0,0,0,0,0,0};
for (int i = 0; i < n; ++i)
c = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c); // dependent chain
if (threadIdx.x == 0) out[blockIdx.x] = c[0];
}
```
```sh
clang++ -x hip -O3 --offload-arch=gfx1151 --cuda-device-only -S wmma_bank.hip -o wmma_bank.s
grep v_wmma wmma_bank.s
```
### Actual output
```asm
v_wmma_f32_16x16x16_f16 v[1:8], v[9:16], v[17:24], v[1:8]
```
Operand banks:
| operand | start reg | bank (`reg % 4`) |
| -------------- | --------- | ---------------- |
| A (`src0`) | v9 | **1** |
| B (`src1`) | v17 | **1** |
| C (`src2`/dst) | v1 | **1** |
All three operands are on bank 1 → +2 cycle penalty.
### Expected output
Operands on three distinct banks, e.g. A on bank 0, B on bank 1, C on bank 2
(inserting small gaps such as `v8` / `v17`):
```asm
v_wmma_f32_16x16x16_f16 v[18:25], v[0:7], v[9:16], v[18:25]
; C bank2 A bank0 B bank1 C bank2
```
## Measured impact (gfx1151)
A fully dependent chain of 256 back-to-back WMMAs, timed with the
`HW_REG_SHADER_CYCLES` counter, single wave of 32 lanes:
| kernel | cyc/WMMA |
| -------------------------------------- | ---------- |
| `__builtin_amdgcn_wmma_*` (this issue) | **34.02** |
| hand-written asm, operands on banks 0/1/2 | **32.00** |
Delta: **+2.0 cyc/WMMA (~6%)**, exactly the bank-collision penalty. The two kernels
are otherwise identical (same chain, same inputs, same accumulator).
## Suggested fix
Teach the AMDGPU register allocator / operand-assignment for the matrix
(`v_wmma_*` / `v_swmma_*`) instructions to prefer assigning `src0`, `src1` and
`src2`/`dst` to distinct VGPR banks (start register `% 4` pairwise distinct),
inserting alignment gaps where free registers allow. Because each operand is a
multiple of 4 VGPRs wide, consecutive allocation is the worst case; even a hint
that biases the three operands onto different `% 4` residues recovers the full
throughput.
## Notes
- A self-contained timing reproducer (intrinsic vs. distinct-bank inline asm,
measured with `HW_REG_SHADER_CYCLES`) is included at the end of this report.
- The same 8-VGPR-operand / `% 4` reasoning applies to the other gfx11/gfx12 WMMA
variants whose operands are multiples of 4 registers wide; only the f16→f32
variant is shown here.
## Full timing reproducer
<details>
<summary><code>wmma_bank_repro.hip</code> — times the intrinsic kernel vs. a distinct-bank inline-asm kernel (click to expand)</summary>
```c++
/**
* Minimal reproducer: hipcc does not place WMMA operands on distinct VGPR banks.
*
* On RDNA3 / RDNA3.5 (gfx11), v_wmma_f32_16x16x16_f16 reads three operands
* (A = src0, B = src1, C = src2 = dst). The hardware can issue one back-to-back
* (dependent) WMMA every 32 cycles ONLY when the three operands start in three
* different VGPR banks (bank = first_vgpr_number % 4). If any two operands share
* a bank, the instruction pays a +2 cycle penalty (34 cycles per WMMA).
*
* Each f16 16x16x16 operand occupies 8 VGPRs, and 8 % 4 == 0, so when the
* register allocator lays operands out consecutively (A at vN, B at vN+8,
* C at vN+16, ...) ALL THREE land on the SAME bank. The compiler never inserts
* the small gap needed to separate them. This reproducer demonstrates that the
* intrinsic-generated kernel runs at ~34 cyc/WMMA while a hand-written inline-asm
* kernel with operands on distinct banks runs at the ideal 32 cyc/WMMA.
*
* What this program does:
* 1. Disassembles its own embedded gfx code object and prints the v_wmma ops of
* both kernels, annotating the VGPR bank of each operand (manual inspection).
* 2. Times a dependent chain of 256 WMMAs (back to back) for each kernel using
* the HW_REG_SHADER_CYCLES counter and reports cycles-per-WMMA.
*
* Build & run: ./wmma_bank_repro.sh
* NOT linked against rocWMMA — uses the __builtin_amdgcn_wmma_* intrinsic directly.
*/
#include <hip/hip_runtime.h>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <algorithm>
typedef __fp16 half16 __attribute__((ext_vector_type(16))); // WMMA A / B operand: 8 VGPRs
typedef float float8 __attribute__((ext_vector_type(8))); // WMMA C operand: 8 VGPRs
#define HIP_CHECK(cmd) do { \
hipError_t e = (cmd); \
if (e != hipSuccess) { \
fprintf(stderr, "HIP error '%s' at %s:%d\n", \
hipGetErrorString(e), __FILE__, __LINE__); \
exit(1); \
} \
} while (0)
#define WMMA_CHAIN 256 // Number of WMMAs back to back
// =============================================================================
// Kernel A: WMMA via the compiler intrinsic (register allocation up to hipcc).
//
// A and B are loaded from distinct global buffers so the compiler cannot merge
// them (CSE would otherwise collapse two identical fills into one register and
// hide the layout). C is a single accumulator reused by every WMMA, so the chain
// is fully dependent and cycles/WMMA == back-to-back issue latency.
// =============================================================================
__global__ __launch_bounds__(32, 8)
void wmma_intrinsic(float* __restrict__ result,
const half16* __restrict__ ga,
const half16* __restrict__ gb,
int iterations,
unsigned* __restrict__ cycles) {
half16 a = ga[threadIdx.x];
half16 b = gb[threadIdx.x];
float8 c;
#pragma unroll
for (int i = 0; i < 8; i++) c[i] = 0.0f;
unsigned start, stop;
// Snapshot the shader cycle counter, then pin the WMMA chain *after* the read:
// the empty asm makes `a` artificially depend on `start`, so the compiler
// cannot sink the counter read into the middle of the chain.
asm volatile("s_getreg_b32 %0, hwreg(HW_REG_SHADER_CYCLES, 0, 32)"
: "=s"(start) :: "memory");
asm volatile("" : "+v"(a) : "s"(start));
#pragma unroll 1
for (int it = 0; it < iterations; it++) {
#pragma unroll
for (int i = 0; i < WMMA_CHAIN; i++)
c = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
}
// Second snapshot. Feeding c[0] in as an input forces the last WMMA to retire
// before the counter is read.
asm volatile("s_getreg_b32 %0, hwreg(HW_REG_SHADER_CYCLES, 0, 32)"
: "=s"(stop) : "v"(c[0]) : "memory");
unsigned elapsed = (stop - start) & 0xfffff; // counter is <32 bits; mask the low 20
if (threadIdx.x == 0) {
result[blockIdx.x] = c[0];
cycles[blockIdx.x] = elapsed;
}
}
// =============================================================================
// Kernel B: identical dependent WMMA chain, hand-written in inline asm with the
// operands placed on three DISTINCT VGPR banks:
// A : v[0:7] bank 0 (0 % 4)
// B : v[9:16] bank 1 (9 % 4) gap at v8
// C : v[18:25] bank 2 (18 % 4) gap at v17 (holds threadIdx.x)
// This is the layout the compiler *should* produce; it runs at 32 cyc/WMMA.
// =============================================================================
__global__ __launch_bounds__(32, 8)
void wmma_asm_bank_distinct(float* __restrict__ result, int iterations,
unsigned* __restrict__ cycles) {
(void)result; (void)iterations; (void)cycles;
asm volatile(
// kernarg: result @0x0 (8B), iterations @0x8 (4B), cycles @0x10 (8B)
"s_load_b32 s3, s[0:1], 0x8\n\t"
"s_waitcnt lgkmcnt(0)\n\t"
"s_cmp_lt_i32 s3, 1\n\t"
"s_cbranch_scc1 2f\n\t"
"v_mov_b32 v17, v0\n\t" // save threadIdx.x out of the way (gap between B and C)
// A = v[0:7] = packed f16 1.0
"v_mov_b32 v0, 0x3c003c00\n\t" "v_mov_b32 v1, 0x3c003c00\n\t"
"v_mov_b32 v2, 0x3c003c00\n\t" "v_mov_b32 v3, 0x3c003c00\n\t"
"v_mov_b32 v4, 0x3c003c00\n\t" "v_mov_b32 v5, 0x3c003c00\n\t"
"v_mov_b32 v6, 0x3c003c00\n\t" "v_mov_b32 v7, 0x3c003c00\n\t"
"v_mov_b32 v8, 0\n\t" // gap (bank 0) / store offset
// B = v[9:16] = packed f16 1.0
"v_mov_b32 v9, 0x3c003c00\n\t" "v_mov_b32 v10, 0x3c003c00\n\t"
"v_mov_b32 v11, 0x3c003c00\n\t" "v_mov_b32 v12, 0x3c003c00\n\t"
"v_mov_b32 v13, 0x3c003c00\n\t" "v_mov_b32 v14, 0x3c003c00\n\t"
"v_mov_b32 v15, 0x3c003c00\n\t" "v_mov_b32 v16, 0x3c003c00\n\t"
// C = v[18:25] = 0
"v_mov_b32 v18, 0\n\t" "v_mov_b32 v19, 0\n\t" "v_mov_b32 v20, 0\n\t" "v_mov_b32 v21, 0\n\t"
"v_mov_b32 v22, 0\n\t" "v_mov_b32 v23, 0\n\t" "v_mov_b32 v24, 0\n\t" "v_mov_b32 v25, 0\n\t"
"s_getreg_b32 s7, hwreg(HW_REG_SHADER_CYCLES, 0, 32)\n\t"
"1:\n\t"
"s_add_i32 s3, s3, -1\n\t"
"s_cmp_lg_u32 s3, 0\n\t"
".rept 256\n\t"
"v_wmma_f32_16x16x16_f16 v[18:25], v[0:7], v[9:16], v[18:25]\n\t"
".endr\n\t"
"s_cbranch_scc1 1b\n\t"
"s_getreg_b32 s8, hwreg(HW_REG_SHADER_CYCLES, 0, 32)\n\t"
"s_sub_u32 s7, s8, s7\n\t"
"s_and_b32 s7, s7, 0xfffff\n\t"
"s_load_b64 s[4:5], s[0:1], 0x0\n\t"
"s_load_b64 s[10:11], s[0:1], 0x10\n\t"
"s_lshl_b32 s6, s2, 2\n\t"
"s_waitcnt lgkmcnt(0)\n\t"
"s_add_u32 s4, s4, s6\n\t"
"s_addc_u32 s5, s5, 0\n\t"
"s_add_u32 s10, s10, s6\n\t"
"s_addc_u32 s11, s11, 0\n\t"
"s_mov_b32 s6, exec_lo\n\t"
"v_cmpx_eq_u32_e32 0, v17\n\t"
"global_store_b32 v8, v18, s[4:5]\n\t"
"v_mov_b32 v26, s7\n\t"
"global_store_b32 v8, v26, s[10:11]\n\t"
"s_mov_b32 exec_lo, s6\n\t"
"2:\n\t"
"s_endpgm\n\t"
:
:
: "v0","v1","v2","v3","v4","v5","v6","v7","v8",
"v9","v10","v11","v12","v13","v14","v15","v16","v17",
"v18","v19","v20","v21","v22","v23","v24","v25","v26",
"s3","s4","s5","s6","s7","s8","s10","s11","vcc_lo","exec_lo","memory"
);
__builtin_unreachable();
}
// -----------------------------------------------------------------------------
// Host side
// -----------------------------------------------------------------------------
int main() {
int dev = 0;
hipDeviceProp_t prop;
HIP_CHECK(hipGetDeviceProperties(&prop, dev));
printf("GPU: %s (%s)\n\n", prop.name, prop.gcnArchName);
// ---- Time the dependent WMMA chains ------------------------------------
printf("=================================================================\n");
printf(" Back-to-back dependent WMMA latency (1 wave, 32 threads)\n");
printf(" chain = %d WMMAs/inner iter; cyc/WMMA = elapsed / (iters * %d)\n",
WMMA_CHAIN, WMMA_CHAIN);
printf("=================================================================\n");
const int iterations = 64; // 64 * 256 = 16384 WMMAs; counter stays < 2^20
float* d_result;
unsigned* d_cyc;
half16* d_a;
half16* d_b;
HIP_CHECK(hipMalloc(&d_result, sizeof(float)));
HIP_CHECK(hipMalloc(&d_cyc, sizeof(unsigned)));
HIP_CHECK(hipMalloc(&d_a, 32 * sizeof(half16)));
HIP_CHECK(hipMalloc(&d_b, 32 * sizeof(half16)));
// Fill A/B with f16 1.0 on the host (values are irrelevant to timing).
std::vector<half16> ha(32), hb(32);
for (int t = 0; t < 32; t++)
for (int i = 0; i < 16; i++) { ha[t][i] = (__fp16)1.0f; hb[t][i] = (__fp16)1.0f; }
HIP_CHECK(hipMemcpy(d_a, ha.data(), 32 * sizeof(half16), hipMemcpyHostToDevice));
HIP_CHECK(hipMemcpy(d_b, hb.data(), 32 * sizeof(half16), hipMemcpyHostToDevice));
const double wmmas = (double)iterations * WMMA_CHAIN;
// Take the minimum over several launches to denoise the ~2-clock counter
// granularity.
auto run_min = [&](const char* label, int reps, auto launch) {
unsigned best = ~0u;
launch(); // warmup
HIP_CHECK(hipDeviceSynchronize());
for (int r = 0; r < reps; r++) {
launch();
HIP_CHECK(hipDeviceSynchronize());
unsigned h = 0;
HIP_CHECK(hipMemcpy(&h, d_cyc, sizeof(unsigned), hipMemcpyDeviceToHost));
if (h < best) best = h;
}
float h_res = 0.0f;
HIP_CHECK(hipMemcpy(&h_res, d_result, sizeof(float), hipMemcpyDeviceToHost));
printf(" %-28s %9u cycles %6.2f cyc/WMMA (result=%.2e)\n",
label, best, best / wmmas, h_res);
return best / wmmas;
};
double cyc_intrin = run_min("intrinsic (compiler regalloc)", 20, [&]{
wmma_intrinsic<<<1, 32>>>(d_result, d_a, d_b, iterations, d_cyc);
});
double cyc_asm = run_min("i<truncated>Please see the issue for the entire body.
![]()
_______________________________________________
llvm-bugs mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-bugs
