boringmorning · Sep 21, 2026 · Aug 11, 2026 · Sep 9, 2026 · Sep 9, 2026 · Sep 10, 2026
@@ -356,7 +356,10 @@ def makeValidMatrixInstructions():
     # 0: disable
     # 1: prefetch one load tile (MTxDepthU) ahead of PrefetchGlobalRead
     # 2: prefetch two load tiles (MTxDepthU) ahead of PrefetchGlobalRead
-    # Currently we do not support GSU, StaggerU, StreamK and general batch. May remove these limitations in the future.
+    # Currently we do not support StaggerU (forced off), general batch, 6-bit float,
+    # or Stream-K other than DP-first (StreamK==3). GSU is supported; with a
+    # workgroup cluster it forces GlobalSplitUWorkGroupMappingRoundRobin on so the
+    # cluster's peers share a K chunk. May remove these limitations in the future.
     "PrefetchGL2": [0, 1, 2],
     # MatrixInstruction Only
     # If set ClusterLocalRead, each iteration dedicated vgprBuffer for localRead

@@ -1,16 +1,26 @@
 from ..Component import GL2Prefetch
 from ..Common import INDEX_CHARS
-from typing import Mapping
-from rocisa.code import Module
+from typing import Mapping, Optional
+from rocisa.code import Module, Label
 from rocisa.instruction import SMulI32, SAddU64, VMovB32, VAddU32, VAddCOU32, \
     VAddCCOU32, VAddNCU64, VLShiftRightB32, VMulLOU32, VMulHIU32, GlobalPrefetchB8, \
-    VCmpGtU32, VCndMaskB32, SSubI32, SMovB32, SAddU32, SAddCU32
+    VCmpGtU32, VCndMaskB32, SSubI32, SMovB32, SAddU32, SAddCU32, SAndB32, SBranch, \
+    SCBranchSCC1, SCMovB32, SLShiftRightB32
 from rocisa.container import sgpr, vgpr, RegisterContainer, VCC, GLOBALModifiers, ContinuousRegister
 from rocisa.functions import vectorMultiply64Bpe, scalarMultiplyBpe, vectorStaticDivideAndRemainder, \
     scalarStaticRemainder
 from rocisa.enum import TemporalHint, CacheScope
 from math import log2, ceil
 
+# Bit 15 of the packed GSU kernel argument selects how the summation loop is cut
+# up between the GSU groups, and the two layouts need different start offsets and
+# strides:
+#   GSUC == 0: the groups interleave every DepthU, so group g starts at iteration
+#              g and then steps GSU iterations at a time.
+#   GSUC == 1: each group owns a contiguous run, so group g starts after every
+#              lower group's run and then steps one iteration at a time.
+GSUC_BIT = 0x8000
+
 class GL2PrefetchLoad(GL2Prefetch):
     asmCaps = {"HasGlobalPrefetch": True}
     globalModifiers = GLOBALModifiers(th=TemporalHint.TH_NT, scope=CacheScope.SCOPE_SE)
@@ -55,7 +65,81 @@ def init(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping):
         tp["gl2nc"] = tp["gl2ncp"] * tp["gl2ncc"]
         tp["gl2nl"] = max(1, ceil(tp["gl2nc"] / numCooperativeThreads))
 
