Repository navigation
fix(hipblaslt): wait for complex B operands before MFMA - #14330
Merged
Merged
Conversation
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.
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
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
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
nakajee
reviewed
Oct 27, 2026
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
approved these changes
Oct 27, 2026
nakajee
left a comment
Contributor
There was a problem hiding this comment.
Thanks for your update.
Looks good.
Contributor
TheRock Submodule Bump ActivityNewest first
|
1 task
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 fromlgkmcnt(11)tolgkmcnt(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.