From d060fd4150a0d2f68924d1b6a6889b609e02a254 Mon Sep 17 00:00:00 2001 From: aliry95amd Date: Mon, 7 Jul 2025 11:09:20 -0500 Subject: [PATCH] Reduce No. SGPRs for StreamK by packing into one SGPR --- .../Tensile/Components/Signature.py | 5 ++-- .../tensilelite/Tensile/Components/StreamK.py | 27 ++++++++++++------- .../tensilelite/Tensile/KernelWriter.py | 5 ++-- .../Source/lib/source/ContractionSolution.cpp | 7 +++-- 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py b/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py index d02bff803072..cad676c8f1c5 100644 --- a/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py +++ b/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py @@ -218,10 +218,9 @@ def __call__(self, writer) -> SignatureBase: signature.addArg("SKItersPerWG", SVK.SIG_VALUE, "u32") userArgumentsInfo.gemmArgumentSize += 36 if kernel["StreamK"] >= 2: # Two-tile SK - signature.addArg("skGrid", SVK.SIG_VALUE, "u32") - signature.addArg("skTiles", SVK.SIG_VALUE, "u32") + signature.addArg("skGridAndTiles", SVK.SIG_VALUE, "u32") signature.addArg("skExtraIters", SVK.SIG_VALUE, "u32") - userArgumentsInfo.gemmArgumentSize += 12 + userArgumentsInfo.gemmArgumentSize += 8 # "dpTilesPerWG" if kernel["ProblemType"]["UseScaleAB"]: diff --git a/projects/hipblaslt/tensilelite/Tensile/Components/StreamK.py b/projects/hipblaslt/tensilelite/Tensile/Components/StreamK.py index 231968c11f9d..4c5ef4be115c 100644 --- a/projects/hipblaslt/tensilelite/Tensile/Components/StreamK.py +++ b/projects/hipblaslt/tensilelite/Tensile/Components/StreamK.py @@ -61,6 +61,7 @@ def __call__(self, writer, kernel): module = Module("XCCMapping On") with writer.allocTmpSgpr(4) as tmpSgprRes: + skGrid = tmpSgprRes.idx sXCC = tmpSgprRes.idx sGridC = tmpSgprRes.idx + 1 sGridF = tmpSgprRes.idx + 2 @@ -74,11 +75,12 @@ def __call__(self, writer, kernel): sTmpRes = ContinuousRegister(idx=sTmp, size=2) # sGridC = ceil(grid / xccm) - module.add(SAddU32(dst=sgpr(sGridC), src0=sgpr("skGrid"), src1=hex(kernel["StreamKXCCMapping"] - 1), comment="ceil(grid/xccm)")) + module.add(SLShiftRightB32(dst=sgpr(skGrid), shiftHex=hex(16), src=sgpr("skGridAndTiles"), comment="Get skGrid")) + module.add(SAddU32(dst=sgpr(sGridC), src0=sgpr(skGrid), src1=hex(kernel["StreamKXCCMapping"] - 1), comment="ceil(grid/xccm)")) module.add(scalarStaticDivideAndRemainder(qReg=sGridC, rReg=-1, dReg=sGridC, divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes, doRemainder=0)) # sGridF = floor(grid / xccm) # sGridM = grid % xccm - module.add(scalarStaticDivideAndRemainder(qReg=sGridF, rReg=sGridM, dReg="skGrid", divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes)) + module.add(scalarStaticDivideAndRemainder(qReg=sGridF, rReg=sGridM, dReg=skGrid, divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes)) # sXCC = wg0 % xccm # sqtmp is temp register for quotient for non-power-of-2 case # sqtmp overlaps temp registers, works in this case and output is discarded @@ -1662,7 +1664,8 @@ def preLoop(self, writer, kernel): # clamp to end of sk iterations # TODO maybe remove clamp, since extra iters code should guarantee total iterations match sTmp = writer.sgprPool.checkOut(1, "TotalSKIters", preventOverflow=False) - module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("skTiles"), src1=sgpr("ItersPerTile"), comment="Total SK iters")) + module.add(SAndB32(dst=sgpr(sTmp), src0=sgpr("skGridAndTiles"), src1=hex(65535), comment="Get skTiles")) + module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(sTmp), src1=sgpr("ItersPerTile"), comment="Total SK iters")) module.add(SMinU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(sTmp), comment="Cap ending iter at total SK iters")) writer.sgprPool.checkIn(sTmp) # check if this WG has no work to do @@ -1689,11 +1692,13 @@ def graWorkGroup(self, writer, kernel, tPA, tPB): # Increment StreamK iteration # If moving from SK to DP, next iteration is first DP # sTmp = offset to first DP tile - module.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr("skTiles"), src1=sgpr("ItersPerTile"), comment="Offset to first DP tile")) + module.add(SAndB32(dst=sgpr(sTmp+3), src0=sgpr("skGridAndTiles"), src1=hex(65535), comment="Get skTiles")) + module.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr("ItersPerTile"), comment="Offset to first DP tile")) module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr("StreamKIdx"), src1=sgpr("ItersPerTile"), comment="WG tile offset")) module.add(SAddU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr(sTmp+1), comment="DP start offset + WG offset")) # If already in DP, add dpShift - module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr("skGrid"), src1=sgpr("ItersPerTile"), comment="DP iterations shift")) + module.add(SLShiftRightB32(dst=sgpr(sTmp+1), shiftHex=hex(16), src=sgpr("skGridAndTiles"), comment="Get skGrid")) + module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("ItersPerTile"), comment="DP iterations shift")) module.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("StreamKIter"), comment="Add DP shift")) # Save DP iter in sTmp module.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr("StreamKIterEnd"), comment="Check if in SK or DP section")) @@ -1760,7 +1765,8 @@ def preLoop(self, writer, kernel): module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIdx"), src1=sgpr("ItersPerTile"), comment="DP starting iteration (case: DP work to do)")) module.add(SMovB32(dst=sgpr("StreamKIterEnd"), src=sgpr("TotalIters"), comment="DP ending iteration (case: only DP work to do)")) sTmp = writer.sgprPool.checkOut(1, "TotalSKIters", preventOverflow=False) - module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("skTiles"), src1=sgpr("ItersPerTile"), comment="Total SK iters")) + module.add(SAndB32(dst=sgpr(sTmp), src0=sgpr("skGridAndTiles"), src1=hex(65535), comment="Get skTiles")) + module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(sTmp), src1=sgpr("ItersPerTile"), comment="Total SK iters")) module.add(SCmpLtU32(src0=sgpr(sTmp), src1=sgpr("TotalIters"), comment="Check if there are DP tiles to do")) module.add(SCBranchSCC1(labelName=skInitDone.getLabelName(), comment="Done init")) writer.sgprPool.checkIn(sTmp) @@ -1784,7 +1790,8 @@ def preLoop(self, writer, kernel): # clamp to end of sk iterations # TODO maybe remove clamp, since extra iters code should guarantee total iterations match sTmp = writer.sgprPool.checkOut(1, "TotalSKIters", preventOverflow=False) - module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("skTiles"), src1=sgpr("ItersPerTile"), comment="Total SK iters")) + module.add(SAndB32(dst=sgpr(sTmp), src0=sgpr("skGridAndTiles"), src1=hex(65535), comment="Get skTiles")) + module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(sTmp), src1=sgpr("ItersPerTile"), comment="Total SK iters")) module.add(SMinU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(sTmp), comment="Cap ending iter at total SK iters")) writer.sgprPool.checkIn(sTmp) @@ -1805,10 +1812,12 @@ def graWorkGroup(self, writer, kernel, tPA, tPB): skUpdateDone = Label("SK_UpdateDone", "") # sTmp+3 = Offset to first SK tile - module.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr("skTiles"), src1=sgpr("ItersPerTile"), comment="Total SK iters")) + module.add(SAndB32(dst=sgpr(sTmp+3), src0=sgpr("skGridAndTiles"), src1=hex(65535), comment="Get skTiles")) + module.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr("ItersPerTile"), comment="Total SK iters")) module.add(SSubU32(dst=sgpr(sTmp+3), src0=sgpr("TotalIters"), src1=sgpr(sTmp+3), comment="Offset to first SK tile")) # If in DP, add dpShift - module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr("skGrid"), src1=sgpr("ItersPerTile"), comment="DP iterations shift")) + module.add(SLShiftRightB32(dst=sgpr(sTmp+1), shiftHex=hex(16), src=sgpr("skGridAndTiles"), comment="Get skGrid")) + module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("ItersPerTile"), comment="DP iterations shift")) module.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("StreamKIter"), comment="Add DP shift")) # if sTmp+1 < sTmp+3, continue DP (add dpShift) module.add(SCmpLtU32(src0=sgpr(sTmp+1), src1=sgpr(sTmp+3), comment="Check if still in DP section")) diff --git a/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py b/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py index 657324b03d93..b24f8d83dd8b 100644 --- a/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py +++ b/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py @@ -4686,11 +4686,10 @@ def readWriteVectors(mat, vw, kernel): self.defineSgpr("SKItersPerWG", 1) self.states.numSgprStreamK += 9 if kernel["StreamK"] >= 2: # Two-tile SK - self.defineSgpr("skGrid", 1) - self.defineSgpr("skTiles", 1) + self.defineSgpr("skGridAndTiles", 1) self.defineSgpr("skExtraIters", 1) # self.defineSgpr("dpTilesPerWG", 1, kernarg=True) - self.states.numSgprStreamK += 3 + self.states.numSgprStreamK += 2 if kernel["LocalWriteUseSgprA"]: self.defineSgpr("LocalWriteAddrA", 1) diff --git a/projects/hipblaslt/tensilelite/Tensile/Source/lib/source/ContractionSolution.cpp b/projects/hipblaslt/tensilelite/Tensile/Source/lib/source/ContractionSolution.cpp index 1e58902d91e5..5b3f3935a802 100644 --- a/projects/hipblaslt/tensilelite/Tensile/Source/lib/source/ContractionSolution.cpp +++ b/projects/hipblaslt/tensilelite/Tensile/Source/lib/source/ContractionSolution.cpp @@ -766,9 +766,12 @@ namespace TensileLite uint32_t skItersPerWG = skTiles * itersPerTile / skGrid; uint32_t skExtraIters = skTiles * itersPerTile % (skGrid); + // Pack skGrid and skTiles into a single uint32_t such that the upper 16 bits + // represent skGrid and the lower 16 bits represent skTiles + uint32_t skGridAndTiles = (skGrid <<16) | (skTiles & 0xFFFF); + args.template append("SKItersPerWG", skItersPerWG); - args.template append("skGrid", skGrid); - args.template append("skTiles", skTiles); + args.template append("skGridAndTiles", skGridAndTiles); args.template append("skExtraIters", skExtraIters); } }