Repository navigation
fix(tensilelite): wait for B LDS reads TN XF32 custom kernel - #14533
Merged
Merged
Conversation
…TN XF32 custom kernel The gfx950 custom kernel Custom_Cijk_Alik_Bljk_S_MX_B_..._MT256x256x32 emulates XF32 with BF16 MFMAs. At the end of its MAINLOOP macro it issues the eight next-iteration B LDS reads, then converts those FP32 values to BF16 (and rewrites the temporaries with v_dot2c) before its s_waitcnt lgkmcnt(0). The conversions can therefore read registers whose ds_read_b128 results have not returned. Move the wait ahead of the conversion block, matching the NN and NT sibling kernels.
github-actions
Bot
added
project: hipblaslt
project: hipsparselt
ci:hipsparselt-fast
labels
⏳ Waiting on Required Checks⏳ Required checks have not all concluded yet. This bot run is exiting without failing so those checks can report independently.
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
⏳ Waiting for required checks to start or finish:
|
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project check 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 #14533 +/- ##
========================================
Coverage 77.29% 77.29%
========================================
Files 2889 2889
Lines 484711 484711
Branches 73195 73195
========================================
+ Hits 374634 374635 +1
Misses 85355 85355
+ Partials 24722 24721 -1
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
newling
marked this pull request as ready for review
October 28, 2026 17:37
newling
changed the title
nakajee
approved these changes
Oct 29, 2026
nakajee
left a comment
nakajee
left a comment
Contributor
There was a problem hiding this comment.
Thanks for your fix.
The change looks reasonable.
Contributor
TheRock Submodule Bump ActivityNewest first
|
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-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 intensilelite/tensilelite/CustomKernels/tensile/.The defect scan prompted by 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 withv_cvt_pk_bf16_f32, and computes the low-order remainder withv_dot2c_f32_bf16, which rewrites the same temporaries. LDS reads complete asynchronously and are tracked by thelgkmcntcounter, 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
MAINLOOPmacro, the kernel issues the eight next-iteration B reads,LR B 0throughLR B 7, intovgprValuB_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 32v_dot2cinstructions and issues another MFMA. Only then does it reach the wait meant for those reads:A conversion can therefore read a temporary whose LDS result has not arrived, and
v_dot2ccan overwrite a temporary that a pending read will still write. The detector reports the first register of each of the eight reads,v192throughv220, on both SIMD-specific loop paths. Its warnings requirelgkmcntat 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 B0immediately 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.MAINLOOPis 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 and after gists. The diff gist 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 and #14330, which change TensileLite's scheduler and cannot reach hand-written custom kernels.
Test Plan
Follow the reproduction recipe to generate baseline and repaired kernels from the same one-solution gfx950 logic entry with
TensileCreateLibrary, then run them throughtensilelite-clientunder 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
developkernel is instruction-identical to the CI-library kernel that produced the original report. It reproduces that report's signature: 24 VGPR read races onv192–v220across waves 0 and 1, with required thresholds of at most 4 to 7.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 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.
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 coverv196throughv220and the second loop path.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
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.