fix(tensilelite): guard against unsafe BufferStore dispatch on oversized D - #12195
Open
tony-davis wants to merge 4 commits into
Open
tony-davis wants to merge 4 commits into
tony-davis wants to merge 4 commits into
Conversation
…zed D (ROCM-31016) Fix BufferStoreOffsetLimitCheck in ContractionProblemPredicates.hpp, which was meant to guard BufferStore=True solutions against outputs whose store byte offset can exceed the fixed-base post-loop SRD's addressable range (allocPostLoopSrd in KernelWriterAssembly.py never re-bases that SRD per workgroup), but had a formula bug: it capped the checked extent to min(MacroTile1, size[1]), which only bounds a single output tile's worth of offset. Any problem whose full column extent (stride[1] * size[1]) exceeds the addressable range while a single tile's worth does not is wrongly reported as supported, which is exactly the large-M, modest-N bf16 shape family in GitHub ROCm/hipBLASLt#2299 / ROCM-31016. On such shapes, affected kernels silently drop stores instead of erroring. Confirmed on real MI350X (gfx950) hardware: an MT16x16x256 BufferStore=True kernel silently dropped ~99.9996% of a 16 GiB bf16 D (~32K of ~8.59B elements actually written) at a shape the buggy predicate reported as supported. Changes: - Check the tensor's true worst-case reachable byte offset, stride[1] * size[1] (uncapped), instead of the tile-capped formula. - Use the 0xfffff000 threshold matching KernelWriterAssembly.py's actual BufferOOB sentinel, instead of the previous 2^32 constant (~4 KiB looser than the real hardware ceiling). - No kernel-generator or dispatch-layer changes: CompoundPredicates (Contractions.py) already auto-attaches this predicate to every BufferStore=True solution at library-export time, and it is already evaluated by the standard solution-selection path. Adds regression coverage: - 3 new unit tests in tensilelite/tests/Predicates_test.cpp exercising the fixed predicate directly (reproduces bug, ordinary shape, boundary case). - New client gtest, buffer_store_offset_guard_gtest.cpp, exercising the real dispatch path on gfx950 hardware against a small in-range shape and the 16 GiB out-of-range repro shape. JIRA ID : ROCM-31016 Co-authored-by: Cursor <cursoragent@cursor.com>
github-actions
Bot
added
project: hipblaslt
project: hipsparselt
ci:hipsparselt-fast
labels
✅ 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. |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The predicate may still allow the documented corruption case, and the regression test has unresolved cleanup, error-handling, and coverage gaps.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR updates the TensileLite BufferStore offset predicate and adds regression coverage for oversized output tensors.
Changes:
- Checks the full
Dextent against the0xfffff000limit. - Adds predicate unit tests.
- Adds and registers a gfx950 dispatch-validation test.
File summaries
| File | Review summary |
|---|---|
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp |
Adds predicate tests. Moderate finding: add coverage distinguishing the 0xfffff000 threshold from 2^32. Nit: correct the fail-before test-result description. |
projects/hipblaslt/tensilelite/include/Tensile/ContractionProblemPredicates.hpp |
Critical finding: the documented 2 GiB corruption case may remain eligible under the current threshold; reconcile the guard with the actual limit. |
projects/hipblaslt/clients/tests/src/CMakeLists.txt |
Registers the new client test. |
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp |
Moderate findings: fix partial-allocation and resource cleanup, check setup and verification statuses, add coverage for the 2.54 GiB case, and provide affected gfx942 coverage. |
Review details
Suppressed comments (6)
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:143
- An allocation failure returns the default outcome without freeing buffers already allocated (for example, A when B or D allocation fails) or indicating setup failure. The tests then interpret this as a missing algorithm or a failed safe dispatch, so a gfx950 host with insufficient free memory fails for an unrelated reason and leaks device memory. Clean up partial allocations and skip/report OOM separately.
if(hipMalloc(&A, (size_t)M * K * sizeof(*A)) != hipSuccess)
return outcome;
if(hipMalloc(&B, (size_t)K * N * sizeof(*B)) != hipSuccess)
return outcome;
if(hipMalloc(&D, (size_t)M * N * sizeof(*D)) != hipSuccess)
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:180
- These return values are ignored, and
prefis uninitialized untilhipblasLtMatmulPreferenceCreatesucceeds. If creation or attribute setup fails, the helper passes an invalid preference tohipblasLtMatmulAlgoGetHeuristicand can misreport setup failure as a dispatch result; check each status and propagate the failure or skip.
hipblasLtMatmulPreference_t pref;
hipblasLtMatmulPreferenceCreate(&pref);
hipblasLtMatmulPreferenceSetAttribute(
pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &wsize, sizeof(wsize));
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:265
- This device test only exercises a D larger than 16 GiB. It therefore cannot catch the linked issue's 2.54 GiB reproducer, which the current
0xfffff000predicate still accepts and may dispatch; the test can pass via the zero-algorithm branch while the original corruption remains. Add the 2 GiB-boundary shape (or an equivalent below-4 GiB case) and require rejection or a bit-exact result.
constexpr int64_t M = 33554560, N = 256, K = 48;
static_assert(M % 256 != 0, "matches the shape family from the PR's repro session");
static_assert((uint64_t)M * N * 2 > 0xfffff000ull,
"D's true byte extent must exceed the BufferOOB sentinel");
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:232
- This macro unconditionally skips gfx942, even though the production predicate is shared across architectures and this file documents that gfx942 uses the same
allocPostLoopSrdpath. The only dispatch-level regression test therefore provides no evidence for that affected architecture; add a gfx942 shared-CI execution or a separately scoped test before counting the defect as covered.
if(gpuArchFamily() != "gfx950") \
GTEST_SKIP() << "ROCM-31016 was only reproduced/validated on gfx950, " \
"not " \
<< gpuArchFamily(); \
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp:196
- The PR description says all three new predicate tests fail before the fix, but these sanity cases do not distinguish the formulas: the 1024x1024 case is below both limits, and the boundary case also passes the old
min(MacroTile1, N)check. Please correct the Test Result to report one fail-before regression and two pass-before controls.
TEST(Predicates, BufferStoreOffsetLimitCheck_OrdinaryProblem_StillAccepted)
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp:229
- This test only asserts a value below the new ceiling. It still passes if the implementation retains the old
2^32threshold, because the large-M case is also above2^32; none of these tests detects the threshold half of the production change. Add an exact or just-over-0xfffff000case that remains below2^32and assert rejection.
// Boundary check: a D just under the BufferOOB sentinel must still pass.
constexpr size_t n = 256;
constexpr size_t m = (0xfffff000ull / 2 / n) - 1; // comfortably under the line
- Files reviewed: 4/4 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
tony-davis
changed the title
tony-davis
marked this pull request as draft
September 16, 2026 17:46
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #12195 +/- ##
========================================
Coverage 70.29% 70.29%
========================================
Files 2812 2812
Lines 462769 462782 +13
Branches 68122 68125 +3
========================================
+ Hits 325271 325284 +13
Misses 113943 113943
Partials 23555 23555
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
…re-basing The prior comment claimed the post-loop store SRD base is never re-based per workgroup. That's incorrect: computeStoreSrdStart does re-base the SRD base per workgroup using safe 64-bit arithmetic. The actual mechanism is narrower: num_records is a 32-bit field derived from D's full byte extent, independent of how the base address is managed, and overflows/wraps for large enough D. Update the comment to describe that mechanism accurately. Credit to Koji Nakajima for catching the inaccuracy. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
|
Just to check @tony-davis: this is draft, but you still want a review? Are you looking for a "general approach" review? |
- Predicates_test.cpp: add a test that isolates the 0xfffff000 vs 2^32
threshold change from the min()-cap extent-formula fix (macroTile1 ==
N, so only the threshold constant differs between old and new code).
- buffer_store_offset_guard_gtest.cpp:
- Correct the header comment's stale claim that the SRD base is never
re-based per workgroup (same correction as the prior predicate
comment fix).
- Rewrite runConstantMatmul with scope-guarded cleanup so every early
return frees exactly what was allocated up to that point instead of
leaking partially-allocated buffers/handles/layouts/descriptors.
- Add the missing hipblasLt*Destroy calls for the handle, matrix
layouts, matmul descriptor, and preference, none of which were
previously released.
- Check the previously-ignored return status of preference creation
and of the verification scan's malloc/memset/memcpy calls; propagate
any failure through a new MatmulOutcome::verified field so a
verification-side failure can no longer be silently read as "nothing
wrong" (unwrittenCount/wrongCount staying at 0).
Re-verified on real MI350X (gfx950) hardware after these changes:
both buffer_store_offset_guard_gtest cases still pass, and the full
tensilelite-tests suite is 451/453 (same 2 pre-existing, unrelated
skips as develop).
Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
Author
|
Thanks for the thorough review! Replied inline on the three threaded comments. The rest of the findings were inside the collapsed review summary rather than separate threads, so replying here:
Re-verified all of the above on real MI350X hardware after the changes: both |
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
Unresolved review comments remain on test scalability, resource handling, boundary coverage, documentation, and gfx942 coverage.
Review details
Suppressed comments (5)
Previously missed (2) — in code that hasn't changed since the last review.
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:163
- The oversized case is expected to return zero algorithms, but this helper allocates and initializes roughly 19 GiB (A plus the 16 GiB D) before issuing the heuristic query. Every corrected run therefore pays a large allocation/fill cost just to select no kernel, and an allocation failure is returned as the same default
found=0outcome; move the selection/layout phase before these allocations and propagate resource failures separately.
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp:137 - This comment still attributes the bug to a base address that is never advanced per workgroup, but
computeStoreSrdStartdoes rebase the SRD base for each workgroup. The remaining limitation is the 32-bitnum_recordsceiling; please remove the stale base-address claim so the regression test documents the actual failure mode.
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:114
- On the pre-fix path, about 8.6 billion D elements remain NaN, so this executes a global 64-bit atomic increment at the same address for nearly every element. That single-counter hotspot can make the intentionally failing regression exceed the 600-second hipblaslt-test alarm; use per-block counters followed by a reduction (or another bounded verification) instead.
if(isnan(v))
atomicAdd(unwrittenCount, 1ull);
projects/hipblaslt/clients/tests/src/buffer_store_offset_guard_gtest.cpp:288
- This regression is skipped on every architecture except gfx950, while the PR identifies gfx942 as using the affected SRD path and the common predicate is attached to every BufferStore solution. The Test Result only reports local MI350X execution, so shared-lane evidence for the defect fix and coverage of gfx942 are missing; please run/report the affected shared lanes or add an explicit waiver explaining why gfx950-only coverage is sufficient.
if(gpuArchFamily() != "gfx950") \
GTEST_SKIP() << "ROCM-31016 was only reproduced/validated on gfx950, " \
"not " \
<< gpuArchFamily(); \
} while(0)
projects/hipblaslt/tensilelite/include/Tensile/ContractionProblemPredicates.hpp:1583
< BufferOOBBytesmakes the exact sentinel value a rejection boundary, but the new tests cover only one value below it and a range above it. Add a case whose computed extent equals0xfffff000(for bf16, M=8,388,600 and N=256) to lock down this fencepost and prevent an accidental<=change from re-admitting the unsafe boundary.
return multiplyElementSize(problem.d().strides()[1] * problem.d().sizes()[1],
problem.d().elementBytes())
< BufferOOBBytes;
- Files reviewed: 4/4 changed files
- Comments generated: 0 new
- Review effort level: Lite
…ause Revert BufferStoreOffsetLimitCheck's extent formula to the original min(MacroTile1, size[1]) cap; the SRD base is re-based per workgroup along N, so the full D extent was never the right quantity to check. Keep the corrected 0xfffff000 threshold constant. Add StreamKWorkgroupNumberCheck, a new predicate that bounds the tile-scaled grid Stream-K's launch-time fallbacks can produce at MAX_WORKGROUP_NUMBER (2^24), mirroring WorkgroupNumberCheck for ordinary kernels. Confirmed on gfx950 hardware that this fallback grid, not the SRD extent, is what silently drops D past this boundary. Update Predicates_test.cpp and buffer_store_offset_guard_gtest.cpp to match. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
One or more issues must be addressed before approval.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (4)
Previously missed (1) — in code that hasn't changed since the last review.
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp:203
- With
macroTile1set to 64, this predicate checksm * 64 * 2(about 1 GiB), whilemwas derived usingn == 256(about 4 GiB). Consequently this test never exercises the claimed near-BufferOOBboundary and would not catch a threshold regression. Passnto the predicate, or derivemfrom the actual macro-tile width.
projects/hipblaslt/tensilelite/Tensile/Contractions.py:554
- This adds the predicate only to future-generated solution metadata. The checked-in gfx950 logic under
library/src/amd_detail/rocblaslt/src/Tensile/Logiccontains Stream-K solutions but noStreamKWorkgroupNumberCheck, so existing shipped logic can still select the overflowing one-workgroup-per-tile fallback; the added client sanity test does not exercise that path. Please regenerate/update the production logic with this change, or land this explicitly together with the runtime regeneration follow-up so the reported corruption is actually prevented.
if isStreamK:
# Stream-K's normal grid is CU-scaled, not tile-scaled, so it is
# exempt from WorkgroupNumberCheck above. But several launch-time
# fallbacks (workspace too small for the ideal partial-tile
# reduction, the tree-fixup 24-bit bounds guard, a fixed-grid
# debug override, ...) hand Stream-K a one-workgroup-per-tile grid
# instead, same shape as the tile-scaled grid WorkgroupNumberCheck
# already bounds. Bound that worst case too.
rv += [cls('StreamKWorkgroupNumberCheck',
value=[state["MacroTile0"], state["MacroTile1"]])]
projects/hipblaslt/tensilelite/include/Tensile/ContractionProblemPredicates.hpp:1568
- The new 4-GiB threshold is not synchronized with all checked-in kernels this predicate can cover: the gfx950 custom Stream-K kernels referenced by production logic still define
BufferOOBas0x80000000(for example,tensilelite/Tensile/CustomKernels/Custom_Cijk_Ailk_Bljk_S_MX_B_BIAS_HA_S_SAV_NTD_SK3_UserArgs_MT256x256x32_MI16x16x1_shortname0_gfx950.s:439). Those kernels suppress stores starting at 2 GiB, so accepting offsets in the 2–4 GiB range can still silently drop output. Regenerate/update the affected custom kernels or keep the host predicate threshold tied to the actual sentinel used by each kernel family.
static constexpr uint64_t BufferOOBBytes = 0xfffff000ull;
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp:312
- The diagnostic string is grammatically incorrect: “fallback grids one workgroup” should say that the fallback “uses one workgroup per tile.”
"a Stream-K launch-time fallback grids one workgroup per tile.";
- Files reviewed: 6/6 changed files
- Comments generated: 2
- Review effort level: Lite
Comment on lines
+545
to
+549
| if isStreamK: | ||
| # Stream-K's normal grid is CU-scaled, not tile-scaled, so it is | ||
| # exempt from WorkgroupNumberCheck above. But several launch-time | ||
| # fallbacks (workspace too small for the ideal partial-tile | ||
| # reduction, the tree-fixup 24-bit bounds guard, a fixed-grid |
Comment on lines
+1688
to
+1692
| return static_cast<size_t>( | ||
| std::ceil(static_cast<float>(problem.freeSizeA(0)) / value[0])) | ||
| * static_cast<size_t>( | ||
| std::ceil(static_cast<float>(problem.freeSizeB(0)) / value[1])) | ||
| * problem.batchSize(0); |
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 : ROCM-31016
Motivation
On gfx950,
hipblasLtMatmulcan silently drop stores to part of the output tensorDfor certain bf16 problems, instead of erroring. No fault is raised; the untouched elements simply keep whatever was already in the destination buffer. This is GitHub ROCm/hipBLASLt#2299 and is tracked as ROCM-31016.This description supersedes the original one. The initial investigation suspected a single root cause (the buffer-store SRD's
num_recordsceiling) and a broad fix (check D's full byte extent). @nakajee flagged that fix as too broad (it would reject many large GEMMs that actually run correctly) and pushed for a closer look. A deeper investigation traced the confirmed hardware reproduction to a different, Stream-K-specific bug. Only the second one is what the confirmed reproduction shape actually exercises.Issue 1:
BufferStoreOffsetLimitCheck's threshold constantThe post-loop store SRD each
BufferStore=Truekernel programs bounds every store with a 32-bitnum_recordsfield, programmed at0xfffff000(~4 GiB - 4 KiB, theBufferOOBsentinel inKernelWriterAssembly.py'sallocPostLoopSrd). The SRD base is re-based per workgroup along the N dimension (computeStoreSrdStart), so only oneMacroTile1's worth of column extent needs to fit undernum_records, not the full D extent; the existing predicate'smin(MacroTile1, size[1])formula was already correct. Its threshold constant, however, was the old2^32boundary rather than the tighter0xfffff000sentinelBufferOOBmoved to. This PR corrects only the constant.In practice, realistic
MacroTile1values are far too small for this predicate to ever reject a real shape end-to-end, so this is a low-risk, narrow correction rather than the primary fix.Issue 2 (the real bug): Stream-K grid-dimension overflow
resolveStreamKSettings()/getSKGridImpl()(ContractionSolution.cpp) has several launch-time fallbacks (workspace too small for the ideal partial-tile reduction, a tree-fixup 24-bit bounds guard, a fixed-grid debug override) that hand a Stream-K kernel a one-workgroup-per-tile grid (skGrid = tiles) instead of its normal CU-scaled grid.WorkgroupNumberCheckalready bounds exactly this kind of tile-scaled grid atMAX_WORKGROUP_NUMBER(2^24) for ordinary kernels, but Stream-K solutions are unconditionally exempted from it, because Stream-K's normal grid isn't tile-scaled; the exemption didn't account for these fallbacks.When the resulting tile count is itself large enough (~2^24), the actual kernel launch's 32-bit work-item count (
workGroupSize * numWorkGroups) overflows. Confirmed directly viaTENSILE_DB=0x200000's Stream-K launch summary and an exhaustivegetAllAlgossweep on real MI350X (gfx950) hardware:tiles == 2^24exactly.N=256sotiles == M) either has its launch rejected outright (invalid argument, since the 32-bit work-item count wraps to exactly 0 at the boundary) or runs a much smaller grid than intended, leaving most ofDunwritten. Confirmed at 16 GiB (~99.9996% of D unwritten) down to ~9.25 GiB.num_recordssentinel from Issue 1), ran correctly. This is what falsified the original "any D extent pastnum_records" theory.Technical Details
BufferStoreOffsetLimitCheck: reverted the extent formula tomin(MacroTile1, size[1]); kept the corrected0xfffff000threshold.StreamKWorkgroupNumberCheck(new predicate): boundsceil(freeSizeA/MacroTile0) * ceil(freeSizeB/MacroTile1) * batchSizeatMAX_WORKGROUP_NUMBER(2^24) for Stream-K solutions, mirroringWorkgroupNumberCheck's existing formula and threshold. Generated byContractions.py'sCompoundPredicatesfor every Stream-K solution (the complement of the existing exemption).Important caveat:
StreamKWorkgroupNumberCheckis a brand-new predicate type. Existing, already-tuned Tensile logic files do not reference it, so this PR alone does not retroactively protect any already-shipped gfx950 (or other Stream-K-capable architecture) kernel; that requires the production logic files to be regenerated by whoever owns that tuning/release pipeline. I'd like to talk about filing a follow-up ticket to track that regeneration; happy to file it once we agree on ownership/scope.Test Plan
tensilelite/tests/Predicates_test.cpp: 3 tests for the correctedBufferStoreOffsetLimitCheckthreshold/formula, 4 new tests forStreamKWorkgroupNumberCheckusing the exact hardware-confirmed boundary (tiles == 2^24accepted, one confirmed-broken shape past it rejected, an ordinary sanity case, and a batch-multiplier case).clients/tests/src/buffer_store_offset_guard_gtest.cpp: sanity case only (see caveat above on why the oversized-D case isn't covered end-to-end here).getAllAlgoshardware sweeps on real MI350X, bisecting the exact safe/broken boundary across shapes from 4 GiB to 16 GiB.tensilelite-testssuite run locally on MI350X (gfx950): 454 pass, same 2 pre-existing skips asdevelop.Test Result
tensilelite-tests: 454 of 456 pass (2 pre-existing, unrelated skips), including all 7 predicate tests above.hipblaslt-test'sBufferStoreOffsetGuard_pre_checkin.SafeOutputDispatchesAndIsExactpasses on real gfx950 hardware.tiles == M):Submission Checklist
Risk level
Low. Solution-selection change only (predicate formulas/thresholds and one new predicate type); no generated kernel code changes.
BufferStoreOffsetLimitCheck's corrected threshold only tightens an already-narrow, rarely-triggered check.StreamKWorkgroupNumberCheckonly affects newly-generated/regenerated logic files (see caveat above); no currently-shipped solution's dispatch outcome changes until that regeneration happens.Related
allocPostLoopSrdper workgroup)