feat(tensilelite): factored 2-D StreamK cluster mode for gfx1250

[jaopaulolc]

JIRA ID : AIHPBLAS-3929

Motivation

The gfx1250 StreamK spatial B-multicast (#9603) and K-split cluster reduction (#9611) are two ways to use one hardware workgroup cluster, and on their own PRs they are mutually exclusive. This PR unifies them into a single factored 2-D cluster mode so one kernel can do both. The cluster is described entirely by its shape ClusterDim = [Cs, Ck] (C = Cs * Ck): Cs spatial B-multicast peers × Ck K-split reduction peers.

Cumulative stacked PR (base develop): for context it also contains #9599, #9603 and #9611; the work new to this PR is the factored mode. The earlier three are reviewed via their own PRs.

Technical Details

Fully ClusterDim-driven, param-free. The StreamKClusterKSplit and StreamKClusterReduction user/benchmark parameters are removed. Cluster behavior is derived solely from the shape:

• Cs = ClusterDim[0], Ck = ClusterDim[1], C = Cs*Ck; StreamKMulticast = (Cs > 1) and StreamKClusterReduction = (Ck > 1) are derived-only internal flags. The three canonical expressions:
  • [C,1] → pure multicast (Cs=C, Ck=1) — 1-D launch, byte-identical to the shipped path;
  • [1,C] → pure reduction (Cs=1, Ck=C) — 2-D launch (matches feat(tensilelite): add StreamK WG-cluster partial reduction for gfx1250 #9611);
  • [Cs,Ck] both > 1 → factored — genuine 2-D cluster (B-multicast along Cs and K-split reduction along Ck).
• Decode/mask: the within-cluster rank decodes "k fastest" (k = StreamKIdx & (Ck-1), s = (StreamKIdx & (C-1)) >> log2(Ck)) after the 2-D StreamKIdx fold. The factored B-multicast mask is maskB_base << k (Cs bits at stride Ck) with a self-only fallback for a partial / M-unaligned cluster. The reduction reuses the whole-cluster cluster-scope -3 split barrier unchanged, with the handshake fences escalated to SCOPE_SYS; owner/peer accounting falls out of skSplit = Ck.
• Host (ContractionSolution.cpp/.hpp): getSKGridImpl unified to skGrid = roundUp(Ck*tiles, C); genuine 2-D launch grid [skGrid/Ck, Ck, 1] with Ck taken from clusterDim.y; ClusterDimCheck aligns on Cs (value[3]) and pins the N-tile divisor value[4]=1 for a 2-D cluster; ClusterReductionIterCheck balances Ck (value=[DepthU, Ck]). The streamKClusterKSplit SizeMapping field + its serialization are removed.
• Validation: a single _validateStreamKClusterShape enforces Cs and Ck each a power of two with C = Cs*Ck ∈ [2,16]; the per-axis multicast/reduction validators compose (the historic mutual-exclusion xor is converted into orthogonal composable axes, not silently dropped).

Test Plan

• Param-free CPU unit tests (derivation, validation matrix, predicate values, clean 2-D emission) across test_streamk_factored_cluster.py, test_streamk_cluster_reduction.py, test_streamk_multicast.py.
• Verify shipped multicast / coop-load / PGR2 gfx1250 _codegen char tests stay byte-identical (run without --snapshot-update); regenerate reduction + factored goldens intentionally.
• Enumerate migrated client configs (err==0, 0 DID_NOT_SATISFY_ASSERTS).
• Real-HW / FFM validation of the new ClusterDim expressions (see below).

Test Result

• 77 CPU unit tests pass across the three StreamK test files.
• Multicast char snapshots byte-identical (multicast + coop-load + PGR2, 5 tests / 2 snapshots pass without --snapshot-update). Reduction + factored goldens regenerated intentionally: the delta is purely the basename solution-hash re-hash from the config expression change ([C,1]+params → [1,C]/[2,2]); every node keeps err: 0, node counts unchanged, no nodes added/removed.
• Config enumeration clean: the 4 migrated gfx1250 client configs derive 120 solutions total (multicast [2,1]/[4,1]/[8,1] → Cs>1,Ck==1; reduction [1,2]/[1,4]/[1,8] → Cs==1,Ck>1), all err==0, correct derived flags, 0 rejections / 0 DID_NOT_SATISFY_ASSERTS.
• SCOPE_SYS cluster-reduction handshake fence preserved (unchanged from feat(tensilelite): add StreamK WG-cluster partial reduction for gfx1250 #9611); the edited host code is clang-format-clean; stinkytofu (rocisa) is untouched.

Config before → after

Mode	Before	After
Pure multicast	ClusterDim=[C,1], StreamKClusterReduction: [0]	ClusterDim=[C,1] (derived: Cs>1, Ck==1)
Pure reduction	ClusterDim=[C,1], StreamKClusterReduction: [1]	ClusterDim=[1,C] (derived: Cs==1, Ck>1)
Factored	ClusterDim=[4,1] + StreamKClusterKSplit: [2]	ClusterDim=[2,2] (derived: Cs>1, Ck>1)

NEW config expressions requiring USER gfx1250 HW re-validation

The following are new expressions of already-validated mechanics but change the launch/derivation, so they need real-HW confirmation:

• Pure reduction [1,C] (MXFP4 + MXFP8) — same reduction mechanics as feat(tensilelite): add StreamK WG-cluster partial reduction for gfx1250 #9611 but now expressed as [1,C] with the genuine 2-D launch (previously [C,1]).
• Factored [Cs,Ck] beyond [2,2] — [2,2] is covered by CPU unit + codegen characterization; larger factorings (e.g. [2,4], [4,2], [2,8], [8,2]) exercising the true composed Cs×Ck path are not yet HW-validated (no client config sets a factored shape > [2,2]).

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Risk level

Medium. [C,1] pure multicast is byte-identical to the shipped path; [1,C] pure reduction matches #9611; the genuinely-new composed Cs×Ck factoring (both axes > 1) is HW-validated only at [2,2] and otherwise covered by CPU unit + codegen characterization pending real-HW confirmation.

Update — folded-in symmetric-cluster-barrier reduction fix (HW validation pending)

Same root cause and fix as folded into #9611, applied here to the factored superset (which includes the reduction path). On gfx1250 the [1,C] reduction (and the reduction axis of a factored [Cs,Ck] cluster, Ck > 1) dropped exactly one peer's partial for C >= 4 at the grid tail because the reduction barrier handshake was asymmetric (peers arrived-and-exited; only the owner waited), so the owner could sum a peer slot before that cross-WGP partial was globally visible across the partitioned L2. Fence escalation alone did not fix it — the defect was structural.

Folded in (commits fb2e2c9 → c3adb03 → 8acc195):

• Symmetric cluster barrier: every member (owner and every peer) both s_barrier_signal -3 (arrive) and s_barrier_wait -3 (wait) on the same cluster barrier (peer-side clusterReduceWait added after clusterReduceSignal); the interim C <= 2 restriction removed so the barrier carries peer→owner sync for all C. Global-flag handshake retained only for a cluster straddling the SK grid boundary; owner/peers gate on the identical predicate (deadlock invariant preserved).
• 64-bit workspace slot offset via s_mul_hi_u32, scoped to the cluster-reduction/factored path (_streamKClusterReductionEnabled, derived Ck > 1) so the pure-multicast [C,1] and non-cluster StreamK paths keep the original 32-bit codegen byte-for-byte.

Local validation (CPU, this fix)

• Unit pytest Tensile/Tests/unit (factored + reduction + multicast suites incl. test_streamk_factored_cluster.py): 134 passed.
• Characterization (run without --snapshot-update): 23 passed / 8 snapshots. Multicast (+ coop-load + PGR2) char snapshots byte-identical; reduction + factored char err == 0.
• Multicast byte-identity confirmed: emitted multicast asm byte-identical to the shipped baseline (STINKY_TOTAL_INST_BYTES unchanged across all 8 kernels; pre-vs-post diff collapses to the emitter's nondeterministic label/schedule floor).
• Config zero-skip enumeration: sk_mxf4gemm_cluster_reduction.yaml, sk_mxf8gemm_cluster_reduction.yaml, sk_mxf4gemm_cluster_factored.yaml, sk_mxf8gemm_cluster_factored.yaml → 18/18 solutions each, 0 skipped, 0 DID_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/factored path" note in the previous update. The 64-bit workspace slot-offset math in the SHARED computeWorkspaceSrd helper (the s_mul_hi_u32 high word folded into SrdWS+1) is now emitted universally — pure-multicast [C,1]/Ck==1, cluster reduction, factored [Cs,Ck], and any non-cluster StreamK that uses this helper — rather than gated behind _streamKClusterReductionEnabled (derived Ck > 1).

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) as partialTileSize = MacroTile0*MacroTile1*bpe * skGrid, and computeWorkspaceSrd addresses slot sPartialIdx ∈ [0, skGrid) at byte offset MacroTile0*MacroTile1*bpe * sPartialIdx. The pure-multicast [C,1] path emits and reads this exact workspace (partial-tile write + owner fixup read), so for a large SK grid the 32-bit s_mul_i32 product wraps past 2³² and the peer-write / owner-read SRD aliases the wrong workspace slot — the same silent corruption as the reduction/factored paths. Concrete example: a 256×256 tile with fp32 partials is 256 KiB/slot, so skGrid ≳ 2³²/262144 = 16384 slots makes offBytes*sPartialIdx ≥ 4 GiB and overflows a 32-bit offset (cluster C = Cs*Ck multiplies the effective slot count, which is why big-C multicast on large problems is the realistic trigger).

Technical Details

Un-gated the one hunk (removed the if _streamKClusterReductionEnabled(...) / else split; always emit the 64-bit s_mul_i32 + s_mul_hi_u32 + s_add_u32/s_addc_u32 sequence). This intentionally drops the multicast byte-identity claimed in the previous update (one extra SGPR + one s_mul_hi_u32 per workspace-SRD setup); byte-identity is deliberately no longer a goal.

Test Result (CPU, single-process)

• Unit + char (streamk / cluster / multicast / reduction / factored suites): 258 passed, 8 snapshots passed. The multicast {basename, err} char goldens are unchanged and still err == 0 (a {basename, err} digest does not move when only instructions are added), so no multicast snapshot was regenerated.
• Confirmed the emitted multicast [C,1], reduction [1,C], and factored [Cs,Ck] kernels all now carry the 64-bit offset (s_mul_hi_u32 … // … high word … for 64-bit SRD folded into SrdWS+1).
• Config zero-skip enumeration: factored streamk_factored_cluster.yaml 4/4, reduction streamk_cluster_reduction.yaml 8/8, multicast streamk_cluster_multicast.yaml 4/4 kernels → err == 0, 0 DID_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).

Update 3 — folded in the ForceDPOnly 2-D [2,2] dual-operand multicast cluster (HW-validated on probe; PR re-validation pending)

Folded the HW-validated (gfx1250) ForceDPOnly 2-D dual-operand multicast capability from the probe branch (users/jolabega/streamk-forcedp-2d-probe, commit a99551ebab) onto this factored-cluster head. The probe was originally based on the older #9612 head (3463fc554f); it is applied here onto the current base so it composes with — and preserves — the symmetric cluster-barrier reduction fix and the universal 64-bit workspace slot-offset (Updates 1 & 2) and the factored path.

What it adds

A genuine 2-D HW cluster ClusterDim = [Cs, Ck] = [2, 2] on a StreamK==3 StreamKForceDPOnly (dense data-parallel, no K-split reduction) kernel where both operands are TDM-multicast:

• Cs = ClusterDim[0] — X-peers on M-adjacent output tiles reuse B (exactly as the shipped 1-D [C,1] ForceDPOnly multicast);
• Ck = ClusterDim[1] — Y-peers on N-adjacent output tiles reuse A (A is genuinely multicast — new).

Here Ck is an N-tiling / A-multicast axis, NOT a K-split reduction axis. This is orthogonal to the factored [Cs,Ck] K-split cluster and is disambiguated purely structurally (StreamKForceDPOnly + both ClusterDim axes > 1), so no new serialized/derived flag is needed.

• Common/Utilities.py: new streamKForceDP2DMulticast structural detector.
• Components/ClusterLoad.py: for ForceDP-2D use the dense A peer count (aPeers = Ck) so maskA multicasts A across the Ck/Y peers (0x5 for [2,2]); maskB unchanged (0x3). [C,1], factored, and reduction paths keep self-only A.
• Components/StreamK.py: new 2-D DP tile-index fold StreamKIdx = batch*(nWG0*nWG1) + WorkGroup1*nWG0 + WorkGroup0 (added as a leading branch before the preserved factored "k fastest" fold) so skIndexToWG lands X-peers M-adjacent (B) and Y-peers N-adjacent (A); streamKMulticastMaskPredicate short-circuits (keep dense masks; no preLoop overwrite).
• SolutionStructs/Solution.py: do not derive StreamKClusterReduction for ForceDP-2D (Cs>1 still derives StreamKMulticast).
• Contractions.py ClusterDimCheck: keep ClusterDim[1] as the N-tile divisor for ForceDP-2D (nWG_y % Ck == 0) instead of pinning to 1.
• src/ContractionSolution.cpp: full M x N tile launch grid [nWG0, nWG1, batch] and getSKGridImpl skGrid = tiles for ForceDP-2D; the factored/reduction 2-D launch [skGrid/Ck, Ck, 1] + roundup preserved in the else branch.

Tests folded in

• Tests/unit/test_streamk_forcedp_2d_probe.py (detector + [2,2] mask math).
• Tests/unit/characterization/_codegen/test_streamk_forcedp_2d_gfx1250_char.py + designed golden _designed/gfx1250/streamk_forcedp_2d.yaml.
• HW-runnable client config renamed from the probe's _probe_forcedp_2d.yaml to Tests/common/streamk/gfx1250/core/sk_mxf4_force_dp_only_cluster_2d_multicast.yaml (consistent with the sibling sk_* naming; collectible by create_tests.py).

Local validation (CPU, single-process, this fold, on the current base)

• Unit: 165 passed across the streamk / cluster / factored / multicast / reduction / ForceDPOnly suites incl. the new test_streamk_forcedp_2d_probe.py (10).
• Characterization (run without --snapshot-update): new ForceDP-2D char passes; the multicast / multicast-PGR2 / coop-load / reduction / factored gfx1250 goldens are byte-identical (8 snapshots, no regeneration) — the ForceDP-2D path is properly gated so it does not perturb the shipped kernels; broader streamk char (r2/r3/r7) 15 passed. No snapshot regeneration was required.
• Config zero-skip enumeration (gfx1250): reduction streamk_cluster_reduction.yaml 8/8, factored streamk_factored_cluster.yaml 4/4, multicast streamk_cluster_multicast.yaml 4/4, and ForceDP-2D streamk_forcedp_2d.yaml 1/1 → err == 0, 0 DID_NOT_SATISFY_ASSERTS.
• ForceDP-2D [2,2] emission confirmed on the current base: maskA = 0x5 and maskB = 0x3 both s_lshl_b32-set and OR'd onto the tdmA/tdmB (and MXSA/MXSB) descriptors (A genuinely multicast); the 2-D DP StreamKIdx = batch*(nWG0*nWG1) + N*nWG0 + M fold is present; the cluster split-barrier is balanced (s_barrier_signal -3 ×4 / s_barrier_wait -3 ×5, matching the shipped 1-D ForceDPOnly multicast); the factored K-split decode/shift are absent.

FFM/functional-sim is historically green regardless of HW correctness and is treated as a regression check only. The ForceDPOnly 2-D [2,2] dual-operand multicast was HW-validated on the probe branch; re-validation on this PR base is still pending (owner: user).

[therock-pr-bot]

✅ All Policy Checks Passed

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

🎉 All policy checks passed!

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[codecov]

Codecov Report

❌ Patch coverage is 77.72512% with 94 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ipblaslt/tensilelite/Tensile/Components/StreamK.py	75.13%	46 Missing and 2 partials ⚠️
...lt/tensilelite/Tensile/SolutionStructs/Solution.py	82.11%	10 Missing and 7 partials ⚠️
.../hipblaslt/tensilelite/src/ContractionSolution.cpp	0.00%	16 Missing and 1 partial ⚠️
...ects/hipblaslt/tensilelite/Tensile/KernelWriter.py	60.00%	2 Missing and 2 partials ⚠️
...blaslt/tensilelite/Tensile/KernelWriterAssembly.py	69.23%	3 Missing and 1 partial ⚠️
...lelite/Tensile/Components/Subtile/SubtileGREmit.py	0.00%	3 Missing ⚠️
...aslt/tensilelite/Tensile/Components/ClusterLoad.py	98.86%	0 Missing and 1 partial ⚠️

❌ 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    #9612      +/-   ##
===========================================
+ Coverage    64.17%   65.76%   +1.60%     
===========================================
  Files         2763     2765       +2     
  Lines       450335   450958     +623     
  Branches     66254    66367     +113     
===========================================
+ Hits        288962   296560    +7598     
+ Misses      140483   133571    -6912     
+ Partials     20890    20827      -63     

Flag	Coverage Δ		*Carryforward flag
TensileLite	34.38% <ø> (+0.03%)	⬆️	Carriedforward from c730f1a
TensileLite-CPP	38.29% <0.00%> (+0.02%)	⬆️	Carriedforward from c730f1a
TensileLite-Unit	42.21% <77.73%> (+6.69%)	⬆️	Carriedforward from c730f1a
hipBLAS	90.62% <ø> (ø)		Carriedforward from c730f1a
hipBLASLt	34.63% <ø> (ø)
hipCUB	82.68% <ø> (ø)		Carriedforward from c730f1a
hipDNN	86.34% <ø> (ø)		Carriedforward from c730f1a
hipFFT	48.90% <ø> (ø)		Carriedforward from c730f1a
hipRAND	76.12% <ø> (ø)		Carriedforward from c730f1a
hipSOLVER	69.18% <ø> (ø)		Carriedforward from c730f1a
hipSPARSE	86.27% <ø> (ø)		Carriedforward from c730f1a
rocBLAS	47.95% <ø> (ø)		Carriedforward from c730f1a
rocFFT	48.51% <ø> (ø)		Carriedforward from c730f1a
rocRAND	57.01% <ø> (ø)		Carriedforward from c730f1a
rocSOLVER	76.84% <ø> (ø)		Carriedforward from c730f1a
rocSPARSE	72.49% <ø> (ø)		Carriedforward from c730f1a
rocThrust	91.36% <ø> (ø)		Carriedforward from c730f1a

*This pull request uses carry forward flags. Click here to find out more.

Files with missing lines	Coverage Δ
...slt/tensilelite/Tensile/Common/GlobalParameters.py	92.49% <ø> (+5.14%)	⬆️
.../hipblaslt/tensilelite/Tensile/Common/Utilities.py	100.00% <ø> (+0.39%)	⬆️
...aslt/tensilelite/Tensile/Common/ValidParameters.py	90.20% <ø> (ø)
...rojects/hipblaslt/tensilelite/Tensile/Component.py	61.48% <100.00%> (+7.00%)	⬆️
...pblaslt/tensilelite/Tensile/Components/__init__.py	100.00% <ø> (ø)
...ects/hipblaslt/tensilelite/Tensile/Contractions.py	78.82% <100.00%> (+2.90%)	⬆️
...aslt/tensilelite/Tensile/SolutionStructs/Naming.py	82.58% <100.00%> (+0.41%)	⬆️
...aslt/tensilelite/Tensile/Components/ClusterLoad.py	98.86% <98.86%> (ø)
...lelite/Tensile/Components/Subtile/SubtileGREmit.py	57.38% <0.00%> (-0.10%)	⬇️
...ects/hipblaslt/tensilelite/Tensile/KernelWriter.py	55.21% <60.00%> (+31.45%)	⬆️
... and 4 more

... and 24 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[therock-pr-bot]

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

• python -m pip install pre-commit
• pre-commit install
• pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.