[hipblaslt] update tensilelite code-gen part for 64bit offset support to add arguments appropriately and to use only two extra SGPRs

Author: shbae

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp

@@ -131,6 +131,8 @@ rocblaslt_status rocblaslt_matmul_impl(const rocblaslt_handle       handle,
     hipDataType        scale_type    = matmul_descr->scale_type;
 
     // Others
+    // Use strided_batch=true for kernel selection (StridedBatched=true kernels with SupportUserArgs)
+    // The actual batch mode is tracked via problem.batchMode() for argument passing
     bool strided_batch = true;
     bool grouped_gemm  = false;
 
@@ -343,6 +345,8 @@ rocblaslt_status rocblaslt_gemm_create_cpp_impl(const rocblaslt_handle
     void*              amaxD         = matmul_descr->amaxD;
 
     // Others
+    // Use strided_batch=true for kernel selection (StridedBatched=true kernels with SupportUserArgs)
+    // The actual batch mode is tracked via problem.batchMode() for argument passing
     bool strided_batch = true;
     bool grouped_gemm  = false;
 

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

@@ -163,15 +163,6 @@ def __call__(self, writer) -> SignatureBase:
         if kernel["ProblemType"]["Sparse"]:
             signature.addArg("MetaData", SVK.SIG_GLOBALBUFFER, "void" , "generic")
 
-        # Batch offset support for general batched mode (pointer array)
-        # Only for non-strided, non-grouped GEMM
-        if not kernel["ProblemType"]["StridedBatched"] and not kernel["ProblemType"]["GroupedGemm"]:
-            signature.addArg("batchOffsetD", SVK.SIG_VALUE, "u64")
-            signature.addArg("batchOffsetC", SVK.SIG_VALUE, "u64")
-            signature.addArg("batchOffsetA", SVK.SIG_VALUE, "u64")
-            signature.addArg("batchOffsetB", SVK.SIG_VALUE, "u64")
-            userArgumentsInfo.gemmArgumentSize += 32  # 4 offsets * 8 bytes each
-
         if kernel["StreamK"] > 0 and kernel["StreamKAtomic"] == 0:
             signature.addArg("AddressWS", SVK.SIG_GLOBALBUFFER, cptValueType, "generic")
             signature.addArg("AddressFlags", SVK.SIG_GLOBALBUFFER, dstValueType, "generic")
@@ -235,6 +226,16 @@ def __call__(self, writer) -> SignatureBase:
                 signature.addArg("skTiles",                        SVK.SIG_VALUE, "u32")
                 userArgumentsInfo.gemmArgumentSize += 8
 
+        # Batch offset support for general batched mode (pointer array)
+        # Placed after core GEMM args (strides, alpha/beta, StreamK)
+        # Enabled for all SupportUserArgs kernels (offsets are 0 when not in General Batched mode)
+        if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
+            signature.addArg("batchOffsetD", SVK.SIG_VALUE, "u64")
+            signature.addArg("batchOffsetC", SVK.SIG_VALUE, "u64")
+            signature.addArg("batchOffsetA", SVK.SIG_VALUE, "u64")
+            signature.addArg("batchOffsetB", SVK.SIG_VALUE, "u64")
+            userArgumentsInfo.gemmArgumentSize += 32  # 4 offsets * 8 bytes each
+
         if kernel["ProblemType"]["UseScaleAB"]:
             signature.addArg("AddressScaleA", SVK.SIG_GLOBALBUFFER, cptValueType, "generic")
             signature.addArg("AddressScaleB", SVK.SIG_GLOBALBUFFER, cptValueType, "generic")

projects/hipblaslt/tensilelite/Tensile/KernelWriter.py

@@ -8460,11 +8460,8 @@ def vgprAllocationImplSubtile():
       self.defineSgpr("StridesMetadata", self.states.m.numSgprStrides)
 
     # Batch offset support for general batched GEMM (pointer array mode)
-    if not kernel["ProblemType"]["StridedBatched"] and not kernel["ProblemType"]["GroupedGemm"]:
-      self.defineSgpr("BatchOffsetD", 2)  # 64-bit offset requires 2 SGPRs
-      self.defineSgpr("BatchOffsetC", 2)  # 64-bit offset requires 2 SGPRs
-      self.defineSgpr("BatchOffsetA", 2)  # 64-bit offset requires 2 SGPRs
-      self.defineSgpr("BatchOffsetB", 2)  # 64-bit offset requires 2 SGPRs
+    # Offsets are loaded on-demand from kernel arguments to avoid using persistent SGPRs
+    # No SGPR allocation here - offsets loaded directly when applying to addresses
 
     # for packed batches without stride restrictions need to do something different here
     assert sorted(kernel["PackedC0IdxChars"]+kernel["PackedC1IdxChars"]) == \

projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py

@@ -1898,9 +1898,13 @@ def loadBatchedAddress(self, kernel, Batch, tmpSgprResource: ContinuousRegister)
           module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address"))
           module.add(SLoadB64(dst=sgpr("AddressD", 2), base=sgpr("AddressD",2), soffset=sgpr(tmpSgpr), comment="load global buffer D address"))
           # Apply batch offset to AddressD for general batched mode
-          if not kernel["ProblemType"]["StridedBatched"] and not kernel["ProblemType"]["GroupedGemm"]:
-            module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array before updating offset"))
-            module.add(SAddU64(dst=sgpr("AddressD", 2), src0=sgpr("AddressD", 2), src1=sgpr("BatchOffsetD", 2), comment="add batch offset to D address"))
+          # Load offset on-demand from kernel args (no persistent SGPRs needed)
+          if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
+            with self.allocTmpSgpr(2) as tmpSgprOffset:
+              module.add(self.argLoader.loadKernArg(tmpSgprOffset.idx, "KernArgAddress", sgprOffset=hex(self.states.batchOffsetDKernArgOffset), dword=2))
+              module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address and Batch Offset Loads"))
+              module.add(SAddU32(dst=sgpr("AddressD+0"), src0=sgpr("AddressD+0"), src1=sgpr(tmpSgprOffset.idx), comment="add batch offset to D address (low)"))
+              module.add(SAddCU32(dst=sgpr("AddressD+1"), src0=sgpr("AddressD+1"), src1=sgpr(tmpSgprOffset.idx+1), comment="add batch offset to D address (high)"))
 
       # Only load C buffer address if Beta is used and potentially non-zero
       if kernel["ProblemType"]["UseBeta"]:
@@ -1913,9 +1917,13 @@ def loadBatchedAddress(self, kernel, Batch, tmpSgprResource: ContinuousRegister)
             module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address"))
             module.add(SLoadB64(dst=sgpr("AddressC", 2), base=sgpr("AddressC",2), soffset=sgpr(tmpSgpr), comment="load global buffer C address"))
             # Apply batch offset to AddressC for general batched mode
-            if not kernel["ProblemType"]["StridedBatched"] and not kernel["ProblemType"]["GroupedGemm"]:
-              module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array before updating offset"))
-              module.add(SAddU64(dst=sgpr("AddressC", 2), src0=sgpr("AddressC", 2), src1=sgpr("BatchOffsetC", 2), comment="add batch offset to C address"))
+            # Load offset on-demand from kernel args (no persistent SGPRs needed)
+            if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
+              with self.allocTmpSgpr(2) as tmpSgprOffset:
+                module.add(self.argLoader.loadKernArg(tmpSgprOffset.idx, "KernArgAddress", sgprOffset=hex(self.states.batchOffsetCKernArgOffset), dword=2))
+                module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address and Batch Offset Loads"))
+                module.add(SAddU32(dst=sgpr("AddressC+0"), src0=sgpr("AddressC+0"), src1=sgpr(tmpSgprOffset.idx), comment="add batch offset to C address (low)"))
+                module.add(SAddCU32(dst=sgpr("AddressC+1"), src0=sgpr("AddressC+1"), src1=sgpr(tmpSgprOffset.idx+1), comment="add batch offset to C address (high)"))
 
         module.add(endCheckLabel)
 
@@ -1935,11 +1943,19 @@ def loadBatchedAddress(self, kernel, Batch, tmpSgprResource: ContinuousRegister)
         module.add(SLoadB64(dst=sgpr("AddressA", 2), base=sgpr("AddressA",2), soffset=sgpr(tmpSgpr), comment="load global buffer A address"))
         module.add(SLoadB64(dst=sgpr("AddressB", 2), base=sgpr("AddressB",2), soffset=sgpr(tmpSgpr), comment="load global buffer B address"))
         # Apply batch offset to AddressA and AddressB for general batched mode
-        if not kernel["ProblemType"]["StridedBatched"] and not kernel["ProblemType"]["GroupedGemm"]:
-          module.add(SWaitCnt(kmcnt=1, comment="Wait for the AddressA Load from the Pointer Array before updating offset"))
-          module.add(SAddU64(dst=sgpr("AddressA", 2), src0=sgpr("AddressA", 2), src1=sgpr("BatchOffsetA", 2), comment="add batch offset to A address"))
-          module.add(SWaitCnt(kmcnt=0, comment="Wait for the AddressB Load from the Pointer Array before updating offset"))
-          module.add(SAddU64(dst=sgpr("AddressB", 2), src0=sgpr("AddressB", 2), src1=sgpr("BatchOffsetB", 2), comment="add batch offset to B address"))
+        # Load offsets on-demand from kernel args (no persistent SGPRs needed)
+        if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
+          with self.allocTmpSgpr(2) as tmpSgprOffset:
+            # Load and apply offset A
+            module.add(self.argLoader.loadKernArg(tmpSgprOffset.idx, "KernArgAddress", sgprOffset=hex(self.states.batchOffsetAKernArgOffset), dword=2))
+            module.add(SWaitCnt(kmcnt=1, comment="Wait for the AddressA Load from the Pointer Array"))
+            module.add(SAddU32(dst=sgpr("AddressA+0"), src0=sgpr("AddressA+0"), src1=sgpr(tmpSgprOffset.idx), comment="add batch offset to A address (low)"))
+            module.add(SAddCU32(dst=sgpr("AddressA+1"), src0=sgpr("AddressA+1"), src1=sgpr(tmpSgprOffset.idx+1), comment="add batch offset to A address (high)"))
+            # Load and apply offset B (reuse same tmpSgprOffset)
+            module.add(self.argLoader.loadKernArg(tmpSgprOffset.idx, "KernArgAddress", sgprOffset=hex(self.states.batchOffsetBKernArgOffset), dword=2))
+            module.add(SWaitCnt(kmcnt=0, comment="Wait for the AddressB and Batch Offset Loads"))
+            module.add(SAddU32(dst=sgpr("AddressB+0"), src0=sgpr("AddressB+0"), src1=sgpr(tmpSgprOffset.idx), comment="add batch offset to B address (low)"))
+            module.add(SAddCU32(dst=sgpr("AddressB+1"), src0=sgpr("AddressB+1"), src1=sgpr(tmpSgprOffset.idx+1), comment="add batch offset to B address (high)"))
 
     module.add(endCheckLabel)
 
@@ -1966,17 +1982,22 @@ def getKernelArgLoadModule(self, kernel, sgprStartIdx, numsOfLoad, preloadNum):
         sgprOffset += 4
         self.argLoader.setOffset(sgprOffset)
 
-    # Load batch offset arguments for general batched mode
-    if not kernel["ProblemType"]["StridedBatched"] and not kernel["ProblemType"]["GroupedGemm"]:
+    # Batch offset arguments for general batched mode
+    # Store kernarg offsets for on-demand loading (no persistent SGPRs needed)
+    if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
       sgprOffset = self.argLoader.getOffset()
-      kernelArgs.add(self.argLoader.loadKernArg("BatchOffsetD", "KernArgAddress", sgprOffset=hex(sgprOffset), dword=2))
+      self.states.batchOffsetDKernArgOffset = sgprOffset
       sgprOffset += 8   # 64-bit = 8 bytes
-      kernelArgs.add(self.argLoader.loadKernArg("BatchOffsetC", "KernArgAddress", sgprOffset=hex(sgprOffset), dword=2))
+      self.states.batchOffsetCKernArgOffset = sgprOffset
       sgprOffset += 8
-      kernelArgs.add(self.argLoader.loadKernArg("BatchOffsetA", "KernArgAddress", sgprOffset=hex(sgprOffset), dword=2))
+      self.states.batchOffsetAKernArgOffset = sgprOffset
       sgprOffset += 8
-      kernelArgs.add(self.argLoader.loadKernArg("BatchOffsetB", "KernArgAddress", sgprOffset=hex(sgprOffset), dword=2))
+      self.states.batchOffsetBKernArgOffset = sgprOffset
       sgprOffset += 8
+      self.argLoader.setOffset(sgprOffset)
+      # Debug: write kernel arg offsets to file for verification
+      with open("/tmp/kernel_batch_offset_debug.txt", "a") as f:
+        f.write(f"{self.states.kernelName}: D@{self.states.batchOffsetDKernArgOffset}, C@{self.states.batchOffsetCKernArgOffset}, A@{self.states.batchOffsetAKernArgOffset}, B@{self.states.batchOffsetBKernArgOffset}\n")
 
     return kernelArgs
 
@@ -4424,7 +4445,15 @@ def computeLoadSrd(self, kernel, tP, tc, indices, bpe):
     moduleLoadGeneralBatch.add(SAddU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr("Address%s+0"%tc), comment="Offsetting to the location [Lower half of address]"))
     moduleLoadGeneralBatch.add(SAddCU32(dst=sgpr(stmp+1), src0=sgpr("Address%s+1"%tc), src1=0, comment="Offsetting to the location [Higher half of address]"))
     moduleLoadGeneralBatch.add(SLoadB64(dst=sgpr("Srd%s"%tc, 2), base=sgpr(stmp, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array"))
-    moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array"))
+    # Load and apply batch offset for General Batched GEMM
+    if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
+      batchOffsetKernArgOffset = self.states.batchOffsetAKernArgOffset if tc == "A" else self.states.batchOffsetBKernArgOffset
+      moduleLoadGeneralBatch.add(SLoadB64(dst=sgpr(stmp, 2), base=sgpr("KernArgAddress", 2), soffset=hex(batchOffsetKernArgOffset), comment="Load batchOffset%s from kernel args"%tc))
+      moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for Matrix Address and Batch Offset Loads"))
+      moduleLoadGeneralBatch.add(SAddU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=sgpr(stmp+0), comment="Add batch offset to %s address (low)"%tc))
+      moduleLoadGeneralBatch.add(SAddCU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr("Srd%s+1"%tc), src1=sgpr(stmp+1), comment="Add batch offset to %s address (high)"%tc))
+    else:
+      moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array"))
     if self.states.groOffsetInMacroTile and ((tc == "A" and not kernel["enableTDMA"]) or (tc == "B" and not kernel["enableTDMB"])):
       prePad1 = int(self.states.srdShiftLeft[tc] * tP["bpeGR"]) # leave room in case we have to pointer shift
       moduleLoadGeneralBatch.add(SSubU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=prePad1, comment="pre-pad to make room for possible pointer shift"))
@@ -12422,9 +12451,18 @@ def computeStoreSrdStart(self, kernel, srdTcList: list, sgprBpeList = [], useSiz
               module.add(SAddU32(dst=sgpr(tmpS0), src0=sgpr(tmpS0), src1=sgpr("Address%s+0"%mat), comment="Offsetting to the location [Lower half of address]"))
               module.add(SAddCU32(dst=sgpr(tmpS1), src0=sgpr("Address%s+1"%mat), src1=0, comment="Offsetting to the location [Higher half of address]"))
               module.add(SLoadB64(dst=sgpr(tmpS0, 2), base=sgpr(tmpS0, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array"))
-              module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array"))
-              module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("Srd%s+0"%mat), src1=sgpr(tmpS0), comment="Offsetting within the Batch Matrix [Lower half of address]"))
-              module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Offsetting within the Batch Matrix [Higher half of address]")) 
+              # Load and apply batch offset for General Batched GEMM (C or D matrix)
+              if kernel["ProblemType"]["SupportUserArgs"] and not kernel["ProblemType"]["GroupedGemm"]:
+                batchOffsetKernArgOffset = self.states.batchOffsetCKernArgOffset if mat == "C" else self.states.batchOffsetDKernArgOffset
+                module.add(SLoadB64(dst=sgpr("Srd%s"%mat, 2), base=sgpr("KernArgAddress", 2), soffset=hex(batchOffsetKernArgOffset), comment="Load batchOffset%s from kernel args"%mat))
+                module.add(SWaitCnt(kmcnt=0, comment="Wait for Matrix Address and Batch Offset Loads"))
+                # Add loaded matrix address to SRD
+                module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("Srd%s+0"%mat), src1=sgpr(tmpS0), comment="Add matrix address to SRD (low)"))
+                module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Add matrix address to SRD (high)"))
+              else:
+                module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array"))
+                module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("Srd%s+0"%mat), src1=sgpr(tmpS0), comment="Offsetting within the Batch Matrix [Lower half of address]"))
+                module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Offsetting within the Batch Matrix [Higher half of address]"))
               module.add(generalBatchedGemmLoad_End)
           module.addSpaceLine()