feat(tensilelite): StreamK cooperative cluster loads

projects/hipblaslt/tensilelite/Tensile/Common/GlobalParameters.py

@@ -571,6 +571,9 @@
     {"StreamKAtomic": [0]},
     {"StreamKXCCMapping": [0]},
     {"StreamKFixupTreeReduction": [0]},
+    # NOTE: StreamKMulticast is a derived-only internal state key (like
+    # ClusterBarrier), auto-enabled by Solution.py for StreamK==3 + ClusterDim
+    # clusters; it is deliberately NOT a benchmark/default parameter here.
     {"DebugStreamK": [0]},
     {"DebugPersistentKernelLoopForever": [False]},
     {"ActivationFused": [True]},
@@ -612,6 +615,9 @@
     # [1, 1] disables clustering. Non-[1, 1] enables Multicast so workgroups within
     # a cluster can share data loaded via TDM-multicast, reducing redundant global reads.
     {"ClusterDim": [[1, 1]]},
+    # Multicast tri-state (see ValidParameters). Default -1 = legacy auto, so
+    # every existing YAML (which omits Multicast) derives identically to before.
+    {"Multicast": [-1]},
     {"HalfPLR": [0]},
     {"TDMIterateMode": [0]}
 ]

projects/hipblaslt/tensilelite/Tensile/Common/ValidParameters.py

@@ -837,6 +837,16 @@ def makeValidMatrixInstructions():
     # 0: use linear reduction
     # 1: use tree reduction
     "StreamKFixupTreeReduction": [0, 1],
+    # NOTE: StreamKMulticast (the gfx1250 StreamK DP cooperative cluster-load /
+    # TDM B-multicast fast path) is intentionally NOT a valid/benchmark
+    # parameter. It is a DERIVED-ONLY internal state key (see the ClusterBarrier
+    # precedent): Solution.assignProblemIndependentDerivedParameters auto-enables
+    # it for StreamK==3 + ClusterDim != [1,1] (the "bare StreamK cluster"
+    # collapse), and _validateStreamKMulticast then hard-rejects any cluster
+    # config that cannot satisfy its constraints. It must not be user/YAML-
+    # settable, so it has no entry here (checkParametersAreValid would otherwise
+    # accept it as a fork/constant param). It still serializes to C++ via
+    # Contractions.SizeMapping (streamKMulticast, read with d.get()).
     # Debug settings for stream-k kernels to disable parts of the kernel
     #   Bit 0: Don't generate fixup code
     #   Bit 1: Don't generate write to partials code
@@ -1117,6 +1127,14 @@ def makeValidMatrixInstructions():
     # Cluster dimension. Clusters have up to 16 work-groups in a cluster, but each work-group in a
     # cluster runs on a separate WGP.
     "ClusterDim": validClusterDimensions,
+    # Multicast (cluster load) control. Decouples the TDM-multicast opt-in from
+    # ClusterDim so barrier-only clustering and cooperative-load clustering
+    # compose independently. Tri-state:
+    #   -1 = auto (legacy: ClusterDim != [1,1] implies Multicast, except the
+    #        StreamK cluster paths); reproduces historic behavior exactly.
+    #    0 = force multicast off.
+    #    1 = force multicast on (independent of the ClusterDim coupling).
+    "Multicast": [-1, 0, 1],
     # Enable PLR 0.5 to save vgprs
     # 0: Disabled
     # 1: Use PLR 0.5 for A

projects/hipblaslt/tensilelite/Tensile/Component.py

