Skip to content

Reduce No. SGPRs for StreamK by packing into one SGPR #520

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
aliry95amd wants to merge 1 commit into from
Closed

Reduce No. SGPRs for StreamK by packing into one SGPR #520

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions projects/hipblaslt/tensilelite/Tensile/Components/Signature.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]:
Expand Down
27 changes: 18 additions & 9 deletions projects/hipblaslt/tensilelite/Tensile/Components/StreamK.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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"))
Expand Down Expand Up @@ -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)
Expand All @@ -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)

Expand All @@ -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"))
Expand Down
5 changes: 2 additions & 3 deletions projects/hipblaslt/tensilelite/Tensile/KernelWriter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 5 additions & 2 deletions projects/hipblaslt/tensilelite/Tensile/Source/lib/source/ContractionSolution.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<uint32_t>("SKItersPerWG", skItersPerWG);
args.template append<uint32_t>("skGrid", skGrid);
args.template append<uint32_t>("skTiles", skTiles);
args.template append<uint32_t>("skGridAndTiles", skGridAndTiles);
args.template append<uint32_t>("skExtraIters", skExtraIters);
}
}
Expand Down