feat(tensilelite): StreamK cooperative cluster loads#9603
Open
jaopaulolc wants to merge 10 commits into
Open
Conversation
jaopaulolc
added
the
rocm:gemm algos
jaopaulolc
changed the base branch from
users/jolabega/clusterload-tristate-multicast
to
develop
July 20, 2026 19:56
✅ All Policy Checks Passed
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
jaopaulolc
changed the title
Codecov Report❌ Patch coverage is ❌ Your project check has failed because the head coverage (76.84%) 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 #9603 +/- ##
===========================================
+ Coverage 69.41% 69.61% +0.20%
===========================================
Files 2770 2765 -5
Lines 451826 450822 -1004
Branches 66440 66350 -90
===========================================
+ Hits 313602 313807 +205
+ Misses 117205 116403 -802
+ Partials 21019 20612 -407
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
jaopaulolc
force-pushed
the
users/jolabega/streamk-cluster-multicast
branch
from
July 20, 2026 21:51
d5d9692 to
e8b9c0c
Compare
|
🎉 All checks passed! This PR is ready for review. |
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
jaopaulolc
force-pushed
the
users/jolabega/streamk-cluster-multicast
branch
from
July 22, 2026 15:48
3db2ccf to
39772ee
Compare
…ulticast
JIRA ID: AIHPBLAS-3929
The TDM multicast ("cluster load") mask machinery was duplicated inline across
KernelWriter, KernelWriterAssembly, and Subtile/SubtileGREmit, and Multicast was
implicitly coupled to ClusterDim != [1, 1]. That coupling made it impossible to
compose barrier-only clustering and cooperative-load clustering independently, or
to reuse the mask attach at a new load site.
- New ClusterLoad component (Components/ClusterLoad.py, registered via
Component.py + Components/__init__.py). ClusterLoadTDM centralizes the mask
value computation (computeMasks), the MulticastMask* SGPR declare/undeclare
(declareSgprs / undeclareSgprs), the topology decision (combined vs split
A/B masks), and the descriptor attach at each load site (applyToDescriptor).
It is a behavior-preserving extraction: every method emits byte-identical
assembly to the previous inline code, and computeMasks receives the exact
SGPR operands the caller already holds so register indices are unchanged.
Selection is capability-based (HasTDM + TDMInst == 3), like TensorDataMoverLoad.
- KernelWriter / KernelWriterAssembly / SubtileGREmit now route mask declare,
undeclare, compute, and descriptor attach through ClusterLoadTDM.find(...).
- Tri-state Multicast parameter (ValidParameters + GlobalParameters default -1):
-1 = auto (legacy: ClusterDim != [1,1] implies Multicast, minus the StreamK
cluster paths), 0 = force off, 1 = force on. Default -1 reproduces the historic
derivation exactly, so every existing YAML (which omits Multicast) derives
byte-identically.
- Solution.py derives Multicast/ClusterBarrier from the tri-state, keying the
legacy auto coupling and the ClusterBarrier gate on StreamK == 0.
- Common/Utilities.clusterEnabled() helper for the ClusterDim != [1,1] test.
- Unit: test_cluster_load_component.py (component find/declare/compute/attach),
test_multicast_tristate.py (-1/0/1 derivation, legacy equivalence).
Low. Behavior-preserving refactor; the tri-state defaults to -1 (legacy auto),
so emitted assembly is unchanged for all existing configs.
…icast) Add the gfx1250 StreamK DP cooperative B-multicast fast path (StreamKMulticast), auto-derived from StreamK=3 + ClusterDim on top of the ClusterLoad component + tri-state Multicast foundation. - Collapse: StreamK==3 && ClusterDim!=[1,1] auto-enables StreamKMulticast (derived-only internal state; no YAML opt-in), forcing Multicast on and enabling the cluster-scope barrier handshake (ClusterBarrier). - StreamK.py: multicast mask predicate (clusterMulticastValid), DP->SK boundary mask clear, and a prologue wave-0 s_barrier_signal -3 that pairs the InsertClusterBarrierPass first-load wait (balanced signal/wait counts). - KernelWriter.py: postMainLoopBarrierCheckAndReset preserves cluster-scope -3 barriers. - ClusterLoad.py: re-add the split-mask hooks (usesCombinedMask / maskSgprName) for StreamKMulticast. - _validateStreamKMulticast: SK3 + [C,1] pow2 + gfx1250/HasTDM/TDMInst=3 + XCC=0; Multicast==1 force-on hard-guard. - Host: streamKMulticast field threaded through SizeMapping / C++ structs; grid multiple-of-C guard + getSKGridImpl ceil(tiles/C)*C override; SK4/SK5 + ClusterDim reject. Mutual-exclusion with the barrier-only cluster reduction is deferred to the following stacked PR. JIRA ID : AIHPBLAS-3929
… the zero-iteration skip path The StreamKMulticast prologue emits a wave-0-elected cluster-scope arrive before the first cooperative-multicast load. Its only matching cluster-scope wait was the first-load wait, which sits after the last-iteration guard's long branch to PrefetchGlobalLastIterEnd. On the zero-full-iteration path (reachable when K is not a whole multiple of DepthU) that branch skips the first-load wait, leaving the prologue arrive unmatched and the cluster-scope barrier unbalanced on that edge. Add a matching cluster wait on the skip edge (guarded by StreamKMulticast) so the arrive is consumed on every control-flow path, preserving whole-cluster barrier symmetry. The wait branches over on scc0 (>=1 full iteration, where the first-load wait pairs the arrive) and executes an all-waves cluster wait on scc1 (zero iterations); it leaves scc intact so the subsequent long branch is unaffected. Update the three gfx1250 cluster char-test balance checks: the skip-edge wait and the first-load wait are mutually exclusive at runtime but both emitted statically, so the static cluster-wait count is now exactly one greater than the signal count (one prologue arrive consumed by exactly one of the two waits).
…multicast - Option A prologue prefetch handshake bracketed with a cluster-scope barrier round (StreamK.py / KernelWriter.py StreamKMulticast plumb). - Remove the PGR<=1 enumeration guard so PGR2 multicast solutions build. - InsertClusterBarrierPass: drain the cooperative broadcast tensor group with both a pre-round drain and a producer-side drain (after the cooperative tensor_load group), gated on StreamKMulticast && PGR>=2. - PrefetchGlobalRead: [1,2] in the mxf4/mxf8 cluster-multicast configs plus a new PGR2 characterization test (+ snapshot + designed yaml). GPU-validated on gfx1250. JIRA ID : N/A
jaopaulolc
force-pushed
the
users/jolabega/streamk-cluster-multicast
branch
from
July 22, 2026 17:28
39772ee to
049da7b
Compare
…ulticast Condense verbose prose that duplicates the design docs (ClusterLoad component docstrings, StreamK multicast prologue helpers, _validateStreamKMulticast) to concise summaries with a "See docs/design/..." pointer, and strip explanatory comments from the client/characterization test YAMLs (copyright/SPDX retained; parsed-YAML unchanged). No codegen change: char snapshots and targeted unit tests unchanged.
jaopaulolc
marked this pull request as ready for review
July 22, 2026 18:15
… multicast Extend PrefetchGlobalRead to [1, 2] in sk_mxf4_force_dp_only_cluster_multicast.yaml so the forced DP-only cooperative B-multicast path is exercised at PGR2 (the B0-validated double-buffer broadcast + producer drain) as well as PGR1. No validator guard blocked PGR2; solutions enumerate err==0.
…st reject branches Add unit coverage for previously-untested branches surfaced by codecov: - ClusterLoad.computeMasks sparse-metadata path (Sparse==1 follows-A, Sparse==2 follows-B) and undeclareSgprs metadata SGPR free (ClusterLoad.py 90%->95%). - _validateStreamKMulticast reject branches unreachable via config derivation (StreamK!=3, StreamKXCCMapping!=0, non-gfx1250 ISA, missing HasTDM, missing HasClusterBarrier), driven directly like test_accept_pgr2.
…rain Reflow the producerDrainAnchors.push_back ternary to the repo clang-format style (Google/IndentWidth 4/ColumnLimit 100). Whitespace only; no behavior change.
cycheng
reviewed
Jul 24, 2026
| /// edge, so the publishing workgroup barrier at the next loop head correctly | ||
| /// orders it for the consuming waves. Matches PGR1, which drains each load in | ||
| /// its own iteration. | ||
| void insertProducerTensorDrainBefore(IRBase* anchor, AsmIRBuilder& irBuilder, GfxArchID archId) { |
Contributor
There was a problem hiding this comment.
@briannwu , could you help review joao's ClusterBarrier changes to see if that make sense to you?
The shared computeWorkspaceSrd helper addressed the StreamK partial-sum
workspace with a 32-bit SMulI32 slot*stride product. The per-slot stride
is MacroTile0*MacroTile1*bpe and the addressed slot index ranges over the
StreamK slot count (the partials workspace is partialTileSize ==
tileSize * skGrid, host ContractionSolution::partialTileSize), so for a
large SK grid the product exceeds 2^32, silently wraps, and the peer
write / owner read SRD aliases the wrong workspace slot.
The overflow depends only on the tile stride and the slot count, not on
the cluster mode, so the multicast [C,1] path -- which emits and reads
this exact workspace via partialsWriteProcedure / fixup -- can overflow
on large problems just like any other StreamK path. Compute the high
word with SMulHIU32 and fold it (plus the lo-add carry) into SrdWS+1,
emitted universally for every StreamK path (multicast, non-cluster, and
any future cluster mode). Adds one SGPR + one s_mul_hi_u32 per
workspace-SRD setup; the {basename, err} char goldens are unchanged and
still emit err==0. Real-HW validation pending (owner: user).
JIRA ID : N/A
1 task
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
2 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-3929
Motivation
gfx1250 StreamK GEMMs that co-locate consecutive data-parallel (DP) workgroups in a 1-D WG cluster (
ClusterDim = [C, 1]) process M-adjacent tiles that share the same B (N-block) over full K. Loading B once and TDM-multicasting it to the whole cluster (while A stays per-workgroup) removes redundant B traffic. This PR adds that cooperative-load fast path (StreamKMulticast) on top of the reusableClusterLoadcomponent + tri-stateMulticastfoundation from #9599.Technical Details
StreamK==3 && ClusterDim!=[1,1]auto-enablesStreamKMulticastinassignProblemIndependentDerivedParameters. It is a derived-only internal state key (likeClusterBarrier), never a YAML/benchmark parameter. The collapse forcesMulticaston and enables the cluster-scope barrier handshake (ClusterBarrier). The bare index-only StreamK cluster state no longer exists.StreamK.py). NewStreamKTwoTileDPFirstmethods:streamKMulticastMaskPredicate(runtimeclusterMulticastValid:nWG0 % Calignment + fully-populated cluster, else B loads normally) andstreamKMulticastBoundaryClear(drops the B broadcast at the DP->SK boundary), plus the cluster-scope arrive/wait helpers that keep theInsertClusterBarrierPasssignal/wait counts balanced:streamKMulticastPrologueSignal(wave-0 prologues_barrier_signal -3),streamKMulticastProloguePrefetchHandshake(brackets the PGR>=2 double-buffer prologue prefetch load that the generic per-load bracketing's backward scan cannot reach), andstreamKMulticastZeroIterClusterWait(consumes the prologue arrive on the zero-full-iteration skip path when K is not a whole multiple of DepthU).InsertClusterBarrierPass(shared/stinkytofu) emits a producer-side cooperative-load drain (s_wait_tensorcnt 0immediately after the cooperativetensor_load_to_ldsgroup, gatedstreamKMulticast && PGR>=2) so the peer-produced B broadcast is retired before the back edge and published coherently by the next loop-head barrier. This makes PGR2 numerically correct on the DP cooperative-multicast path (validated on gfx1250 B0). The multicast client configs sweepPrefetchGlobalRead: [1, 2].KernelWriter.py.postMainLoopBarrierCheckAndResetpreserves cluster-scope-3barriers (only workgroup-scope barriers are rebuilt).ClusterLoad.py. Re-adds the split-mask hooks (usesCombinedMask/maskSgprName) soStreamKMulticastbindsMulticastMaskA(self) on A andMulticastMaskB(broadcast) on B._validateStreamKMulticast: SK3,ClusterDim=[C,1](pow2 2..16), gfx1250HasTDM/HasClusterBarrier/TDMInst=3,XCCMapping=0, not atomic; plus aMulticast==1force-on hard-guard. SK4/SK5 +ClusterDimis rejected.ContractionSolution).streamKMulticastthreaded throughSizeMapping+ C++ structs/serialization; grid multiple-of-C guard and agetSKGridImplmulticast override (ceil(tiles/C)*C, keeps every launched HW cluster full for the single DP round; trailing partial clusters disabled by the runtime predicate).Test Plan
test_streamk_multicast.py,test_streamk_cluster_sk45_reject.py,test_multicast_tristate.py,test_cluster_load_component.py,test_multicast_legacy_coercion.py._codegensuite (non-multicast goldens byte-identical) +test_streamk_cluster_multicast_gfx1250_char.py,test_streamk_cluster_coop_load_gfx1250_char.py,test_streamk_cluster_multicast_pgr2_gfx1250_char.py.sk_mxf8gemm_cluster_multicast,sk_mxf4gemm_cluster_multicast,sk_mxf4_force_dp_only_cluster_multicast.Test Result
_codegengoldens byte-identical. Multicast char kernels emit a balanced cluster barrier (s_barrier_signal -3count ==s_barrier_wait -3count).Notes
Rebased onto latest
develop. Comment and test-YAML footprint trimmed (verbose prose condensed toSee docs/design/...pointers; the design docdocs/design/cluster-load-component-and-streamk-multicast.mdis included on this PR). Cumulative over #9599.Submission Checklist
Risk level
Low–moderate.
StreamKMulticastis derived-only and gated to gfx1250 SK3 clusters; existing YAMLs derive byte-identically (non-multicast goldens unchanged). The PGR2 enablement adds a narrowly-gated (streamKMulticast && PGR>=2) producer-side drain in the shared cluster-barrier pass that is inert for every other path. Base isdevelop; the branch is cumulative over #9599.Update — 64-bit workspace slot-offset fix (applies to the multicast path; HW validation pending)
Added the 64-bit workspace slot-offset computation to the shared
computeWorkspaceSrdhelper so every StreamK path that addresses the partial-sum workspace — including this PR's cluster multicast[C,1]path and non-cluster StreamK — emits a 64-bit offset instead of a 32-bit one. (This mirrors the now-universal fix landed on the stacked reduction/factored PRs #9611/#9612; on this branch the 64-bit hunk did not previously exist, so it is added here, un-guarded.)Motivation — could the overflow happen on the multicast path? Yes.
computeWorkspaceSrdsetSrdWS = AddressWS + offBytes * sPartialIdxwith a 32-bits_mul_i32, whereoffBytes = MacroTile0*MacroTile1*bpe(per-slot tile stride) andsPartialIdx ∈ [0, skGrid). The partials workspace is sizedpartialTileSize = offBytes * skGrid(hostContractionSolution::partialTileSize). The overflow depends only on the tile stride and the StreamK slot count, not on the cluster mode: the multicast[C,1]kernel emits and reads this exact workspace (partial-tile write viapartialsWriteProcedure+ owner fixup read), so for a large SK gridoffBytes * sPartialIdxexceeds 2³², the 32-bit product silently wraps, and the peer-write / owner-read SRD aliases the wrong workspace slot. Concrete example: a 256×256 tile with fp32 partials is 256 KiB/slot, soskGrid ≳ 2³²/262144 = 16384slots overflows a 32-bit offset (clusterCmultiplies the effective slot count, so big-Cmulticast on large problems is the realistic trigger).Technical Details
Compute the high word with
s_mul_hi_u32and fold it (plus the lo-add carry) intoSrdWS+1(s_mul_i32low +s_mul_hi_u32high +s_add_u32/s_addc_u32). Emitted for all StreamK paths; adds one SGPR + ones_mul_hi_u32per workspace-SRD setup. This changes the multicast codegen (an extra instruction), which is intended.Test Result (CPU, single-process; reused main worktree venv + editable rocisa .so)
{basename, err}char goldens are unchanged and stillerr == 0(a{basename, err}digest does not move when only instructions are added), so no snapshot was regenerated.[C,1]kernel now carries the 64-bit offset (s_mul_hi_u32 … // partials tile offset (high word) for 64-bit SRDfolded intoSrdWS+1).streamk_cluster_multicast.yaml4/4,streamk_cluster_multicast_pgr2.yaml4/4,streamk.yaml1/1 kernels →err == 0, 0DID_NOT_SATISFY_ASSERTS, 64-bit SRD present in every kernel.FFM/functional-sim is historically green regardless of HW correctness and is treated as a regression check only. Real gfx1250 hardware validation is still pending (owner: user).