+    def isGSUEnabled(self, kernel: Mapping) -> bool:
+        """True when the kernel emits the GSUOn paths, so GSU/GSUSumIdx are live."""
+        return kernel["GlobalSplitU"] > 0 or kernel["GlobalSplitU"] == -1
+
+    def calculateGSUIterOffset(self, writer: "KernelWriterAssembly", kernel: Mapping, \
+                               dstSgprIdx: int, tmpSgprRes: ContinuousRegister) -> Module:
+        """Unroll iteration at which this workgroup's GSU chunk starts.
+
+        One unroll iteration is one DepthU step for every tensor, so the result is
+        tensor independent and callers scale it by each tensor's own per-iteration
+        byte increment. That keeps the chunk-layout math in one place instead of
+        repeating it per tensor and per TLU/MX/metadata layout.
+
+        Clobbers GSUSumIdx+1, which the GSU component also uses as scratch;
+        computeLoadSrd and calculateLoopNumIterGsu both recompute it later.
+        """
+        mod = Module("gl2 prefetch GSU start iteration")
+        depthU: int = kernel["DepthU"]
+        gsucLabel = Label(writer.labels.getNameInc("GL2PrefetchGSUC"), "")
+        gsucLabelEnd = Label(writer.labels.getNameInc("GL2PrefetchGSUC_End"), "")
+
+        mod.addComment("gl2 prefetch GSU start iteration")
+        mod.add(SAndB32(dst=sgpr(dstSgprIdx), src0=sgpr("GSU"), src1=hex(GSUC_BIT), \
+            comment="SCC = (GSUC == 1) ?"))
+        mod.add(SCBranchSCC1(labelName=gsucLabel.getLabelName(), comment="branch if GSUC == 1"))
+        mod.add(SMovB32(dst=sgpr(dstSgprIdx), src=sgpr("GSUSumIdx"), \
+            comment="interleaved chunks: startIter = GSUSumIdx"))
+        mod.add(SBranch(gsucLabelEnd.getLabelName()))
+        mod.add(gsucLabel)
+        mod.add(SLShiftRightB32(dst=sgpr(dstSgprIdx), shiftHex=int(log2(depthU)), src=sgpr("SizesSum"), \
+            comment="numIter = SizesSum / DepthU(%u)" % depthU))
+        mod.add(writer.calculateLoopNumIterOffsetGsu(kernel, dstSgprIdx, tmpSgprRes))
+        mod.add(SMovB32(dst=sgpr(dstSgprIdx), src=sgpr(tmpSgprRes.idx), \
+            comment="contiguous chunks: startIter = accumulated iters of lower groups"))
+        mod.add(gsucLabelEnd)
+        return mod
+
+    def applyGSUChunk(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping, \
+                      gsuIterSgpr: int, baseSgprIdx: int, tmpSgprIdx: int, tmpVgprIdx: int) -> Module:
+        """Move the prefetch base onto this workgroup's GSU chunk and widen the step.
+
+        Both are multiples of the one-DepthU increment setIncrement produced, so
+        the layout only has to be decoded once (calculateGSUIterOffset). The start
+        offset must consume the unscaled increment, so the scaling happens after
+        it and before the PGR pre-skip, which already steps by whole chunks.
+        """
+        mod = Module("gl2 prefetch GSU chunk offset")
+        tc: str = tp["tensorChar"]
+        incName: str = f"GL2PrefetchInc{tc}"
+
+        mod.addComment(f"gl2 prefetch GSU chunk offset of {tc}")
+        mod.addModuleAsFlatItems(writer.s_mul_u64_u32(
+            sgpr(tmpSgprIdx), sgpr(tmpSgprIdx + 1),
+            sgpr(gsuIterSgpr), sgpr(incName),
+            tmpVgprIdx, comment="gsuOffset = startIter * inc"))
+        mod.add(SAddU64(sgpr(baseSgprIdx, 2), sgpr(baseSgprIdx, 2), sgpr(tmpSgprIdx, 2), \
+            comment="skip to this WG's GSU chunk"))
+        # Widen the step to the chunk stride. Kept 32-bit to mirror GlobalReadIncs
+        # on the real load path (GSU.graIncrements), which the prefetch has to
+        # track: a stride that overflows 32 bits is already broken there.
+        mod.add(SAndB32(dst=sgpr(tmpSgprIdx), src0=sgpr("GSU"), src1=writer.gsuMaskHex(kernel), \
+            comment="Restore GSU"))
+        mod.add(SAndB32(dst=sgpr(tmpSgprIdx + 1), src0=sgpr("GSU"), src1=hex(GSUC_BIT), \
+            comment="SCC = (GSUC == 1) ?"))
+        mod.add(SCMovB32(dst=sgpr(tmpSgprIdx), src=1, comment="stride stays DepthU if GSUC == 1"))
+        mod.add(SMulI32(sgpr(incName), sgpr(incName), sgpr(tmpSgprIdx), \
+            comment="addr increment *= GSU chunk stride"))
+        return mod
+
     def setIncrement(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping) -> Module:
+        """Bytes the prefetch address advances for one DepthU step along K.
+
+        This is the *unscaled* step. Under GSU, applyGSUChunk widens it to the
+        workgroup's chunk stride once the start offset has consumed it.
+        """
         mod = Module()
         tc: str = tp["tensorChar"]
         tIdx: int = tp['idx']
