Modified vgpr/sgpr occupancy limit method at the Global Write stage.

 * Add refineOccupancy to set the proper occupancy limit.
 * Allocated the required SGPRs needed by storeRemapAddStore() before calculating occupancy.

Author: vin-huang

tensilelite/Tensile/KernelWriterAssembly.py

@@ -10253,6 +10253,8 @@ def cleanupGlobalWrite(self, kernel):
       self.vgprPool.checkIn(self.vgprs.storeRemapCoord0)
       self.vgprPool.checkIn(self.vgprs.storeRemapCoord1)
       self.vgprPool.checkIn(self.vgprs.storeRemapOffsetCoord1)
+      for v in self.vgprs.storeRemapAS:
+        self.vgprPool.checkIn(v)
     if kernel["BufferStore"]:
       self.vgprPool.checkIn(self.vgprs.cinRowPtr)
       self.vgprPool.checkIn(self.vgprs.coutRowPtrD)
@@ -10374,10 +10376,7 @@ def storeRemapAddStore(self, kernel, tmpVgpr, tmpS01, edge, StoreRemapLastBatch)
     rpe = self.states.bpeCexternal / self.states.bpr
     rpv = rpe * gwvw
 
-    # num registers to check out
-    storeRegs = []
-    for i in range(0, nElements, gwvw):
-      storeRegs.append(self.vgprPool.checkOutAligned(int(rpv), int(rpv), "store element d"))
+    storeRegs = self.vgprs.storeRemapAS
     src = vgpr(self.vgprs.storeRemapLR)
     for rIdx, i in enumerate(range(0, nElements, gwvw)):
       offset = self.storeRemapLrOffset * bpe * (i//gwvw)
@@ -10482,8 +10481,6 @@ def storeRemapAddStore(self, kernel, tmpVgpr, tmpS01, edge, StoreRemapLastBatch)
 
     module.addSpaceLine()
     self.vgprPool.checkIn(vTmp)
-    for v in storeRegs:
-      self.vgprPool.checkIn(v)
 
     #Data exchange between different waves
     #Make sure LDS reads are finished of all waves
@@ -10628,6 +10625,12 @@ def storeRemapComputeStoreVgprs(self, kernel):
 
       self.vgprPool.checkIn(tmpV0)
 
+      nElements = kernel["MacroTile0"]*kernel["MatrixInstN"]//kernel["MIWaveGroup"][0]//self.states.kernel["WavefrontSize"]
+      rpe = self.states.bpeCexternal / self.states.bpr
+      rpv = rpe * gwvw
+      self.vgprs.storeRemapAS = []
+      for i in range(0, nElements, gwvw):
+        self.vgprs.storeRemapAS.append(self.vgprPool.checkOutAligned(int(rpv), int(rpv), "store element d"))
     return module
 
   ##############################################################################
@@ -11620,77 +11623,19 @@ def getMBSKGSUTotal(self, kernel):
     GSUtotal = max(2,GSUtotal)
     return GSUtotal
 
-  ##############################################################################
-  # globalWriteElementBatch :
-  ##############################################################################
-  def globalWriteElementBatch(self, kernel, tPA, tPB, activation, \
-                              applyAlpha, beta, edge, atomic, \
-                              vectorWidths, elements, activationLabelList, \
-                              tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, \
-                              actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, \
-                              edgeModule, writeLabels, endLabel, \
-                              edge_mode_pos, currentInstLength, \
-                              idx0, idx1, idx2, idxMN, vectorDataTypes, factorDims):
-    factorDim = factorDims[idx2]
-    edgeModule.add(writeLabels[beta][edge][factorDim][idxMN])
-    if idx2 == 0:
-      edge_mode_pos = len(edgeModule.items())
-
-    # for storeRemap edge case, non-beta still can enable vector stores
-    if kernel["StoreRemapVectorWidth"] and not beta:
-      edgeI = False
-    else:
-      edgeI = edge
-    #edgeI = True  # set to True to disable vector stores
-    gwvw = vectorWidths[edgeI]
-
-    #print "globalWriteElements: edge=", edge, "beta=", beta, "atomic=", atomic
-
-    ########################################
-    # Calculate Vgprs for Write Batching
-    ########################################
-    self.vgprPool.resetOccupancyLimit()
-    self.sgprPool.resetOccupancyLimit()
-
-    # Temporarily grow pool for sgpr
-    sgprList = []
-    if kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel':
-      sgprList.append(self.sgprPool.checkOut(1, preventOverflow=False))
-      sgprList.append(self.sgprPool.checkOut(1, preventOverflow=False))
-      sgprList.append(self.sgprPool.checkOut(1, preventOverflow=False))
-      sgprList.append(self.sgprPool.checkOutAligned(2, 2, preventOverflow=False))
-      sgprList.append(self.sgprPool.checkOutAligned(2, 2, preventOverflow=False))
-      sgprList.append(self.sgprPool.checkOutAligned(4, 4, preventOverflow=False))
-      for s in sgprList:
-        self.sgprPool.checkIn(s)
-    if actPCMaxTempSgpr > 0:
-      self.sgprPool.checkIn(self.sgprPool.checkOutAligned(actPCMaxTempSgpr, 2 if actPCMaxTempSgpr > 1 else 1, preventOverflow=False))
-
-    tmpVgprDynamic = None
-    tmpVgprDynamicSize  = 0
-    tmpVgprDynamicAlign = 0
-    if kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel':
-      GSUTotal = self.getMBSKGSUTotal(kernel)
-      vgprMbsk = (GSUTotal-1) * gwvw * max(1, kernel["ProblemType"]["DestDataType"].numRegisters())
-      tmpVgprDynamicSize  = vgprMbsk
-      tmpVgprDynamicAlign = 4
-    if tmpVgprDynamicSize > 0:
-      tmpVgprDynamic = ContinuousRegister(idx=self.vgprPool.checkOutAligned(tmpVgprDynamicSize, tmpVgprDynamicAlign), size=tmpVgprDynamicSize)
-
-    ss = StoreState(self, kernel, gwvw, edge, beta, atomic, elements[edgeI], vectorDataTypes, dim=factorDim)
-
-    def setOccupancy():
-      # Use VGPR up to next occupancy threshold:
-      maxVgprs, occupancy = self.getMaxRegsForOccupancy(kernel["NumThreads"], self.vgprPool.size(), self.sgprPool.size(), \
-                                                        self.getLdsSize(kernel), self.agprPool.size(), self.states.doubleVgpr)
-      # Set occupancy limit for register pools
-      # TODO: Support gfx12
-      if kernel["ISA"][0] != 12:
-        self.vgprPool.setOccupancyLimit(self.states.regCaps["MaxVgpr"], self.states.regCaps["PhysicalMaxVgpr"] // occupancy)
-        self.sgprPool.setOccupancyLimit(self.states.regCaps["MaxSgpr"], self.states.regCaps["PhysicalMaxSgpr"] // occupancy)
-      return maxVgprs, occupancy
-
-    maxVgprs, occupancy = setOccupancy()
+  def setOccupancy(self, kernel):
+    # Use VGPR up to next occupancy threshold:
+    maxVgprs, occupancy = self.getMaxRegsForOccupancy(kernel["NumThreads"], self.vgprPool.size(), self.sgprPool.size(), \
+                                                      self.getLdsSize(kernel), self.agprPool.size(), self.states.doubleVgpr)
+    # Set occupancy limit for register pools
+    # TODO: Support gfx12
+    if kernel["ISA"][0] != 12:
+      self.vgprPool.setOccupancyLimit(self.states.regCaps["MaxVgpr"], self.states.regCaps["PhysicalMaxVgpr"] // occupancy)
+      self.sgprPool.setOccupancyLimit(self.states.regCaps["MaxSgpr"], self.states.regCaps["PhysicalMaxSgpr"] // occupancy)
+    return maxVgprs, occupancy
+
+  def refineOccupancy(self, kernel, atomic, elements, actPCMaxTempSgpr, \
+                      edgeI, gwvw, maxVgprs, ss):
     # Get estimated numVgprAvailable
     # print("Max vgprs =", maxVgprs, self.vgprPool.size(), self.vgprPool.availableBlock(ss.numVgprsPerElement, ss.align))
     numVgprAvailable = self.vgprPool.availableBlockMaxVgpr(maxVgprs, ss.numVgprsPerElement, ss.align)
@@ -11711,12 +11656,12 @@ def setOccupancy():
           % (minElements,ss.numVgprsPerElement))
       self.vgprPool.growPool(0, minElements, ss.numVgprsPerElement, \
         "grow-pool for GlobalWrite")
-      maxVgprs, occupancy = setOccupancy()
+      maxVgprsN, occupancy = self.setOccupancy(kernel)
+      if maxVgprs != maxVgprsN:
+        #print("refineOccupancy maxVgprs, new", maxVgprsN, "old", maxVgprs)
+        return self.refineOccupancy(kernel, atomic, elements, actPCMaxTempSgpr, edgeI, gwvw, maxVgprsN, ss)
       numVgprAvailable = self.vgprPool.available()
 
-    # set atomicW after we potentially resize GWVW
-    atomicW = min(gwvw, self.getVectorAtomicWidth(kernel))
-
     # print("NumVgprAvailable", numVgprAvailable)
     if ss.numVgprsPerElement:
       numElementsPerBatch = numVgprAvailable // ss.numVgprsPerElement
@@ -11761,6 +11706,7 @@ def setOccupancy():
 
     # check best numElementsPerBatch to handle a column block
     # elements of column block must be multiple size of numElementsPerBatch
+    nBatchesPerRow = 0
     if kernel["StoreRemapVectorWidth"]:
       firstRow = [e for e in elements[edgeI] if e[0]==0 and e[2]==0] # format for element = (tt1, tt0, vc1, vc0)
       # find the largest factor and smaller than numElementPerBatch
@@ -11785,19 +11731,106 @@ def setOccupancy():
     totalNeededVgpr = ss.numVgprsPerElement * numElementsPerBatch
     # print("Available vgprs =", numVgprAvailable, "Needed vgprs =", totalNeededVgpr, "pool size =", self.vgprPool.size())
     if numVgprAvailable < totalNeededVgpr:
+      self.vgprPool.resetOccupancyLimit()
       print2("info: growing pool += %d * %d for GlobalWrite\n" \
           % (numBatches,ss.numVgprsPerElement))
       availableBlock = min(0, self.vgprPool.available() - numVgprAvailable)
       self.vgprPool.growPool(0, totalNeededVgpr + availableBlock, 1, "grow-pool for GlobalWrite")
+      maxVgprsN, occupancy = self.setOccupancy(kernel)
+      if maxVgprs != maxVgprsN:
+        #print("refineOccupancy maxVgprs, new", maxVgprsN, "old", maxVgprs)
+        return self.refineOccupancy(kernel, atomic, elements, actPCMaxTempSgpr, edgeI, gwvw, maxVgprsN, ss)
+  
     # # Get true numVgprAvailable
     # numVgprAvailable = self.vgprPool.availableBlock(ss.numVgprsPerElement, ss.align)
     # print("Available vgprs =", numVgprAvailable, "pool size =", self.vgprPool.size())
 
     numSgprs = ss.cfg.numTempSgprPerBatch + ss.cfg.numMaskSgprPerBatch + ss.cfg.numMaskSgprPerElement * numElementsPerBatch
 
+    if actPCMaxTempSgpr:
+      numSgprs = max(actPCMaxTempSgpr, numSgprs)
+
+    self.sgprPool.resetOccupancyLimit()
+    self.sgprPool.checkIn(self.sgprPool.checkOutAligned(numSgprs, 2, preventOverflow=False))
+    maxVgprsN, occupancy = self.setOccupancy(kernel)
+    if maxVgprs != maxVgprsN:
+      #print("refineOccupancy maxVgprs, new", maxVgprsN, "old", maxVgprs)
+      return self.refineOccupancy(kernel, atomic, elements, actPCMaxTempSgpr, edgeI, gwvw, maxVgprsN, ss)
+    return numElementsPerBatch, nBatchesPerRow, numBatches, numSgprs
+    
+
+  ##############################################################################
+  # globalWriteElementBatch :
+  ##############################################################################
+  def globalWriteElementBatch(self, kernel, tPA, tPB, activation, \
+                              applyAlpha, beta, edge, atomic, \
+                              vectorWidths, elements, activationLabelList, \
+                              tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, \
+                              actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, \
+                              edgeModule, writeLabels, endLabel, \
+                              edge_mode_pos, currentInstLength, \
+                              idx0, idx1, idx2, idxMN, vectorDataTypes, factorDims):
+    factorDim = factorDims[idx2]
+    edgeModule.add(writeLabels[beta][edge][factorDim][idxMN])
+    if idx2 == 0:
+      edge_mode_pos = len(edgeModule.items())
+
+    # for storeRemap edge case, non-beta still can enable vector stores
+    if kernel["StoreRemapVectorWidth"] and not beta:
+      edgeI = False
+    else:
+      edgeI = edge
+    #edgeI = True  # set to True to disable vector stores
+    gwvw = vectorWidths[edgeI]
+
+    #print "globalWriteElements: edge=", edge, "beta=", beta, "atomic=", atomic
+
+    ########################################
+    # Calculate Vgprs for Write Batching
+    ########################################
+    self.vgprPool.resetOccupancyLimit()
+    self.sgprPool.resetOccupancyLimit()
+
+    # Temporarily grow pool for sgpr
+    sgprList = []
+    if kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel':
+      sgprList.append(self.sgprPool.checkOut(1, preventOverflow=False))
+      sgprList.append(self.sgprPool.checkOut(1, preventOverflow=False))
+      sgprList.append(self.sgprPool.checkOut(1, preventOverflow=False))
+      sgprList.append(self.sgprPool.checkOutAligned(2, 2, preventOverflow=False))
+      sgprList.append(self.sgprPool.checkOutAligned(2, 2, preventOverflow=False))
+      sgprList.append(self.sgprPool.checkOutAligned(4, 4, preventOverflow=False))
+      for s in sgprList:
+        self.sgprPool.checkIn(s)
+    if actPCMaxTempSgpr > 0:
+      self.sgprPool.checkIn(self.sgprPool.checkOutAligned(actPCMaxTempSgpr, 2 if actPCMaxTempSgpr > 1 else 1, preventOverflow=False))
+
+    tmpVgprDynamic = None
+    tmpVgprDynamicSize  = 0
+    tmpVgprDynamicAlign = 0
+    if kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel':
+      GSUTotal = self.getMBSKGSUTotal(kernel)
+      vgprMbsk = (GSUTotal-1) * gwvw * max(1, kernel["ProblemType"]["DestDataType"].numRegisters())
+      tmpVgprDynamicSize  = vgprMbsk
+      tmpVgprDynamicAlign = 4
+    if tmpVgprDynamicSize > 0:
+      tmpVgprDynamic = ContinuousRegister(idx=self.vgprPool.checkOutAligned(tmpVgprDynamicSize, tmpVgprDynamicAlign), size=tmpVgprDynamicSize)
+
+    maxVgprs, occupancy = self.setOccupancy(kernel)
+
+    ss = StoreState(self, kernel, gwvw, edge, beta, atomic, elements[edgeI], vectorDataTypes, dim=factorDim)
+
+    actPCMaxTempSgpr_ = None
     if activationLabelList and isInsertActFunctionCallAddrCalc:
       assert activationSetPCStruct, activationEnumStrList and activationLabelList and toActModuleList
-      numSgprs = max(actPCMaxTempSgpr, numSgprs)
+      actPCMaxTempSgpr_ = actPCMaxTempSgpr
+
+    numElementsPerBatch, nBatchesPerRow, numBatches, numSgprs = self.refineOccupancy(kernel, atomic, elements, actPCMaxTempSgpr_, edgeI, gwvw, maxVgprs, ss)
+
+    # set atomicW after we potentially resize GWVW
+    atomicW = min(gwvw, self.getVectorAtomicWidth(kernel))
+
+    if activationLabelList and isInsertActFunctionCallAddrCalc:
       edgeModule.add(self.insertActFunctionCallAddrCalc(activationSetPCStruct.sgprOffsetActivation, \
         gwvw, toActModuleList, activationEnumStrList, activationLabelList, \
         idx0, idx1))