Skip to content

fix(hipblaslt): wait for complex B operands before MFMA - #14330

Merged
newling merged 2 commits into
developfrom
users/newling/fix-gfx942-complex-lds-wait
Oct 28, 2026
Merged

newling merged 2 commits into
developfrom
users/newling/fix-gfx942-complex-lds-wait

Conversation

@newling

@newling newling commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

JIRA ID : AIHPBLAS-3605

Motivation

TensileLite generates matrix-multiplication kernels for hipBLASLt. In a gfx942 complex-float case with M224, N128 and K128, arithmetic consumes an operand before its shared-memory load is guaranteed to finish. rocjitsu, a GPU instruction emulator with race checking, reports three races even though numerical validation passes. The same scheduling defect also reproduces on gfx950.

This changes wait instructions inside the repeated multiplication loop. Arithmetic on the first B operand can still overlap later shared-memory loads, but moving the second wait earlier may reduce that overlap. GPU throughput for this wait placement has not been measured.

Technical Details

The scheduler waits for the first B operand and lets later B loads remain outstanding. It then inserts a second wait after processing the first A tile. Complex multiplication advances through B first and emits four matrix fused multiply-add instructions per complex product, reaching the next B operand before that wait.

Select scheduler index 4 for the second wait in complex kernels. Negation instructions also occupy scheduler slots, making this placement conservative across all conjugation choices. Keep the partial initial wait and the existing real-type index calculation. A/B load order is unchanged.

For the reported gfx942 kernel, the initial wait remains lgkmcnt(3), which allows three shared-memory operations to remain outstanding. The second wait moves before the next B operand and changes from lgkmcnt(11) to lgkmcnt(8). Assembly changes only wait instructions; register and shared-memory allocations are identical.

Generated gfx942 before assembly, after assembly and complete diff show the wait relocation. The gfx950 before, after and diff bundles cover two complex-float and two complex-double configurations, with source hashes and static comparisons.

The generation regression covers complex float and double, gfx942 and gfx950, and all four A/B conjugation combinations. It checks that partial waiting is preserved and that the remaining B loads complete before their operands are consumed.

Test Plan

Run the regression and adjacent generation/wait tests. Generate and assemble both complex precisions on both architectures, then replay the reported gfx942 workload and four retained gfx950 configurations with race checking and numerical validation. Include shorter K values and a gfx942 tail case. Run shared architecture checks and measure GPU throughput before landing.

Test Result

All 16 regression cases pass without skips; six fail with the original scheduler. Eleven adjacent tests and three unchanged assembly snapshots pass. Both complex precisions generate and assemble for gfx942 and gfx950, including additional tile and load-layout configurations.

The gfx942 reproducer changes from three races and 17 warnings to zero. Four gfx950 cases, including complex double, change from 1–3 races each to zero. Every replay confirms the expected kernel/solution, one target dispatch and norm error zero. The gfx942 candidate also passes K32, K48, K136 with a tail, and K256.

The complete race-detection report contains all ten before/after replay logs, code-object hashes, dispatch checks and numerical results.

Replays use regenerated instructions with symbols adapted to retained dispatch metadata. Root hooks, new-test formatting and whitespace checks pass. Shared automated checks for this revision and physical GPU performance measurements remain pending; earlier MI300X timings measured a different wait placement.

Submission Checklist

Risk level

Medium: the changed waits execute inside the repeated multiplication loop and can affect throughput. The candidate preserves partial waiting and resource allocations, while regression and replay results cover the reproduced operand-order defect. GPU performance and integration with current develop remain unqualified.

The partial B-read wait assumes the real-valued operand order. Complex multiplication consumes later B operands before the extra wait, allowing an LDS result to remain pending at its first use.

Exclude complex types from both parts of the optimization and add gfx942 code-generation regressions for all A/B conjugation combinations.
@therock-pr-bot

therock-pr-bot Bot commented Oct 23, 2026 •

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

Check Status Details
📝 PR Description ✅ Pass —
⛔ Forbidden Files ✅ Pass —
🧪 Unit Test ✅ Pass —
🔎 pre-commit ✅ Pass —
🚫 Draft PR 🔜 To Be Enabled —
🚩 Feature Flag 🔜 To Be Enabled —
📊 Code Coverage 🔜 To Be Enabled —
🤖 therock-pr-bot ✅ Pass —

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

