ROCm · bnemanich · May 20, 2026 · May 20, 2026 · May 22, 2026 · May 23, 2026
@@ -77,17 +77,15 @@ def __init__(self, writer, kernel, config,
 
         # Dispatch table — unroll_iter is passed for mfma/lr
         self._dispatch = {
-            'mfma':         lambda em, ui: self.emit_mfma(em.source, ui),
-            'lr':           lambda em, ui: self.emit_lr(em.source, ui),
-            'gr':           lambda em, ui: self.emit_gr(em.source),
-            'wait_gr':      lambda em, ui: self.emit_wait_gr(em.source),
-            'wait_lr':      lambda em, ui: self.emit_wait_lr(),
-            'sync':         lambda em, ui: self.emit_sync(),
-            'lr_inc':       lambda em, ui: self.emit_lr_inc(em.source),
-            'gr_inc':       lambda em, ui: self.emit_gr_inc(em.source),
-            'gr_ptr_inc':   lambda em, ui: self.emit_gr_ptr_inc(em.source),
-            'gr_lds_swap':  lambda em, ui: self.emit_gr_lds_swap(em.source),
-            'skip':         lambda em, ui: self.emit_skip(em.source),
+            'mfma':     lambda em, ui: self.emit_mfma(em.source, ui),
+            'lr':       lambda em, ui: self.emit_lr(em.source, ui),
+            'gr':       lambda em, ui: self.emit_gr(em.source),
+            'wait_gr':  lambda em, ui: self.emit_wait_gr(em.source),
+            'wait_lr':  lambda em, ui: self.emit_wait_lr(),
+            'sync':     lambda em, ui: self.emit_sync(),
+            'lr_inc':   lambda em, ui: self.emit_lr_inc(em.source),
+            'gr_inc':   lambda em, ui: self.emit_gr_inc(em.source),
+            'skip':     lambda em, ui: self.emit_skip(em.source),
         }
 
     def emit_mfma(self, placement, unroll_iter=0):
@@ -224,31 +222,6 @@ def emit_gr_inc(self, source):
         module.add(globalReadLDSBufferSwap(tc, self.writer, self.kernel))
         return list(module.flatitems())
 
-    def emit_gr_ptr_inc(self, source):
-        """Emit globalReadPtrUpdates only (SRD advance, no LDS-buffer toggle).
-
-        Used by preloop's PGR=2 path so the SRD advance can happen
-        unconditionally while the LDS swap (`emit_gr_lds_swap`) is gated
-        by the `SkipOp(LE 1, NLL)` jump over GR(MT 1). See PR #7636
-        nakajee review on the small-counter realign.
-        """
-        tensor = source.tensor
-        tc = {'A': 'A', 'B': 'B', 'SA': 'MXSA', 'SB': 'MXSB'}.get(tensor, tensor)
-        module = Module()
-        if tensor in ('SA', 'SB'):
-            module.add(globalReadScalePtrUpdates(tc, self.writer, self.kernel))
-        else:
-            module.add(globalReadPtrUpdates(tc, self.writer, self.kernel))
-        return list(module.flatitems())
-
-    def emit_gr_lds_swap(self, source):
-        """Emit globalReadLDSBufferSwap only (LDS-buffer toggle, no SRD advance)."""
-        tensor = source.tensor
-        tc = {'A': 'A', 'B': 'B', 'SA': 'MXSA', 'SB': 'MXSB'}.get(tensor, tensor)
-        module = Module()
-        module.add(globalReadLDSBufferSwap(tc, self.writer, self.kernel))
-        return list(module.flatitems())
-
     def emit_skip(self, source):
         """Emit skip guard: compare LoopCounterL and branch."""
         skipLabel = Label(f"SkipTo{source.target}", "")

@@ -115,7 +115,8 @@
     _emitLRDTLInit, _emitLRLDSBufferSwap,
     lraTileAssignment, localReadDoSubtile, localReadDTLInitCommonSwapVgpr,
     localReadLDSBufferSwap, localReadResetOffsetsSubtile,
-    emitSingleDsRead, emitSubtileDsRead, setExecMask,
+    emitSingleDsRead, emitSubtileDsRead,
+    setExecMask,
 )
 from .SubtileScaleEmit import (
     emitScaleGROffset, emitScaleLROffset,
@@ -859,6 +860,88 @@ def deallocVgprTileRegisters_legacy(self, writer, kernel):
         vtiles.regList.pool.checkIn(vtiles.regList.indices[0])
     self.vgprTiles = []
 
+  def initVgprTileSlots(self, writer, kernel):
+    """Create empty `vgprTiles` RegisterTileInfo slots without
+    checking out VGPRs from the pool.
+
+    The tail scaffold's per-mmak interleave alloc/free pattern
+    populates one K-slice of slots per mmak via
+    `allocVgprTileRegistersForMmak` and releases them via
+    `freeVgprTileRegistersForMmak`. Emit code (`emitSubtileDsRead*`,
+    `emitMfmaInstruction`, byte-refine mask chain) reads
+    `vgprTiles[idx].regList.indices[0]`, so the slot list must exist
+    at full length even when only one slice is populated.
+
+    Only supports AB and MXScaleTilePair geometries; CD tiles still
+    use the full `allocVgprTileRegisters_legacy` path.
+    """
+    self.vgprTiles = []
+    numMMATiles = int(self.localMMATileGrid[0] * self.localMMATileGrid[1])
+    isDTile = isinstance(self.geometry, CDTileGeometry)
+    assert not isDTile, "initVgprTileSlots only supports AB / MXScale tiles"
+    for _ in range(numMMATiles):
+      self.vgprTiles.append(RegisterTileInfo(writer.vgprPool, RegisterType.Vgpr))
+
+  def allocVgprTileRegistersForMmak(self, writer, kernel, mmak):
+    """Allocate VGPRs for the K-slice of `vgprTiles` owned by `mmak`.
+
+    For each `mma` index in the M dimension the slice tile id is:
+
+        sId0 = mma // lrSubtileShape[0]
+        sId1 = mmak // lrSubtileShape[1]
+        innerK = mmak % lrSubtileShape[1]
+        tileId = (sId1 * lrLocalGrid0 + sId0) *
+                 (lrSubtileShape[0] * lrSubtileShape[1]) + innerK
+
+    Each unique `tileId` in the slice gets `ceil(mmaTileRegCount)`
+    aligned VGPRs checked out from the writer's vgprPool. Returns the
+    list of allocated tile ids so the caller can pair it with a
+    matching `freeVgprTileRegistersForMmak` call.
+
+    `initVgprTileSlots` must have populated `self.vgprTiles` first.
+    Only supports AB tiles (the MX-scale path's `(2,2)` LR subtile
+    shape interleaves a `_mma0` index that this slice formula does
+    not unroll; tail scales stay allocated up front).
+    """
+    assert not isinstance(self.geometry, MXScaleTilePair), (
+      "allocVgprTileRegistersForMmak does not support MX scale tiles; "
+      "their LR subtile shape interleaves _mma0 in the tile id which "
+      "the slice formula does not cover")
+    ss = self.lr.subtileShape
+    numMmaTilePerSubtile = ss[0] * ss[1]
+    lrLocalGrid0 = self.localMMATileGrid[0] // ss[0]
+    sId1 = mmak // ss[1]
+    innerK = mmak % ss[1]
+    numDword = int(math.ceil(self.mmaTileRegCount))
+    allocated = []
+    seen = set()
+    for mma in range(self.localMMATileGrid[0]):
+      sId0 = mma // ss[0]
+      tileId = (sId1 * lrLocalGrid0 + sId0) * numMmaTilePerSubtile + innerK
+      if tileId in seen:
+        continue
+      seen.add(tileId)
+      vstart = writer.vgprPool.checkOutAligned(numDword, numDword)
+      slot = self.vgprTiles[tileId]
+      slot.regList.indices.clear()
+      for k in range(numDword):
+        slot.append(vstart + k)
+      allocated.append(tileId)
+    return allocated
+
+  def freeVgprTileRegistersForMmak(self, writer, kernel, allocated):
+    """Release VGPRs allocated by `allocVgprTileRegistersForMmak`.
+
+    `allocated` is the list returned by the matching alloc call. The
+    slot's `regList.indices` is cleared so a downstream re-alloc (or
+    a leak assertion) sees an empty slot.
+    """
+    for tileId in allocated:
+      slot = self.vgprTiles[tileId]
+      if slot.regList.indices:
+        writer.vgprPool.checkIn(slot.regList.indices[0])
+        slot.regList.indices.clear()
+
   def __str__(self):
     return (f"TileInfo(tc={self.tc}, geometry={type(self.geometry).__name__}, "
             f"mmaTileShape={self.mmaTileShape}, "

@@ -348,15 +348,7 @@ def __str__(self):
 
 @dataclass
 class GRIncOp(BaseOp):
-    """Pointer update + LDS swap for global reads on a specific tensor.
-
-    The compound op is used in mainloop placements via `assign_grinc_lrinc`,
-    where SRD advance and LDS-buffer toggle always happen together as a
-    pair (one mainloop iter consumes one LDS half and prefetches into the
-    other). Preloop instead splits the two halves via `GRPtrIncOp` +
-    `GRLDSwapOp` so the LDS toggle can sit under the `SkipOp(LE 1, NLL)`
-    guard while the SRD advance still happens unconditionally.
-    """
+    """Pointer update + LDS swap for global reads on a specific tensor."""
     tensor: str = ""
 
     def __post_init__(self):
@@ -366,43 +358,6 @@ def __str__(self):
         return f"gr_inc({self.tensor})"
 
 
-@dataclass
-class GRPtrIncOp(BaseOp):
-    """SRD pointer advance only (no LDS-buffer toggle).
-
-    Used in preloop alongside `GRLDSwapOp` so the SRD advance can happen
-    unconditionally (the tail body needs SRD pointing at the right K for
-    every origCounter value, even when GR(MT 1) is bypassed by
-    `SkipOp(LE 1, NLL)`).
-    """
-    tensor: str = ""
-
-    def __post_init__(self):
-        self.kind = 'gr_ptr_inc'
-
-    def __str__(self):
-        return f"gr_ptr_inc({self.tensor})"
-
-
-@dataclass
-class GRLDSwapOp(BaseOp):
-    """LDS double-buffer toggle for global reads (no SRD advance).
-
-    Paired with `GRPtrIncOp` in preloop so the LDS swap can be placed
-    under the `SkipOp(LE 1, NLL)` guard. Skipping the swap when
-    origCounter <= 1 leaves LWA aligned with the LR buffer half that
-    actually got drained, removing the need for a tail-time XOR-back
-    realign for the small-counter case (per nakajee review on PR #7636).
-    """
-    tensor: str = ""
-
-    def __post_init__(self):
-        self.kind = 'gr_lds_swap'
-
-    def __str__(self):
-        return f"gr_lds_swap({self.tensor})"
-
-
 @dataclass
 class SkipOp(BaseOp):
     """Skip guard: compare LoopCounter and branch."""
@@ -2192,23 +2147,13 @@ def build_preloop(self) -> List[List[List[EmittedModule]]]:
                 SkipOp(compare='LE', value=1, target='NLL'),
             ])
         else:
