feat(tensilelite): add StreamK WG-cluster partial reduction for gfx1250#9611
Open
jaopaulolc wants to merge 16 commits into
Open
feat(tensilelite): add StreamK WG-cluster partial reduction for gfx1250#9611jaopaulolc wants to merge 16 commits into
jaopaulolc wants to merge 16 commits into
Conversation
✅ All Policy Checks Passed
📖 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. |
jaopaulolc
added
the
rocm:gemm algos
1 task
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 #9611 +/- ##
===========================================
+ Coverage 64.17% 69.64% +5.47%
===========================================
Files 2763 2765 +2
Lines 450335 450984 +649
Branches 66254 66378 +124
===========================================
+ Hits 288962 314054 +25092
+ Misses 140483 116404 -24079
+ Partials 20890 20526 -364
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
jaopaulolc
force-pushed
the
users/jolabega/streamk-cluster-reduction
branch
from
July 22, 2026 15:48
67dbde3 to
9011c34
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
StreamK GEMM tiles whose partial results are reduced across workgroups (WGs) synchronize that reduction with a cross-CU global-workspace flag spin-wait. On gfx1250 the WGs that reduce one tile can instead be co-located into a hardware workgroup cluster and synchronized with an intra-cluster split barrier, removing the spin-wait from the fast path. Adds an opt-in solution parameter StreamKClusterReduction. When enabled, a StreamK tile's fixup peers are placed in a single 1-D workgroup cluster (ClusterDim = [C, 1]); the HW WG-id remap (WorkGroup0 = cluster_x*C + wg_x) makes each cluster own the contiguous StreamK index range [c*C, c*C+C), i.e. exactly one tile's peers. - Kernel (Components/StreamK.py): clusterReduceSignal / clusterReduceWait / clusterReduceIntraCheck helpers emit the wave-0-elected s_barrier_signal / s_barrier_wait handshake in the fixup epilogue. A uniform intra-cluster predicate selects the barrier fast path or the retained global-flag path so a cluster never mixes the two. The WG-id reread is skipped under clustering so StreamKIdx stays the global index. - Host (src/ContractionSolution.cpp, ContractionSolution.hpp + serialization): getSKGridImpl reduction override rounds the SK grid to C*tiles (fixed even split, SKItersPerWG = itersPerTile/C) and keeps the launch grid a multiple of C; the two-tile cluster-reduction kernarg block is emitted only when sk.grid == C*tiles; the field is threaded through SizeMapping / Contractions. - Predicate (ContractionProblemPredicates.hpp + serialization + Contractions.py): ClusterReductionIterCheck rejects at selection time any problem where itersPerTile is not a multiple of C (the split barrier would otherwise over-signal); the global-flag reduction is retained as the runtime fallback. - Solution.py: _validateStreamKClusterReduction gates the opt-in (SK3, non-atomic, non-DP-only, [C,1] power-of-two cluster, gfx1250 with HasClusterBarrier, TDMInst != 0). Multicast and the mainloop ClusterBarrier are kept off for a StreamK reduction cluster: its peers iterate different K-splits and never march the mainloop in lockstep. Multicast and cluster reduction are mutually exclusive: the StreamKMulticast auto-derivation is gated on NOT StreamKClusterReduction, and both validators reject the simultaneous combination with a neutral message. This stacks on the reusable ClusterLoad component + tri-state Multicast and the StreamK cooperative multicast; the mutual exclusion is intended to be unified by a later factored combined-cluster mode. Validation: CPU asm-string unit tests (test_streamk_cluster_reduction.py), snapshot codegen characterization (test_streamk_cluster_reduction_gfx1250_char), the ClusterReductionIterCheck predicate characterization and C++ gtests, plus opt-in-off client GEMM YAMLs. The feature is off by default and strictly opt-in.
… 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 is 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 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-reduction
branch
from
July 22, 2026 17:28
9011c34 to
9074f06
Compare
…eduction Condense verbose prose that duplicates the design docs (ClusterLoad docstrings, StreamK multicast/reduce cluster helpers, _validateStreamKMulticast / _validateStreamKClusterReduction) 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
…orce-DP-only client tests Extend PrefetchGlobalRead to [1, 2] in the MX-FP4/FP8 cluster-reduction client configs and sk_mxf4_force_dp_only_cluster_multicast.yaml so both PGR1 and PGR2 are exercised on HW. No validator guard blocks PGR2; solutions enumerate err==0 (reduction: 12 PGR1 + 12 PGR2 each; force-DP-only: 3 + 3).
…for zero-skip coverage
The mxf4/mxf8 cluster-reduction client YAMLs paired big tiles (MT64x64, MT128x128)
and high cluster factors (C up to 8) with low-M / low-K sizes, so 32 of 96
solution x size runs were rejected at selection time by ClusterDimCheck
(ceil(M/MT0) % C != 0) or ClusterReductionIterCheck (ceil(K/DepthU) % C != 0),
surfacing as DID_NOT_SATISFY_ASSERTS with no real coverage.
Redesign both configs so every (tile x C x size) combo runs:
- Drop MT128x128 (MT128 x C8 needs M a multiple of 1024 and is cluster-resource
heavy); keep MT32x16 / MT32x32 / MT64x64.
- Sizes with M and N multiples of 512 and K a multiple of 2048 so
itersPerTile = ceil(K/256) in {8,16} is divisible by C in {2,4,8} and
ceil(M/MT0) is divisible by C for every kept tile:
Exact [512,512,1,2048], [1024,512,1,2048], [512,512,1,4096], [512,1024,2,2048].
- Retain C in {2,4,8}, PrefetchGlobalRead [1,2], a multi-iter K (4096) and a
batched size ([512,1024,2,2048]).
Result: 0 DID_NOT_SATISFY_ASSERTS (72 runs each, was 32/96), verified by
enumerating solutions via the char config harness and evaluating both host
predicates per solution x size, and end-to-end on the FFM client.
…r validator reject branches Add unit coverage for previously-untested branches (codecov gaps): ClusterLoad computeMasks sparse-metadata (Sparse==1/2) + undeclareSgprs metadata; and the direct reject branches of _validateStreamKMulticast and _validateStreamKClusterReduction that are unreachable through config derivation (StreamK!=3, atomic, force-DP-only, XCC=3, non-[C,1]/non-pow2 cluster, non-gfx1250 ISA, missing HasTDM/HasClusterBarrier, TDMInst==0, multicast/reduction combo).
… to SCOPE_SYS The gfx1250 cluster-reduction split-barrier handshake published peer partials with a SCOPE_DEV global_wb and read them after a SCOPE_DEV global_inv. The split barrier orders execution but carries no memory semantics, so on gfx1250's partitioned L2 a SCOPE_DEV writeback was not guaranteed visible to a peer reading on a different partition -> a rare (~0.1%, one boundary cluster) miscompute at the largest grid (512,512,1,2048) on PGR1 (PGR2 masked it via prefetch timing). Escalate only the cluster-reduction release (peer partial writeback) and the owner's post-barrier acquire invalidate to SCOPE_SYS via a new optional `scope` arg on releaseFence/acquireFence; the global-flag fallback and all non-cluster StreamK keep SCOPE_DEV. InsertClusterBarrierPass.cpp (blob cf772a0) untouched. Confirmed on real gfx1250: MXFP4 and MXFP8 cluster reduction pass (the PGR1 (512,512,1,2048) cases go FAILED->PASSED, no regression). Reduction char + unit tests green (30 passed).
…rain Reflow the producerDrainAnchors.push_back ternary to the repo clang-format style (Google/IndentWidth 4/ColumnLimit 100). Whitespace only; no behavior change.
…sterDim-driven) Derive the StreamK workgroup-cluster roles purely from the cluster shape ClusterDim = [Cs, Ck] on StreamK=3 instead of a user opt-in: * StreamKMulticast iff Cs = ClusterDim[0] > 1 ([C,1] = pure multicast) * StreamKClusterReduction iff Ck = ClusterDim[1] > 1 ([1,C] = pure reduction) Remove the StreamKClusterReduction user/benchmark parameter (now a derived-only internal state key, mirroring StreamKMulticast). Pure reduction migrates from the 1-D [C,1] cluster to a genuine 2-D [1,C] cluster: the launch grid becomes [skGrid/Ck, Ck, 1] and the kernel folds the cluster Y rank into the linear index (StreamKIdx = WorkGroup0*Ck + WorkGroup1), reusing the HW-validated 2-D mechanics. The whole-cluster size C = Cs*Ck now drives the split-barrier span, the ClusterReductionIterCheck predicate (Ck = ClusterDim[1]), the ClusterDimCheck N-tile divisor (pinned to 1 for genuine 2-D clusters), and the host grid launch. Pure multicast [C,1] is preserved byte-identically (Ck==1: no 2-D fold, 1-D launch, C = ClusterDim[0]); the SCOPE_SYS reduction fence is unchanged. Migrate reduction client + designed configs [C,1]+StreamKClusterReduction -> [1,C]; drop the now-invalid StreamKClusterReduction knob from multicast configs. Regenerate reduction/ValidParameters/SolutionClass goldens (param removal + [C,1]->[1,C] expression); multicast + StreamK codegen goldens stay byte-identical.
jaopaulolc
added a commit
that referenced
this pull request
Jul 23, 2026
…terDim-driven) Drop the StreamKClusterKSplit and StreamKClusterReduction user/benchmark parameters: the StreamK workgroup cluster is now described entirely by its shape ClusterDim = [Cs, Ck] (C = Cs*Ck). StreamKMulticast (Cs = ClusterDim[0] > 1) and StreamKClusterReduction (Ck = ClusterDim[1] > 1) are derived-only internal flags. The three canonical expressions: * [C,1] -> pure multicast (byte-identical to the shipped 1-D path) * [1,C] -> pure reduction (2-D launch, matches #9611) * [Cs,Ck] -> factored (both>1) (genuine 2-D cluster: B-multicast x K-split) - streamKClusterFactors returns (Cs, Ck, C, is2D) purely from ClusterDim. - Solution: single _validateStreamKClusterShape (Cs, Ck each pow2, C in [2,16]); multicast/reduction validators compose (no longer mutually exclusive). - Contractions: ClusterDimCheck value[3]=Cs, N-divisor value[4] pinned to 1 for genuine 2-D clusters; ClusterReductionIterCheck value=[DepthU, Ck]. - Host (ContractionSolution.cpp/.hpp): 2-D grid [skGrid/Ck, Ck, 1], Ck from clusterDim.y; remove streamKClusterKSplit field + serialization. - ValidParameters/GlobalParameters: remove both user params (derived-only note). - Migrate factored designed YAML [4,1]+KSplit=2 -> [2,2]; reduction designed + client YAMLs [C,1] -> [1,C]; multicast client YAMLs drop StreamKClusterReduction. - Rewrite unit tests param-free; regenerate reduction/factored char goldens (basename rehash only, err==0); multicast char goldens byte-identical. Preserves the SCOPE_SYS cluster-reduction handshake fence.
…1250
Pure StreamK cluster reduction (ClusterDim [1,C]) produced a flaky,
cluster-aligned column-stripe numerical error on gfx1250 for C>=4 (~8-9%
of fresh runs on 256x512x4096 F8 MT32x16 PGR1 SK3). Capturing wrong
elements showed device ~= (C-1)/C * reference (0.751 measured for C=4):
exactly one peer's partial contribution was dropped -- a peer-contribution
visibility bug, not wrong-slot addressing. Both orientations (256x512 vs
512x256) select identical SK scalars (tiles=256, itersPerTile=16, C=4,
grid=1024, exact even split), so the orientation trigger is purely the
spatial WG->WGP mapping, i.e. whether a cluster straddles WGP boundaries.
Root cause: the intra-cluster split-barrier fast path lets the owner clear
a single s_barrier_signal/wait -3 and then read peer partials. That -3
cluster barrier only orders the owner against HW-cluster-scoped peers; a
C>=4 reduction cluster spans multiple WGPs, so the owner proceeds before a
cross-WGP peer's partial store is globally visible and sums a stale
(missing) partial. A SCOPE_SYS fence around the barrier does not help
because the barrier itself does not wait for the peer store.
Fix: restrict the split-barrier fast path to C<=2 (single WGP, validated
race-free); for C>=4 fall back to the proven per-peer global-flag
handshake, where each peer publishes its partial with a release and raises
its own completion flag and the owner spins on each peer flag with an
acquire before reading that slot -- a genuine cross-peer visibility gate.
The owner's global-flag acquire is escalated to SCOPE_SYS for
cluster-reduction kernels to pair with the peer's existing SCOPE_SYS
release across gfx1250's partitioned L2. Both owner wait and peer signal
gate on the identical compile-time predicate, so the cluster never mixes
barrier signalers with flag setters (deadlock invariant preserved).
Also fix a latent 64-bit workspace slot-offset overflow in
computeWorkspaceSrd: MacroTile0*MacroTile1*bpe * PartialIdx can exceed
2^32 for large SK grids, so the 32-bit product wrapped and aliased a wrong
slot; carry the SMulHIU32 high word into SrdWS+1 instead of only the
lo-add carry.
Validation (gfx1250, GPU3, fresh process per rep):
before: 256x512 C4 83/1000 fail, 512x512 C4 31/1000, 512x512 C8 16/500
after (0 fail, 0 hang):
256x512 C4 520/520
512x512 C4 520/520
512x512 C8 520/520
1024x1024 C4 420/420
1024x1024 C8 420/420
hang-check 512x1024 C8 120/120, 1024x1024 C8 120/120 (0 hang)
suites: sk_mxf8gemm_cluster_reduction 72/72 PASS,
sk_mxf4gemm_cluster_reduction 72/72 PASS,
sk_mxf4gemm_cluster_multicast 296/296 PASS (baseline unaffected).
The secondary C=8 oversubscription spin-wait hang did not reproduce at the
tested grids and no new failures were introduced; it remains a pre-existing
persistent-kernel limitation (grid must fit resident capacity) and is left
documented rather than clamped under deadline.
JIRA ID : N/A
…rrier (no flags) The gfx1250 StreamK [1,C] cluster reduction dropped exactly one peer's partial for C >= 4 at the grid tail (device ~= (C-1)/C * reference, a cluster-aligned column stripe confined to the last ~16 tiles / WG 960-1023). Root cause: the reduction cluster barrier handshake was ASYMMETRIC -- each peer only issued s_barrier_signal -3 (arrive) and then branched away / exited, while only the owner issued s_barrier_wait -3. A peer-signal-only arrive is not a genuine cluster synchronization point on gfx1250's partitioned L2: the owner could observe the barrier satisfied and sum a peer slot before that cross-WGP peer's partial was globally visible. Escalating the release/acquire fences to SCOPE_SYS did not fix it because the defect was structural, not a missing drain. Fix: mirror the HW-validated multicast handshake and make the reduction barrier SYMMETRIC -- every cluster member (owner AND every peer) both arrives (s_barrier_signal -3) AND waits (s_barrier_wait -3) on the same cluster barrier, so no peer proceeds past the rendezvous until the whole cluster has arrived and the owner's paired SCOPE_SYS acquire has a real cross-cluster order to observe. Concretely: add the peer-side clusterReduceWait after clusterReduceSignal in partialsWriteProcedure, and remove the C <= 2 restriction in _streamKClusterBarrierFastEnabled so the barrier path (not a per-peer global-flag fallback) carries peer->owner synchronization for ALL C. The 64-bit slot-offset fix (SMulHIU32 in computeWorkspaceSrd) and the SCOPE_SYS release/acquire on the reduction path are retained. The per-tile global-flag handshake remains ONLY for a cluster that straddles the SK grid boundary (clusterReduceIntraCheck == SCC0); both owner and peers gate on that identical uniform predicate, so a cluster never mixes barrier participants with flag setters (deadlock invariant preserved). Validation (gfx1250, physical GPU3), barrier-only kernels confirmed by disassembly (s_barrier_signal/wait -3 present on owner and peer, peer skips the flag store on the intra path, for both C=4 |3 and C=8 |7): 256x512x4096 [1,4]: 1000/1000 pass, 0 fail, 0 hang 512x512x4096 [1,4]: 1000/1000 pass, 0 fail, 0 hang 512x512x4096 [1,8]: 500/500 pass, 0 fail, 0 hang 1024x1024x4096 [1,4]: 500/500 pass, 0 fail, 0 hang 1024x1024x4096 [1,8]: 500/500 pass, 0 fail, 0 hang hang check 512x1024 [1,8], 1024x1024 [1,8] (timeout 30): 0 hang Full suites (fresh codegen from this source): sk_mxf8gemm_cluster_reduction 72/72 PASS, sk_mxf4gemm_cluster_reduction 72/72 PASS, sk_mxf4gemm_cluster_multicast 296/296 PASS (multicast unregressed). CPU unit tests: test_streamk_cluster_reduction 28/28, char emit test pass. JIRA ID : N/A
…/factored The ported symmetric-cluster-barrier fix folded a 64-bit workspace slot offset (SMulHIU32 high word) into the SHARED computeWorkspaceSrd helper, which the pure-multicast ([C,1]) and non-cluster StreamK paths also emit. That perturbed the multicast kernel codegen (extra s_mul_hi_u32 + an SGPR checkout that cascaded into register/schedule changes), so the multicast asm was no longer byte-identical to the shipped baseline. Gate the 64-bit high-word path on _streamKClusterReductionEnabled (the same predicate every other hunk of the fix already uses): the cluster-reduction / factored (derived Ck>1) path -- which drives the large SK grids the 64-bit offset guards against, and is the only path this fix targets -- keeps the HW-validated 64-bit codegen; the multicast and non-cluster paths take the else branch and emit the original 32-bit offset byte-for-byte. Verified: multicast STINKY_TOTAL_INST_BYTES unchanged and the pre-vs-post multicast diff collapses to the emitter's nondeterministic label/schedule floor (identical to a pre-vs-pre baseline), i.e. the multicast instruction stream is byte-identical to the shipped path. JIRA ID : N/A
Un-gate the 64-bit workspace slot-offset computation in the shared
computeWorkspaceSrd helper so the high-word (SMulHIU32) fold into SrdWS+1
is emitted for every StreamK path that addresses the partial-sum
workspace -- multicast ([C,1]), cluster reduction, factored, and
non-cluster StreamK -- not only the cluster-reduction/factored paths.
The overflow this guards against depends only on the per-slot tile
stride (MacroTile0*MacroTile1*bpe) and the StreamK slot count: the
partials workspace is partialTileSize == tileSize * skGrid regardless of
cluster mode, so the multicast path can silently wrap a 32-bit
slot*stride product on large problems exactly like the reduction path,
aliasing the wrong workspace slot. This intentionally changes the
multicast codegen (byte-identity with the prior 32-bit path is
deliberately dropped); the {basename, err} char goldens are unchanged
and still emit err==0. Real-HW validation pending (owner: user).
JIRA ID : N/A
This was referenced Jul 25, 2026
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
1 participant
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
StreamK GEMM tiles whose partial results are reduced across workgroups synchronize that reduction with a cross-CU global-workspace flag spin-wait. On gfx1250 the workgroups that reduce one tile can instead be co-located into a hardware workgroup cluster and synchronized with an intra-cluster split barrier, taking the spin-wait off the fast path. This PR adds that opt-in cluster-reduction fast path.
This is a cumulative stacked PR (base
develop): for context it also contains the reusable ClusterLoad component + tri-state Multicast (#9599) and the StreamK cooperative multicast /StreamKMulticast(#9603); the work new to this PR is the StreamK workgroup-cluster partial reduction. The first two are reviewed via their own PRs.Technical Details
A new opt-in solution parameter
StreamKClusterReductionplaces a StreamK tile's fixup peers into a single 1-D workgroup cluster (ClusterDim = [C, 1]); the HW WG-id remap (WorkGroup0 = cluster_x*C + wg_x) makes each cluster own the contiguous StreamK index range[c*C, c*C+C)— exactly one tile's peer group.Components/StreamK.py):clusterReduceSignal/clusterReduceWait/clusterReduceIntraCheckemit the wave-0-electeds_barrier_signal -3/s_barrier_wait -3handshake in the fixup epilogue. A uniform intra-cluster predicate commits every member of a cluster to either the barrier fast path or the retained global-flag path together, so a cluster never mixes the two.src/ContractionSolution.cpp+ serialization): thegetSKGridImplreduction override rounds the SK grid toC*tiles(fixed even split,SKItersPerWG = itersPerTile/C) and keeps the launch grid a multiple ofC; the two-tile cluster-reduction kernarg block is emitted only whensk.grid == C*tiles; the field is threaded throughSizeMapping/Contractions.Contractions.py+ContractionProblemPredicates.hpp):ClusterReductionIterCheckrejects at selection time any problem whereceil(K/DepthU)is not a multiple ofC(else the split barrier over-signals); the global-flag reduction is retained as the runtime fallback.Solution.py):_validateStreamKClusterReductiongates the opt-in (SK3, non-atomic, non-DP-only,[C,1]power-of-two cluster, gfx1250 withHasClusterBarrier,TDMInst != 0).StreamKMulticastandStreamKClusterReductionare mutually exclusive (incompatible cluster semantics on one 1-D cluster); both the derivation and the validators enforce it. The factored mode in feat(tensilelite): factored 2-D StreamK cluster mode for gfx1250 #9612 relaxes this into composable orthogonal axes.Test Plan
test_streamk_cluster_reduction_gfx1250_char.py) plus the selection-predicate characterization and C++ gtests.ClusterDim ∈ {[2,1],[4,1],[8,1]}andPrefetchGlobalRead ∈ {1, 2}.Test Result
DID_NOT_SATISFY_ASSERTS, was 32/96): tiles MT32x16/MT32x32/MT64x64,ClusterDim ∈ {[2,1],[4,1],[8,1]},PrefetchGlobalRead ∈ {1,2}, overM,Nmultiples of 512 andKmultiples of 2048.SCOPE_DEVpeer-partial writeback was not guaranteed visible to an owner reading across gfx1250's partitioned L2 — a rare (~0.1%, one boundary cluster) miscompute at (512,512,1,2048) on PGR1 (PGR2 masked it via prefetch timing). Fixed here by escalating only the cluster-reduction release/acquire handshake fences toSCOPE_SYS(the global-flag fallback and non-cluster StreamK keepSCOPE_DEV). Confirmed on real gfx1250: MXFP4 and MXFP8 cluster reduction pass (the two PGR1 cases FAILED→PASSED, no regression); reduction char + unit tests green.Submission Checklist
Risk level
Low–Medium. The feature is strictly opt-in (
StreamKClusterReduction, off by default), gated to gfx1250 SK3 clusters, and retains the global-flag reduction as a runtime fallback, so default paths are unaffected. The one prior PGR1 large-size memory-ordering issue is fixed and HW-confirmed.Update — folded-in symmetric-cluster-barrier reduction fix (HW validation pending)
The gfx1250 StreamK
[1,C]cluster reduction dropped exactly one peer's partial forC >= 4at the grid tail (device ≈(C-1)/C * reference, a cluster-aligned column stripe on the last ~16 tiles / WG 960-1023). The previously-notedSCOPE_DEV→SCOPE_SYSfence escalation was necessary but not sufficient: the defect was structural, not a missing drain. The reduction barrier handshake was asymmetric — each peer onlys_barrier_signal -3(arrive) and then branched away/exited, while only the owners_barrier_wait -3, so the owner could observe the barrier satisfied and sum a peer slot before that cross-WGP peer's partial was globally visible across gfx1250's partitioned L2.Folded in (commits
b7467db→3195f0d→4ceefb5):s_barrier_signal -3) and waits (s_barrier_wait -3) on the same barrier — mirroring the HW-validated multicast handshake — so no peer proceeds past the rendezvous until the whole cluster has arrived and the owner's pairedSCOPE_SYSacquire has a real cross-cluster order to observe. Added the peer-sideclusterReduceWaitafterclusterReduceSignalinpartialsWriteProcedureand removed the interimC <= 2fast-path restriction, so the barrier path (not a per-peer global-flag fallback) carries peer→owner sync for all C. The per-tile global-flag handshake remains only for a cluster that straddles the SK grid boundary (clusterReduceIntraCheck == SCC0); owner and peers gate on the identical uniform predicate (deadlock invariant preserved).MacroTile0*MacroTile1*bpe * partialIdxcan exceed 2³² for large SK grids, wrapping a 32-bits_mul_i32and aliasing the wrong slot; the high word is now carried vias_mul_hi_u32intoSrdWS+1. Scoped to the cluster-reduction path so the multicast[C,1]and non-cluster StreamK paths keep the original 32-bit codegen byte-for-byte.Local validation (CPU, this fix)
pytest Tensile/Tests/unit(streamk/cluster suites incl.test_streamk_cluster_reduction.py,test_streamk_multicast.py): 119 passed.--snapshot-update): 22 passed / 7 snapshots. Multicast char snapshots byte-identical; reduction charerr == 0.streamk_cluster_multicast+_pgr2) is byte-identical to the shipped baseline (STINKY_TOTAL_INST_BYTESunchanged across all 8 kernels; pre-vs-post diff collapses to the emitter's nondeterministic label/schedule floor, identical to a pre-vs-pre baseline).sk_mxf4gemm_cluster_reduction.yaml,sk_mxf8gemm_cluster_reduction.yaml→ 18/18 solutions each, 0 skipped, 0DID_NOT_SATISFY_ASSERTS, err == 0.Functional simulation (FFM) is historically green even when HW fails, so it is treated only as a regression check. Real gfx1250 hardware validation is still pending (owner: reviewer/user).
Update 2 — 64-bit workspace-offset fix now applies to ALL StreamK paths (incl. multicast)
Supersedes the "scoped to the cluster-reduction path" note in the previous update. The 64-bit workspace slot-offset math in the SHARED
computeWorkspaceSrdhelper (thes_mul_hi_u32high word folded intoSrdWS+1) is now emitted universally — multicast[C,1], cluster reduction, factored, and any non-cluster StreamK that uses this helper — rather than gated behind_streamKClusterReductionEnabled.Motivation — could the overflow happen on multicast too? Yes.
The overflow depends only on the per-slot tile stride and the StreamK slot count, not on the cluster mode. The partials workspace is sized (host
ContractionSolution::partialTileSize) aspartialTileSize = MacroTile0*MacroTile1*bpe * skGrid, andcomputeWorkspaceSrdaddresses slotsPartialIdx ∈ [0, skGrid)at byte offsetMacroTile0*MacroTile1*bpe * sPartialIdx. The multicast[C,1]path emits and reads this exact workspace (partial-tile write + owner fixup read), so for a large SK grid the 32-bits_mul_i32product wraps past 2³² and the peer-write / owner-read SRD aliases the wrong workspace slot — the same silent corruption as the reduction path. Concrete example: a 256×256 tile with fp32 partials is 256 KiB/slot, soskGrid ≳ 2³²/262144 = 16384slots makesoffBytes*sPartialIdx ≥ 4 GiBand overflows a 32-bit offset (clusterCmultiplies the effective slot count, which is why big-Cmulticast on large problems is the realistic trigger).Technical Details
Un-gated the one hunk (removed the
if _streamKClusterReductionEnabled(...) / elsesplit; always emit the 64-bits_mul_i32+s_mul_hi_u32+s_add_u32/s_addc_u32sequence). This intentionally drops the multicast byte-identity claimed in the previous update (one extra SGPR + ones_mul_hi_u32per workspace-SRD setup); byte-identity is deliberately no longer a goal.Test Result (CPU, single-process)
{basename, err}char goldens are unchanged and stillerr == 0(a{basename, err}digest does not move when only instructions are added), so no multicast 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(all 4 multicast 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).