@@ -312,6 +312,11 @@ class GL2Prefetch(Component):
     GL2 Prefetch
     """
 
+class ClusterLoad(Component):
+    """
+    Cluster (multicast) TDM load: multicast-mask compute + descriptor attach.
+    """
+
 # Importing here allows auto-registry of components in the Components directory.
 # Each file must be listed in __all__ in Components/__init__.py
 # "noqa" prevents linter from complaining here.

projects/hipblaslt/tensilelite/Tensile/Components/ClusterLoad.py

@@ -0,0 +1,175 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+"""Cluster (multicast) TDM load component.
+
+Centralizes the multicast ("cluster load") mask machinery (value compute,
+``MulticastMask*`` SGPR declare/undeclare, combined-vs-split topology decision,
+and per-load-site descriptor attach) that was previously duplicated across
+``KernelWriter``/``KernelWriterAssembly``/``SubtileGREmit``. Behavior-preserving:
+every method emits byte-identical assembly, receiving the SGPR operands the
+caller already holds rather than re-allocating. Capability-selected
+(``HasTDM`` + ``TDMInst == 3``), like ``TensorDataMoverLoad``.
+"""
+
+from ..Component import ClusterLoad
+from ..Common import clusterEnabled
+from typing import Mapping
+from rocisa.code import Module, Label
+from rocisa.container import sgpr
+from rocisa.instruction import SLShiftLeftB32, SMulI32, SBitcmp1B32, SCBranchSCC1, SBranch
+
+
+class ClusterLoadTDM(ClusterLoad):
+    asmCaps = {"HasTDM": True}
+    kernel  = {"TDMInst": 3}
+
+    def __call__(self, writer: "KernelWriterAssembly", kernel: Mapping):
+        # Abstract-satisfying no-op, mirrors TensorDataMoverLoad.__call__.
+        pass
+
+    # -- topology decision ---------------------------------------------------
+
+    def usesCombinedMask(self, kernel: Mapping) -> bool:
+        """True when the single-parity combined ``MulticastMask`` applies.
+
+        Single source of truth for the combined-vs-split decision. Subtile and
+        StreamK DP multicast both need the split A/B masks: subtile issues A and
+        B on every wave (no wave-parity split), and StreamK broadcasts only B
+        across the [C,1] cluster while A stays per-workgroup -- the combined
+        parity mask would be wrong for both.
+        """
+        if kernel.get("StreamKMulticast", 0):
+            return False
+        tdmA: bool = kernel["enableTDMA"]
+        tdmB: bool = kernel["enableTDMB"]
+        return tdmA and tdmB and kernel["NumWaves"] > 1 and not kernel.get("UseSubtileImpl")
+
+    def maskSgprName(self, kernel: Mapping, tc: str, *, subtile: bool = False,
+                     waveSeparated: bool = False) -> str:
+        """Resolve the multicast-mask SGPR name.
+
+        Wave-separated (non-subtile) uses the combined ``"MulticastMask"``;
+        dense/subtile and StreamK multicast use the split ``f"MulticastMask{tc}"``
+        (any ``MXS`` prefix stripped). StreamK forces the split name so B never
+        resolves to the never-declared combined SGPR.
+        """
+        if kernel.get("StreamKMulticast", 0):
+            string = tc.removeprefix("MXS") if tc.startswith("MXS") else tc
+            return f"MulticastMask{string}"
+        if waveSeparated and not subtile:
+            return "MulticastMask"
+        string = tc.removeprefix("MXS") if tc.startswith("MXS") else tc
+        return f"MulticastMask{string}"
+
+    def cooperativeThreadPartition(self, kernel: Mapping, tc: str) -> int:
+        """Cooperating-workgroup count for ``tc``: ClusterDim[1] (A) / [0] (B)."""
+        subTc: str = tc[-1]
+        return kernel["ClusterDim"][1] if subTc == "A" else kernel["ClusterDim"][0]
+
+    # -- SGPR declare / undeclare -------------------------------------------
+
+    def declareSgprs(self, writer: "KernelWriter", kernel: Mapping) -> None:
+        """Allocate the ``MulticastMask*`` SGPRs (lift of KernelWriter)."""
+        if not kernel["Multicast"]:
+            return
+        tdmM: bool = kernel["enableTDMMetadata"]
+        if self.usesCombinedMask(kernel):
+            writer.defineSgpr("MulticastMask", 1)
+        else:
+            writer.defineSgpr("MulticastMaskA", 1)
+            writer.defineSgpr("MulticastMaskB", 1)
+        if tdmM:
+            writer.defineSgpr("MulticastMaskMetadata", 1)
+
+    def undeclareSgprs(self, writer: "KernelWriter", kernel: Mapping) -> Module:
+        """Free the ``MulticastMask*`` SGPRs (lift of KernelWriter)."""
+        mod = Module()
+        if not (kernel["Multicast"] and kernel["TDMInst"] != 0):
+            return mod
+        tdmM: bool = kernel["enableTDMMetadata"]
+        if self.usesCombinedMask(kernel):
+            mod.add(writer.undefineSgpr("MulticastMask"))
+        else:
+            mod.add(writer.undefineSgpr("MulticastMaskA"))
+            mod.add(writer.undefineSgpr("MulticastMaskB"))
+        if tdmM:
+            mod.add(writer.undefineSgpr("MulticastMaskMetadata"))
+        return mod
+
+    # -- mask value computation ---------------------------------------------
+
+    def computeMasks(self, writer: "KernelWriterAssembly", kernel: Mapping, *,
+                     sgprWgX: int, sgprWgY: int, sgprNWgX: int, sTmp: int) -> Module:
+        """Compute the multicast mask value(s) into the ``MulticastMask*`` SGPRs.
+
+        Verbatim lift of the ``defineAndResources`` mask compute; the caller
+        passes the operands it already holds (``sgprWgX``/``sgprWgY``/``sgprNWgX``
+        and ``sTmp`` whose ``+4`` slot is scratch) so the output is byte-identical.
+        """
+        mod = Module()
+        if not kernel["Multicast"]:
+            return mod
+        mod.addComment0("Calculate multicast mask")
+
+        maskA = 1
+        for idx in range(kernel["ClusterDim"][1]):
+            maskA |= (1 << (idx * kernel["ClusterDim"][0]))
+
+        maskB = (1 << kernel["ClusterDim"][0]) - 1
+
+        if kernel["enableTDMMetadata"]:
+            if kernel["ProblemType"]["Sparse"] == 1:
+                mod.add(SLShiftLeftB32(dst=sgpr("MulticastMaskMetadata"), shiftHex=sgpr(sgprWgX), src=hex(maskA),\
+                                        comment="Setting metadata mask (follows sparse A)"))
+            elif kernel["ProblemType"]["Sparse"] == 2:
+                mod.add(SMulI32(dst=sgpr(sTmp+4), src0=sgpr(sgprWgY), src1=sgpr(sgprNWgX),\
+                                comment="Shift factor: wg_y * nwg_x (metadata)"))
+                mod.add(SLShiftLeftB32(dst=sgpr("MulticastMaskMetadata"), shiftHex=sgpr(sTmp+4), src=hex(maskB),\
+                                        comment="Setting metadata mask (follows sparse B)"))
+
+        if self.usesCombinedMask(kernel):
+            setMulticastMaskLblOdd = Label(f"setMulticastMask_OddWave", "")
+            setMulticastMaskLblEven = Label(f"setMulticastMask_EvenWave", "")
+            setMulticastMaskLblEnd = Label(f"setMulticastMaskEnd", "")
+
+            mod.add(SBitcmp1B32(sgpr("WaveIdx"), 0, "Check parity of wId"))
+            mod.add(SCBranchSCC1(setMulticastMaskLblOdd.getLabelName(), "Jump if wId is odd"))
+
+            mod.add(setMulticastMaskLblEven)
+            mod.add(SLShiftLeftB32(dst=sgpr("MulticastMask"), shiftHex=sgpr(sgprWgX), src=hex(maskA),\
+                                    comment="Setting maskA for even wave"))
+            mod.add(SBranch(setMulticastMaskLblEnd.getLabelName()))
+            mod.add(setMulticastMaskLblOdd)
+            mod.add(SMulI32(dst=sgpr(sgprWgY), src0=sgpr(sgprWgY), src1=sgpr(sgprNWgX),\
+                            comment="Shift factor: wg_y * nwg_x"))
+            mod.add(SLShiftLeftB32(dst=sgpr("MulticastMask"), shiftHex=sgpr(sgprWgY), src=hex(maskB),\
+                                    comment="Setting maskB for odd wave"))
+            mod.add(setMulticastMaskLblEnd)
+
+        else:
+            mod.add(SLShiftLeftB32(dst=sgpr("MulticastMaskA"), shiftHex=sgpr(sgprWgX), src=hex(maskA),\
+                                    comment="Setting maskA"))
+
+            mod.add(SMulI32(dst=sgpr(sgprWgY), src0=sgpr(sgprWgY), src1=sgpr(sgprNWgX),\
+                            comment="Shift factor: wg_y * nwg_x"))
+            mod.add(SLShiftLeftB32(dst=sgpr("MulticastMaskB"), shiftHex=sgpr(sgprWgY), src=hex(maskB),\
+                                    comment="Setting maskB"))
+        return mod
+
+    # -- descriptor attach ---------------------------------------------------
+
+    def applyToDescriptor(self, writer: "KernelWriterAssembly", kernel: Mapping,
+                          group1: int | str, tc: str, *, subtile: bool = False,
+                          waveSeparated: bool = False) -> Module:
+        """OR the multicast mask into descriptor ``Group1[word0]``.
+
+        Folds the ``Multicast and enableCluster`` gate, mask-name choice, and the
+        ``SOrB32`` attach; returns an empty ``Module`` when the gate is not met.
+        """
+        from .TensorDataMover import TensorDataMoverLoad
+        mod = Module()
+        if kernel["Multicast"] and clusterEnabled(kernel["ClusterDim"]):
+            mask = self.maskSgprName(kernel, tc, subtile=subtile, waveSeparated=waveSeparated)
+            tdm = TensorDataMoverLoad.find(writer)
+            mod.add(tdm.setMulticastMask(group1, mask, writer))
+        return mod