@@ -73,7 +157,13 @@ def setIncrement(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapp
             mod.add(SMovB32(dst=sgpr(f"GL2PrefetchInc{tc}"), src=round(du * bpe), comment="addr increment"))
         return mod
 
-    def calculateStartAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping) -> Module:
+    def calculateStartAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping, \
+                           gsuIterSgpr: Optional[int] = None) -> Module:
+        """Compute this workgroup's prefetch start addresses.
+
+        gsuIterSgpr holds the shared GSU chunk start iteration from
+        calculateGSUIterOffset, or None when the kernel has no GSU paths.
+        """
         mod = Module()
         globalPrefetchSize: int = writer.states.regCaps["GlobalPrefetchSize"]
         tc: str = tp["tensorChar"]
@@ -223,6 +313,13 @@ def calculateStartAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp
                         sgpr("WorkGroup2"), sgpr(tmpSgprIdx2),
                         tmpVgprIdx, comment="batch offset * wg2"))
                     mod.add(SAddU64(sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx2, 2)))
+            # GSU chunk offset. Must precede the PGR pre-skip: it consumes the
+            # unscaled increment and leaves behind the chunk-strided one that the
+            # pre-skip and every in-loop increment then use.
+            if gsuIterSgpr is not None:
+                mod.add(self.applyGSUChunk(writer, kernel, tp, gsuIterSgpr, \
+                    tmpSgprIdx0, tmpSgprIdx2, tmpVgprIdx))
+
             # skip PGR loads (uses GSU-adjusted increment)
             if kernel["PrefetchGlobalRead"] > 0:
                 if kernel["PrefetchGlobalRead"] > 1:
@@ -265,4 +362,4 @@ def incrementAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Map
             mod.add(VAddCOU32(vgpr(addrName), VCC(), vgpr(addrName), inc))
             mod.add(VAddCCOU32(vgpr(addrNameHi), VCC(), vgpr(addrNameHi), 0, VCC()))
 
