[Bugfix] Sync both wave groups before the gfx950 A16W16 HTI GEMM reads its first and last stages

[jhinpan]

Summary

In the half-tile-interleaved (HTI) path of kernels/gemm/gemm_a16w16_gfx950.py, the first wave group can read LDS that the second group wrote by LDS DMA before the second group has waited for those copies. This PR makes the second group retire those copies before the barrier that releases the read, in the prologue and in the tail. It adds no barrier.

Motivation

The ping-pong schedule runs the second wave group (wid // n_waves == 1) one barrier ahead (line numbers at main):

async_load_b_to_lds(0, 0, 0)      # L1023-1026: every wave copies its part of stage 0
...
if wid // n_waves == 1:           # L1028-1029: second group only
    rocdl.s_barrier()
rocdl.s_barrier()                 # L1031
...
wait_vmcnt_and_barrier(half_ldg_b_iters + half_ldg_a_iters)  # L1037: s_waitcnt vmcnt, then s_barrier
...
b0 = load_b_fragment(0, 0)        # L1043-1044: the first iteration reads stage 0
a0 = load_a_fragment(0, 0)

An s_barrier completes when every wave has arrived at a barrier, whichever one. So the barrier in the first group's wait_vmcnt_and_barrier (L1037) completes together with the second group's s_barrier() at L1031, before the second group runs its own s_waitcnt vmcnt. The first group then reads stage 0, part of which the second group wrote with buffer_load ... lds. vmcnt counts only the issuing wave's copies, and LLVM inserts no wait before s_barrier on gfx950 (it does on gfx908, which lacks back-off barriers), so nothing orders those copies before the read.

The tail has the same order for the stage 1 A half 0: wait_vmcnt_and_barrier(0) (L1106), then load_a_fragment(0, 1) (L1108). The split-K tail starts with wait_vmcnt_and_barrier(0) (L1088), so split-K is exposed only in the prologue. In the main loop, one more s_barrier follows each wait_vmcnt_and_barrier before the data is read (L1063, L1083), and that absorbs the offset.

How we found it: a sync sanitizer prototype (#1237) reports "LDS read before an async write to it is retired and visible" at exactly these three reads, on MI355X at 4096^3 bf16 with the upstream tuned HTI policy. In the 4,096 reports it kept in full (of 327,680), the reader is in waves 2-3 and the writer in waves 4-7. The output check passes with and without a fix: in both places the second group issues the copies before the first group issues the loads that the first group then waits for, so a wrong result needs the earlier copies to complete after those later loads. We have not seen a wrong result on hardware.

Changes

• Prologue: after its extra s_barrier, the second group runs s_waitcnt vmcnt(0), so its stage 0 copies retire before its barrier at L1031, which pairs with the first group's barrier at L1037.
• Tail: before the barrier at L1104, which pairs with the first group's barrier at L1106, the second group runs s_waitcnt vmcnt(half_ldg_b_iters + half_ldg_a_iters). This retires its stage 1 A half 0 copy; the newer B half 1 and A half 1 copies may stay in flight. In the split-K tail the wait finds nothing to wait for.
• The first revision of this PR added an s_barrier for all waves after L1037 and after L1106 instead. That also closes the race, but the first group then waits for the second group's own wait at L1037 before it reads, and it costs 0.8-3.5% on the split-K policies (table below). Keeping that tail barrier out of split-K did not change this (1.006-1.036): the prologue barrier is the cost.
• The main loop and the non-HTI path do not change.

Performance

MI355X, one GPU. The rows are the HTI rows of test_gemm_a16w16_acc_tuned_policies (which uses a bias) and the "Common PR-tuned" and "Exact HTI split-K winners" rows of test_gemm_a16w16_acc_split_k, in bf16 with layout nt.

Test list	M×N×K	Config	Baseline (µs)	This PR / baseline	First revision (two barriers) / baseline
tuned	4096×4096×4096	256×256×64, GM 4, bias	102.94	0.999	1.006
tuned	4096×4096×8192	256×256×64, GM 4, bias	196.30	0.999	1.002
tuned	8192×8192×8192	256×256×64, GM 4, bias	744.18	0.999	1.002
tuned	8192×8192×8192	256×256×64, bias	744.50	0.999	1.001
split-K, PR-tuned	65×96×4096	64×64×256, SK 4	9.24	0.996	1.035
split-K, PR-tuned	96×64×7168	64×64×256, SK 7	9.80	0.996	1.033
split-K, PR-tuned	128×64×4096	64×64×256, SK 8	7.48	0.995	1.021
split-K, PR-tuned	128×128×6144	64×64×256, SK 6	10.00	0.996	1.032
split-K, PR-tuned	65×136×2048	64×128×128, SK 4	8.20	0.990	1.020
split-K, PR-tuned	129×72×1024	128×64×128, SK 2, bias	11.40	1.007	1.028
split-K, PR-tuned	129×136×1536	128×128×128, SK 3, GM 4	10.08	0.996	1.024
split-K, PR-tuned	257×264×512	256×256×64, SK 2	19.88	0.998	1.020
split-K, exact winner	4096×128×6144	64×64×256, SK 2	24.00	0.998	1.013
split-K, exact winner	512×2560×6144	128×128×128, SK 3, GM 4	35.76	0.999	1.008
split-K, exact winner	576×256×4096	64×64×256, SK 4, bias	11.80	1.000	1.031
split-K, exact winner	1024×128×6144	64×64×256, SK 6, bias	12.92	0.997	1.025
split-K, exact winner	1216×64×7168	64×64×256, SK 7	12.04	1.003	1.030
split-K, exact winner	1280×64×4096	64×64×256, SK 8	9.84	1.000	1.012

Geometric mean, this PR / baseline: 0.998 on the 14 split-K rows and 0.999 on the 4 tuned rows. First revision: 1.024 and 1.003.

Method: FlyDSL 0.4.0 wheel from PyPI with kernels/ at v0.4.0, whose HTI kernel is identical to this PR's base, and torch 2.10.0+rocm7.2.4. One process loads every variant and interleaves them: call i runs variant i mod 4 on input set i mod N, with N odd and the sets rotated over 2 GiB, so all variants see the same clock and cache state. Each row takes 500 profiled calls per variant; the table shows the median kernel duration, then the median of three runs. An A/A control with the baseline in all four slots stayed within 0.7% on the split-K rows and 0.1% on the tuned rows. Separate runs per variant varied by up to 6% on the 7-10 µs split-K rows, which is why the variants are interleaved.

Testing

• Unit tests added/updated: a test cannot force the wave order that exposes the race.
• Performance benchmarks run: on MI355X, see Performance.
• Tested on MI300X: this kernel targets gfx950.

On MI355X, tests/kernels/test_gemm_a16w16_gfx950.py without the benchmark-marked cases passes with this change: 879 passed, 48 skipped.

The sanitizer prototype was not rerun on this revision. Instead, a model of the schedule checks the same rule: it executes the HTI section of the kernel source once per wave group, pairs barriers by count, and retires each wave's copies in order at its vmcnt waits. For split-K and non-split-K with 2 to 8 K tiles, it reports exactly the three reads above on main and nothing with this change. It reports them again when either new wait is removed or its count is raised by one. In the gfx950 ISA, each new s_waitcnt sits in the second group's branch before the s_barrier it must precede.

Dependencies

• No new third-party dependencies added

Breaking Changes

None.

[jhinpan]

Converting to draft until the MI355X timing table is in the description.