# Copyright Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT """Instruction emitter for LogicalScheduler. Converts the logical schedule (EmittedModule chains) into concrete GPU instructions by dispatching each opType to its emit method. """ from __future__ import annotations from Tensile.Components.Subtile.Kernel import emitMfmaInstruction from Tensile.Components.Subtile.SubtileGREmit import ( emitSingleBufferLoad, globalReadPtrUpdates, globalReadLDSBufferSwap, ) from Tensile.Components.Subtile.SubtileLREmit import ( emitSingleDsRead, localReadLDSBufferSwap, ) from Tensile.Components.Subtile.SubtileScaleEmit import ( globalReadDoScaleSubtile, globalReadScalePtrUpdates, ) from rocisa.code import Module from rocisa.instruction import ( SWaitCnt, SBarrier, DSLoadB32, SCmpEQU32, SCmpLeU32, SCBranchSCC1, SMovB32, VAddU32, VAndB32, VCmpGEI32, VCmpGTI32, VCmpLeI32, VCmpLtI32, VCndMaskB32, VLShiftLeftB32, VLShiftRightB32, VMovB32, VSubI32, ) from rocisa.instruction import SWaitTensorcnt from rocisa.container import vgpr, sgpr, DSModifiers, ContinuousRegister from rocisa.code import Label class SWaitCntEx(SWaitCnt): """SWaitCnt with adjustVmcnt flag for the instruction scheduler post-pass. isWaitGr is a typed marker the instruction scheduler uses to identify the wait_gr SWaitCnt (instead of substring-matching the instruction comment). """ def __init__(self, adjustVmcnt=True, isWaitGr=True, **kwargs): super().__init__(**kwargs) self._adjustVmcnt = adjustVmcnt self._isWaitGr = isWaitGr @property def adjustVmcnt(self): return self._adjustVmcnt @property def isWaitGr(self): return self._isWaitGr def __deepcopy__(self, memo): return SWaitCntEx( adjustVmcnt=self._adjustVmcnt, isWaitGr=self._isWaitGr, vlcnt=self.vlcnt, vscnt=self.vscnt, dscnt=self.dscnt, kmcnt=self.kmcnt, comment=self.comment) def _zigzag_order(rows, cols): """Return (row, col) pairs in boustrophedon (zigzag/snake) order. Even rows traverse left-to-right, odd rows right-to-left. Row transitions share the column coordinate with the last element of the previous row, so every consecutive pair shares exactly one coordinate -- guaranteeing a VGPR source-cache hit at every step. """ result = [] for r in range(rows): if r % 2 == 0: for c in range(cols): result.append((r, c)) else: for c in range(cols - 1, -1, -1): result.append((r, c)) return result class InstructionEmitter: """Emits GPU instructions for each opType in the LogicalScheduler output. VGPR tile indexing uses placement-level tile maps (tileId → vgprTileId) set by assign_vgpr_tiles(). Per-tensor VGPR tile lists are indexed by vgprTileId. All tensors (A, B, SA, SB) use the same tile-map approach. """ def __init__(self, writer, kernel, config, tileInfoA, tileInfoB, dtileInfo, vgprTilesA, vgprTilesB, scaleTileInfoA=None, scaleTileInfoB=None, vgprTilesSA=None, vgprTilesSB=None, tensorParametersA=None, tensorParametersB=None): self.writer = writer self.kernel = kernel self.config = config self.tileInfoA = tileInfoA self.tileInfoB = tileInfoB self.dtileInfo = dtileInfo self.vgprTilesA = vgprTilesA self.vgprTilesB = vgprTilesB self.vgprTilesSA = vgprTilesSA or [] self.vgprTilesSB = vgprTilesSB or [] self.tensorParametersMap = {} if tensorParametersA is not None: self.tensorParametersMap['A'] = tensorParametersA if tensorParametersB is not None: self.tensorParametersMap['B'] = tensorParametersB # Derived state self.hasScale = scaleTileInfoA is not None and scaleTileInfoB is not None self.subtileShapeK = tileInfoA.subtileShape[1] self.tileInfoMap = {'A': tileInfoA, 'B': tileInfoB} if self.hasScale: self.tileInfoMap['SA'] = scaleTileInfoA self.tileInfoMap['SB'] = scaleTileInfoB # Per-uid K-window: numSubIterK / numUnroll[t]. Single-DU configs # (numUnroll=1) span the full numSubIterK range; multi-DU slices it. self._per_uid_k = { t: config.numSubIterK // config.numUnroll.get(t, 1) for t in ('A', 'B') } if self.hasScale: self._per_uid_k['SA'] = config.numSubIterK // config.numUnroll.get('SA', 1) self._per_uid_k['SB'] = config.numSubIterK // config.numUnroll.get('SB', 1) # 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), 'skip': lambda em, ui: self.emit_skip(em.source), 'mask_k': lambda em, ui: self.emit_mask_k(em.source), 'inline': lambda em, ui: self.emit_inline(em.source), 'gl2_prefetch': lambda em, ui: self.emit_gl2_prefetch(), 'gl2_prefetch_inc': lambda em, ui: self.emit_gl2_prefetch_inc(), } # Sentinel for the long-lived per-lane diff vgpr. Set by # emit_mask_k_init, consumed by every emit_mask_k call in the tail body. self._tail_vDiff = None def emit_mfma(self, placement, unroll_iter=0): """Emit MFMA instructions from MFMAPlacement.""" module = Module() subIterK = placement.subIterK tile_maps = {t: placement.vgpr_tile_maps[t][unroll_iter] for t in placement.vgpr_tile_maps} aTiles = placement.tileA.tileId_list bTiles = placement.tileB.tileId_list hasWmmaSourceCache = self.writer.states.asmCaps.get("HasWMMA_V3", False) if hasWmmaSourceCache and len(aTiles) > 1 and len(bTiles) > 1: # Zigzag ordering: traverse the AxB tile grid in snake order so # every consecutive WMMA pair shares exactly one operand (A or B), # guaranteeing a VGPR source-cache hit at every step. # Always sweep along the longer tile dimension for maximum reuse. if len(aTiles) >= len(bTiles): abPairs = [(aTiles[c], bTiles[r]) for r, c in _zigzag_order(len(bTiles), len(aTiles))] else: abPairs = [(aTiles[r], bTiles[c]) for r, c in _zigzag_order(len(aTiles), len(bTiles))] else: abPairs = [(a, b) for a in aTiles for b in bTiles] for a, b in abPairs: groupA = (a // self.config.lrA.mn) * self.config.lrA.mn groupB = (b // self.config.lrB.mn) * self.config.lrB.mn aTile = self.vgprTilesA[tile_maps['A'][groupA]] bTile = self.vgprTilesB[tile_maps['B'][groupB]] dTile = self.dtileInfo.vgprTiles[a + b * self.dtileInfo.localMMATileGrid[0]] if self.hasScale: scaleGroupA = (a // self.config.lrSA.mn) * self.config.lrSA.mn scaleGroupB = (b // self.config.lrSB.mn) * self.config.lrSB.mn scaleATile = self.vgprTilesSA[tile_maps['SA'][scaleGroupA]] scaleBTile = self.vgprTilesSB[tile_maps['SB'][scaleGroupB]] scaleAVgpr = next(iter(scaleATile)) scaleBVgpr = next(iter(scaleBTile)) mShapeA = self.tileInfoMap['SA'].lrSubtileShape[0] mShapeB = self.tileInfoMap['SB'].lrSubtileShape[0] kShapeA = self.tileInfoMap['SA'].lrSubtileShape[1] kShapeB = self.tileInfoMap['SB'].lrSubtileShape[1] sAsel = (a % mShapeA) + mShapeA * (subIterK % kShapeA) sBsel = (b % mShapeB) + mShapeB * (subIterK % kShapeB) else: scaleAVgpr = scaleBVgpr = -1 sAsel = sBsel = 0 module.add(emitMfmaInstruction( self.writer, self.kernel, aTile, bTile, dTile, dTile, scaleAVgpr=scaleAVgpr, scaleBVgpr=scaleBVgpr, scaleAsel=sAsel, scaleBsel=sBsel, comment=f"MFMA C[{a},{b}] += A[{a},K={subIterK}] * B[{b},K={subIterK}]")) return list(module.flatitems()) def emit_lr(self, placement, unroll_iter=0): """Emit LR (ds_read) instructions from LRPlacement.""" module = Module() tensor = placement.tensor tile_map = placement.vgpr_tile_map[unroll_iter] if placement.vgpr_tile_map else {} if tensor in ('A', 'B'): ti = self.tileInfoMap[tensor] vgprTiles = self.vgprTilesA if tensor == 'A' else self.vgprTilesB lrGran = self.config.lrA if tensor == 'A' else self.config.lrB # `per_uid_k` (= numSubIterK / numUnroll) folds k into a # uid-local subtile coordinate so multi-DU (numUnroll > 1) # uids each iterate [0, per_uid_k). Under single-DU # (numUnroll == 1) per_uid_k == numSubIterK and the modulo # is a no-op for in-window k, but it ALIASES wrap-LRs that # use k >= numSubIterK to address the alternate LDS half # (PGR>=1 prefetch target) back to k=0 — losing the half # offset. Skip the modulo in that case. nUnroll = self.config.numUnroll.get(tensor, 1) per_uid_k = self._per_uid_k[tensor] if nUnroll > 1 else None for tileId in range(placement.tiles.tileId_start, placement.tiles.tileId_end, lrGran.mn): for k in range(placement.tiles.subIterK_start, placement.tiles.subIterK_end, lrGran.k): local_k = (k % per_uid_k) if per_uid_k is not None else k subtileK = local_k // self.subtileShapeK subIterK_within = local_k % self.subtileShapeK dstTile = vgprTiles[tile_map[tileId]] swizzled = self.writer.states.subtileLdsSwizzle module.add(emitSingleDsRead( ti, tileId, subtileK, subIterK_within, dstTile, swizzled=swizzled)) elif tensor in ('SA', 'SB'): tc = 'MXSA' if tensor == 'SA' else 'MXSB' ti = self.tileInfoMap[tensor] lrGran = self.config.lrSA if tensor == 'SA' else self.config.lrSB vgprTilesScale = self.vgprTilesSA if tensor == 'SA' else self.vgprTilesSB for tileId in range(placement.tiles.tileId_start, placement.tiles.tileId_end, lrGran.mn): scaleGroupIdx = tileId // lrGran.mn groupKey = scaleGroupIdx * lrGran.mn kGroupIdx = placement.tiles.subIterK_start // ti.lrSubtileShape[1] numKGroups = ti.lrLocalSubtileGrid[1] dsOffset = int(ti.lrSubtileSize) * (scaleGroupIdx * numKGroups + kGroupIdx) vdst = next(iter(vgprTilesScale[tile_map[groupKey]])) module.add(DSLoadB32( dst=vgpr(vdst), src=vgpr(ti.sharedVgprLROffset[0]), ds=DSModifiers(offset=dsOffset), comment=f"scale{tc}[group{scaleGroupIdx},K={placement.tiles.subIterK_start}]: load 4B from LDS")) return list(module.flatitems()) def emit_gr(self, placement): """Emit GR (buffer_load) instructions from GRPlacement.""" module = Module() tensor = placement.tensor if tensor in ('A', 'B'): ti = self.tileInfoMap[tensor] grGran = self.config.grA if tensor == 'A' else self.config.grB uid_k_base = placement.unrollId * grGran.k for tileId in range(placement.tiles.tileId_start, placement.tiles.tileId_end, grGran.mn): for k in range(placement.tiles.subIterK_start, placement.tiles.subIterK_end, grGran.k): subtileK = (k - uid_k_base) // self.subtileShapeK module.add(emitSingleBufferLoad(ti, self.kernel, tileId, subtileK)) elif tensor in ('SA', 'SB'): tc = 'MXSA' if tensor == 'SA' else 'MXSB' module.add(globalReadDoScaleSubtile(tc, self.writer, self.kernel)) return list(module.flatitems()) def emit_wait_gr(self, source): """Emit SWaitCnt for wait_gr from BaseOp with wait_gr_counts.""" counts = source.wait_gr_counts if counts is None: return [] # force_drain means "wait for every outstanding global read to retire" # (full vmcnt(0) drain); the per-tensor counts are kept only for the # diagnostic comment. See WaitGROp.force_drain in LogicalScheduler.py. force_drain = getattr(source, 'force_drain', False) if self.kernel.get("enableTDMA", False) and self.kernel.get("enableTDMB", False): tdmCnt = 0 if force_drain else (counts.A + counts.B + counts.SA + counts.SB) label = "full drain" if force_drain else "tensor_load_to_lds" return [SWaitTensorcnt(tensorcnt=tdmCnt, comment=f"Wait TDM ({label}): A={counts.A} B={counts.B} SA={counts.SA} SB={counts.SB}")] # TODO. Hardcoded for now, but we should just get this from atomic emit codes (emitSingleBufferLoad, ...) grMap = {'A': max(1,int(1.0/self.tileInfoA.loadRatioGR)), 'B': max(1,int(1.0/self.tileInfoB.loadRatioGR)), 'SA': 1, 'SB': 1} if force_drain: grCnt = 0 label = "full drain" else: grCnt = (counts.A * grMap['A'] + counts.B * grMap['B'] + counts.SA * grMap['SA'] + counts.SB * grMap['SB']) label = "per-subIterK" swait = SWaitCntEx(vlcnt=grCnt, vscnt=-1, adjustVmcnt=source.adjustVmcnt, comment=f"Wait GR ({label}): A={counts.A} B={counts.B} SA={counts.SA} SB={counts.SB}") return [swait] def emit_wait_lr(self): return [SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LR to complete")] def emit_sync(self): return [SBarrier(comment="Barrier")] def emit_inline(self, source): """Emit a writer-built Module supplied by an InlineModuleOp callback.""" if source.build is None: return [] mod = source.build(self) return list(mod.flatitems()) if mod is not None else [] def emit_lr_inc(self, source): """Emit localReadLDSBufferSwap for a single tensor.""" tensor = source.tensor tc = {'A': 'A', 'B': 'B', 'SA': 'MXSA', 'SB': 'MXSB'}.get(tensor, tensor) module = Module() module.add(localReadLDSBufferSwap(tc, self.writer, self.kernel)) return list(module.flatitems()) def emit_gr_inc(self, source): """Emit globalReadPtrUpdates + globalReadLDSBufferSwap for a single tensor.""" 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)) module.add(globalReadLDSBufferSwap(tc, self.writer, self.kernel)) return list(module.flatitems()) def emit_gl2_prefetch(self): """Emit GL2 prefetch loads (global_prefetch_b8) for all tensors.""" writer = self.writer kernel = self.kernel tPA = self.tensorParametersMap['A'] tPB = self.tensorParametersMap['B'] mod = writer.gl2PrefetchIssueLoad(kernel, tPA, tPB) return list(mod.flatitems()) def emit_gl2_prefetch_inc(self): """Emit GL2 prefetch address increment with end-of-K guard. Mirrors the SIA path: when LoopCounterL <= PGR + PrefetchGL2, zero the increment SGPRs so prefetch stops advancing past K. Then advance all GL2 prefetch addresses by the (possibly zeroed) increment. """ writer = self.writer kernel = self.kernel tPA = self.tensorParametersMap['A'] tPB = self.tensorParametersMap['B'] from rocisa.code import Module from rocisa.instruction import SCmpLeU32, SCMovB32 from rocisa.container import sgpr mod = Module("GL2 Prefetch Increment") loopCounter = writer.loopCounter(kernel, writer.states.unrollIdx) pgl = kernel["PrefetchGL2"] pgr = kernel["PrefetchGlobalRead"] mod.add(SCmpLeU32(src0=loopCounter, src1=pgr + pgl, comment=f"counterL <= PGR({pgr})+PGL({pgl})?")) mod.add(SCMovB32(dst=sgpr("GL2PrefetchIncA"), src=0)) mod.add(SCMovB32(dst=sgpr("GL2PrefetchIncB"), src=0)) if kernel["ProblemType"].get("MXBlockA", 0): mod.add(SCMovB32(dst=sgpr("GL2PrefetchIncMXSA"), src=0)) if kernel["ProblemType"].get("MXBlockB", 0): mod.add(SCMovB32(dst=sgpr("GL2PrefetchIncMXSB"), src=0)) mod.add(writer.gl2PrefetchIncrementAddr(kernel, tPA, tPB)) return list(mod.flatitems()) def emit_skip(self, source): """Emit skip guard: compare LoopCounterL and branch.""" labelName = source.target if source.rawLabel else f"SkipTo{source.target}" skipLabel = Label(labelName, "") cmpMap = {"EQ": SCmpEQU32, "LE": SCmpLeU32} cmpCls = cmpMap[source.compare] module = Module() # SCMP inline range is -16..64 (signed); stage non-inline via scratch sgpr. if -16 <= source.value <= 64: module.add(cmpCls( src0=sgpr("LoopCounterL"), src1=source.value, comment=f"LoopCounter {source.compare} {source.value}?")) else: with self.writer.allocTmpSgpr(1, tag="InstructionEmitter_skip_tmpSgpr") as litSgprInfo: litSgpr = litSgprInfo.idx module.add(SMovB32( dst=sgpr(litSgpr), src=hex(source.value), comment=f"stage literal {source.value} (non-inline) for cmp src1")) module.add(cmpCls( src0=sgpr("LoopCounterL"), src1=sgpr(litSgpr), comment=f"LoopCounter {source.compare} {source.value}?")) module.add(SCBranchSCC1( labelName=skipLabel.getLabelName(), comment=source.branchComment or f"skip to {source.target}")) return list(module.flatitems()) def _mfma_K_constants(self): """Constants used by both mask emitters. Returns (numMIInUnroll, dividerFortidInK): * numMIInUnroll = MI_M * MI_K / WavefrontSize — per-lane K chunk consumed by one MFMA call (contiguous K-elements held by each lane). * dividerFortidInK = MI_M — lanes with the same (Serial % WavefrontSize) / MI_M share the same K-position; WavefrontSize / MI_M = number of K-groups in the wave. """ kernel = self.kernel MI_M = kernel["MatrixInstM"] MI_K = kernel["MatrixInstK"] waveSize = kernel["WavefrontSize"] numMIInUnroll = MI_M * MI_K // waveSize dividerFortidInK = MI_M return numMIInUnroll, dividerFortidInK def emit_mask_k_init(self): """Stage the per-subIterK invariants for emit_mask_k. Only two vgprs survive past this routine and into emit_mask_k: * self._tail_vDiff — rem - laneK_0 (for subIterK=n the effective diff is diff - n*MatrixInstK, folded into cmp constants) * self._tail_boundaryMask[4] — BF16 only; precomputed boundary masks indexed by vgpr i, derived from d = rem % 8 """ numMIInUnroll, dividerFortidInK = self._mfma_K_constants() writer = self.writer module = Module() # Compute K per lane (numMIInUnroll, dividerFortidInK are power of 2) waveSize = self.kernel["WavefrontSize"] divLog2 = (dividerFortidInK).bit_length() - 1 mulLog2 = (numMIInUnroll).bit_length() - 1 kReg = writer.vgprPool.checkOut(1, "tail_kReg") module.add(VAndB32( dst=vgpr(kReg), src0=waveSize - 1, src1=vgpr("Serial"), comment=f"tail_kReg = Serial & {waveSize - 1} (Serial % {waveSize})")) module.add(VLShiftRightB32( dst=vgpr(kReg), shiftHex=divLog2, src=vgpr(kReg), comment=f"tail_kReg >>= {divLog2} (tail_kReg / {dividerFortidInK})")) # laneK_0 and diff = rem - laneK_0 are shared across all subIterK. workKVgpr = writer.vgprPool.checkOut(1, "tail_workK") self._tail_vDiff = writer.vgprPool.checkOut(1, "tail_vDiff") loopCounterName = writer.loopCounterName( self.kernel, writer.states.unrollIdx) module.add(VLShiftLeftB32( dst=vgpr(workKVgpr), shiftHex=mulLog2, src=vgpr(kReg), comment=f"laneK_0 = tail_kReg << {mulLog2} (tail_kReg * {numMIInUnroll})")) writer.vgprPool.checkIn(kReg) module.add(VSubI32( dst=vgpr(self._tail_vDiff), src0=sgpr(loopCounterName), src1=vgpr(workKVgpr), comment="diff = rem - laneK_0 (shared across all subIterK)")) writer.vgprPool.checkIn(workKVgpr) self._tail_boundaryMask = None laneSGPRCount = writer.states.laneSGPRCount if self.kernel["ProblemType"]["DataTypeA"].isBFloat16(): halfMaskVgpr = writer.vgprPool.checkOut(1, "tail_halfMask") module.add(VMovB32( dst=vgpr(halfMaskVgpr), src="0x0000FFFF", comment="BF16 half-mask: keep K0 (low 16b), zero K1 (high 16b)")) # Precompute the boundary masks from d = rem % numMIInUnroll. # The boundary-mask pattern (which vgprs are full/half/zero) # depends only on rem % numMIInUnroll: laneK_0 is always a # multiple of numMIInUnroll, and (assuming MatrixInstK is a # multiple of numMIInUnroll) every "boundary" lane in every # subIterK has effective_diff ≡ rem (mod numMIInUnroll). So all # boundary lanes share the same mask pattern (one per vgpr held # per lane: numMIInUnroll // kStride masks). kStride = 2 # BF16: 2 K-elements packed per 32-bit vgpr assert numMIInUnroll % kStride == 0, \ f"numMIInUnroll ({numMIInUnroll}) must be a multiple of kStride ({kStride})" numBoundaryMasks = numMIInUnroll // kStride vDLaneRem = writer.vgprPool.checkOut(1, "tail_vDLaneRem") module.add(VAndB32( dst=vgpr(vDLaneRem), src0=sgpr(loopCounterName), src1=numMIInUnroll - 1, comment=f"d = rem % {numMIInUnroll} (boundary-mask pattern depends only on this)")) self._tail_boundaryMask = [ writer.vgprPool.checkOut(1, f"tail_boundaryMask{i}") for i in range(numBoundaryMasks) ] with writer.allocTmpSgpr(laneSGPRCount, alignment=laneSGPRCount, tag="InstructionEmitter_mask_k_init_tmpSgpr") as tmpSgprInfo: maskSgpr = tmpSgprInfo.idx for i in range(numBoundaryMasks): bm = self._tail_boundaryMask[i] hiBound = i * kStride + kStride # d < hiBound → half-keep loBound = i * kStride + 1 # d < loBound → zero module.add(VCmpLtI32( dst=sgpr(maskSgpr, laneSGPRCount), src0=vgpr(vDLaneRem), src1=hiBound, comment=f"boundary[{i}]: d < {hiBound} ? halfKeep : full")) module.add(VCndMaskB32( dst=vgpr(bm), src0=-1, src1=vgpr(halfMaskVgpr), src2=sgpr(maskSgpr, laneSGPRCount), comment=f"boundaryMask[{i}] = (d<{hiBound}) ? halfKeep : full")) module.add(VCmpLtI32( dst=sgpr(maskSgpr, laneSGPRCount), src0=vgpr(vDLaneRem), src1=loBound, comment=f"boundary[{i}]: d < {loBound} ? 0 : prev")) module.add(VCndMaskB32( dst=vgpr(bm), src0=vgpr(bm), src1=0, src2=sgpr(maskSgpr, laneSGPRCount), comment=f"boundaryMask[{i}] = (d<{loBound}) ? 0 : prev")) writer.vgprPool.checkIn(halfMaskVgpr) writer.vgprPool.checkIn(vDLaneRem) return list(module.flatitems()) def emit_mask_k(self, source): """Per-lane K-mask for one subIterK, applied to A/B tiles via V_AND_B32. Uses diff = rem - laneK_0 from emit_mask_k_init. For subIterK=n the effective per-lane diff is diff - n*MatrixInstK, folded into the cmp immediates (no per-call sub). BF16 (kStride=2 K positions per vgpr) builds a per-vgpr 3-state mask from the precomputed boundary masks (full / boundary[i] / zero). Non-BF16 (e.g. FP4) builds a single 2-state mask shared across all vgprs, assuming rem aligns to the per-lane K stride (true for rem=32 with FP4 MIK=128). A boundary inside a non-BF16 vgpr would need per-byte/nibble handling. """ assert self._tail_vDiff is not None, \ "emit_mask_k_init must run before emit_mask_k" writer = self.writer kernel = self.kernel subIterK = source.subIterK kBaseConst = subIterK * kernel["MatrixInstK"] laneSGPRCount = writer.states.laneSGPRCount isBF16 = kernel["ProblemType"]["DataTypeA"].isBFloat16() kStride = 2 # BF16: 2 elements packed per 32-bit vgpr (low=K0, high=K1) module = Module() def _unique_ids(key): m = source.vgpr_tile_map.get(key, [{}])[0] return sorted(set(m.values())) aIds, bIds = _unique_ids('A'), _unique_ids('B') # All A/B tiles share the same vgprs-per-lane count and K layout, # so the i-th vgpr of every tile gets the same mask. refTiles = self.vgprTilesA if aIds else self.vgprTilesB refIds = aIds or bIds vgprPerInUnroll = len(list(refTiles[refIds[0]])) if refIds else 0 with writer.allocTmpSgpr(laneSGPRCount, alignment=laneSGPRCount, tag="InstructionEmitter_mask_k_tmpSgpr") as tmpSgprInfo: maskSgpr = tmpSgprInfo.idx def _emit_cmp(cmpCls, literal, comment): # VOPC inline range is -16..64; stage out-of-range via scratch sgpr # (e.g. BF16 MI_K=32 subIterK>=2, or FP4 MI_K=128 subIterK>=1). if -16 <= literal <= 64: module.add(cmpCls( dst=sgpr(maskSgpr, laneSGPRCount), src0=vgpr(self._tail_vDiff), src1=literal, comment=comment)) else: with writer.allocTmpSgpr(1, tag="InstructionEmitter_mask_k_litSgprInfo") as litSgprInfo: litSgpr = litSgprInfo.idx module.add(SMovB32( dst=sgpr(litSgpr), src=hex(literal), comment=f"stage literal {literal} (non-inline)")) module.add(cmpCls( dst=sgpr(maskSgpr, laneSGPRCount), src0=vgpr(self._tail_vDiff), src1=sgpr(litSgpr), comment=comment)) if isBF16: # 3-way per (lane, subIterK), 2 cmps shared across all i: # sFull = effective_diff_n >= numMIInUnroll → -1 # sZero = effective_diff_n <= 0 → 0 # else → boundary[i] numMIInUnroll = vgprPerInUnroll * kStride fullLit = kBaseConst + numMIInUnroll - 1 zeroLit = kBaseConst maskVgprs = [writer.vgprPool.checkOut(1, f"mask_k_msk{i}_k{subIterK}") for i in range(vgprPerInUnroll)] _emit_cmp(VCmpGTI32, fullLit, f"sFull: diff > {fullLit} (effective_diff_{subIterK} >= {numMIInUnroll})") for i in range(vgprPerInUnroll): module.add(VCndMaskB32( dst=vgpr(maskVgprs[i]), src0=vgpr(self._tail_boundaryMask[i]), src1=-1, src2=sgpr(maskSgpr, laneSGPRCount), comment=f"mask[{i}] = sFull ? full : boundary[{i}]")) _emit_cmp(VCmpLeI32, zeroLit, f"sZero: diff <= {zeroLit} (effective_diff_{subIterK} <= 0)") for i in range(vgprPerInUnroll): module.add(VCndMaskB32( dst=vgpr(maskVgprs[i]), src0=vgpr(maskVgprs[i]), src1=0, src2=sgpr(maskSgpr, laneSGPRCount), comment=f"mask[{i}] = sZero ? 0 : prev")) else: # 2-state shared mask: diff < kBaseConst+1 → 0, else -1. literal = kBaseConst + 1 sharedMask = writer.vgprPool.checkOut(1, f"mask_k_msk_k{subIterK}") maskVgprs = [sharedMask] * vgprPerInUnroll _emit_cmp(VCmpLtI32, literal, f"mask: diff < {literal} (laneK_{subIterK} >= rem)") module.add(VCndMaskB32( dst=vgpr(sharedMask), src0=-1, src1=0, src2=sgpr(maskSgpr, laneSGPRCount), comment=f"mask = (diff < {literal}) ? 0 : -1")) # MX-scaled paths apply the K mask via the scale lanes (host # zero-pads OOB scale bytes to 0x00, which v_mfma_scale treats # as ~2^-127 and underflows the contribution to zero). Emitting # a second data mask here is redundant and, in practice on # gfx950, the per-lane V_AND against the cndmask-derived mask # corrupts the v_mfma_scale operand bypass and double-zeros # valid residual lanes (~0.5% wrong outputs at K%MIK!=0). # # Skip the data K-mask for the MX paths that the gfx950 corruption # actually hits: all multi-DU MX kernels, and all 8-bit-float MX # (MXFP8) kernels regardless of unroll depth. GPU-validated on # gfx950 (full-element, K%MIK!=0 tails): single-DU MXFP8 fails with # the mask and passes without it, while MXFP4 (4-bit) passes either # way -- so single-DU MXFP4 keeps the mask and stays develop- # identical, but single-DU MXFP8 must skip it like multi-DU does. _emit_mask_multi_du = (self.config.numUnroll.get('A', 1) > 1 or self.config.numUnroll.get('B', 1) > 1) _data_is_8bit_float = kernel["ProblemType"]["DataTypeA"].is8bitFloat() _skip_data_kmask = self.hasScale and (_emit_mask_multi_du or _data_is_8bit_float) if not _skip_data_kmask: for label, ids, tilesDict in (("A", aIds, self.vgprTilesA), ("B", bIds, self.vgprTilesB)): for tid in ids: for i, v in enumerate(list(tilesDict[tid])): module.add(VAndB32( dst=vgpr(v), src0=vgpr(v), src1=vgpr(maskVgprs[i]), comment=f"mask {label}[{i}] (K=[{i*kStride},{i*kStride+kStride-1}])")) # Scale-mask reuse (non-BF16 hasScale only): AND the scale # vgprs with the data mask we just built. Redundant when # host zero-padding leaves OOB scale bytes at 0x00 (0x00 # AND 0 = 0x00) but kept to stay explicit under any input # pattern. scaleStride = self.config.lrSA.k if self.hasScale else 0 if (not isBF16) and self.hasScale \ and (self.vgprTilesSA or self.vgprTilesSB) \ and scaleStride > 0 and (subIterK % scaleStride == 0): for tensor, tilesList in (('SA', self.vgprTilesSA), ('SB', self.vgprTilesSB)): liveIds = sorted(set( source.vgpr_tile_map.get(tensor, [{}])[0].values())) for tid in liveIds: for v in list(tilesList[tid]): module.add(VAndB32( dst=vgpr(v), src0=vgpr(v), src1=vgpr(maskVgprs[0]), comment=f"mask scale vgpr (reuse A/B mask, subIterK={subIterK})")) for m in set(maskVgprs): writer.vgprPool.checkIn(m) return list(module.flatitems()) def emit_mask_k_done(self): """Release the long-lived tail-loop vgprs (scratch ones are released at the end of emit_mask_k_init).""" if getattr(self, "_tail_vDiff", None) is not None: self.writer.vgprPool.checkIn(self._tail_vDiff) self._tail_vDiff = None if getattr(self, "_tail_boundaryMask", None) is not None: for bm in self._tail_boundaryMask: self.writer.vgprPool.checkIn(bm) self._tail_boundaryMask = None return [] def populate(self, emitted, unroll_iter=0): """Walk emitted partitions and fill em.instructions.""" for partition_emitted in emitted: for emitted_group in partition_emitted: for em in emitted_group: handler = self._dispatch.get(em.opType) if handler: em.instructions = handler(em, unroll_iter)