Repository navigation
fix(tensilelite): refuse kernel launches whose grid does not fit in 32 bits - #12283
Conversation
KernelInvocation carries its launch grid in dim3, which is vector3<size_t>, but both launch APIs take 32-bit grid parameters: hipExtModuleLaunchKernel's globalWorkSize arguments are unsigned int, and HIP_LAUNCH_CONFIG's gridDim fields are too. A grid that does not fit was narrowed at the call rather than refused. That gives two outcomes. A work-item count of exactly 2^32 narrows to zero, which HIP rejects, so the call fails loudly. A count above 2^32 narrows to the low 32 bits, so the launch succeeds, a fraction of the intended workgroups run, and every tile the missing workgroups would have written is left holding whatever was there before. Measured on gfx950 at 256 threads, 2^24 + 16 workgroups wrapped to a 16-workgroup launch and left the other 16,777,216 workgroups' worth of D unwritten with no error raised. Validate the grid in SolutionAdapter::launchKernel and return hipErrorInvalidValue when it cannot be expressed. The check runs before the code object is looked up, since a grid that cannot be expressed is impossible whether or not the kernel loads, and that is the last point where the 64-bit values are still intact. The two launch paths narrow different quantities, so each is bounded against the one it actually passes: hipDrvLaunchKernelEx enumerates the grid in workgroups, so it is bounded on numWorkGroups, while hipExtModuleLaunchKernel takes work items, so it is bounded on numWorkItems. A single shared check would wrongly reject cluster launches whose workgroup count fits even though the work-item product does not. LaunchDimGuard_test drives launchKernel directly with a synthetic KernelInvocation, so it needs no GPU work and no kernel to exist. It covers the 2^32 boundary, the silent 2^24 + 16 case, the largest grid that must still pass, an oversized y dimension, and both directions of the cluster path's workgroups-versus-work-items distinction. That last pair needs to know how the library was compiled, so the tests target picks up the cached HIP_HAS_CLUSTER_LAUNCH result rather than repeating the feature test. ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
…imit
Stream-K normally sizes its grid to the compute-unit count, which stays small.
But getSKGridImpl() has launch-time paths that hand it one workgroup per output
tile instead, the dominant one being the tree-fixup bound:
if(itersPerTile >= 65536 || itersPerWG >= 65536
|| (tiles * itersPerTile) >= 16777216)
skGrid = tiles;
Once the tile count reaches 2^24 the third clause holds for any itersPerTile of
1 or more, so it fires for exactly the shapes that cannot survive it.
WorkgroupNumberCheck already bounds a tile-scaled grid for ordinary kernels,
but it exempts Stream-K unconditionally, on the reasoning that a Stream-K grid
is not tile-scaled. That reasoning does not cover these fallbacks, so nothing
bounded the grid at all, and the launch narrowed it instead.
Add StreamKWorkgroupNumberCheck, which bounds
ceil(freeSizeA/MacroTile0) * ceil(freeSizeB/MacroTile1) * batchSize at
UINT32_MAX / 256, so a 256-thread kernel's largest admitted grid is 2^32 - 256
work items. It is named apart from MAX_WORKGROUP_NUMBER despite the equal
magnitude, because that constant guards an unrelated fp32 conversion in the
hand-written custom kernels. The tile count uses integer ceiling division: a
float ceil() loses precision above 2^24 and would undercount tiles at the
boundary this exists to guard.
The bound assumes 256 threads, the dominant Stream-K configuration on gfx950
with 290 distinct macro tiles. It is exact there, conservative at 64 and 128
threads, and too permissive at 512. Across all 2,796 shipped logic files one
family binds tighter, MacroTile 448x128 at 512 threads, and reaching its wrap
point needs roughly 896 GiB of D. The launch guard added in the preceding
commit makes that case a loud failure rather than a silent one.
The predicate is checked for every Stream-K solution rather than only when a
fallback is predicted, because the reduction strategy that decides which path
a launch takes is not available at predicate time.
Contractions.py attaches it on the isStreamK arm that already skips the
ordinary WorkgroupNumberCheck. That arm has no test of its own; the coverage
here is the predicate's behavior, in Predicates_test.cpp.
ROCM-31290
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. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #12283 +/- ##
===========================================
+ Coverage 70.43% 71.68% +1.25%
===========================================
Files 2933 2849 -84
Lines 482485 468310 -14175
Branches 69520 69188 -332
===========================================
- Hits 339819 335694 -4125
+ Misses 118818 108992 -9826
+ Partials 23848 23624 -224
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
The preceding commit noted that the isStreamK arm has no test of its own and that the coverage there is the predicate's behavior in Predicates_test.cpp. This closes that gap: the arm itself is now covered. Add a real gfx950 Stream-K fixture chain to test_contractions_char.py, reusing solution_pool_gfx950.yaml, which is already vendored for the BenchmarkProblems solution-pool tests, so test_problem_predicate_compound_streamk exercises CompoundPredicates()'s isStreamK branch against a real derived solution state rather than a hand-synthesized one. The test asserts both directions of the branch: StreamKWorkgroupNumberCheck is attached, and the ordinary WorkgroupNumberCheck is not. Both assertions were checked by mutation. Removing the emission fails the first; removing the isStreamK guard so both predicates attach fails the second. Log this as DECISIONS.md D24: the first arm of the "Contractions matrix" slice ADR 0007 deferred, landed via an already-vendored fixture instead of the two alternatives that ADR explicitly rejected, vendoring a new fixture or hand-synthesizing derived state. ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
…y computes The comment claimed the helper mirrors ContractionProblemGemm::getNumTiles(sizeMapping, 1). It does not. getNumTiles multiplies over every free and batch index and, when packBatchDims is set, folds batch into the M or N extent before dividing by the macro tile. The helper uses index 0 of each and multiplies batch in afterwards. The two agree for the single-free-index, unpacked-batch case, which is what reaches this predicate today, and WorkgroupNumberCheck makes the same simplification, so the two predicates agree with each other on every shape. That is the property worth recording, because a reader who trusts the word "mirrors" would not re-derive it before changing either one. Comment only; no behavior change. Also makes enableCluster const on the cluster-launch path, matching the definition already used on the other side of the #ifdef. ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
…rid-overflow Resolves a DECISIONS.md conflict: the Gemm-From-Anywhere custom kernel framework (#11155, 1833a35) landed D24 and D25 on develop while this branch was adding its own D24. Renumbered this branch's entry to D26 and placed it after develop's two, which are taken unchanged. Only the numbering collided; the entries are unrelated. Co-authored-by: Cursor <cursoragent@cursor.com>
tony-davis
changed the title
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved direct launch coverage, validation ordering, and multi-batch tile-count issues remain.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds Stream-K grid bounds to prevent 32-bit HIP launch overflow through launch validation and solution-selection predicates.
Changes:
- Adds launch-dimension validation for regular and cluster paths.
- Adds and serializes
StreamKWorkgroupNumberCheck. - Adds unit, characterization, and decision-record coverage.
File summaries
| File | Reviewed changes |
|---|---|
projects/hipblaslt/tensilelite/tests/Predicates_test.cpp |
Tests Stream-K predicate boundaries. |
projects/hipblaslt/tensilelite/tests/LaunchDimGuard_test.cpp |
Tests launch-dimension guards. |
projects/hipblaslt/tensilelite/tests/CMakeLists.txt |
Registers new tests. |
projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/DECISIONS.md |
Documents characterization decisions. |
projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/Contractions/test_contractions_char.py |
Tests predicate attachment. |
projects/hipblaslt/tensilelite/Tensile/Contractions.py |
Attaches the Stream-K predicate. |
projects/hipblaslt/tensilelite/src/hip/HipSolutionAdapter.cpp |
Validates launch dimensions. |
projects/hipblaslt/tensilelite/include/Tensile/Serialization/ContractionPredicates.hpp |
Adds predicate serialization support. |
projects/hipblaslt/tensilelite/include/Tensile/ContractionProblemPredicates.hpp |
Implements Stream-K tile bounds. |
Review details
- Files reviewed: 9/9 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.
… count Three findings from the Copilot review of #12283, all verified against the source before acting. Move the launch-dimension check ahead of FindCodeObject. The comment claimed the check ran "before the code object is looked up", but FindCodeObject() was several lines above it; only the getKernel() function lookup came after. The check is now the first thing in launchKernel, so an unexpressible grid costs no module load. This also places it ahead of the m_debugSkipLaunch early return, so that debug knob now reports the invalid grid instead of returning success. Count every batch index in StreamKWorkgroupNumberCheck::tiles. getNumTiles() multiplies over the whole m_batchIndices loop while this took batchSize(0), matching WorkgroupNumberCheck and the LeadingFree*SizesGreaterOrEqual predicates. Those carry an assert(batchIndices().size() <= 1) that compiles out of a release build, and undercounting admits a grid that then wraps, which is the failure this predicate exists to prevent. Accumulating can only raise the count, so it can only reject more. Covered by StreamKWorkgroupNumberCheck_EveryBatchIndexCounted, built on a two-batch-index problem: all batch indices give 2^26 tiles and reject, batchSize(0) alone gives 2^20 and would admit. Narrow the claim about hand-written custom kernels. rocblaslt's rocRoller custom kernels call hipExtModuleLaunchKernel directly from library/src/amd_detail/rocblaslt/src/rocroller/custom_kernels.cpp and never reach SolutionAdapter, so this guard does not make that path loud. That path builds its grid from uint32_t tile counts rather than narrowing a 64-bit dim3, so its failure mode differs from the one guarded here; the comment now says so rather than overclaiming coverage. ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
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 (2)
projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/Contractions/test_contractions_char.py:157
- The Test Plan says the
Contractions.pyisStreamKarm is not covered, but this newly added test explicitly exercises that arm and verifies the predicate swap. Please update the PR description so the stated coverage and review risk match the actual tests.
def test_problem_predicate_compound_streamk(problem_type_streamk, solution_state_streamk):
# test_problem_predicate_compound above never exercises
# CompoundPredicates()'s isStreamK branch (Contractions.py): the vendored
# gfx942 fixture's solution has no StreamK field. This drives the same
# call from a real Stream-K solution's derived state instead, which must
# skip the ordinary WorkgroupNumberCheck and attach
# StreamKWorkgroupNumberCheck in its place.
tags = {
p.tag
for p in C.ProblemPredicate.CompoundPredicates(solution_state_streamk, problem_type_streamk)
}
assert "StreamKWorkgroupNumberCheck" in tags
projects/hipblaslt/tensilelite/tests/LaunchDimGuard_test.cpp:20
- The statement that
SolutionAdapter::launchKernelis where every launch passes through is inaccurate: this repository also has directhipExtModuleLaunchKernelcalls for custom kernels, as the implementation's caveat notes. Please limit this test comment to the ordinary SolutionAdapter path so it does not imply coverage of those unguarded launches.
// SolutionAdapter::launchKernel is the one place every launch passes through,
// so that is where the refusal lives.
- Files reviewed: 9/9 changed files
- Comments generated: 1
- Review effort level: Lite
Copilot review follow-up on #12283, and a correction to f86cab9. getNumTiles() multiplies every free size before dividing by the macro tile (the m_freeIndicesA / m_freeIndicesB loops in ContractionProblem.cpp), while tiles() divided freeSizeA(0) alone. The product of the free sizes is larger than any one of them, so this under-counted, which is the unsafe direction: the predicate admitted grids that then wrap. Measured on a two-free-index contraction with freeSizeA {4096, 4096} and freeSizeB 256 at a 16x16 macro tile. The real count is 16,777,216 tiles, one past the bound, and the predicate computed 4,096 and accepted it. f86cab9 fixed the same mistake for batch indices but asserted this divergence could only over-count. That was wrong, and the comment said so in the source. Both the code and the comment are corrected here: tiles() now accumulates every free and batch index, leaving packBatchDims as the only divergence from getNumTiles, which does over-count and so only rejects more. Covered by StreamKWorkgroupNumberCheck_EveryFreeIndexCounted, sized to land exactly one tile past the bound so it pins the boundary as well as the accumulation. Mutation-confirmed: reverting to freeSizeA(0)/freeSizeB(0) fails it and no other test. ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
… count Copilot review follow-up on #12283. The comment said the packBatchDims divergence from getNumTiles could only over-count. That holds for modes 1 and 2, where batch folds into one extent before the macro-tile divide, but not for mode 3, where getNumTiles folds batch into both M and N and so scales by batch squared while this predicate scales by batch. One batch of 2 at a 16x16 tile gives getNumTiles 1024 against 512 here, an under-count. No code change. The predicate is serialized with MacroTile0/MacroTile1 only and never sees sizeMapping, so it cannot read the packing mode, and "conservatively reject mode 3" is not available to it either. Mirroring getNumTiles means carrying the mode in the predicate's serialized value, across the Python emission, the C++ struct and its serialization. That is not worth doing for a mode that does not occur: packBatchDims is 0 in all 77,421 solutions of the shipped gfx950 library. The launch guard added in this PR is the backstop, since it bounds the grid it is actually handed rather than a predicted one, so an under-count here surfaces as hipErrorInvalidValue instead of silent truncation. The comment now enumerates all four modes and says which one is unsafe, rather than making a blanket claim. This is the second time that blanket claim has been wrong, after the free-index case in 55c37f9. ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
|
@copilot resolve the merge conflicts in this pull request |
…ng/rocm-31290-streamk-grid-overflow # Conflicts: # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/DECISIONS.md Co-authored-by: tony-davis <7938412+tony-davis@users.noreply.github.com>
Merged |
|
@copilot resolve the merge conflicts in this pull request |
…ng/rocm-31290-streamk-grid-overflow # Conflicts: # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/DECISIONS.md Co-authored-by: tony-davis <7938412+tony-davis@users.noreply.github.com>
Merged |
newling
left a comment
newling
left a comment
There was a problem hiding this comment.
Nice find. I had a discussion with an AI about this, and learned quite a lot in the process.
While discussing with the AI, I had 2 ideas. The first was, isn't the best safety net position where the 64-bit integer is cast to a 32-bit integer? That resulted in this branch, which might be complementary to this: https://github.com/newling/rocm-libraries/tree/users/newling/launch-grid-narrowing-checks
The second was, the logic for stream-k in the case where workspace is too small and it falls back to 1-tile per workgroup, we might be able to unify the logic with the non-stream-k case. This resulted in this: https://github.com/newling/rocm-libraries/tree/users/newling/streamk-launch-limits which looks slightly simpler.
I will generate an AI review now.
(my summary of the PR, not part of review per se) Essentially the constraint imposed by the hip APIs is that the number of threads (or workgroups, depending on the API) in each dimension of a grid must be less then 2**32. Which when considered there are more than just 1 (x) grid is really ample! Unfortunately tensilielite squeezes all threads into the x dimension, and so the constraint becomes more than just theoretical. A further waste is that we know the number of threads will be divided by the number of threads per workgroup, so we're self-imposing a constraint on the GEMM sizes we can run that is narrower than it needs to be.... but, that is too much to try and fix here. This PR is (correctly) about making sure we fit within the constraints.
newling
left a comment
newling
left a comment
There was a problem hiding this comment.
The only request in the above is for an additional test, which I think is worth looking at. Seems like something a mutation test might catch...
I'm approving this as is as it fixes a real issue without adding too much complexity. But I've left a lot of loose threads in my review, so am also happy to re-review this if you decide to deviate significantly from the current design @tony-davis
|
Thanks James! I really appreciate you digging into this. Your summary is spot on too; squeezing everything into x makes the limit a lot more real than it needs to be. On your first idea, I think that's the same check this PR already has, just in a different spot. I put it at the top of On your second idea, you're right, and it made me realize I missed that That said, your branch feels like more than this fix needs. So here's what I'd propose for this PR: drop Your full rework would then make a great follow-up. It's the more precise version, since it only rejects a Stream-K solution when the fallback actually happens. That wins back the tuned candidates the conservative bound gives up. I'd want to see a timing check on it, though. It adds a second Does that split sound OK to you? |
…ng/rocm-31290-streamk-grid-overflow
|
Quick update, and a correction to my last comment. Until I read your AI review, I hadn't caught that #12204 (merged Sep 18) already fixed the tree-fixup fallback. When the tile count hits the thread-aware cap, it now launches So I think the right shape for this PR is smaller still: just the launch guard. The remaining ways to get an oversized grid (the insufficient-workspace fallback, That leaves your |
…nch guard #12204, merged after this PR was opened, changed the tree-fixup fallback in getSKGridImpl(). Once the tile count reaches (UINT32_MAX + 1) / threadsPerWorkGroup, it now launches cuCount * occupancy instead of one workgroup per tile. That was the main path StreamKWorkgroupNumberCheck existed to keep selection away from. With it fixed, the predicate rejects every Stream-K solution at or above that tile count, including ones that now run correctly whenever the workspace fits, so it throws away working candidates. Remove the predicate, its serialization and Python emission, its gtests, the characterization test for the isStreamK arm, and the DECISIONS.md entry. This also drops the float-precision test that review found never crossed the boundary. The launch guard in SolutionAdapter::launchKernel stays and is now the whole fix. Paths that still hand Stream-K one workgroup per output tile (the insufficient-workspace fallback in resolveStreamKSettings(), K == 0, the data-parallel debug knob) and the skFixedGrid / skGridMultiplier overrides fail with hipErrorInvalidValue instead of launching a truncated grid. Its comments now say that instead of pointing at the removed predicate. A selection-time check that resolves the grid that will actually launch, workspace fallback included, is the better fix on the selection side. It is left to a follow-up; see the review discussion on #12283. tensilelite-tests on gfx950 (MI355X): 500 ran, 497 passed, 3 skipped (the same three noted in the PR description). ROCM-31290 Co-authored-by: Cursor <cursoragent@cursor.com>
tony-davis
changed the title
TheRock Submodule Bump ActivityNewest first
|
JIRA ID : ROCM-31290
Motivation
hipBLASLt launches each GEMM kernel with a grid, the number of parallel work items the hardware should run.
KernelInvocationholds that grid in a 64-bit type, but both launch entry points take 32-bit grid parameters.hipExtModuleLaunchKerneldeclares itsglobalWorkSizearguments asunsigned int, and the cluster path'sgridDimfields are the same width. A grid too large to express was narrowed at the call rather than refused.That produces two outcomes. A work-item count of exactly 2^32 narrows to zero, and HIP rejects the launch with an error. A count above 2^32 narrows to its low 32 bits, a small non-zero number, so the launch succeeds while only a fraction of the intended workgroups run. The output tiles the missing workgroups would have written keep whatever the buffer held before, and
hipblasLtMatmulreturns success.Stream-K kernels reach this most easily. They normally size their grid to the compute-unit count, but several paths hand them one workgroup per output tile instead, and solution selection's
LaunchLimitscheck exempts Stream-K.Technical Details
SolutionAdapter::launchKernelnow validates the grid and returnshipErrorInvalidValuewhen it does not fit, with a diagnostic naming the offending dimensions. The check runs first, before the code object is looked up, since that is the last point where the 64-bit values are still intact.The two launch paths narrow different quantities, so each is bounded against the one it actually passes.
hipDrvLaunchKernelExtakes the grid in workgroups and is bounded onnumWorkGroups.hipExtModuleLaunchKerneltakes work items and is bounded onnumWorkItems. A single shared check would wrongly reject a cluster launch whose workgroup count fits even though the work-item product does not.#12204 already fixed the most common way to get there, the tree-fixup fallback in
getSKGridImpl(), which now launchescuCount * occupancyonce the tile count reaches the thread-aware cap. The paths that remain include the insufficient-workspace fallback inresolveStreamKSettings(), K == 0, the data-parallel debug knob, and theskFixedGrid/skGridMultiplieroverrides. The guard turns all of them from silent truncation into a loud error.Earlier revisions of this PR also added a selection-time predicate,
StreamKWorkgroupNumberCheck. It was removed in review. After #12204 it rejected Stream-K solutions that now run correctly, and a check that resolves the grid that will actually launch is a better fit for the selection side. That is left to a follow-up.This changes host-side launch validation only. No kernel code or solution selection changes.
Test Plan
tensilelite-tests,tensilelite/tests/LaunchDimGuard_test.cpp. Six cases that driveSolutionAdapter::launchKerneldirectly with a syntheticKernelInvocation, so no GPU work is dispatched and no kernel needs to exist. They cover the 2^32 boundary, the silent case above it, the largest grid that must still pass, an oversized y dimension, and both directions of the cluster path's workgroups-versus-work-items distinction.tensilelite-testssuite on gfx950 to check nothing else moved.Not covered: rocblaslt's rocRoller custom kernels call
hipExtModuleLaunchKerneldirectly and never reachSolutionAdapter. That path builds its grid fromuint32_ttile counts, so its failure mode is arithmetic overflow rather than narrowing, which is tracked separately as ROCM-31230.Test Result
tensilelite-testson an MI355X (gfx950), MessagePack configuration: 500 ran, 497 passed, 3 skipped. The three skips are unrelated and also skip ondevelop:RangeLibraryTest.SpecificSizes/0,MXDataGenFP4Test.ZeroFrequencyWithinBounds/3, andDataInitializationReferenceRecompute.Gfx1250ReusesPreswizzledMXAcrossSolutions.For context on impact, a hardware sweep against the TheRock
10.1.0a20260822library (which predates #12204) found 39 Stream-K algorithms at M = 19,398,656, N = 256 that returned success with part of the output never written, 22 of them leaving exactly 2^32 elements unwritten. #12204 fixes most of those at that shape. The sweep was not rerun against a current library.Submission Checklist
Risk level
Low (2/5). The guard is a few integer comparisons at the top of the launch path, and it only changes behavior for grids that could not have run correctly. Solution selection is unchanged.