-        return mod
+        return mod
@@ -3679,6 +3679,17 @@ def noLoadLoopBody( self, kernel, tensorParametersA, tensorParametersB, pack, pa
             self.codes.localWriteMXSB = Module()
           self.codes.globalReadMetadata = StructuredModule() # empty
 
+        # The gl2 prefetch modules are built for the unroll loop (_loopBody) and
+        # would otherwise be re-emitted here by makeSchedule, since nothing else
+        # rebuilds them. There is no load left for this workgroup to run ahead of
+        # once it reaches a no-load loop, and the address is only kept in range by
+        # the guards on the sites that do have one: the pre-loop skips itself when
+        # counterL <= PGR, and the unroll loop is only entered above that. Emitting
+        # them here reaches the start address without either guard, which for a
+        # short GSU chunk is PGR iterations past the end of K.
+        self.codes.gl2PrefetchIncrement = Module()
+        self.codes.gl2Prefetch = Module()
+
         callMakeSchedule = not isNGLL or kernel["ExpandPointerSwap"] or UnrollLoopSwapGlobalReadOrder or isDTVAB or \
           (kernel["PrefetchGlobalRead"] >= 3 and isNGLL) or \
           (self.states.doPackPreSchedulingNextLoop and isNGLL)

@@ -20753,9 +20753,21 @@ def gl2PrefetchCalcAddr(self, kernel, tPA, tPB) -> Module:
     if kernel["enableTDMMetadata"]:
       tpList.append(tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"])
 
-    for tp in tpList:
-      mod.add(comp.setIncrement(self, kernel, tp))
-      mod.add(comp.calculateStartAddr(self, kernel, tp))
+    if not comp.isGSUEnabled(kernel):
+      for tp in tpList:
+        mod.add(comp.setIncrement(self, kernel, tp))
+        mod.add(comp.calculateStartAddr(self, kernel, tp))
+      return mod
+
+    # The GSU chunk starts at the same unroll iteration for every tensor, so derive
+    # it once here and let each tensor scale it by its own per-iteration increment.
+    with self.allocTmpSgpr(3, tag="gl2PrefetchCalcAddr_gsu") as tmpSgprRes:
+      gsuIterSgpr = tmpSgprRes.idx
+      offsetTmp = ContinuousRegister(idx=tmpSgprRes.idx + 1, size=2)
+      mod.add(comp.calculateGSUIterOffset(self, kernel, gsuIterSgpr, offsetTmp))
+      for tp in tpList:
+        mod.add(comp.setIncrement(self, kernel, tp))
+        mod.add(comp.calculateStartAddr(self, kernel, tp, gsuIterSgpr))
     return mod
 
   def gl2PrefetchIssueLoad(self, kernel, tPA, tPB) -> Module:

@@ -6016,12 +6016,29 @@ def calcEpilogueTurns(factorDims: List) -> int:
         return
       # TODO: support staggerU if needed
       _disableRuntimeStaggerU(state)
-      # TODO: support GSU if needed
-      state["InternalSupportParams"]["SupportUserGSU"] = False
-      if state["GlobalSplitU"] > 1 or state["GlobalSplitU"] == -1:
-        reject(state, printRejectionReason, "Currently PrefetchGL2 does not support GSU")
+      # A cluster's workgroups cooperate on one folded prefetch footprint, taking
+      # their slot from WorkGroup{i} % ClusterDim (Components/GL2Prefetch.py), so
+      # they have to agree on the K chunk. The cluster owns ClusterDim[1]
+      # consecutive raw y values, and the default mapping makes the group the
+      # fast axis of that y (GSUSumIdx = wg1 % GSU), which spreads peers across
+      # chunks and folds their divided WorkGroup1 onto fewer tiles. Round-robin
+      # makes it the slow axis instead (GSUSumIdx = wg1 / NumWorkGroups1) so a
+      # cluster shares one group and spans distinct tiles. Force it whenever a
+      # cluster is live rather than keying off the tuned GlobalSplitU, since
+      # SupportUserGSU is left on above and GSU can arrive at runtime; both
+      # branches are always emitted and picked off the GSU sgpr, so this only
+      # moves the host-side default and is inert at GSU<=1.
+      # Residual: a cluster straddling a tilesN boundary still splits across two
+      # groups, the same perf-only boundary-cluster caveat GL2Prefetch.init notes
+      # for padded WGs -- the prefetch only warms cache, so a coverage gap costs
+      # bandwidth, never correctness.
+      if state["ClusterDim"] != [1, 1]:
+        state["GlobalSplitUWorkGroupMappingRoundRobin"] = True
+      # 256 bytes is not multiple of 6 bits, causing math calculations errors
+      if state["ProblemType"]["DataTypeA"].is6bitFloat() or state["ProblemType"]["DataTypeB"].is6bitFloat():
+        reject(state, printRejectionReason, "PrefetchGL2 does not support 6-bit float")
         return
-      if state["StreamK"] != 0 and state["StreamK"] != 3:
+      if state["StreamK"] not in [0, 3]:
         reject(state, printRejectionReason, "PrefetchGL2 only supports DP-first (StreamK==3) Stream-K")
         return
       if state["ProblemType"]["Batched"] and not state["ProblemType"]["StridedBatched"]:

@@ -202,7 +202,7 @@ BenchmarkProblems:
         - PrefetchGL2: [1,2]
         - HalfPLR: [2]
         - StoreRemapVectorWidth: [0]
-        - GlobalSplitU: [1,3]
+        - GlobalSplitU: [1,2,3]
         - GlobalSplitUAlgorithm: [MultipleBuffer]
         - GlobalReadVectorWidthA: [-1]
         - GlobalReadVectorWidthB: [-1]
@@ -218,7 +218,7 @@ BenchmarkProblems:
         - WavefrontSize: [32]
         - ExpandPointerSwap: [false]
         - GroupLoadStore: [False]
-        - ClusterDim: [[4,4]]
+        - ClusterDim: [[1, 1],[2, 2],[4,4]]
       BenchmarkForkParameters:
       JoinParameters:
       BenchmarkJoinParameters:
@@ -278,7 +278,7 @@ BenchmarkProblems:
         - WavefrontSize: [32]
         - ExpandPointerSwap: [false]
         - GroupLoadStore: [False]
-        - ClusterDim: [[4,4]]
+        - ClusterDim: [[1, 1],[2, 2],[4,4]]
       BenchmarkForkParameters:
       JoinParameters:
       BenchmarkJoinParameters: