feat(tensilelite): reusable ClusterLoad component and tri-state Multicast#9599
Open
jaopaulolc wants to merge 7 commits into
Open
feat(tensilelite): reusable ClusterLoad component and tri-state Multicast#9599jaopaulolc wants to merge 7 commits into
jaopaulolc wants to merge 7 commits into
Conversation
1 task
github-actions
Bot
added
project: hipblaslt
project: hipsparselt
ci:hipsparselt-fast
labels
✅ All Policy Checks Passed
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
jaopaulolc
added
the
rocm:gemm algos
|
🎉 All checks passed! This PR is ready for review. |
jaopaulolc
force-pushed
the
users/jolabega/clusterload-tristate-multicast
branch
from
July 20, 2026 17:46
ddd66d8 to
0bd356e
Compare
1 task
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #9599 +/- ##
===========================================
+ Coverage 69.50% 69.65% +0.15%
===========================================
Files 2770 2765 -5
Lines 452498 451197 -1301
Branches 66620 66464 -156
===========================================
- Hits 314492 314270 -222
+ Misses 117526 116662 -864
+ Partials 20480 20265 -215
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
This was referenced Jul 20, 2026
…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.
jaopaulolc
force-pushed
the
users/jolabega/clusterload-tristate-multicast
branch
from
July 22, 2026 17:28
30bbc1c to
4e3c6bd
Compare
Condense the verbose module/method docstrings on the reusable ClusterLoad component to concise summaries. Behavior-preserving: no code change, unit tests unchanged.
Add unit coverage for the previously-untested enableTDMMetadata path in the ClusterLoad component: computeMasks Sparse==1 (follows-A) / Sparse==2 (follows-B) metadata mask emission and undeclareSgprs metadata SGPR free.
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 cluster-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
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
The TDM multicast ("cluster load") mask machinery was duplicated inline across
KernelWriter,KernelWriterAssembly, andSubtile/SubtileGREmit, andMulticastwas implicitly coupled toClusterDim != [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. This PR extracts that machinery into a reusable component and decouples the multicast opt-in fromClusterDimvia a tri-state parameter.This is PR (i) of a stacked series ("reusable ClusterLoad component + tri-state Multicast"): a byte-exact foundation with no new StreamK cluster feature. It supersedes the earlier #9570 (recreated on a clean, rebased single-commit branch).
Technical Details
ClusterLoadcomponent (Components/ClusterLoad.py,ClusterLoadTDM; registered viaComponent.py+Components/__init__.py). It centralizes the mask value computation (computeMasks), theMulticastMask*SGPR declare/undeclare (declareSgprs/undeclareSgprs), the combined-vs-split A/B topology decision (usesCombinedMask), and the descriptor attach at each load site (applyToDescriptor). It is a byte-exact TDM cluster-load extraction: every method emits assembly identical to the previous inline code, andcomputeMasksreceives the exact SGPR operands the caller already holds so register indices are unchanged. Selection is capability-based (HasTDM+TDMInst == 3), likeTensorDataMoverLoad. It contains no StreamK-cluster symbols (grep -c StreamKMulticast== 0).ClusterLoadTDM.find(...).Multicast(ValidParameters=[-1, 0, 1],GlobalParametersdefault-1):-1= auto (legacy:int(ClusterDim != [1,1] and StreamK == 0)),0= force off,1= force on.Multicastis now an int end-to-end (derivation emits0/1). A pre-load coercion (coerceLegacyMulticastType, scoped to the serialized-state loading path) normalizes the legacy serializedbool(false→0/true→1) for already-shipped library-logic YAMLs, so the strict type gate and emitted msgpack both see an int (avoids astd::bad_castat C++ deserialization). Default-1reproduces the historic ClusterDim-coupled derivation exactly, so every existing YAML (which omitsMulticast) derives byte-identically.Multicastis excluded from the kernel name (SolutionStructs/Naming.pydiscards it): it is a derivation/runtime knob that does not alter emitted assembly, so kernel names stay byte-identical to develop (no_Mtoken) and all_codegenassembly goldens are unchanged.Test Plan
Ran the previously-affected characterization suites, the full
_codegenassembly golden suite, and the new unit gates (CPU-only, gfx942/gfx1250 fixtures):SolutionArms,SolutionClass,ValidParameters,test_PrefetchAcrossPersistent.py,characterization/_codegen,test_cluster_load_component.py,test_multicast_legacy_coercion.py,test_multicast_tristate.py.Test Result
All green: ≈416 passed, 100
_codegensnapshots byte-identical (kernel names unchanged). Develop-owned goldens/tests touched are purely the intentional new intMulticastparam:SolutionArmschar asserts (is True/False→== 1/0) plus_resetrestoring the-1auto sentinel so re-derivation exercises the legacy path.ValidParametersroster/structure snapshots regenerated (purely+Multicast: [-1, 0, 1]).test_PrefetchAcrossPersistent.py(asmCapsHasTDM/kernelTDMInston the PAP-TDM writer mock soClusterLoadTDM.find()matches).FFM: the component is exercised by non-StreamK cluster-load configs (
gemm/gfx12/tdm_multicast_gfx1250.yaml,sparse/gfx1250/spmm_tdm_multicast.yaml) — asm-confirmed multicast masks, clean.gfx1250: (
gemm/gfx12/tdm_multicast_gfx1250.yaml,sparse/gfx1250/spmm_tdm_multicast.yamlpass.Submission Checklist
Risk level
Low. Behavior-preserving refactor plus an opt-in
Multicastparameter that defaults to-1(legacy auto);Multicastis excluded from the kernel name, so emitted assembly and kernel names are byte-identical to develop for all existing configs.Update — 64-bit workspace slot-offset fix (applies to the StreamK cluster/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/tri-state-multicast 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 PRs #9611/#9612 and #9603; on this branch the 64-bit hunk did not previously exist, so it is added here, un-guarded.)Reachability note: although this PR is described as the cluster-load / tri-state-multicast foundation, its
Tensile/Components/StreamK.pyretains the full StreamK partials path —computeWorkspaceSrdis called from six sites (partials write, deferred write, and owner fixup read), and the branch ships a StreamK cluster designed config +test_streamk_cluster_gfx1250_char.pythat emit it. So the 32-bit workspace-address overflow is genuinely reachable here (not pure infra), and the fix is applied rather than skipped.Motivation — the overflow
computeWorkspaceSrdsetSrdWS = AddressWS + offBytes * sPartialIdxwith a 32-bits_mul_i32, whereoffBytes = MacroTile0*MacroTile1*bpe(per-slot tile stride) andsPartialIdx ∈ [0, skGrid); the partials workspace ispartialTileSize = offBytes * skGrid(hostContractionSolution::partialTileSize). For a large SK grid the product exceeds 2³², wraps, and the peer-write / owner-read SRD aliases the wrong workspace slot. It depends only on the tile stride and slot count, not the cluster mode. 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).Technical Details
Compute the high word with
s_mul_hi_u32and fold it (plus the lo-add carry) intoSrdWS+1. Emitted for all StreamK paths; adds one SGPR + ones_mul_hi_u32per workspace-SRD setup.Test Result (CPU, single-process; reused main worktree venv + editable rocisa .so)
{basename, err}char golden is unchanged and stillerr == 0(a{basename, err}digest does not move when only instructions are added), so no snapshot was regenerated.s_mul_hi_u32 … // partials tile offset (high word) for 64-bit SRDfolded intoSrdWS+1).streamk_cluster.yaml1/1,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).