@therock-pr-bot

therock-pr-bot Bot commented Oct 23, 2026 •

Copy link
Copy Markdown

🎉 All checks passed! This PR is ready for review.

@codecov

codecov Bot commented Oct 24, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

❌ Your project status has failed because the head coverage (50.81%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files
@@             Coverage Diff             @@
##           develop   #14330      +/-   ##
===========================================
+ Coverage    75.79%   76.10%   +0.31%     
===========================================
  Files         2880     2886       +6     
  Lines       478823   484467    +5644     
  Branches     71949    73115    +1166     
===========================================
+ Hits        362888   368663    +5775     
+ Misses       91072    90457     -615     
- Partials     24863    25347     +484     
Flag Coverage Δ *Carryforward flag
TensileLite-CPP 50.81% <ø> (+1.87%) ⬆️
TensileLite-Unit 84.06% <100.00%> (+0.38%) ⬆️
hipBLAS 90.62% <ø> (ø) Carriedforward from fe92436
hipBLASLt 68.99% <ø> (+4.67%) ⬆️
hipCUB 82.68% <ø> (ø) Carriedforward from fe92436
hipDNN 86.85% <ø> (ø) Carriedforward from fe92436
hipFFT 42.56% <ø> (ø) Carriedforward from fe92436
hipRAND 76.12% <ø> (ø) Carriedforward from fe92436
hipSOLVER 69.67% <ø> (ø) Carriedforward from fe92436
hipSPARSE 86.99% <ø> (ø) Carriedforward from fe92436
rocBLAS 68.73% <ø> (ø) Carriedforward from fe92436
rocFFT 45.68% <ø> (ø) Carriedforward from fe92436
rocRAND 57.31% <ø> (ø) Carriedforward from fe92436
rocSOLVER 76.74% <ø> (ø) Carriedforward from fe92436
rocSPARSE 76.86% <ø> (ø) Carriedforward from fe92436
rocThrust 93.15% <ø> (ø) Carriedforward from fe92436

*This pull request uses carry forward flags. Click here to find out more.

Files with missing lines Coverage Δ
.../hipblaslt/tensilelite/tensilelite/KernelWriter.py 84.89% <100.00%> (+0.30%) ⬆️

... and 73 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@newling
newling marked this pull request as ready for review October 25, 2026 20:34
@newling
newling requested a review from a team as a code owner October 25, 2026 20:34
@newling
newling requested review from hcman2 and nakajee October 27, 2026 01:37
Move the second complex B wait to scheduler index 4, before a later B operand is consumed. Negation instructions can only delay that use, so this preserves partial waiting without conjugation bookkeeping.

Extend the generation regression to complex float and double on gfx942 and gfx950. Regenerated controls remove the reproduced races on both architectures while retaining numerical results and resource allocations.

@nakajee nakajee left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your update.
Looks good.

@newling
newling enabled auto-merge (squash) October 28, 2026 03:16
@newling
newling merged commit 7570194 into develop Oct 28, 2026
165 of 171 checks passed
@newling
newling deleted the users/newling/fix-gfx942-complex-lds-wait branch October 28, 2026 06:15
@assistant-librarian

Copy link
Copy Markdown
Contributor

TheRock Submodule Bump Activity

Newest first

newling added a commit that referenced this pull request Oct 29, 2026
JIRA ID : AIHPBLAS-5076

## Motivation

hipBLASLt ships three hand-written gfx950 assembly kernels for FP32
GEMMs that use `F32XdlMathOp: X` (XF32):
`Custom_Cijk_{Ailk_Bljk,Ailk_Bjlk,Alik_Bljk}_S_MX_B_BIAS_HA_S_SAV_NTD_SK3_UserArgs_MT256x256x32_MI16x16x1_shortname0_gfx950`.
They split each FP32 input into BF16 parts and feed BF16 matrix
instructions (MFMA). Unlike TensileLite's generated kernels, their
instruction schedule, including every wait, is written by hand in
`tensilelite/tensilelite/CustomKernels/tensile/`.

The defect scan prompted by
[AIHPBLAS-5076](https://amd-hub.atlassian.net/browse/AIHPBLAS-5076)
included a gfx950 race-detector sweep of hipBLASLt solution indices
under the rocjitsu GPU emulator. It found that the TN kernel
(`Alik_Bljk`) reads registers before the local data share (LDS) loads
that fill them are guaranteed to complete. Both shipped solution indices
that select this kernel report 24 register races and memory-wait
warnings at M=N=K=256, while numerical validation still passes. The NN
and NT siblings run clean. This PR moves one existing wait so the TN
kernel waits for its B operand before converting it, as the siblings
already do.

## Technical Details

Each main-loop iteration reads FP32 tiles from LDS into temporary
registers with `ds_read_b128`, converts them to BF16 with
`v_cvt_pk_bf16_f32`, and computes the low-order remainder with
`v_dot2c_f32_bf16`, which rewrites the same temporaries. LDS reads
complete asynchronously and are tracked by the `lgkmcnt` counter, so
every instruction that reads or overwrites those temporaries must first
wait until the read has returned.

**BEFORE — the B conversions precede their wait.** Near the end of the
`MAINLOOP` macro, the kernel issues the eight next-iteration B reads,
`LR B 0` through `LR B 7`, into `vgprValuB_T0_I0` (`v192`–`v223`). The
last wait before them, `s_waitcnt lgkmcnt(5) // Wait for all LR A0 to
finish`, retires only the A reads. After about 65 more instructions, the
macro converts all 32 B temporaries, rewrites them with 32 `v_dot2c`
instructions and issues another MFMA. Only then does it reach the wait
meant for those reads:

```asm
LR B 7 \usePLR            // eighth ds_read_b128 into vgprValuB_T0_I0+28..31
...                       // A conversions and 21-22 MFMAs, no LDS wait
v_cvt_pk_bf16_f32 v[vgprValuB_X0_I0+0+0], v[vgprValuB_T0_I0+0+0], v[vgprValuB_T0_I0+0+1]
...                       // 15 more conversions, 32 v_dot2c, 4 more conversions
v_mfma_f32_16x16x32_bf16 ...   // mfmaIndex:189
s_waitcnt lgkmcnt(0) // Wait for all LR B0 to finish
```

A conversion can therefore read a temporary whose LDS result has not
arrived, and `v_dot2c` can overwrite a temporary that a pending read
will still write. The detector reports the first register of each of the
eight reads, `v192` through `v220`, on both SIMD-specific loop paths.
Its warnings require `lgkmcnt` at most 7 for the first read and at most
4 for the fourth: the number of younger reads that may still be pending.

**AFTER — wait once, before the first reader.** The NN and NT kernels
place `s_waitcnt lgkmcnt(0) // Wait for LR B0` immediately before their
B conversions. This PR moves the TN kernel's existing wait to the same
place, ahead of the conversion block. No LDS operation is issued between
the old and new positions, so the later conversions that relied on the
old position remain covered. `MAINLOOP` is expanded four times: the
even- and odd-SIMD main loops, the last full iteration without global
loads, and the single-iteration path, where the conversions are compiled
out. Each expansion changes by that one moved instruction. Instruction
count (26,228), registers, LDS allocation and every other wait are
unchanged.

The complete hand-written source and kernel disassembly are in the
[before](https://gist.github.com/9046d6ae97b74319d395c327418d4724) and
[after](https://gist.github.com/77af1eef91145f10c7fddbe8ef4fdbdc) gists.
The [diff
gist](https://gist.github.com/c77390cdf533af011d691f4052c3affb) holds
the source diff and a machine-instruction diff showing exactly the four
moved waits. Use **Raw** to download the complete large files.

**Performance implications.** After the last B read, the moved wait
still follows the same 64–66 instructions, including 21–22 MFMAs. That
independent work provides an opportunity to overlap the LDS reads with
computation. If the wait stalls, it delays the B conversions only until
the data they consume has arrived. The previous placement saved that
time only by reading data the hardware did not guarantee to be present.
The loop keeps its partial A waits, global-load waits, barriers and
instruction order. Physical GPU timing was not measured, and emulator
wall times are not performance evidence.

This differs from the generator fixes
[#14147](#14147) and
[#14330](#14330), which
change TensileLite's scheduler and cannot reach hand-written custom
kernels.

## Test Plan

Follow the [reproduction
recipe](https://gist.github.com/12f1e2186fdcb357d6095e1ca5c72ac7) to
generate baseline and repaired kernels from the same one-solution gfx950
logic entry with `TensileCreateLibrary`, then run them through
`tensilelite-client` under rocjitsu's race plugin with full-output
numerical validation. The recipe specifies the tool versions, uses
portable path variables, selects reduced-library solution 0, and checks
the expected kernel dispatch, races and wait warnings as well as
numerical validation. Cover a single iteration, short and long K, a K
tail, multiple Stream-K tiles, M/N edge tiles and batch 2. Confirm that
the regenerated baseline matches the kernel in which the race was first
found.

## Test Result

The regenerated `develop` kernel is instruction-identical to the
CI-library kernel that produced the original report. It reproduces that
report's signature: 24 VGPR read races on `v192`–`v220` across waves 0
and 1, with required thresholds of at most 4 to 7.

| Shapes (M×N×K, batch) | Before: races / warnings | After: races /
warnings |
| --- | --- | --- |
| 256×256×32 (one iteration) | 0 / 0 | 0 / 0 |
| 256×256×64 | 8 / 17 | 0 / 0 |
| 256×256×96, ×256, ×300 | 24 / 17 each | 0 / 0 |
| 256×256×4096 | 24 / 272 | 0 / 0 |
| 512×512×512 | 24 / 136 | 0 / 0 |
| 768×512×1024 | 24 / 408 | 0 / 0 |
| 300×260×200 | 24 / 68 | 0 / 0 |
| 256×256×256, batch 2 | 24 / 34 | 0 / 0 |

All 20 runs exit 0, pass numerical validation and dispatch this kernel
exactly once. The detector deduplicates races, so the count is not the
number of affected lanes. Functional emulation passing numerically does
not show whether the race corrupts results on hardware; the tracked
pending reads establish the missing dependency. The [race evidence
gist](https://gist.github.com/052449eb88bb3555dadb3360e7d8e001) holds
all 20 logs, the original CI report and the matrix summary. It also
holds the one-solution logic entry, client configuration, detector
profile, and source/tool/code-object hashes. Device code was assembled
for gfx950 with ROCm 7.1 amdclang (Clang 20), code-object version 5.
rocjitsu is at
[99b05c7d01c](ROCm/rocm-systems@99b05c7).

The defect scan and targeted baseline/repaired replays provide the race
evidence for this fix. A dedicated race-check job will follow later. Its
behavioral check will select the affected TN XF32 kernel and reject race
records and memory-wait warnings, since the hazardous baseline already
passes numerical validation.

The required Component CI, Math CI, Multi-Arch CI and pre-commit checks
pass. The existing gfx950 race-check job also passes, but its general
workloads do not select this kernel. The gfx1250 FFM and two aggregate
coverage checks remain red; those statuses are separate from the
targeted baseline/repaired results above.

**Race log, BEFORE.** The first complete record from the original
report, solution index 77224 at M=N=K=256. The first arrow marks the
pending `ds_read_b128`; the second marks the conversion that reads its
destination. The detector printed the elision marker itself. The other
23 records cover `v196` through `v220` and the second loop path.

```text
RACE kernel=Custom_Cijk_Alik_Bljk_S_MX_B_BIAS_HA_S_SAV_NTD_SK3_UserArgs_MT256x256x32_MI16x16x1_shortname0_gfx950 symbol=Custom_Cijk_Alik_Bljk_S_MX_B_BIAS_HA_S_SAV_NTD_SK3_UserArgs_MT256x256x32_MI16x16x1_shortname0_gfx950 dispatch=3 type=VGPR access=read reg=192 wave=0 lane=0 wg=0,0,0 conflict=unknown
Race on VGPR v192 [workgroup (0, 0, 0), wave 0, lane 0]
  ==>  0x719a87f16b34  ds_read_b128 v[192:195], v17  ; <-- wave 0
       0x719a87f16b3c  v_sub_f32_e32 v165, v165, v155
       0x719a87f16b40  v_mfma_f32_16x16x32_bf16 acc[236:239], v[142:145], v[42:45], acc[236:239]
       0x719a87f16b48  v_cvt_f32_bf16_e32 v156, v21
       0x719a87f16b4c  v_sub_f32_e32 v166, v166, v156
       0x719a87f16b50  v_mfma_f32_16x16x32_bf16 acc[236:239], v[138:141], v[42:45], acc[236:239]
       0x719a87f16b58  v_cvt_f32_bf16_sdwa v157, v21 dst_sel:DWORD dst_unused:UNUSED_PRESERVE src0_sel:WORD_1
       0x719a87f16b60  v_sub_f32_e32 v167, v167, v157
       0x719a87f16b64  v_mfma_f32_16x16x32_bf16 acc[144:147], v[114:117], v[54:57], acc[144:147]
       0x719a87f16b6c  v_cvt_pk_bf16_f32 v22, v160, v161
       0x719a87f16b74  v_cvt_pk_bf16_f32 v23, v162, v163
       0x719a87f16b7c  v_mfma_f32_16x16x32_bf16 acc[144:147], v[118:121], v[50:53], acc[144:147]
       0x719a87f16b84  v_cvt_pk_bf16_f32 v24, v164, v165
       0x719a87f16b8c  v_cvt_pk_bf16_f32 v25, v166, v167
       0x719a87f16b94  v_mfma_f32_16x16x32_bf16 acc[144:147], v[114:117], v[50:53], acc[144:147]
       0x719a87f16b9c  s_waitcnt vmcnt(63) expcnt(7) lgkmcnt(5)
       ... 112 instructions elided ...
       0x719a87f16e98  v_mfma_f32_16x16x32_bf16 acc[244:247], v[142:145], v[58:61], acc[244:247]
       0x719a87f16ea0  v_cvt_f32_bf16_sdwa v155, v44 dst_sel:DWORD dst_unused:UNUSED_PRESERVE src0_sel:WORD_1
       0x719a87f16ea8  v_sub_f32_e32 v189, v189, v155
       0x719a87f16eac  v_mfma_f32_16x16x32_bf16 acc[244:247], v[138:141], v[58:61], acc[244:247]
       0x719a87f16eb4  v_cvt_f32_bf16_e32 v156, v45
       0x719a87f16eb8  v_sub_f32_e32 v190, v190, v156
       0x719a87f16ebc  v_mfma_f32_16x16x32_bf16 acc[248:251], v[138:141], v[70:73], acc[248:251]
       0x719a87f16ec4  v_cvt_f32_bf16_sdwa v157, v45 dst_sel:DWORD dst_unused:UNUSED_PRESERVE src0_sel:WORD_1
       0x719a87f16ecc  v_sub_f32_e32 v191, v191, v157
       0x719a87f16ed0  v_mfma_f32_16x16x32_bf16 acc[248:251], v[142:145], v[66:69], acc[248:251]
       0x719a87f16ed8  v_cvt_pk_bf16_f32 v46, v184, v185
       0x719a87f16ee0  v_cvt_pk_bf16_f32 v47, v186, v187
       0x719a87f16ee8  v_mfma_f32_16x16x32_bf16 acc[248:251], v[138:141], v[66:69], acc[248:251]
       0x719a87f16ef0  s_sub_u32 s8, s8, 1
       0x719a87f16ef4  v_cvt_pk_bf16_f32 v49, v190, v191
  ==>  0x719a87f16efc  v_cvt_pk_bf16_f32 v82, v192, v193  ; <-- wave 0 lane 0
END_RACE
```

**Race log, AFTER.** Every shape above, including M=N=K=256, completes
full-output validation with no race records and no memory-wait warnings.

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

## Risk level

Low (2/5): one wait moves earlier in a single shipped hand-written
gfx950 kernel. That can add a stall only when B data is still in flight,
and it cannot remove a dependency. Race-detector and numerical replays
cover the affected loop paths and shapes. Physical GPU timing is
unmeasured, and a dedicated shared-CI race regression will follow later.


[AIHPBLAS-5076]:
https://amd-hub.atlassian.net/browse/AIHPBLAS-5076?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants