[tensilelite] Support TDM iterate mode

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

@@ -541,7 +541,8 @@
     # TDM-multicast data. Requires ClusterDim != [1, 1] and TDMInst != 0;
     # False: standard per-WG barriers, no inter-WG synchronization.
     {"ClusterBarrier": [ False ]},
-    {"HalfPLR": [0]}
+    {"HalfPLR": [0]},
+    {"TDMIterateMode": [0]}
 ]
 
 # dictionary of defaults comprised of default option for each parameter

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

@@ -107,6 +107,7 @@ def getRequiredParametersMin() -> set:
         'TransposeLDSMetadata',
         'TDMInst',
         "TDMSplit",
+        "TDMIterateMode",
         "MXScaleFormat",
         "MXLoadInst",
         'SwInstructionPrefetch',

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

@@ -1045,7 +1045,14 @@ def makeValidMatrixInstructions():
     # 1: Use PLR 0.5 for A
     # 2: Use PLR 0.5 for B
     # 3: Use PLR 0.5 for both A and B
-    "HalfPLR": [0, 1, 2, 3]
+    "HalfPLR": [0, 1, 2, 3],
+    # Enable iterate-mode TDM
+    # -1: Auto. Enable per-tensor when LBSPP > 1024 B (exceeds pad_interval encoding).
+    # 0: Disabled
+    # 1: Use iterate-mode for A
+    # 2: Use iterate-mode for B
+    # 3: Use iterate-mode for both A and B
+    "TDMIterateMode": [-1, 0, 1, 2, 3]
 }
 
 newMIValidParameters = {

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

@@ -255,21 +255,21 @@ def initOperands(self, group0: int | str, group1: int | str, group2: Optional[in
 
         return mod
 
-    def setDataType(self, dtype: DataType, group1: str | int, isMetadata: bool = False) -> Module:
-        mod = Module()
-        dataSizeOp = None
-
+    @staticmethod
+    def dataSizeShift(dtype: DataType, isMetadata: bool = False) -> int:
         if isMetadata or dtype.isInt8() or dtype.is8bitFloat() or dtype.isFloat4() or dtype.is6bitFloat():
-            dataSizeOp = 0
-        elif dtype.isBFloat16() or dtype.isHalf():
-            dataSizeOp = 1
-        elif dtype.isSingle():
-            dataSizeOp = 2
-        elif dtype.isDouble():
-            dataSizeOp = 3
-
-        assert dataSizeOp is not None
+            return 0
+        if dtype.isBFloat16() or dtype.isHalf():
+            return 1
+        if dtype.isSingle():
+            return 2
+        if dtype.isDouble():
+            return 3
+        raise AssertionError(f"unsupported dtype for TDM data_size: {dtype}")
 
+    def setDataType(self, dtype: DataType, group1: str | int, isMetadata: bool = False) -> Module:
+        mod = Module()
+        dataSizeOp = self.dataSizeShift(dtype, isMetadata)
         mod.add(SAndB32(sgpr(group1), sgpr(group1), hex(0xFFFCFFFF), "Reset data_size"))
         mod.add(SOrB32(sgpr(group1), sgpr(group1), hex(dataSizeOp << 16), f"Set data_size to {dataSizeOp}"))
         return mod
@@ -311,8 +311,12 @@ def incrementGlobalAddr(self, group0: int | str, sgprIncrement: int | str) -> Mo
 
     def setIterationEnabled(self, group1, enabled: bool) -> Module:
         mod = Module()
-        mask = 1 << 19 if enabled else 0xFFF7FFFF
-        mod.add(SAndB32(sgpr(group1), sgpr(group1), hex(mask)))
+        if enabled:
+            mod.add(SOrB32(sgpr(group1), sgpr(group1), hex(1 << 19),
+                           "set iterate_enable (D# Group 1 bit 19)"))
+        else:
+            mod.add(SAndB32(sgpr(group1), sgpr(group1), hex(0xFFF7FFFF),
+                            "clear iterate_enable (D# Group 1 bit 19)"))
         return mod
 
     def resetTensorDimForTail(self, group1: int | str, sgprTail: int, tdmDescIdx: int, writer: "KernelWriterAssembly", constShifter: int=0, isMXS: bool=False, isSparseTrack: bool=False) -> Module:

projects/hipblaslt/tensilelite/Tensile/SolutionStructs/Solution.py

@@ -2542,11 +2542,50 @@ def depthUIteration(
         state["_DepthUMXSB"] = depthUB // state["ProblemType"]["MXBlockB"]
       state["_DepthUMetadata"] = depthUM# internal
 
-      # Auto-derived VW must keep LdsBlockSizePerPad <= 1024 bytes
-      # (TDM pad_interval cap of 7). Halve VW until it fits.
+      # fp6 doesn't support LDS padding yet.
+      for tc in ["A", "B"]:
+        if state["ProblemType"]["MacDataType%s" % tc].is6bitFloat() and (
+            state["LdsPad%s" % tc] != 0 or state["LdsBlockSizePerPad%s" % tc] != 0):
+          reject(state, printRejectionReason,
+                 f"fp6 MacDataType{tc}: LdsPad{tc} and LdsBlockSizePerPad{tc} must be 0")
+          return
+
+      iterModeMask = state["TDMIterateMode"]
       if state["TDMInst"] and state["EnableMatrixInstruction"] and not state["ProblemType"]["Sparse"]:
+        # Stage 1: decide iterate-mode per tensor.
+        if iterModeMask == -1:
+          state.pop("_TDMIterateModeA", None)
+          state.pop("_TDMIterateModeB", None)
+          for tc in ["A", "B"]:
+            if not state["UnrollMajorLDS%s" % tc]:
+              continue
+            vw = state["VectorWidth%s" % tc]
+            bpe_tc = state["ProblemType"]["MacDataType%s" % tc].numBytes()
+            lbspp = roundUpToNearestMultiple(int(state["_DepthU%s" % tc] * bpe_tc * vw), 256)
+            if lbspp > 1024:
+              state["_TDMIterateMode%s" % tc] = True
+          # Resolve -1 into the concrete mask for kernel naming. backupValues
+          # restores -1 before each DepthU candidate, so this overwrite is safe.
+          state["TDMIterateMode"] = (1 if state.get("_TDMIterateModeA", False) else 0) | \
+                                    (2 if state.get("_TDMIterateModeB", False) else 0)
+        else:
+          if iterModeMask & 1:
+            if not state["UnrollMajorLDSA"]:
+              reject(state, printRejectionReason, "TDMIterateMode bit for A set but UnrollMajorLDSA is False")
+              return
+            state["_TDMIterateModeA"] = True
+          if iterModeMask & 2:
+            if not state["UnrollMajorLDSB"]:
+              reject(state, printRejectionReason, "TDMIterateMode bit for B set but UnrollMajorLDSB is False")
+              return
+            state["_TDMIterateModeB"] = True
+
+        # Stage 2: for non-iterate tensors, halve auto-derived VW until LBSPP
+        # fits the pad_interval 1024 B limit.
         multiple = 256
         for tc in ["A", "B"]:
+          if state.get("_TDMIterateMode%s" % tc, False):
+            continue
           if not state.get("_inputVW%s_was_auto" % tc, False):
             continue
           if not state["UnrollMajorLDS%s" % tc]:
@@ -2557,15 +2596,13 @@ def depthUIteration(
           tmpBpe = state["ProblemType"]["MacDataType%s" % tc].numBytes()
           depthU_tc = state["_DepthU%s" % tc]
           origVw = vw
-          if tmpBpe == 0.75:
-              continue # FP6 not yet supported; skip this tc, still process the other
           while vw > 1:
             candidate = roundUpToNearestMultiple(int(depthU_tc * tmpBpe * vw), multiple)
             dwords = candidate // 4
             if dwords > 0 and (dwords & (dwords - 1)) == 0:
               pad_interval = int(math.log2(dwords)) - 1
               if pad_interval <= 7:
-                break  # this VW produces a valid LdsBlockSizePerPad
+                break  # current VW fits pad_interval encoding
             vw //= 2
           if vw != origVw:
             state["VectorWidth%s" % tc] = vw
@@ -2770,15 +2807,20 @@ def checkLdsBlockSizePerPadForTDM(ldsBlockSizePerPadA: int, ldsBlockSizePerPadB:
           pads = {"A": ldsBlockSizePerPadA, "B": ldsBlockSizePerPadB, "MXSA": ldsBlockSizePerPadMXSA, "MXSB": ldsBlockSizePerPadMXSB}
           for tc, val in pads.items():
             if val == 0: continue
+            # A/B in iterate-mode bypass the pad_interval encoding; skip their
+            # check. MXSA/MXSB do not support iterate-mode, so their LBSPP
+            # must still satisfy the pad_interval constraints.
+            if tc in ("A", "B") and state.get("_TDMIterateMode%s" % tc, False):
+              continue
             dwords = val // 4
             if dwords == 0 or (dwords & (dwords - 1)) != 0:
               reject(state, printRejectionReason, f"LdsBlockSizePerPad{tc}={val}: val//4={dwords} must be a positive power of 2 for TDM hardware encoding")
               return
             pad_interval = TensorDataMoverLoad.calPadInterval(val)
             if pad_interval > 7:
-              # MXSA/MXSB share the host A/B operand's VectorWidth.
-              vwName = "VectorWidthA" if tc.endswith("A") else "VectorWidthB"
-              reject(state, printRejectionReason, f"pad_interval=(log2(LdsBlockSizePerPad//4)-1)={pad_interval} should be smaller than or equal to 7 for ldsBlockSizePerPad{tc}={val}. Please reduce DepthU or {vwName}")
+              reject(state, printRejectionReason,
+                     f"LdsBlockSizePerPad{tc}={val} exceeds TDM padding (1024B). "
+                     f"Set TDMIterateMode or reduce DepthU / VectorWidth.")
 
       def calcMXSLdsBlockSizePerPad(tc: str, lrvw: int) -> int:
         LdsBlockSizePerPad = state["LdsBlockSizePerPad%s"%tc]

projects/hipblaslt/tensilelite/Tensile/Tests/common/gemm/gfx12/mxf4_gfx1250.yaml

@@ -64,12 +64,13 @@ BenchmarkProblems:
           - [16, 16, 128, 1,  1,  1, 1,  2, 2]
           - [16, 16, 128, 1,  1,  1, 2,  2, 1]
           - [16, 16, 128, 1,  1,  1, 4,  2, 1]
-          - [16, 16, 128, 1,  1,  1, 4,  4, 1]
+          - [16, 16, 128, 1,  1,  1, 8,  4, 1]
           - [16, 16, 128, 1,  1,  2, 2,  2, 2]
           - [16, 16, 128, 1,  1,  4, 4,  2, 2]
           - [16, 16, 128, 1,  1,  8, 8,  2, 2]
-        - DepthU: [128, 256]
+        - DepthU: [128, 256, 512]
         - TDMInst: [3]
+        - TDMIterateMode: [0, 3]
         - LocalReadVectorWidth: [32]
         - ScheduleIterAlg: [0]
         - PrefetchGlobalRead: [2]
@@ -364,11 +365,11 @@ BenchmarkProblems:
         - GlobalReadVectorWidthA: [32]
         - GlobalReadVectorWidthB: [32]
         - InnerUnroll: [1]
-        - LdsBlockSizePerPadA: [0]
-        - LdsBlockSizePerPadB: [0]
+        - LdsBlockSizePerPadA: [-1]
+        - LdsBlockSizePerPadB: [-1]
         - TransposeLDS: [1]
-        - LdsPadA: [0]
-        - LdsPadB: [0]
+        - LdsPadA: [-1]
+        - LdsPadB: [-1]
         - WaveSeparateGlobalReadA: [0]
         - WaveSeparateGlobalReadB: [0]
         - 1LDSBuffer: [0]
@@ -435,11 +436,11 @@ BenchmarkProblems:
         - GlobalReadVectorWidthA: [32]
         - GlobalReadVectorWidthB: [32]
         - InnerUnroll: [1]
-        - LdsBlockSizePerPadA: [0]
-        - LdsBlockSizePerPadB: [0]
+        - LdsBlockSizePerPadA: [-1]
+        - LdsBlockSizePerPadB: [-1]
         - TransposeLDS: [1]
-        - LdsPadA: [0]
-        - LdsPadB: [0]
+        - LdsPadA: [-1]
+        - LdsPadB: [-1]
         - WaveSeparateGlobalReadA: [0]
         - WaveSeparateGlobalReadB: [0]
         - 1LDSBuffer: [0]

projects/hipblaslt/tensilelite/Tensile/Tests/common/gemm/gfx12/mxf8_gfx1250.yaml

@@ -66,6 +66,7 @@ BenchmarkProblems:
           - [16, 16, 128, 1,  1,  8, 4,  2, 2]
         - DepthU: [256]
         - TDMInst: [3]
+        - TDMIterateMode: [0, 3]
         - TransposeLDS: [-1]
         - LdsPadA: [-1]
         - LdsPadB: [-1]
@@ -141,6 +142,7 @@ BenchmarkProblems:
         - HalfPLR: [0,1,2,3]
         - DepthU: [256]
         - TDMInst: [3]
+        - TDMIterateMode: [-1,0,1,2,3]
         - TransposeLDS: [-1]
         - LdsPadA: [-1]
         - LdsPadB: [-1]
@@ -184,4 +186,4 @@ BenchmarkProblems:
           - Exact: [257, 256, 1, 2048] #edge-m
           - Exact: [256, 257, 1, 2048] #edge-n
         - ActivationArgs:
-          - [Enum: none]
+          - [Enum: none]

projects/hipblaslt/tensilelite/Tensile/Tests/common/gemm/gfx12/tdm_gfx1250.yaml

@@ -49,7 +49,7 @@ BenchmarkProblems:
         - MatrixInstruction:
           - [16, 16, 32, 1,  1,  1, 1,  1, 1 ]
           - [16, 16, 32, 1,  1,  1, 1,  2, 2 ]
-          # - [16, 16, 32, 1,  1,  8, 8,  2, 2 ]
+          - [16, 16, 32, 1,  1,  8, 4,  2, 2 ]
           # - [16, 16, 32, 1,  1,  1, 1,  4, 1 ]
           # - [16, 16, 32, 1,  1,  2, 2,  2, 2 ]
           # - [16, 16, 32, 1,  1,  2, 2,  2, 2]
@@ -61,6 +61,7 @@ BenchmarkProblems:
         - LdsBlockSizePerPadB: [-1]
         - LdsPadMetadata: [0]
         - TDMInst: [3]
+        - TDMIterateMode: [-1, 0, 3]
         - LDSTrInst: [False]
         - StaggerU: [0, 4, 16]
         - StaggerUMapping: [0, 1]
@@ -72,8 +73,8 @@ BenchmarkProblems:
         - GlobalSplitUAlgorithm: [MultipleBuffer]
         - GlobalReadVectorWidthA: [-1]
         - GlobalReadVectorWidthB: [-1]
-        - VectorWidthA: [1]
-        - VectorWidthB: [1]
+        - VectorWidthA: [-1]
+        - VectorWidthB: [-1]
         - LocalReadVectorWidth: [8]
         - 1LDSBuffer: [0]
         - SourceSwap: [0]
@@ -216,6 +217,7 @@ BenchmarkProblems:
         - LdsBlockSizePerPadB: [-1]
         - LdsPadMetadata: [0]
         - TDMInst: [3]
+        - TDMIterateMode: [-1, 0, 2]
         - LDSTrInst: [True]
         - StaggerU: [0, 4, 16]
         - StaggerUMapping: [0, 1]
@@ -424,14 +426,16 @@ BenchmarkProblems:
         - MatrixInstruction:
           - [16, 16, 4, 1,  1,  2, 2,  1, 1 ]
           - [16, 16, 4, 1,  1,  1, 1,  2, 2 ]
-        - DepthU: [16]
+          - [16, 16, 4, 1,  1,  4, 4,  1, 1 ]
+        - DepthU: [16, 128]
         - TransposeLDS: [-1]
         - LdsPadA: [-1]
         - LdsPadB: [-1]
         - LdsBlockSizePerPadA: [-1]
         - LdsBlockSizePerPadB: [-1]
         - LdsPadMetadata: [0]
         - TDMInst: [3]
+        - TDMIterateMode: [0, 3]
         - LDSTrInst: [False]
         - StaggerU: [0, 4, 16]
         - StaggerUMapping: [0, 1]
@@ -443,8 +447,8 @@ BenchmarkProblems:
         - GlobalSplitUAlgorithm: [MultipleBuffer]
         - GlobalReadVectorWidthA: [-1]
         - GlobalReadVectorWidthB: [-1]
-        - VectorWidthA: [1]
-        - VectorWidthB: [1]
+        - VectorWidthA: [-1]
+        - VectorWidthB: [-1]
         #- LocalReadVectorWidth: [8]
         - 1LDSBuffer: [0]
         - SourceSwap: [0]
@@ -467,5 +471,6 @@ BenchmarkProblems:
           - Exact: [1, 1, 1, 127] # 7DU + tail
           - Exact: [64, 64, 1, 128]
           - Exact: [64, 64, 1, 129]
+          - Exact: [64, 64, 1, 512]
         - ActivationArgs:
           - [Enum: none]