feat(tensilelite): add StreamK WG-cluster partial reduction for gfx1250

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

@@ -571,6 +571,11 @@
     {"StreamKAtomic": [0]},
     {"StreamKXCCMapping": [0]},
     {"StreamKFixupTreeReduction": [0]},
+    # NOTE: StreamKMulticast and StreamKClusterReduction are derived-only internal
+    # state keys (like ClusterBarrier), auto-enabled by Solution.py purely from the
+    # cluster shape ClusterDim = [Cs, Ck] for StreamK==3 (StreamKMulticast iff Cs>1,
+    # StreamKClusterReduction iff Ck>1); they are deliberately NOT benchmark/default
+    # parameters here. See docs/design/streamk-wg-clusters.md.
     {"DebugStreamK": [0]},
     {"DebugPersistentKernelLoopForever": [False]},
     {"ActivationFused": [True]},
@@ -612,6 +617,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/Utilities.py

@@ -365,6 +365,28 @@ def clusterEnabled(clusterDim):
     """True when a workgroup cluster is requested (ClusterDim [x, y] is not [1, 1])."""
     return (clusterDim[0] * clusterDim[1]) != 1
 
+def streamKClusterFactors(d):
+    """Return (Cs, Ck, C, is2D) for a StreamK workgroup cluster (ClusterDim-driven).
+
+    The StreamK cluster is fully described by ClusterDim = [Cs, Ck]; there are no
+    user factoring/reduction knobs. Cs = ClusterDim[0] is the spatial B-multicast
+    axis (X), Ck = ClusterDim[1] is the K-split reduction axis (Y), and the total
+    cluster is C = Cs * Ck. ``is2D`` is True exactly when Ck > 1, i.e. when the
+    launch grid must be genuinely 2-D ([skGrid/Ck, Ck, 1]) and the linear StreamK
+    index folds the cluster Y rank in (StreamKIdx = WorkGroup0*Ck + WorkGroup1).
+
+    Config expressions:
+      * [C, 1] -> Cs=C, Ck=1  : pure multicast   (1-D launch, byte-identical)
+      * [1, C] -> Cs=1, Ck=C  : pure reduction    (2-D launch)
+      * [Cs,Ck]-> both > 1     : factored          (2-D launch; factored branch only)
+
+    ``d`` may be a kernel or a solution ``state`` dict; both expose "ClusterDim".
+    See docs/design/streamk-wg-clusters.md.
+    """
+    cd = d["ClusterDim"]
+    cs, ck = cd[0], cd[1]
+    return cs, ck, cs * ck, (ck > 1)
+
 def log2(x):
     return int(log(x, 2) + 0.5)
 

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

@@ -837,6 +837,20 @@ def makeValidMatrixInstructions():
     # 0: use linear reduction
     # 1: use tree reduction
     "StreamKFixupTreeReduction": [0, 1],
+    # NOTE: StreamKMulticast and StreamKClusterReduction (the gfx1250 StreamK DP
+    # cooperative cluster-load / TDM B-multicast fast path and the cluster split-
+    # barrier K-reduction fast path) are intentionally NOT valid/benchmark
+    # parameters. They are DERIVED-ONLY internal state keys (see the ClusterBarrier
+    # precedent): Solution.assignProblemIndependentDerivedParameters auto-enables
+    # them for StreamK==3 purely from the cluster shape ClusterDim = [Cs, Ck]
+    # (StreamKMulticast iff Cs = ClusterDim[0] > 1, StreamKClusterReduction iff
+    # Ck = ClusterDim[1] > 1) -- so [C,1] = pure multicast, [1,C] = pure reduction.
+    # _validateStreamKMulticast / _validateStreamKClusterReduction then hard-reject
+    # any cluster config that cannot satisfy their constraints. They must not be
+    # user/YAML-settable, so they have no entry here (checkParametersAreValid would
+    # otherwise accept them as a fork/constant param). They still serialize to C++
+    # via Contractions.SizeMapping (streamKMulticast / streamKClusterReduction, read
+    # with d.get()). See docs/design/streamk-wg-clusters.md.
     # 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 +1131,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