fix vgpr allocation for half/bfloat16 types
@@ -1178,14 +1178,12 @@ class KernelWriterAssembly(KernelWriter):
valuBlocks = (1+kernel["PrefetchLocalRead"]) * kernel["InnerUnroll"]
if kernel["MatrixInstruction"]:
- self.numVgprValuAPerBlock = kernel["ThreadTileA"]*tPA["bpe"]//self.bpr
- #self.numVgprValuBPerBlock = kernel["ThreadTileB"]*tPB["bpe"]//self.bpr
- #self.numVgprValuBPerBlock = kernel["MacroTile1"] // (4 * kernel["MatrixInstN"] * kernel["MatrixInstB"]) # BBlocks
- self.numVgprValuBPerBlock = (kernel["ThreadTileB"] * tPA["bpe"]) // (kernel["MatrixInstN"] * self.bpr) # ABlocks
+ # NOTE: ThreadTileA/B in MatrixInstruction context is a bit ambiguous
+ self.numVgprValuAPerBlock = kernel["ThreadTileA"]*tPA["bpe"] * kernel["VectorWidth"]//self.bpr
+ self.numVgprValuBPerBlock = (kernel["ThreadTileB"] * tPA["bpe"] * kernel["VectorWidth"]) // (kernel["MatrixInstN"] * self.bpr) # ABlocks
else:
self.numVgprValuAPerBlock = kernel["ThreadTileA"]*tPA["bpe"]//self.bpr
self.numVgprValuBPerBlock = kernel["ThreadTileB"]*tPB["bpe"]//self.bpr
-
if kernel["ProblemType"]["DataType"].isBFloat16() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
self.numVgprValuAPerBlock = self.numVgprValuAPerBlock * 2
self.numVgprValuBPerBlock = self.numVgprValuBPerBlock * 2
