feat(tensilelite): reusable ClusterLoad component and tri-state Multicast#9599
Open
jaopaulolc wants to merge 2 commits into
Open
feat(tensilelite): reusable ClusterLoad component and tri-state Multicast#9599jaopaulolc wants to merge 2 commits into
jaopaulolc wants to merge 2 commits into
Conversation
1 task
✅ 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. |
…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 20, 2026 17:46
ddd66d8 to
0bd356e
Compare
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 #9599 +/- ##
===========================================
+ Coverage 64.16% 64.56% +0.40%
===========================================
Files 2757 2758 +1
Lines 449555 449605 +50
Branches 66143 66150 +7
===========================================
+ Hits 288443 290261 +1818
+ Misses 139965 138783 -1182
+ Partials 21147 20561 -586
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
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
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.Multicastderivation andClusterLoad.pyreference no StreamK-cluster symbols; the auto (-1) path is exactlyint(ClusterDim != [1,1] and StreamK == 0), andClusterBarrieris gated onStreamK == 0.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.