-            # PGR=2: split the n→n+1 GR_INC into ptr-inc (always fires;
-            # the tail body relies on Srd<tc> being advanced by one DU
-            # regardless of origCounter) and lds-swap (fires only when
-            # we will actually use the swapped LDS half, i.e. when
-            # GR(MT 1) is not bypassed by `SkipOp(LE 1, NLL)`). For
-            # origCounter==1 the swap is skipped, leaving LWA aligned
-            # with the LR buffer half that NLL drains from, so the
-            # subtile tail scaffold no longer needs a small-counter
-            # XOR-back realign. See nakajee PR #7636 review comment.
             emitted = self._to_emitted([
                 *self._make_gr_all_tensors(0, all_tiles),
-                *self._make_depops_all_tensors(GRPtrIncOp),
+                *self._make_depops_all_tensors(GRIncOp),
                 WaitGROp(wait_gr_counts=WaitGRCounts()),
                 SyncOp(),
                 *self._make_lr_all_tensors(lr_tiles),
                 SkipOp(compare='LE', value=1, target='NLL'),
-                *self._make_depops_all_tensors(GRLDSwapOp),
                 *self._make_preloop_mt1_grs(),
                 SkipOp(compare='LE', value=2, target='NGLL'),
             ])

@@ -639,6 +639,7 @@ def emitSubtileDsRead(writer, kernel, tileInfo, subtileId):
 
   return module
 
+
 ##################################################
 # Subroutine to generate LR load code
 # Initial idea: maybe store asm in modules in a separate obj?