feat(tensilelite): StreamK single-hop work stealing#8442
Merged
bnemanich merged 53 commits intoJul 24, 2026
Conversation
jaopaulolc
force-pushed
the
users/jolabega/dyn-streamk-queue-work-stealing-rebased
branch
from
June 15, 2026 17:48
3ba0a69 to
350d1de
Compare
jaopaulolc
changed the title
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #8442 +/- ##
===========================================
+ Coverage 69.50% 69.57% +0.07%
===========================================
Files 2770 2770
Lines 452498 452684 +186
Branches 66620 66649 +29
===========================================
+ Hits 314492 314929 +437
+ Misses 117526 117388 -138
+ Partials 20480 20367 -113
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
Add StreamKWorkStealing, a codegen-time solution parameter (not a kernel
argument) that enables single-hop next-neighbor work stealing in the
dynamic-queue StreamK fetch, for both SK4 (StreamKDynamic) and the dynamic
sub-path of SK5 (StreamKHybrid). When off (default) the emitted assembly is
identical to baseline.
When a workgroup's home queue (one of 8 hardcoded queues) is drained, it makes
exactly one atomic attempt to claim the structural-extra tile of the next
queue ((queueIdx+1) & 0x7), gated so it only fires when there is a tile
remainder and the neighbor owns an extra the home does not. To keep tile
ownership unique while stealing, the home queue's atomic auto-reset is disabled
whenever a remainder exists; the last workgroup restores the 8 queue counters
plus a completion counter (placed at synchronizer offset 0x80, clear of both
the 256B-strided queue counters and the partials/fixup flag region at 0x800)
in kernelEnd. The host already zeroes the whole Synchronizer buffer once at
handle creation, so no host-side workspace changes are needed. A lost steal
race leaves the index out of range and is absorbed by the existing
valid-work-item check.
Shared helpers on the StreamK base are reused by SK4 and SK5 (parameterized on
the skGrid/SKGrid SGPR). SK5's kernelEnd reset is gated on the hybrid mode bit
so it is a no-op in the static (SK3) sub-path. StreamKWorkStealing is rejected
unless StreamK in {4,5} and is incompatible with StreamKAtomic.
Unit tests (test_streamk_work_stealing.py): assert the StreamKWorkStealing
solution parameter exists, that the steal/no-reset/kernelEnd-reset assembly is
emitted by the new StreamK helpers, and that Solution validation rejects the
parameter outside StreamK in {4,5} and with StreamKAtomic.
GPU YAML tests (sk_dynamic_work_stealing.yaml for SK4, sk_hybrid_work_stealing.yaml
for SK5): fork StreamKWorkStealing over [0,1] (and StreamKHybridMode over [0,1]
for SK5) with NumElementsToValidate=-1 over aligned and unaligned/batched shapes
(640x640, 384x384x3) so the steal path and the kernelEnd counter-reset across
repeated launches are exercised. Validated on gfx950: all off/on and hybrid-mode
variants PASSED.
jaopaulolc
force-pushed
the
users/jolabega/dyn-streamk-queue-work-stealing-rebased
branch
from
June 18, 2026 01:08
8be6664 to
948a5a1
Compare
jaopaulolc
marked this pull request as ready for review
June 18, 2026 03:04
jaopaulolc
added
the
rocm:gemm algos
Correct bool/int type mismatches flagged by the strict corpus-clean validator: BoundsCheck and PrefetchLocalRead must be int, MIArchVgpr must be bool. Resolves the 2 corpus-clean failures for the new sk_dynamic_work_stealing / sk_hybrid_work_stealing tests.
…-streamk-queue-work-stealing-rebased
…ealing Adding StreamKWorkStealing to the required (min naming) parameter set inserts the SKWS0 token into the verbose solution name and perturbs the kernel identity-hash basenames, consistent with the rest of the StreamK family (SKA/SKFTR/SKFDPO/SKXCCM). Regenerate the syrupy golden snapshots to match. Diff is pure naming/identity churn -- no err-code, instruction count, or emitted-assembly changes.
…kernel end StreamK work-stealing (WS) kernels coordinate through device-global work-queue and completion counters. On the deferred-store / K-split epilogue path the kernel ended via SEndpgm without emitting the StreamK end-of-kernel reset (only the non-deferred functionEnd path emitted it). The counters therefore kept their drained terminal value across launches, so the first launch did the real work and launches 2..N found the queue exhausted and exited as no-ops. This produced physically impossible benchmark numbers (per-launch time collapsing below the arithmetic floor, us < min_us, super-peak Gflops) and would silently no-op repeated real-world calls until the buffer was re-zeroed. - Fix 1 (root cause, KernelWriter.py): emit skComponent.kernelEnd() in the hasAnyDeferred epilogue branch before SEndpgm, mirroring functionEnd. WS gating lives inside kernelEnd, so non-WS kernels are unaffected. - Fix 2 (defense-in-depth, tensile_host.cpp, default OFF): optional per-launch hipMemsetAsync of the Synchronizer for StreamK-WS solutions, gated by env TENSILE_STREAMK_WS_RESET_SYNCHRONIZER (off by default; Fix 1 is sufficient). Validated on gfx950/MI355X with a --streamk-ws-mode only library: per-launch time is now flat at the physical ~4.56 ms across iters (was collapsing toward sub-us), adaptive timing is sub-peak with us >= min_us, and --verify at iters=10 passes (norm_error ~1e-5) on WS-selecting shapes. (cherry picked from commit 58f9489)
(cherry picked from commit 2c5a52c)
(cherry picked from commit 65c65bb)
…anch The cherry-picked WS counter-reset fix included an optional host-side per-launch Synchronizer memset gated on solution->sizeMapping.streamKWorkStealing. That SizeMapping field does not exist on this branch (work-stealing is not exposed via the hipBLASLt API / to clients here), so the host block failed to compile (no member named 'streamKWorkStealing' in 'TensileLite::SizeMapping'). Remove the host-side block entirely. The kernel-side reset (KernelWriter.py, emitting the StreamK kernel-end reset on the deferred-store epilogue path) is the actual fix and is retained; it is sufficient on its own.
…-streamk-queue-work-stealing-rebased # Conflicts: # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_emit_bigfiles_char.ambr # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_emit_gfx1250_char.ambr # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_r3_datamover_char.ambr # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_r3_streamk_gfx1250_char.ambr
…count Add a NumXCD arch capability (populated in rocisa initArchCaps) mirroring origami get_default_num_xcds. StreamK codegen now derives the work-queue count from archCaps["NumXCD"] instead of a hardcoded 8 table; the value is a compile-time constant and asserted to be a power of two. The host derives the same count via origami and rejects at runtime when NUM_XCD is non-power-of-two or does not match the arch's XCD count; the work-stealing workspace is sized from it. Generated assembly is byte-identical for gfx942/gfx950.
… cache-line log2 Trim over-verbose and stale StreamK work-stealing comments and docstrings; fix the stale _WS_NUM_QUEUES_BY_ISA reference in ContractionSolution.cpp to point at _wsQueueConstants / archCaps["NumXCD"]; clarify that gfx942 covers both MI300X (8 XCDs) and MI300A (6 XCDs) with MI300A excluded at runtime by the host guard; compute _WS_CACHE_LINE_LOG2 from a named 256B constant. No logic change; emitted asm byte-identical (gfx942 and gfx950, SK4 and SK5, StreamKWorkStealing 0 and 1).
…line size Derive the per-queue counter stride from the hardware cache-line size (128 B for gfx942/gfx950) via a new origami accessor (hardware_t::get_default_cache_line_bytes), mirrored into a rocisa CacheLineBytes arch-cap. Stride == cache-line size so each per-XCD atomic counter occupies its own line (no false sharing), replacing the hardcoded 256 B over-alignment. Halves the per-queue stride/workspace; queue-index math (count/mask/log2 from NumXCD) is unchanged.
… dead code Behavior-neutral cleanup from a multi-model code review; no host logic changes and emitted kernel assembly is byte-identical for gfx942/gfx950. - collapse dead CacheLineBytes ternary (rocisa hardware_caps.hpp) and the all-128 origami get_default_cache_line_bytes switch to a plain return 128 - correct the origami hardware.hpp and ValidParameters StreamKWorkStealing docstrings to match the shipped single-hop next-neighbor implementation - trim the over-verbose StreamK banner/helper docstrings, the kernelEnd reset comments, the Solution DebugStreamK reject comment, and DECISIONS D17 - remove session/history spill from StreamK comments and the streamk work-stealing YAML headers (no test params/problem sizes changed) - delete two ghost-guarding unit tests that asserted absence of symbols from an abandoned in-PR design; de-session-spill remaining test docstrings
…wn XCD count and gfx1250 (1) Treat unknown/null NUM_XCD (missing analyticalHardware) as unsupported for the dynamic-queue path and fail early rather than under-allocating the counter workspace. streamKDynamicQueueUnsupported now returns true for unknown hardware so the softwarePredicate excludes the dynamic-queue solution and a non-dynamic-queue solution serves the GEMM; solve() throws an explicit error early if the baked queue count / NUM_XCD is unknown (0) before workspace sizing. (2) Reject StreamKWorkStealing on gfx1250, which lacks scalar atomics (HasSAtomic=false) for the steal path (streamKWorkStealingSteal emits s_atomic_inc with no vector fallback). gtest (CuCount_test) and unit test (test_streamk_work_stealing) updated. No codegen change for gfx942/gfx950 (asm byte-identical).
…-streamk-queue-work-stealing-rebased # Conflicts: # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_emit_bigfiles_char.ambr # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_emit_gfx1250_char.ambr # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_r3_datamover_char.ambr # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/_codegen/__snapshots__/test_r3_streamk_gfx1250_char.ambr
This _codegen golden did not textually conflict during the merge (only develop touched it, via #9227/#9217), so it auto-merged to develop's content. But the branch's StreamKWorkStealing -> required-parameter promotion (D17) also churns this kernel's basename hash, leaving the golden stale. Diff is basename-hash-only; err:0 unchanged, no instruction-count or emitted-assembly changes.
jaopaulolc
enabled auto-merge (squash)
July 20, 2026 16:29
AlexBrownAMD
approved these changes
Jul 20, 2026
jaopaulolc
changed the title
…-streamk-queue-work-stealing-rebased # Conflicts: # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/SolutionClass/__snapshots__/test_solution_class_char.ambr
…-streamk-queue-work-stealing-rebased # Conflicts: # projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/SolutionClass/__snapshots__/test_solution_class_char.ambr
bnemanich
disabled auto-merge
July 24, 2026 20:23
bnemanich
deleted the
users/jolabega/dyn-streamk-queue-work-stealing-rebased
branch
July 24, 2026 20:24
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
6 participants
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-3646
Motivation
In the dynamic-queue StreamK schedulers, a workgroup whose home queue drains early sits idle while sibling queues still have structural-extra tiles to process. This load imbalance is most visible on tile counts that do not divide evenly across the queues. This PR lets a drained workgroup steal a single neighboring tile so the imbalance is smoothed out.
Technical Details
Adds single-hop next-neighbor work stealing to the dynamic-queue StreamK fetch, controlled by a new codegen-time Tensile solution parameter
StreamKWorkStealing(∈[0, 1], default0). It applies to both SK4 (StreamKDynamic) and SK5 (StreamKHybrid's dynamic sub-path).When a workgroup's home queue (1 of 8) drains, it makes exactly one atomic attempt to claim the structural-extra tile of queue
(queueIdx+1) & 0x7, gated so it only fires when a tile remainder exists and the neighbor owns an extra the home does not. Home-queue atomic auto-reset is disabled while stealing (to keep tile ownership unique under races); the last workgroup restores the queue + completion counters inkernelEnd.Test Plan
-m unit, TensileLite): fullTensile/Tests/unit/suite, including the characterization.ambrgoldens and the strict-corpus gate (test_input_yaml_corpus_is_strict_clean), plustest_streamk_work_stealing.py(asserts the parameter exists, the steal / no-reset / kernelEnd-reset assembly is emitted by the new helpers, and Solution validation rejects the parameter outsideStreamK ∈ {4,5}and withStreamKAtomic).NumElementsToValidate=-1(full-element validation):sk_dynamic_work_stealing.yaml(SK4): forkStreamKWorkStealing∈[0,1].sk_hybrid_work_stealing.yaml(SK5): forkStreamKWorkStealing∈[0,1]×StreamKHybridMode∈[0,1].640x640(25 tiles → steal fires) and batched384x384x3, with repeated launches to exercise thekernelEndcounter reset.Test Result
Revalidated on top of the latest
develop(post-merge):.ambrbyte-stable on re-run).sk_dynamic_work_stealing.yamlandsk_hybrid_work_stealing.yamlboth PASSED (clientExit=0), allStreamKWorkStealingon/off ×StreamKHybridModeon/off variants validating with full-element checking; broader StreamK GPU family 22 passed / 27 skipped (non-gfx950 archs) / 0 failed.Submission Checklist
Risk level
Low — the entire feature is behind a codegen-time toggle that defaults to off, so generated kernels and runtime behavior are unchanged unless
StreamKWorkStealing=1is explicitly selected.