# Copyright Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT ################################################################################ # LR (local read) emit and alloc dispatch. # # singledispatch over LR tag sentinels (LRTag_1x2, LRTag_TLU1, etc.). # ABLRTile calls these via self.config.tag as the dispatch key. # # Structure: # 1. Dispatch bases — @singledispatch declarations # 2. Implementations — logic functions decorated with @register ################################################################################ from functools import singledispatch from math import prod from rocisa.code import Module from rocisa.container import DSModifiers, EXEC, vgpr, sgpr from rocisa.enum import RegisterType from rocisa.instruction import ( DSLoadB128, SMovB32, SMovB64, VAddU32, VAndB32, VMovB32, VXorB32, VLShiftLeftB32, VLShiftRightB32, VMulLOU32, VPermlane16SwapB32, ) from .SubtileGeometry import ( LRTag_1x1, LRTag_1x2, LRTag_TLU1, ) from .SubtileScaleEmit import emitScaleLRLDSSwap ################################################################################ # 1. Dispatch bases ################################################################################ @singledispatch def _emitLocalReadOffset(tag, tile, ti, writer, kernel): raise NotImplementedError(f"emitLocalReadOffset not implemented for {type(tag).__name__}") @singledispatch def _emitLocalRead(tag, tile, ti, writer, kernel): raise NotImplementedError(f"emitLocalRead not implemented for {type(tag).__name__}") @singledispatch def _allocLROffsetRegisters(tag, tile, ti, writer, kernel): raise NotImplementedError(f"allocLROffsetRegisters not implemented for {type(tag).__name__}") @singledispatch def _deallocLROffsetRegisters(tag, tile, ti, writer, kernel): raise NotImplementedError(f"deallocLROffsetRegisters not implemented for {type(tag).__name__}") @singledispatch def _emitLRDTLInit(tag, tile, ti, writer, kernel): raise NotImplementedError(f"emitLRDTLInit not implemented for {type(tag).__name__}") @singledispatch def _emitLRLDSBufferSwap(tag, tile, ti, writer, kernel): raise NotImplementedError(f"emitLRLDSBufferSwap not implemented for {type(tag).__name__}") # Stubs for tags not yet implemented. _stub = lambda tag, tile, ti, writer, kernel: None _emitLocalReadOffset.register(LRTag_TLU1)(_stub) _emitLocalRead.register(LRTag_TLU1)(_stub) _allocLROffsetRegisters.register(LRTag_TLU1)(_stub) _deallocLROffsetRegisters.register(LRTag_TLU1)(_stub) _emitLRDTLInit.register(LRTag_TLU1)(_stub) _emitLRLDSBufferSwap.register(LRTag_TLU1)(_stub) ################################################################################ # Helpers ################################################################################ def _setExecMask(module, writer, maskLo, maskHi): """Set EXEC mask to a 64-bit immediate value.""" tmpSgpr = writer.sgprPool.checkOutAligned(2, 2, "setExecMask tmpSgpr", False) module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(maskLo), comment="exec mask lo")) module.add(SMovB32(dst=sgpr(tmpSgpr+1), src=hex(maskHi), comment="exec mask hi")) module.add(SMovB64(dst=EXEC(), src=sgpr(tmpSgpr, 2), comment="Set exec mask")) writer.sgprPool.checkIn(tmpSgpr) setExecMask = _setExecMask ################################################################################ # 2. Implementations ################################################################################ # --- LR offset emit (TLU=0) -------------------------------------------------- @_emitLocalReadOffset.register(LRTag_1x1) @_emitLocalReadOffset.register(LRTag_1x2) def _emitLROffset_TLU0(tag, tile, ti, writer, kernel): """LR offset for row-major (TLU=0) subtile with swizzling. Ported from legacy lraTileAssignment + _computeLROffset + _applyWavePartitionLROffset. Operates on a single tensor component (A or B). The LDS read layout uses MFMA register mapping: lane16 = laneId % instM (M row within MMA tile) lane16Group = laneId // instM (K column group) Steps: 1. Compute lane16 and lane16Group from Serial. 2. Apply rotation and swizzling to colOffset (de-swizzle to match GR's LDS layout). 3. Compute rowOffset = lane16 * subIterKBytes. 4. For each ds_read within the subtile: offset = (colOffset + advance) % blockSize * loadWidth + rowOffset. 5. Apply wave partition offset (shift LR offsets by wave's LDS region). """ return Module(f"LR Offset 1x2 ({ti.tc})") # STUB module = Module(f"LR Offset 1x2 ({ti.tc})") tc = ti.tc wavesize = kernel["WavefrontSize"] subIterKBytes = ti.subIterKBytes loadWidth = ti.loadWidthLR mi_m = ti.mmaTileShape[0] ldsRowBankSize = writer.states.archCaps["LDSBankCount"] * writer.states.archCaps["LDSBankWidth"] numRowsPerLDSBanks = ldsRowBankSize // subIterKBytes blockSize = subIterKBytes // loadWidth numMFMACols = int(ti.mmaTileShape[1] * ti.bpe) // loadWidth wg_m = ti.waveGroupSize numWaves = ti.numWaves waves_coop = numWaves // wg_m tmpVgpr = writer.vgprPool.checkOut(5, tag="_emitLROffset_TLU0_tmpVgpr") lane16 = tmpVgpr lane16Group = tmpVgpr + 1 rotation = tmpVgpr + 2 rowOffset = tmpVgpr + 3 colOffset = tmpVgpr + 4 # --- 1. lane16 and lane16Group from Serial --- module.add(VAndB32(dst=vgpr(lane16Group), src0=vgpr("Serial"), src1=wavesize-1, comment=f"{tc}: laneId")) module.add(VLShiftRightB32(dst=vgpr(lane16Group), shiftHex=hex(mi_m.bit_length()-1), src=vgpr(lane16Group), comment=f"{tc}: lane16Group = laneId // {mi_m}")) module.add(VAndB32(dst=vgpr(lane16), src0=vgpr("Serial"), src1=mi_m-1, comment=f"{tc}: lane16 = laneId %% {mi_m}")) # --- 2. Swizzling: rotation + permlane16 de-swizzle --- module.addComment0(f"{tc}: LR swizzling") # ldsRowId = lane16 // numRowsPerLDSBanks module.add(VLShiftRightB32(dst=vgpr(rotation), shiftHex=hex(numRowsPerLDSBanks.bit_length()-1), src=vgpr(lane16), comment=f"{tc}: lds_row_id")) # rotation = (ldsRowId // 2) * 2 module.add(VLShiftRightB32(dst=vgpr(rotation), shiftHex=hex(1), src=vgpr(rotation), comment=f"{tc}: ldsRowId // 2")) module.add(VLShiftLeftB32(dst=vgpr(rotation), shiftHex=hex(1), src=vgpr(rotation), comment=f"{tc}: (ldsRowId // 2) * 2")) # colOffset = (rotation + lane16Group) % blockSize module.add(VAddU32(dst=vgpr(colOffset), src0=vgpr(rotation), src1=vgpr(lane16Group), comment=f"{tc}: rotation + lane16Group")) module.add(VAndB32(dst=vgpr(colOffset), src0=vgpr(colOffset), src1=hex(blockSize-1), comment=f"{tc}: %% blockSize")) # Permlane16 swap to match GR's quad_perm swizzle pattern _setExecMask(module, writer, 0x33333333, 0x33333333) module.add(VPermlane16SwapB32(dst=vgpr(colOffset), src=vgpr(colOffset), comment=f"{tc}: de-swizzle")) _setExecMask(module, writer, -1, -1) # --- 3. rowOffset = lane16 * subIterKBytes --- module.add(VLShiftLeftB32(dst=vgpr(rowOffset), shiftHex=hex(subIterKBytes.bit_length()-1), src=vgpr(lane16), comment=f"{tc}: row = lane16 * {subIterKBytes}")) # --- 4. Compute LR offsets for each ds_read within the subtile --- # offset[0] = colOffset * loadWidth + rowOffset # offset[i] = ((colOffset + i * numMFMACols) % blockSize) * loadWidth + rowOffset module.add(VMovB32(dst=vgpr(tile.sharedVgprLROffset[0]), src=vgpr(colOffset), comment=f"{tc}: LR offset 0 col")) for i in range(1, ti.numLRPerSubtile): module.add(VAddU32(dst=vgpr(tile.sharedVgprLROffset[i]), src0=vgpr(tile.sharedVgprLROffset[i-1]), src1=hex(numMFMACols), comment=f"{tc}: advance col for MFMA {i}")) module.add(VAndB32(dst=vgpr(tile.sharedVgprLROffset[i]), src0=vgpr(tile.sharedVgprLROffset[i]), src1=hex(blockSize-1), comment=f"{tc}: col %% blockSize")) for i in range(ti.numLRPerSubtile): module.add(VLShiftLeftB32(dst=vgpr(tile.sharedVgprLROffset[i]), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(tile.sharedVgprLROffset[i]), comment=f"{tc}: col * {loadWidth}")) module.add(VAddU32(dst=vgpr(tile.sharedVgprLROffset[i]), src0=vgpr(tile.sharedVgprLROffset[i]), src1=vgpr(rowOffset), comment=f"{tc}: row + col")) writer.vgprPool.checkIn(tmpVgpr) # --- 5. Wave partition: shift LR offsets by wave's LDS region --- # Each wave reads from a different partition of LDS along the tc's own wave-group axis. # Guard: wg_m > 1 ensures tc's own axis has multiple waves (for A: wg_m, for B: wg_n). # Without this guard, a 1x4 WG would wrongly treat A's 4 N-waves as M-partitions. if waves_coop > 1 and wg_m > 1: # Each wave reads from a different M partition. The A LDS region has size # MT * subIterKBytes, split into wg_m partitions (one per M-direction wave). # B uses the same stride since B partition also maps 1:1 to M-direction waves. MT = ti.globalMMATileGrid[0] * ti.mmaTileShape[0] sInterval = MT * subIterKBytes // wg_m waveId = writer.vgprPool.checkOut(1, tag="_emitLROffset_TLU0_waveId") module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment=f"{tc}: waveId")) if tc == 'A': module.add(VAndB32(dst=vgpr(waveId), src0=hex(waves_coop - 1), src1=vgpr(waveId), comment=f"{tc}: waveId %% {waves_coop}")) else: module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(waves_coop.bit_length()-1), src=vgpr(waveId), comment=f"{tc}: waveId // {waves_coop}")) tmpSgpr = writer.sgprPool.checkOut(1, tag="_emitLROffset_TLU0_tmpSgpr") module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(sInterval), comment=f"{tc}: LR partition stride")) module.add(VMulLOU32(dst=vgpr(waveId), src1=vgpr(waveId), src0=sgpr(tmpSgpr))) for i in range(ti.numLRPerSubtile): module.add(VAddU32(dst=vgpr(tile.sharedVgprLROffset[i]), src0=vgpr(tile.sharedVgprLROffset[i]), src1=vgpr(waveId), comment=f"{tc}: + wave partition")) writer.vgprPool.checkIn(waveId) writer.sgprPool.checkIn(tmpSgpr) elif wg_m > 1: # waves_coop == 1 but wg_m > 1: each wave owns a separate LDS region MT = ti.globalMMATileGrid[0] * ti.mmaTileShape[0] sInterval = MT * subIterKBytes // (numWaves) waveId = writer.vgprPool.checkOut(1, tag="_emitLROffset_TLU0_waveId") module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment=f"{tc}: waveId")) tmpSgpr = writer.sgprPool.checkOut(1, tag="_emitLROffset_TLU0_tmpSgpr") module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(sInterval), comment=f"{tc}: LR partition stride")) module.add(VMulLOU32(dst=vgpr(waveId), src1=vgpr(waveId), src0=sgpr(tmpSgpr))) for i in range(ti.numLRPerSubtile): module.add(VAddU32(dst=vgpr(tile.sharedVgprLROffset[i]), src0=vgpr(tile.sharedVgprLROffset[i]), src1=vgpr(waveId), comment=f"{tc}: + wave partition")) writer.vgprPool.checkIn(waveId) writer.sgprPool.checkIn(tmpSgpr) # --- 6. Add global LDS start offset for B (B data follows A in LDS) --- ldsStartOffset = getattr(writer, f'ldsStartOffset{tc}', 0) if ldsStartOffset: stmp = writer.sgprPool.checkOut(1, tag="_emitLROffset_TLU0_stmp") module.add(SMovB32(dst=sgpr(stmp), src=ldsStartOffset, comment=f"{tc}: ldsStartOffset")) for i in range(ti.numLRPerSubtile): module.add(VAddU32(dst=vgpr(tile.sharedVgprLROffset[i]), src0=vgpr(tile.sharedVgprLROffset[i]), src1=sgpr(stmp), comment=f"{tc}: + LDS offset")) writer.sgprPool.checkIn(stmp) return module # --- LR alloc/dealloc (LRTag_1x2) ------------------------------------------- @_allocLROffsetRegisters.register(LRTag_1x1) @_allocLROffsetRegisters.register(LRTag_1x2) def _allocLROffsetRegs_1x2(tag, tile, ti, writer, kernel): """Allocate LR offset registers for row-major (TLU=0) 1x2 subtile shape. Two register groups are allocated: 1. sharedVgprLROffset[]: one VGPR per ds_read within a subtile. numLRPerSubtile = ceil(lrSubtileSize / (loadWidthLR * waveSize)). Each VGPR holds the per-lane byte offset into LDS for one ds_read_b128. 2. sharedVgprLROffsetSwap[]: same count, used for double-buffering. While one set is in use for the current iteration's LR, the other holds pre-computed offsets for the next iteration. """ tile.sharedVgprLROffset = [] tile.sharedVgprLROffsetSwap = [] for i in range(ti.numLRPerSubtile): tile.sharedVgprLROffset.append(writer.vgprPool.checkOut(1, tag="_allocLROffsetRegs_1x2_sharedVgprLROffset")) tile.sharedVgprLROffsetSwap.append(writer.vgprPool.checkOut(1, tag="_allocLROffsetRegs_1x2_sharedVgprLROffsetSwap")) @_deallocLROffsetRegisters.register(LRTag_1x1) @_deallocLROffsetRegisters.register(LRTag_1x2) def _deallocLROffsetRegs_1x2(tag, tile, ti, writer, kernel): """Deallocate LR offset registers.""" if isinstance(tile.sharedVgprLROffset, list): for voff in tile.sharedVgprLROffset: writer.vgprPool.checkIn(voff) tile.sharedVgprLROffset = [] if isinstance(tile.sharedVgprLROffsetSwap, list): for voff in tile.sharedVgprLROffsetSwap: writer.vgprPool.checkIn(voff) tile.sharedVgprLROffsetSwap = [] # --- LR load emit (LRTag_1x2) ----------------------------------------------- @_emitLocalRead.register(LRTag_1x1) @_emitLocalRead.register(LRTag_1x2) def _emitLR_1x2(tag, tile, ti, writer, kernel): return Module(f"LR Load 1x2 ({ti.tc})") # STUB """Emit ds_read_b128 for all subtiles in the local grid. For each subtile (sId0, sId1), for each MMA tile in K (subtileShape[1]): - addrVgpr = sharedVgprLROffset[mfmaId] (per-lane LDS byte offset) - ds_offset = subtile position in LDS (constant immediate) - dst = vgprTiles[tileIdx] (destination register tile) The tile index mapping: for subtile at linearId with numLRPerSubtile reads, tileIdx = linearId * numLRPerSubtile + mfmaId This assumes non-interleaved layout (subtileShape[0]=1 for 1x2). """ module = Module(f"LR Load 1x2 ({ti.tc})") tc = ti.tc # TODO: Remove legacy TileInfo dependency after full migration. # Uses legacy's grid/sizes/vgprTiles because TileInfo's expanded subtileShape # doesn't match the LDS layout computed from legacy values. legacyTi = getattr(writer.states, tc.lower()).tileInfo subtileSize = int(legacyTi.subtileSize) for i in range(int(legacyTi.localSubtileGrid[0])): for j in range(int(legacyTi.localSubtileGrid[1])): for du in range(int(legacyTi.subtileShape[1])): mfmaId = du addrVgpr = tile.sharedVgprLROffset[mfmaId] # DS offset: subtile position in LDS offset = i * subtileSize + j * int(legacyTi.globalSubtileGrid[0]) * subtileSize # Destination tile register tileIdx = ti.lrTileIndexForSubtile(i, j, mfmaId) dstTile = ti.vgprTiles[tileIdx] dstVgpr = dstTile.regList.indices[0] numRegs = len(dstTile.regList.indices) module.add(DSLoadB128( dst=vgpr(dstVgpr, numRegs), src=vgpr(addrVgpr), ds=DSModifiers(offset=offset), comment=f"LR {tc}[{i},{j}] k={du}") ) return module # --- LR DTL init (LRTag_1x2) ------------------------------------------------ @_emitLRDTLInit.register(LRTag_1x1) @_emitLRDTLInit.register(LRTag_1x2) def _emitLRDTLInit_1x2(tag, tile, ti, writer, kernel): return Module(f"LR DTL Init ({ti.tc})") # STUB """Compute swap VGPRs for LR double-buffering. For each sharedVgprLROffset[i], computes the corresponding swap offset: swap[i] = XOR(offset[i], offset[i] + ldsTotalSize) This mask toggles the LR read between the two LDS buffer halves. """ module = Module(f"LR DTL Init ({ti.tc})") stmp = writer.sgprPool.checkOut(1, tag="_emitLRDTLInit_1x2_stmp") module.add(SMovB32(dst=sgpr(stmp), src=writer.ldsTotalSize, comment=f"{ti.tc}: ldsTotalSize for swap")) for i in range(len(tile.sharedVgprLROffset)): vOff = tile.sharedVgprLROffset[i] vSwap = tile.sharedVgprLROffsetSwap[i] module.add(VAddU32(dst=vgpr(vSwap), src0=vgpr(vOff), src1=sgpr(stmp), comment=f"{ti.tc}: offset + ldsTotalSize")) module.add(VXorB32(dst=vgpr(vSwap), src0=vgpr(vOff), src1=vgpr(vSwap), comment=f"{ti.tc}: swap mask = XOR")) writer.sgprPool.checkIn(stmp) return module # --- LR LDS buffer swap (LRTag_1x2) ----------------------------------------- @_emitLRLDSBufferSwap.register(LRTag_1x1) @_emitLRLDSBufferSwap.register(LRTag_1x2) def _emitLRLDSSwap_1x2(tag, tile, ti, writer, kernel): """Toggle LR read offsets between double-buffer halves. XOR each sharedVgprLROffset with its swap mask to flip to the other buffer. """ module = Module() module.addComment0("Emit code to swap %s LR vgpr offsets"%ti.tc) for i in range(len(tile.sharedVgprLROffset)): vOff = tile.sharedVgprLROffset[i] vSwap = tile.sharedVgprLROffsetSwap[i] module.add(VXorB32(dst=vgpr(vOff), src0=vgpr(vOff), src1=vgpr(vSwap), comment="")) return module ################################################################################ # Legacy LR emit functions (moved from SubtileBasedKernel.py) ################################################################################ def _computeLROffset(module, tileInfo, colOffset, rowOffset, swizzled): tc = tileInfo.tc subIterKBytes = tileInfo.subIterKBytes loadWidth = tileInfo.loadWidthLR numMFMACols = int(tileInfo.mmaTileShape[1] * tileInfo.bpe) // loadWidth # TN case only # Without LDS swizzling (e.g. TDM), the full DepthU tile is contiguous in LDS, # so the K-row is depthUBytes wide. With swizzling, GR writes individual # subtile K-groups, so the effective K-row is subIterKBytes. ldsKBytes = subIterKBytes if swizzled else tileInfo.depthUBytes blockSize = ldsKBytes // loadWidth # Each ds_load_b128 fills REGS_PER_DS_READ VGPRs. Tiles with more VGPRs # (e.g. 8-VGPR wave32 BF16 or wave64 FP8) need multiple reads. Consecutive # LR offset entries advance by colsPerRead = numMFMACols / numReadsForTile # so entries within the same MMA tile cover equal K sub-portions. REGS_PER_DS_READ = loadWidth // 4 numReadsForTile = tileInfo.geometry.lr.mmaLayout.vgprs // REGS_PER_DS_READ colsPerRead = numMFMACols // numReadsForTile module.add(VMovB32(dst=vgpr(tileInfo.sharedVgprLROffset[0]), src=vgpr(colOffset), comment="%s: laneId"%tc)) for vgprId in range(1, len(tileInfo.sharedVgprLROffset)): module.add(VAddU32(dst=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src0=vgpr(tileInfo.sharedVgprLROffset[vgprId-1]), src1=hex(colsPerRead), comment="%s: colOffset for read %u"%(tc, vgprId))) module.add(VAndB32(dst=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src0=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src1=hex(blockSize-1), comment="%s: colOffset = colOffset %% block_size"%tc)) for vgprId in range(0, len(tileInfo.sharedVgprLROffset)): module.add(VLShiftLeftB32(dst=vgpr(tileInfo.sharedVgprLROffset[vgprId]), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(tileInfo.sharedVgprLROffset[vgprId]), comment="%s: colOffset*loadWidth"%tc)) module.add(VAddU32(dst=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src0=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src1=vgpr(rowOffset), comment="%s: row + col"%tc)) def _applyWavePartitionLROffset(module, writer, kernel, tileInfo): """Apply wave-based partition offset to LR offsets. loadRatioGR >= 2.0: no partition needed, contiguous subtiles (1x4 for A , 4x1 for B) loadRatioGR == 1.0: 2x2 config, each wave loads half of the subtile loadRatioGR == 0.5: 4x1 for A , 1x4 for B. Split in 4 subtiles groups """ tc = tileInfo.tc # TDM handles wave partitioning via descriptors # For single-wave, TDM puts all data at the wave's LDS base -- no partition needed. # For multi-wave, each wave's TDM writes to a different LDS region, so LR # offsets must include a per-wave partition offset. if kernel.get("enableTDM%s" % tc, False): numWaves = prod(kernel["MIWaveGroup"]) if numWaves == 1: return # Multi-wave TDM: add per-wave LDS offset based on axis position wgM, wgN = kernel["MIWaveGroup"] numWavesThisAxis = wgM if tc == 'A' else wgN if numWavesThisAxis <= 1: return # this tensor's axis is not split wavesize = kernel["WavefrontSize"] du = kernel["DepthU"] mt = kernel["MacroTile0"] if tc == 'A' else kernel["MacroTile1"] bpe = tileInfo.bpe waveId = writer.vgprPool.checkOut(1) module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment="waveId")) # Decompose to axis component if tc == 'A' and wgN > 1: module.add(VAndB32(dst=vgpr(waveId), src0=hex(wgM - 1), src1=vgpr(waveId), comment="waveIdM = waveId %% %d" % wgM)) elif tc == 'B' and wgM > 1: module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wgM.bit_length()-1), src=vgpr(waveId), comment="waveIdN = waveId / %d" % wgM)) # LDS offset per wave = waveId_axis * (mt / numWavesThisAxis * (du*bpe + pad)) rowBytes = int(du * bpe) + int(getattr(tileInfo, "ldsRowPadBytes", 0)) ldsPerWave = int(mt // numWavesThisAxis) * rowBytes tmpSgpr = writer.sgprPool.checkOut(1) module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(ldsPerWave), comment="LDS bytes per wave for %s" % tc)) module.add(VMulLOU32(dst=vgpr(waveId), src1=vgpr(waveId), src0=sgpr(tmpSgpr), comment="waveOffset")) for vgprId in range(len(tileInfo.sharedVgprLROffset)): module.add(VAddU32(dst=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src0=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src1=vgpr(waveId), comment="%s: TDM wave partition LR offset" % tc)) writer.vgprPool.checkIn(waveId) writer.sgprPool.checkIn(tmpSgpr) return if tileInfo.loadRatioGR >= 2.0: return wavesize = kernel["WavefrontSize"] subIterKBytes = tileInfo.subIterKBytes loadWidth = tileInfo.loadWidthGR waveId = writer.vgprPool.checkOut(1, tag="_applyWavePartitionLROffset_waveId") module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment="waveId")) partitionOffset = tileInfo.mmaTileShape[0] * tileInfo.localSubtileGrid[0] numRowsPerWave = wavesize // (subIterKBytes // loadWidth) if tileInfo.loadRatioGR == 1.0: mWaves = kernel["MIWaveGroup"][0] if tc == 'A': module.add(VAndB32(dst=vgpr(waveId), src0=hex(mWaves - 1), src1=vgpr(waveId), comment="%s: waveId %% %d"%(tc, mWaves))) else: module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(mWaves.bit_length()-1), src=vgpr(waveId), comment="%s: waveId / %d"%(tc, mWaves))) sInterval = partitionOffset * subIterKBytes elif tileInfo.loadRatioGR == 0.5: sInterval = partitionOffset * subIterKBytes else: raise NotImplementedError("Unsupported loadRatioGR for wave partition: %s"%str(tileInfo.loadRatioGR)) if sInterval == 0: writer.vgprPool.checkIn(waveId) return tmpSgpr = writer.sgprPool.checkOut(1, tag="_applyWavePartitionLROffset_tmpSgpr") module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(sInterval), comment="%s: interleave stride"%tc)) module.add(VMulLOU32(dst=vgpr(waveId), src1=vgpr(waveId), src0=sgpr(tmpSgpr), comment="")) for vgprId in range(len(tileInfo.sharedVgprLROffset)): module.add(VAddU32(dst=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src0=vgpr(tileInfo.sharedVgprLROffset[vgprId]), src1=vgpr(waveId), comment="%s: wave partition LR offset"%tc)) writer.vgprPool.checkIn(waveId) writer.sgprPool.checkIn(tmpSgpr) ################################################## # Subroutine to generate LR offset calculation code # def lraTileAssignment(writer, kernel): return _lraTileAssignment_legacy(writer, kernel) def _lraWavePartitioning_legacy(module, writer, kernel): tileInfoA = writer.states.a.tileInfo tileInfoB = writer.states.b.tileInfo _applyWavePartitionLROffset(module, writer, kernel, tileInfoA) _applyWavePartitionLROffset(module, writer, kernel, tileInfoB) def _lraTileAssignment_fp8_legacy(writer, kernel, module): """FP8 LR offset: block-swap + wave de-rotation for MFMA 16x16x128. Two ds_read_b128 per MFMA (numLRPerSubtile=2), using complementary block assignments to achieve zero LDS bank conflicts: finalColId = (lane16Group + 2*(lane16 >> 3)) % 4 [undo GR wave rotation] colOffset_0 = finalColId + swap_bit * 4 colOffset_1 = colOffset_0 ^ 4 where: swap_bit = (lane16 >> 1) & 1 The rotation 2*(lane16>>3) undoes the GR step 2 wave K_group rotation: waves with waveId&1==1 (M-rows 8..15) wrote with rotation=2; lane16>=8 reads them back with de-rotation=2. Together they achieve zero bank conflicts. """ tileInfoA = writer.states.a.tileInfo tileInfoB = writer.states.b.tileInfo subIterKBytes = tileInfoA.subIterKBytes wavesize = kernel["WavefrontSize"] mi_m = tileInfoA.mmaTileShape[0] loadWidth = tileInfoA.loadWidthLR tmpVgpr = writer.vgprPool.checkOut(6, tag="_lraTileAssignment_fp8_legacy_tmpVgpr") lane16, lane16Group, scratch, rowOffset, colOffset0, colOffset1 = range(tmpVgpr, tmpVgpr + 6) module.add(VAndB32(dst=vgpr(lane16), src0=vgpr("Serial"), src1=mi_m-1, comment="lane16 = laneId % 16")) module.add(VAndB32(dst=vgpr(lane16Group), src0=vgpr("Serial"), src1=wavesize-1, comment="laneId")) module.add(VLShiftRightB32(dst=vgpr(lane16Group), shiftHex=hex(mi_m.bit_length()-1), src=vgpr(lane16Group), comment="lane16Group = laneId // 16")) module.add(VLShiftRightB32(dst=vgpr(scratch), shiftHex=hex(3), src=vgpr(lane16), comment="lane16 >> 3 (1 if M-row >= 8)")) module.add(VLShiftLeftB32(dst=vgpr(scratch), shiftHex=hex(1), src=vgpr(scratch), comment="rotation = 2 * (lane16 >> 3)")) module.add(VAddU32(dst=vgpr(colOffset0), src0=vgpr(lane16Group), src1=vgpr(scratch), comment="lane16Group + rotation")) module.add(VAndB32(dst=vgpr(colOffset0), src0=vgpr(colOffset0), src1=hex(3), comment="finalColId = (lane16Group + rotation) % 4")) module.add(VLShiftRightB32(dst=vgpr(scratch), shiftHex=hex(1), src=vgpr(lane16), comment="lane16 >> 1")) module.add(VAndB32(dst=vgpr(scratch), src0=vgpr(scratch), src1=hex(1), comment="swap_bit")) module.add(VLShiftLeftB32(dst=vgpr(scratch), shiftHex=hex(2), src=vgpr(scratch), comment="swap_val = swap_bit * 4")) module.add(VAddU32(dst=vgpr(colOffset0), src0=vgpr(colOffset0), src1=vgpr(scratch), comment="colOffset_0 = finalColId + swap_val")) module.add(VXorB32(dst=vgpr(colOffset1), src0=vgpr(colOffset0), src1=hex(4), comment="colOffset_1 = colOffset_0 ^ 4")) module.add(VLShiftLeftB32(dst=vgpr(rowOffset), shiftHex=hex(subIterKBytes.bit_length()-1), src=vgpr(lane16), comment=f"rowOffset = lane16 * {subIterKBytes}")) for tileInfo in [tileInfoA, tileInfoB]: module.add(VLShiftLeftB32(dst=vgpr(tileInfo.sharedVgprLROffset[0]), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(colOffset0), comment=f"{tileInfo.tc}: col0 * {loadWidth}")) module.add(VAddU32(dst=vgpr(tileInfo.sharedVgprLROffset[0]), src0=vgpr(tileInfo.sharedVgprLROffset[0]), src1=vgpr(rowOffset), comment=f"{tileInfo.tc}: offset[0]")) if len(tileInfo.sharedVgprLROffset) > 1: module.add(VLShiftLeftB32(dst=vgpr(tileInfo.sharedVgprLROffset[1]), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(colOffset1), comment=f"{tileInfo.tc}: col1 * {loadWidth}")) module.add(VAddU32(dst=vgpr(tileInfo.sharedVgprLROffset[1]), src0=vgpr(tileInfo.sharedVgprLROffset[1]), src1=vgpr(rowOffset), comment=f"{tileInfo.tc}: offset[1]")) writer.vgprPool.checkIn(tmpVgpr) _lraWavePartitioning_legacy(module, writer, kernel) stmp = writer.sgprPool.checkOut(1, tag="_lraTileAssignment_legacy_stmp") module.add(SMovB32(dst=sgpr(stmp), src=writer.ldsStartOffsetB, comment="ldsStartOffsetB")) for vgprId in range(len(tileInfoB.sharedVgprLROffset)): module.add(VAddU32(dst=vgpr(tileInfoB.sharedVgprLROffset[vgprId]), src0=sgpr(stmp), src1=vgpr(tileInfoB.sharedVgprLROffset[vgprId]), comment="B matrix offset in LDS")) writer.sgprPool.checkIn(stmp) return module def _lraTileAssignment_legacy(writer, kernel): module = Module() module.addComment0("LR Offset Calculation for Subtile Based Tiling") tileInfoA = writer.states.a.tileInfo tileInfoB = writer.states.b.tileInfo if tileInfoA.bpe == 1: # FP8: block-swap swizzle, no VPermlane16Swap return _lraTileAssignment_fp8_legacy(writer, kernel, module) subIterKBytes = tileInfoA.subIterKBytes wavesize = kernel["WavefrontSize"] mi_m = tileInfoA.mmaTileShape[0] loadWidth = tileInfoA.loadWidthLR ldsRowBankSize = writer.states.archCaps["LDSBankCount"] * writer.states.archCaps["LDSBankWidth"] # With LDS swizzling (gfx950), K-row is one subtile group; without, full DepthU. ldsKBytes = subIterKBytes if writer.states.subtileLdsSwizzle else tileInfoA.depthUBytes padBytes = int(getattr(tileInfoA, "ldsRowPadBytes", 0)) ldsRowStride = ldsKBytes + padBytes numRowsPerLDSBanks = ldsRowBankSize // ldsKBytes blockSize = ldsKBytes // loadWidth tmpVgpr = writer.vgprPool.checkOut(6, tag="_lraTileAssignment_legacy_tmpVgpr") lane16, lane16Group, rotation, rowOffset, colOffset = range(tmpVgpr, tmpVgpr + 5) module.add(VAndB32(dst=vgpr(lane16Group), src0=vgpr("Serial"), src1=wavesize-1, comment="laneId")) module.add(VLShiftRightB32(dst=vgpr(lane16Group), shiftHex=hex(mi_m.bit_length()-1), src=vgpr(lane16Group), comment="lane16Group")) module.add(VAndB32(dst=vgpr(lane16), src0=vgpr("Serial"), src1=mi_m-1, comment="laneId %% 16")) module.add(VMovB32(dst=vgpr(colOffset), src=vgpr(lane16Group), comment="colOffset = lane16Group")) if writer.states.subtileLdsSwizzle: module.add(VLShiftRightB32(dst=vgpr(rotation), shiftHex=hex(numRowsPerLDSBanks.bit_length()-1), src=vgpr(lane16), comment="lds_row_id")) module.add(VLShiftRightB32(dst=vgpr(rotation), shiftHex=hex(1), src=vgpr(rotation), comment="(lds_row_id //2 )")) module.add(VLShiftLeftB32(dst=vgpr(rotation), shiftHex=hex(1), src=vgpr(rotation), comment="rotation=(lds_row_id //2) * 2")) module.add(VAddU32(dst=vgpr(colOffset), src0=vgpr(rotation), src1=vgpr(lane16Group), comment="colOffset = rotation + lane16Group")) setExecMask(module, writer, 0x33333333, 0x33333333) module.add(VPermlane16SwapB32(dst=vgpr(colOffset), src=vgpr(colOffset), comment="apply swizzling")) setExecMask(module, writer, -1, -1) module.add(VAndB32(dst=vgpr(colOffset), src0=vgpr(colOffset), src1=hex(blockSize-1), comment="colOffset = colOffset %% blockSize")) # Without swizzling, the LDS M-row stride is depthUBytes (contiguous K row). # With swizzling, GR writes individual subtile K-groups, so subIterKBytes applies. # TDM pad adds 16B per row, breaking pow2; fall back to VMul when padded. if padBytes == 0: module.add(VLShiftLeftB32(dst=vgpr(rowOffset), shiftHex=hex(ldsRowStride.bit_length()-1), src=vgpr(lane16), comment="offsetRow = %d*lane16" % ldsRowStride)) else: module.add(VMulLOU32(dst=vgpr(rowOffset), src0=hex(ldsRowStride), src1=vgpr(lane16), comment="offsetRow = %d*lane16" % ldsRowStride)) _computeLROffset(module, tileInfoA, colOffset, rowOffset, writer.states.subtileLdsSwizzle) _computeLROffset(module, tileInfoB, colOffset, rowOffset, writer.states.subtileLdsSwizzle) writer.vgprPool.checkIn(tmpVgpr) _lraWavePartitioning_legacy(module, writer, kernel) for vgprId in range(len(tileInfoB.sharedVgprLROffset)): module.add(VAddU32(dst=vgpr(tileInfoB.sharedVgprLROffset[vgprId]), src0=writer.ldsStartOffsetB, src1=vgpr(tileInfoB.sharedVgprLROffset[vgprId]), comment="B matrix offset in LDS")) return module def localReadResetOffsetsSubtile(writer, kernel): module = Module() module.addComment0("REMOVE WHEN IMPLEMNTED: Placeholder for subtile based LR offset reset code") for i in range(8): module.addComment("") return module def emitSingleDsRead(tileInfo, sId0, sId1, subIterK, dstTile, swizzled=True): """Emit DSLoadB128 instruction(s) for one MMA tile within a subtile. For wave32 tiles with 8 VGPRs, emits two DSLoadB128 instructions (each loading 4 VGPRs) since ds_load_b256 is not available. Args: tileInfo: TileInfo (for subtileSize, loadRatioGR, sharedVgprLROffset, tc) sId0: Subtile row index (used for offset computation) subIterK: subIterK index within the subtile (maps to mfmaC; subtileShape[0]=1 so mfmaR=0) dstTile: RegisterTileInfo \u2014 destination vgpr tile for the load swizzled: If True, LDS uses swizzled subtile layout; if False, contiguous K-row layout Returns a Module. For tiles with numRegs > 4 (e.g. FP8 8-VGPR tiles), emits multiple ds_read_b128 instructions (one per 4 VGPRs), each using the next sharedVgprLROffset entry. """ REGS_PER_DS_READ = tileInfo.loadWidthLR // 4 # load width in bytes / 4 bytes per VGPR # du maps to mfmaC, mfmaR is always 0 (subtileShape[0]=1) mfmaId = tileInfo.getSubtileShapeLinearId(subIterK, 0) if swizzled: # Swizzled: GR writes individual subtile K-groups into LDS. offsetStride = int(tileInfo.subtileSize) offset = sId0 * offsetStride + sId1 * int(tileInfo.globalSubtileGrid[0]) * offsetStride else: # Non-swizzled: full DepthU tile is contiguous in LDS with K as the fast # dimension. Each M-row is depthUBytes wide. A subtile row covers # subtileShape[0] * instM M-rows, so stride = that * depthUBytes. instM = int(tileInfo.mmaTileShape[0]) instK = int(tileInfo.mmaTileShape[1]) subtileShapeM = int(tileInfo.subtileShape[0]) subtileShapeK = int(tileInfo.subtileShape[1]) depthUBytes = int(tileInfo.depthUBytes) # Add padding rowPadBytes = getattr(tileInfo, "ldsRowPadBytes", 0) rowStride = depthUBytes + rowPadBytes offsetStride = subtileShapeM * instM * rowStride offset = sId0 * offsetStride + sId1 * subtileShapeK * instK * int(tileInfo.bpe) dstVgpr = dstTile.regList.indices[0] numRegs = len(dstTile.regList.indices) numReadsForTile = numRegs // REGS_PER_DS_READ module = Module() for readIdx in range(numReadsForTile): addrVgpr = tileInfo.sharedVgprLROffset[mfmaId * numReadsForTile + readIdx] module.add(DSLoadB128( dst=vgpr(dstVgpr + readIdx * REGS_PER_DS_READ, REGS_PER_DS_READ), src=vgpr(addrVgpr), ds=DSModifiers(offset=offset), comment="Subtile%s[%u, %u] subIterK=%u read=%u" % (tileInfo.tc, sId0, sId1, subIterK, readIdx))) return module def emitSubtileDsRead(writer, kernel, tileInfo, subtileId): module = Module() sId0 = subtileId[0] sId1 = subtileId[1] REGS_PER_DS_READ = tileInfo.loadWidthLR // 4 # load width in bytes / 4 bytes per VGPR offsetStride = int(tileInfo.subtileSize) offset = sId0 * offsetStride + sId1 * int(tileInfo.globalSubtileGrid[0]) * offsetStride lrOffsetIdx = 0 for du in range(tileInfo.subtileShape[1]): mfmaId = tileInfo.getSubtileShapeLinearId(du, 0) tileIdx = tileInfo.lrTileIndexForSubtile(sId0, sId1, mfmaId) dstTile = tileInfo.vgprTiles[tileIdx] dstVgpr = dstTile.regList.indices[0] numRegs = len(dstTile.regList.indices) # Each tile may need multiple ds_read_b128 when numRegs > 4 (e.g. FP8 8-vgpr tiles). # Each read uses the next sharedVgprLROffset entry. numReadsForTile = numRegs // REGS_PER_DS_READ for readIdx in range(numReadsForTile): addrVgpr = tileInfo.sharedVgprLROffset[lrOffsetIdx] module.add(DSLoadB128( dst=vgpr(dstVgpr + readIdx * REGS_PER_DS_READ, REGS_PER_DS_READ), src=vgpr(addrVgpr), ds=DSModifiers(offset=offset), comment="Subtile%s[%u, %u] subIterK=%u read=%u" % (tileInfo.tc, sId0, sId1, du, readIdx))) lrOffsetIdx += 1 return module ################################################## # Subroutine to generate LR load code # Initial idea: maybe store asm in modules in a separate obj? # def localReadDoSubtile(tc, writer, kernel): module = Module() tileInfo = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo for i in range(tileInfo.localSubtileGrid[0]): for j in range(tileInfo.localSubtileGrid[1]): module.add(emitSubtileDsRead(writer, kernel, tileInfo, [i, j])) return module def localReadDTLInitCommonSwapVgpr(writer, kernel): module = Module() atile = writer.states.a.tileInfo btile = writer.states.b.tileInfo stmp = writer.sgprPool.checkOut(1, tag="_localReadDTLInitCommonSwapVgpr_stmp") module.add(SMovB32(dst=sgpr(stmp), src=writer.ldsTotalSize, comment="Store Total Lds Size for one buffer")) for i in range(len(atile.sharedVgprLROffset)): vgprId = atile.sharedVgprLROffset[i] vgprSwapId = atile.sharedVgprLROffsetSwap[i] module.add(VAddU32(dst=vgpr(vgprSwapId), src0=vgpr(vgprId), src1=sgpr(stmp), comment="")) module.add(VXorB32(dst=vgpr(vgprSwapId), src0=vgpr(vgprId), src1=vgpr(vgprSwapId), comment="")) for i in range(len(btile.sharedVgprLROffset)): vgprId = btile.sharedVgprLROffset[i] vgprSwapId = btile.sharedVgprLROffsetSwap[i] module.add(VAddU32(dst=vgpr(vgprSwapId), src0=vgpr(vgprId), src1=sgpr(stmp), comment="")) module.add(VXorB32(dst=vgpr(vgprSwapId), src0=vgpr(vgprId), src1=vgpr(vgprSwapId), comment="")) writer.sgprPool.checkIn(stmp) return module ################################################## # Subroutine to generate DTL M0 LDS buffer swap # def localReadLDSBufferSwap(tc, writer, kernel): if tc in ['A', 'B']: ti_ = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo return ti_.emitLRLDSBufferSwap(writer, kernel) else: ti_ = writer.states.mxsa.tileInfo if tc == 'MXSA' else writer.states.mxsb.tileInfo return emitScaleLRLDSSwap(ti_, writer, kernel)