################################################################################ # # Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # ################################################################################ from rocisa import rocIsa, countInstruction, countGlobalRead, countSMemLoad, findInstCount from rocisa.asmpass import getActFuncModuleName, getActFuncBranchModuleName from rocisa.code import KernelBody, Label, Macro, Module, RegSet, SrdUpperValue, \ StructuredModule, TextBlock, ValueEndif, ValueIf, ValueSet, SignatureBase from rocisa.container import DSModifiers, SDWAModifiers, VOP3PModifiers, True16Modifiers, \ MUBUFModifiers, SMEMModifiers, EXEC, VCC, RegisterContainer, \ DPPModifiers, vgpr, sgpr, accvgpr, mgpr, ContinuousRegister, \ HWRegContainer, GLOBALModifiers, MemTokenData from rocisa.instruction import SGetPositivePCOffset, SLongBranch, SLongBranchPositive, SLongBranchNegative, SCLongBranchScc0, SCLongBranchScc1, SCLongBranchVccnz, \ SMulInt64to32, VCvtBF16toFP32, SBfeU32 from rocisa.functions import vectorStaticDivide, vectorStaticRemainder, vectorUInt32CeilDivideAndRemainder, \ vectorStaticDivideAndRemainder, scalarStaticDivideAndRemainder, scalarStaticCeilDivide, \ scalarStaticRemainder, scalarUInt32DivideAndRemainder, sMagicDiv, vectorStaticMultiply, \ vectorStaticMultiplyAdd, scalarStaticMultiply64, BranchIfZero, BranchIfNotZero, DSInit, \ ArgumentLoader, scalarMultiplyBpe, scalarMultiply64Bpe, vectorMultiplyBpe, vectorMultiply64Bpe from rocisa.enum import InstType, SelectBit, CacheScope, HighBitSel, TemporalHint, NonVolatile from rocisa.macro import MacroVMagicDiv, PseudoRandomGenerator from . import CUSTOM_KERNEL_PATH from rocisa.instruction import BranchInstruction, BufferLoadB128, BufferLoadB32, \ BufferLoadB16, BufferLoadU16, BufferLoadB64, BufferLoadB96, BufferLoadD16B16, BufferLoadD16HIB16, BufferLoadD16HIU8, \ BufferLoadD16U8, BufferStoreB128, BufferStoreB16, BufferStoreB32, BufferStoreB64, \ BufferStoreB8, BufferStoreD16HIB16, CommonInstruction, DSBPermuteB32, DSLoadB128, \ DSLoadB16, DSLoadB32, DSLoadB64, DSLoadU16, DSStoreB128, DSStoreB16, DSStoreB32, \ DSStoreB64, DSStoreB8, DSStoreInstruction, FlatLoadB128, FlatLoadB32, FlatLoadB64, \ FlatLoadD16B16, FlatLoadD16HIB16, FlatStoreB128, FlatStoreB32, FlatStoreB64, \ FlatStoreD16B16, FlatStoreD16HIB16, GlobalReadInstruction, MXMFMAInstruction, MFMAInstruction, MUBUFReadInstruction, \ MacroInstruction, SAShiftRightI32, SAbsI32, SAddCU32, SAddI32, SAddU32, SAddU64, SAndB32, \ SAndB64, SAndN2B32, SAtomicDec, SBarrier, SBfmB32, SBitcmp1B32, SBranch, SCBranchSCC0, \ SCBranchSCC1, SCBranchVCCNZ, SCBranchVCCZ, SCMovB32, SCSelectB32, SCSelectB64, SCmpEQI32, \ SCmpEQU32, SCmpEQU64, SCmpGeI32, SCmpGeU32, SCmpGtI32, SCmpGtU32, SCmpKEQU32, \ SCmpKGeU32, SCmpKGtU32, SCmpKLGU32, SCmpLeI32, SCmpLeU32, SCmpLgU32, SCmpLtU32, SCmpLtI32, \ SEndpgm, SFf1B32, SGetRegB32, SFlbitI32B32, SLShiftLeft2AddU32, SLShiftLeftB32, SLShiftLeftB64, SLShiftRightB32, \ SLShiftRightB64, SLoadB32, SLoadB64, SMFMAInstruction, SMemLoadInstruction, SMaxU32, SMinI32, \ SMinU32, SMovB32, SMovB64, SMulHIU32, SMulI32, SNop, SOrB32, SOrSaveExecB32, \ SOrSaveExecB64, SSExtI16toI32, SSetPCB64, SSetRegIMM32B32, SSetPrior, SSubBU32, SSubI32, SSubU32, SSubU64, SSetVgprMsb,\ SWaitCnt, SWaitAlu, SXorB32, VAShiftRightI32, VAccvgprReadB32, VAccvgprWrite, VAccvgprWriteB32, \ VAdd3U32, VAddCCOU32, VAddCOU32, VAddF32, VAddF64, VAddLShiftLeftU32, VAddU32, VAndB32, \ VBfeU32, VCmpEQI32, VCmpEQU32, VCmpGEI32, VCmpGEU32, VCmpGtU32, VCmpGTI32, VCmpLeI32, VCmpLtI32, \ VCmpLtU32, VCmpNeU64, VCmpUF32, VCmpXGeU32, VCmpXLtU32, VCmpXLtU64, VCndMaskB32, VCvtF16toF32, VCvtI32toF32, \ VCvtF32toF16, VCvtFP8toF32, VCvtInstruction, VCvtPkF32toBF16, VCvtPkF32toBF8, \ VCvtPkF32toFP8, VCvtPkFP8toF32, VCvtSRF32toBF8, VCvtSRF32toFP8, VCvtScaleFP8toF16, \ VCvtScalePkF16toBF8, VCvtScalePkF16toFP8, VCvtScalePkFP8toF16, VLShiftLeftB32, \ VLShiftLeftB64, VLShiftRightB32, VLShiftRightB64, VMadU32U24, VMaxF32, VMinI32, VMovB32, VMovB64, VMulF32, \ VMulHIU32, VMulLOU32, VMulPKF32S, VMulU32U24, VNotB32, VOrB32, VPackF16toB32, \ VPrngB32, VReadfirstlaneB32, VReadlaneB32, VSubF32, VSubI32, VSubU32, VXorB32, GlobalLoadTR8B64, GlobalLoadTR16B128, \ GlobalLoadB32, GlobalLoadB64, GlobalLoadB96, GlobalLoadB128, GlobalLoadD16B16, GlobalLoadD16HIB16, \ GlobalLoadD16U8, GlobalLoadD16HIU8, \ GlobalStoreB8, GlobalStoreB16, GlobalStoreB32, GlobalStoreB64, GlobalStoreB128, GlobalStoreD16HIB16 from .Component import Component, TensorDataMover, GL2Prefetch from .Components.TensorDataMover import TensorDataMoverLoad from .Components.GL2Prefetch import GL2PrefetchLoad from .Components.GlobalWriteBatch import GlobalWriteBatchWriter from .KernelWriterModules import * from .AsmMemoryHelpers import dsStore, dsLoad, _vgprOffset from .SolutionStructs import isPackedIndex from .AsmStoreState import StoreState, VectorDataTypes from .Activation import ActivationType from .CustomKernels import isCustomKernelConfig from .Common import roundUp, log2, ceilDivide, choose_multiplier, wmmaV3InputVgprLayout, clusterEnabled from .OccupancyMeasure import compute_occupancy_from_asm_source, _arch_caps_for_kernel from rocisa.instruction import ECvtF16toF32, ECvtF32toF16, ECvtPkFP8toF32 from Tensile.Common import print2, printExit, printWarning, INDEX_CHARS, DebugConfig, DataDirection, isSubtileMultiDU from Tensile.Components.NonTemporal import decodeNonTemporal, forceCoherentNonTemporal from Tensile.Common.DataType import DataType from Tensile.Common.RegisterPool import RegisterPool, allocTmpGpr, allocTmpGprList from .Components.WorkGroupMappingAlgos import DefaultWGM, wgmXCC, SpaceFillingCurveWalk from Tensile.KernelWriter import KernelWriter, ABMatrixInfo from Tensile.SolutionStructs.Naming import getKernelFileBase from Tensile.Toolchain.Component import Assembler def _cacheHintTensor(tc): return "D" if tc == "TD" else tc def _temporalHint(kernel, tc): return TemporalHint(kernel.get("TemporalHint%s"%_cacheHintTensor(tc), 0)) def _nonVolatile(kernel, tc): return NonVolatile(kernel.get("NonVolatile%s"%_cacheHintTensor(tc), 0)) from math import ceil, floor, log, prod from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, field from functools import lru_cache from typing import List, Mapping, NamedTuple, Optional, Tuple, Union import os @dataclass class TailOptParams: idx: int = 0 behavior: str = "" jumpLabel: Label = None kLabelsList: List[Label] = None tmpVgpr: int = None kSgpr: int = None vgprSlot: List[int] = None loadNum: int = 0 periodParam: List[int] = None preDirectToLdsLoads: int = 0 firstLoop: int = 0 finalLoop: int = 0 @dataclass class GprInfo: addrVgpr: List[str] = field(default_factory=list) offset: List[str] = field(default_factory=list) @dataclass class GlobalReadGprRecord: a: GprInfo = field(default_factory=GprInfo) b: GprInfo = field(default_factory=GprInfo) ################################################################################ # Assembly Kernel ################################################################################ class KernelWriterAssembly(KernelWriter): ############################################################################## # Init ############################################################################## def __init__( self, assembler: Assembler, debugConfig: DebugConfig, ): super(KernelWriterAssembly, self).__init__(assembler, debugConfig) self.globalread_gpr_record = GlobalReadGprRecord() def _getCustomKernelSource(self, kernel, CustomKernelDirectory): kernelName = getKernelFileBase(self.debugConfig.splitGSU, kernel) with open(os.path.join(CustomKernelDirectory, (kernelName + ".s"))) as f: rocmVersion = self.assembler.rocm_version if not (rocmVersion.major >= 6 and rocmVersion.patch >= 32650): code = [] for line in f.readlines(): if "amdhsa_user_sgpr_kernarg_preload" not in line: code.append(line) code = "".join(code) else: code = f.read() self.tPA = {} self.tPB = {} self.states.kernel = kernel self.states.language = "ASM" self.states.version = kernel["ISA"] return code def getSourceFileString(self, kernel) -> Tuple[int, str]: assert kernel["KernelLanguage"] == "Assembly" # Skip if .o files will have already been built for this file if kernel.duplicate: self.language = "ASM" return (0, "") # should this be an non zero number isCustom = isCustomKernelConfig(kernel) try: if isCustom: code = self._getCustomKernelSource(kernel, CUSTOM_KERNEL_PATH) # The occupancy formula in compute_occupancy_from_resources is # gfx9/wave64-specific; only override CUOccupancy for gfx9 custom # kernels. Non-gfx9 (gfx10/11/12, wave32) custom kernels leave # CUOccupancy at its default (-1) rather than emit a wrong value (Fix #2). _isa = kernel.get("ISA", [9, 0, 8]) if _isa[0] == 9: # Use live rocisa caps so custom-kernel occupancy matches the codegen path. # For custom kernels, setRocIsa() is called (not setKernel()), so # self.states.regCaps/archCaps are not initialized — query the singleton directly. _ti = rocIsa.getInstance() _regCaps = _ti.getRegCaps() _archCaps = _ti.getArchCaps() # Some codegen paths (e.g. grouped_gemm_ck custom kernels) reach here with # a rocisa singleton whose getRegCaps()/getArchCaps() return an incomplete or # empty dict (setKernel was never called for this ISA in that context). # Accessing missing keys would raise KeyError and break codegen. Guard: # try the live singleton caps first; on any KeyError/AttributeError fall back # to the static hardware table. If neither source has the caps, skip the # override and leave CUOccupancy at its default (-1). try: _rocisa_caps = ( _regCaps["MaxVgpr"] * (2 if _archCaps.get("ArchAccUnifiedRegs") else 1), _regCaps["PhysicalMaxSgpr"], _archCaps["DeviceLDS"], _archCaps["MaxWavesPerSimd"], ) except (KeyError, AttributeError): # Singleton caps incomplete: fall back to static table (or None → skip). _rocisa_caps = _arch_caps_for_kernel(kernel) occ = compute_occupancy_from_asm_source(kernel, code, arch_caps=_rocisa_caps) if occ is not None: kernel["CUOccupancy"] = occ else: code = self._getKernelSource(kernel) errcode = 0 except RuntimeError as e: printWarning(f"Failed to generate assembly source code for {kernel}: {e}") code = "" errcode = -2 return (errcode, code) def getSgprOccupancy(self, sgprs): return self.states.regCaps["PhysicalMaxSgpr"]//sgprs def getVgprOccupancy(self, numThreads, vgprs, doubleVgpr=False): multiplier = int(ceil(max(numThreads, 256) / 256.0)) # example: wg=512 multiplier=2, 1024=4 maxOccupancy = self.states.archCaps["MaxWavesPerSimd"]//multiplier vgprAllocateAligned = 4 if not doubleVgpr else 8 totalVgprs = self.states.regCaps["MaxVgpr"] if not doubleVgpr else self.states.regCaps["MaxVgpr"]*2 vgprsAligned = int(ceil(vgprs/vgprAllocateAligned))*vgprAllocateAligned vgprsAligned *= multiplier if vgprsAligned > totalVgprs: return 0 elif vgprsAligned < 1: return maxOccupancy occupancy = min(totalVgprs//vgprsAligned, maxOccupancy) #print("vgprs = ", vgprs, "vgprsAligned = ", vgprsAligned, "unifiedVgprRegs = " ,unifiedVgprRegs, "Occupancy = ", occupancy) return occupancy ######################################## def getOccupancy(self, numThreads, vgprs, sgprs, ldsSize, accvgprs=0, doubleVgpr=False): deviceLdsSize = self.states.archCaps["DeviceLDS"] ldsLimitedOccupancy = self.getLdsLimitedOccupancy(deviceLdsSize, ldsSize) if not doubleVgpr: vgprLimitedOccupancy = self.getVgprOccupancy(numThreads, vgprs, doubleVgpr) accvgprLimitedOccupancy = self.getVgprOccupancy(numThreads, accvgprs, doubleVgpr) else: vgprLimitedOccupancy = self.getVgprOccupancy(numThreads, ceil(vgprs/8)*8+accvgprs, doubleVgpr) accvgprLimitedOccupancy = vgprLimitedOccupancy sgprLimitedOccupancy = self.getSgprOccupancy(sgprs) return min(ldsLimitedOccupancy, vgprLimitedOccupancy, accvgprLimitedOccupancy, sgprLimitedOccupancy) # TODO: also consider sgpr def getMaxRegsForOccupancy(self, numThreads, vgprs, sgprs, ldsSize, accvgprs=0, doubleVgpr=False): lastVgprs = vgprs considerAccVgprs = 0 if not doubleVgpr else accvgprs totalVgprs = self.states.regCaps["MaxVgpr"] if not doubleVgpr else self.states.regCaps["MaxVgpr"]*2 initOccupancy = self.getOccupancy(numThreads, vgprs, sgprs, ldsSize, accvgprs, doubleVgpr) if initOccupancy == 0: return lastVgprs, 1 def getVgpr(vgpr, doubleVgpr): return vgpr if not doubleVgpr else ceil(vgpr/8)*8 while (getVgpr(vgprs, doubleVgpr) + considerAccVgprs) < totalVgprs and vgprs < self.states.regCaps["MaxVgpr"]: vgprs += 1 if self.getVgprOccupancy(numThreads, getVgpr(vgprs, doubleVgpr) + considerAccVgprs, doubleVgpr) >= initOccupancy: lastVgprs = vgprs next else: break return lastVgprs, initOccupancy @staticmethod def getLdsLimitedOccupancy(deviceLdsSize, ldsSize): if ldsSize == 0: # No LDS usage: LDS is not the binding constraint. # Return a large sentinel so other limits (VGPR, wave cap) win in min(). return deviceLdsSize // 256 # As ldsSize gets large, rounding might push us slightly higher than deviceLdsSize. # Clamp at deviceLdsSize ldsSize = min(ldsSize + 255, deviceLdsSize) & 0xffffff00 # 256-byte granularity ldsLimitedOccupancy = deviceLdsSize//ldsSize return ldsLimitedOccupancy @staticmethod def getLdsSize(kernel): ldsSize = kernel["LdsNumBytes"] return ldsSize ######################################## def sizeRef(self, idx): """ Return sgpr() or const with the specified size See above definitions for how these are mapped to Free or Sum sizes based on the problem definition. """ idxChar= INDEX_CHARS[idx] return sgpr("Size%s"%idxChar) def loopChar(self, kernel, loopIdx): loopDim = kernel["ProblemType"]["IndicesSummation"][loopIdx] return INDEX_CHARS[loopDim] def loopSizeRef(self, kernel, loopIdx): loopDim = kernel["ProblemType"]["IndicesSummation"][loopIdx] return self.sizeRef(loopDim) def loopCounterName(self, kernel, loopIdx): return "LoopCounter%s"%(self.loopChar(kernel, loopIdx)) def loopCounter(self, kernel, loopIdx): """ Return loopCounter for loopIdx wrapped in "SGPR" syntax loop idx is 0...unrollIdx """ return sgpr(self.loopCounterName(kernel,loopIdx)) def checkLastIter(self, kernel, comment="at last iteration?"): """ Return last iteration of unroll loop. """ return SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), src1=0, comment=comment) def isConstUnitStride(self, stride): if isinstance(stride, RegisterContainer): return False return stride.startswith("const") def useConstSgprGlobalReadIncsForTc(self, tc: str) -> bool: if tc == "A": return self.states.a.useConstSgprGlobalReadIncs if tc == "B": return self.states.b.useConstSgprGlobalReadIncs if tc == "MXSA": return self.states.mxsa.useConstSgprGlobalReadIncs if tc == "MXSB": return self.states.mxsb.useConstSgprGlobalReadIncs return False def isTdmWaveSeparated(self, kernel) -> bool: return kernel["enableTDMA"] and kernel["enableTDMB"] and kernel["NumWaves"] > 1 ######################################## def strideRef(self, tc, dim): """ Return sgpr with specified stride or define starting with const if constant. dim is index 0...max indices and is in global index space. """ problemType = self.states.kernel["ProblemType"] if tc in ['A','B', "MXSA", "MXSB", "Metadata"]: if not problemType["UseInitialStridesAB"] and \ dim == problemType["IndexAssignments%s"%tc][0]: return ("constStride%s%s"%(tc,self.states.indexChars[dim])) else: return sgpr("Stride%s%s"%(tc,self.states.indexChars[dim])) elif tc in ['E','D','C','TD']: if not problemType["UseInitialStridesCD"] and dim == 0: return ("constStride%s%s"%(tc,self.states.indexChars[dim])) else: return sgpr("Stride%s%s"%(tc,self.states.indexChars[dim])) else: raise ValueError("unexpected tensorChar='%s' in stride function"%tc) ############################################################################## # Find Memory Instruction For Width and Stride ############################################################################## def findMemoryInstructionForWidthStride(self, width, strides, combine, \ instructions): for i in range(0, len(instructions)): instruction = instructions[i] numAddresses = instruction.numAddresses numOffsets = instruction.numOffsets offsetMultiplier = instruction.offsetMultiplier blockWidth = instruction.blockWidth valid = True if width < blockWidth: valid = False if ((width / blockWidth) != floor(width / blockWidth)): valid = False if combine: # try to combine ops if numOffsets > 0: # if inst combines using offsets for stride in strides: if stride % offsetMultiplier != 0: valid = False else: # don't try to combine ops if numOffsets > 1 or numAddresses > 1: valid = False if valid: return i else: continue printWarning("Could not find valid memory instruction for width=%f" % width) return len(instructions) ############################################################################## # Select Memory Instruction # when selecting instruction, need to support stride in both dims ############################################################################## def selectMemoryInstruction(self, operation, # ReadGlobal, WriteLocal, ReadLocal width, # num registers 1 chunk write2, # Para, Perp, None para2, # NumLoadsPara >= 2 perp2, # NumLoadsPerp >= 2 strides ): #instructions = self.memoryArchitecture[operation] instructions = self.memoryInstructions[operation] # try to combine if (write2 == "Coalesced" and para2) \ or (write2 == "Perpendicular" and perp2): instructionIdx = self.findMemoryInstructionForWidthStride( \ width, strides, True, instructions) # don't or can't combine else: instructionIdx = self.findMemoryInstructionForWidthStride( \ width, strides, False, instructions) if instructionIdx < len(instructions): # found return instructionIdx else: raise RuntimeError("Could not find valid memory instruction for operation=%s, width=%f, kernel=%s" %(operation, width, self.states.kernelName)) def selectTransposedDSReadInstuctionIdx(self, numReturnedRegisters: int, bpe: int): insts = self.memoryInstructions["TrLocalRead"] return next(i for i, inst in enumerate(insts) if inst.bpe == bpe and inst.blockWidth == numReturnedRegisters) def initGlobalReadMemoryInstruction(self, instructions, tP, bpr): # globalRead instruction; no flat_load2_* globalReadWidth = float(tP["nrcv"]*tP["bpeGR"])/bpr globalRead2Coalesced = tP["nrc"] > 1 globalRead2Perpendicular = tP["nrp"] > 1 globalReadInstructionIdx = self.selectMemoryInstruction("GlobalRead", globalReadWidth, \ False, \ globalRead2Coalesced, globalRead2Perpendicular, [] ) tP["globalReadInstruction"] = instructions["GlobalRead"][globalReadInstructionIdx] def initLocalWriteMemoryInstruction(self, instructions, kernel, tP, bpr): ######################################## # localWrite instruction # for local, tile->para, unroll->perp # wtc = writeTileDimComponents localWriteWidth = tP["nwcv"]*tP["bpeDS"]/bpr localWrite2Coalesced = tP["nrc"]>1 or tP["wtc"] localWrite2Perpendicular = tP["nrp"]>1 # localWrite stride tile if tP["tlu"]: if tP["wtc"]: localWriteStrideTile = 1 else: localWriteStrideTile = kernel[tP["lsc"]] else: localWriteStrideTile = kernel[tP["lsp"]] localWriteStrideTile = int(localWriteStrideTile*tP["bpeDS"])//bpr # localWrite stride unroll if tP["tlu"]: localWriteStrideUnroll = kernel[tP["lsc"]]*kernel[tP["mt"]] else: if tP["wtc"]: localWriteStrideUnroll = 1*kernel[tP["mt"]] else: localWriteStrideUnroll = kernel[tP["lsc"]]*kernel[tP["mt"]] localWriteStrideUnroll = \ int(localWriteStrideUnroll*tP["bpeDS"])//bpr localWriteInstructionIdx = self.selectMemoryInstruction("LocalWrite", localWriteWidth, \ False, \ localWrite2Coalesced, localWrite2Perpendicular, [localWriteStrideTile, localWriteStrideUnroll] ) tP["localWrite2Coalesced"] = localWrite2Coalesced tP["localWrite2Perpendicular"] = localWrite2Perpendicular tP["localWriteStrideTile"] = localWriteStrideTile tP["localWriteStrideUnroll"] = localWriteStrideUnroll tP["localWriteInstruction"] = instructions["LocalWrite"][localWriteInstructionIdx] def initLocalReadMemoryInstruction(self, instructions, kernel, tP, bpr): tChar = tP["tensorChar"] if kernel["UnrollMajorLDS%s"%tChar]: if tChar == "A": localReadWidth = (self.states.lrvwUnrollA * tP["bpeDS"]) / bpr elif tChar == "B": localReadWidth = (self.states.lrvwUnrollB * tP["bpeDS"]) / bpr elif tChar == "MXSA": if self.states.asmCaps["HasWMMA_V3"] and kernel["MXScaleFormat"] == "InMemorySwizzle": localReadWidth = (self.states.lrvwUnrollMXSA * kernel["VectorWidthA"] * tP["bpeDS"]) / bpr else: localReadWidth = (self.states.lrvwUnrollMXSA * tP["bpeDS"]) / bpr elif tChar == "MXSB": if self.states.asmCaps["HasWMMA_V3"] and kernel["MXScaleFormat"] == "InMemorySwizzle": localReadWidth = (self.states.lrvwUnrollMXSB * kernel["VectorWidthB"] * tP["bpeDS"]) / bpr else: localReadWidth = (self.states.lrvwUnrollMXSB * tP["bpeDS"]) / bpr elif tChar == "Metadata": localReadWidth = (self.states.lrvwUnrollMetadata * tP["bpeDS"]) / bpr else: raise Exception(f"unsupport tc %s{tChar}") else: if tChar == "A": localReadWidth = (self.states.lrvwTileA * tP["bpeDS"]) / bpr elif tChar == "B": localReadWidth = (self.states.lrvwTileB * tP["bpeDS"]) / bpr elif tChar == "MXSA": localReadWidth = (self.states.lrvwTileMXSA * tP["bpeDS"]) / bpr elif tChar == "MXSB": localReadWidth = (self.states.lrvwTileMXSB * tP["bpeDS"]) / bpr elif tChar == "Metadata": localReadWidth = (self.states.lrvwTileMetadata * tP["bpeDS"]) / bpr else: raise Exception(f"unsupport tc %s{tChar}") #localReadStridePerpendicular = 0 localRead2Perpendicular = False localReadStrideCoalesced = int(kernel[tP["tt"]] * tP["bpeDS"]) // bpr localRead2Coalesced = False tP["enableLDSTr"] = kernel["enableLDSTr%s"%tChar] lrInstPoolName = "LocalRead" if tP["enableLDSTr"]: lrInstPoolName = "TrLocalRead" maxTrLoadNumReturnedVgpr = 4 if self.states.asmCaps["HasLDSTrB128B16"] else 2 if tP["bpeDS"] in (0.5, 1): maxTrLoadNumReturnedVgpr = 2 elif tP["bpeDS"] == 0.75: maxTrLoadNumReturnedVgpr = 3 elif tP["bpeDS"] != 2: assert False, f"Unhandled bpeDS: {tP['bpeDS']}" loadNumReturnVgrp = min(int(localReadWidth*kernel["MIInputPerThread"]), maxTrLoadNumReturnedVgpr) if tChar != "Metadata" else 2 localReadInstructionIdx = self.selectTransposedDSReadInstuctionIdx(loadNumReturnVgrp, tP["bpeDS"]) else: localReadInstructionIdx = self.selectMemoryInstruction("LocalRead", localReadWidth, \ False, \ localRead2Coalesced, localRead2Perpendicular, [localReadStrideCoalesced] ) tP["localRead2Coalesced"] = localRead2Coalesced tP["localRead2Perpendicular"] = localRead2Perpendicular tP["localReadStrideCoalesced"] = localReadStrideCoalesced tP["localReadInstruction"] = instructions[lrInstPoolName][localReadInstructionIdx] def allocTmpSgpr(self, num: int, alignment=None, tag=None): def overflowListener(e): self.states.overflowedResources = 2 if self.db["AssertOnSgprOverflow"]: raise e tmpSgpr = allocTmpGpr(self.sgprPool, num, self.states.regCaps["MaxSgpr"], alignment, tag, overflowListener) return tmpSgpr def allocTmpSgprList(self, nums: List[int], alignments: Optional[List[int]]=None, tag=None): def overflowListener(e): self.states.overflowedResources = 2 if self.db["AssertOnSgprOverflow"]: raise e tmpSgpr = allocTmpGprList(self.sgprPool, nums, self.states.regCaps["MaxSgpr"], alignments, tag, overflowListener) return tmpSgpr def defineMultiVgprIndex(self, names: List[str], numVgprs: List[int], align=1): assert(len(names) == len(numVgprs)) vgprIdxVec = self.vgprPool.checkOutMulti(numVgprs, align, tags=names) return vgprIdxVec def defineSgprIdx(self, name, numSgprs, align=1): if numSgprs == 0: return sgprIdx = self.sgprPool.checkOutAligned(numSgprs, align, tag=name, preventOverflow=False) #self.sgprIdx = roundUpToNearestMultiple(self.sgprIdx,align) #print (name, "->", self.sgprIdx, "+", numSgprs) self.sgprs[name] = sgprIdx return sgprIdx def defineSgpr(self, name, numSgprs, align=1): if numSgprs == 0: return Module() # check if previous define sgprs are being used.. for s in self.states.freeSgprVarPool: self.setSgprToInUseState(s) ret = RegSet("s", "sgpr"+name, self.defineSgprIdx(name, numSgprs, align)) for s in self.states.freeSgprVarPool: self.setSgprToFreeState(s) return ret def defineMultiSgprIndex(self, names: List[str], numSgprs: List[int], align=1): assert(len(names) == len(numSgprs)) sgprIdxVec = self.sgprPool.checkOutMulti(numSgprs, align, tags=names) #self.sgprIdx = roundUpToNearestMultiple(self.sgprIdx,align) #print (name, "->", self.sgprIdx, "+", numSgprs) for idx, name in enumerate(names): self.sgprs[name] = sgprIdxVec[idx] return sgprIdxVec def setSgprToInUseState(self, name): self.sgprPool.removeFromCheckOut(self.sgprs[name]) return RegSet("s", "sgpr"+name, self.sgprs[name]) def undefineSgpr(self, name): if name not in self.sgprs: # TODO: Remove guard after full TileInfo migration return ValueSet(name="sgpr"+name, value="UNDEF", format = -1) self.sgprPool.checkIn(self.sgprs[name]) # undefine a sgpr string twice will cause compiler error. # User must not add the UNDEF code module except it is the last one. return ValueSet(name="sgpr"+name, value="UNDEF", format = -1) def setSgprToFreeState(self, name): self.sgprPool.addFromCheckOut(self.sgprs[name]) # undefine a sgpr string twice will cause compiler error. # Must call setSgprToInUseState again before calling setSgprToFreeState. return ValueSet(name="sgpr"+name, value="UNDEF", format = -1) def addSgprVarToPool(self, name): if name not in self.sgprs.keys() or (name in self.states.freeSgprVarPool): return self.states.freeSgprVarPool.add(name) ret = self.setSgprToFreeState(name) return ret def removeSgprVarFromPool(self, name): if name not in self.sgprs.keys() or (name not in self.states.freeSgprVarPool): return self.states.freeSgprVarPool.remove(name) ret = self.setSgprToInUseState(name) return ret def removeGRSrdVariableSgprsFromPool(self, kernel): module = Module("RemoveGRSrdSgprsFromPool") if kernel["BufferLoad"]: # resource descriptor (SRD) A and B, must be aligned on 4-SGPR boundary self.removeSgprVarFromPool("SrdA") self.removeSgprVarFromPool("SrdB") if kernel["ProblemType"]["MXBlockA"]: self.removeSgprVarFromPool("SrdMXSA") if kernel["ProblemType"]["MXBlockB"]: self.removeSgprVarFromPool("SrdMXSB") if self.states.use64bShadowLimit: self.removeSgprVarFromPool("ShadowLimitA") self.removeSgprVarFromPool("ShadowLimitB") if self.states.use64bShadowLimitMX: if kernel["ProblemType"]["MXBlockA"]: self.removeSgprVarFromPool("ShadowLimitMXSA") if kernel["ProblemType"]["MXBlockB"]: self.removeSgprVarFromPool("ShadowLimitMXSB") self.removeSgprVarFromPool("WrapUA") self.removeSgprVarFromPool("WrapUB") if kernel["ProblemType"]["MXBlockA"]: self.removeSgprVarFromPool("WrapUMXSA") if kernel["ProblemType"]["MXBlockB"]: self.removeSgprVarFromPool("WrapUMXSB") return module def undefineSubtileMainLoopSgprs(self, kernel): """Undefine SGPRs used only during the main loop that are not needed in the post-loop. Called for subtile kernels after deallocOffsetRegisters and before endSummation/post-loop.""" module = Module("UndefineSubtileMainLoopSgprs") sgprsToUndefine = [ "LocalWriteBaseAddrA", "LocalWriteBaseAddrB", "LocalWriteBaseAddrMXSA", "LocalWriteBaseAddrMXSB", "SwapA", "SwapB", "SwapMXSA", "SwapMXSB", ] for name in sgprsToUndefine: if name in self.sgprs: module.add(self.undefineSgpr(name)) return module def removeGROffsetsVariableSgprsFromPool(self, kernel): module = Module("RemoveGROffsetSgprsFromPool") # Wave-separated TDM copies A/B (and MX) incs into tdm*Incs during setup; main loop # uses tdm*Incs only. GlobalReadIncs* are released afterward — do not pin them here, # UNLESS stagger code needs them (calculateStagger references GlobalReadIncs for all tensors). needsStaggerSgprs = kernel["StaggerU"] > 0 or kernel["InternalSupportParams"]["SupportCustomStaggerU"] if not self.isTdmWaveSeparated(kernel) or needsStaggerSgprs: self.removeSgprVarFromPool("GlobalReadIncsA") self.removeSgprVarFromPool("GlobalReadIncsB") if kernel["ProblemType"]["MXBlockA"]: self.removeSgprVarFromPool("GlobalReadIncsMXSA") if kernel["ProblemType"]["MXBlockB"]: self.removeSgprVarFromPool("GlobalReadIncsMXSB") return module def releaseGlobalReadIncsSgprsAfterTdmWaveSep(self, kernel): """Return GlobalReadIncs* SGPRs to the pool for wave-separated TDM. Call immediately after tdmSetupIncrementWaveSeparated when staggerUCode is false. When staggerUCode is true, defer until after removeStaggerAB (stagger still uses GlobalReadIncs*; see calculateStagger). """ module = Module("ReleaseGlobalReadIncsAfterTdmWaveSep") if not self.isTdmWaveSeparated(kernel): return module if self.states.a.numSgprGlobalReadIncs > 0: module.add(self.addSgprVarToPool("GlobalReadIncsA")) if self.states.b.numSgprGlobalReadIncs > 0: module.add(self.addSgprVarToPool("GlobalReadIncsB")) if kernel["ProblemType"]["MXBlockA"] and self.states.mxsa.numSgprGlobalReadIncs > 0: module.add(self.addSgprVarToPool("GlobalReadIncsMXSA")) if kernel["ProblemType"]["MXBlockB"] and self.states.mxsb.numSgprGlobalReadIncs > 0: module.add(self.addSgprVarToPool("GlobalReadIncsMXSB")) return module def defineVariableSgprs(self, kernel): module = Module("DefineVariableSgpr") #------------------------ # Registers defined below this point are not available in the post-loop # Post-loop is after tail loop exits, ie the store code. # (we reclaim them to use as temps, typically for execmasks) # Mostly impacts flat kernels and GSU edge since these need SGPR # for conditionals # For subtile kernels, TDM descriptor allocation is deferred to after # StreamK preLoop so freed SK-constant slots can serve as temps first. if not kernel.get("UseSubtileImpl"): module.addModuleAsFlatItems(self.defineTdmSgprs(kernel)) if kernel["BufferLoad"]: # resource descriptor (SRD) A and B, must be aligned on 4-SGPR boundary if not kernel["enableTDMA"]: module.add(self.defineSgpr("SrdA", 4, 4)) self.addSgprVarToPool("SrdA") if not kernel["enableTDMB"]: module.add(self.defineSgpr("SrdB", 4, 4)) self.addSgprVarToPool("SrdB") if kernel["ProblemType"]["MXBlockA"] and not kernel["enableTDMA"]: module.add(self.defineSgpr("SrdMXSA", 4, 4)) self.addSgprVarToPool("SrdMXSA") if kernel["ProblemType"]["MXBlockB"] and not kernel["enableTDMB"]: module.add(self.defineSgpr("SrdMXSB", 4, 4)) self.addSgprVarToPool("SrdMXSB") if not kernel["enableTDMMetadata"] and kernel["ProblemType"]["Sparse"]: module.add(self.defineSgpr("SrdMetadata", 4, 4)) if self.states.use64bShadowLimit: if not kernel["enableTDMA"]: module.add(self.defineSgpr("ShadowLimitA", 2, 2)) self.addSgprVarToPool("ShadowLimitA") if not kernel["enableTDMB"]: module.add(self.defineSgpr("ShadowLimitB", 2, 2)) self.addSgprVarToPool("ShadowLimitB") if not kernel["enableTDMMetadata"] and kernel["ProblemType"]["Sparse"]: module.add(self.defineSgpr("ShadowLimitMetadata", 2, 2)) if self.states.use64bShadowLimitMX: if kernel["ProblemType"]["MXBlockA"] and not kernel["enableTDMA"]: module.add(self.defineSgpr("ShadowLimitMXSA", 2, 2)) self.addSgprVarToPool("ShadowLimitMXSA") if kernel["ProblemType"]["MXBlockB"] and not kernel["enableTDMB"]: module.add(self.defineSgpr("ShadowLimitMXSB", 2, 2)) self.addSgprVarToPool("ShadowLimitMXSB") if self.states.staggerUCode: module.add(self.defineSgpr("StaggerUIter", 1)) # stagger loop iterations, used for various iter counts in the code isDTVAorB = (kernel["DirectToVgprA"] != kernel["DirectToVgprB"]) # only one of them is enabled if kernel["PrefetchGlobalRead"] >= 2 and isDTVAorB: # PGR2 + DTVA or B (only 1 side), need separate StaggerUIter for DTV load module.add(self.defineSgpr("StaggerUIterDTV", 1)) # stagger loop iterations, used for various iter counts in the code wrapAlignment = 2 if self.states.asmCaps["s_sub_u64"] and self.states.asmCaps["HasWMMA_V3"] else 1 module.add(self.defineSgpr("WrapUA", 2, wrapAlignment)) # Bytes to add to SrdA to reset address from N-1 iter to AddressA module.add(self.defineSgpr("WrapUB", 2, wrapAlignment)) # Bytes to add to SrdB to reset address from N-1 iter to AddressB if kernel["ProblemType"]["MXBlockA"]: module.add(self.defineSgpr("WrapUMXSA", 2, wrapAlignment)) # Bytes to add to SrdA to reset address from N-1 iter to AddressMXSA if kernel["ProblemType"]["MXBlockB"]: module.add(self.defineSgpr("WrapUMXSB", 2, wrapAlignment)) # Bytes to add to SrdA to reset address from N-1 iter to AddressMXSB if kernel["ProblemType"]["Sparse"]: module.add(self.defineSgpr("WrapUMetadata", 2, wrapAlignment)) # Bytes to add to SrdMetadata to reset address from N-1 iter to AddressMetadata self.addSgprVarToPool("WrapUA") self.addSgprVarToPool("WrapUB") if self.states.a.numSgprGlobalReadIncs > 0: module.add(self.defineSgpr("GlobalReadIncsA", self.states.a.numSgprGlobalReadIncs)) if kernel["NumWaves"] < 2: self.addSgprVarToPool("GlobalReadIncsA") if kernel["ProblemType"]["MXBlockA"] and self.states.mxsa.numSgprGlobalReadIncs > 0: module.add(self.defineSgpr("GlobalReadIncsMXSA", self.states.mxsa.numSgprGlobalReadIncs)) if kernel["NumWaves"] < 2: self.addSgprVarToPool("GlobalReadIncsMXSA") if self.states.b.numSgprGlobalReadIncs > 0: module.add(self.defineSgpr("GlobalReadIncsB", self.states.b.numSgprGlobalReadIncs)) if kernel["NumWaves"] < 2: self.addSgprVarToPool("GlobalReadIncsB") if kernel["ProblemType"]["MXBlockB"] and self.states.mxsb.numSgprGlobalReadIncs > 0: module.add(self.defineSgpr("GlobalReadIncsMXSB", self.states.mxsb.numSgprGlobalReadIncs)) if kernel["NumWaves"] < 2: self.addSgprVarToPool("GlobalReadIncsMXSB") if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.add(self.defineSgpr("GlobalReadIncsMetadata", self.states.m.numSgprGlobalReadIncs)) if self.states.IncLdsBufSwitch: module.add(self.defineSgpr("LDSBufferReadInc", 1)) module.add(self.defineSgpr("LDSBufferWriteInc", 1)) needPackK16 = False needPackK8Lw = False needPackK8Hi = False if kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16(): if self.states.lrvwTileA > 1 or self.states.lrvwTileB > 1: needPackK16 = True if self.states.lrvwTileMetadata > 1: needPackK8Lw = True if self.states.lrvwTileMetadata > 2: needPackK8Hi = True elif kernel["ProblemType"]["DataType"].isInt8() or kernel["ProblemType"]["DataType"].is8bitFloat(): if self.states.lrvwTileA > 1 or self.states.lrvwTileB > 1 or self.states.lrvwTileMetadata > 1: needPackK8Lw = True if self.states.lrvwTileA > 2 or self.states.lrvwTileB > 2 or self.states.lrvwTileMetadata > 2: needPackK8Hi = True # need extra sgpr for packing metadata tPackM = "M" if needPackK16 and needPackK8Lw else "" if needPackK16: module.add(self.defineSgpr("PackKForV0", 1)) module.add(self.defineSgpr("PackKForV1", 1)) if kernel["StreamK"] != 0: self.states.nonPostLoopSgpr.append("PackKForV0") self.states.nonPostLoopSgpr.append("PackKForV1") if needPackK8Lw: module.add(self.defineSgpr("PackKFor%sV0"%tPackM, 1)) module.add(self.defineSgpr("PackKFor%sV1"%tPackM, 1)) if kernel["StreamK"] != 0 and not needPackK16 : self.states.nonPostLoopSgpr.append("PackKForV0") self.states.nonPostLoopSgpr.append("PackKForV1") if needPackK8Hi: module.add(self.defineSgpr("PackKFor%sV2"%tPackM, 1)) module.add(self.defineSgpr("PackKFor%sV3"%tPackM, 1)) if kernel["StreamK"] != 0 and not needPackK16: self.states.nonPostLoopSgpr.append("PackKForV2") self.states.nonPostLoopSgpr.append("PackKForV3") if kernel["ProblemType"]["StochasticRounding"] and not self.states.asmCaps["v_prng_b32"]: module.add(self.defineSgpr("RNDSeed", 1)) if kernel["_UseSgprForGRO"]: needFirstSgprOffset = kernel["DirectToLdsA"] and kernel["UseInstOffsetForGRO"] numberOfSgpr = self.states.a.numVgprGlobalReadOffsets if needFirstSgprOffset else (self.states.a.numVgprGlobalReadOffsets-1) if numberOfSgpr > 0 and not kernel["enableTDMA"]: module.add(self.defineSgpr("ScalarGlobalReadOffsetA", numberOfSgpr)) needFirstSgprOffset = kernel["DirectToLdsB"] and kernel["UseInstOffsetForGRO"] numberOfSgpr = self.states.b.numVgprGlobalReadOffsets if needFirstSgprOffset else (self.states.b.numVgprGlobalReadOffsets-1) if numberOfSgpr > 0 and not kernel["enableTDMB"]: module.add(self.defineSgpr("ScalarGlobalReadOffsetB", numberOfSgpr)) if kernel["ProblemType"]["MXBlockA"]: numberOfSgpr = self.states.mxsa.numVgprGlobalReadOffsets if needFirstSgprOffset else (self.states.mxsa.numVgprGlobalReadOffsets-1) if numberOfSgpr > 0: module.add(self.defineSgpr("ScalarGlobalReadOffsetMXSA", numberOfSgpr)) if kernel["ProblemType"]["MXBlockB"]: numberOfSgpr = self.states.mxsb.numVgprGlobalReadOffsets if needFirstSgprOffset else (self.states.mxsb.numVgprGlobalReadOffsets-1) if numberOfSgpr > 0: module.add(self.defineSgpr("ScalarGlobalReadOffsetMXSB", numberOfSgpr)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: needFirstSgprOffset = kernel["DirectToLdsMetadata"] and kernel["UseInstOffsetForGRO"] numberOfSgpr = self.states.m.numVgprGlobalReadOffsets if needFirstSgprOffset else (self.states.m.numVgprGlobalReadOffsets-1) if numberOfSgpr > 0 and not kernel["enableTDMMetadata"]: module.add(self.defineSgpr("ScalarGlobalReadOffsetMetadata", numberOfSgpr)) # debug flag to allocate dummy / unused sgpr # useful when comparing code that adds new kernel arguments to see what # was actually changed numDummySgpr= 0 for i in range(numDummySgpr): module.add(self.defineSgpr("DummySgpr%d"%i, 1)) if kernel["PrefetchGL2"]: module.add(self.defineSgpr("GL2PrefetchIncA", self.states.rpgo)) module.add(self.defineSgpr("GL2PrefetchIncB", self.states.rpgo)) if kernel["ProblemType"]["MXBlockA"]: module.add(self.defineSgpr("GL2PrefetchIncMXSA", self.states.rpgo)) if kernel["ProblemType"]["MXBlockB"]: module.add(self.defineSgpr("GL2PrefetchIncMXSB", self.states.rpgo)) if kernel["enableTDMMetadata"]: module.add(self.defineSgpr("GL2PrefetchIncMetadata", self.states.rpgo)) if self.sgprPool.size() > self.states.regCaps["MaxSgpr"]: print ("warning: Number of defined SGPRS (%d) overflowed max SGPRS (%d)." \ % (self.sgprPool.size(), self.states.regCaps["MaxSgpr"])) # End of define sgprs #------------------------ ######################################################### # Below calculates the number of sgprs needed not in epilogue ######################################################### self.states.numStoreSgprNames2 = [] self.states.numStoreSgprNameSizes2 = [] storeSgprLoad2 = 0 if kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel' or kernel["AdaptiveGemmGSUA"] == 1: storeSgprLoad2 += self.states.rpga self.states.numStoreSgprNames2.append("AddressTD") self.states.numStoreSgprNameSizes2.append(self.states.rpga) if kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel' or kernel["AdaptiveGemmGSUA"] == 1: storeSgprLoad2 += self.states.rpga self.states.numStoreSgprNames2.append("Synchronizer") self.states.numStoreSgprNameSizes2.append(self.states.rpga) self.states.numStoreSgprToLoad2 = storeSgprLoad2 return module ############################################################################## def defineTdmSgprs(self, kernel): """Allocate TDM descriptor SGPRs. Extracted so subtile can defer this.""" module = Module("DefineTdmSgprs") if kernel["enableTDMA"]: module.add(self.defineSgpr("tdmAGroup0", 4, 4)) module.add(self.defineSgpr("tdmAGroup1", 8, 4)) # WS mode aliases B's descriptor onto A's SGPRs, so A's Group2 # must exist whenever either tile is iterate-mode. if (kernel.get("_TDMIterateModeA", False) or (kernel.get("_TDMIterateModeB", False) and kernel["NumWaves"] > 1)): # Group 2: iter_count / lds_inc / global_inc when iterate_enable=1. module.add(self.defineSgpr("tdmAGroup2", 4, 4)) # Group 3 aliased to Group 2 (4-op instruction form requires Group 3 to # be a valid SGPR name; null is rejected by the assembler). module.add(RegSet("s", "sgprtdmAGroup3", "sgprtdmAGroup2")) if kernel["ProblemType"]["MXBlockA"]: module.add(self.defineSgpr("tdmMXSAGroup0", 4, 4)) module.add(self.defineSgpr("tdmMXSAGroup1", 8, 4)) if kernel["enableTDMB"]: # Alias B descriptor onto A for multi-wave to reduce SGPR pressure. # Subtile uses separate descriptors -- deferred allocation provides # enough SGPR headroom, and separate descriptors avoid the reinit # overhead before each tensor_load_to_lds. if kernel["NumWaves"] > 1 and not kernel.get("UseSubtileImpl"): module.add(RegSet("s", "sgprtdmBGroup0", "sgprtdmAGroup0")) module.add(RegSet("s", "sgprtdmBGroup1", "sgprtdmAGroup1")) if kernel.get("_TDMIterateModeA", False) or kernel.get("_TDMIterateModeB", False): module.add(RegSet("s", "sgprtdmBGroup2", "sgprtdmAGroup2")) module.add(RegSet("s", "sgprtdmBGroup3", "sgprtdmAGroup2")) if kernel["ProblemType"]["MXBlockB"]: module.add(RegSet("s", "sgprtdmMXSBGroup0", "sgprtdmMXSAGroup0")) module.add(RegSet("s", "sgprtdmMXSBGroup1", "sgprtdmMXSAGroup1")) else: module.add(self.defineSgpr("tdmBGroup0", 4, 4)) module.add(self.defineSgpr("tdmBGroup1", 8, 4)) if kernel.get("_TDMIterateModeB", False): module.add(self.defineSgpr("tdmBGroup2", 4, 4)) module.add(RegSet("s", "sgprtdmBGroup3", "sgprtdmBGroup2")) if kernel["ProblemType"]["MXBlockB"]: module.add(self.defineSgpr("tdmMXSBGroup0", 4, 4)) module.add(self.defineSgpr("tdmMXSBGroup1", 8, 4)) if kernel["enableTDMA"] and kernel["enableTDMB"] and kernel["NumWaves"] > 1: module.add(self.defineSgpr("tdmABIncs", 1)) # Multi-wave TDMSplit recomputes the LDS/global split increments (and the # dim1 H0/H1 boundaries) transiently at point of use (see # _tdmSplitMultiWaveInc); nothing is persisted here. if kernel["ProblemType"]["MXBlockA"] and kernel["ProblemType"]["MXBlockB"]: module.add(self.defineSgpr("tdmMXSAMXSBIncs", 1)) # Single-wave (NumWaves == 1): descriptors are independent, so TDMSplit # uses per-tensor increment SGPRs instead of the aliased AB pair. if (kernel["enableTDMA"] and kernel["enableTDMB"] and kernel["NumWaves"] == 1 and kernel["TDMSplit"] and not kernel["ProblemType"]["Sparse"]): module.add(self.defineSgpr("tdmAGlobalSplitIncs", 1)) module.add(self.defineSgpr("tdmALdsSplitIncs", 1)) module.add(self.defineSgpr("tdmBGlobalSplitIncs", 1)) module.add(self.defineSgpr("tdmBLdsSplitIncs", 1)) if kernel["enableTDMMetadata"]: module.add(self.defineSgpr("tdmMetadataGroup0", 4, 4)) module.add(self.defineSgpr("tdmMetadataGroup1", 8, 4)) # TDM LDS address tracking for aliased descriptors with subtile double-buffering if kernel["UseSubtileImpl"] and kernel["enableTDMA"] and kernel["enableTDMB"]: module.add(self.defineSgpr("tdmLdsAddrA", 1)) module.add(self.defineSgpr("tdmLdsAddrB", 1)) module.add(self.defineSgpr("tdmLdsSwapMaskA", 1)) module.add(self.defineSgpr("tdmLdsSwapMaskB", 1)) return module def functionSignature(self) -> SignatureBase: """ Function Signature called after rest of code """ signatureClass = Component.Signature.find(self) signature = signatureClass(self) return signature # This function is to generate vgprValuA/B_T reg definition for X32F emulation # This needs to be called again for TailLoop to force enable tranpose code def macroAndSetF32XEmuTregSingle(self, kernel, sAorB: ABMatrixInfo, tP) -> Module: tc = tP["tensorChar"] moduleVgprMacroValu_T = Module("VALU%s T Vgpr Macro"%tc) # PLR index: from X0 to X (at most) -> VGPRs will be duplicated LoopIters times (at most) # eg, if LoopIters = 4, there would be at most 4*VGPRs PLR = self.states.numVgprBuffer numBi = PLR if sAorB.numVgprEmu > 0 and sAorB.numVgprValu > 0: # Do not generate vgprValuA if numVgprValuA is 0 ri = 0 numBiFactor = numBi for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): if sAorB.TF32EmuInterleaveTreg: # TF32EmuInterleaveTreg case, allocate half of ValuA/B_X. Number of block is same as ValuA/B_X. ri2 = ri//2 else: # allocate same amount as ValuA/B_X. Only 1 block # Need to allocate actual Vreg only for bi == 0. # Reuse same location for T1,... ri2 = 0 moduleVgprMacroValu_T.add(RegSet("v", "vgprValu%s_T%u_I%u"%(tc, bi,iui), sAorB.startVgprCvt, ri2)) ri += sAorB.numVgprValuPerBlock return moduleVgprMacroValu_T def macroAndSetF32XEmuTreg(self, kernel, tPA, tPB) -> Module: self.moduleVgprMacroValuA_T = self.macroAndSetF32XEmuTregSingle(kernel, self.states.a, tPA) self.moduleVgprMacroValuB_T = self.macroAndSetF32XEmuTregSingle(kernel, self.states.b, tPB) # MFMA case if kernel["UseMFMAF32XEmulation"]: self.moduleVgprMacroValuB_T.add(RegSet("v", "IdentityMatrix", self.states.startVgprIdentityMatrix, 0)) def _emitGL2PrefetchAddrRegSets(self, module, kernel, tPA, tPB): """Emit RegSet declarations for GL2 prefetch address vgprs (A, B, and MX scale variants).""" if not kernel["PrefetchGL2"]: return for label, tP, stateObj in [("A", tPA, self.states.a), ("B", tPB, self.states.b)]: for i in range(tP["gl2nl"]): module.add(RegSet("v", f"vgprGL2PrefetchAddr{label}_{i}", stateObj.startVgprGL2PrefetchAddr + i * self.states.rpga)) for mxKey, tP, label, stateObj in [("MXBlockA", tPA, "MXSA", self.states.mxsa), ("MXBlockB", tPB, "MXSB", self.states.mxsb)]: if kernel["ProblemType"][mxKey]: mx = tP["MX"] for i in range(mx["gl2nl"]): module.add(RegSet("v", f"vgprGL2PrefetchAddr{label}_{i}", stateObj.startVgprGL2PrefetchAddr + i * self.states.rpga)) if kernel["enableTDMMetadata"]: tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] for i in range(tPM["gl2nl"]): module.add(RegSet("v", f"vgprGL2PrefetchAddrMetadata_{i}", self.states.m.startVgprGL2PrefetchAddr + i * self.states.rpga)) def macroAndSet(self, kernel, tPA, tPB) -> Module: module = Module("MacroNSet") module.add(MacroVMagicDiv(kernel["MagicDivAlg"])) tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] tPMXSA = tPA["MX"] if kernel["ProblemType"]["MXBlockA"] else None tPMXSB = tPB["MX"] if kernel["ProblemType"]["MXBlockB"] else None PLR = self.states.numVgprBuffer numBi = PLR ######################################## # VGPR Macros ######################################## def macroAndSetImplClassic(): module.addComment2("VGPR Assignments for MX") module.add(RegSet("v", "vgprMXSBase", 0)) moduleVgprMacroMXS = Module("VALU/G2L Vgpr Macro") moduleVgprMacroValuMXSDummy = Module("VALU Dummy Vgpr Macro") moduleVgprMacroValuMXSA = Module("VALUA Vgpr Macro") moduleVgprMacroValuMXSB = Module("VALUB Vgpr Macro") moduleVgprMacroValuMXSAPack = Module("VALUA Pack Vgpr Macro") moduleVgprMacroValuMXSBPack = Module("VALUB Pack Vgpr Macro") moduleVgprMacroG2LMXSA = Module("G2LA Vgpr Macro") moduleVgprMacroG2LMXSB = Module("G2LB Vgpr Macro") if bool(kernel["ProblemType"]["MXBlockA"]) ^ bool(kernel["ProblemType"]["MXBlockB"]): moduleVgprMacroValuMXSDummy.add(RegSet("v", "vgprValuMXSDummy", "vgprMXSBase", 0)) if kernel["ProblemType"]["MXBlockA"]: ri = 0 if self.states.mxsa.numVgprValu > 0: # Do not generate vgprValuMXSA if numVgprValuA is 0 numBiFactor = numBi if kernel["DirectToVgprA"] and (self.states.packDTVA or self.states.convDTVA): # DirectToVgpr case, we need LoopIters * 2 buffers numBiFactor = kernel["LoopIters"] * 2 if self.states.lrvwTileMXSA > 1: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSA_X0_I0_BASE", "vgprMXSBase", self.states.mxsa.startVgprValu)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMXSA.add(RegSet("v", "vgprValuMXSA_X%u_I%u"%(bi,iui), "vgprValuMXSA_X0_I0_BASE", ri)) ri += self.states.mxsa.numVgprValuPerBlock if not kernel["UnrollMajorLDSA"]: ri = 0 ri = 0 if not kernel["UnrollMajorLDSA"]: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSA_X0_I0_D0_PACK", "vgprMXSBase", self.states.mxsa.startVgprValuPack)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): for data in range(0,kernel["MIInputPerThreadMXSA"]): moduleVgprMacroValuMXSAPack.add(RegSet("v", "vgprValuMXSA_X%u_I%u_D%u"%(bi,iui,data),"vgprValuMXSA_X0_I0_D0_PACK", ri)) ri += ceil(kernel["VectorWidthMXSA"] / self.states.bpr) * kernel["MIWaveTileA"] // kernel["VectorWidthMXSA"] else: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSA_X0_I0_BASE", "vgprMXSBase", self.states.mxsa.startVgprValu)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMXSA.add(RegSet("v", "vgprValuMXSA_X%u_I%u"%(bi,iui), "vgprValuMXSA_X0_I0_BASE", ri)) ri += self.states.mxsa.numVgprValuPerBlock ri = 0 if not kernel["UnrollMajorLDSA"]: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSA_X0_I0_D0_PACK", "vgprMXSBase", self.states.mxsa.startVgprValuPack)) for data in range(1,int(self.states.bpr)): for bi in range(0,numBiFactor): # buffer indices if bi % self.states.numVgprBufferPackMXSA == 0: ri = (data-1) * kernel["InnerUnroll"] * self.states.numVgprBufferPackMXSA * self.states.mxsa.numVgprValuPerBlock for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMXSAPack.add(RegSet("v", "vgprValuMXSA_X%u_I%u_D%u"%(bi,iui,data),"vgprValuMXSA_X0_I0_D0_PACK", ri)) ri += self.states.mxsa.numVgprValuPerBlock if kernel["ProblemType"]["MXBlockB"]: ri = 0 if self.states.mxsb.numVgprValu > 0: # Do not generate vgprValuMXSB if numVgprValuB is 0 numBiFactor = numBi if kernel["DirectToVgprB"] and (self.states.packDTVB or self.states.convDTVB): # DirectToVgpr case, we need LoopIters * 2 buffers numBiFactor = kernel["LoopIters"] * 2 if self.states.lrvwTileMXSB > 1: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSB_X0_I0_BASE", "vgprMXSBase", self.states.mxsb.startVgprValu)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMXSB.add(RegSet("v", "vgprValuMXSB_X%u_I%u"%(bi,iui), "vgprValuMXSB_X0_I0_BASE", ri)) ri += self.states.mxsb.numVgprValuPerBlock if not kernel["UnrollMajorLDSB"]: ri = 0 ri = 0 if not kernel["UnrollMajorLDSB"]: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSB_X0_I0_D0_PACK", "vgprMXSBase", self.states.mxsb.startVgprValuPack)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): for data in range(0,kernel["MIInputPerThreadMXSB"]): moduleVgprMacroValuMXSBPack.add(RegSet("v", "vgprValuMXSB_X%u_I%u_D%u"%(bi,iui,data), "vgprValuMXSB_X0_I0_D0_PACK", ri)) ri += ceil(kernel["VectorWidthMXSB"] / self.states.bpr) * kernel["MIWaveTileB"] // kernel["VectorWidthMXSB"] else: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSB_X0_I0_BASE", "vgprMXSBase", self.states.mxsb.startVgprValu)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMXSB.add(RegSet("v", "vgprValuMXSB_X%u_I%u"%(bi,iui), "vgprValuMXSB_X0_I0_BASE", ri)) ri += self.states.mxsb.numVgprValuPerBlock ri = 0 if not kernel["UnrollMajorLDSB"]: moduleVgprMacroMXS.add(RegSet("v", "vgprValuMXSB_X0_I0_D0_PACK", "vgprMXSBase", self.states.mxsb.startVgprValuPack)) for data in range(1,int(self.states.bpr)): for bi in range(0,numBiFactor): # buffer indices if bi % self.states.numVgprBufferPackMXSB == 0: ri = (data-1) * kernel["InnerUnroll"] * self.states.numVgprBufferPackMXSB * self.states.mxsb.numVgprValuPerBlock for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMXSBPack.add(RegSet("v", "vgprValuMXSB_X%u_I%u_D%u"%(bi,iui,data), "vgprValuMXSB_X0_I0_D0_PACK", ri)) ri += self.states.mxsb.numVgprValuPerBlock if kernel["ProblemType"]["MXBlockA"] and (self.states.mxsa.startVgprG2L is not None): moduleVgprMacroMXS.add(RegSet("v", "vgprG2LMXSA_BASE", "vgprMXSBase", self.states.mxsa.startVgprG2L)) if kernel["ProblemType"]["MXBlockB"] and (self.states.mxsb.startVgprG2L is not None): moduleVgprMacroMXS.add(RegSet("v", "vgprG2LMXSB_BASE", "vgprMXSBase", self.states.mxsb.startVgprG2L)) if kernel["ProblemType"]["MXBlockA"]: if not kernel["DirectToLdsMXSA"] or self.do["KeepDirectToLdsAlloc"]: moduleVgprMacroG2LMXSA.add(RegSet("v", "vgprG2LMXSA", "vgprG2LMXSA_BASE", 0)) if kernel["DirectToVgprMXSA"]: # additional definition G2LA2 for swapping register sets moduleVgprMacroG2LMXSA.add(RegSet("v", "vgprG2LMXSA2", "vgprG2LMXSA_BASE", self.states.mxsa.numVgprG2LAllocated//2)) if kernel["ProblemType"]["MXBlockB"]: if not kernel["DirectToLdsMXSB"] or self.do["KeepDirectToLdsAlloc"]: moduleVgprMacroG2LMXSB.add(RegSet("v", "vgprG2LMXSB", "vgprG2LMXSB_BASE", 0)) if kernel["DirectToVgprB"]: # additional definition G2LB2 for swapping register sets moduleVgprMacroG2LMXSB.add(RegSet("v", "vgprG2LMXSB2", "vgprG2LMXSB_BASE", self.states.mxsb.numVgprG2LAllocated//2)) self.moduleVgprMacroMXS = moduleVgprMacroMXS self.moduleVgprMacroValuMXSDummy = moduleVgprMacroValuMXSDummy self.moduleVgprMacroValuMXSA = moduleVgprMacroValuMXSA self.moduleVgprMacroValuMXSB = moduleVgprMacroValuMXSB self.moduleVgprMacroValuMXSAPack = moduleVgprMacroValuMXSAPack self.moduleVgprMacroValuMXSBPack = moduleVgprMacroValuMXSBPack self.moduleVgprMacroG2LMXSA = moduleVgprMacroG2LMXSA self.moduleVgprMacroG2LMXSB = moduleVgprMacroG2LMXSB module.addComment2("VGPR Macro Assignments for MX") module.add(self.moduleVgprMacroMXS) module.add(moduleVgprMacroValuMXSDummy) module.add(self.moduleVgprMacroValuMXSA) module.add(self.moduleVgprMacroValuMXSB) module.add(self.moduleVgprMacroValuMXSAPack) module.add(self.moduleVgprMacroValuMXSBPack) module.add(self.moduleVgprMacroG2LMXSA) module.add(self.moduleVgprMacroG2LMXSB) module.addComment2("VGPR Assignments") module.addComment0("ValuC range: [%u-%u), %s"%(self.states.c.startVgprValu, self.states.c.startVgprValu+self.states.c.numVgprValu, \ "serializedStore enabled" if self.states.serializedStore else "")) module.add(RegSet("v", "vgprValuC", self.states.c.startVgprValu)) module.addComment0("ValuA/B Xn=PLR buffer idx, In=InnerUnroll idx") # PLR index: from X0 to X (at most) -> VGPRs will be duplicated LoopIters times (at most) # eg, if LoopIters = 4, there would be at most 4*VGPRs moduleVgprMacro = Module("VALU/G2L Vgpr Macro") moduleVgprMacroValuA = Module("VALUA Vgpr Macro") moduleVgprMacroValuB = Module("VALUB Vgpr Macro") moduleVgprMacroValuAPack = Module("VALUA Pack Vgpr Macro") moduleVgprMacroValuBPack = Module("VALUB Pack Vgpr Macro") moduleVgprMacroValuM = Module("VALUMetadata Vgpr Macro") moduleVgprMacroValuMPack = Module("VALUMetadata Pack Vgpr Macro") moduleVgprMacroG2LA = Module("G2LA Vgpr Macro") moduleVgprMacroG2LB = Module("G2LB Vgpr Macro") module.add(RegSet("v", "vgprBase", self.states.startVgpr)) ri = 0 if self.states.a.numVgprValu > 0: # Do not generate vgprValuA if numVgprValuA is 0 numBiFactor = numBi if kernel["DirectToVgprA"] and (self.states.packDTVA or self.states.convDTVA): # DirectToVgpr case, we need LoopIters * 2 buffers numBiFactor = kernel["LoopIters"] * 2 if self.states.lrvwTileA > 1: moduleVgprMacro.add(RegSet("v", "vgprValuA_X0_I0_BASE", "vgprBase", self.states.a.startVgprValu - self.states.startVgpr)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuA.add(RegSet("v", "vgprValuA_X%u_I%u"%(bi,iui), "vgprValuA_X0_I0_BASE", ri)) ri += self.states.a.numVgprValuPerBlock if tPA["bpe"] < 4 and not kernel["UnrollMajorLDSA"] and not (kernel["UsePLRPack"] and self.states.numItersPLR): ri = 0 ri = 0 if tPA["bpe"] < 4 and not kernel["UnrollMajorLDSA"] and not kernel["enableLDSTrA"]: moduleVgprMacro.add(RegSet("v", "vgprValuA_X0_I0_D0_PACK", "vgprBase", self.states.a.startVgprValuPack - self.states.startVgpr)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): for data in range(0,kernel["MIInputPerThreadA"]): moduleVgprMacroValuAPack.add(RegSet("v", "vgprValuA_X%u_I%u_D%u"%(bi,iui,data),"vgprValuA_X0_I0_D0_PACK", ri)) ri += ceil(kernel["VectorWidthA"] * tPA["bpe"] / self.states.bpr) * kernel["MIWaveTileA"] // kernel["VectorWidthA"] if (kernel["UsePLRPack"] and self.states.numItersPLR): ri = 0 else: moduleVgprMacro.add(RegSet("v", "vgprValuA_X0_I0_BASE", "vgprBase", self.states.a.startVgprValu - self.states.startVgpr)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuA.add(RegSet("v", "vgprValuA_X%u_I%u"%(bi,iui), "vgprValuA_X0_I0_BASE", ri)) ri += self.states.a.numVgprValuPerBlock if kernel["HalfPLRA"]: ri = 0 for bi in range(3): # buffer indices moduleVgprMacroValuA.add(RegSet("v", "vgprValuA_G%u"%(bi), "vgprValuA_X0_I0_BASE", ri)) ri += self.states.a.numVgprValuPerBlock // 2 ri = 0 if tPA["bpe"] < 4 and not kernel["UnrollMajorLDSA"] and not kernel["enableLDSTrA"]: moduleVgprMacro.add(RegSet("v", "vgprValuA_X0_I0_D0_PACK", "vgprBase", self.states.a.startVgprValuPack - self.states.startVgpr)) for data in range(1,int(self.states.bpr/tPA["bpeDS"])): for bi in range(0,numBiFactor): # buffer indices if bi % self.states.numVgprBufferPackA == 0: ri = (data-1) * kernel["InnerUnroll"] * self.states.numVgprBufferPackA * self.states.a.numVgprValuPerBlock for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuAPack.add(RegSet("v", "vgprValuA_X%u_I%u_D%u"%(bi,iui,data),"vgprValuA_X0_I0_D0_PACK", ri)) ri += self.states.a.numVgprValuPerBlock ri = 0 if self.states.b.numVgprValu > 0: # Do not generate vgprValuB if numVgprValuB is 0 numBiFactor = numBi if kernel["DirectToVgprB"] and (self.states.packDTVB or self.states.convDTVB): # DirectToVgpr case, we need LoopIters * 2 buffers numBiFactor = kernel["LoopIters"] * 2 if self.states.lrvwTileB > 1: moduleVgprMacro.add(RegSet("v", "vgprValuB_X0_I0_BASE", "vgprBase", self.states.b.startVgprValu - self.states.startVgpr)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuB.add(RegSet("v", "vgprValuB_X%u_I%u"%(bi,iui), "vgprValuB_X0_I0_BASE", ri)) ri += self.states.b.numVgprValuPerBlock if (tPB["bpe"] < 4 and not kernel["UnrollMajorLDSB"]) and not (kernel["UsePLRPack"] and self.states.numItersPLR): ri = 0 ri = 0 if tPB["bpe"] < 4 and not kernel["UnrollMajorLDSB"] and not kernel["enableLDSTrB"]: moduleVgprMacro.add(RegSet("v", "vgprValuB_X0_I0_D0_PACK", "vgprBase", self.states.b.startVgprValuPack - self.states.startVgpr)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): for data in range(0,kernel["MIInputPerThreadB"]): moduleVgprMacroValuBPack.add(RegSet("v", "vgprValuB_X%u_I%u_D%u"%(bi,iui,data), "vgprValuB_X0_I0_D0_PACK", ri)) ri += ceil(kernel["VectorWidthB"] * tPB["bpe"] / self.states.bpr) * kernel["MIWaveTileB"] // kernel["VectorWidthB"] if (kernel["UsePLRPack"] and self.states.numItersPLR): ri = 0 else: moduleVgprMacro.add(RegSet("v", "vgprValuB_X0_I0_BASE", "vgprBase", self.states.b.startVgprValu - self.states.startVgpr)) for bi in range(0,numBiFactor): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuB.add(RegSet("v", "vgprValuB_X%u_I%u"%(bi,iui), "vgprValuB_X0_I0_BASE", ri)) ri += self.states.b.numVgprValuPerBlock if kernel["HalfPLRB"]: ri = 0 for bi in range(3): # buffer indices moduleVgprMacroValuB.add(RegSet("v", "vgprValuB_G%u"%(bi), "vgprValuB_X0_I0_BASE", ri)) ri += self.states.b.numVgprValuPerBlock // 2 ri = 0 if tPB["bpe"] < 4 and not kernel["UnrollMajorLDSB"] and not kernel["enableLDSTrB"]: moduleVgprMacro.add(RegSet("v", "vgprValuB_X0_I0_D0_PACK", "vgprBase", self.states.b.startVgprValuPack - self.states.startVgpr)) for data in range(1,int(self.states.bpr/tPB["bpeDS"])): for bi in range(0,numBiFactor): # buffer indices if bi % self.states.numVgprBufferPackB == 0: ri = (data-1) * kernel["InnerUnroll"] * self.states.numVgprBufferPackB * self.states.b.numVgprValuPerBlock for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuBPack.add(RegSet("v", "vgprValuB_X%u_I%u_D%u"%(bi,iui,data), "vgprValuB_X0_I0_D0_PACK", ri)) ri += self.states.b.numVgprValuPerBlock # T reg definition for F32XEmu self.macroAndSetF32XEmuTreg(kernel, tPA, tPB) if kernel["ConvertAfterDS"]: cvtTemp = max(self.states.a.startVgprValuCvtTemp, self.states.b.startVgprValuCvtTemp) if (cvtTemp != -1): moduleVgprMacro.add(RegSet("v", "vgprCvtTemp", "vgprBase", cvtTemp - self.states.startVgpr)) if kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["UseBias"] and (kernel["ProblemType"]["BiasSrc"] == "A" or kernel["ProblemType"]["BiasSrc"] == "B"): moduleVgprMacro.add(RegSet("v", "vgprValuSum", "vgprBase", self.states.bias.startVgprValu - self.states.startVgpr)) if kernel["ProblemType"]["Sparse"]: if kernel["DirectToVgprSparseMetadata"]: moduleVgprMacro.add(RegSet("v", "vgprValuMetadata_X0_I0_BASE", "vgprBase", self.states.m.startVgprValu - self.states.startVgpr)) moduleVgprMacroValuM.add(RegSet("v", "vgprValuMetadata", "vgprValuMetadata_X0_I0_BASE", 0)) else: ri = 0 if self.states.m.numVgprValu > 0: # Do not generate vgprValu if numVgprValu is 0 if self.states.lrvwTileMetadata > 1: moduleVgprMacro.add(RegSet("v", "vgprValuMetadata_X0_I0_BASE", "vgprBase", self.states.m.startVgprValu - self.states.startVgpr)) for bi in range(0,PLR): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuM.add(RegSet("v", "vgprValuMetadata_X%u_I%u"%(bi,iui), "vgprValuMetadata_X0_I0_BASE", ri)) ri += self.states.m.numVgprValuPerBlock if not kernel["UnrollMajorLDSMetadata"]: ri = 0 ri = 0 if not kernel["UnrollMajorLDSMetadata"] and not kernel["enableLDSTrMetadata"]: miWaveTile = kernel["MIWaveTileB"] if kernel["ProblemType"]["Sparse"] == 2 else kernel["MIWaveTileA"] moduleVgprMacro.add(RegSet("v", "vgprValuMetadata_X0_I0_D0_PACK", "vgprBase", self.states.m.startVgprValuPack - self.states.startVgpr)) for data in range(0,kernel["MIInputPerThreadMetadata"]): for bi in range(0,PLR): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMPack.add(RegSet("v", "vgprValuMetadata_X%u_I%u_D%u"%(bi,iui,data), "vgprValuMetadata_X0_I0_D0_PACK", ri)) ri += ceil(kernel["VectorWidthMetadata"] * tPM["bpe"] / self.states.bpr) * miWaveTile // kernel["VectorWidthMetadata"] else: moduleVgprMacro.add(RegSet("v", "vgprValuMetadata_X0_I0_BASE", "vgprBase", self.states.m.startVgprValu - self.states.startVgpr)) for bi in range(0,PLR): # buffer indices for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuM.add(RegSet("v", "vgprValuMetadata_X%u_I%u"%(bi,iui), "vgprValuMetadata_X0_I0_BASE", ri)) ri += self.states.m.numVgprValuPerBlock if not kernel["UnrollMajorLDSMetadata"] and not kernel["enableLDSTrMetadata"]: moduleVgprMacro.add(RegSet("v", "vgprValuMetadata_X0_I0_D0_PACK", "vgprBase", self.states.m.startVgprValuPack - self.states.startVgpr)) for data in range(1,kernel["MIInputPerThreadMetadata"]): for bi in range(0,PLR): # buffer indices if bi % self.states.numVgprBufferPackMetadata == 0: ri = (data-1) * kernel["InnerUnroll"] * self.states.numVgprBufferPackMetadata * kernel["MIWaveTileMetadata"] for iui in range(0, kernel["InnerUnroll"]): moduleVgprMacroValuMPack.add(RegSet("v", "vgprValuMetadata_X%u_I%u_D%u"%(bi,iui,data),"vgprValuMetadata_X0_I0_D0_PACK", ri)) ri += kernel["MIWaveTileMetadata"] if not kernel["LocalWriteUseSgprA"] and self.states.a.numVgprLocalWriteAddr > 0: module.add(RegSet("v", "vgprLocalWriteAddrA", \ self.states.a.startVgprLocalWriteAddr)) if self.states.a.numVgprLocalWriteAddr > 1: module.add(RegSet("v", "vgprLocalWriteAddrOverhangA", \ self.states.a.startVgprLocalWriteAddr+1)) if kernel["ProblemType"]["MXBlockA"]: if not kernel["LocalWriteUseSgprMXSA"] and self.states.mxsa.numVgprLocalWriteAddr > 0: module.add(RegSet("v", "vgprLocalWriteAddrMXSA", \ self.states.mxsa.startVgprLocalWriteAddr)) if self.states.mxsa.numVgprLocalWriteAddr > 1: module.add(RegSet("v", "vgprLocalWriteAddrOverhangMXSA", \ self.states.mxsa.startVgprLocalWriteAddr+1)) if not kernel["LocalWriteUseSgprB"] and self.states.b.numVgprLocalWriteAddr > 0: module.add(RegSet("v", "vgprLocalWriteAddrB", \ self.states.b.startVgprLocalWriteAddr)) if self.states.b.numVgprLocalWriteAddr > 1: module.add(RegSet("v", "vgprLocalWriteAddrOverhangB", \ self.states.b.startVgprLocalWriteAddr+1)) if kernel["ProblemType"]["MXBlockB"]: if not kernel["LocalWriteUseSgprMXSB"] and self.states.mxsb.numVgprLocalWriteAddr > 0: module.add(RegSet("v", "vgprLocalWriteAddrMXSB", \ self.states.mxsb.startVgprLocalWriteAddr)) if self.states.mxsb.numVgprLocalWriteAddr > 1: module.add(RegSet("v", "vgprLocalWriteAddrOverhangMXSB", \ self.states.mxsb.startVgprLocalWriteAddr+1)) if self.states.m.numVgprLocalWriteAddr > 0: module.add(RegSet("v", "vgprLocalWriteAddrMetadata", \ self.states.m.startVgprLocalWriteAddr)) if self.states.m.numVgprLocalWriteAddr > 1: module.add(RegSet("v", "vgprLocalWriteAddrOverhangMetadata", \ self.states.m.startVgprLocalWriteAddr+1)) if kernel["BufferLoad"]: module.add(RegSet("v", "vgprGlobalReadOffsetA", \ self.startVgprGlobalReadOffsetA)) if kernel["ProblemType"]["MXBlockA"]: module.add(RegSet("v", "vgprGlobalReadOffsetMXSA", \ self.startVgprGlobalReadOffsetMXSA)) module.add(RegSet("v", "vgprGlobalReadOffsetB", \ self.startVgprGlobalReadOffsetB)) if kernel["ProblemType"]["MXBlockB"]: module.add(RegSet("v", "vgprGlobalReadOffsetMXSB", \ self.startVgprGlobalReadOffsetMXSB)) if kernel["ProblemType"]["Sparse"]: module.add(RegSet("v", "vgprGlobalReadOffsetMetadata", \ self.startVgprGlobalReadOffsetMetadata)) else: module.add(RegSet("v", "vgprGlobalReadAddrA", \ self.startVgprGlobalReadAddressesA)) if kernel["ProblemType"]["MXBlockA"]: module.add(RegSet("v", "vgprGlobalReadAddrMXSA", \ self.startVgprGlobalReadAddressesMXSA)) module.add(RegSet("v", "vgprGlobalReadAddrB", \ self.startVgprGlobalReadAddressesB)) if kernel["ProblemType"]["MXBlockB"]: module.add(RegSet("v", "vgprGlobalReadAddrMXSB", \ self.startVgprGlobalReadAddressesMXSB)) if self.states.a.startVgprG2L is not None: moduleVgprMacro.add(RegSet("v", "vgprG2LA_BASE", "vgprBase", self.states.a.startVgprG2L - self.states.startVgpr)) if self.states.b.startVgprG2L is not None: moduleVgprMacro.add(RegSet("v", "vgprG2LB_BASE", "vgprBase", self.states.b.startVgprG2L - self.states.startVgpr)) if not kernel["DirectToLdsA"] or self.do["KeepDirectToLdsAlloc"]: moduleVgprMacroG2LA.add(RegSet("v", "vgprG2LA", "vgprG2LA_BASE", 0)) if kernel["DirectToVgprA"]: # additional definition G2LA2 for swapping register sets moduleVgprMacroG2LA.add(RegSet("v", "vgprG2LA2", "vgprG2LA_BASE", self.states.a.numVgprG2LAllocated//2)) if not kernel["DirectToLdsB"] or self.do["KeepDirectToLdsAlloc"]: moduleVgprMacroG2LB.add(RegSet("v", "vgprG2LB", "vgprG2LB_BASE", 0)) if kernel["DirectToVgprB"]: # additional definition G2LB2 for swapping register sets moduleVgprMacroG2LB.add(RegSet("v", "vgprG2LB2", "vgprG2LB_BASE", self.states.b.numVgprG2LAllocated//2)) if kernel["UnrollLoopSwapGlobalReadOrder"] and not kernel["DirectToLdsA"] and not kernel["DirectToLdsB"]: if kernel["ULSGRODoubleG2L"] == 0: moduleVgprMacroG2LA.add(RegSet("v", "vgprG2LB2", "vgprG2LA_BASE", 0)) moduleVgprMacroG2LA.add(RegSet("v", "vgprG2LA2", "vgprG2LA_BASE", self.states.b.numVgprG2LAllocated)) else: moduleVgprMacroG2LA.add(RegSet("v", "vgprG2LA2", "vgprG2LA_BASE", self.states.a.numVgprG2LAllocated)) moduleVgprMacroG2LB.add(RegSet("v", "vgprG2LB2", "vgprG2LB_BASE", self.states.b.numVgprG2LAllocated)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: moduleVgprMacro.add(RegSet("v", "vgprG2LMetadata", "vgprBase", self.states.m.startVgprG2L - self.states.startVgpr)) if ((tPA["bpe"] < 4 and not kernel["UnrollMajorLDSA"]) or \ (tPB["bpe"] < 4 and not kernel["UnrollMajorLDSB"]) or \ (kernel["ProblemType"]["Sparse"] and not kernel["UnrollMajorLDSMetadata"] and (kernel["MIInputPerThreadMetadata"] == 4))) \ and (kernel["ProblemType"]["DataType"].isInt8() or kernel["ProblemType"]["DataType"].is8bitFloat()) or \ (self.states.asmCaps["HasSWMMAC_gfx1250"] and kernel["ProblemType"]["Sparse"] and not kernel["UnrollMajorLDSMetadata"]): moduleVgprMacro.add(RegSet("v", "vgprPackTemp", "vgprBase", self.states.a.startVgprValuPackTemp - self.states.startVgpr)) if self.states.globalReadIncsUseVgpr: module.add(RegSet("v", "vgprGlobalReadIncsA", \ self.startVgprGlobalReadIncsA)) if kernel["ProblemType"]["MXBlockA"]: module.add(RegSet("v", "vgprGlobalReadIncsMXSA", \ self.startVgprGlobalReadIncsMXSA)) module.add(RegSet("v", "vgprGlobalReadIncsB", \ self.startVgprGlobalReadIncsB)) if kernel["ProblemType"]["MXBlockB"]: module.add(RegSet("v", "vgprGlobalReadIncsMXSB", \ self.startVgprGlobalReadIncsMXSB)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.add(RegSet("v", "vgprGlobalReadIncsMetadata", \ self.startVgprGlobalReadIncsMetadata)) if self.states.a.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrA", \ self.states.a.startVgprLocalReadAddr)) if kernel["ProblemType"]["MXBlockA"]: if self.states.mxsa.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrMXSA", \ self.states.mxsa.startVgprLocalReadAddr)) if self.states.b.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrB", \ self.states.b.startVgprLocalReadAddr)) if kernel["ProblemType"]["MXBlockB"]: if self.states.mxsb.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrMXSB", \ self.states.mxsb.startVgprLocalReadAddr)) if self.states.IncLdsBufSwitch: # 3 or more LDS buffer case. Need to keep original vgprLocalReadAdd if self.states.a.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrOrigA", \ self.states.a.startVgprLocalReadAddrOrig)) if kernel["ProblemType"]["MXBlockA"]: if self.states.mxsa.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrOrigMXSA", \ self.states.mxsa.startVgprLocalReadAddrOrig)) if self.states.b.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrOrigB", \ self.states.b.startVgprLocalReadAddrOrig)) if kernel["ProblemType"]["MXBlockB"]: if self.states.mxsb.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrOrigMXSB", \ self.states.mxsb.startVgprLocalReadAddrOrig)) if self.states.m.numVgprLocalReadAddr > 0: module.add(RegSet("v", "vgprLocalReadAddrMetadata", \ self.states.m.startVgprLocalReadAddr)) if self.states.a.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrA", \ self.states.a.startVgprLocalReadSwapAddr)) if self.states.b.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrB", \ self.states.b.startVgprLocalReadSwapAddr)) if self.states.mxsa.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrMXSA", \ self.states.mxsa.startVgprLocalReadSwapAddr)) if self.states.mxsb.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrMXSB", \ self.states.mxsb.startVgprLocalReadSwapAddr)) if self.states.m.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrMetadata", \ self.states.m.startVgprLocalReadSwapAddr)) if kernel["ProblemType"]["MXBlockA"]: if self.states.mxsa.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrMXSA", \ self.states.mxsa.startVgprLocalReadSwapAddr)) if kernel["ProblemType"]["MXBlockB"]: if self.states.mxsb.numVgprLocalReadSwapAddr > 0: module.add(RegSet("v", "vgprLocalReadSwapAddrMXSB", \ self.states.mxsb.startVgprLocalReadSwapAddr)) if self.states.a.numVgprLocalWriteSwapAddr > 0: module.add(RegSet("v", "vgprLocalWriteSwapAddrA", \ self.states.a.startVgprLocalWriteSwapAddr)) if self.states.b.numVgprLocalWriteSwapAddr > 0: module.add(RegSet("v", "vgprLocalWriteSwapAddrB", \ self.states.b.startVgprLocalWriteSwapAddr)) if self.states.m.numVgprLocalWriteSwapAddr > 0: module.add(RegSet("v", "vgprLocalWriteSwapAddrMetadata", \ self.states.m.startVgprLocalWriteSwapAddr)) if kernel["ProblemType"]["MXBlockA"]: if self.states.mxsa.numVgprLocalWriteSwapAddr > 0: module.add(RegSet("v", "vgprLocalWriteSwapAddrMXSA", \ self.states.mxsa.startVgprLocalWriteSwapAddr)) if kernel["ProblemType"]["MXBlockB"]: if self.states.mxsb.numVgprLocalWriteSwapAddr > 0: module.add(RegSet("v", "vgprLocalWriteSwapAddrMXSB", \ self.states.mxsb.startVgprLocalWriteSwapAddr)) if kernel["ProblemType"]["OutputAmaxD"]: module.add(RegSet("v", "vgprAmaxOut", self.startVgprAmaxOut)) module.add(RegSet("v", "vgprAmaxOutB", self.startVgprAmaxOutB)) if kernel["ProblemType"]["DataType"].isDoubleComplex() and kernel["MIArchVgpr"]: module.add(RegSet("v", "vgprAlphaTmp", \ self.states.startVgprAlphaTmp)) self._emitGL2PrefetchAddrRegSets(module, kernel, tPA, tPB) module.add(RegSet("v", "vgprSerial", self.states.startVgprSerial)) if self.debugConfig.debugKernel: module.add(RegSet("v", "vgprAddressDbg", \ self.states.startVgprAddressDbg)) #module.addComment0("Occu: %u waves/simd" % self.numWavesPerSimd ) # module.addComment0("Num VGPR=%u"%self.vgprPool.size()) # module.addComment0("Num AccVGPR=%u"%self.agprPool.size()) self.moduleVgprMacro = moduleVgprMacro self.moduleVgprMacroValuA = moduleVgprMacroValuA self.moduleVgprMacroValuB = moduleVgprMacroValuB self.moduleVgprMacroValuAPack = moduleVgprMacroValuAPack self.moduleVgprMacroValuBPack = moduleVgprMacroValuBPack self.moduleVgprMacroValuM = moduleVgprMacroValuM self.moduleVgprMacroValuMPack = moduleVgprMacroValuMPack self.moduleVgprMacroG2LA = moduleVgprMacroG2LA self.moduleVgprMacroG2LB = moduleVgprMacroG2LB module.addComment2("VGPR Macro Assignments") module.add(self.moduleVgprMacro) module.add(self.moduleVgprMacroValuA) module.add(self.moduleVgprMacroValuA_T) module.add(self.moduleVgprMacroValuB) module.add(self.moduleVgprMacroValuB_T) module.add(self.moduleVgprMacroValuAPack) module.add(self.moduleVgprMacroValuBPack) module.add(self.moduleVgprMacroValuM) module.add(self.moduleVgprMacroValuMPack) module.add(self.moduleVgprMacroG2LA) module.add(self.moduleVgprMacroG2LB) def macroAndSetVgprImplSubtile(): module.add(RegSet("v", "vgprSerial", self.states.startVgprSerial)) #self.vgprPool.remove(self.states.startVgprSerial, 1) #module.addComment0("Need %u vgprs for GR A"%(self.states.a.tileInfo.numGRPerSubtile)) self._emitGL2PrefetchAddrRegSets(module, kernel, tPA, tPB) return if not kernel["UseSubtileImpl"]: macroAndSetImplClassic() else: macroAndSetVgprImplSubtile() ######################################## # SGPR Macros ######################################## module.addComment2("SGPR Assignments") # Emit declarations for all sgprs allocated with defineSgpr # in the order they were declared for skey in self.sgprs: module.add(RegSet("s", "sgpr"+skey, self.sgprs[skey])) # module.addComment0("max SGPR=%u"%self.sgprPool.size()) if self.states.streamK.emitsParallelReductionSgprAliases: module.addSpaceLine() module.addComment0("StreamK Parallel Reduction Assignments") module.add(RegSet("s", "sgprSkSplit", "sgprskTiles", 0)) module.add(RegSet("s", "sgprSkPartialIdx", "sgprBeta", 0)) if kernel["StreamK"] == 5: # SK5 hybrid: the kernel only declares the 6 SK3-named kernarg slots # (see KernelWriter.py SK5 defineSgpr block). The SK4 code path still # reads via SK4 names (TotalItems, SKTiles, SKSplit, SKItersPerWI, # SKGrid); emit RegSet aliases so those names resolve to the matching # SK3-named slot. Slot 0 (ItersPerTile) is shared and needs no alias. module.addComment0("SK5 hybrid: SK4 reader name -> SK3 primary slot") module.add(RegSet("s", "sgprTotalItems", "sgprMagicNumberItersPerTile", 0)) module.add(RegSet("s", "sgprSKTiles", "sgprMagicShiftItersPerTile", 0)) module.add(RegSet("s", "sgprSKSplit", "sgprSKItersPerWG", 0)) module.add(RegSet("s", "sgprSKItersPerWI", "sgprskGrid", 0)) module.add(RegSet("s", "sgprSKGrid", "sgprskTiles", 0)) # SK5 hybrid: the SK3-only and SK4-only persistent SGPRs are mutually # exclusive at runtime (StreamKHybridMode selects one path), so the # SK3-only iter SGPRs are RegSet-aliased onto the SK4-only idx SGPRs # (defined via the for-skey RegSet loop above) instead of taking their # own physical registers. Drops 2 real SGPRs vs the naive union. module.addComment0("SK5 hybrid: SK3-only iter SGPR -> SK4-only idx slot (mutually exclusive)") module.add(RegSet("s", "sgprStreamKIter", "sgprStreamKTileIdx", 0)) module.add(RegSet("s", "sgprStreamKIterEnd", "sgprStreamKPartialIdx", 0)) elif kernel["StreamK"] == 4: module.add(RegSet("s", "sgprSkPartialIdx", "sgprBeta", 0)) module.addSpaceLine() module.addComment0("Size Assignments") problemType = kernel["ProblemType"] for idx in range(max(problemType["IndexAssignmentsA"] + problemType["IndexAssignmentsB"])+1): idxChar= INDEX_CHARS[idx] if idx in problemType["IndicesFree"] or idx in problemType["IndicesBatch"]: idxType="Free" elif idx in problemType["IndicesSummation"]: idxType="Sum" idx = idx - problemType["NumIndicesC"] else: raise ValueError("unexpected index type in size assignments") module.add(RegSet("s", "sgprSize%s"%(idxChar), \ "sgprSizes%s"%idxType, idx)) module.addSpaceLine() module.addComment0("Stride Assignments") for tc in ('D','C'): for idx in range(0, problemType["NumIndicesC"]): i = idx idxChar= self.states.indexChars[idx] if i == 0 and not kernel["ProblemType"]["UseInitialStridesCD"]: module.add(ValueSet("constStride%s%s"%(tc,idxChar), 1)) else: if not kernel["ProblemType"]["UseInitialStridesCD"]: i = i-1 module.add(RegSet("s", "sgprStride%s%s"%(tc,idxChar), \ "sgprStrides%s"%tc, i)) tcList = ["A", "B"] if kernel["ProblemType"]["Sparse"]: tcList.append("Metadata") if kernel["ProblemType"]["MXBlockA"]: tcList += ["MXSA"] if kernel["ProblemType"]["MXBlockB"]: tcList += ["MXSB"] for tc in tcList: for i, idx in enumerate(problemType["IndexAssignments%s"%tc]): idxChar= self.states.indexChars[idx] if i == 0 and not kernel["ProblemType"]["UseInitialStridesAB"]: module.add(ValueSet("constStride%s%s"%(tc,idxChar), 1)) else: if not kernel["ProblemType"]["UseInitialStridesAB"]: i = i-1 module.add(RegSet("s", "sgprStride%s%s"%(tc,idxChar), \ "sgprStrides%s"%tc, i)) module.addSpaceLine() module.add(ValueSet("MT0", kernel["MacroTile0"])) module.add(ValueSet("MT1", kernel["MacroTile1"])) module.add(ValueSet("DepthU", kernel["DepthU"])) module.addComment0("Number of elements to shift-left SRD") module.add(ValueSet("SrdShiftLeftA", self.states.srdShiftLeft['A'])) if kernel["ProblemType"]["MXBlockA"]: module.add(ValueSet("SrdShiftLeftMXSA", self.states.srdShiftLeft['MXSA'])) module.add(ValueSet("SrdShiftLeftB", self.states.srdShiftLeft['B'])) if kernel["ProblemType"]["MXBlockB"]: module.add(ValueSet("SrdShiftLeftMXSB", self.states.srdShiftLeft['MXSB'])) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.add(ValueSet("SrdShiftLeftMetadata", self.states.srdShiftLeft["Metadata"])) if kernel["BufferLoad"] or kernel["BufferStore"]: module.addComment0("2GB limit - set offsets to -1 to exceed this and clamp") module.add(ValueSet("BufferLimit", 0xffffffff, format=1)) # Use 2^31 for BufferOOB behavior. # set BufferOOB to NumRecords field of SRD. # if thread is OOB, we can invalid it by setting vOffset to BufferOOB module.add(ValueSet("BufferOOB", 0xfffff000, format=1)) # Set BufferOOB to 0xfffff000 instead of 0xffffffff (reserving 4KB) to prevent # overflows from instruction offsets (max 4095) during memory access. srdUpperValue = SrdUpperValue(self.states.version) module.addComment2("Bits 127:96 of SRD.\n" + srdUpperValue.desc()) module.add(ValueSet("Srd127_96", srdUpperValue.getValue(), format=1)) if ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1)): module.addComment0("MT offset for 64b address (=MT0*MT1*bpeC)") reductionOffset = int(kernel["MacroTile0"]*kernel["MacroTile1"]*self.states.bpeCinternal) reductionOffsetLow32 = reductionOffset & 0xFFFFFFFF reductionOffsetHigh32 = (reductionOffset >> 32) & 0xFFFFFFFF module.add(ValueSet("MTOffset", reductionOffsetLow32, format=1)) module.add(ValueSet("MTOffsetH32", reductionOffsetHigh32, format=1)) if self.states.IncLdsBufSwitch or self.states.useCommonSgprSwap: module.addComment0("%d LDS Blocks for PGR %d"%(self.states.numLDSBlk, kernel["PrefetchGlobalRead"])) module.add(ValueSet("LdsOneBlockSize", kernel["LdsOffsetA_Blk"])) module.add(ValueSet("LdsBlockEndSize", kernel["LdsOffsetA_Blk"] * self.states.numLDSBlk)) ######################################## # Global Offsets ######################################## # justOffset32 means we should only write the 32-bit offset # This is used in Buffer addressing modes. # Flat addressing modes expect the GLOBAL_OFFSET to initialize a full 64-bit address GOList = [ \ ("C", list(range(0, kernel["ProblemType"]["NumIndicesC"])), kernel["BufferStore"], None, False), \ ("A", kernel["ProblemType"]["IndexAssignmentsA"], kernel["BufferLoad"], tPA, False), \ ("B", kernel["ProblemType"]["IndexAssignmentsB"], kernel["BufferLoad"], tPB, False) ] if kernel["ProblemType"]["MXBlockA"]: GOList.append(("MXSA", kernel["ProblemType"]["IndexAssignmentsMXSA"], kernel["BufferLoad"], tPMXSA, False)) if kernel["ProblemType"]["MXBlockB"]: GOList.append(("MXSB", kernel["ProblemType"]["IndexAssignmentsMXSB"], kernel["BufferLoad"], tPMXSB, False)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: GOList.append(("Metadata", kernel["ProblemType"]["IndexAssignmentsMetadata"], kernel["BufferLoad"], tPM, False)) if kernel["ProblemType"]["SwizzleTensorA"]: GOList.append(("A", kernel["ProblemType"]["IndexAssignmentsA"], kernel["BufferLoad"], tPA, True)) if kernel["ProblemType"]["SwizzleTensorB"]: GOList.append(("B", kernel["ProblemType"]["IndexAssignmentsB"], kernel["BufferLoad"], tPB, True)) for (tc, indices, justOffset32, tP, isSwizzled) in GOList: # BufferStore does not use this macro so don't generate it: # Subtile impl does not use these macros either if (tc == "C" and kernel["BufferStore"]) or kernel["UseSubtileImpl"]: continue # function name and comment suffix_tc = tc + "_SWIZZLED" if isSwizzled else tc module.addComment1("Global Offset %s"%suffix_tc) numDim = len(indices) idxChars = [] for i in indices: idxChars.append(self.states.indexChars[i]) # macro declaration calcDims = [] # dimensions which are participating in the address calc (ignores other summation) mirrorSumDims = [] macroArgs = [] for i in range(0, numDim): if tc == 'C': useInitialStrides = kernel["ProblemType"]["UseInitialStridesCD"] idxChar = self.states.indexChars[i] else: useInitialStrides = kernel["ProblemType"]["UseInitialStridesAB"] idxChar = self.states.indexChars[tP['ia'][i]] # tile index or unroll vgpr or summation # other summation (other than unroll) are included in the GLOBAL_OFFSET macro but not used in address calc if tc in ('A','C') and indices[i] == kernel["ProblemType"]["Index0"] \ or tc in ('B','C', "Metadata") and indices[i] == kernel["ProblemType"]["Index1"] \ or indices[i] == kernel["ProblemType"]["IndexUnroll"]: macroArgs.append("vgprOffset%s:req" % idxChars[i]) calcDims.append(i) elif indices[i] in kernel["ProblemType"]["IndicesSummation"]: # other summation index (not unroll) if tc in ('A', 'B', "Metadata") and indices[i] in kernel["ProblemType"]["MirrorDims%s" % tc]: mirrorSumDims.append(i) continue else: # other batch or free index if isPackedIndex(kernel, indices[i]): calcDims.append(i) macroArgs.append("vgprOffset%s:req" % idxChars[i]) elif not justOffset32: # buffer/justOffset32 scalars are included in SRD not the offset, so skip here calcDims.append(i) macroArgs.append("sgprOffset%s:req" % idxChars[i]) macro = Macro("GLOBAL_OFFSET_%s" % suffix_tc, ["vgprAddr:req", *macroArgs, "vgprTmp:req"]) # Each index may be skipped, scaled by stride, or unscaled # If destLo is unset, no accumulation is necessary. # if the first index (i==0) is unscaled (UseInitialStrides), # it can be combined at the next update or moved at end # (if there is no next update) pendingOffset = None # offset pending for accumulation offsetIsVgpr = False # True if the source is VGPR ; False if SGPR destLo = None # true for first addr calc. In this case, we can directly write addr # rather than accumulating through a tmp writeDirectToAddr = 1 # mirror other summation indices for i in mirrorSumDims: if writeDirectToAddr: dest = vgpr("Addr+0", isMacro=True) needAdd = 0 # don't need add since writing address directly. writeDirectToAddr = 0 else: dest = vgpr("Tmp+0", isMacro=True) needAdd = 1 macro.add(VSubU32(dst=dest, \ src0=sgpr("Size%s"%INDEX_CHARS[indices[i]]), \ src1=1, \ comment="mirror %s%s 1"%(tc, INDEX_CHARS[indices[i]]))) macro.add(VMulLOU32(dst=dest, \ src0=dest, \ src1=self.strideRef(tc, indices[i]), \ comment="mirror %s%s 2"%(tc, INDEX_CHARS[indices[i]]))) if needAdd: writeDirectToAddr = 0 # safety net, once we write address can't directly overwrite it later destLo = vgpr("Addr+0", isMacro=True) destHi = vgpr("Addr+1", isMacro=True) srcLo = pendingOffset if pendingOffset else destLo srcHi = 0 if pendingOffset else destHi macro.add(VAddCOU32(dst=destLo, \ dst1=VCC(), \ src0=srcLo, \ src1="v[\\vgprTmp+0]", \ comment="accumulate %s lower"%idxChar)) for i in calcDims: # should have eliminated these above idx = indices[i] isMirrorIdx = tc in ('A', 'B', "Metadata") and idx in kernel["ProblemType"]["MirrorDims%s" % tc] assert not (idx in kernel["ProblemType"]["IndicesSummation"] and idx != kernel["ProblemType"]["IndexUnroll"]) if indices[i] == kernel["ProblemType"]["Index0"] \ or indices[i] == kernel["ProblemType"]["Index1"] \ or indices[i] == kernel["ProblemType"]["IndexUnroll"]: offsetIsVgpr = True # other c index sgpr (free or batch) elif indices[i] < kernel["ProblemType"]["NumIndicesC"]: if isPackedIndex(kernel, indices[i]): offsetIsVgpr = True else: offsetIsVgpr = False else: assert(0) # no other type allowed # offset is VGPR or SGPR string to use for the offset if offsetIsVgpr: offset = "v[\\vgprOffset%s]" % idxChars[i] else: offset = "s[\\sgprOffset%s]" % idxChars[i] # macro.addComment0("dim%s pendingOffset=%s offset=%s offsetIsVgpr=%s" \ # % (self.states.indexChars[indices[i]], pendingOffset, offset, offsetIsVgpr)) needAdd = 0 # should be indices[i]?? if i==0 and not useInitialStrides: # slide into next address calc - can do addr = pendingOffset + nextAddrCalc pendingOffset = offset writeDirectToAddr = 0 else: # tile index or unroll vgpr if offsetIsVgpr: if writeDirectToAddr: destLo = vgpr("Addr+0", isMacro=True) destHi = vgpr("Addr+1", isMacro=True) needAdd = 0 # don't need add since writing address directly. writeDirectToAddr = 0 else: destLo = vgpr("Tmp+0", isMacro=True) destHi = vgpr("Tmp+1", isMacro=True) needAdd = 1 if isMirrorIdx: macro.add(VSubI32( dst=vgpr("Tmp+0", isMacro=True), src0=sgpr("Size%s"%INDEX_CHARS[idx]), \ src1=offset, \ comment="mirror %s%s 1"%(tc, INDEX_CHARS[indices[i]]))) macro.add(VSubI32(\ dst=vgpr("Tmp+0", isMacro=True), src0="v[\\vgprTmp+0]", \ src1=1, \ comment="mirror %s%s 2"%(tc, INDEX_CHARS[indices[i]]))) offset = "v[\\vgprTmp+0]" # offset * stride if isSwizzled: # Swizzle directly uses the offset base that has already been multiplied by the stride. macro.add(VMovB32(dst=destLo, src=offset)) else: macro.add(VMulLOU32(dst=destLo, src0=self.strideRef(tc, indices[i]), \ src1=offset, \ comment="mul d%u lower"%i)) if not justOffset32: macro.add(VMulHIU32(dst=destHi, src0=self.strideRef(tc, indices[i]), \ src1=offset, \ comment="mul d%u upper"%i)) else: # offset is SGPR: assert not isMirrorIdx if not justOffset32: # buffer mode (aka justOffset32) does scalars into SRD not offset macro.add(VMovB32(dst=vgpr("Tmp+2", isMacro=True), src="s[\\sgprOffset%s]"%idxChars[i], \ comment="sgprOffset -> vgprTmp+2")) # offset * stride macro.add(VMulLOU32(dst=vgpr("Tmp+0", isMacro=True), \ src0=self.strideRef(tc, indices[i]), src1="v[\\vgprTmp+2]", \ comment="other stride mul d%u lower"%i)) macro.add(VMulHIU32(dst=vgpr("Tmp+1", isMacro=True), \ src0=self.strideRef(tc, indices[i]), src1="v[\\vgprTmp+2]", \ comment="mul d%u upper"%i)) needAdd = 1 if needAdd: writeDirectToAddr = 0 # safety net, once we write address can't directly overwrite it later destLo = vgpr("Addr+0", isMacro=True) destHi = vgpr("Addr+1", isMacro=True) # addr += offset * stride (lo) : accumulate just-computed address term into addr srcLo = pendingOffset if pendingOffset else destLo srcHi = 0 if pendingOffset else destHi macro.add(VAddCOU32(dst=destLo, dst1=VCC(), \ src0=srcLo, src1="v[\\vgprTmp+0]", \ comment="accumulate %s lower"%idxChar)) # addr += offset * stride (hi) if not justOffset32: macro.add(VAddCCOU32(dst=vgpr("Addr+1", isMacro=True), dst1=VCC(), \ src0="v[\\vgprTmp+1]", src1=srcHi, src2=VCC(), \ comment="accumulate %s upper"%idxChar)) pendingOffset = None # pendingOffset but never got a chance to apply it, # need to just add an explicit move or add: # this can happen for small-order tensors if pendingOffset != None: destLo = vgpr("Addr+0", isMacro=True) if writeDirectToAddr: macro.add(VMovB32(dst=destLo, src=offset, comment="setup d0 lower")) if not justOffset32: macro.add(VMovB32(dst=vgpr("Addr+1", isMacro=True), src=0, comment="d0 upper")) else: macro.add(VAddCOU32(dst=destLo, dst1=VCC(), \ src0=destLo, src1=pendingOffset, \ comment="accumulate final pendingOffset")) if tP != None and kernel["BufferLoad"] and self.states.srdShiftLeft[tc]: macro.add(VAddU32(dst=vgpr("Addr+0", isMacro=True), \ src0=hex(self.states.srdShiftLeft[tc]), \ src1="v[\\vgprAddr+0]", \ comment="add prepad for pointer shift")) module.add(macro) if kernel["ProblemType"]["StochasticRounding"] and not self.states.asmCaps["v_prng_b32"] : module.add(PseudoRandomGenerator()) if not kernel["EnableMatrixInstruction"]: # Macro MAC PLR = kernel["PrefetchLocalRead"] \ if kernel["PrefetchLocalRead"] < kernel["LoopIters"] \ else kernel["LoopIters"] - 1 for m in range(0, 1+PLR): macro = Macro("MAC_%ux%u_X%u" % (kernel["ThreadTile0"], kernel["ThreadTile1"], m), []) component = Component.MAC.find(self) if not component: printExit("Assembly doesn't support datatype %s" % kernel["ProblemType"]["DataType"]) innerModule = component(self, tPA, tPB, m, kernel["InnerUnroll"]) for item in innerModule.items(): macro.add(item) module.add(macro) # dot2: oneIUI for tail loop if kernel["UseDotInstruction"]: macro = Macro("MAC_%ux%u_X%u_OneIUI" % (kernel["ThreadTile0"],kernel["ThreadTile1"], 0), [""]) component = Component.MAC.find(self) if not component: printExit("Assembly doesn't support datatype %s" % kernel["ProblemType"]["DataType"]) innerModule = component(self, tPA, tPB, 0, 1) for item in innerModule.items(): macro.add(item) module.add(macro) module.setNoOpt(True) return module def checkResources(self, kernel, mkb: KernelBody): # register allocation totalVgprs = self.vgprPool.size() totalAgprs = self.agprPool.size() totalSgprs = self.sgprPool.size() mkb.setGprs(totalVgprs=totalVgprs, totalAgprs=totalAgprs, totalSgprs=totalSgprs) if self.vgprPool.size() > self.states.regCaps["MaxVgpr"]: self.states.overflowedResources = 1 elif self.sgprPool.size() > self.states.regCaps["MaxSgpr"]: self.states.overflowedResources = 2 # TODO: Add target occupancy or kept the occupancy settings from globalWriteBatch kernel["CUOccupancy"] = self.getOccupancy(kernel["NumThreads"], self.vgprPool.size(), self.sgprPool.size(), \ self.getLdsSize(kernel), self.agprPool.size(), self.states.doubleVgpr) if kernel["_ScheduleIterAlg"] == 2 and kernel["CUOccupancy"] < 2: self.states.overflowedResources = 6 if self.states.asmCaps["HasWMMA"]: # RDNA: workgroup executes on a WGP (2 CUs, 4 SIMDs) physicalMaxVgprPerWGP = self.states.regCaps["PhysicalMaxVgprCU"] * 2 vgprBudgetPerThread = physicalMaxVgprPerWGP // kernel["NumThreads"] else: vgprBudgetPerThread = self.states.regCaps["PhysicalMaxVgprCU"] // kernel["NumThreads"] numWorkGroupsPerExecUnit = vgprBudgetPerThread // mkb.getNextFreeVgpr() if numWorkGroupsPerExecUnit < 1: self.states.overflowedResources = 4 if self.states.invalidLSUCode: self.states.overflowedResources = 7 self.vgprPool.checkFinalState() if self.states.overflowedResources: mkb.body.add(ValueEndif("overflowed resources")) if self.states.overflowedResources == 1: msg = "too many vgprs" elif self.states.overflowedResources == 2: msg = "too many sgprs" elif self.states.overflowedResources == 3: msg = "half store requires at least two elements per batch" elif self.states.overflowedResources == 4: msg = "Occupancy limit" elif self.states.overflowedResources == 5: if kernel["DirectToLdsA"] or kernel["DirectToLdsB"]: msg = "cannot schedule local read with DirectToLds" else: msg = "reading and writing LDS at same time require 2 LDS buffer" elif self.states.overflowedResources == 6: msg = "SIA2 better with occupancy 2" elif self.states.overflowedResources == 7: msg = "invalid LSU code due to assertion fail" elif self.states.overflowedResources == 8: msg = "not enough LDS space" else: msg = "unknown" if self.debugConfig.printSolutionRejectionReason: printWarning("%s overflowed resources. errorCode=%d, msg=\"%s\", vgprs=%u, sgprs=%u" \ % (self.states.kernelName, self.states.overflowedResources, msg, \ self.vgprPool.size(), self.sgprPool.size())) mkb.body.add(SEndpgm(comment="overflowed resources"), 0) mkb.body.add(ValueIf(value="0"), 1) ############################################################################## # updateOccupancyFromMaxVgpr ############################################################################## def updateOccupancyFromMaxVgpr(self, kernel, mkb, max_vgpr: int) -> None: """Update CUOccupancy and the kernel descriptor after rocIsaPass. rocIsaPass builds a register interference graph and computes the highest VGPR index actually referenced by instructions; the result is passed in as max_vgpr (= max_vgpr_index + 1 from rocIsaPassResult.maxVgpr). When that is lower than the checkResources pool estimate, the kernel descriptor and CUOccupancy are updated so .amdhsa_next_free_vgpr reflects actual usage. Only runs on ArchAccUnifiedRegs ISAs (gfx90a/gfx942/gfx950). max_vgpr: Pre-computed max-VGPR count from rocIsaPassResult.maxVgpr. When <= 0 (not supplied by the pass), the update is skipped. """ if not self.states.archCaps.get("ArchAccUnifiedRegs"): return if max_vgpr <= 0: return # Acc VGPRs are tracked by a separate pool and are always fully populated # by MFMA instructions; rocIsaPass skips acc registers in its graph, so # keep the pool size as the authoritative AGPR count (unchanged from before). pool_agprs = self.agprPool.size() pool_total = int(ceil(self.vgprPool.size() / 8.0)) * 8 + self.agprPool.size() max_vgpr_total = int(ceil(max_vgpr / 8.0)) * 8 + pool_agprs if max_vgpr_total >= pool_total: return mkb.setGprs(totalVgprs=max_vgpr, totalAgprs=pool_agprs, totalSgprs=self.sgprPool.size()) kernel["CUOccupancy"] = self.getOccupancy( kernel["NumThreads"], max_vgpr, self.sgprPool.size(), self.getLdsSize(kernel), pool_agprs, self.states.doubleVgpr) ############################################################################## # code phrase for load batched address from array of buffer pointer ############################################################################## def loadBatchedAddress(self, kernel, Batch, tmpSgprResource: ContinuousRegister): tmpSgpr = tmpSgprResource.idx laneSC = tmpSgprResource.size module = Module("loadBatchedAddress %s" % Batch) module.addSpaceLine() # handle Batch C/D if not kernel["_GlobalAccumulation"]: for idx in kernel["ProblemType"]["IndicesBatch"]: if not isPackedIndex(kernel,idx): module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address")) module.add(SLoadB64(dst=sgpr("AddressD", 2), base=sgpr("AddressD",2), soffset=sgpr(tmpSgpr), comment="load global buffer D address")) # Only load C buffer address if Beta is used and potentially non-zero if kernel["ProblemType"]["UseBeta"]: endCheckLabel = Label(self.labels.getName(f"label_skip_c_buffer_deref_{Batch}"), "") module.add(BranchIfZero("Beta", kernel["ProblemType"]["ComputeDataType"].toEnum(), tmpSgpr, laneSC, endCheckLabel, \ kernel['WavefrontSize'])) for idx in kernel["ProblemType"]["IndicesBatch"]: if not isPackedIndex(kernel,idx): module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address")) module.add(SLoadB64(dst=sgpr("AddressC", 2), base=sgpr("AddressC",2), soffset=sgpr(tmpSgpr), comment="load global buffer C address")) module.add(endCheckLabel) #handle Batch A/B endCheckLabel = Label(self.labels.getName(f"label_skip_ab_buffer_deref_{Batch}"), "") module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(1), comment="check summation size")) for i in range(0, self.states.numSgprSizesSum): module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr("SizesSum+%u"%(i)), src1=sgpr(tmpSgpr), comment="check summation size")) module.add(SCmpEQU32(src0=sgpr(tmpSgpr), src1=0, comment="skip buffer deref is size of summation is 0")) module.add(SCBranchSCC1(labelName=endCheckLabel.getLabelName(), comment="skip buffer deref is size of summation is 0")) module.add(BranchIfZero("Alpha", kernel["ProblemType"]["ComputeDataType"].toEnum(), tmpSgpr, laneSC, endCheckLabel, \ kernel['WavefrontSize'])) module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address")) for idx in kernel["ProblemType"]["IndicesBatch"]: if not isPackedIndex(kernel,idx): module.add(SLoadB64(dst=sgpr("AddressA", 2), base=sgpr("AddressA",2), soffset=sgpr(tmpSgpr), comment="load global buffer A address")) module.add(SLoadB64(dst=sgpr("AddressB", 2), base=sgpr("AddressB",2), soffset=sgpr(tmpSgpr), comment="load global buffer B address")) module.add(endCheckLabel) return module ############################################################################## def getKernelArgLoadModule(self, kernel, sgprStartIdx, numsOfLoad, preloadNum): kernelArgs = Module("load arguments") kernelArgs.addComment1("Load Kernel Args") if self.debugConfig.debugKernel: kernelArgs.add(self.argLoader.loadKernArg("AddressDbg", "KernArgAddress", dword=2)) self.argLoader.resetOffset() kernelArgs.addModuleAsFlatItems(self.argLoader.loadAllKernArg(sgprStartIdx, "KernArgAddress", numsOfLoad, preloadNum)) if kernel["ProblemType"]["UseScaleAB"] == "Scalar": sgprOffset = self.argLoader.getOffset() for preloadScale, name in zip([self.states.preloadScaleA, self.states.preloadScaleB], ['A','B']): if preloadScale: kernelArgs.add(self.argLoader.loadKernArg("AddressScale%s"%name, "KernArgAddress", sgprOffset=hex(sgprOffset), dword=2)) sgprOffset += (self.states.rpga * self.states.bpr) if kernel["ExpertSchedulingMode"] > 0 and kernel["ESMRuntimeGate"]: sgprOffset = self.argLoader.getOffset() kernelArgs.add(self.argLoader.loadKernArg(self.states.esmRuntimeFlagSgpr, "KernArgAddress", sgprOffset=hex(sgprOffset), dword=1)) sgprOffset += 4 self.argLoader.setOffset(sgprOffset) return kernelArgs def localReadAddresses(self, kernel, tPA, tPB, tPM): module = Module("Local Read Addresses") #################################### # Local Read Addresses #################################### module.addComment2("Local Read Addresses") # tile assignments module.addComment1("local read addresses: tile assignments a/b") module.add(self.lraTileAssignment(kernel, tPA, tPB)) # final offsets module.addComment1("local read addresses: final offsets a") module.add(self.lraFinalOffset(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.addComment1("local read addresses: final offsets mxsa") module.add(self.lraFinalOffset(kernel, tPA["MX"])) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.addComment1("local read addresses: final offsets metadata") module.add(self.lraFinalOffset(kernel, tPM)) if kernel["ProblemType"]["MXBlockB"]: module.addComment1("local read addresses: final offsets mxsb") module.add(self.lraFinalOffset(kernel, tPB["MX"])) module.addComment1("local read addresses: final offsets b") module.add(self.lraFinalOffset(kernel, tPB)) # declare addresses module.addComment1("local read addresses: declare addresses a") module.add(self.lraDeclareAddresses(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.addComment1("local read addresses: declare addresses mxsa") module.add(self.lraDeclareAddresses(kernel, tPA["MX"])) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.addComment1("local read addresses: declare addresses metadata") module.add(self.lraDeclareAddresses(kernel, tPM)) if kernel["ProblemType"]["MXBlockB"]: module.addComment1("local read addresses: declare addresses mxsb") module.add(self.lraDeclareAddresses(kernel, tPB["MX"])) module.addComment1("local read addresses: declare addresses b") module.add(self.lraDeclareAddresses(kernel, tPB)) if self.states.a.numVgprLocalReadAddr > 0: module.add(self.lraSwapAddressesForDTLPad(kernel, tPA)) module.add(self.lraAddressesInitFor3LDSBlk(kernel, tPA, False, True)) if self.states.mxsa.numVgprLocalReadAddr > 0: module.add(self.lraSwapAddressesForDTLPad(kernel, tPA["MX"])) module.add(self.lraAddressesInitFor3LDSBlk(kernel, tPA["MX"], False, True)) if self.states.b.numVgprLocalReadAddr > 0: module.add(self.lraSwapAddressesForDTLPad(kernel, tPB)) module.add(self.lraAddressesInitFor3LDSBlk(kernel, tPB, False, True)) if self.states.mxsb.numVgprLocalReadAddr > 0: module.add(self.lraSwapAddressesForDTLPad(kernel, tPB["MX"])) module.add(self.lraAddressesInitFor3LDSBlk(kernel, tPB["MX"], False, True)) if self.states.m.numVgprLocalReadAddr > 0: module.add(self.lraSwapAddressesForDTLPad(kernel, tPM)) return module def localWriteAddresses(self, kernel, tPA, tPB, tPM): module = Module("Local Write Addresses") #################################### # Local Write Addresses #################################### module.addComment2("Local Write Addresses") # tile assignments if not kernel["enableTDMA"]: module.add(self.lwaTileAssignment(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.add(self.lwaTileAssignment(kernel, tPA["MX"])) if not kernel["enableTDMB"]: module.add(self.lwaTileAssignment(kernel, tPB)) if kernel["ProblemType"]["MXBlockB"]: module.add(self.lwaTileAssignment(kernel, tPB["MX"])) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"] and not kernel["enableTDMMetadata"]: module.add(self.lwaTileAssignment(kernel, tPM)) # unroll assignments if not kernel["enableTDMA"]: module.add(self.lwaUnrollAssignment(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.add(self.lwaUnrollAssignment(kernel, tPA["MX"])) if not kernel["enableTDMB"]: module.add(self.lwaUnrollAssignment(kernel, tPB)) if kernel["ProblemType"]["MXBlockB"]: module.add(self.lwaUnrollAssignment(kernel, tPB["MX"])) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"] and not kernel["enableTDMMetadata"]: module.add(self.lwaUnrollAssignment(kernel, tPM)) # first offsets module.addComment1("local write addresses: first offset a") if not kernel["enableTDMA"]: module.add(self.lwaFirstOffset(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.addComment1("local write addresses: first offset mxsa") module.add(self.lwaFirstOffset(kernel, tPA["MX"])) if not kernel["enableTDMB"]: if kernel["ProblemType"]["MXBlockB"]: module.addComment1("local write addresses: first offset mxsb") module.add(self.lwaFirstOffset(kernel, tPB["MX"])) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"] and not kernel["enableTDMMetadata"]: module.addComment1("local write addresses: first offset metadata") module.add(self.lwaFirstOffset(kernel, tPM)) module.addComment1("local write addresses: first offset b") if not kernel["enableTDMB"]: module.add(self.lwaFirstOffset(kernel, tPB)) return module def remapWgSerial(self, kernel, earlyStop=True): module = Module("RemapWgSerial") ######## # remap wg serial to wg0,wg1,wg2 ######## # FIXME: Here does not support UseBatch: False if "WorkGroup2" in self.sgprs: with self.allocTmpSgpr(2, tag="remapWgSerial_tmpSgpr") as tmpSgpr: module.addComment1("remap wg from 1D(idxWG012) to 3D(wg2,wg1,wg0)") module.addComment0("wg2 = idxWG012 * smallMagicNumber(1/(numWG0*numWG1))") tmpVgpr = self.vgprPool.checkOut(2, tag="remapWgSerial_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) module.add(SMulI32(dst=sgpr(tmpSgpr.idx), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"))) if kernel["GlobalSplitU"] != 0: module.add(SAndB32(dst=sgpr(tmpSgpr.idx+1), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SMulI32(dst=sgpr(tmpSgpr.idx), src0=sgpr(tmpSgpr.idx), src1=sgpr(tmpSgpr.idx+1))) module.add(scalarUInt32DivideAndRemainder(qReg=tmpSgpr.idx, dReg="WorkGroup0", divReg=tmpSgpr.idx, rReg=tmpSgpr.idx+1,\ tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=False)) module.add(SMovB32(dst=sgpr("WorkGroup2"), src=sgpr(tmpSgpr.idx))) # Convert 1D to 2D sgprWGM = "WGM" module.addComment0("idxWG01 = idxWG012 - wg2 * numWG0 * numWG1") module.add(SMulI32(dst=sgpr(tmpSgpr.idx), src0=sgpr("NumWorkGroups1"), src1=sgpr("NumWorkGroups0"))) module.add(SMulI32(dst=sgpr(tmpSgpr.idx), src0=sgpr(tmpSgpr.idx), src1=sgpr("WorkGroup2"))) if kernel["GlobalSplitU"] != 0: module.add(SMulI32(dst=sgpr(tmpSgpr.idx), src0=sgpr(tmpSgpr.idx), src1=sgpr(tmpSgpr.idx+1))) module.add(SSubU32(dst=sgpr("WorkGroup0"), src0=sgpr("WorkGroup0"), src1=sgpr(tmpSgpr.idx))) module.addComment0("wg1 = idxWG01 * smallMagicNumber(1/numWG0)") module.add(scalarUInt32DivideAndRemainder(qReg=tmpSgpr.idx, dReg="WorkGroup0", divReg="NumWorkGroups0", rReg=tmpSgpr.idx+1,\ tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=False)) self.vgprPool.checkIn(tmpVgpr) module.add(SMovB32(dst=sgpr("WorkGroup1"), src=sgpr(tmpSgpr.idx))) module.addComment0("wg0 = idxWG01 - wg1 * numWG0") module.add(SMulI32(dst=sgpr(tmpSgpr.idx), src0=sgpr("WorkGroup1"), src1=sgpr("NumWorkGroups0"))) module.add(SSubU32(dst=sgpr("WorkGroup0"), src0=sgpr("WorkGroup0"), src1=sgpr(tmpSgpr.idx))) # early stop if wgIdx exceed wg needed if earlyStop: module.addComment1("Early stop if wg exceed") module.add(SCmpGeU32(src0=sgpr("WorkGroup2"), src1=sgpr("SizesFree+2"))) label_EarlyStop = Label(self.labels.getNameInc("EarlyStop_if_wg_exceed"), "") label_nonEarlyStop = Label(self.labels.getNameInc("NoEarlyStop_wgExceed"), "") module.add(SCBranchSCC0(labelName=label_nonEarlyStop.getLabelName())) module.add(label_EarlyStop) module.add(SEndpgm()) module.add(label_nonEarlyStop) return module def disableWmmaArbStall(self) -> Module: mod = Module() if self.states.archCaps["HasWmmaArbStallBit"]: mod.add(SSetRegIMM32B32(HWRegContainer(reg="26", value=[4, 1]), src=1, comment="Disable WMMA arb stall")) return mod def defineAndResources(self, kernel, tPA, tPB, tPM): module = Module("allocateResources") module.add(self.macroAndSet(kernel, tPA, tPB)) module.addComment2("Allocate Resources") moduleArgs = Module("load arguments") moduleRegInit = Module("Init regs") tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] sgprNumsOfGemm = None if self.do["PreLoop"]: # Need to guard again since some defined sgprs are added into sgprPool if self.states.numSgprPreload > 0: while(1): tmpSgpr = self.sgprPool.checkOut(1, tag="defineAndResources_preloadGuard_tmpSgpr", preventOverflow=False) if tmpSgpr >= self.states.archCaps["MaxSgprPreload"]: self.sgprPool.checkIn(tmpSgpr) break self.states.preloadGuard.append(tmpSgpr) ### temp sgpr for groupedgemm ### # can be start from sgpr_preload_end sgprNumsOfGemm = self.sgprPool.checkOut(1, tag="defineAndResources_sgprNumsOfGemm", preventOverflow=False) self.kernArgOffset = 0 self.argLoader = ArgumentLoader() self.externalArgLoader = ArgumentLoader() ######################################## # Common parameters sgprArgType = self.sgprPool.checkOut(1, tag="defineAndResources_sgprArgType", preventOverflow=False) sgprPackedArgs = self.sgprPool.checkOut(1, tag="defineAndResources_sgprPackedArgs", preventOverflow=False) tmpSgprNumWorkGroups = self.sgprPool.checkOut(1, tag="defineAndResources_tmpSgprNumWorkGroups", preventOverflow=False) ######################################## # kernel args parameters load = self.states.numSgprToLoad sgprStart = self.sgprs["SizesFree"] if self.states.numSgprPreload == 0 or self.states.archCaps["SgprPreloadPad"]: commonArgs = Module("load arguments") commonArgs.addComment1("Load num of Gemms") commonArgs.add(self.argLoader.loadKernArg(sgprNumsOfGemm, "KernArgAddress", 0, dword=1)) # Load combined internal arguments commonArgs.addComment1("Load packed kernel args (StaggerU/GSU)") commonArgs.add(self.argLoader.loadKernArg(sgprPackedArgs, "KernArgAddress", 4, dword=1)) commonArgs.addComment1("Load WGM data") commonArgs.add(self.argLoader.loadKernArg("WGM", "KernArgAddress", 8, dword=1)) commonArgs.addComment1("Load num of WGs") commonArgs.add(self.argLoader.loadKernArg(tmpSgprNumWorkGroups, "KernArgAddress", 12, dword=1)) ######################################## # load ws/ user args hbmArgs = Module("load HBM arguments") hbmArgs.addComment1("Load address of kernel arguments") hbmArgs.add(self.argLoader.loadKernArg("KernArgAddress", "KernArgAddress", self.states.userArgsInfo.commonArgsSize, dword=2)) moduleArgs.addModuleAsFlatItems(deepcopy(commonArgs)) moduleArgs.add(SWaitCnt(kmcnt=0, comment="load args")) moduleArgs.add(SLShiftRightB32(dst=sgpr(sgprArgType), shiftHex=hex(30), src=sgpr(sgprNumsOfGemm), comment="Get arg type")) moduleArgs.add(SAndB32(dst=sgpr(sgprNumsOfGemm), src0=hex(0x3FFFFFFF), src1=sgpr(sgprNumsOfGemm), comment="Get nums of gemm")) if ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1)): extReadEpilogueLabeltmp = Label(label=self.labels.getNameInc("LoadExternalEpilogueStruct"), comment="") moduleArgs.addComment0("Check if custom structure pointer is null") if kernel["ProblemType"]["SupportUserArgs"]: moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=2, comment="ArgType == 2 ?")) moduleArgs.add(SCBranchSCC0(labelName=extReadEpilogueLabeltmp.getLabelName())) moduleArgs.addComment1("Grouped Gemm: Load address of external kernel arguments") moduleArgs.add(self.argLoader.loadKernArg("AddressTD", "KernArgAddress", hex(self.states.userArgsInfo.commonArgsSize+16), dword=2)) moduleArgs.add(self.argLoader.loadKernArg("Synchronizer", "KernArgAddress", hex(self.states.userArgsInfo.commonArgsSize+8), dword=2)) moduleArgs.add(extReadEpilogueLabeltmp) #moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=(0), comment="Is kernel args")) labelHBM = Label("HBMArgs", comment="") labelLoadEnd = Label("LoadArgsEnd", comment="") # Routing General Batched GEMM to Strided Batched GEMM path Bypass_ArgType3_to_ArgType0_Instance1 = Label("Bypass_ArgType3_to_ArgType0_Instance1", comment="") moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=(3), comment="Is kernel argType == 3")) moduleArgs.add(SCBranchSCC1(labelName=Bypass_ArgType3_to_ArgType0_Instance1.getLabelName())) moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=(0), comment="Is kernel args")) moduleArgs.add(SCBranchSCC0(labelName=labelHBM.getLabelName())) moduleArgs.add(Bypass_ArgType3_to_ArgType0_Instance1) moduleArgs.add(SAddU32(dst=sgpr("KernArgAddress"), src0=sgpr("KernArgAddress"), src1=hex(self.states.userArgsInfo.commonArgsSize), comment="Shift common args")) moduleArgs.add(SAddCU32(dst=sgpr("KernArgAddress+1"), src0=sgpr("KernArgAddress+1"), src1=0)) if kernel["ExpertSchedulingMode"] > 0 and kernel["ESMRuntimeGate"]: self.states.esmRuntimeFlagSgpr = self.sgprPool.checkOut(1, tag="defineAndResources_esmRuntimeFlagSgpr", preventOverflow=False) self.states.esmRuntimeFlagVgpr = self.vgprPool.checkOut(1, tag="defineAndResources_esmRuntimeFlagVgpr") moduleArgs.addModuleAsFlatItems(self.getKernelArgLoadModule(kernel, sgprStart, load, 0)) if self.states.numSgprPreload > 0: moduleArgs.add(SWaitCnt(kmcnt=0, comment="preload")) moduleArgs.add(SBranch(labelName=labelLoadEnd.getLabelName())) moduleArgs.add(labelHBM) moduleArgs.addModuleAsFlatItems(deepcopy(hbmArgs)) moduleArgs.add(SWaitCnt(kmcnt=0, comment="wait for args to load")) moduleArgs.add(labelLoadEnd) if self.states.numSgprPreload > 0: if self.states.archCaps["SgprPreloadPad"]: common_kern_entry = Label(label="common_kernel_entry", comment="for both preload/non-preload common code") #For groupgemm, the preload happened prior to this stage moduleArgs.add(SBranch(common_kern_entry.getLabelName())) # jump to common path total_inst_dwords = 0 for inst in moduleArgs.items(): if isinstance(inst, (BranchInstruction, SWaitCnt, CommonInstruction)): total_inst_dwords = total_inst_dwords + 1 elif isinstance(inst, (SMemLoadInstruction)): total_inst_dwords = total_inst_dwords + 2 assert total_inst_dwords <= 64 moduleArgs.addComment1("pad %u snops to satisfy 0x100 code size for Preload Backward Compatibility Prologue" % (64 - total_inst_dwords)) for i in range(64 - total_inst_dwords): moduleArgs.add(SNop(waitState=0, comment="")) moduleArgs.add(Label("Preload_Offset_Start", "")) # Common args preload preloadSgprStartIdx = self.states.rpga moduleArgs.add(SAndB32(dst=sgpr(sgprNumsOfGemm), src0=hex(0x3FFFFFFF), src1=sgpr(preloadSgprStartIdx), comment="Get nums of gemm")) moduleArgs.add(SLShiftRightB32(dst=sgpr(sgprArgType), shiftHex=hex(30), src=sgpr(preloadSgprStartIdx), comment="Get arg type")) if ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1)): extReadEpilogueLabeltmp = Label(label=self.labels.getNameInc("LoadExternalEpilogueStruct"), comment="") moduleArgs.addComment0("Check if custom structure pointer is null") if kernel["ProblemType"]["SupportUserArgs"]: moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=2, comment="ArgType == 2 ?")) moduleArgs.add(SCBranchSCC0(labelName=extReadEpilogueLabeltmp.getLabelName())) moduleArgs.add(SMovB32(dst=sgpr("Synchronizer+1"), src=sgpr(preloadSgprStartIdx+7), comment="Load Synchronizer data")) moduleArgs.add(SMovB32(dst=sgpr("Synchronizer"), src=sgpr(preloadSgprStartIdx+6), comment="Load Synchronizer data")) moduleArgs.add(SMovB32(dst=sgpr("AddressTD+1"), src=sgpr(preloadSgprStartIdx+9), comment="Load AddressTD data")) moduleArgs.add(SMovB32(dst=sgpr("AddressTD"), src=sgpr(preloadSgprStartIdx+8), comment="Load AddressTD data")) moduleArgs.add(extReadEpilogueLabeltmp) moduleArgs.add(SMovB32(dst=sgpr(sgprPackedArgs), src=sgpr(preloadSgprStartIdx+1), comment="Preload internal args")) # Routing the General Batched GEMM to Strided Batched GEMM path Bypass_ArgType3_to_ArgType0_Instance2 = Label("Bypass_ArgType3_to_ArgType0_Instance2", comment="") moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=(3), comment="Is kernel argType == 3")) moduleArgs.add(SCBranchSCC1(labelName=Bypass_ArgType3_to_ArgType0_Instance2.getLabelName())) moduleArgs.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=(0), comment="Is kernel args")) preloadLabelHBM = Label("Preload_HBMArgs", comment="") perloadLabelLoadEnd = Label("Preload_LoadArgsEnd", comment="") moduleArgs.add(SCBranchSCC0(labelName=preloadLabelHBM.getLabelName())) moduleArgs.add(Bypass_ArgType3_to_ArgType0_Instance2) moduleArgs.add(SAddU32(dst=sgpr("KernArgAddress"), src0=sgpr("KernArgAddress"), src1=hex(self.states.userArgsInfo.commonArgsSize), comment="Shift common args")) moduleArgs.add(SAddCU32(dst=sgpr("KernArgAddress+1"), src0=sgpr("KernArgAddress+1"), src1=0)) self.argLoader.resetOffset() moduleArgs.addModuleAsFlatItems(self.getKernelArgLoadModule(kernel, sgprStart, load, self.states.numSgprPreload - self.states.userArgsInfo.commonArgsNum)) i = self.states.userArgsInfo.commonArgsNum while i < self.states.numSgprPreload: dstIndex = sgprStart+i-self.states.userArgsInfo.commonArgsNum srcIndex = preloadSgprStartIdx+i if (i + 1 < self.states.numSgprPreload) and (dstIndex % 2 == 0) and (srcIndex % 2 == 0): moduleArgs.add(SMovB64(dst=sgpr(sgprStart+i-self.states.userArgsInfo.commonArgsNum, 2), src=sgpr(preloadSgprStartIdx+i, 2), comment="move preload data to correct sgpr")) i += 2 else: moduleArgs.add(SMovB32(dst=sgpr(sgprStart+i-self.states.userArgsInfo.commonArgsNum), src=sgpr(preloadSgprStartIdx+i), comment="move preload data to correct sgpr")) i += 1 moduleArgs.add(SBranch(labelName=perloadLabelLoadEnd.getLabelName())) moduleArgs.add(preloadLabelHBM) moduleArgs.add(SMovB64(dst=sgpr("KernArgAddress", 2), src=sgpr(preloadSgprStartIdx+4, 2), comment="Load address of kernel arguments")) moduleArgs.add(perloadLabelLoadEnd) moduleArgs.add(SMovB32(dst=sgpr("WGM"), src=sgpr(preloadSgprStartIdx+2), comment="Preload internal args2")) moduleArgs.add(SMovB32(dst=sgpr(tmpSgprNumWorkGroups), src=sgpr(preloadSgprStartIdx+3), comment="Load num of WGs")) # add common kern entry label if self.states.archCaps["SgprPreloadPad"]: moduleRegInit.add(common_kern_entry) if not self.states.archCaps["WorkGroupIdFromTTM"]: for i in range(kernel["ProblemType"]["NumIndicesC"]): moduleRegInit.add(SMovB32(dst=sgpr("WorkGroup0+%u"%i), src=sgpr(preloadSgprStartIdx+self.states.numSgprPreload+i), \ comment="restore workgroup id")) moduleRegInit.add(SAndB32(dst=sgpr("StaggerU"), src0=sgpr(sgprPackedArgs), src1=hex(0xFFFF0000), comment="Restore StaggerU related vars")) moduleRegInit.add(SLShiftRightB32(dst=sgpr("StaggerU"), shiftHex=hex(16), src=sgpr("StaggerU"))) if kernel["GlobalSplitU"] != 0 or kernel["AdaptiveGemmNTAB"] != 0: moduleRegInit.add(SAndB32(dst=sgpr("GSU"), src0=sgpr(sgprPackedArgs), src1=hex(0xFFFF), comment="Restore GSUConfig and GSU")) if kernel["ProblemType"]["SupportUserArgs"]: moduleRegInit.add(SMovB32(dst=sgpr("ArgType"),src=sgpr(sgprArgType))) self.sgprPool.checkIn(sgprPackedArgs) if kernel["StorePriorityOpt"]: moduleRegInit.add(SSetPrior(prior=3, comment="optimization store")) if self.do["PreLoop"]: if self.db["InitSgpr"] & 0x1: moduleRegInit.addComment1("Init SGPRs") for i in range(self.states.firstInitSgpr, self.sgprPool.size()): moduleRegInit.add(SMovB32(dst=sgpr(i), src=hex(self.consts.initSgprValue), comment="InitSgpr&0x1")) moduleRegInit.addSpaceLine() if self.db["InitVgpr"] & 0x1: moduleRegInit.addComment1("Init VGPRs") for i in range(1, self.states.totalVgprs): moduleRegInit.add(VMovB32(dst=vgpr(i), src=hex(self.consts.initVgprValue), comment="InitVgpr&0x1")) moduleRegInit.addSpaceLine() # set m0 moduleRegInit.add(SMovB32(dst=mgpr(0), src=hex(kernel["LdsNumBytes"]), comment="LDS clamp at %u bytes"%(kernel["LdsNumBytes"]))) # set Serial id vgpr moduleRegInit.add(VMovB32(dst=vgpr("Serial"), src=vgpr(0), comment="thread serial id")) if self.states.kernel["WavefrontSize"] == 32: moduleRegInit.add(SMovB32(dst=VCC(setHi=True), src=0, comment="Ensure hi bits are zero")) # Initialize WaveIdx whenever TDM is enabled, independent of WorkGroupIdFromTTM. # TDM helpers (initTDMDescriptor*, tdmGlobalOffset*, tdmSetupIncrement*, # TensorDataMover wave-separated paths) read sgpr("WaveIdx") unconditionally, # so it must be initialized for every TDM kernel. tdmA: bool = kernel["enableTDMA"] tdmB: bool = kernel["enableTDMB"] if tdmA or tdmB: with self.allocTmpSgpr(1, tag="defineAndResources_tmpSgprRes") as tmpSgprRes: wavelen: int = kernel["WavefrontSize"] sgprIdx: int = tmpSgprRes.idx moduleRegInit.add(VReadfirstlaneB32(sgpr(sgprIdx), vgpr("Serial"), "first tId")) moduleRegInit.add(SLShiftRightB32(sgpr("WaveIdx"), ceil(log2(wavelen)), sgpr(sgprIdx), "wId=fTid // wavelen")) # init workgroup id from ttmp if self.states.archCaps["WorkGroupIdFromTTM"]: if not clusterEnabled(kernel["ClusterDim"]): moduleRegInit.addComment1("Init workgroup id from ttmp") moduleRegInit.add(SMovB32(dst=sgpr("WorkGroup0"), src="ttmp9")) moduleRegInit.add(SAndB32(dst=sgpr("WorkGroup1"), src0=hex(0xFFFF), src1="ttmp7")) moduleRegInit.add(SLShiftRightB32(dst=sgpr("WorkGroup2"), shiftHex=hex(0x10), src="ttmp7")) else: module.addComment1("Init workgroup id from ttmp with cluster remap") with self.allocTmpSgpr(5, tag="defineAndResources_tmpSgprInfo") as tmpSgprInfo: sTmp = tmpSgprInfo.idx label_enable_cluster = Label(label="EnableCluster", comment="") label_calculate_workgroup_done = Label(label="RemapWorkGroupDone", comment="") moduleRegInit.add(SGetRegB32(dst=sgpr(sTmp+0), src=HWRegContainer(reg="HW_REG_IB_STS2", \ value=[6, 4]), comment="cluster_id")) moduleRegInit.add(SCmpEQU32(src0=sgpr(sTmp+0), src1=0, comment="cluster_id == 0x0 ?")) moduleRegInit.add(SCBranchSCC0(label_enable_cluster.getLabelName())) moduleRegInit.add(SMovB32(dst=sgpr("WorkGroup0"), src="ttmp9", comment="")) moduleRegInit.add(SAndB32(dst=sgpr("WorkGroup1"), src0=hex(0xFFFF), src1="ttmp7", comment="")) moduleRegInit.add(SLShiftRightB32(dst=sgpr("WorkGroup2"), shiftHex=hex(0x10), src="ttmp7")) moduleRegInit.add(SBranch(label_calculate_workgroup_done.getLabelName())) moduleRegInit.add(label_enable_cluster) moduleRegInit.add(SMovB32(dst=sgpr(sTmp+0), src="ttmp6", comment="Read TTMP6 register")) moduleRegInit.add(SMovB32(dst=sgpr(sTmp+1), src="ttmp7", comment="Read TTMP7 register, \ cluster_z | cluster_y.")) moduleRegInit.add(SBfeU32(dst=sgpr(sTmp+2), src0=sgpr(sTmp+0), src1=((4 << 16) | 4), \ comment="Etract wg_y.")) moduleRegInit.add(SBfeU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+0), src1=((4 << 16) | 16), \ comment="Etract nwg_y. Value is nwg_y - 1")) moduleRegInit.add(SAddU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=1, comment="")) moduleRegInit.add(SAndB32(dst=sgpr(sTmp+4), src0=sgpr(sTmp+1), src1=hex(0xFFFF), \ comment="Etract cluster_y.")) moduleRegInit.add(SMulI32(dst=sgpr("WorkGroup1"), src0=sgpr(sTmp+4), src1=sgpr(sTmp+3),\ comment="cluster_y * nwg_y")) moduleRegInit.add(SAddU32(dst=sgpr("WorkGroup1"), src0=sgpr("WorkGroup1"), \ src1=sgpr(sTmp+2), \ comment="WorkGroup1 = (cluster_y * nwg_y) + wg_y")) moduleRegInit.add(SBfeU32(dst=sgpr("WorkGroup2"), src0=sgpr(sTmp+1), src1=((16 << 16) | 16), \ comment="Etract cluster_z.")) moduleRegInit.add(SAndB32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+0), src1=hex(0xF), \ comment="Etract wg_x.")) moduleRegInit.add(SBfeU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+0), src1=((4 << 16) | 12), \ comment="Etract nwg_x. Value is nwg_x - 1")) moduleRegInit.add(SAddU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=1, comment="")) moduleRegInit.add(SMulI32(dst=sgpr("WorkGroup0"), src0="ttmp9", src1=sgpr(sTmp+3),\ comment="cluster_x * nwg_x")) moduleRegInit.add(SAddU32(dst=sgpr("WorkGroup0"), src0=sgpr("WorkGroup0"), \ src1=sgpr(sTmp+1), \ comment="WorkGroup0 = (cluster_x * nwg_x) + wg_x")) moduleRegInit.add(SBfeU32(dst=sgpr(sTmp+4), src0=sgpr(sTmp+0), src1=((4 << 16) | 8), \ comment="Etract wg_z.")) moduleRegInit.add(SBfeU32(dst=sgpr(sTmp+0), src0=sgpr(sTmp+0), src1=((4 << 16) | 20), \ comment="Etract nwg_z. Value is nwg_z - 1")) moduleRegInit.add(SAddU32(dst=sgpr(sTmp+0), src0=sgpr(sTmp+0), src1=1, comment="")) moduleRegInit.add(SMulI32(dst=sgpr("WorkGroup2"), src0=sgpr("WorkGroup2"), src1=sgpr(sTmp+0),\ comment="cluster_z * nwg_z")) moduleRegInit.add(SAddU32(dst=sgpr("WorkGroup2"), src0=sgpr("WorkGroup2"), \ src1=sgpr(sTmp+4), \ comment="WorkGroup2 = (cluster_z * nwg_z) + wg_z")) moduleRegInit.add(label_calculate_workgroup_done) if kernel["Multicast"]: moduleRegInit.addComment0("Calculate multicast mask") sgprWgX = sTmp+1 sgprWgY = sTmp+2 sgprNWgX = sTmp+3 maskA = 1 for idx in range(kernel["ClusterDim"][1]): maskA |= (1 << (idx * kernel["ClusterDim"][0])) maskB = (1 << kernel["ClusterDim"][0]) - 1 if kernel["enableTDMMetadata"]: if kernel["ProblemType"]["Sparse"] == 1: moduleRegInit.add(SLShiftLeftB32(dst=sgpr("MulticastMaskMetadata"), shiftHex=sgpr(sgprWgX), src=hex(maskA),\ comment="Setting metadata mask (follows sparse A)")) elif kernel["ProblemType"]["Sparse"] == 2: moduleRegInit.add(SMulI32(dst=sgpr(sTmp+4), src0=sgpr(sgprWgY), src1=sgpr(sgprNWgX),\ comment="Shift factor: wg_y * nwg_x (metadata)")) moduleRegInit.add(SLShiftLeftB32(dst=sgpr("MulticastMaskMetadata"), shiftHex=sgpr(sTmp+4), src=hex(maskB),\ comment="Setting metadata mask (follows sparse B)")) if tdmA and tdmB and kernel["NumWaves"] > 1 and not kernel.get("UseSubtileImpl"): setMulticastMaskLblOdd = Label(f"setMulticastMask_OddWave", "") setMulticastMaskLblEven = Label(f"setMulticastMask_EvenWave", "") setMulticastMaskLblEnd = Label(f"setMulticastMaskEnd", "") moduleRegInit.add(SBitcmp1B32(sgpr("WaveIdx"), 0, "Check parity of wId")) moduleRegInit.add(SCBranchSCC1(setMulticastMaskLblOdd.getLabelName(), "Jump if wId is odd")) moduleRegInit.add(setMulticastMaskLblEven) moduleRegInit.add(SLShiftLeftB32(dst=sgpr("MulticastMask"), shiftHex=sgpr(sgprWgX), src=hex(maskA),\ comment="Setting maskA for even wave")) moduleRegInit.add(SBranch(setMulticastMaskLblEnd.getLabelName())) moduleRegInit.add(setMulticastMaskLblOdd) moduleRegInit.add(SMulI32(dst=sgpr(sgprWgY), src0=sgpr(sgprWgY), src1=sgpr(sgprNWgX),\ comment="Shift factor: wg_y * nwg_x")) moduleRegInit.add(SLShiftLeftB32(dst=sgpr("MulticastMask"), shiftHex=sgpr(sgprWgY), src=hex(maskB),\ comment="Setting maskB for odd wave")) moduleRegInit.add(setMulticastMaskLblEnd) else: moduleRegInit.add(SLShiftLeftB32(dst=sgpr("MulticastMaskA"), shiftHex=sgpr(sgprWgX), src=hex(maskA),\ comment="Setting maskA")) moduleRegInit.add(SMulI32(dst=sgpr(sgprWgY), src0=sgpr(sgprWgY), src1=sgpr(sgprNWgX),\ comment="Shift factor: wg_y * nwg_x")) moduleRegInit.add(SLShiftLeftB32(dst=sgpr("MulticastMaskB"), shiftHex=sgpr(sgprWgY), src=hex(maskB),\ comment="Setting maskB")) # SrdD can be used as temp sgprs for a bit if self.states.doShadowInit and kernel["BufferStore"]: self.addSgprVarToPool("SrdD") waitForScaleAB = False moduleScaleAB = Module("Load ScaleAB") if kernel["ProblemType"]["UseScaleAB"] == "Scalar": for preloadScale, name in zip([self.states.preloadScaleA, self.states.preloadScaleB], ['A','B']): if preloadScale: waitForScaleAB = True moduleScaleAB.add(SMovB32(dst=sgpr("Scale%s"%name), src=1.0 , comment="init as 1" )) label = Label(self.labels.getNameInc("Scale%sValid"%name), "") moduleScaleAB.add(BranchIfZero("AddressScale%s"%name, DataType('int64').toEnum(), -1, int(kernel["WavefrontSize"]/32), label, kernel["WavefrontSize"])) # load scale data moduleScaleAB.add(SLoadB32(dst=sgpr("Scale%s"%name), base=sgpr("AddressScale%s"%name,2), soffset=0, comment="load scale%s"%name)) moduleScaleAB.add(label) moduleWg = Module("Calculate Workgroup") # C regs are not used during initialization so mark them as available - # we will claim then just before the start of the unroll loop: if not kernel["UseSubtileImpl"]: if self.states.lastValuMXSAB: self.vgprPool.add(0 , \ self.states.lastValuMXSAB, "ValuMXSAB") # Add as available moduleWg.addComment0("init: add vgpr [%u...%u) to pool" % \ (self.states.mxsa.startVgprValu, self.states.lastValuMXSAB+self.states.mxsa.startVgprValu)) self.vgprPool.add(self.states.a.startVgprValu , \ self.states.lastValuAB - self.states.a.startVgprValu , "ValuAB") # Add as available moduleWg.addComment0("init: add vgpr [%u...%u) to pool" % \ (self.states.a.startVgprValu, self.states.lastValuAB+self.states.a.startVgprValu)) self.vgprPool.add(self.states.c.startVgprValu, \ self.states.c.numVgprValu, "ValuC-Block") # Add as available moduleWg.addComment0("init: add vgpr [%u...%u) to pool" % \ (self.states.c.startVgprValu, self.states.c.startVgprValu+self.states.c.numVgprValu)) numAccvgprs = self.states.totalAgprs self.agprPool.add(0, numAccvgprs, "ValuC-Block") moduleWg.addComment0("init: add agpr [%u...%u) to pool" % \ (0, numAccvgprs)) if kernel["StreamK"] == 0 and not kernel["UseSubtileImpl"]: moduleWg.add(self.localReadAddresses(kernel, tPA, tPB, tPM)) moduleWg.add(self.localWriteAddresses(kernel, tPA, tPB, tPM)) def waitForArgsToLoad(): if kernel["ProblemType"]["SupportUserArgs"]: moduleWg.add(SWaitCnt(kmcnt=0, comment="wait for %u/%u bytes of kern args" % \ (self.argLoader.getOffset() - (self.states.numSgprPreload*4), self.externalArgLoader.getOffset()))) else: moduleWg.add(SWaitCnt(kmcnt=0, comment="wait for %u bytes of kern args" % \ (self.argLoader.getOffset() - (self.states.numSgprPreload*4)))) if kernel["ExpertSchedulingMode"] > 0 and kernel["ESMRuntimeGate"]: moduleWg.add(VMovB32(dst=vgpr(self.states.esmRuntimeFlagVgpr), src=sgpr(self.states.esmRuntimeFlagSgpr), comment="move ESM runtime flag sgpr -> vgpr")) self.sgprPool.checkIn(self.states.esmRuntimeFlagSgpr) moduleWg.addModuleAsFlatItems(moduleScaleAB) def calculateWG(): #### calculate numWorkGroup #### qReg = self.vgprPool.checkOut(4, tag="calculateWG_qReg") dReg = qReg + 1 divReg = qReg + 2 rReg = qReg + 3 moduleWg.add(VMovB32(dst=vgpr(divReg), src="MT0", comment="set MT0 into sgpr")) moduleWg.add(VMovB32(dst=vgpr(dReg), src=sgpr("SizesFree+0"), comment="set Free0 size")) moduleWg.add(vectorUInt32CeilDivideAndRemainder(qReg=qReg, dReg=dReg, divReg=divReg, rReg=rReg, doRemainder=False)) moduleWg.add(VMovB32(dst=vgpr(divReg), src="MT1", comment="set MT1 into sgpr")) moduleWg.add(VMovB32(dst=vgpr(dReg), src=sgpr("SizesFree+1"), comment="set Free1 size")) moduleWg.add(VReadfirstlaneB32(dst=sgpr("NumWorkGroups0"), src=vgpr(qReg), comment="set back to numWorkGroup0")) moduleWg.add(vectorUInt32CeilDivideAndRemainder(qReg=qReg, dReg=dReg, divReg=divReg, rReg=rReg, doRemainder=False)) if self.states.archCaps["TransOpWait"]: moduleWg.add(SNop(waitState=0, comment="1 wait states")) moduleWg.add(VReadfirstlaneB32(dst=sgpr("NumWorkGroups1"), src=vgpr(qReg), comment="set back to numWorkGroup1")) self.vgprPool.checkIn(qReg) if self.states.numSgprPreload > 0: calculateWG() waitForArgsToLoad() else: waitForArgsToLoad() calculateWG() if not kernel["ProblemType"]["StridedBatched"]: with self.allocTmpSgpr(self.states.laneSGPRCount, tag="calculateWG_tmpSgpr") as tmpSgpr: moduleWg.add(self.loadBatchedAddress(kernel, "WorkGroup2", tmpSgpr)) moduleWg.add(SWaitCnt(kmcnt=0, comment="wait global buffer address ready")) elif waitForScaleAB: moduleWg.add(SWaitCnt(kmcnt=0, comment="wait for scaleA/B to load")) labelMultiGemm = Label(label="MultiGemm", comment="") labelMultiGemmEnd = Label(label="MultiGemmEnd", comment="") module.add(moduleArgs) module.add(moduleRegInit) # Reorder WGIDs module.add(wgmXCC(self, kernel, tmpSgprNumWorkGroups)) self.sgprPool.checkIn(tmpSgprNumWorkGroups) tmpSgprNumWorkGroups = None # General Batched GEMM Routed to Strided Batched GEMM path ArgType3_Routed_To_ArgType0 = Label(label="ArgType3_Routed_To_ArgType0", comment="") module.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=3)) module.add(SCBranchSCC1(labelName=ArgType3_Routed_To_ArgType0.getLabelName())) module.add(SCmpEQU32(src0=sgpr(sgprArgType), src1=0)) self.sgprPool.checkIn(sgprArgType) sgprArgType = None # Cannot be used after this point module.add(SCBranchSCC0(labelName=labelMultiGemm.getLabelName())) module.add(ArgType3_Routed_To_ArgType0) module.add(deepcopy(moduleWg)) if kernel["StreamK"] == 0: if clusterEnabled(kernel["ClusterDim"]): module.add(SBranch(labelName=labelMultiGemmEnd.getLabelName(), comment="Already using 3D WorkGroups, skip remap")) module.add(self.remapWgSerial(kernel, earlyStop=False)) module.add(SBranch(labelName=labelMultiGemmEnd.getLabelName())) module.add(labelMultiGemm) # Return the sgprs cause after this point preloads ends for i in self.states.preloadGuard: self.sgprPool.checkIn(i) self.states.preloadGuard = [] numStoreSgprToLoad = self.states.numStoreSgprToLoad if kernel["ProblemType"]["UseScaleAB"] == "Scalar": if self.states.preloadScaleA: numStoreSgprToLoad += 2 if self.states.preloadScaleB: numStoreSgprToLoad += 2 ###### GroupedGemm ############ ###### # linear search ###### with self.allocTmpSgpr(8, 2, tag="defineAndResources_tmpSgpr") as tmpSgpr: tmpSgprM = self.sgprs["SizesFree"] tmpSgprN = tmpSgprM+1 tmpSgprB = tmpSgprN+1 tmpSgprArgAddress0 = tmpSgpr.idx tmpSgpr0 = tmpSgpr.idx + 2 tmpSgprNumWG0 = tmpSgpr.idx + 4 tmpSgprNumWG1 = tmpSgpr.idx + 5 tmpSgprAddrM = tmpSgpr.idx + 6 tmpSgprAccumTiles = tmpSgpr.idx + 7 tmpSgprLoopCounter = self.sgprs["NumWorkGroups0"] tmpSgprArgOffsett = self.sgprs["NumWorkGroups1"] # offset KernArgAddress to address of M extValidLabel = Label(label="IsExternalValid", comment="") extValidLabelEnd = Label(label="IsExternalValidEnd", comment="") if kernel["ProblemType"]["SupportUserArgs"]: module.addComment1("Check if custom structure pointer is null") module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=2, comment="ArgType == 2 ?")) module.add(SCBranchSCC1(labelName=extValidLabel.getLabelName(), comment="branch if ArgType == 2")) if ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == 'MultipleBufferSingleKernel' or kernel["AdaptiveGemmGSUA"] == 1)): module.add(SMovB32(dst=sgpr(tmpSgprArgOffsett), src=(self.argLoader.getOffset() + (numStoreSgprToLoad * 4) + (self.states.numSgprAddressGSUSync)*4), comment="KernArgAddressOffset")) else: module.add(SMovB32(dst=sgpr(tmpSgprArgOffsett), src=(self.argLoader.getOffset() + (numStoreSgprToLoad * 4)), comment="KernArgAddressOffset")) module.add(SMulI32(dst=sgpr(tmpSgprAddrM), src0=sgpr(sgprNumsOfGemm), src1=4)) # offset wgTable module.add(SMovB64(dst=sgpr(tmpSgprArgAddress0,2), src=sgpr("KernArgAddress",2))) module.add(SBranch(extValidLabelEnd.getLabelName())) module.add(extValidLabel) module.add(SMovB32(dst=sgpr(tmpSgprArgOffsett), src=self.states.userArgsInfo.totalSize)) module.add(SMovB32(dst=sgpr(tmpSgprAddrM), src=0)) module.add(SMovB64(dst=sgpr(tmpSgprArgAddress0,2), src=sgpr("KernArgAddress",2))) module.add(extValidLabelEnd) else: module.add(SMovB32(dst=sgpr(tmpSgprArgOffsett), src=(self.argLoader.getOffset() + (numStoreSgprToLoad * 4)))) module.add(SMulI32(dst=sgpr(tmpSgprAddrM), src0=sgpr(sgprNumsOfGemm), src1=4)) # offset wgTable module.add(SMovB64(dst=sgpr(tmpSgprArgAddress0,2), src=sgpr("KernArgAddress",2))) # prefetch 1 arg load module.addComment1("Grouped Gemm:: prefetch 1 arg load") module.add(SMovB32(dst=sgpr(tmpSgprLoopCounter), src=1)) module.add(SMovB32(dst=sgpr(tmpSgprAccumTiles), src=0)) module.add(self.argLoader.loadKernArg(tmpSgprM, tmpSgprArgAddress0, sgpr(tmpSgprAddrM), dword=4)) #module.add(SCmpKEQU32(src=sgpr(sgprNumsOfGemm), simm16=1, comment="if gemm_count is 1?")) module.add(self.getSCMPKInstruction("EQU32", sgprNumsOfGemm, 1, comment="if gemm_count is 1?")) label_noLoadLoop = Label("wgTable_noLoadLoop", "") module.add(SCBranchSCC1(labelName=label_noLoadLoop.getLabelName())) # Start to search module.addComment1("Grouped Gemm:: accumulate numTiles for each gemm") module.addComment0("Grouped Gemm:: loop start") label_Loop_gemm_count = Label("Loop_GemmCount", "") module.add(label_Loop_gemm_count) module.add(SWaitCnt(kmcnt=0)) # calculate numTiles regStateRes = ContinuousRegister(idx=tmpSgpr0, size=2) module.add(scalarStaticCeilDivide(qReg=sgpr(tmpSgprNumWG0), dReg=sgpr(tmpSgprM), divisor=kernel["MacroTile0"], tmpSgprRes=regStateRes)) module.add(scalarStaticCeilDivide(qReg=sgpr(tmpSgprNumWG1), dReg=sgpr(tmpSgprN), divisor=kernel["MacroTile1"], tmpSgprRes=regStateRes)) # accumulate tiles of each gemm module.add(SMulI32(dst=sgpr(tmpSgprNumWG0), src0=sgpr(tmpSgprNumWG0), src1=sgpr(tmpSgprNumWG1))) module.add(SMulI32(dst=sgpr(tmpSgprNumWG0), src0=sgpr(tmpSgprNumWG0), src1=sgpr(tmpSgprB))) if kernel["GlobalSplitU"] != 0: module.add(SAndB32(dst=sgpr(tmpSgprNumWG1), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SMulI32(dst=sgpr(tmpSgprNumWG0), src0=sgpr(tmpSgprNumWG0), src1=sgpr(tmpSgprNumWG1))) module.add(SAddU32(dst=sgpr(tmpSgprAccumTiles), src0=sgpr(tmpSgprAccumTiles), src1=sgpr(tmpSgprNumWG0))) # check wgIndex >= AccumTiles? module.add(SCmpLtU32(src0=sgpr("WorkGroup0"), src1=sgpr(tmpSgprAccumTiles))) label_FOUND = Label("FOUND", "") module.add(SCBranchSCC1(labelName=label_FOUND.getLabelName())) gsuComponent = Component.GSU.find(self) module.add(gsuComponent.defineAndResources(self, kernel, tmpSgpr0, tmpSgprM, tmpSgprN, tmpSgprNumWG0, tmpSgprAccumTiles)) module.add(SAddU32(dst=sgpr(tmpSgprAddrM), src0=sgpr(tmpSgprAddrM), src1=sgpr(tmpSgprArgOffsett))) module.add(self.argLoader.loadKernArg(tmpSgprM, tmpSgprArgAddress0, sgpr(tmpSgprAddrM), dword=4)) module.add(SAddU32(dst=sgpr(tmpSgprLoopCounter), src0=sgpr(tmpSgprLoopCounter), src1=1)) # loop gemm count module.add(SCmpLtU32(src0=sgpr(tmpSgprLoopCounter), src1=sgpr(sgprNumsOfGemm))) module.add(SCBranchSCC1(labelName=label_Loop_gemm_count.getLabelName())) # noLoadLoop module.addComment1("Grouped Gemm:: noLoadLoop") module.add(label_noLoadLoop) module.add(SWaitCnt(kmcnt=0)) # calculate numTiles regStateRes = ContinuousRegister(idx=tmpSgpr0, size=2) module.add(scalarStaticCeilDivide(qReg=sgpr(tmpSgprNumWG0), dReg=sgpr(tmpSgprM), divisor=kernel["MacroTile0"], tmpSgprRes=regStateRes)) module.add(scalarStaticCeilDivide(qReg=sgpr(tmpSgprNumWG1), dReg=sgpr(tmpSgprN), divisor=kernel["MacroTile1"], tmpSgprRes=regStateRes)) # accumulate tiles of each gemm tmpSgprGSU = tmpSgpr.idx module.add(SMulI32(dst=sgpr(tmpSgprNumWG0), src0=sgpr(tmpSgprNumWG0), src1=sgpr(tmpSgprNumWG1))) module.add(SMulI32(dst=sgpr(tmpSgprNumWG0), src0=sgpr(tmpSgprNumWG0), src1=sgpr(tmpSgprB))) if kernel["GlobalSplitU"] != 0: module.add(SAndB32(dst=sgpr(tmpSgprGSU), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SMulI32(dst=sgpr(tmpSgprNumWG0), src0=sgpr(tmpSgprNumWG0), src1=sgpr(tmpSgprGSU))) module.add(SAddU32(dst=sgpr(tmpSgprAccumTiles), src0=sgpr(tmpSgprAccumTiles), src1=sgpr(tmpSgprNumWG0))) # gemmIndex found tmpSgprWgLeft = tmpSgpr.idx tmpSgprGemmIdxLeft = tmpSgprWgLeft + 1 module.addComment1("Grouped Gemm:: gemmIndex found") module.add(label_FOUND) module.add(SSubU32(dst=sgpr(tmpSgprGemmIdxLeft), src0=sgpr(tmpSgprLoopCounter), src1=1)) module.add(SSubU32(dst=sgpr(tmpSgprWgLeft), src0=sgpr(tmpSgprAccumTiles), src1=sgpr(tmpSgprNumWG0))) ######## # load arguments of gemm ######## module.add(SSubU32(dst=sgpr("WorkGroup0"), src0=sgpr("WorkGroup0"), src1=sgpr(tmpSgprWgLeft))) extLabel = Label(label="LoadExternalStruct", comment="") extLabelEnd = Label(label="LoadExternalStructEnd", comment="") if kernel["ProblemType"]["SupportUserArgs"]: module.addComment0("Check if custom structure pointer is null") module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=2, comment="ArgType == 2 ?")) module.add(SCBranchSCC1(labelName=extLabel.getLabelName(), comment="branch if ArgType == 2")) module.addComment1("Grouped Gemm: offset argument address to gemm") module.addComment0("Grouped Gemm: offset address from wg_table_start to args_start") module.add(SLShiftLeft2AddU32(dst=sgpr("KernArgAddress"), src0=sgpr(sgprNumsOfGemm), src1=sgpr("KernArgAddress"))) module.add(SAddCU32(dst=sgpr("KernArgAddress+1"), src0=sgpr("KernArgAddress+1"), src1=0)) module.addComment0("Grouped Gemm: offset address from args_start to gemm_start") if ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == 'MultipleBufferSingleKernel' or kernel["AdaptiveGemmGSUA"] == 1)): module.add(SMulI32(dst=sgpr(tmpSgprGemmIdxLeft), src0=sgpr(tmpSgprGemmIdxLeft),\ src1=(self.argLoader.getOffset() + (numStoreSgprToLoad * 4) + (self.states.numSgprAddressGSUSync)*4), comment="KernArgAddressOffset")) else: module.add(SMulI32(dst=sgpr(tmpSgprGemmIdxLeft), src0=sgpr(tmpSgprGemmIdxLeft),\ src1=(self.argLoader.getOffset() + (numStoreSgprToLoad * 4)), comment="KernArgAddressOffset")) module.add(SAddU32(dst=sgpr("KernArgAddress"), src0=sgpr("KernArgAddress"), src1=sgpr(tmpSgprGemmIdxLeft))) module.add(SAddCU32(dst=sgpr("KernArgAddress+1"), src0=sgpr("KernArgAddress+1"), src1=0)) module.add(self.getKernelArgLoadModule(kernel, sgprStart, load, 4)) if kernel["ProblemType"]["SupportUserArgs"]: module.add(SBranch(extLabelEnd.getLabelName())) module.add(extLabel) module.addComment0("Grouped Gemm: offset address from args_start to gemm_start") # Currently a magic number cause the structure is fixed, should the structure gen by python so we can know the size? module.add(SMulI32(dst=sgpr(tmpSgprGemmIdxLeft), src0=sgpr(tmpSgprGemmIdxLeft),src1=self.states.userArgsInfo.totalSize)) module.add(SAddU32(dst=sgpr("KernArgAddress"), src0=sgpr("KernArgAddress"), src1=sgpr(tmpSgprGemmIdxLeft))) module.add(SAddCU32(dst=sgpr("KernArgAddress+1"), src0=sgpr("KernArgAddress+1"), src1=0)) moduleExternalArgs = Module("Load external Arguments") # Here alpha and beta in user args are fixed sizes, so we need to exclude beta and read it with a different offset load = load - self.states.numSgprBeta moduleExternalArgs.addModuleAsFlatItems(self.externalArgLoader.loadAllKernArg(sgprStart, "KernArgAddress", load, 4)) offset = self.externalArgLoader.getOffset() + self.states.bpr * (self.states.userArgsInfo.alphaMaxRegisterSize - self.states.numSgprAlpha) self.externalArgLoader.setOffset(offset) if kernel["ProblemType"]["UseBeta"]: moduleExternalArgs.addComment("Read Beta") moduleExternalArgs.addModuleAsFlatItems(self.externalArgLoader.loadAllKernArg(self.sgprs["Beta"], "KernArgAddress", self.states.numSgprBeta)) offset = self.externalArgLoader.getOffset() + self.states.bpr * (self.states.userArgsInfo.betaMaxRegisterSize - self.states.numSgprBeta) else: # Even when not using Beta, we need to skip over the Beta argument space offset = self.externalArgLoader.getOffset() + self.states.bpr * self.states.userArgsInfo.betaMaxRegisterSize if kernel["ProblemType"]["UseScaleAB"] == "Scalar": sgprOffset = self.externalArgLoader.getOffset() for preloadScale, name in zip([self.states.preloadScaleA, self.states.preloadScaleB], ['A','B']): if preloadScale: moduleExternalArgs.add(self.externalArgLoader.loadKernArg("AddressScale%s"%name, "KernArgAddress", sgprOffset=hex(sgprOffset), dword=2)) sgprOffset += self.states.userArgsInfo.scaleASize self.externalArgLoader.setOffset(offset) module.add(moduleExternalArgs) module.add(extLabelEnd) # Update label labels = [] for inst in moduleWg.items(): if isinstance(inst, Label): self.labels.getNameInc(inst.label) labels.append([inst.label, self.labels.getNameInc(inst.label)]) inst.label = labels[-1][1] for inst in moduleWg.items(): if isinstance(inst, BranchInstruction): removeIdx = -1 for idx, label in enumerate(labels): if inst.labelName == Label(label=label[0],comment="").getLabelName(): inst.labelName = Label(label=label[1],comment="").getLabelName() removeIdx = idx if removeIdx != -1: del labels[removeIdx] module.add(moduleWg) earlyReturnModule = Module("Early stop if N(SizeFreeJ) == 0") earlyReturnModule.addComment1("Early stop if N(SizeFreeJ) == 0") earlyReturnModule.add(SCmpEQU32(sgpr("SizeJ"), 0)) earlyReturnLabel = Label("EarlyStop_if_N_is_0", "") noEarlyReturnLabel = Label("NoEarlyStop_N0", "") earlyReturnModule.add(SCBranchSCC0(noEarlyReturnLabel.getLabelName())) earlyReturnModule.add(earlyReturnLabel) earlyReturnModule.add(SEndpgm()) earlyReturnModule.add(noEarlyReturnLabel) module.add(earlyReturnModule) if kernel["StreamK"] == 0: if clusterEnabled(kernel["ClusterDim"]): module.add(SBranch(labelName=labelMultiGemmEnd.getLabelName(), comment="Already using 3D WorkGroups, skip remap")) module.add(self.remapWgSerial(kernel)) module.addSpaceLine() module.add(labelMultiGemmEnd) # CheckIn temp sgprs if sgprNumsOfGemm: self.sgprPool.checkIn(sgprNumsOfGemm) sgprNumsOfGemm = None # gfx1250 moves SK constants to VGPRs and fully frees their SGPR slots # before defineVariableSgprs so those slots can be reused. if kernel["StreamK"] and self.isStreamKConstantsToVgprEnabled(kernel): module.add(self.moveStreamKConstantsToVgpr(kernel)) # define the rest of sgprs module.addModuleAsFlatItems(self.defineVariableSgprs(kernel)) needPackK16 = False needPackK8Lw = False needPackK8Hi = False if kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16(): if self.states.lrvwTileA > 1 or self.states.lrvwTileB > 1: needPackK16 = True if self.states.lrvwTileMetadata > 1: needPackK8Lw = True if self.states.lrvwTileMetadata > 2: needPackK8Hi = True elif kernel["ProblemType"]["DataType"].isInt8() or kernel["ProblemType"]["DataType"].is8bitFloat(): if self.states.lrvwTileA > 1 or self.states.lrvwTileB > 1 or self.states.lrvwTileMetadata > 1: needPackK8Lw = True if self.states.lrvwTileA > 2 or self.states.lrvwTileB > 2 or self.states.lrvwTileMetadata > 2: needPackK8Hi = True # need extra sgpr for packing metadata tPackM = "M" if needPackK16 and needPackK8Lw else "" if needPackK16: module.add(SMovB32(dst=sgpr("PackKForV0"), src="0x05040100", comment="")) module.add(SMovB32(dst=sgpr("PackKForV1"), src="0x07060302", comment="")) if needPackK8Lw: module.add(SMovB32(dst=sgpr("PackKFor%sV0"%tPackM), src="0x0c0c0400", comment="")) module.add(SMovB32(dst=sgpr("PackKFor%sV1"%tPackM), src="0x0c0c0501", comment="")) if needPackK8Hi: module.add(SMovB32(dst=sgpr("PackKFor%sV2"%tPackM), src="0x0c0c0602", comment="")) module.add(SMovB32(dst=sgpr("PackKFor%sV3"%tPackM), src="0x0c0c0703", comment="")) # self.states.groOffsetInMacroTile == 1 case, subtract pre-pad here if self.states.groOffsetInMacroTile and not kernel["UseSubtileImpl"]: # Added logic to check for Pointer Array case (ArgType==3) and not prepad the double pointer addresses Skip_Address_Prepad_For_Pointer_Array = Label(label="Skip_Address_Prepad_For_Pointer_Array", comment="Skip pre-padding of address for pointer array case") if kernel["ProblemType"]["SupportUserArgs"]: module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=3, comment="ArgType == 3 for General Batched GEMM")) module.add(SCBranchSCC1(labelName=Skip_Address_Prepad_For_Pointer_Array.getLabelName())) if not kernel["enableTDMA"]: prePad = int(self.states.srdShiftLeft["A"] * tPA["bpeGR"]) # leave room in case we have to pointer shift module.add(SSubU32(dst=sgpr("AddressA+0"), src0=sgpr("AddressA+0"), src1=prePad, comment="pre-pad to make room for possible pointer shift")) module.add(SSubBU32(dst=sgpr("AddressA+1"), src0=sgpr("AddressA+1"), src1=0, comment="pre-pad to make room for possible pointer shift")) if not kernel["enableTDMB"]: prePad = int(self.states.srdShiftLeft["B"] * tPB["bpeGR"]) # leave room in case we have to pointer shift module.add(SSubU32(dst=sgpr("AddressB+0"), src0=sgpr("AddressB+0"), src1=prePad, comment="pre-pad to make room for possible pointer shift")) module.add(SSubBU32(dst=sgpr("AddressB+1"), src0=sgpr("AddressB+1"), src1=0, comment="pre-pad to make room for possible pointer shift")) if not kernel["enableTDMA"] and kernel["ProblemType"]["MXBlockA"]: prePad = int(self.states.srdShiftLeft["MXSA"]) # leave room in case we have to pointer shift module.add(SSubU32(dst=sgpr("AddressMXSA+0"), src0=sgpr("AddressMXSA+0"), src1=prePad, comment="pre-pad to make room for possible pointer shift")) module.add(SSubBU32(dst=sgpr("AddressMXSA+1"), src0=sgpr("AddressMXSA+1"), src1=0, comment="pre-pad to make room for possible pointer shift")) if not kernel["enableTDMB"] and kernel["ProblemType"]["MXBlockB"]: prePad = int(self.states.srdShiftLeft["MXSB"]) # leave room in case we have to pointer shift module.add(SSubU32(dst=sgpr("AddressMXSB+0"), src0=sgpr("AddressMXSB+0"), src1=prePad, comment="pre-pad to make room for possible pointer shift")) module.add(SSubBU32(dst=sgpr("AddressMXSB+1"), src0=sgpr("AddressMXSB+1"), src1=0, comment="pre-pad to make room for possible pointer shift")) if not kernel["enableTDMMetadata"] and kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: prePad = int(self.states.srdShiftLeft["Metadata"] * tPM["bpe"]) # leave room in case we have to pointer shift module.add(SSubU32(dst=sgpr("AddressMetadata+0"), src0=sgpr("AddressMetadata+0"), src1=prePad, comment="pre-pad to make room for possible pointer shift")) module.add(SSubBU32(dst=sgpr("AddressMetadata+1"), src0=sgpr("AddressMetadata+1"), src1=0, comment="pre-pad to make room for possible pointer shift")) module.add(Skip_Address_Prepad_For_Pointer_Array) # Check alpha == 0, is done before kernel body # so if alpha/beta=Half, they haven't been converted to f32 # This means we can use ComputeDataType as AlphaType (even +"HPA") if self.do["ApplyAlpha"]: if self.states.useBias == DataDirection.WRITE and (kernel["ProblemType"]["BiasSrc"] == "A" or kernel["ProblemType"]["BiasSrc"] == "B"): # Temporarily turn off "Set summation dim=0 if Alpha == 0" cause reduction. pass else: module.addComment1("Short circuit condition if Alpha == 0, then sumDims=0") endCheckLabel = Label("AlphaNonZero", "") module.add(BranchIfNotZero("Alpha", kernel["ProblemType"]["ComputeDataType"].toEnum(), endCheckLabel)) # Conditional set summation dimensions to 0 on SCC==1 for i in range(0, self.states.numSgprSizesSum): module.add(SMovB32(dst=sgpr("SizesSum+%u"%(i)), src=0, comment="Set summation dim=0 if Alpha == 0")) # Short circuit for stream-k is handled in the stream-k component by skipping partial tiles and setting loop counter to 0 # Jump here if alpha is non-zero module.add(endCheckLabel) if kernel["MagicDivAlg"]==2: for idxChar in sorted(set(kernel["PackedC0IdxChars"][:-1] + kernel["PackedC1IdxChars"][:-1])): module.add(SLShiftRightB32(dst=sgpr("MagicAbitSize%s"%idxChar), src=sgpr("MagicShiftSize%s"%idxChar), shiftHex=31, comment="extract abit")) module.add(SAndB32(dst=sgpr("MagicShiftSize%s"%idxChar), src0=sgpr("MagicShiftSize%s"%idxChar), src1=hex(0x7fffffff), comment="remove abit")) ######################################## # Debug Buffer if self.debugConfig.debugKernel: module.addComment1("Debug Buffer") # nwg0 FIXME use NumWorkGroups0 nwg0 = self.vgprPool.checkOut(1, tag="debugBuffer_nwg0") tmpVgpr = self.vgprPool.checkOutAligned(2, 2, tag="debugBuffer_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) module.addComment("nwg0 = (size%s + MT%s - 1) / MT%s;" \ % (self.states.tileChar0, self.states.tileChar0, self.states.tileChar0)) module.add(VMovB32(dst=vgpr(tmpVgpr), src=hex(kernel["MacroTile0"]-1), comment="MT0-1")) module.add(VAddCOU32(dst=vgpr(nwg0), dst1=VCC(), src0=sgpr("SizesFree+0"), \ src1=vgpr(tmpVgpr), comment="%s = size0+MT0-1"%vgpr(nwg0))) module.add(vectorStaticDivide(nwg0, nwg0, kernel["MacroTile0"], tmpVgprRes)) self.vgprPool.checkIn(tmpVgpr) self.nipt = 16 # num integers per thread v = self.vgprPool.checkOut(3, tag="debugBuffer_v") module.add(VMovB32(dst=vgpr(v), src=sgpr("WorkGroup0"), comment="%s=wg0"%vgpr(v) )) module.add(VMovB32(dst=vgpr(v+1), src=sgpr("WorkGroup1"), comment="%s=wg1"%vgpr(v+1) )) module.add(VMulLOU32(dst=vgpr(v+1), src0=vgpr(v+1), src1=vgpr(nwg0), \ comment="%s=wg1*nwg0"%vgpr(v+1) )) module.add(VAddCOU32(dst=vgpr(v), dst1=VCC(), src0=vgpr(v), src1=vgpr(v+1), \ comment="%s=wg1*nwg0+wg0"%vgpr(v) )) with self.allocTmpSgpr(1, tag="debugKernel_tmpSgprInfo") as tmpSgprInfo: module.add(vectorStaticMultiply(vgpr(v), vgpr(v), kernel["NumThreads"], tmpSgprInfo)) module.add(VAddCOU32(dst=vgpr(v), dst1=VCC(), src0=vgpr(v), src1=vgpr("Serial"), \ comment="%s=tid+NT*(wg1*nwg0+wg0)=serial"%vgpr(v) )) module.add(VMulLOU32(dst=vgpr(v), src0=hex(self.nipt*4), src1=vgpr(v), \ comment="%s=serial*nipt*4"%vgpr(v) )) module.add(VMovB32(dst=vgpr(v+1), src=0)) module.add(VAddCOU32(dst=vgpr("AddressDbg"), dst1=VCC(), src0=sgpr("AddressDbg"), src1=vgpr(v), \ comment="%s=AddrD* + serial*nipt*4"%vgpr("AddressDbg") )) module.add(VMovB32(dst=vgpr(v+2), src=sgpr("AddressDbg+1"), comment="%s=AddressD1"%vgpr(v+2) )) module.add(VAddCCOU32(dst=vgpr("AddressDbg+1"), dst1=VCC(), src0=vgpr(v+2), \ src1=vgpr(v+1), src2=VCC(), comment="%s=AddrD* + serial*nipt*4"%vgpr("AddressDbg") )) module.add(SMovB32(dst=sgpr("DebugKernelItems"), src=0)) self.vgprPool.checkIn(v) self.vgprPool.checkIn(nwg0) if self.db["InitLds"]: tmp = ContinuousRegister(idx=self.vgprPool.checkOut(2, tag="initLds_tmpVgpr"), size=2) module.add(DSInit(tmp, kernel["NumThreads"], kernel["LdsNumBytes"], self.consts.initLdsValue)) self.vgprPool.checkIn(tmp.idx) return module def extractPackedCoord1ToRowStart(self, kernel, packedC1, packedCoordVgpr, storeChar): if kernel["ProblemType"]["UseE"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0): printExit("extractPackedCoord1ToRowStart doe not support with output E.") # calculate packed rowStart vgpr # vgprTmp assignments: # - tmp+0 is the incoming packed coordinate 1, used on replay too # - tmp+1 is DIV output # - tmp+2 is scratch # - tmp+3 holds thread rowStart free1 offset module = Module("extractPackedCoord1ToRowStart") tmpV0 = self.vgprPool.checkOut(4, tag="extractPackedCoord1ToRowStart_tmpV0") tmpV1 = tmpV0 + 1 tmpV2 = tmpV0 + 2 tmpV3 = tmpV0 + 3 module.add(VMovB32(dst=vgpr(tmpV0), src=vgpr(packedCoordVgpr), comment="copy coord1 then unpack")) for i,idx in enumerate(packedC1[:-1]): idxChar= INDEX_CHARS[idx] module.addComment0("extract %s"%self.sizeRef(idx)) if self.states.asmCaps["HasVgprMSB"]: module.add(VMagicDiv(kernel["MagicDivAlg"], tmpV1, vgpr(tmpV0), sgpr("MagicNumberSize%s"%idxChar), \ sgpr("MagicShiftSize%s"%idxChar), sgpr("MagicAbitSize%s"%idxChar))) else: module.add(MacroInstruction(name="V_MAGIC_DIV", \ args=[tmpV1, vgpr(tmpV0), sgpr("MagicNumberSize%s"%idxChar), \ sgpr("MagicShiftSize%s"%idxChar), (sgpr("MagicAbitSize%s"%idxChar) if kernel["MagicDivAlg"]==2 else "0")])) module.add(VMulLOU32(dst=vgpr(tmpV2), src0=vgpr(tmpV1), src1=self.sizeRef(idx), comment="remainder part 1")) module.add(VSubU32(dst=vgpr(tmpV2), src0=vgpr(tmpV0), src1=vgpr(tmpV2), comment="remainder part 2")) if i==0: module.add(VMulLOU32(dst=vgpr(tmpV3), src0=vgpr(tmpV2), \ src1=self.strideRef(storeChar, idx), comment="addrCalc <- scaled extracted dim")) else: module.add(VMulLOU32(dst=vgpr(tmpV2), src0=vgpr(tmpV2), \ src1=self.strideRef(storeChar, idx), comment="scale extracted dim")) module.add(VAddU32(dst=vgpr(tmpV3), src0=vgpr(tmpV3), \ src1=vgpr(tmpV2), comment="addrCalc += scaled extracted dim ")) if i < len(packedC1)-2: module.add(VMovB32(dst=vgpr(tmpV0), src=vgpr(tmpV1), \ comment="Copy remaining bits for next divide")) module.addComment0("extract final %s"%self.sizeRef(packedC1[-1])) module.add(VMulLOU32(dst=vgpr(tmpV2), src0=vgpr(tmpV1), \ src1=self.strideRef(storeChar, packedC1[-1]), comment="scale final extracted dim")) module.add(VAddU32(dst=vgpr(self.vgprs.coutRowPtrD), src0=vgpr(tmpV3), \ src1=vgpr(tmpV2), comment="rowStart += scaled extracted dim ")) self.vgprPool.checkIn(tmpV0) return module ############################################################################## # Global Read Addresses: WorkGroup ############################################################################## def graWorkGroup(self, kernel, tPA, tPB): module = Module("graWorkGroup") module.addComment0("graWorkGroup mapping") skComponent = Component.StreamK.find(self) module.add(skComponent.graWorkGroup(self, kernel, tPA, tPB)) gsuComponent = Component.GSU.find(self) module.add(gsuComponent.graWorkGroup(self, kernel)) # Start of WGM algos sgprWGM = "WGM" if len(kernel["SpaceFillingAlgo"]): module.addComment0("Start of Generic WGM algo") self.states.WGMTransformLevels = len(kernel["SpaceFillingAlgo"]) module.add(SpaceFillingCurveWalk(self, kernel, sgprWGM)) else: module.add(DefaultWGM(self, kernel, sgprWGM)) return module def graMetadataTileAssignment(self, kernel, tP): module = Module("graMetadataTileAssignment") module.addComment0("calculate metadata gra tile assignment") if kernel["DirectToVgprSparseMetadata"]: # alloc vgpr wReg = self.vgprPool.checkOut(1,"wReg") # quotient tReg = self.vgprPool.checkOut(1,"tReg") # remainder tmpVgpr = self.vgprPool.checkOutAligned(2,2,"tmpVgpr") dummy = self.vgprPool.checkOut(1,"dummy") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) # get constant parameter waveWidth = kernel["WavefrontSize"] # parameter for get each type index num1DBlocks = kernel["MatrixInstBM"] num1DWaves = kernel["MIWaveGroup"][0] dividedForWaveId = waveWidth # strider for each type of index offsetWave = kernel["MatrixInstM"] * num1DBlocks with self.allocTmpSgpr(1, tag="graMetadataTileAssignment_tmpSgprInfo") as tmpSgprInfo: # tile offset module.add(vectorStaticRemainder(dummy, tReg, "Serial", waveWidth, tmpVgprRes, tmpSgprInfo, \ comment="0. thread id in wave: wtid = tid %% wavelength(%u)" % waveWidth)) module.add(vectorStaticRemainder(dummy, tReg, tReg, kernel["MatrixInstM"], tmpVgprRes, tmpSgprInfo, \ comment="1. tile offset: nIdx = wtid %% MI_N(%u)" % kernel["MatrixInstM"])) # wave offset if num1DWaves > 1: module.add(vectorStaticDivide(wReg, "Serial", dividedForWaveId, tmpVgprRes, \ "2. wave offset in N dimen: wtid = tid / dividedForWaveId(%u)" % dividedForWaveId)) module.add(vectorStaticRemainder(dummy, wReg, wReg, num1DWaves, tmpVgprRes, tmpSgprInfo, \ "2. wave offset in M dimen: wtid0 = wtid / num1DWaves(%u)" % num1DWaves)) module.add(vectorStaticMultiplyAdd(vgpr("GlobalReadOffsetMetadata"), vgpr(wReg), offsetWave, vgpr(tReg), tmpSgprInfo, \ "2. wave offset in M dimen: wOffset = wtid0 * W0 offset(%u); 2. tile coord = tileOffset + wOffset" % offsetWave)) else: module.add(VMovB32(vgpr("GlobalReadOffsetMetadata"), vgpr(tReg), \ comment="2. tile coord = tileOffset")) for graIdx in range (1, kernel["MIWaveTile"][0]): MIWaveGroup0ShapeSize = kernel["MatrixInstM"] * kernel["MatrixInstBM"] * kernel["MIWaveGroup"][0] vgprGro = "GlobalReadOffsetMetadata+%u"%(graIdx) module.add(VMovB32(vgpr(vgprGro), hex(MIWaveGroup0ShapeSize*graIdx), comment="7. coord offset of WaveTile %u"%graIdx)) module.add(VAddU32(vgpr(vgprGro), vgpr(vgprGro), vgpr("GlobalReadOffsetMetadata"), \ "3. final coord0: tile coord + waveTile coord")) # release register self.vgprPool.checkIn(tReg) self.vgprPool.checkIn(wReg) self.vgprPool.checkIn(tmpVgpr) self.vgprPool.checkIn(dummy) return module ############################################################################## # Global Read Addresses: Tile Assignment A/B # global read addresses: tile offset assignment (message from .s) ############################################################################## def graTileAssignment(self, kernel, tP): module = Module("graTileAssignment") tc = tP["tensorChar"] tReg = tP["gpr"]["lwoT"] module.addComment0("graTileAssignment%s = %s" % (tc, vgpr(tReg))) if self.states.groOffsetInMacroTile: tReg2 = tReg # treg2 and treg same register and value - we store the 'static' # part of the address calculation in the SRD to maximize the # range of the 32-bit GRO else: tReg2 = self.vgprPool.checkOut(1, 'treg2', self.states.preventVgprOverflowDuringNewTile) with self.allocTmpSgpr(1, tag="graTileAssignment_tmpSgprInfo") as tmpSgprInfo: if not self.states.groOffsetInMacroTile: tmpVgpr = self.vgprPool.checkOut(1, 'graTA vgpr', self.states.preventVgprOverflowDuringNewTile) # Buffer Load will set the SRD to start of the MacroTile # So don't add the static wg-related component here - save for later. module.add(vectorStaticMultiply(vgpr(tmpVgpr), sgpr(tP["wg"]), kernel[tP["mt"]], tmpSgprInfo)) # workgroup module.add(VAddCOU32(dst=vgpr(tReg2), dst1=VCC(), src0=vgpr(tmpVgpr), \ src1=vgpr(tReg), comment="gro%s-tile = serialTile*VW + (wg%s*MT%s)" \ % (tc, tc, tc) )) self.vgprPool.checkIn(tmpVgpr) tP["gpr"]["tReg"] = tReg2 return Module("graTileAssignment (Empty)") if self.dontAppendCode else module ############################################################################## # Global Read Addresses: Unroll Assignment ############################################################################## def graUnrollAssignment(self, kernel, tP): module = Module("graUnrollAssignment") # note groOffsetInMacroTile rolls these into SRD so don't change here: if not self.states.groOffsetInMacroTile: gsuOffset = self.vgprPool.checkOut(1, "gsuOffset", self.states.preventVgprOverflowDuringNewTile) module.add(VMovB32(dst=vgpr(gsuOffset), src=sgpr("GSUSumIdx"), comment="=gsuSumIdx")) with self.allocTmpSgpr(1, tag="graUnrollAssignment_tmpSgprInfo") as tmpSgprInfo: # graUnrollAssignment += gsuSumIdx*DepthU module.add(vectorStaticMultiply(vgpr(gsuOffset), vgpr(gsuOffset), kernel["DepthU"], tmpSgprInfo)) module.add(VAddCOU32(dst=vgpr(tP["gpr"]["uReg"]), dst1=VCC(), \ src0=vgpr(gsuOffset), src1=vgpr(tP["gpr"]["uReg"]), \ comment="graUnrollAssignment += gsuOffset")) self.vgprPool.checkIn(gsuOffset) else: module.addComment0(str(vgpr(tP["gpr"]["uReg"]))) return Module("graUnrollAssignment (Empty)") if self.dontAppendCode else module ############################################################################## # Global Read Addresses: Other Free Assignments ############################################################################## def graOtherFreeAssignments(self): module = Module("graOtherFreeAssignments") module.addComment0("s[sgprWorkGroup2]") return module ############################################################################## # Global Read Addresses: Other Summation Assignments ############################################################################## def graOtherSummationAssignments(self, kernel): module = Module("graOtherSummationAssignments") for i in range(0,kernel["ProblemType"]["NumIndicesSummation"]-1): index = i module.add(ValueSet(name="globalReadOffsetA%s" % self.states.indexChars[index], value=0)) module.add(ValueSet(name="globalReadOffsetB%s" % self.states.indexChars[index], value=0)) return module def alignTo(self, dstBase: Union[int, str], srcBase: Union[int, str], alignment: int) -> Module: assert (alignment & (alignment - 1)) == 0 and f"Alignment must be power of two" module = Module() module.addComment(f"Align to {alignment}") module.add(SAddU32(sgpr(dstBase), sgpr(srcBase), alignment-1)) module.add(SLShiftRightB32(dst=sgpr(dstBase), src=sgpr(dstBase), shiftHex=log2(alignment))) module.add(SLShiftLeftB32(dst=sgpr(dstBase), src=sgpr(dstBase), shiftHex=log2(alignment))) return module def setTo(self, dstBase: Union[int, str], srcBase: Union[int, str]) -> Module: module = Module() module.add(SMovB32(dst=sgpr(dstBase), src=sgpr(srcBase), comment="Set to %s"%(srcBase))) return module ############################################################################## # Global Read Addresses: Tile Offsets A/B ############################################################################## def graTileOffsets(self, kernel, tP, margin=-1): module = Module("graTileOffsets") tc = tP["tensorChar"] tP["vgprPackedOffsets"] = None tP["vgprTileOffsetsCheckOut"] = False tP["numVgprTileOffsets"] = 0 if kernel["_UseSgprForGRO"]: # Let the vgprTileOffsets checkin handle tReg later since these are same vgpr tP["vgprTileOffsets"] = tP["gpr"]["tReg"] else: numTileOffsets = tP["nrt"] tP["vgprTileOffsets"] = self.vgprPool.checkOut(numTileOffsets, "vgprTileOffsets", self.states.preventVgprOverflowDuringNewTile) tP["vgprTileOffsetsCheckOut"] = True v = tP["vgprTileOffsets"] numExtraPackedOffsetsPerTile = len(tP["PackedIndices"])-1 if numExtraPackedOffsetsPerTile: tP["vgprPackedOffsets"] = self.vgprPool.checkOut(numExtraPackedOffsetsPerTile * numTileOffsets, "vgprPackedOffsets", self.states.preventVgprOverflowDuringNewTile) strideIdx = tP["lsc"] if tP["tlu"] else tP["lsp"] stride = kernel[strideIdx] # adjustment for DirectToVgpr + tlu=False + VW > 1 case strideInterleave = False if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%c"%tc] and (not tP["tlu"]) and kernel["VectorWidth%s"%tc] > 1: strideInterleave = True stride = stride * kernel["VectorWidth%s"%tc] - (kernel["VectorWidth%s"%tc] - 1) strideMask = (kernel["VectorWidth%s"%tc] - 1) if tP["isM"] and not margin == -1: # margin is the number of how many continuous element need to be read # shift metadata global read offset to align A's global read offset. module.add(VLShiftRightB32(dst=vgpr(v), shiftHex=hex(log2(margin)), src=vgpr(tP["gpr"]["tReg"]), comment="gro%s%s_%u /= %d"%(tP["tensorChar"], tP["tileChar"], 0, margin))) module.add(VLShiftLeftB32(dst=vgpr(v), shiftHex=hex(log2(margin)), src=vgpr(v), comment="gro%s%s_%u *= %d"%(tP["tensorChar"], tP["tileChar"], 0, margin))) else: if not tP["isSwizzled"]: module.add(VMovB32(dst=vgpr(v), src=vgpr(tP["gpr"]["tReg"]), comment="gro%s%s_%u"%(tP["tensorChar"], tP["tileChar"], 0) )) else: lsu = kernel["LocalSplitU"] # localSplitU if tP["isA"]: WvG_MorN = kernel["MIWaveGroup"][0] swzMorN = kernel["MatrixInstM"] elif tP["isB"]: WvG_MorN = kernel["MIWaveGroup"][1] swzMorN = kernel["MatrixInstN"] swzStride = tP["swizzleK"] swzBlockSize = swzMorN * swzStride vw = kernel[f"VectorWidth{tc}"] kPack = tP["swizzlePackK"] laneSize = int(kernel["MatrixInstK"] / 4) * kPack # the size of one swizzle's lane with self.allocTmpSgpr(2, tag="graTileOffsets_tmpSgprInfo") as tmpSgprInfo: swzBlkVWSizeSgpr = tmpSgprInfo.idx swzStridePerWave = swzBlkVWSizeSgpr + 1 module.addComment0("SWZ-%s: calculation of gro%s%s_%u"%(tc, tP["tensorChar"], tP["tileChar"], 0)) module.add(VMovB32(dst=vgpr(v), src=vgpr(tP["gpr"]["tReg"]))) swzBlkVWSizeVgpr = self.vgprPool.checkOut(1, tag="graTileOffsets_swzBlkVWSizeVgpr") module.add(SMovB32(dst=sgpr(swzBlkVWSizeSgpr), src=hex(swzBlockSize * vw), \ comment="swizzled block * VW = (MI_MorN(%u) * MI_K(%u) * kPack) * VW(%u)" %(swzMorN, int(swzStride / kPack), vw))) module.add(VMovB32(dst=vgpr(swzBlkVWSizeVgpr), src=sgpr(swzBlkVWSizeSgpr))) module.add(VMulU32U24(dst=vgpr(v), src0=vgpr(v), src1=vgpr(swzBlkVWSizeVgpr))) # LSU part if lsu > 1: tmpVgprRes = None wave_id = self.vgprPool.checkOut(1, tag="graTileOffsets_wave_id") # quotient # constant du = kernel["DepthU"] lsuStride = du // lsu numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] # codes module.add(vectorStaticDivide(wave_id, "Serial", kernel["WavefrontSize"] * numWaves, tmpVgprRes, comment="LSU offset: Get LSU wave_id")) module.add(VLShiftLeftB32(dst=vgpr(wave_id), shiftHex=hex(log2(lsuStride*swzMorN)), src=vgpr(wave_id), comment="LSU offset: LSU_wave_id*MI_MorN(%u)*lsuStrideK(%u)"%(swzMorN, lsuStride))) module.add(VAddU32(dst=vgpr(v), src0=vgpr(wave_id), src1=vgpr(v), comment="tileOffset += LSU offset")) self.vgprPool.checkIn(wave_id) # don't emit the code if we don't need to do more than one nrt if tP["nrt"] > 1: module.addComment0("swzStridePerWave = (number of swizzle block in K) * WaveGroup_MorN") module.addComment(f"Align to {swzStride}") module.add(SAddU32(sgpr(swzStridePerWave), sgpr("SizesSum"), swzStride-1)), module.add(SLShiftRightB32(dst=sgpr(swzStridePerWave), src=sgpr(swzStridePerWave), shiftHex=hex(log2(swzStride)), comment="numKr = DimK / swizzleK")), module.add(SMulI32(dst=sgpr(swzStridePerWave), src0=hex(WvG_MorN), src1=sgpr(swzStridePerWave), comment="numKr *= MI_WaveGroup, wave-M (SWZ-A) or wave-N (SWZ-B)")) module.addComment0("swzBlkVWOffset = swzBlkWvGSize - laneSize * (VW - 1)") module.add(SMulI32(dst=sgpr(swzBlkVWSizeSgpr), src0=sgpr(swzStridePerWave), src1=sgpr(swzBlkVWSizeSgpr), \ comment="swzBlkWvGSize = numKr * (swzBlockSize * VW)")) module.add(SSubU32(dst=sgpr(swzBlkVWSizeSgpr), src0=sgpr(swzBlkVWSizeSgpr), src1=(laneSize * (vw - 1)), \ comment="swzBlkVWOffset = swzBlkWvGSize - laneSize * (VW - 1)")) module.add(VMovB32(dst=vgpr(swzBlkVWSizeVgpr), src=sgpr(swzBlkVWSizeSgpr)) ) for l in range(1, tP["nrt"]): strideValue = stride if strideInterleave and (l & strideMask) != 0: strideValue = 1 if not tP["isSwizzled"]: module.add(VAddCOU32(dst=vgpr(v+l), dst1=VCC(), src0=strideValue, \ src1=vgpr(v+l-1), comment="gro%s%s_%u += %s"%(tP["tensorChar"], tP["tileChar"], l, strideIdx) )) # swizzle else: # VW > 1 if (strideInterleave and (l & strideMask) != 0): module.add(VAddCOU32(dst=vgpr(v+l), dst1=VCC(), src0=laneSize, \ src1=vgpr(v+l-1), comment="SWZ-%s: gro%s%s_%u"%(tc, tP["tensorChar"], tP["tileChar"], l) )) # VW == 1 else: module.add(VAddCOU32(dst=vgpr(v+l), dst1=VCC(), src0=vgpr(swzBlkVWSizeVgpr), \ src1=vgpr(v+l-1), comment="SWZ-%s: gro%s%s_%u"%(tc, tP["tensorChar"], tP["tileChar"], l) )) if tP["isSwizzled"]: self.vgprPool.checkIn(swzBlkVWSizeVgpr) # TODO- check for swizzle if numExtraPackedOffsetsPerTile: tmpV = self.vgprPool.checkOutAligned(2,2,"packTmp", self.states.preventVgprOverflowDuringNewTile) for l in range(0, tP["nrt"]): lastGroVgpr = vgpr(v+l) lastGroIdx = tP["PackedIndices"][0] module.addSpaceLine() for p in range(0, numExtraPackedOffsetsPerTile): groIdx = tP["PackedIndices"][p+1] groChar = INDEX_CHARS[tP["PackedIndices"][p+1]] groVgpr = vgpr(tP["vgprPackedOffsets"] + l*numExtraPackedOffsetsPerTile + p) pChar = INDEX_CHARS[tP["PackedIndices"][p]] if self.states.asmCaps["HasVgprMSB"]: module.add(VMagicDiv(kernel["MagicDivAlg"], tmpV, lastGroVgpr, sgpr("MagicNumberSize%s"%pChar), \ sgpr("MagicShiftSize%s"%pChar), sgpr("MagicAbitSize%s"%pChar))) else: module.add(MacroInstruction(name="V_MAGIC_DIV", \ args=[tmpV, lastGroVgpr, sgpr("MagicNumberSize%s"%pChar), \ sgpr("MagicShiftSize%s"%pChar), (sgpr("MagicAbitSize%s"%pChar) if kernel["MagicDivAlg"]==2 else "0")] )) module.add(VMovB32(dst=groVgpr, src=vgpr(tmpV), comment="extract gro%s%s_%u (%s)"%(tc,groChar,l,groVgpr))) module.add(VMulLOU32(dst=vgpr(tmpV), src0=groVgpr, src1=sgpr("SizesFree+%u"%lastGroIdx), comment="remainder part 1")) module.add(VSubU32(dst=lastGroVgpr, src0=lastGroVgpr, src1=vgpr(tmpV), \ comment="remove extracted bits from gro%s%s_%u (%s)"%(tc, INDEX_CHARS[lastGroIdx], l, lastGroVgpr))) lastGroVgpr = groVgpr lastGroIdx = groIdx self.vgprPool.checkIn(tmpV) # groOffsetInMacroTile uses same register for both of these, don't free it here: if tP["gpr"]["lwoT"] != tP["gpr"]["tReg"] : self.vgprPool.checkIn(tP["gpr"]["tReg"]) tP["gpr"]["tReg"] = None return Module("graTileOffsets (Empty)") if self.dontAppendCode else module ############################################################################## # Global Read Addresses: Unroll Offsets A/B ############################################################################## def graUnrollOffsets(self, kernel, tP): module = Module("graUnrollOffsets") tc = tP["tensorChar"] if kernel["_UseSgprForGRO"]: tP["gpr"]["unrollOffsets"] = tP["gpr"]["uReg"] # swizzle elif tP["isSwizzled"]: numUnrollOffsets = tP["nru"] tP["gpr"]["unrollOffsets"] = self.vgprPool.checkOut(numUnrollOffsets, "unrollOffsets", self.states.preventVgprOverflowDuringNewTile) v = tP["gpr"]["unrollOffsets"] # Please note that the following description uses swizzle A as an example. # The notation of A and M should be replaced with B and N when referring to swizzle B. # # swizzle offset formula = F(row, col) + swzBlkOffsetM(tid) (unit: element) # row, col: local index in a swizzled block # row = (tid / swzM) * laneSize # col = (tid % (swzM / VWA)) * VWA # laneSize: number of elements loaded by a single lane. # F(row, col): swizzled offset when vector width is 1. # F(row, col) = row % laneSize + col * laneSize + (row / laneSize) * swzM * laneSize # swzBlkOffsetM(tid) = [tid / (swzM / VWA)] * swzSizeK # swzSizeK = numSwzBlkK * numElmInSwzBlk # numSwzBlkK: number of swizzled block in K dimension. # numElmInSwzBlk: number of element in a swizzled block. if tP["isA"]: swzMorN = kernel["MatrixInstM"] elif tP["isB"]: swzMorN = kernel["MatrixInstN"] swzStride = tP["swizzleK"] vw = kernel[f"VectorWidth{tc}"] kPack = tP["swizzlePackK"] laneSize = int(kernel["MatrixInstK"] / 4) * kPack # the size of one swizzle's lane numElmInSwzBlk = swzMorN * swzStride # Calculate local index in a swizzled block row = self.vgprPool.checkOut(1, tag="graUnrollOffsets_row") col = self.vgprPool.checkOut(1, tag="graUnrollOffsets_col") tmpVgpr = self.vgprPool.checkOut(1, tag="graUnrollOffsets_tmpVgpr") module.addComment0("SWZ-%s: r = swzRow = (tid / swzMorN(%u)) * laneSize(%u)"%(tc, swzMorN, laneSize)) module.add(VAndB32(dst=vgpr(v), src0=vgpr("Serial"), src1=(kernel["WavefrontSize"]-1), comment="tid")) module.add(VLShiftRightB32(dst=vgpr(row), shiftHex=hex(int(log(swzMorN, 2))), src=vgpr(v))) module.add(VLShiftLeftB32(dst=vgpr(row), shiftHex=hex(int(log(laneSize, 2))), src=vgpr(row))) module.addComment0("SWZ-%s: c = swzCol = [tid mod (swzMorN(%u) / VW(%u))] * VW(%u)"%(tc, swzMorN, vw, vw)) module.add(VAndB32(dst=vgpr(col), src0=vgpr(v), src1=int(swzMorN/vw)-1)) if vw > 1: module.add(VLShiftLeftB32(dst=vgpr(col), shiftHex=hex(int(log(vw, 2))), src=vgpr(col))) module.addComment0("SWZ-%s: F = r mod laneSize + c * laneSize + (r/laneSize) * swzMorN * laneSize"%tc) module.add(VLShiftRightB32(dst=vgpr(tmpVgpr), shiftHex=hex(int(log(laneSize, 2))), src=vgpr(row), comment="tmp = r / laneSize")) module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr), shiftHex=hex(int(log(swzMorN * laneSize, 2))), src=vgpr(tmpVgpr), comment="tmp *= swzMorN * laneSize")) module.add(VMadU32U24(dst=vgpr(tmpVgpr), src0=vgpr(col), src1=laneSize, src2=vgpr(tmpVgpr), comment="tmp += c * laneSize")) module.add(VAndB32(dst=vgpr(row), src0=vgpr(row), src1=laneSize-1, comment="tmp2 = row mod laneSize")) if vw == 1: module.add(VAddU32(dst=vgpr(v), src0=vgpr(tmpVgpr), src1=vgpr(row), comment="offset = tmp + tmp2")) else: fVgpr = self.vgprPool.checkOut(1, tag="graUnrollOffsets_fVgpr") module.add(VAddU32(dst=vgpr(fVgpr), src0=vgpr(tmpVgpr), src1=vgpr(row), comment="F = tmp + tmp2")) module.addComment0("SWZ-%s: swzSizeK = numSwzBlkK * numElmInSwzBlk"%tc) swzSizeK = self.vgprPool.checkOut(1, tag="graUnrollOffsets_swzSizeK") tmpSgpr = self.sgprPool.checkOut(1, tag="graUnrollOffsets_tmpSgpr") module.add(SAddU32(sgpr(tmpSgpr), sgpr("SizesSum"), swzStride - 1, comment="Align to %u"%swzStride)) module.add(SLShiftRightB32(dst=sgpr(tmpSgpr), src=sgpr(tmpSgpr), shiftHex=hex(log2(swzStride)), comment="numKr = DimK / %u"%swzStride)) module.add(SMulI32(dst=sgpr(tmpSgpr, 1), src0=sgpr(tmpSgpr, 1), src1=hex(numElmInSwzBlk))) module.add(VMovB32(dst=vgpr(swzSizeK), src=sgpr(tmpSgpr, 1))) self.sgprPool.checkIn(tmpSgpr) module.addComment0("SWZ-%s: swizzle offset = F + [(tid / (swzMorN(%u) / VW(%u))) mod VW] * swzSizeK"%(tc, swzMorN, vw)) module.add(VLShiftRightB32(dst=vgpr(tmpVgpr), shiftHex=hex(int(log((swzMorN/vw), 2))), src=vgpr(v), comment="tmp = tid / (swzMorN / VW)")) module.add(VAndB32(dst=vgpr(tmpVgpr), src0=vgpr(tmpVgpr), src1=vw-1, comment="tmp = tmp mod VW")) module.add(VMadU32U24(dst=vgpr(v), src0=vgpr(tmpVgpr), src1=vgpr(swzSizeK), src2=vgpr(fVgpr), comment="offset = tmp * swzSizeK + F")) self.vgprPool.checkIn(fVgpr) self.vgprPool.checkIn(swzSizeK) self.vgprPool.checkIn(row) self.vgprPool.checkIn(col) self.vgprPool.checkIn(tmpVgpr) for l in range(1, tP["nru"]): module.add(VAddCOU32(dst=vgpr(v+l), dst1=VCC(), src0=hex(numElmInSwzBlk), \ src1=vgpr(v+l-1), comment="gro%s%s_%u"%(tP["tensorChar"], self.states.unrollChar, l))) else: numUnrollOffsets = tP["nru"] tP["gpr"]["unrollOffsets"] = self.vgprPool.checkOut(numUnrollOffsets, "unrollOffsets", self.states.preventVgprOverflowDuringNewTile) v = tP["gpr"]["unrollOffsets"] strideIdx = (tP["lsp"] if tP["tlu"] else tP["lsc"]) stride = kernel[strideIdx] if (tc in ("A", "B", "MXSA", "MXSB")) and kernel["DirectToVgpr%s"%tc] and kernel["LocalSplitU"] > 1: # DTV + LSU case, we need to divide stride by LSU stride = stride // kernel["LocalSplitU"] prevStride = 0 totalStride = 0 dtvKInterval = self.states.dtvKIntervalA if tP["isA"] else self.states.dtvKIntervalB module.add(VMovB32(dst=vgpr(v), src=vgpr(tP["gpr"]["uReg"]), comment="gro%s%s_%u"%(tP["tensorChar"], self.states.unrollChar, 0))) for l in range(1, tP["nru"]): totalStride += stride if dtvKInterval > 1: # DirectToVgpr + k interval > 1 case, stride * dtvKInterval is added every dtvKInterval. # Add mod in mod != 0 case totalStride = stride * (l - (l % dtvKInterval)) + (l % dtvKInterval) currStride = totalStride - prevStride prevStride = totalStride module.add(VAddCOU32(dst=vgpr(v+l), dst1=VCC(), src0=currStride, \ src1=vgpr(v+l-1), comment="gro%s%s_%u + %s"%(tP["tensorChar"], self.states.unrollChar, l, strideIdx))) #self.vgprPool.checkIn(tP["gpr"]["uReg"]) return Module("graUnrollOffsets (Empty)") if self.dontAppendCode else module ############################################################################## # Global Read Addresses: Shift A/B # See if the load (including vw) will extend past the 'free' dim of the # tensor. If so clip to the last legal value which is inside the array ############################################################################## def graMetadataShift(self, kernel, tP): module = Module("graMetadataShift") if kernel["DirectToVgprSparseMetadata"]: margin = tP["glvw"] if tP["rtv"] else 1 # Subtract the static component from SizesFree: with self.allocTmpSgpr(2+self.states.laneSGPRCount, tag="graMetadataShift_tmpSgprInfo") as tmpSgprInfo: edgeSgpr = tmpSgprInfo.idx shiftSgpr = edgeSgpr+1 laneMaskSgpr = edgeSgpr+2 module.add(SMulI32(dst=sgpr(edgeSgpr), src0=sgpr(tP["wg"]), src1=kernel[tP["mt"]], comment="WorkGroup[01] * MT")) module.add(SSubI32(dst=sgpr(edgeSgpr), src0=self.sizeRef(tP["idx"]), src1=sgpr(edgeSgpr), comment="edge = Size%s - WG*MT"%(tP["tileChar"]))) module.add(SAndB32(dst=sgpr(shiftSgpr), src0=sgpr(edgeSgpr), src1=hex(margin-1), comment="edge size & (glvw-1)")) module.add(SSubU32(dst=sgpr(shiftSgpr), src0=hex(margin), src1=sgpr(shiftSgpr), comment="shift size = glvw - (edge size & (glvw-1))")) module.add(SAndN2B32(dst=sgpr(edgeSgpr), src0=sgpr(edgeSgpr), src1=hex(margin-1), comment="edgeCoord = edge & !(glvw-1)")) # apply shiftPointer into vgpr offset shiftedCoord = self.vgprPool.checkOut(1, "shiftedCoord", self.states.preventVgprOverflowDuringNewTile) for graIdx in range (0, kernel["MIWaveTile"][0]): vgprGro = "GlobalReadOffsetMetadata+%u"%(graIdx) # check if in shift area module.add(VCmpLeI32(dst=sgpr(laneMaskSgpr,self.states.laneSGPRCount), src0=sgpr(edgeSgpr), src1=vgpr(vgprGro), comment="edgeCoord <= coord")) # calculate shifted coord module.add(VSubI32(dst=vgpr(shiftedCoord), src0=vgpr(vgprGro), src1=sgpr(shiftSgpr), comment="shiftedCoord = coord - shift size")) # apply shift if condition is true module.add(VCndMaskB32(dst=vgpr(vgprGro), src0=vgpr(vgprGro), src1=vgpr(shiftedCoord), src2=sgpr(laneMaskSgpr,self.states.laneSGPRCount), comment="coord = (cond) ? shifted coord : ori coord")) self.vgprPool.checkIn(shiftedCoord) return module ############################################################################## def graShift(self, kernel, tP, margin=-1): # graShift requires a vgpr for each address component (so each component # can be examined and shifted if necessary) - therefore does not work # with UseSgprForGRO. assert(not kernel["_UseSgprForGRO"]), "%s"%self.states.kernelName module = Module("graShift") #tc = tP["tensorChar"] # edge value marginO = margin # for the edge case, using A's margin to instead Metadata's margin, # otherwise, loaded data of A and Metadata will not match if margin == -1: margin = tP["glvw"] if tP["rtv"] else 1 edge = self.vgprPool.checkOut(1, "edge", self.states.preventVgprOverflowDuringNewTile) with self.allocTmpSgpr(1, tag="graShift_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx if self.states.groOffsetInMacroTile: # Subtract the static component from SizesFree: module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tP["wg"]), src1=kernel[tP["mt"]], comment="WorkGroup[01] * MT")) module.add(SSubU32(dst=sgpr(tmpSgpr), src0=self.sizeRef(tP["idx"]), src1=sgpr(tmpSgpr), \ comment="edge = Size%s - WG*MT"%(tP["tileChar"]))) # use math here to use unsigned (to increase range) # - add srdShiftLeft to tmpSgpr - ensure it is always positive # - below add srdShiftLeft to a tmp copy of the offset used for the compare # edge = (Size - WG*MT) - margin = the last valid load position that won't cause OOB # offset = the current load position for this thread # so if offset is larger than edge, we go back to the edge position module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=margin, comment="edge -= margin(%u)"%(margin))) module.add(VMovB32(dst=vgpr(edge), src=sgpr(tmpSgpr), comment="edge vgpr = Size%s- WG*MT - margin(%u)"%(tP["tileChar"], margin) )) #shiftedEdge = self.vgprPool.checkOut(1, "shiftedEdge", self.states.preventVgprOverflowDuringNewTile) #module.add(VAddCOU32(dst=vgpr(shiftedEdge), dst1=VCC(), src0=vgpr(edge), src1=self.states.srdShiftLeft[tc], # comment="shiftedEdge = edge + srdShiftLeft({})".format(self.states.srdShiftLeft[tc]))) else: module.add(SSubU32(dst=sgpr(tmpSgpr), src0=self.sizeRef(tP["idx"]), src1=margin, \ comment="edge = Size%s-%u"%(tP["tileChar"], margin) )) module.add(VMovB32(dst=vgpr(edge), src=sgpr(tmpSgpr), \ comment="edge vgpr = Size%s-%u"%(tP["tileChar"], margin) )) # shift offsets vSrc = tP["vgprTileOffsets"] vDst = tP["vgprTileOffsets"] with self.allocTmpSgpr(self.states.laneSGPRCount, tag="graShift_tmpSgprInfo2") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx for l in range(0, tP["nrt"]): # compare cmpCommentText = "offset < edge" if self.states.groOffsetInMacroTile: module.add(VMinI32(dst=vgpr(vDst+l), src0=vgpr(edge), src1=vgpr(vSrc+l), comment="offset = (%s) ? offset(v%u) : edge(v%u)"%(cmpCommentText, vSrc+l, edge))) else: module.add(VCmpLtU32(dst=sgpr(tmpSgpr,self.states.laneSGPRCount), src0=vgpr(vSrc+l), src1=vgpr(edge), comment="shiftedOffset < shiftedEdge")) # shift module.add(VCndMaskB32(dst=vgpr(vDst+l), src0=vgpr(edge), src1=vgpr(vSrc+l), src2=sgpr(tmpSgpr,self.states.laneSGPRCount), comment="offset = (%s) ? offset(v%u) : edge(v%u)"%(cmpCommentText, vSrc+l, edge))) # For metadata and using A's margin, shift extra tail offset if tP["isM"] and not marginO == -1: module.add(VAndB32(dst=vgpr(edge), src0=(margin-1), src1=vgpr(tP["gpr"]["tReg"]), comment="shifTailOffstet = tailOffset %% %d"%(margin))) for l in range(0, tP["nrt"]): module.add(VAddU32(dst=vgpr(vDst+l), src0=vgpr(edge), src1=vgpr(vSrc+l), comment="offset += shifTailOffstet")) self.vgprPool.checkIn(edge) return module ############################################################################## def graShiftMX(self, kernel, tP, tPref): # graShift requires a vgpr for each address component (so each component # can be examined and shifted if necessary) - therefore does not work # with UseSgprForGRO. assert(not kernel["_UseSgprForGRO"]), "%s"%self.states.kernelName if not self.states.groOffsetInMacroTile: raise Exception(f"unsupport non groOffsetInMacroTile") module = Module("graShift") margin = tPref["glvw"] if tPref["rtv"] else 1 offsetGrpVgpr = self.vgprPool.checkOut(1, "offset", self.states.preventVgprOverflowDuringNewTile) with self.allocTmpSgpr(self.states.laneSGPRCount+1,2, tag="graShiftMX_tmpSgprInfo") as tmpSgprInfo: edgeGrpSgpr = tmpSgprInfo.idx edgeSgpr = tmpSgprInfo.idx + self.states.laneSGPRCount shiftSgpr = tmpSgprInfo.idx + self.states.laneSGPRCount module.add(SMulI32(dst=sgpr(edgeSgpr), src0=sgpr(tP["wg"]), src1=kernel[tP["mt"]], comment="WorkGroup[01] * MT")) module.add(SSubU32(dst=sgpr(edgeSgpr), src0=self.sizeRef(tP["idx"]), src1=sgpr(edgeSgpr), \ comment="edge = Size%s - WG*MT"%(tP["tileChar"]))) module.add(SLShiftRightB32(dst=sgpr(edgeGrpSgpr), shiftHex=hex(log2(margin)), src=sgpr(edgeSgpr), comment="edge group")) module.add(SAndB32(dst=sgpr(shiftSgpr), src0=sgpr(edgeSgpr), src1=(margin-1), comment="shift value = (GRVW - edge) % (GRVW-1)")) module.add(SSubU32(dst=sgpr(shiftSgpr), src0=margin, src1=sgpr(shiftSgpr), comment="shift value = (GRVW - edge) % (GRVW-1)")) module.add(SAndB32(dst=sgpr(shiftSgpr), src0=sgpr(shiftSgpr), src1=(margin-1), comment="shift value = (GRVW - edge) % (GRVW-1)")) vSrc = tP["vgprTileOffsets"] vDst = tP["vgprTileOffsets"] for l in range(0, tP["nrt"]): module.add(VLShiftRightB32(dst=vgpr(offsetGrpVgpr), shiftHex=hex(log2(margin)), src=vgpr(vSrc+l), comment="offsetGrp == edgeGrp ? shift : normal")) module.add(VCmpEQU32(dst=sgpr(edgeGrpSgpr, self.states.laneSGPRCount), src0=sgpr(edgeGrpSgpr), src1=vgpr(offsetGrpVgpr))) module.add(VSubU32(dst=vgpr(offsetGrpVgpr), src0=vgpr(vSrc+l), src1=sgpr(shiftSgpr), comment="offsetGrp == edgeGrp ? shift : normal")) module.add(VCndMaskB32(dst=vgpr(vDst+l), src0=vgpr(vSrc+l), src1=vgpr(offsetGrpVgpr), src2=sgpr(edgeGrpSgpr, self.states.laneSGPRCount))) self.vgprPool.checkIn(offsetGrpVgpr) return module ############################################################################## # Global Read Addresses: Final Offsets metadata ############################################################################## def graMetadataFinalOffsets(self, kernel, tP): module = Module("graMetadataFinalOffsets") module.addComment1("calculate metadata gra final offset") # alloc vgpr kReg = self.vgprPool.checkOut(1,"kReg") # remainder tmpVgpr = self.vgprPool.checkOutAligned(2,2,"tmpVgpr") dummy = self.vgprPool.checkOut(1,"dummy") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) # get constant parameter tc = tP["tensorChar"] waveWidth = kernel["WavefrontSize"] inputPerThread = kernel["MIInputPerThread"] # parameter for get each type index dividendForKId = kernel["MatrixInstM"] unrollSummation = [ i for i in tP["ia"] if i in kernel["ProblemType"]["IndicesSummation"] ] # strider for each type of index strideTile = self.sizeRef(unrollSummation[-1]) strideK = inputPerThread with self.allocTmpSgpr(1, tag="graMetadataFinalOffsets_tmpSgprInfo") as tmpSgprInfo: # unroll offset module.add(vectorStaticRemainder(dummy, kReg, "Serial", waveWidth, tmpVgprRes, tmpSgprInfo, \ "0. thread id in wave: wtid = tid %% wavelength(%u)" % waveWidth)) module.add(vectorStaticDivide(kReg, kReg, dividendForKId, tmpVgprRes, \ "1. unroll offset: kIdx = wtid / (MIN(%u) )" % (kernel["MatrixInstN"]))) module.add(vectorStaticMultiply(vgpr(kReg), vgpr(kReg), strideK, tmpSgprInfo, \ "1. unroll offset: grKOffset = kIdx * mStride(%u)" % strideK)) # Calculate final element offset for graIdx in range (0, kernel["MIWaveTile"][0]): vgprGro = "GlobalReadOffsetMetadata+%u"%(graIdx) module.add(VMulLOU32(vgpr(vgprGro), vgpr(vgprGro), strideTile, \ "2. tile offset: tile coord * sizeL")) module.add(VAddU32(vgpr(vgprGro), vgpr(vgprGro), vgpr(kReg), \ "2. final global read offset: fgrOffset = tile Offset + unroll Offset")) # elements to bytes module.add(vectorStaticDivide(vgprGro, vgprGro, 8, tmpVgprRes, \ " 3. bytes offset : bnIdx = global read elememt offset / 8")) module.add(Label("graFinalMeta", "")) # release register self.vgprPool.checkIn(kReg) self.vgprPool.checkIn(tmpVgpr) self.vgprPool.checkIn(dummy) return module ############################################################################## # Global Read Addresses: Final Offsets A/B ############################################################################## def graFinalOffsets(self, kernel, tP): module = Module("graFinalOffsets") tc = tP["tensorChar"] tmp = self.vgprPool.checkOut(3, "tmp", self.states.preventVgprOverflowDuringNewTile) graIdx = 0 swapPerpPara = (((tc in ("A", "B", "MXSA", "MXSB")) and kernel["DirectToVgpr%s"%tc]) and (not tP["tlu"]) and tP["nrp"] > 1) # both UseSgprForGRO and DTVA/B are enabled if ((tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc]) and kernel["_UseSgprForGRO"]: if tP["tlu"]: sizeOfaPart = kernel["LocalReadVectorWidth%s"%tc] numOfVW = tP["nrp"]//sizeOfaPart for sPerp in range(0, tP["nrpv"]*numOfVW): for perp in range(0, sizeOfaPart): for para in range(0, tP["nrc"]): for sPara in range(0, tP["nrcv"]//tP["nrcvpi"]): # single loop singleModule, graIdx = self.graFinalOffsetsSingleLoop(kernel, tP, tc, tmp, graIdx, perp, sPerp, para, sPara) module.add(singleModule) else: numOfVW = kernel["MIWaveTile%c"%tc]//kernel["VectorWidth%s"%tc] if kernel["VectorWidth%c"%tc]>1 else 1 sizeOfaPart = tP["nrp"]//numOfVW # swap para and perp for para in range(0, tP["nrc"]): for sPara in range(0, int(tP["nrcv"]/tP["nrcvpi"])): for sPerp in range(0, tP["nrpv"]*numOfVW): for perp in range(0, sizeOfaPart): # single loop singleModule, graIdx = self.graFinalOffsetsSingleLoop(kernel, tP, tc, tmp, graIdx, perp, sPerp, para, sPara) module.add(singleModule) # DTVA/B always go this way, including swizzled elif (not swapPerpPara): if kernel["UseGeneralizedNLCOne%s"%tc] and not self.states.inTailLoop: module.add(self.graFinalOffsetsSingleLoopGNLC(kernel, tP, tc)) else: module.addComment0("=============================================================") for perp in range(0, tP["nrp"]): for sPerp in range(0, tP["nrpv"]): for para in range(0, tP["nrc"]): for sPara in range(0, tP["nrcv"]//tP["nrcvpi"]): # single loop singleModule, graIdx = self.graFinalOffsetsSingleLoop(kernel, tP, tc, tmp, graIdx, perp, sPerp, para, sPara) module.add(singleModule) module.addComment0("=============================================================") else: # swap para and perp for para in range(0, tP["nrc"]): for sPara in range(0, int(tP["nrcv"]/tP["nrcvpi"])): for perp in range(0, tP["nrp"]): for sPerp in range(0, tP["nrpv"]): # single loop singleModule, graIdx = self.graFinalOffsetsSingleLoop(kernel, tP, tc, tmp, graIdx, perp, sPerp, para, sPara) module.add(singleModule) self.vgprPool.checkIn(tP["gpr"]["lwoT"]) tP["gpr"]["lwoT"] = None # if kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1: self.vgprPool.checkIn(tP["gpr"]["uReg2"]) tP["gpr"]["uReg2"] = None self.vgprPool.checkIn(tP["gpr"]["uReg"]) tP["gpr"]["uReg"] = None if "subIterReg" in tP["gpr"]: if tP["gpr"]["subIterReg"] is not None: self.vgprPool.checkIn(tP["gpr"]["subIterReg"]) tP["gpr"]["subIterReg"] = None if tP["vgprTileOffsetsCheckOut"]: self.vgprPool.checkIn(tP["vgprTileOffsets"]) tP["vgprTileOffsets"] = None tP["vgprTileOffsetsCheckOut"] = False # _UseSgprForGRO uses same vgpr for ureg and tP["gpr"]["unrollOffsets"] so # let checkin(ureg) do the checkin # vgprTileOffsets is renamed version of treg/lwo so checkin here if not kernel["_UseSgprForGRO"]: self.vgprPool.checkIn(tP["gpr"]["unrollOffsets"]) tP["gpr"]["unrollOffsets"] = None if tP["vgprPackedOffsets"] != None: self.vgprPool.checkIn(tP["vgprPackedOffsets"]) tP["vgprPackedOffsets"] = None self.vgprPool.checkIn(tmp) return Module("Global Read Addresses: Final Offsets A/B (Empty)") if self.dontAppendCode else module ############################################################################## # Global Read Addresses: Final Offsets A/B (single loop) ############################################################################## def graFinalOffsetsSingleLoopGNLC(self, kernel, tP, tc, margin = -1): module = Module() if margin == -1: margin = tP["glvw"] if tP["rtv"] else 1 module.addComment("Using GLNC for %s"%tc) groVgpr0 = "GlobalReadOffset%s+%u" % (tc, 0) parDimSize = kernel["MacroTile%s"%tc] if kernel["ProblemType"]["TLU%s"%tc] == 1 else kernel["DepthU"] numThreadsCoalesced = (parDimSize // kernel["GlobalReadVectorWidth%s"%tc]) numThreadsPerMI = max((kernel["MatrixInstM"] if kernel["ProblemType"]["TLU%s"%tc] == 1 \ else kernel["MatrixInstK"]) // kernel["GlobalReadVectorWidth%s"%tc],\ 1) numThreadGroupsPerParDim = numThreadsCoalesced // numThreadsPerMI module.addComment0("NumThreadsCoalesced%s = %u, %u total threads, %u thread groups"%( \ tc, numThreadsCoalesced, kernel["NumThreads"], numThreadGroupsPerParDim)) module.add(VMovB32(dst=vgpr(groVgpr0), src=vgpr("Serial"))) for perp in range(1, tP["ntpl"]): groVgpr = "GlobalReadOffset%s+%u" % (tc, perp) groVgprPrev = "GlobalReadOffset%s+%u" % (tc, perp - 1) strideLoad = kernel["NumThreads"] # stride between consecutive loads module.add(VAddU32(dst=vgpr(groVgpr), src0=strideLoad, src1=vgpr(groVgprPrev), comment=" = vgprSerial + %u * %u"%(perp, strideLoad))) # TODOBS: Only use perperm for pow2 cases for now.. abMatrixInfo = self.states.a if tc == 'A' else self.states.b perpStride = abMatrixInfo.gNLCPerpStride permBlock = abMatrixInfo.gNLCPermBlock usePerpPerm = perpStride > 1 perpBlockSize = abMatrixInfo.gRDtlSwizzlePerpBlockSize tmpv = self.vgprPool.checkOutAligned(2,2, tag="graFinalOffsetsSingleLoopGNLC_tmpv") tmpv2 = self.vgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_tmpv2") tmps = self.sgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_tmps") tmps2 = self.sgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_tmps2") divsor = numThreadsCoalesced useMagicDiv = divsor > 1 and (divsor & (divsor-1)) != 0 if useMagicDiv: maxDividend = kernel["NumThreads"] * tP["ntpl"] # Verified limits for magic div algo assert divsor < 512 and maxDividend < 64 * 1024 Nbits = 32 # If divsor is even, compute largest odd value that is multiple of divsor. # divsor = 2^ee * divsor2 divsor2, cc,ee = divsor, 1, 0 while divsor2 % 2 == 0: divsor2 //= 2 cc *= 2 ee += 1 mm, shPost, l = choose_multiplier(divsor2, Nbits - ee, 16) module.add(SMovB32(dst=sgpr(tmps), src=mm, comment="Used in magic div algo, multiplicative constant for 1/%u"%divsor2)) module.add(SMovB32(dst=sgpr(tmps2), src=divsor2)) if kernel["EdgeType"] == "ShiftPtr" and kernel["ProblemType"]["TLU%s"%tc] == 1: tmpSgpr = self.sgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_tmpSgpr") module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tP["wg"]), src1=kernel[tP["mt"]], comment="WorkGroup[01] * MT")) module.add(SSubU32(dst=sgpr(tmpSgpr), src0=self.sizeRef(tP["idx"]), src1=sgpr(tmpSgpr), \ comment="edge = Size%s - WG*MT"%(tP["tileChar"]))) module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=margin, comment="edge -= margin(%u)"%(margin))) for perp in range(0, tP["ntpl"]): strideChar = 'L' if tc == 'A' else 'K' grov = "GlobalReadOffset%s+%u" % (tc, perp) # Compute division # tmpv = vgprSerial // divsor if useMagicDiv: if ee > 0: module.add(VLShiftRightB32(dst=vgpr(tmpv2), shiftHex=ee, src=vgpr(grov), comment="division")) module.add(VMulLOU32(dst=vgpr(tmpv), src0=sgpr(tmps), src1=vgpr(tmpv2), \ comment="division" )) module.add(VMulHIU32(dst=vgpr(tmpv+1), src0=sgpr(tmps), src1=vgpr(tmpv2), \ comment="division" )) module.add(VLShiftRightB64(dst=vgpr(tmpv,2), shiftHex=(Nbits - ee + shPost), src=vgpr(tmpv,2), comment="division")) else: module.add(VLShiftRightB32(dst=vgpr(tmpv), shiftHex=log2(divsor), src=vgpr(grov), comment="division")) # Compute remainder # tmpv2 = vgprSerial % divsor if useMagicDiv: module.add(VMulLOU32(dst=vgpr(tmpv2), src0=sgpr(tmps2), src1=vgpr(tmpv))) module.add(VLShiftLeftB32(dst=vgpr(tmpv2), shiftHex=ee, src=vgpr(tmpv2), comment="remainder")) module.add(VSubU32(dst=vgpr(tmpv2), src0=vgpr(grov), src1=vgpr(tmpv2))) else: module.add(VAndB32(dst=vgpr(tmpv2), src0=hex(divsor - 1), src1=vgpr(grov))) # Permute logic perp dim if usePerpPerm: reMap0 = self.vgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_reMap0") reMap1 = self.vgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_reMap1") module.addComment0("Computing strided(%u) perp indicies"%perpStride) module.add(VAndB32(dst=vgpr(reMap0), src0=(permBlock // perpStride - 1), src1=vgpr(tmpv), comment="r0 = I %% (%u // %u)"%(permBlock, perpStride))) module.add(VLShiftLeftB32(dst=vgpr(reMap0), shiftHex=log2(perpStride), src=vgpr(reMap0), comment="r0 = %u * r0"%(perpStride))) module.addComment0("Computing r1 = (I %% %u) // (%u // %u)"%(permBlock, permBlock, perpStride)) module.add(VAndB32(dst=vgpr(reMap1), src0=(permBlock - 1), src1=vgpr(tmpv), comment="r1 = I %% (%u)"%(permBlock))) module.add(VLShiftRightB32(dst=vgpr(reMap1), shiftHex=log2(permBlock // perpStride), src=vgpr(reMap1), comment="r1 = (r1) // (%u // %u)"%(permBlock, perpStride))) module.add(VAddU32(dst=vgpr(reMap0), src0=vgpr(reMap0), src1=vgpr(reMap1), comment="r0 = r0 + r1" )) module.add(VLShiftRightB32(dst=vgpr(reMap1), shiftHex=log2(permBlock), src=vgpr(tmpv), comment="r1 = I // %u"%(permBlock))) module.add(vectorStaticMultiplyAdd(vgpr(tmpv), vgpr(reMap1), permBlock, vgpr(reMap0), None)) module.addComment0("Done computing strided(%u) perp indices"%perpStride) self.vgprPool.checkIn(reMap0) self.vgprPool.checkIn(reMap1) applyRotation = perpBlockSize > 0 bS = abMatrixInfo.gRDtlSwizzleParaBlockSize if applyRotation: # ntc = numThreadsCoalecsed # tmpv2 = serial % ntc # tmpv = serial / ntc tmpv3 = self.vgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_tmpv3") tmpv4 = self.vgprPool.checkOut(1, tag="graFinalOffsetsSingleLoopGNLC_tmpv4") # For TLU=0, apply an index rotation (with wrap-around) in the parallel dimension for a fix blocksize (=MI_K) # in the summation dimension. # ex: Given thread ID indices: [0,1,2,3,4,5,6,7, 8,9,10,..., 16,...., 24,..., 32,...., 40,....] # where threads 0-7 loads row/col 1, threads 8-15 loads row/col 2, ... # we rotate the indices as follows: # [0,1,2,3,4,5,6,7, 15,8,9,10,11,12,13,14, 22,23,16,17,18,19,20,21, ...] # # rotated index is computed by: new_idx = ((bS - off) + idx) % bS where off is the amount to rotate # module.addComment0("Apply index rotation in parallel dim") # Computes {0,1,2,..bS-1, 0,1,2,..bS-1, ...} module.add(VAndB32(dst=vgpr(tmpv3), src0=hex(bS - 1), src1=vgpr(tmpv2))) module.add(VLShiftRightB32(dst=vgpr(tmpv4), shiftHex=log2(perpBlockSize), src=vgpr(tmpv))) module.add(VAndB32(dst=vgpr(tmpv4), src0=hex(bS - 1), src1=vgpr(tmpv4))) module.add(VSubU32(dst=vgpr(tmpv4), src0=bS, src1=vgpr(tmpv4))) module.add(VAddU32(dst=vgpr(tmpv3), src0=vgpr(tmpv2), src1=vgpr(tmpv4))) module.add(VAndB32(dst=vgpr(tmpv3), src0=hex(bS - 1), src1=vgpr(tmpv3))) # Compute subgroup ids of each {0,...bS-1} module.add(VLShiftRightB32(dst=vgpr(tmpv4), shiftHex=log2(bS), src=vgpr(tmpv2))) with self.allocTmpSgpr(1, tag="graFinalOffsetsSingleLoopGNLC_tmpSgprInfo") as tmpSgprInfo: module.add(vectorStaticMultiplyAdd(vgpr(tmpv2), vgpr(tmpv4), bS, vgpr(tmpv3), tmpSgprInfo)) self.vgprPool.checkIn(tmpv3) self.vgprPool.checkIn(tmpv4) stride = "Strides%s"%(tc) module.add(VLShiftLeftB32(dst=vgpr(grov), shiftHex=log2(kernel["GlobalReadVectorWidth%s"%tc]), src=vgpr(tmpv2))) module.add(VMulLOU32(dst=vgpr(tmpv), src0=sgpr(stride), src1=vgpr(tmpv))) if kernel["EdgeType"] == "ShiftPtr" and kernel["ProblemType"]["TLU%s"%tc] == 1: module.add(VMinI32(dst=vgpr(grov), src0=sgpr(tmpSgpr), src1=vgpr(grov), comment="")) module.add(VAddU32(dst=vgpr(grov), src0=vgpr(tmpv), src1=vgpr(grov), \ comment="final" )) module.add(vectorMultiplyBpe(grov, grov, tP["bpeGR"])) ptrshift = int(self.states.srdShiftLeft[tc] * tP["bpeGR"]) module.add(VAddU32(dst=vgpr(grov), src0=ptrshift , src1=vgpr(grov), \ comment="ptr-shift" )) self.vgprPool.checkIn(tmpv) self.vgprPool.checkIn(tmpv2) self.sgprPool.checkIn(tmps) self.sgprPool.checkIn(tmps2) if kernel["EdgeType"] == "ShiftPtr" and kernel["ProblemType"]["TLU%s"%tc] == 1: self.sgprPool.checkIn(tmpSgpr) return module ############################################################################## # Global Read Addresses: Final Offsets A/B (single loop) ############################################################################## def graFinalOffsetsSingleLoop(self, kernel, tP, tc, tmp, graIdx, perp, sPerp, para, sPara): module = Module("graFinalOffsetsSingleLoop") problemType = kernel["ProblemType"] tVW = 1 tVS = 0 uVW = 1 uVS = 0 # single loop start # vgpr assignments if tP["tlu"]: vgprTile = tP["vgprTileOffsets"] + para*tVW + sPara*tVS vgprUnroll = tP["gpr"]["unrollOffsets"] + perp*uVW + sPerp*uVS else: vgprTile = tP["vgprTileOffsets"] + perp*tVW + sPara*tVS vgprUnroll = tP["gpr"]["unrollOffsets"] + para*uVW + sPerp*uVS if graIdx==0 or not kernel["_UseSgprForGRO"]: # emit global offset macro # TODO -refactor this and macro def to pass all indices, use the ones we need if kernel["BufferLoad"]: if tP["isSwizzled"]: bfName = "GLOBAL_OFFSET_%s_SWIZZLED" % tP["tensorChar"] else: bfName = "GLOBAL_OFFSET_%s" % tP["tensorChar"] bfArgs = ["vgprGlobalReadOffset%s+%u"%(tP["tensorChar"], graIdx)] else: bfName = "GLOBAL_OFFSET_%s" % tP["tensorChar"] bfArgs = ["vgprGlobalReadAddr%s+%u"%(tP["tensorChar"], graIdx)] packedIter = 0 #iterator through ia iaToGpr = [None] * problemType["TotalIndices"] for i in tP["ia"]: if i < problemType["NumIndicesC"]: if i == tP["tileIdx"]: iaToGpr[i] = vgprTile bfArgs.append( "%u" % iaToGpr[i] ) else: if isPackedIndex(kernel,i): iaToGpr[i] = tP["vgprPackedOffsets"] + \ (vgprTile-tP["vgprTileOffsets"])*(len(tP["PackedIndices"])-1) + \ packedIter bfArgs.append( "%u" % (iaToGpr[i]) ) packedIter += 1 else: # just a group index if not kernel["BufferLoad"]: # buffer load adds these to SRD not the GLOBAL_OFFSET here bfArgs.append( "sgprWorkGroup%u"%i ) else: # summation index if i == problemType["IndexUnroll"]: iaToGpr[i] = vgprUnroll bfArgs.append( "%u" % iaToGpr[i] ) # other summation indices are ignored bfArgs.append( "%u" % tmp ) bfComment = "gRO%s_%u_%u_%u_%u" % (tP["tensorChar"], para, sPara, perp, sPerp) if self.states.asmCaps["HasVgprMSB"]: module.add(self.globalOffset(kernel, tP, tc, bfArgs, bfComment)) else: module.add(MacroInstruction(name=bfName, args=bfArgs, comment=bfComment)) dest = f'GlobalReadOffset{tP["tensorChar"]}+{graIdx}' if kernel["BufferLoad"]: module.add(vectorMultiplyBpe(dest, dest, tP["bpeGR"])) else: startVgpr = self.startVgprGlobalReadAddressesA if tc == 'A' else self.startVgprGlobalReadAddressesB destVgpr = startVgpr + graIdx module.add(vectorMultiply64Bpe(destVgpr, destVgpr, tP["bpeGR"], tmp)) with self.allocTmpSgpr(2, tag="graFinalOffsetsSingleLoop_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx # modify start if (not kernel["_UseSgprForGRO"]) and kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"]: # add room for instruction offset groVgpr = "GlobalReadOffset%s+%u" % (tP["tensorChar"], graIdx) module.add(SMovB32(dst=sgpr(tmpSgpr), src=self.buff_load_inst_offset_max)) module.add(VAddU32(dst=vgpr(groVgpr), src0=vgpr(groVgpr), src1=sgpr(tmpSgpr), comment="shift for UseInstOffsetForGRO")) ldsInc = (self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * kernel["GlobalReadVectorWidth%s"%tc] * tP["bpeGR"] if kernel["LdsBlockSizePerPad%s"%tc] != 0: ldsInc += (ldsInc // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeGR"] else: padInterval = (self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * self.states.bpr ldsInc += (ldsInc // padInterval) * kernel["LdsPad%s"%tc] * tP["bpeGR"] # buffer_load only support 12 bit instruction offset # we have to increase m0 if offset is larger thant 12 bits # so only keep 12 bit offset and subtract it on global address # global address will add back by buffer_load instruction offset ldsInc = (ldsInc * graIdx) % self.buff_load_inst_offset_max if (ldsInc != 0): module.add(SMovB32(dst=sgpr(tmpSgpr), src=int(ldsInc))) module.add(VSubU32(dst=vgpr(groVgpr), src0=vgpr(groVgpr), src1=sgpr(tmpSgpr), comment="sub offset for buffer_load instoffset")) def computeScalarGroImpl(scalarGro): # this needs unroll stride in some cases and free stride in others # if we have multiple free strides - what is expected behavior? # could just extract the first free dimension from A? stride1 = "Stride%s%s"%(tc,self.states.indexChars[tP["idx"]]) if tP["tlu"]: tileStride = kernel[tP["lsc"]] * (para*tVW + sPara*tVS) unrollStride = 0 if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc]: unrollStride = perp*uVW + (kernel[tP["lsp"]] * sPerp * kernel["LocalReadVectorWidth%s"%tc]) else: unrollStride = kernel[tP["lsp"]] * (perp*uVW + sPerp*uVS) unrollSummation = [ i for i in tP["ia"] if i in problemType["IndicesSummation"] ] strideU = "Stride%s%s"%(tc,self.states.indexChars[unrollSummation[-1]]) module.add(SMulI32(dst=sgpr(scalarGro), src0=sgpr(strideU), src1=unrollStride, \ comment="compute offset diff (scaled unrollDim)")) if tileStride: module.add(SAddU32(dst=sgpr(scalarGro), src0=sgpr(scalarGro), src1=tileStride, \ comment="compute offset diff (tileDim)")) else: tileStride = 0 if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc] and kernel["VectorWidth%s"%tc] > 1: tileStride = perp*tVW + (kernel[tP["lsp"]] * sPerp * kernel["VectorWidth%s"%tc]) else: tileStride = kernel[tP["lsp"]] * (perp*tVW + sPara*tVS) unrollStride = kernel[tP["lsc"]] * (para*uVW + sPerp*uVS) strideF = "Stride%s%s"%(tc,self.states.indexChars[tP['tileIdx']]) module.add(SMulI32(dst=sgpr(scalarGro), src0=sgpr(strideF), src1=tileStride, \ comment="compute offset diff (scaled tileDim)")) if unrollStride: module.add(SAddU32(dst=sgpr(scalarGro), src0=sgpr(scalarGro), src1=unrollStride, \ comment="compute offset diff (unrollDim)")) # Using offsets so GRO holds a byte offset not an element offset # So scale here before comparison: module.add(scalarMultiplyBpe(scalarGro, scalarGro, float(tP["bpeGR"]))) if kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"]: # add room for instruction offset module.add(SAddU32(dst=sgpr(scalarGro), src0=sgpr(scalarGro), src1=self.buff_load_inst_offset_max, comment="shift for UseInstOffsetForGRO")) ldsInc = (self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * kernel["GlobalReadVectorWidth%s"%tc] * tP["bpeGR"] if kernel["LdsBlockSizePerPad%s"%tc] != 0: ldsInc += (ldsInc // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeGR"] else: padInterval = (self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * self.states.bpr ldsInc += (ldsInc // padInterval) * kernel["LdsPad%s"%tc] * tP["bpeGR"] # buffer_load only support 12 bit instruction offset # we have to increase m0 if offset is larger thant 12 bits # so only keep 12 bit offset and subtract it on global address # global address will add back by buffer_load instruction offset ldsInc = (ldsInc * graIdx) % self.buff_load_inst_offset_max if (ldsInc != 0): module.add(SSubU32(dst=sgpr(scalarGro), src0=sgpr(scalarGro), src1=int(ldsInc), comment="sub offset for buffer_load instoffset")) if self.states.checkGRO: # Debug mode to verify that the computed offsets are offset by the expected scalar print(tc, "tileStride=", tileStride, "unrollStride=", unrollStride, \ "stride=%s"%(stride1)) module.add(self.getVectorDiffAssert(vgpr("GlobalReadOffset%s+%u"%(tc,0)), \ vgpr("GlobalReadOffset%s+%u"%(tc,graIdx)), \ sgpr(scalarGro))) needFirstSgprOffset = kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"] # TODO: Check correctness of MXS on TDM if (kernel["_UseSgprForGRO"] or self.states.checkGRO) and (needFirstSgprOffset or graIdx > 0) and ("MXS" in tc or not kernel["enableTDM%s"%tc]): # compute offsets for scalar global read offsets: if kernel["_UseSgprForGRO"]: tmpIdx = graIdx if needFirstSgprOffset else graIdx-1 scalarGro = "ScalarGlobalReadOffset%s+%u"%(tc, tmpIdx) computeScalarGroImpl(scalarGro) else: # TODO: need for investagation for replacing by allocTmpSgpr with self.allocTmpSgpr(1, tag="computeScalarGroImpl_tmpSgprInfo") as tmpSgprInfo: scalarGro = tmpSgprInfo.idx computeScalarGroImpl(scalarGro) graIdx += self.states.rpgo if kernel["BufferLoad"] else self.states.rpga return module, graIdx def computeMetaDataSrd(self, kernel, tP, tc, indices): module = Module("computeMetaDataSrd") wroteTileStart = True with self.allocTmpSgpr(2 + 2 + 2, tag="computeMetaDataSrd_tmpSgprInfo") as tmpSgprInfo: stmp = tmpSgprInfo.idx gsuoffset = stmp tensorSize = stmp tileStart = stmp+2 blockOffset = stmp+4 actualBatchSize = stmp+5 #for broadcast actualBatchIndex = stmp+5 #for broadcast module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart+0), sgpr(tileStart+1), sgpr(tP["wg"]), kernel[tP["mt"]], comment="WorkGroup[01] * MT")) unrollSummation = [ i for i in tP["ia"] if i in kernel["ProblemType"]["IndicesSummation"] ] module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart), sgpr(tileStart+1), sgpr(tileStart+0), self.sizeRef(unrollSummation[-1]), \ comment="scaled tile-offset by Summation size")) depthU = kernel["DepthU"] gsucLabel = Label(label=self.labels.getNameInc("GSUC_M"), comment="") gsucLabelEnd = Label(label=self.labels.getNameInc("GSUC_M_End"), comment="") module.add(SAndB32(dst=sgpr(tmpSgprInfo.idx), src0=sgpr("GSU"), src1=hex(0x8000), comment="SCC = (GSUC == 1) ?")) gsuOffsetStr = "gsuOffset = DepthU*GSUSumIdx" module.add(SCBranchSCC1(labelName=gsucLabel.getLabelName(), comment="branch if GSUC == 1")) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(stmp+0), sgpr(stmp+1), depthU, sgpr("GSUSumIdx"), comment=gsuOffsetStr)) module.add(SBranch(gsucLabelEnd.getLabelName())) module.add(gsucLabel) gsuOffsetStr = "gsuOffset = DepthU*accumulatedNumOfLoopCounterL" loopCounterName = self.loopCounterName(kernel, self.states.unrollIdx) module.add(SLShiftRightB32(dst=sgpr(loopCounterName), src=sgpr("SizesSum"), shiftHex=log2(depthU), \ comment="s[%s] = s[sgprSizesSum] / %s"%(loopCounterName, depthU))) module.add(self.calculateLoopNumIterOffsetGsu(kernel, loopCounterName, tmpSgprInfo)) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(stmp+0), sgpr(stmp+1), sgpr(stmp+0), depthU, comment=gsuOffsetStr)) module.add(gsucLabelEnd) module.add(SAddU32(dst=sgpr(tileStart+0), src0=sgpr(tileStart+0), src1=sgpr(stmp+0), comment="accum GsuOffet term to tilestart")) module.add(SAddCU32(dst=sgpr(tileStart+1), src0=sgpr(tileStart+1), src1=sgpr(stmp+1), comment="accum GsuOffet term to tilestart")) sizeIndex = [ dim for dim in tP["ia"] ] assert(len(sizeIndex) >= 2) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(blockOffset), sgpr(blockOffset+1), self.sizeRef(sizeIndex[0]), self.sizeRef(sizeIndex[1]), \ comment="calculate metadata tensor size")) for dim in sizeIndex[2:]: module.add(SCmpEQU32(src0=sgpr("Stride%s%s"%(tc,self.states.indexChars[tP['ia'][dim]])), src1=0, comment="broadcast %s?"%tc)) module.add(SCSelectB32(dst=sgpr(actualBatchSize), src0=hex(1) , src1=self.sizeRef(sizeIndex[dim]), comment="set batchSize as 1 for boardcast %s"%tc)) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tensorSize), sgpr(tensorSize+1), sgpr(blockOffset), sgpr(actualBatchSize), \ comment="calculate metadata tensor size")) module.add(SCSelectB32(dst=sgpr(actualBatchIndex), src0=0 , src1=sgpr("WorkGroup2"), comment="set batchIndex as 0 for boardcast %s"%tc)) if self.states.use64bShadowLimit: limitTmp0 = "ShadowLimitMetadata" limitTmp1 = "ShadowLimitMetadata+1" else: limitTmp0 = stmp+0 limitTmp1 = stmp+1 module.add(SSubU32(sgpr(limitTmp0), sgpr(tensorSize), sgpr(tileStart+0), "sub tileStart")) module.add(SSubBU32(sgpr(limitTmp1), sgpr(tensorSize+1), sgpr(tileStart+1), "sub tileStart")) if self.states.use64bShadowLimit: module.add(SLShiftRightB64(sgpr(limitTmp0,2), hex(log2(8)), sgpr(limitTmp0,2), "Set limit to use bytes")) module.add(SCmpEQU32(sgpr(limitTmp1), 0, "are we within 2^32?")) module.add(SCSelectB32(sgpr("SrdMetadata+2"), sgpr(limitTmp0), "BufferLimit", "Move shadow to real if we are within 2^32")) else: module.add(SLShiftRightB32(sgpr("SrdMetadata+2"), hex(log2(8)), sgpr(limitTmp0), "Set limit to use bytes")) numDim = len(indices) wg=2 for i in range(0, numDim): idx = indices[i] if not ( idx == kernel["ProblemType"]["Index0"] \ or idx == kernel["ProblemType"]["Index1"] \ or idx in kernel["ProblemType"]["IndicesSummation"] \ or isPackedIndex(kernel, idx)): assert(wg==2) if not wroteTileStart: module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart+0), sgpr(tileStart+1), sgpr(blockOffset), sgpr("WorkGroup2"), comment="block offset*WG")) wroteTileStart = True else: module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(stmp+0), sgpr(stmp+1), sgpr(blockOffset), sgpr(actualBatchIndex), comment="block offset*WG")) module.add(SAddU32(sgpr(tileStart+0), sgpr(tileStart+0), sgpr(stmp+0), "accum wg term to tilestart")) module.add(SAddCU32(sgpr(tileStart+1), sgpr(tileStart+1), sgpr(stmp+1), "accum wg term to tilestart")) if wroteTileStart: module.add(SLShiftRightB64(sgpr(tileStart,2), hex(log2(8)), sgpr(tileStart,2), "Set limit to use bytes")) module.add(SAddU32(sgpr("SrdMetadata+0"), sgpr("AddressMetadata+0"), sgpr(tileStart+0), "SRD base = Address+ tileStart0")) module.add(SAddCU32(sgpr("SrdMetadata+1"), sgpr("AddressMetadata+1"), sgpr(tileStart+1), "SRD base = Address+ tileStart1")) else: module.add(SMovB64(sgpr("SrdMetadata+0", 2), sgpr("AddressMetadata+0", 2), "init SRD base address" )) module.add(SMovB32(sgpr("SrdMetadata+3"), "Srd127_96", "Set bits 127_96 in SRD")) return module ############################################################################## # Add the constant offsets to the specified srd. # Srd is set to point to the base of the tile. All offsets except lowest-order # 2d dims are computed into the SRD. # GRO are offset from the tile SRD and the first GRO will be 0 # Only called for BufferLoad=1 (or eventually BufferStore=1) ############################################################################## def computeLoadSrd(self, kernel, tP, tc, indices, bpe): module = Module("computeLoadSrd") moduleLoadGeneralBatch = Module("computeLoadSrd-GeneralBatch") moduleLoadStridedBatch = Module("computeLoadSrd-StridedBatch") use64bShadowLimit = self.states.use64bShadowLimitMX if tc in ["MXSA", "MXSB"] else self.states.use64bShadowLimit isgfx950 = kernel["ISA"][:2] == (9, 5) # An MX scale tensor uses the swizzled SRD-limit math when MXScaleFormat # selects a swizzled layout: HostPreSwizzle (gfx950) or InMemorySwizzle # (gfx1250 TDM-populated). NoSwizzle MX scales fall through to the # standard tensor2dSize SRD math. mxScaleFormat = kernel.get("MXScaleFormat", "NoSwizzle") isMxSwizzledScaleLayout = ("MXS" in tc) and mxScaleFormat in ("InMemorySwizzle", "HostPreSwizzle") # UseSubtileImpl uses a tile-boundary fixed Srd+2 for both MX scale and data A/B. # This avoids 32-bit overflow when computing the full tensor2dSize (N*K or M*K > 2^32). useSubtile = bool(kernel.get("UseSubtileImpl")) useFixedSrd2 = useSubtile isPreShuffledAB = tc in ("A", "B") and kernel["ProblemType"].get("SwizzleTensor%s" % tc, False) isSwizzledSubtile = (isMxSwizzledScaleLayout or isPreShuffledAB) and useSubtile if isMxSwizzledScaleLayout: useFixedSrd2 = True tcab = "A" if tc == "MXSA" else "B" mxBlock = kernel["ProblemType"]["MXBlock%s"%tcab] swizzleSize0 = 32 # M,N direction swizzleSize1 = 256 # K direction swizzleBlockSize = swizzleSize0 * swizzleSize1 // mxBlock else: if isSwizzledSubtile: swizzleSize0 = 16 # M,N direction swizzleSize1 = 32 # K direction for MXFP4 (TODO: use bpe to support different dataTypes) else: swizzleSize0 = 1 # M,N direction swizzleSize1 = 1 # K direction swizzleBlockSize = swizzleSize0 * swizzleSize1 allocateTensor2dSize = use64bShadowLimit and not useFixedSrd2 numDim = len(indices) with self.allocTmpSgpr(2 + 2 + (0 if allocateTensor2dSize else 2), tag="computeLoadSrd_tmpSgprInfo") as tmpSgprInfo: stmp = tmpSgprInfo.idx tileStart = stmp+2 if use64bShadowLimit: tensor2dSize0 = "ShadowLimit%s+0"%tc tensor2dSize1 = "ShadowLimit%s+1"%tc elif not useFixedSrd2: tensor2dSize0 = stmp+4 tensor2dSize1 = stmp+5 wroteTileStart = False #--- # Compute tileStart #elements from the 2D array start # Add tile (and unroll if GSU) component into SRD - SRD will point to beginning of the macro-tile: if self.states.groOffsetInMacroTile: # packed modes can't use this mode, and code here assumes 1 index. assert(len(kernel["PackedC0IndicesX"])==1) assert(len(kernel["PackedC1IndicesX"])==1) wroteTileStart = True #tP['ia'][1] # This is guaranteed to fit in 32-bit since the WG*MT is a number of elements in some unsigned direction: if useFixedSrd2: # UseSubtileImpl fixedSrd2 case (including swizzle and nonSwizzle): tile start uses roundup(MT/swizzleSize0) mt = roundUp(kernel[tP["mt"]] / swizzleSize0) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart+0), sgpr(tileStart+1), sgpr(tP["wg"]), mt, comment="WorkGroup[01] * roundup(MT/%u)"%swizzleSize0)) else: module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart+0), sgpr(tileStart+1), sgpr(tP["wg"]), kernel[tP["mt"]], comment="WorkGroup[01] * MT")) strideF = self.strideRef(tc, tP['tileIdx']) if not self.isConstUnitStride(strideF): if useFixedSrd2: # Tile-boundary SRD+2 for UseSubtileImpl (unified for MX scale and data A/B). # Avoids 32-bit overflow from computing full tensor2dSize when N*K or M*K > 2^32. # # tileStart is in block units (roundUp(MT/swizzleSize0)). # numLine = min(roundUp(size/swizzleSize0) - tileStart_blk, roundUp(MT/swizzleSize0)) - 1 # Srd+2 = numLine * stride_bytes + swizzleBlockSize*(DepthU/swizzleSize1) # # Key: numLine/numElems <= MT (compile-time), so the multiply stays in 32 bits. mt_units = mt # roundUp(MT/swizzleSize0), compile-time extra_bytes = swizzleBlockSize * (kernel["DepthU"] // swizzleSize1) for i in range(0, numDim): idx = indices[i] if idx == kernel["ProblemType"]["Index0"] or idx == kernel["ProblemType"]["Index1"]: size = self.sizeRef(idx) if isSwizzledSubtile: # tileStart already in block units (WG * roundUp(MT/swizzleSize0)) module.add(SAddU32(dst=sgpr(stmp+0), src0=size, src1=(swizzleSize0 - 1), comment="size + %u - 1"%swizzleSize0)) module.add(SLShiftRightB32(dst=sgpr(stmp+0), src=sgpr(stmp+0), shiftHex=log2(swizzleSize0), comment="roundup(size/%u)"%swizzleSize0)) module.add(SSubU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr(tileStart+0), comment="numBlkToEnd = roundUp(size/%u) - tileStart_blk"%swizzleSize0)) else: # tileStart in element units (WG * MT); no block rounding needed module.add(SSubU32(dst=sgpr(stmp+0), src0=size, src1=sgpr(tileStart+0), comment="numToEnd = size - WG*MT")) module.add(SMinU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=mt_units, comment="min (numBlkToEnd, roundup(MT/%u))"%swizzleSize0)) module.add(SSubU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=1, comment="numLine = min - 1 (0-based index)")) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(stmp+0), sgpr(stmp+1), sgpr(stmp+0), \ strideF, comment="numLine * stride")) if isMxSwizzledScaleLayout: module.add(SAddU32(dst=sgpr("Srd%s+2"%tc), src0=sgpr(stmp+0), src1=extra_bytes, comment="buffer_load limit for %s"%tc)) else: # (numLine * stride + DepthU) * bpe -- mirrors scale path structure module.add(SAddU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=extra_bytes, comment="+ DepthU (one K step)")) module.add(scalarMultiplyBpe("Srd%s+2"%tc, stmp+0, float(tP["bpeGR"]), comment="buffer_load limit for %s (tile-boundary, avoids 32-bit overflow)"%tc)) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart), sgpr(tileStart+1), sgpr(tileStart+0), \ strideF, comment="tlu=0, scaled tile-offset by stride")) skComponent = Component.StreamK.find(self) module.add(skComponent.computeLoadSrd(self, kernel, tP, stmp)) gsuComponent = Component.GSU.find(self) module.add(gsuComponent.computeLoadSrd(self, kernel, tP, stmp, tileStart)) # Output : tileStart[0:1] have offset in elements from the 2D start of the tile. # if groOffsetInMacroTile=1, 2DStart + tileStart gives the the start of the macro-tile; # This is used to compute the limit. # Later we modify tileStart to include batch and higher-order dims and add this to SRD. #--- # Compute BUFFER Limit: prePad = int(self.states.srdShiftLeft[tc] * tP["bpeGR"]) # leave room in case we have to pointer shift if not wroteTileStart: module.add(SMovB64(dst=sgpr(tileStart, 2), src=0, comment="set default tileStart")) #Calculate tensor 2d size # For UseSubtileImpl kernels (MX and non-MX), useFixedSrd2=True so tensor2dSize is not needed. if not useFixedSrd2: if use64bShadowLimit or ((not use64bShadowLimit) and tensor2dSize0 % 2 == 0): module.add(SMovB64(dst=sgpr(tensor2dSize0, 2), src=0x1, comment="Init tensor size")) else: module.add(SMovB32(dst=sgpr(tensor2dSize0), src=0x1, comment="Init tensor size")) module.add(SMovB32(dst=sgpr(tensor2dSize1), src=0x0, comment="init tensor size")) for i in range(0, numDim): if useFixedSrd2: # fixed Srd2 case, skip tensor2dSize0/1 calculation continue idx = indices[i] if idx == kernel["ProblemType"]["Index0"] \ or idx == kernel["ProblemType"]["Index1"] \ or idx in kernel["ProblemType"]["IndicesSummation"] \ or isPackedIndex(kernel, idx): stride = self.strideRef(tc,idx) size = self.sizeRef(idx) # The sizeL of a structure sparsity 2:4 matrix is half of the dense matrix. if (idx in kernel["ProblemType"]["IndicesSummation"]): if ((tP["isA"] and kernel["ProblemType"]["Sparse"] == 1) or \ (tP["isB"] and kernel["ProblemType"]["Sparse"] == 2) or \ (tP["isM"])) : divider = 8 if tP["isM"] else 2 module.add(SLShiftRightB32(dst=sgpr(stmp), src=size, shiftHex=hex(int(log(divider,2))), comment="(size/%u)"%divider)) module.add(SSubU32(dst=sgpr(stmp), src0=sgpr(stmp), src1=0x1, comment="(size/%u-1)"%divider)) elif tc == "MXSA": mxBlock = kernel["ProblemType"]["MXBlockA"] module.add(SLShiftRightB32(dst=sgpr(stmp), src=size, shiftHex=log2(mxBlock), comment="(size/%d-1)" %mxBlock)) module.add(SSubU32(dst=sgpr(stmp), src0=sgpr(stmp), src1=0x1, comment="(size/%d-1)" %mxBlock)) elif tc == "MXSB": mxBlock = kernel["ProblemType"]["MXBlockB"] module.add(SLShiftRightB32(dst=sgpr(stmp), src=size, shiftHex=log2(mxBlock), comment="(size/%d-1)" %mxBlock)) module.add(SSubU32(dst=sgpr(stmp), src0=sgpr(stmp), src1=0x1, comment="(size/%d-1)" %mxBlock)) elif tP["isSwizzled"]: module.addModuleAsFlatItems(self.alignTo(stmp, "SizeL", tP["swizzleK"])) module.add(SSubU32(dst=sgpr(stmp), src0=sgpr(stmp), src1=1, comment="SWZ-%s align: (sizeL-1)"%tc)) else: module.add(SSubU32(dst=sgpr(stmp), src0=size, src1=0x1, comment="(size-1)")) else: if tP["isSwizzled"]: if tP["isA"] and idx == kernel["ProblemType"]["Index0"]: module.addModuleAsFlatItems(self.alignTo(stmp, "SizeI", 16)) module.add(SSubU32(dst=sgpr(stmp), src0=sgpr(stmp), src1=1, comment="SWZ-%s align: (sizeM-1)"%tc)) elif tP["isB"] and idx == kernel["ProblemType"]["Index1"]: module.addModuleAsFlatItems(self.alignTo(stmp, "SizeJ", 16)) module.add(SSubU32(dst=sgpr(stmp), src0=sgpr(stmp), src1=1, comment="SWZ-%s align: (sizeN-1)"%tc)) else: module.add(SSubU32(dst=sgpr(stmp), src0=size, src1=0x1, comment="(size-1)")) else: module.add(SSubU32(dst=sgpr(stmp), src0=size, src1=0x1, comment="(size-1)")) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(stmp), sgpr(stmp+1), stride, \ sgpr(stmp), comment="stride x (size-1)")) module.add(SAddU32(dst=sgpr(tensor2dSize0), src0=sgpr(tensor2dSize0), src1=sgpr(stmp+0), comment="sum tensor size")) module.add(SAddCU32(dst=sgpr(tensor2dSize1), src0=sgpr(tensor2dSize1), src1=sgpr(stmp+1), comment="sum tensor size")) # skip ShadowLimit and Srd+2 calculation here in useFixedSrd2 case if not useFixedSrd2: if use64bShadowLimit: limitTmp0 = "ShadowLimit%s+0"%tc limitTmp1 = "ShadowLimit%s+1"%tc else: limitTmp0 = stmp+0 limitTmp1 = stmp+1 module.add(SSubU32(dst=sgpr(limitTmp0), src0=sgpr(tensor2dSize0), src1=sgpr(tileStart+0), comment="sub tileStart")) module.add(SSubBU32(dst=sgpr(limitTmp1), src0=sgpr(tensor2dSize1), src1=sgpr(tileStart+1), comment="sub tileStart")) if use64bShadowLimit: # Set initial buffer limit # if the limit is >64bit, incrementSrd decrements the shadow as the SRD increments, # and when we get within 32-bit we start to step down the SRD # if the limit is <32bits, set it accurately here: # Note lshl_b64 the higher-numbered SGPR has the upper 32-bits module.add(scalarMultiply64Bpe("ShadowLimit%s"%tc, "ShadowLimit%s"%tc, tP["bpeGR"], stmp, "Set limit to use bytes")) if prePad: module.add(SAddU32(dst=sgpr("ShadowLimit%s+0"%tc), src0=sgpr("ShadowLimit%s+0"%tc), src1=prePad, comment="extend limit for pre-pad")) module.add(SAddCU32(dst=sgpr("ShadowLimit%s+1"%tc), src0=sgpr("ShadowLimit%s+1"%tc), src1=0, comment="extend limit for pre-pad")) if kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"]: module.add(SAddU32(dst=sgpr("ShadowLimit%s+0"%tc), src0=sgpr("ShadowLimit%s+0"%tc), src1=self.buff_load_inst_offset_max, comment="extend limit for directToLDS instruction offset")) module.add(SAddCU32(dst=sgpr("ShadowLimit%s+1"%tc), src0=sgpr("ShadowLimit%s+1"%tc), src1=0, comment="extend limit for directToLDS instruction offset")) module.add(SCmpEQU32(src0=sgpr("ShadowLimit%s+1"%tc), src1=0, comment="are we within 2^32?")) module.add(SCSelectB32(dst=sgpr("Srd%s+2"%tc), src0=sgpr("ShadowLimit%s+0"%tc), src1="BufferLimit", comment="Move shadow to real if we are within 2^32")) else: # put limit directly into SRD: module.add(scalarMultiplyBpe("Srd%s+2"%tc, stmp, float(tP["bpeGR"]), comment="Set limit to use bytes")) if prePad: module.add(SAddU32(dst=sgpr("Srd%s+2"%tc), src0=sgpr("Srd%s+2"%tc), src1=prePad, comment="extend limit for pre-pad")) # Apply any high-order address components to the tileStart and eventually the SRD - batch idx for batched gemm wg=2 # TODO - refactor since only WG2 is supported and this is always batch for i in range(1, numDim): idx = indices[i] if idx == kernel["ProblemType"]["Index0"] \ or idx == kernel["ProblemType"]["Index1"] \ or idx in kernel["ProblemType"]["IndicesSummation"] \ or isPackedIndex(kernel, idx): continue # these will be captured in GRO not the SRD (or other summations are always 0) else: assert(wg==2) # can only have one wg2 with a batch. Other dimensions should be packed into wg0/wg1 stride = "Stride%s%s"%(tc,self.states.indexChars[tP['ia'][i]]) stridedBatchedGemmLoad = Label(label=self.labels.getNameInc("StridedBatchedGemmLoad"+tc), comment="Computing the Batch Matrix's base address for Strided Batched GEMM") stridedBatchedGemmLoad_End = Label(label=self.labels.getNameInc("StridedBatchedGemmLoad"+tc+"_End"), comment="End Computing the Batch Matrix's base address for Strided Batched") if kernel["ProblemType"]["SupportUserArgs"]: module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=3, comment="ArgType == 3 for General Batched GEMM")) module.add(SCBranchSCC0(labelName=stridedBatchedGemmLoad.getLabelName())) else: module.add(SBranch(labelName=stridedBatchedGemmLoad.getLabelName())) moduleLoadGeneralBatch.add(SMulI32(dst=sgpr(stmp+0), src0=8, src1=sgpr("WorkGroup2"), comment="Compute Offset into Pointer Array")) if not wroteTileStart: moduleLoadStridedBatch.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tileStart+0), sgpr(tileStart+1), sgpr(stride), sgpr("WorkGroup2"), comment="Stride*WG")) wroteTileStart = True else: moduleLoadStridedBatch.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(stmp+0), sgpr(stmp+1), sgpr(stride), sgpr("WorkGroup2"), comment="Stride*WG")) moduleLoadStridedBatch.add(SAddU32(dst=sgpr(tileStart+0), src0=sgpr(tileStart+0), src1=sgpr(stmp+0), comment="accum wg term to tilestart")) moduleLoadStridedBatch.add(SAddCU32(dst=sgpr(tileStart+1), src0=sgpr(tileStart+1), src1=sgpr(stmp+1), comment="accum wg term to tilestart")) wg+=1 moduleLoadGeneralBatch.add(SCmpEQU32(src0=sgpr("SizesSum"), src1=hex(0), comment="Don't dereference Pointer array if SizesSum == 0")) moduleLoadGeneralBatch.add(SCBranchSCC1(labelName=stridedBatchedGemmLoad_End.getLabelName())) moduleLoadGeneralBatch.add(SAddU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr("Address%s+0"%tc), comment="Offsetting to the location [Lower half of address]")) moduleLoadGeneralBatch.add(SAddCU32(dst=sgpr(stmp+1), src0=sgpr("Address%s+1"%tc), src1=0, comment="Offsetting to the location [Higher half of address]")) moduleLoadGeneralBatch.add(SLoadB64(dst=sgpr("Srd%s"%tc, 2), base=sgpr(stmp, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array")) moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) if self.states.groOffsetInMacroTile and ((tc == "A" and not kernel["enableTDMA"]) or (tc == "B" and not kernel["enableTDMB"])): prePad1 = int(self.states.srdShiftLeft[tc] * tP["bpeGR"]) # leave room in case we have to pointer shift moduleLoadGeneralBatch.add(SSubU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=prePad1, comment="pre-pad to make room for possible pointer shift")) moduleLoadGeneralBatch.add(SSubBU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr("Srd%s+1"%tc), src1=0, comment="pre-pad to make room for possible pointer shift")) # Add the tile start to the SRD if wroteTileStart: moduleLoadStridedBatch.add(scalarMultiply64Bpe(tileStart, tileStart, tP["bpeGR"], stmp, "tileStart")) moduleLoadStridedBatch.add(SAddU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Address%s+0"%tc), src1=sgpr(tileStart+0), comment="SRD base = Address+ tileStart0")) moduleLoadStridedBatch.add(SAddCU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr("Address%s+1"%tc), src1=sgpr(tileStart+1), comment="SRD base = Address+ tileStart1")) moduleLoadGeneralBatch.add(scalarMultiply64Bpe(tileStart, tileStart, tP["bpeGR"], stmp, "tileStart")) moduleLoadGeneralBatch.add(SAddU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr(tileStart+0), src1=sgpr("Srd%s+0"%tc), comment="SRD base = Address+ tileStart0")) moduleLoadGeneralBatch.add(SAddCU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr(tileStart+1), src1=sgpr("Srd%s+1"%tc), comment="SRD base = Address+ tileStart1")) else: moduleLoadStridedBatch.add(SMovB64(dst=sgpr("Srd%s"%tc, 2), src=sgpr("Address%s"%tc, 2), comment="init SRD base address")) moduleLoadGeneralBatch.add(SBranch(labelName = stridedBatchedGemmLoad_End.getLabelName())) moduleLoadGeneralBatch.add(stridedBatchedGemmLoad) module.add(moduleLoadGeneralBatch) # Logic for General Batched GEMM comes first module.add(moduleLoadStridedBatch) # Logic for Strided Batched GEMM comes second module.add(stridedBatchedGemmLoad_End) # self.states.groOffsetInMacroTile == 1 case, pre-pad is already subtracted from AddressA/B if prePad and self.states.groOffsetInMacroTile == 0: module.add(SSubU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=prePad, comment="pre-pad to make room for possible pointer shift")) module.add(SSubBU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr("Srd%s+1"%tc), src1=0, comment="pre-pad to make room for possible pointer shift")) if kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"]: module.add(SSubU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=self.buff_load_inst_offset_max, comment="make room for directToLDS instruction offset")) module.add(SSubBU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr("Srd%s+1"%tc), src1=0, comment="make room for directToLDS instruction offset")) module.add(SMovB32(dst=sgpr("Srd%s+3"%tc), src="Srd127_96", comment="Set bits 127_96 in SRD")) if use64bShadowLimit and not useFixedSrd2: module.add(self.shiftSrd(tc)) #if tP["isB"]: # module.add(self.getCmpAssert(self.asmAssert.ne, sgpr("WorkGroup1"), 0xA)) return module ############################################################################## # Tighten Srd{A,B,MXSA,MXSB}+2 (OOB limit) before the tail loop: # Srd+2 -= (DepthU - roundUp(K_rem, K_pad)) * bpe_K_byteStride # K_pad = 1 for A/B, 256 for MX scale (matches host rearrangePaddedMXScaleLayout). # Precondition: upstream Srd+2 is computeLoadSrd's tile-boundary limit # (useFixedSrd2 + non-isConstUnitStride); PGR/mainloop GR_INC only touch Srd+0/+1. # tPs must share kPad and bpe_K_byteStride (delta emitted once, reused). ############################################################################## def computeTailLoopSrdLimit(self, kernel, tPs): module = Module("computeTailLoopSrdLimit") if not tPs: return module if not kernel.get("UseSubtileImpl"): return module if self.states.groOffsetInMacroTile != 1: return module # TDM uses tensor_load_to_lds descriptors, not SRD buffer loads hasTDM = kernel.get("enableTDMA") and kernel.get("enableTDMB") # Filter out A/B tensors when TDM is active (SrdA/SrdB not allocated) if hasTDM: tPs = [tP for tP in tPs if tP["tensorChar"] not in ("A", "B")] if not tPs: return module MX_PAD_K = 256 depthU = int(kernel["DepthU"]) # Gather per-tc (kPad, bpeForLimit). All tPs must agree, otherwise bail. tcKpadBpeList = [] for tP in tPs: tc = tP["tensorChar"] if tc not in ("A", "B", "MXSA", "MXSB"): return module strideF = self.strideRef(tc, tP['tileIdx']) if self.isConstUnitStride(strideF): return module isMx = tc in ("MXSA", "MXSB") if isMx: tcab = "A" if tc == "MXSA" else "B" mxBlock = int(kernel["ProblemType"].get("MXBlock%s" % tcab, 0)) if mxBlock <= 0: return module # When DepthU <= MX_PAD_K, roundUp(K_rem, 256) >= DepthU so delta <= 0 # always; host padding alone already covers the over-read. if depthU <= MX_PAD_K: return module mxScaleFormat = kernel.get("MXScaleFormat", "NoSwizzle") isMxSwizzledScale = mxScaleFormat in ("InMemorySwizzle", "HostPreSwizzle") swizzleSize0 = 32 if isMxSwizzledScale else 1 if swizzleSize0 % mxBlock != 0: return module bytesPerKElement = swizzleSize0 // mxBlock # 1 for gauntlet mxBlock=32 kPad = MX_PAD_K bpeForLimit = float(bytesPerKElement) else: kPad = 1 bpeForLimit = float(tP["bpeGR"]) tcKpadBpeList.append((tc, kPad, bpeForLimit)) # Joint emission requires uniform kPad/bpe across all tPs (true for A/B # symmetric configs and for MXSA/MXSB scale pairs in the gauntlet). kPad = tcKpadBpeList[0][1] bpeForLimit = tcKpadBpeList[0][2] if any(kp != kPad or bp != bpeForLimit for _, kp, bp in tcKpadBpeList): return module tcs = [t for t, _, _ in tcKpadBpeList] isMx = tcs[0] in ("MXSA", "MXSB") loopCounterName = self.loopCounterName(kernel, self.states.unrollIdx) with self.allocTmpSgpr(1, tag="computeTailLoopSrdLimit_tmpSgprInfo") as tmpSgprInfo: stmp = tmpSgprInfo.idx module.addComment1( "Tighten Srd{%s}+2 for tail loop: -= (DepthU - roundUp(K_rem,%u)) * %s" % (",".join(tcs), kPad, "bpe" if not isMx else ("bpkE=%u" % int(bpeForLimit)))) if kPad > 1: padMaskInv = (~(kPad - 1)) & 0xffffffff module.add(SAddU32(dst=sgpr(stmp), src0=sgpr(loopCounterName), src1=(kPad - 1), comment="K_rem + (K_pad-1) for roundUp")) module.add(SAndB32(dst=sgpr(stmp), src0=sgpr(stmp), src1=hex(padMaskInv), comment="K_rem_aligned = roundUp(K_rem, %u)" % kPad)) module.add(SSubU32(dst=sgpr(stmp), src0=depthU, src1=sgpr(stmp), comment="delta_K = DepthU - K_rem_aligned (>= 0)")) else: # delta_K = DepthU - K_rem (full tail, including odd trailing element # so the boundary DTL load doesn't need to bump Srd+2). module.add(SSubU32(dst=sgpr(stmp), src0=depthU, src1=sgpr(loopCounterName), comment="delta_K = DepthU - K_rem")) # delta_bytes = delta_K * bpe_K_byteStride (handles bpe=0.5 via SShr 1) module.add(scalarMultiplyBpe(stmp, stmp, bpeForLimit, comment="delta_bytes")) for tc in tcs: module.add(SSubU32(dst=sgpr("Srd%s+2" % tc), src0=sgpr("Srd%s+2" % tc), src1=sgpr(stmp), comment="Srd%s+2 -= delta_bytes (tail-loop tight)" % tc)) return module ############################################################################## # Combined A+B tail-loop boundary DTL load. # # For bf16 + subtile + odd K_rem, the regular dwordx4 GR at the K-boundary # straddles an OOB dword; HW zeroes the dword and the trailing 16-bit K # element is lost. This helper discovers the lane that owns that element # (per-lane swizzled vaddrs live in sharedVgprGROffset[i]) and issues a # single-lane buffer_load_ushort lds:1 into the correct LDS slot for both # A and B back-to-back under a single skip label. ############################################################################## def tailLoopBoundaryDtlLoadAB(self, kernel, tPA, tPB): module = Module("tailLoopBoundaryDtlLoadAB") # Gate: both tensors must be eligible. If not, fall back to per-tensor # emission (preserves the existing per-tensor early-out semantics). def _eligible(tP): tc = tP["tensorChar"] if not kernel.get("UseSubtileImpl"): return False if tc not in ("A", "B"): return False if not kernel["ProblemType"]["DataType"].isBFloat16(): return False if self.states.groOffsetInMacroTile != 1: return False strideF = self.strideRef(tc, tP['tileIdx']) if self.isConstUnitStride(strideF): return False tileInfo = self.states.a.tileInfo if tc == 'A' else self.states.b.tileInfo if not tileInfo.localSubtilesRegister: return False # TDM path: no per-lane GR offset registers allocated regList = tileInfo.localSubtilesRegister[0] if len(regList) > 0 and not regList.is_sgpr: return False return True assert _eligible(tPA) and _eligible(tPB), \ "tailLoopBoundaryDtlLoadAB requires bf16/subtile eligibility for both A and B" # Both eligible — emit shared K math once, then per-tensor M/load blocks. loopCounterName = self.loopCounterName(kernel, self.states.unrollIdx) waveSize = kernel["WavefrontSize"] depthU = kernel["DepthU"] laneMaskCount = self.states.laneSGPRCount # K-side geometry is identical for A and B in the bf16 subtile path. bpe = int(tPA["bpeGR"]) elemsPerLane = 16 // bpe # bf16 -> 8 tileInfoA = self.states.a.tileInfo subtileKElems = int(tileInfoA.subtileShape[1]) * int(tileInfoA.mmaTileShape[1]) def _isPow2(n): return n > 0 and (n & (n - 1)) == 0 assert _isPow2(subtileKElems) skipLabel = Label(self.labels.getNameInc("tailBoundarySkipAB"), "") module.addComment1("Tail-loop boundary DTL load (A+B)") # ── Shared block: parity gate + K-derived values ────────────────────────── # Reserve sgprs that must outlive the per-tensor block. # sIntra : intra_K_byte (shared by A and B) # sOffK : offsetK_bytes (shared by A and B) # sKm1 : K_rem - 1 (consumed per-tensor for alignedK_within_subtile) # Parity scratch is short-lived and goes into the per-tensor inner alloc # below, so SCC is set in the same allocation as the skip target. with self.allocTmpSgpr(3, tag="tailLoopBoundaryDtlLoadAB_kHoldInfo") as kHoldInfo: sIntra = kHoldInfo.idx + 0 sOffK = kHoldInfo.idx + 1 sKm1 = kHoldInfo.idx + 2 # Parity gate (uses sKm1 as throwaway scratch before its real value). module.add(SAndB32(dst=sgpr(sKm1), src0=sgpr(loopCounterName), src1=hex(1), comment="K_rem & 1 (SCC=1 if odd)")) module.add(SCBranchSCC0(labelName=skipLabel.getLabelName(), comment="K_rem even -> skip boundary load")) # K_rem_m1, intra_K_byte, offsetK_bytes (shared for A and B) module.add(SSubU32(dst=sgpr(sKm1), src0=sgpr(loopCounterName), src1=1, comment="K_rem_m1 = K_rem - 1")) module.add(SAndB32(dst=sgpr(sIntra), src0=sgpr(sKm1), src1=hex(elemsPerLane - 1), comment="K_rem_m1 & (elemsPerLane-1)")) module.add(SLShiftLeftB32(dst=sgpr(sIntra), shiftHex=log2(bpe), src=sgpr(sIntra), comment="intra_K_byte")) module.add(SAndB32(dst=sgpr(sOffK), src0=sgpr(sKm1), src1=hex(0xFFFFFFFF & ~(subtileKElems - 1)), comment="K_rem_m1 & ~(subtileKElems-1)")) module.add(SLShiftLeftB32(dst=sgpr(sOffK), shiftHex=log2(bpe), src=sgpr(sOffK), comment="offsetK_bytes")) # alignedK_within_subtile = (K_rem_m1 & withinSubMask) * bpe. # withinSubMask = (subtileKElems-1) & ~(elemsPerLane-1) == # (K_rem_m1 & (subtileKElems-1)) with the low elemsPerLane bits dropped. # That equals (K_rem_m1 - K_subtile_floor) with low intra bits dropped, # i.e. sIntra is the dropped part. So: # alignedK_within_subtile = (K_rem_m1 mod subtileKElems) * bpe - sIntra # which is exactly the value used in sTarget — but to keep this change # mechanically equivalent to the prior code we just compute it inline # per-tensor when we build sTarget. # ── Per-tensor block ──────────────────────────────────────────────── for tP in (tPA, tPB): tc = tP["tensorChar"] tileInfo = self.states.a.tileInfo if tc == 'A' else self.states.b.tileInfo regList = tileInfo.localSubtilesRegister[0] hasSoffsetSgpr = (len(regList) > 0) soffsetSgprIdx = regList.indices[0] if hasSoffsetSgpr else None strideF = self.strideRef(tc, tP['tileIdx']) mt = kernel[tP["mt"]] tileIdx = tP["tileIdx"] mElemsPerSubtileRow = int(tileInfo.subtileShape[0]) * int(tileInfo.mmaTileShape[0]) assert _isPow2(mElemsPerSubtileRow) withinSubMask = (subtileKElems - 1) & ~(elemsPerLane - 1) numGR = int(tileInfo.numGRPerSubtile) assert numGR >= 1 and _isPow2(numGR) and (mElemsPerSubtileRow % numGR == 0) mPerGRLoad = mElemsPerSubtileRow // numGR subtileSizeBytes = mElemsPerSubtileRow * subtileKElems * bpe subtileOffsetBytes = subtileSizeBytes // numGR with self.allocTmpSgpr(11, 2, tag="tailLoopBoundaryDtlLoadAB_tmpSgprInfo") as tmpSgprInfo: stmp = tmpSgprInfo.idx # stmp+0 : sTarget -> sLaneId # stmp+1 : alignedK_within_subtile / scratch / m0 base # stmp+2 : numLine / numLine_within_subtile # stmp+4..5 : v_cmp result pair (also sMask reuse before clobber) # stmp+7 : offsetM_glb_bytes # stmp+8 : offsetM_lds_bytes # stmp+9 : offsetM_floor / sChunkSel / sM0Step # stmp+10: sAnyMatch # numLine = min(SizeI/J - WG*MT, MT) - 1 module.add(SMulI32(dst=sgpr(stmp+2), src0=sgpr(tP["wg"]), src1=mt, comment="WG%s * MT(%u)" % (tP["wg"][-1], mt))) module.add(SSubU32(dst=sgpr(stmp+2), src0=self.sizeRef(tileIdx), src1=sgpr(stmp+2), comment="numToEnd = Size%s - WG*MT" % self.states.indexChars[tileIdx])) module.add(SMinU32(dst=sgpr(stmp+2), src0=sgpr(stmp+2), src1=mt, comment="min(numToEnd, MT)")) module.add(SSubU32(dst=sgpr(stmp+2), src0=sgpr(stmp+2), src1=1, comment="numLine = min - 1 (0-based)")) # alignedK_within_subtile = (K_rem_m1 & withinSubMask) * bpe # We re-derive it from the shared sIntra + sOffK chain: # K_rem_m1 = sOffK/bpe (high) + (alignedK)/bpe + sIntra/bpe (low) # but cheapest is to just compute it fresh from sKm1 — sKm1 is dead # outside this inner block. Reuse stmp+1. module.add(SAndB32(dst=sgpr(stmp+1), src0=sgpr(sKm1), src1=hex(withinSubMask), comment="K_rem_m1 & withinSubMask (=%d)" % withinSubMask)) module.add(SLShiftLeftB32(dst=sgpr(stmp+1), shiftHex=log2(bpe), src=sgpr(stmp+1), comment="alignedK_within_subtile")) # offsetM_floor = numLine & ~(mEPSR-1) module.add(SAndB32(dst=sgpr(stmp+9), src0=sgpr(stmp+2), src1=hex(0xFFFFFFFF & ~(mElemsPerSubtileRow - 1)), comment="offsetM_floor = numLine & ~(mEPSR-1)")) # offsetM_glb_bytes = offsetM_floor * strideF * bpe module.add(SMulI32(dst=sgpr(stmp+7), src0=sgpr(stmp+9), src1=strideF, comment="offsetM_floor * strideF (elements)")) module.add(scalarMultiplyBpe(stmp+7, stmp+7, float(bpe), comment="offsetM_glb_bytes = * bpe")) # LDS layout (matches emitSingleBufferLoad): # m0_subtile_base = (sId0 + sId1*globalSubtileGrid[0]) * subtileSize # where sId0 = offsetM_floor/mEPSR, sId1 = K_floor/subtileKElems, # and subtileSize = mEPSR * subtileKElems * bpe. # Therefore: # M-row part : offsetM_floor * subtileKElems * bpe # K-band part: kSubtileIdx * globalSubtileGrid[0] * subtileSizeBytes mRowStrideBytes = subtileKElems * bpe kBandStrideBytes = int(tileInfo.globalSubtileGrid[0]) * subtileSizeBytes # offsetM_lds_bytes (M-row part) = offsetM_floor * subtileKElems * bpe if _isPow2(mRowStrideBytes): module.add(SLShiftLeftB32(dst=sgpr(stmp+8), shiftHex=log2(mRowStrideBytes), src=sgpr(stmp+9), comment="offsetM_lds (M-row) = offsetM_floor * subtileKElems*bpe")) else: module.add(SMulI32(dst=sgpr(stmp+8), src0=sgpr(stmp+9), src1=hex(mRowStrideBytes), comment="offsetM_lds (M-row) = offsetM_floor * subtileKElems*bpe")) # Add K-band contribution: kSubtileIdx * globalSubtileGrid[0] * subtileSize. # kSubtileIdx = sOffK / (subtileKElems*bpe); stmp+9 is free here # (chunkSel overwrites it below). if kBandStrideBytes != 0: module.add(SLShiftRightB32(dst=sgpr(stmp+9), shiftHex=log2(subtileKElems * bpe), src=sgpr(sOffK), comment="kSubtileIdx = offsetK_bytes >> log2(subtileKElems*bpe)")) if _isPow2(kBandStrideBytes): module.add(SLShiftLeftB32(dst=sgpr(stmp+9), shiftHex=log2(kBandStrideBytes), src=sgpr(stmp+9), comment="kBand_bytes = kSubtileIdx * GSG[0] * subtileSize")) else: module.add(SMulI32(dst=sgpr(stmp+9), src0=sgpr(stmp+9), src1=hex(kBandStrideBytes), comment="kBand_bytes = kSubtileIdx * GSG[0] * subtileSize")) module.add(SAddU32(dst=sgpr(stmp+8), src0=sgpr(stmp+8), src1=sgpr(stmp+9), comment="offsetM_lds_bytes += kBand_bytes")) # numLine_within_subtile = numLine & (mEPSR - 1) module.add(SAndB32(dst=sgpr(stmp+2), src0=sgpr(stmp+2), src1=hex(mElemsPerSubtileRow - 1), comment="numLine_within_subtile")) sChunkSel = stmp + 9 if numGR > 1: module.add(SLShiftRightB32(dst=sgpr(sChunkSel), shiftHex=log2(mPerGRLoad), src=sgpr(stmp+2), comment="chunkSel = numLine_ws >> log2(mPerGRLoad)")) module.add(SAndB32(dst=sgpr(sChunkSel), src0=sgpr(sChunkSel), src1=hex(numGR - 1), comment="chunkSel &= (numGR-1)")) # sTarget = numLine_ws * stride * bpe + alignedK - soffset module.add(SMulI32(dst=sgpr(stmp+0), src0=sgpr(stmp+2), src1=strideF, comment="numLine_ws * stride (elements)")) module.add(scalarMultiplyBpe(stmp+0, stmp+0, float(bpe), comment="-> bytes")) module.add(SAddU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr(stmp+1), comment="+ alignedK_within_subtile")) if hasSoffsetSgpr: module.add(SSubU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr(soffsetSgprIdx), comment="- soffset (per-lane vaddr space)")) # Precompute the m0/vaddr add-bases shared across all chunks. # sM0Base = intra + offsetK + offsetM_lds # sVaddrBase = intra + offsetK + offsetM_glb # Reuse stmp+1 / stmp+7 to avoid extra sgprs (alignedK is dead now; # offsetM_glb_bytes folds into sVaddrBase in place). sM0Base = stmp + 1 sVaddrBase = stmp + 7 # sM0Base = intra + offsetM_lds_bytes # (offsetM_lds_bytes already includes the K-band term; # sOffK belongs only in sVaddrBase — m0 encodes K-bands via the # GSG[0]*subtileSize stride, not via sId1*subtileKElems*bpe.) module.add(SAddU32(dst=sgpr(sM0Base), src0=sgpr(sIntra), src1=sgpr(stmp+8), comment="sM0Base = intra + offsetM_lds_bytes")) module.add(SAddU32(dst=sgpr(sVaddrBase), src0=sgpr(sVaddrBase), src1=sgpr(sIntra), comment="sVaddrBase = offsetM_glb + intra")) module.add(SAddU32(dst=sgpr(sVaddrBase), src0=sgpr(sVaddrBase), src1=sgpr(sOffK), comment=" + offsetK")) sCmpLo = stmp + 4 sCmpHi = stmp + 5 sLaneId = stmp + 0 sAnyMatch = stmp + 10 sMask = stmp + 4 sM0Step = stmp + 9 mubuf = MUBUFModifiers(offen=True, offset12=0, glc=False, slc=False, nt=False, lds=True) ldSoffset = sgpr(soffsetSgprIdx) if hasSoffsetSgpr else 0 vTmp = self.vgprPool.checkOut(1, "tailBoundaryVTmp%s" % tc) if numGR == 1: vaddrVgpr = tileInfo.sharedVgprGROffset[0] else: vMerged = self.vgprPool.checkOut(1, "tailBoundaryVMerged%s" % tc) module.add(VMovB32(dst=vgpr(vMerged), src=vgpr(tileInfo.sharedVgprGROffset[0]), comment="vMerged = chunk-0 vaddrs")) for i in range(1, numGR): module.add(SCmpEQU32(src0=sgpr(sChunkSel), src1=hex(i), comment="chunkSel == %d ?" % i)) module.add(SCSelectB64(dst=sgpr(sMask, laneMaskCount), src0=-1, src1=0, comment="sMask = all-ones if chunk %d" % i)) module.add(VCndMaskB32(dst=vgpr(vMerged), src0=vgpr(vMerged), src1=vgpr(tileInfo.sharedVgprGROffset[i]), src2=sgpr(sMask, laneMaskCount), comment="vMerged = chunk %d vaddrs" % i)) vaddrVgpr = vMerged sM0StepIsZero = (numGR == 1) if not sM0StepIsZero: module.add(SLShiftLeftB32(dst=sgpr(sM0Step), shiftHex=log2(subtileOffsetBytes), src=sgpr(sChunkSel), comment="sM0Step = chunkSel * subtileOffsetBytes")) # v_cmp + ff1 -> sLaneId (wave64 only for now). assert waveSize == 64 module.add(VCmpEQU32(dst=sgpr(sCmpLo, 2), src0=sgpr(stmp+0), src1=vgpr(vaddrVgpr), comment="find lane whose vaddr == sTarget")) module.add(SOrB32(dst=sgpr(sAnyMatch), src0=sgpr(sCmpLo), src1=sgpr(sCmpHi), comment="sAnyMatch = vcc_lo | vcc_hi")) # Use stmp+2 as scratch for ff1(hi)+32 (numLine_ws is dead). module.add(SFf1B32(dst=sgpr(stmp+2), src=sgpr(sCmpHi), comment="ff1(vcc_hi)")) module.add(SAddU32(dst=sgpr(stmp+2), src0=sgpr(stmp+2), src1=32, comment="+ 32 (hi-half lane offset)")) module.add(SFf1B32(dst=sgpr(sLaneId), src=sgpr(sCmpLo), comment="ff1(vcc_lo)")) module.add(SCmpEQI32(src0=sgpr(sLaneId), src1=-1, comment="SCC=1 if lo had no match")) module.add(SCSelectB32(dst=sgpr(sLaneId), src0=sgpr(stmp+2), src1=sgpr(sLaneId), comment="sLaneId = lo==-1 ? hi+32 : lo")) # m0 = LocalWriteBaseAddr + sLaneId*16 + sM0Base (+ sM0Step) # # sM0Base already holds intra + offsetK + offsetM_lds; one add for # sLaneId*16, optional add for sM0Step, then fold LWBA into m0. # We do sLaneId<<4 into a scratch (stmp+4 is now dead — sCmpLo). sShifted = stmp + 4 module.add(SLShiftLeftB32(dst=sgpr(sShifted), shiftHex=log2(16), src=sgpr(sLaneId), comment="sLaneId * 16 (per-lane LDS stride)")) module.add(SAddU32(dst=sgpr(sShifted), src0=sgpr(sShifted), src1=sgpr(sM0Base), comment="+ sM0Base (intra + offsetK + offsetM_lds)")) if not sM0StepIsZero: module.add(SAddU32(dst=sgpr(sShifted), src0=sgpr(sShifted), src1=sgpr(sM0Step), comment="+ sM0Step")) module.add(SAddU32(dst=mgpr(0), src0=sgpr("LocalWriteBaseAddr%s" % tc), src1=sgpr(sShifted), comment="m0 = LocalWriteBaseAddr%s + ..." % tc)) # vaddr = v[vaddrVgpr][sLaneId] + sVaddrBase sVaddrTmp = stmp + 4 # reused (sShifted dead) module.add(VReadlaneB32(dst=sgpr(sVaddrTmp), src0=vgpr(vaddrVgpr), src1=sgpr(sLaneId), comment="sTarget = v[vaddr][sLaneId]")) module.add(SAddU32(dst=sgpr(sVaddrTmp), src0=sgpr(sVaddrTmp), src1=sgpr(sVaddrBase), comment="+ sVaddrBase (intra + offsetK + offsetM_glb)")) module.add(VMovB32(dst=vgpr(vTmp), src=sgpr(sVaddrTmp), comment="vaddr")) # EXEC gate, DTL load, wait, restore EXEC module.add(SCmpEQU32(src0=sgpr(sAnyMatch), src1=0, comment="SCC=1 if no matching lane on this wave")) module.add(SCSelectB64(dst=EXEC(), src0=0, src1=1, comment="EXEC = match ? 1 : 0")) module.add(BufferLoadU16(dst=None, vaddr=vgpr(vTmp), saddr=sgpr("Srd%s" % tc, 4), soffset=ldSoffset, mubuf=mubuf, comment="boundary 16-bit DTL load -> LDS@m0")) module.add(SMovB64(dst=EXEC(), src=-1, comment="restore EXEC = full")) if numGR > 1: self.vgprPool.checkIn(vMerged) self.vgprPool.checkIn(vTmp) module.add(skipLabel) return module ############################################################################## # Global Read Addresses: Addresses A/B ############################################################################## def graAddresses(self, kernel, tP): module = Module("graAddresses") tc = tP["tensorChar"] graIdx = 0 if tP["isSwizzled"]: # "StrideA0I" or "StrideB1J" strideName = "Stride%s%s"%(tc,self.states.indexChars[tP["idx"]]) module.addModuleAsFlatItems(self.setTo(strideName, "SizeL")) module.addModuleAsFlatItems(self.alignTo(strideName, strideName, tP["swizzleK"])) if kernel["BufferLoad"]: # maxAddrSgpr = size[n] * stride[n-1] module.addComment0("max read offset = size[n] * stride[n-1]") module.add(self.computeLoadSrd(kernel, tP, tc, kernel["ProblemType"]["IndexAssignments%s"%tc], tP["bpeGR"])) if kernel["ProblemType"]["Sparse"] and kernel["DirectToVgprSparseMetadata"]: if (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]): module.add(self.computeMetaDataSrd(kernel, tP, tc, kernel["ProblemType"]["IndexAssignments%s"%tc])) else: tmp = self.vgprPool.checkOut(2, tag="graAddresses_tmp", preventOverflow=self.states.preventVgprOverflowDuringNewTile) skComponent = Component.StreamK.find(self) module.add(skComponent.graAddresses(self, kernel, tP, tmp)) for perp in range(0, tP["nrp"]): for sPerp in range(0, tP["nrpv"]): for para in range(0, tP["nrc"]): for sPara in range(0, tP["nrcv"]//tP["nrcvpi"]): comment = "gRA%s_%u_%u_%u_%u = addr%s+grO%s_%u_%u_%u_%u" \ % (tP["tensorChar"], para, sPara, perp, sPerp, \ tP["tensorChar"], tP["tensorChar"], \ para, sPara, perp, sPerp ) module.add(VAddCOU32( dst=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ src1=vgpr(tmp+0), \ comment=comment+" (lower)")) module.add(VAddCCOU32( dst=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ src1=vgpr(tmp+1), \ src2=VCC(), \ comment=comment+" (upper)")) graIdx += self.states.rpga self.vgprPool.checkIn(tmp) return module ############################################################################## # Global Read Addresses: Increments # Define graIncrements, called once for each summation ############################################################################## def graIncrements(self, kernel, loopIdx, tP): module = Module("graIncrements") tc = tP["tensorChar"] dimIdx = kernel["ProblemType"]["IndicesSummation"][loopIdx] # dimension index loopChar = self.states.indexChars[dimIdx] stride = self.strideRef(tc, dimIdx) isMirrorIdx = dimIdx in kernel["ProblemType"]["MirrorDims%s"%tc] #print (tc, ": loopIdx=", loopIdx, "dimIdx=", dimIdx, "strideIdx=", strideIdx) assert(self.states.unrollIdx == kernel["ProblemType"]["NumIndicesSummation"]-1) if loopIdx==self.states.unrollIdx: if self.states.globalReadIncsUseVgpr: with self.allocTmpSgpr(3, tag="graIncrements_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx gsuSgpr = tmpSgpr + 2 du = kernel["_DepthU%s"%tc] duBpe = int(du * tP["bpeGR"]) module.add(SAndB32(dst=sgpr(tmpSgpr), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SMulI32(dst=sgpr(gsuSgpr), src0=sgpr(tmpSgpr), src1=duBpe, comment="GSU*_DepthUTc*Bpe")) module.add(SAndB32(dst=sgpr(tmpSgpr), src0=sgpr("GSU"), src1=hex(0x8000), comment="SCC = (GSUC == 1) ?")) module.add(SCMovB32(dst=sgpr(gsuSgpr), src=duBpe, comment="DepthU*Bpe if GSUC = 1")) module.add(SMulI32(dst=sgpr(tmpSgpr+0), src0=sgpr(gsuSgpr), src1=stride, \ comment="incr%s%s = %s*DepthU*bpeGR (unrollIdx)"%(tc, loopChar, stride) )) # TODO - this should be mul-H?? module.add(SMovB32( dst=sgpr(tmpSgpr+1), \ src=0, \ comment="(carry)")) module.add(VMovB32( dst=vgpr("GlobalReadIncs%s+%u+0"%(tc, 2*loopIdx)), \ src=sgpr(tmpSgpr+0))) module.add(VMovB32( dst=vgpr("GlobalReadIncs%s+%u+1"%(tc, 2*loopIdx)), \ src=sgpr(tmpSgpr+1))) else: # not globalReadIncsUseVgpr, ie use SGPR gsuComponent = Component.GSU.find(self) module.add(gsuComponent.graIncrements(self, kernel, loopIdx, tP)) else: # other summation if self.states.globalReadIncsUseVgpr: printExit("NumIndicesSummation=%u not yet supported in assembly unless globalReadIncsUseVgpr==0" \ % kernel["ProblemType"]["NumIndicesSummation"] ) else: graInc = "GlobalReadIncs%s+%u"%(tc, loopIdx) # subtract increments done by the inner iterations # may be negative: loopIdxPrev = loopIdx + 1 dimIdxPrev = kernel["ProblemType"]["IndicesSummation"][loopIdxPrev] # dimension index loopCharPrev = self.states.indexChars[dimIdxPrev] stridePrev = self.strideRef(tc, dimIdxPrev) isMirrorIdxPrev = dimIdxPrev in kernel["ProblemType"]["MirrorDims%s"%tc] module.addComment1("compute globalReadInc for higher-level loop") with self.allocTmpSgpr(3, tag="graIncrements_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx # Summations always appear in both A and B, can compute number of iterations just once: if loopIdxPrev==self.states.unrollIdx: loopCounterName= self.loopCounterName(kernel, self.states.unrollIdx) if tP["isA"]: quotient = loopCounterName dividend = "SizesSum+%u"%self.states.unrollIdx divisor = kernel["DepthU"] module.add(scalarStaticDivideAndRemainder(quotient, None, dividend, \ divisor, tmpSgprInfo, 0)) gsuComponent = Component.GSU.find(self) if kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1: module.add(gsuComponent.calculateLoopNumIterGsu(self, kernel, loopCounterName, tmpSgprInfo)) module.add(gsuComponent.graIncrementsRestore(self, kernel, loopCounterName)) module.add(SMulI32(dst=sgpr(graInc), src0=stridePrev, src1=sgpr(loopCounterName), \ comment="tmp <- stride%s%s * myWgUnrollIters" %(tc, loopCharPrev))) else: module.add(SMulI32(dst=sgpr(graInc), src0=stridePrev, src1=self.sizeRef(dimIdxPrev), \ comment="tmp <- stride%s%s * size%s%s" %(tc, loopCharPrev, tc, loopCharPrev))) # subtract amount that previous inner loop will have already incremented: # graInc is used as temp for the prev loop calc if isMirrorIdx and isMirrorIdxPrev: module.add(SSubI32(dst=sgpr(graInc), \ src0=sgpr(graInc), \ src1=stride, \ comment="incr%s%s = - stride%s%s"%(tc, loopChar, tc, loopChar) )) elif isMirrorIdx: module.add(SAddI32(dst=sgpr(graInc), \ src0=stride, \ src1=sgpr(graInc), \ comment="incr%s%s = stride%s%s + "%(tc, loopChar, tc, loopChar) )) module.add(SSubI32(dst=sgpr(graInc), \ src0=0, \ src1=sgpr(graInc), \ comment="incr%s%s = - (stride%s%s + )"%(tc, loopChar, tc, loopChar) )) elif isMirrorIdxPrev: module.add(SAddI32(dst=sgpr(graInc), \ src0=stride, \ src1=sgpr(graInc), \ comment="incr%s%s = stride%s%s + "%(tc, loopChar, tc, loopChar) )) else: module.add(SSubI32(dst=sgpr(graInc), \ src0=stride, \ src1=sgpr(graInc), \ comment="incr%s%s = stride%s%s - "%(tc, loopChar, tc, loopChar) )) module.add(SLShiftLeftB32( dst=sgpr(graInc), \ src=sgpr(graInc), \ shiftHex="BpeGR%sLog2"%tc, comment="<- scale by bpeDS")) return Module("graIncrements (Empty)") if self.dontAppendCode else module ############################################################################## # Local Write Addresses: Tile Assignment A/B ############################################################################## def lwaTileAssignment(self, kernel, tP): module = Module("lwaTileAssignment") tc = tP["tensorChar"] divisorName = tP["lvc"] divisor = kernel[divisorName] # DTV case, use tlu path isDTVAB = (tc in ("A", "B", "MXSA", "MXSB")) and kernel["DirectToVgpr%s"%tc] isTr = (tc == "A" or tc == "B") and kernel["enableGLTr%s"%tc] isSwizzledOrTr = tP["isSwizzled"] or isTr swizzledOrTrName = "" if tP["isSwizzled"]: swizzledOrTrName = "SwizzleTensor%s" % tc elif isTr: swizzledOrTrName = "GLTr%s"%tc # swizzled goes to else if (tP["tlu"] or isDTVAB) and (not isSwizzledOrTr): rReg = self.vgprPool.checkOut(1, "lwaTA rReg0", self.states.preventVgprOverflowDuringNewTile) # tile = serial%divisor qReg = self.vgprPool.checkOut(1, "lwaTA qReg0", self.states.preventVgprOverflowDuringNewTile) # unroll = serial/divisor tReg = rReg uReg = qReg tOpStr = "%" uOpStr = "/" else: qReg = self.vgprPool.checkOut(1, 'lwaTA qReg1', self.states.preventVgprOverflowDuringNewTile) # tile = serial/divisor rReg = self.vgprPool.checkOut(1, 'lwaTA rReg1', self.states.preventVgprOverflowDuringNewTile) # unroll = serial%divisor tReg = qReg uReg = rReg tOpStr = "/" uOpStr = "%" module.addComment0("%s = %u" % (divisorName, kernel[divisorName])) module.addComment0("%s = %s-unroll = serial%s%s" \ % (vgpr(uReg), tc, uOpStr, divisorName) ) tmpVgpr = self.vgprPool.checkOutAligned(2, 2, 'lwaTA vgpr', self.states.preventVgprOverflowDuringNewTile) tmpVgprRes = ContinuousRegister(tmpVgpr, 2) dividendReg = "Serial" # local serial if kernel["WaveSeparateGlobalRead%s"%tc]: dividendReg = self.vgprPool.checkOut(1, "idInWave", self.states.preventVgprOverflowDuringNewTile) dummy = self.vgprPool.checkOut(1, "dummy", self.states.preventVgprOverflowDuringNewTile) with self.allocTmpSgpr(1, tag="lwaTileAssignment_tmpSgprInfo") as tmpSgprInfo: module.add(vectorStaticRemainder(dummy, dividendReg, "Serial", kernel["WavefrontSize"], tmpVgprRes, tmpSgprInfo)) # store DirectToVgpr K interval for later use dtvKInterval = 1 # swizzle if isSwizzledOrTr: module.addComment0("TileAssignment for DirectToVgpr%s and %s" % (tc, swizzledOrTrName)) module.add(vectorStaticDivideAndRemainder(qReg, rReg, dividendReg, kernel["WavefrontSize"], tmpVgprRes)) # Calc numKr if tP["isSwizzled"]: tmp = self.sgprPool.checkOut(1, tag="lwaTileAssignment_isSwizzled_tmp") numKr = sgpr(tmp) swzStride = tP["swizzleK"] module.addComment(f"Align to {swzStride}") module.add(SAddU32(numKr, sgpr("SizesSum"), swzStride-1)) module.add(SLShiftRightB32(dst=numKr, shiftHex=hex(log2(swzStride)), src=numKr, comment="%s: numKr = DimK / %s"%(swizzledOrTrName, swzStride))) elif isTr: numKr = kernel["MatrixInstK"] // tP["glvw"] tmp = self.vgprPool.checkOut(2, tag="lwaTileAssignment_isTr_tmp") with self.allocTmpSgpr(1, tag="lwaTileAssignment_isTr_tmpSgprInfo") as tmpSgprInfo: module.add(vectorStaticDivide(tmp, rReg, int(tP["bpeGR"]*8), tmpVgprRes, comment="%s: %s = %s / %s"%(swizzledOrTrName, vgpr(tmp), vgpr(rReg), int(tP["bpeGR"]*8)))) module.add(vectorStaticMultiply(vgpr(tmp), vgpr(tmp), int(tP["bpeGR"]*4), tmpSgprInfo, comment="%s: %s = %s * %s"%(swizzledOrTrName, vgpr(tmp), vgpr(tmp), int(tP["bpeGR"]*4)))) module.add(vectorStaticRemainder(tmp+1, rReg, rReg, int(tP["bpeGR"]*4), tmpVgprRes, tmpSgprInfo, comment="%s: %s = %s %% %s"%(swizzledOrTrName, vgpr(tmp), vgpr(rReg), int(tP["bpeGR"]*4)))) module.add(VAddU32(dst=vgpr(rReg), src0=vgpr(tmp), src1=vgpr(rReg), comment="%s: %s = %s + %s"%(swizzledOrTrName, vgpr(rReg), vgpr(rReg), vgpr(tmp)))) WvG_M = kernel["MIWaveGroup"][0] if tP["isA"]: module.add(VAndB32(dst=vgpr(qReg), src0=hex(WvG_M-1), src1=vgpr(qReg), comment="%s: wave_id (along_M) %%= MIWG[0]"%(swizzledOrTrName))) module.add(VMulU32U24(dst=vgpr(qReg), src0=numKr, src1=vgpr(qReg), comment="%s: wave_id (along_M) *= numKr"%(swizzledOrTrName))) elif tP["isB"]: # NB: # Calc of w_id is: /= MIWG[0], not %= MIWG[1] module.add(VLShiftRightB32(dst=vgpr(qReg), shiftHex=log2(WvG_M), src=vgpr(qReg), comment="%s: wave_id (along_N) /= MIWG[0]"%(swizzledOrTrName))) if kernel["LocalSplitU"] > 1: # LSU for swizzleB: make sure wave_id is in the range of MIWG[1], so need extra %=MIWG[1] WvG_N = kernel["MIWaveGroup"][1] module.add(VAndB32(dst=vgpr(qReg), src0=hex(WvG_N-1), src1=vgpr(qReg), comment="%s: LSU Case: wave_id (along_N) %%= MIWG[1]"%(swizzledOrTrName))) module.add(VMulU32U24(dst=vgpr(qReg), src0=numKr, src1=vgpr(qReg), comment="%s: wave_id (along_N) *= numKr"%(swizzledOrTrName))) if tP["isSwizzled"]: self.sgprPool.checkIn(tmp) elif isTr: module.add(VBfeU32(dst=vgpr(tmp), src0=vgpr(dividendReg), src1=int(tP["bpeGR"])+1, src2=1, comment="%s: offset for the right half of the tile"%(swizzledOrTrName))) module.add(VAddU32(dst=vgpr(qReg), src0=vgpr(tmp), src1=vgpr(qReg), comment="%s: wave_id += offset for the right half of the tile"%(swizzledOrTrName))) self.vgprPool.checkIn(tmp) elif isDTVAB: # offset calculation for DirectToVgpr # call function from LraTileAssignmentMFMA for DirectToVgpr module.addComment0("TileAssignment for DirectToVgpr%s" % tc) component = Component.LraTileAssignment.find(self) module.add(component.LraTileAssignmentCode(self, kernel, tP, tReg, uReg, tmpVgprRes, dividendReg=dividendReg, isDTVAB=True)) # The other side of lrvw if tP["isA"]: # the other is B tluOther = kernel["ProblemType"]["TLUB"] if tluOther: lrvwOther = self.states.lrvwTileB else: lrvwOther = self.states.lrvwUnrollB elif tP["isB"]: # the other is A tluOther = kernel["ProblemType"]["TLUA"] if tluOther: lrvwOther = self.states.lrvwTileA else: lrvwOther = self.states.lrvwUnrollA elif tP["isMXSA"]: # the other is B tluOther = kernel["ProblemType"]["TLUB"] if tluOther: lrvwOther = self.states.lrvwTileMXSB else: lrvwOther = self.states.lrvwUnrollMXSB elif tP["isMXSB"]: # the other is A tluOther = kernel["ProblemType"]["TLUA"] if tluOther: lrvwOther = self.states.lrvwTileMXSA else: lrvwOther = self.states.lrvwUnrollMXSA else: raise RuntimeError(f"invalid tc {tc}") if lrvwOther >= 2 and (not tluOther) and tP["tlu"]: # DirectToVgpr + LocalReadVectorWidth>=2 case, multiply qReg by lrvwOther dtvKInterval = int(lrvwOther) if tluOther and tP["tlu"]: # DirectToVgpr + both TLU case, multiply qReg by kernel["MIInputPerThread"] dtvKInterval = kernel["MIInputPerThread"] module.add(vectorStaticMultiply(vgpr(qReg), vgpr(qReg), dtvKInterval, None)) # DTV+localSplitU case. Calculate LSU offset here if kernel["LocalSplitU"] > 1: # allocate resources wave_id = self.vgprPool.checkOut(1, tag="lwaTileAssignment_LocalSplitU_wave_id") # quotient # constant lsu = kernel["LocalSplitU"] du = kernel["_DepthU%s"%tc] lsuStride = du // lsu numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] # generate instruction module.add(vectorStaticDivide(wave_id, "Serial", kernel["WavefrontSize"] * numWaves, tmpVgprRes, comment="LSU offset: Get LSU wave_id")) module.add(VMovB32(dst=vgpr(tmpVgpr), src=hex(lsuStride))) module.add(VMulLOU32(dst=vgpr(wave_id), src0=vgpr(tmpVgpr), src1=vgpr(wave_id), \ comment="LSU offset: lsuoffset = wave_id*lsuStride(%u)" % (lsuStride))) module.add(VAddU32(dst=vgpr(qReg), src0=vgpr(wave_id), src1=vgpr(qReg), \ comment="LSU Offset: offset += lsuoffset" )) self.vgprPool.checkIn(wave_id) else: module.add(vectorStaticDivideAndRemainder(qReg, rReg, dividendReg, divisor, tmpVgprRes)) if kernel["WaveSeparateGlobalRead%s"%tc] == 1: with self.allocTmpSgpr(1, tag="lwaTileAssignment_WaveSeparateGlobalRead1_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx module.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr), src=vgpr("Serial"), comment="WaveIdxWavefrontWidth")) module.add(SLShiftRightB32(dst=sgpr(tmpSgpr), src=sgpr(tmpSgpr), shiftHex=hex(log2(kernel["WavefrontSize"])), comment="WaveId")) module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=(kernel[tP["lsp"]] * tP["nrp"]), \ comment="Each wave loads continuous lsp(%u)*nrp(%u) columns" % (kernel[tP["lsp"]], tP["nrp"]))) module.add(VAddU32(dst=vgpr(qReg), src0=sgpr(tmpSgpr), src1=vgpr(qReg), \ comment="Add back to column index")) self.vgprPool.checkIn(dividendReg) self.vgprPool.checkIn(dummy) elif kernel["WaveSeparateGlobalRead%s"%tc] == 2: module.add(VLShiftRightB32(vgpr(dividendReg), hex(log2(kernel["WavefrontSize"])), vgpr("Serial"), "WaveID")) module.add(VMovB32(vgpr(dummy), kernel["NumLoadsPerpendicular%s"%tc]*kernel["NumThreads"]//kernel["WavefrontSize"], comment="Global Read Wave: add back to cloumn index")) module.add(VMulLOU32(vgpr(qReg), vgpr(dummy), vgpr(qReg), "Global Read Wave: add back to cloumn index")) module.add(VAddU32(vgpr(qReg), vgpr(dividendReg), vgpr(qReg), "Global Read Wave: add back to cloumn index")) self.vgprPool.checkIn(dividendReg) self.vgprPool.checkIn(dummy) with self.allocTmpSgpr(1, tag="lwaTileAssignment_tmpSgprInfo") as tmpSgprInfo: if tP["glvw"] > 1: if tP["tlu"]: module.addComment0("tile *= glvw") module.add(vectorStaticMultiply(vgpr(tReg), vgpr(tReg), tP["glvw"], tmpSgprInfo)) else: module.addComment0("unroll *= glvw") module.add(vectorStaticMultiply(vgpr(uReg), vgpr(uReg), tP["glvw"], tmpSgprInfo)) uReg2 = self.vgprPool.checkOut(1, "uReg2", self.states.preventVgprOverflowDuringNewTile) module.add(VMovB32(dst=vgpr(uReg2), src=vgpr(uReg), comment="copy for GlobalSplitU")) tP["gpr"]["uReg2"] = uReg2 tP["gpr"]["lwoT"] = tReg tP["gpr"]["uReg"] = uReg self.vgprPool.checkIn(tmpVgpr) # store DirectToVgpr K interval for later use if tP["isA"]: self.states.dtvKIntervalA = dtvKInterval elif tP["isB"]: self.states.dtvKIntervalB = dtvKInterval return module ############################################################################## # Local Write Addresses: Unroll Assignment A/B ############################################################################## def lwaUnrollAssignment(self, kernel, tP): module = Module("lwaUnrollAssignment") uReg = tP["gpr"]["uReg2"]# if kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1 else "uReg"] module.addComment0("lwaUnrollAssignment%s = %s" % (tP["tensorChar"], vgpr(uReg))) return module ############################################################################## # Local Write Addresses: First Offset A/B ############################################################################## def lwaFirstOffset(self, kernel, tP): module = Module("lwaFirstOffset") tc = tP["tensorChar"] LdsPad = kernel["LdsPad%s"%tc] if kernel["LdsBlockSizePerPad%s"%tc] == 0 else 0 #"lwFOA = lwA%s + lwA%s*MT%s" \ # % (tP["tileChar"], self.states.unrollChar, tP["tileChar"]) uReg = tP["gpr"]["uReg2"] # if kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1 else "uReg"] if kernel["LocalWriteUseSgpr%s"%tc]: destVgpr = self.vgprPool.checkOut(1, "destVgpr", self.states.preventVgprOverflowDuringNewTile) else: destVgpr = "LocalWriteAddr%s"%tc if kernel["UnrollMajorLDS%s" % tc]: lds_stride = kernel["_DepthU%s"%tc] + LdsPad module.add(VMulU32U24(dst=vgpr(destVgpr), src0=hex(lds_stride), src1=vgpr(tP["gpr"]["lwoT"]), \ comment="lw%s%s**(DepthU_Compute + PAD)"%(tc, self.states.unrollChar))) module.add(VAddU32(dst=vgpr(destVgpr), src0=vgpr(uReg), src1=vgpr(destVgpr), \ comment="lwFO%s = (lw%s%s + lw%s%s*(DepthU+PAD))" % (tc, tc, tc, tc, self.states.unrollChar) )) else: lds_stride = kernel["MacroTile%s"%tc] + LdsPad module.add(VMulU32U24(dst=vgpr(destVgpr), src0=hex(lds_stride), src1=vgpr(uReg), \ comment="lw%s%s**(MT%s + PAD)"%(tc, self.states.unrollChar, tc))) module.add(VAddU32(dst=vgpr(destVgpr), src0=vgpr(tP["gpr"]["lwoT"]), src1=vgpr(destVgpr), \ comment="lwFO%s = (lw%s%s + lw%s%s*(MT%s+PAD))" % (tc, tc, tc, tc, self.states.unrollChar, tP["tileChar"]) )) module.add(vectorMultiplyBpe(destVgpr, destVgpr, tP["bpeDS"])) # LdsBlockSizePerPad: add padding if kernel["LdsBlockSizePerPad%s"%tc] != 0 and kernel["LdsPad%s"%tc] != 0: tmpVgpr = self.vgprPool.checkOutAligned(2, 2, tag="lwaFirstOffset_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) module.add(vectorStaticDivide(tmpVgpr, destVgpr, kernel["LdsBlockSizePerPad%s"%tc], tmpVgprRes, \ "padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) with self.allocTmpSgpr(1, tag="lwaFirstOffset_tmpSgprInfo") as tmpSgprInfo: module.add(vectorStaticMultiplyAdd(vgpr(destVgpr), vgpr(tmpVgpr), int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), vgpr(destVgpr), tmpSgprInfo, \ "padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) self.vgprPool.checkIn(tmpVgpr) if tc in ("A", "B", "MXSA", "MXSB", "Metadata"): if kernel["LdsOffset%s" % tc] != 0: module.add(VAddCOU32( dst=vgpr(destVgpr), \ dst1=VCC(), \ src0=hex(kernel["LdsOffset%s" % tc]), \ src1=vgpr(destVgpr), \ comment="lwFO%s = lw%s + lw%s*MT%s + LDS_OFFSET_%s=%u" % (tc, tP["tileChar"], \ self.states.unrollChar, tP["tileChar"], tc, kernel["LdsOffset%s"%tc]) )) numLwa = 0 if tP["isA"]: numLwa = self.states.a.numVgprLocalWriteAddr elif tP["isB"]: numLwa = self.states.b.numVgprLocalWriteAddr elif tP["isM"]: numLwa = self.states.m.numVgprLocalWriteAddr elif tP["isMXSA"]: numLwa = self.states.mxsa.numVgprLocalWriteAddr elif tP["isMXSB"]: numLwa = self.states.mxsb.numVgprLocalWriteAddr else: raise RuntimeError(f"invalid tc {tc}") for i in range(1, numLwa): module.add(VAddU32(dst=vgpr("LocalWriteAddr%s+%u"%(tc,i)), src0=i * self.states.regCaps["maxLDSConstOffset"], src1= vgpr(destVgpr), \ comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024) )) #LSC_ * LSP_ numBytesPerElement = tP["bpeGR"] validWIPerLoad = kernel[tP["lsc"]] * kernel[tP["lsp"]]//tP["glvw"] validBytesPerLoad = int(kernel[tP["lsc"]] * kernel[tP["lsp"]] * numBytesPerElement) maxBytesPerLoad = int(kernel["NumThreads"] * tP["glvw"] * numBytesPerElement) if kernel["WaveSeparateGlobalRead%s"%tc] == 1: validBytesPerLoad *= (kernel["NumThreads"] // self.states.kernel["WavefrontSize"]) elif kernel["WaveSeparateGlobalRead%s"%tc] == 2: if kernel["ProblemType"]["TLU%s"%tc]: validBytesPerLoad *= (kernel["DepthU"] // kernel["NumLoadsPerpendicular%s"%tc] // (kernel["NumThreads"] // kernel["WavefrontSize"])) else: validBytesPerLoad *= (kernel["MacroTile%s"%tc] // kernel["NumLoadsPerpendicular%s"%tc] // (kernel["NumThreads"] // kernel["WavefrontSize"])) isDTVAB = ((tP["isA"] or tP["isB"] or tP["isMXSA"] or tP["isMXSB"]) and kernel["DirectToVgpr%s"%tc]) # For GNLC we don't need to check these asserts since num threads coalesced may not divide num threads assert (validBytesPerLoad <= maxBytesPerLoad) or isDTVAB or kernel["UseGeneralizedNLCOne%s"%tc] assert (kernel[tP["lsc"]] * kernel[tP["lsp"]] % tP["glvw"] == 0) or isDTVAB or kernel["UseGeneralizedNLCOne%s"%tc] if validBytesPerLoad != maxBytesPerLoad: with self.allocTmpSgpr(1, tag="lwaFirstOffset_tmpSgprInfo2") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx module.add(SMovB32(dst=sgpr(tmpSgpr), src=validWIPerLoad, \ comment="lsc*lsp=%u*%u"%(kernel[tP["lsc"]],kernel[tP["lsp"]] ))) module.add(VCmpLtU32( dst=VCC(), \ src0=vgpr("Serial"), \ src1=sgpr(tmpSgpr), \ comment="fractional: ensure tid < global read tile elements")) tmpVgpr = self.vgprPool.checkOut(1, "tmpVgpr", self.states.preventVgprOverflowDuringNewTile) module.add(VMovB32(dst=vgpr(tmpVgpr), src=hex(self.consts.ldsOOB))) module.add(VCndMaskB32( dst=vgpr(destVgpr), \ src0=vgpr(tmpVgpr), \ src1=vgpr(destVgpr), \ comment="Mask load so out-of-gr-tile bounds returns 0")) self.vgprPool.checkIn(tmpVgpr) if kernel["LocalWriteUseSgpr%s"%tc]: if not kernel["UseGeneralizedNLCOne%s"%tc]: # TODO: Can refactor code above to Compute this directly: if self.states.archCaps["CrosslaneWait"]: module.add(SNop(waitState=0, comment="1 wait states required before reading vgpr by lane")) module.add(VReadfirstlaneB32( dst=sgpr("LocalWriteAddr%s"%tc), \ src=vgpr(destVgpr), \ comment="Copy lds write address VGPR to SGPR")) else: tmpv = self.vgprPool.checkOut(1, tag="lwaFirstOffset_tmpv") module.add(VLShiftRightB32(dst=vgpr(tmpv), shiftHex=log2(kernel["WavefrontSize"]), src=vgpr("Serial"), comment="Compute waveID")) if self.states.archCaps["CrosslaneWait"]: module.add(SNop(waitState=0, comment="1 wait states required before reading vgpr by lane")) module.add(VReadfirstlaneB32( dst=sgpr("LocalWriteAddr%s"%tc), \ src=vgpr(tmpv), \ comment="Copy lds write address VGPR to SGPR")) module.add(SMulI32(dst=sgpr("LocalWriteAddr%s"%tc), src0=sgpr("LocalWriteAddr%s"%tc), \ src1=int((kernel["WavefrontSize"] * kernel["GlobalReadVectorWidth%s"%tc]+kernel["LdsPad%s"%tc]) * tP["bpeGR"]) )) if tc in ('A', 'B') and kernel["LdsOffset%s" % tc] != 0: module.add(SAddU32(dst=sgpr("LocalWriteAddr%s"%tc), src0=sgpr("LocalWriteAddr%s"%tc), \ src1=kernel["LdsOffset%s" % tc] )) self.vgprPool.checkIn(tmpv) self.vgprPool.checkIn(destVgpr) if kernel["StoreSwapAddr"] or self.states.useCommonSgprSwap: if kernel["LocalWriteUseSgpr%s"%tc]: if self.states.useCommonSgprSwap: # Need only once. Generate the code for "A" only if tc == "A": module.add(SMovB32(dst=sgpr("SwapCommon"), src=0, comment="Initialize SwapCommon")) else: # needed for the VReadfirstlaneB32 in the prior code block if self.states.archCaps["CrosslaneWait"]: module.add(SNop(waitState=0, comment="1 wait states")) module.add(SAddU32(dst=sgpr("Swap%s"%tc), src0=sgpr("LocalWriteAddr%s"%tc), src1=kernel["LdsOffsetA_Blk"], comment="Calculate starting lds addr of second buffer")) module.add(SXorB32(dst=sgpr("Swap%s"%tc), src0=sgpr("Swap%s"%tc), src1=sgpr("LocalWriteAddr%s"%tc), comment="xor both lds buffer offsets to enable swapping")) else: module.add(VAddU32(dst=vgpr("LocalWriteSwapAddr%s"%tc), src0=kernel["LdsOffsetA_Blk"], src1=vgpr("LocalWriteAddr%s"%tc), \ comment="starting lds addr of second buffer" )) module.add(VXorB32(dst=vgpr("LocalWriteSwapAddr%s"%tc), \ src0=vgpr("LocalWriteSwapAddr%s"%tc), \ src1=vgpr("LocalWriteAddr%s"%tc), \ comment="xor both lds offsets to enable swapping")) return Module("lwaUnrollAssignment (Empty)") if self.dontAppendCode or isDTVAB else module ############################################################################## # Local Read Addresses: Tile Assignment ############################################################################## def lraTileAssignment(self, kernel, tPA, tPB): module = Module("lraTileAssignment") component = Component.LraTileAssignment.find(self) tPMXSA = tPA["MX"] if kernel["ProblemType"]["MXBlockA"] else None tPMXSB = tPB["MX"] if kernel["ProblemType"]["MXBlockB"] else None tP0 = tPA if tPB["tile01Idx"] else tPB tPMXS0 = tPMXSA if tPB["tile01Idx"] else tPMXSB tP1 = tPB if tPB["tile01Idx"] else tPA tPMXS1 = tPMXSB if tPB["tile01Idx"] else tPMXSA if component: # do not generate local read code if DirectToVgpr is enabled tc = tP0["tensorChar"] if not kernel["DirectToVgpr%s"%tc]: module.add(component(self, kernel, tP0)) if tPMXS0: module.add(component(self, kernel, tPMXS0)) # do not generate local read code if DirectToVgpr is enabled tc = tP1["tensorChar"] if not kernel["DirectToVgpr%s"%tc]: module.add(component(self, kernel, tP1)) if tPMXS1: module.add(component(self, kernel, tPMXS1)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] module.add(component(self, kernel, tPM)) else: assert 0, "Compoment lraTileAssignment Not Found" return module ############################################################################## # Local Read Addresses: Final Offset A/B ############################################################################## def lraFinalOffset(self, kernel, tP): module = Module("lraFinalOffset") tc = tP["tensorChar"] # do not generate local read code if DirectToVgpr is enabled if (tc in ("A", "MXSA", "B", "MXSB")) and kernel["DirectToVgpr%s"%tc]: return Module("lraFinalOffset (Empty)") if kernel["EnableMatrixInstruction"]: # allocate resources wave_id = self.vgprPool.checkOut(1, tag="lraFinalOffset_wave_id") # quotient rReg = self.vgprPool.checkOut(1, tag="lraFinalOffset_rReg") # remainder, unused here tmpVgpr = self.vgprPool.checkOutAligned(2, 2, tag="lraFinalOffset_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) # constant tc = tP["tensorChar"] tile01 = tP["tile01Idx"] LdsPad = kernel["LdsPad%s" % tc] if kernel["LdsBlockSizePerPad%s" % tc] == 0 else 0 mtAddPad = kernel["MacroTile%u" % tile01] + LdsPad umlds = kernel["UnrollMajorLDS%s" % tc] lsu = kernel["LocalSplitU"] du = kernel["_DepthU%s"%tc] lsuStride = du // lsu numWavesPerLsu = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] mxScaleFormat = kernel.get("MXScaleFormat", "NoSwizzle") isMxSwizzled = ("MXS" in tc) and mxScaleFormat in ("InMemorySwizzle", "HostPreSwizzle") # generate instruction module.add(vectorStaticDivide(wave_id, "Serial", kernel["WavefrontSize"], tmpVgprRes)) module.add(vectorStaticDivide(wave_id, wave_id, numWavesPerLsu, tmpVgprRes, comment="LSU offset: Get LSU wave_id")) with self.allocTmpSgpr(1, tag="lraFinalOffset_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx if umlds == False or isMxSwizzled: module.add(SMovB32(dst=sgpr(tmpSgpr), src=mtAddPad*lsuStride, \ comment="LSU offset: stride = lsuStride(%u)*(MT%u(%u) + PAD%u(%u))" % (lsuStride,tile01, kernel["MacroTile%u" % tile01], tile01, LdsPad))) else: module.add(SMovB32(dst=sgpr(tmpSgpr), src=lsuStride, \ comment="LSU offset: stride = lsuStride(%u) when umlds==True" % (lsuStride))) module.add(VMulLOU32(dst=vgpr(wave_id), src0=sgpr(tmpSgpr), src1=vgpr(wave_id), \ comment="LSU offset: lsuoffset = wave_id*lsuStride*(MT%u+PAD)"%tile01)) # final offset finalVgpr = "LocalReadAddr%s"%tc module.add(VAddU32(dst=vgpr(finalVgpr), src0=vgpr(wave_id), src1=vgpr(tP["gpr"]["lro"]), \ comment="Final Offset: offset = (lro%s+lsuoffset)*bpeDS" % tile01 )) module.add(vectorMultiplyBpe(finalVgpr, finalVgpr, tP["bpeDS"])) # LdsBlockSizePerPad: add padding if kernel["LdsBlockSizePerPad%s"%tc] != 0 and kernel["LdsPad%s"%tc] !=0: module.add(vectorStaticDivide(rReg, "LocalReadAddr%s"%tc, kernel["LdsBlockSizePerPad%s"%tc], tmpVgprRes, \ "Final Offset: padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) with self.allocTmpSgpr(1, tag="lraFinalOffset_tmpSgprInfo2") as tmpSgprInfo: module.add(vectorStaticMultiplyAdd(vgpr("LocalReadAddr%s"%tc), vgpr(rReg), int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), vgpr("LocalReadAddr%s"%tc), tmpSgprInfo, \ "Final Offset: padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) # footprintPacked seg-interleave: add the component jump (wtid0*(fA+fB), bytes) AFTER re-pad # so A1/B1 pack exactly at the previous tile's end regardless of padA vs padB. if "segWaveByteOff" in tP.get("gpr", {}): module.add(VAddU32(dst=vgpr("LocalReadAddr%s"%tc), src0=vgpr(tP["gpr"]["segWaveByteOff"]), \ src1=vgpr("LocalReadAddr%s"%tc), comment="seg interleave: += wtid0*(fA+fB) (post-pad)")) self.vgprPool.checkIn(tP["gpr"].pop("segWaveByteOff")) # release resources self.vgprPool.checkIn(tmpVgpr) self.vgprPool.checkIn(wave_id) self.vgprPool.checkIn(rReg) self.vgprPool.checkIn(tP["gpr"]["lro"]) else: # dot2: support WaveSplitk # constant tile01 = tP["tile01Idx"] LdsPad = kernel["LdsPad%s" % tc] if kernel["LdsBlockSizePerPad%s" % tc] == 0 else 0 divisor = kernel["SubGroup0"] * kernel["SubGroup1"] mtAddPad = kernel["MacroTile%u" % tile01] + LdsPad lrvw = self.states.lrvwUnrollA if tc == 'A' else self.states.lrvwUnrollB # final offset finalVgpr = vgpr("LocalReadAddr%s"%tc) with self.allocTmpSgpr(1, tag="lraFinalOffset_tmpSgprInfo3") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx if kernel["UseDotInstruction"]: kidx = self.vgprPool.checkOut(1, tag="lraFinalOffset_kidx") # remainder module.add(VAndB32(dst=vgpr(kidx), src0=(kernel["NumWaveSplitK"]-1), src1=vgpr("Serial"), comment="kidx = Serial % NumWaveSplitK")) module.add(vectorStaticMultiply(vgpr(kidx), vgpr(kidx), lrvw, tmpSgprInfo, \ "*= lrvw")) # Final offset module.add(VAddLShiftLeftU32(dst=finalVgpr, shiftHex=hex(log2(tP["bpe"])), src0=vgpr(kidx), src1=vgpr(tP["gpr"]["lro"]), \ comment="Final Offset: add padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) self.vgprPool.checkIn(kidx) else: sgid = self.vgprPool.checkOut(1, tag="lraFinalOffset_sgid") # quotient vtmp = self.vgprPool.checkOut(1, tag="lraFinalOffset_vtmp") # tmp vCont = ContinuousRegister(vtmp, 1) module.add(vectorStaticDivide(sgid, "Serial", divisor, vCont, \ "LSU offset: sgid = Serial / subGroup(%u)" % divisor)) module.add(vectorStaticMultiply(vgpr(sgid), vgpr(sgid), mtAddPad, tmpSgprInfo, \ "LSU offset: lsuoffset = sgid*(MT%u+PAD)"%tile01)) # module.add(SMovB32(dst=sgpr(tmpSgpr), src=mtAddPad*lsuStride, \ # comment="LSU offset: stride = lsuStride(%u)*(MT%u(%u) + PAD%u(%u))" % (lsuStride,tile01, kernel["MacroTile%u" % tile01], tile01, LdsPad))) module.add(vectorStaticMultiply(vgpr(tP["gpr"]["lro"]), vgpr(tP["gpr"]["lro"]), kernel["VectorWidthB"], tmpSgprInfo, \ "Final Offset: lr%sOffset * VW" % tc)) module.add(VAddLShiftLeftU32(dst=finalVgpr, shiftHex=hex(log2(tP["bpe"])), src0=vgpr(sgid), src1=vgpr(tP["gpr"]["lro"]), \ comment="Final Offset: add padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) self.vgprPool.checkIn(sgid) self.vgprPool.checkIn(vtmp) # release resources self.vgprPool.checkIn(tP["gpr"]["lro"]) # LdsBlockSizePerPad: add padding if kernel["LdsBlockSizePerPad%s"%tc] != 0 and kernel["LdsPad%s"%tc] !=0: with self.allocTmpSgpr(1, tag="lraFinalOffset_tmpSgprInfo4") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx rReg = self.vgprPool.checkOut(1, tag="lraFinalOffset_rReg") # remainder, unused here module.add(vectorStaticDivide(rReg, "LocalReadAddr%s"%tc, kernel["LdsBlockSizePerPad%s"%tc], tmpSgpr, \ "Final Offset: padding %u per block %u" % (kernel["LdsPad%s"%tc], kernel["LdsBlockSizePerPad%s"%tc]))) module.add(vectorStaticMultiplyAdd(vgpr("LocalReadAddr%s"%tc), vgpr(rReg), int(kernel["LdsPad%s"%tc] * tP["bpe"]), vgpr("LocalReadAddr%s"%tc), tmpSgprInfo, \ "Final Offset: padding %u per block %u" % (int(kernel["LdsPad%s"%tc] * tP["bpeDS"]), kernel["LdsBlockSizePerPad%s"%tc]))) self.vgprPool.checkIn(rReg) return module ############################################################################## # Local Read Addresses offset conversion for DTL + NLC > 1 ############################################################################## def lraOffsetConversionForDTLandNLC(self, kernel, tP, offset_val, generateAsm=False, \ finalVgpr=None, tmp1=None, tmp2=None): module = Module("lraOffsetConversionForDTLandNLC") # another address conversion for DirectToLds + NumLoadsCoalesced > 1 divisorName = tP["lvc"] divisor = kernel[divisorName] width = kernel["WavefrontSize"] if tP["tlu"] else kernel["DepthU"] if divisor < width: # DirectToLds + above conditions, rotate offset_val bits to adjust LDS offset lowerScale = tP["nrc"] upperScale = (kernel["WavefrontSize"] // divisor) # bit rotation necessary only when nrc > 1 if lowerScale > 1: tile01 = tP["tile01Idx"] rightShift = int(log2(lowerScale)) # assuming power of 2 leftShift = int(log2(upperScale)) # assuming power of 2 line = kernel["MacroTile%u" % tile01] if tP["tlu"] else kernel["DepthU"] ldsLineSize = int(line * tP["bpe"]) // lowerScale maskBitsLow = (lowerScale - 1) * ldsLineSize maskBitsHigh = (upperScale - 1) * lowerScale * ldsLineSize maskBitsAll = (maskBitsLow | maskBitsHigh) # offset_val conversion low = offset_val & maskBitsLow high = offset_val & maskBitsHigh low <<= leftShift high >>= rightShift val = low | high offset_val = (offset_val & (~maskBitsAll)) | val # generate asm code if generateAsm: with self.allocTmpSgpr(1, tag="lraOffsetConversionForDTLandNLC_tmpSgprInfo") as tmpSgprInfo: tmpSgpr2 = tmpSgprInfo.idx module.add(VAndB32(dst=vgpr(tmp1), src0=hex(maskBitsLow), src1=finalVgpr, \ comment="Offset rotation for DirectToLds + %s > 1"%tP["lsc"])) module.add(VAndB32(dst=vgpr(tmp2), src0=hex(maskBitsHigh), src1=finalVgpr)) module.add(VLShiftLeftB32(dst=vgpr(tmp1), shiftHex=hex(leftShift), src=vgpr(tmp1))) module.add(VLShiftRightB32(dst=vgpr(tmp2), shiftHex=hex(rightShift), src=vgpr(tmp2))) module.add(VOrB32(dst=vgpr(tmp1), src0=vgpr(tmp1), src1=vgpr(tmp2))) module.add(SMovB32(dst=sgpr(tmpSgpr2), src=hex(maskBitsAll))) module.add(VNotB32(dst=vgpr(tmp2), src=sgpr(tmpSgpr2))) module.add(VAndB32(dst=finalVgpr, src0=vgpr(tmp2), src1=finalVgpr)) module.add(VOrB32(dst=finalVgpr, src0=vgpr(tmp1), src1=finalVgpr)) return module, offset_val ############################################################################## # Local Read Addresses: Declare Addresses A/B ############################################################################## def lraDeclareAddresses(self, kernel, tP): module = Module("lraDeclareAddresses") tc = tP["tensorChar"] # no need to generate add code if LdsOffset is 0 or DirectToVgprB if (tc in ("A", "B", "MXSA", "MXSB")) and kernel["DirectToVgpr%s"%tc]: module = Module("lraDeclareAddresses (Empty)") elif (kernel["LdsOffset%s"%tc] != 0) or \ (kernel.get("LDSSegmentInterleave") == 1 and tc in ("B", "MXSA", "MXSB")): _segOff = kernel["LDSSegInterleaveOffsets"] if kernel.get("LDSSegmentInterleave") == 1 else {} if tc == "B" and kernel.get("LDSSegmentInterleave") == 1: _ldsBase = _segOff["ldsBaseB"] elif tc in ("MXSA", "MXSB") and _segOff.get("ldsBase" + tc) is not None: _ldsBase = _segOff["ldsBase" + tc] # relocated MX scale base else: _ldsBase = kernel["LdsOffset%s"%tc] module.add(VAddCOU32(dst=vgpr("LocalReadAddr%s+0"%tc), dst1=VCC(), src0=hex(_ldsBase), src1=vgpr("LocalReadAddr%s+0"%tc), comment=" += LdsOffset%s (lower)"%tc)) if tP["isA"]: numVgpr = self.states.a.numVgprLocalReadAddr elif tP["isMXSA"]: numVgpr = self.states.mxsa.numVgprLocalReadAddr elif tP["isMXSB"]: numVgpr = self.states.mxsb.numVgprLocalReadAddr elif tP["isB"]: numVgpr = self.states.b.numVgprLocalReadAddr elif tP["isM"]: numVgpr = self.states.m.numVgprLocalReadAddr else: raise Exception(f"unsupport tc %s{tc}") for i in range(1, numVgpr): module.add(VAddU32(dst=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"],i)), src0=i * self.states.regCaps["maxLDSConstOffset"], \ src1= vgpr("LocalReadAddr%s+0"%tP["tensorChar"]), comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024) )) return module def lraSwapAddressesForDTLPad(self, kernel, tP): module = Module("lraSwapAddressesForDTLPad") tc = tP["tensorChar"] if kernel["StoreSwapAddr"]: module.add(VAddU32(dst=vgpr("LocalReadSwapAddr%s"%tc), src0=kernel["LdsOffsetA_Blk"], src1=vgpr("LocalReadAddr%s"%tc), \ comment="Calculate starting lds addr of second buffer" )) module.add(VXorB32(dst=vgpr("LocalReadSwapAddr%s"%tc), \ src0=vgpr("LocalReadSwapAddr%s"%tc), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="xor both lds buffer offsets to enable swapping")) return module def lraAddressesInitFor3LDSBlk(self, kernel, tP, initSreg, initVreg): module = Module("lraAddressesInitForPGR3") tc = tP["tensorChar"] if self.states.IncLdsBufSwitch: # initialize inc register for LocalReadAddress if initSreg: module.add(SMovB32( dst=sgpr("LDSBufferReadInc"), \ src=0, \ comment="init LRAddr inc Sgpr")) # 3 or more LDS block case, need to keep original LocalReadAddr # need to copy to LocalReadAddrOrig after all init calculation for LocalReadAddr is done if initVreg: module.add(VMovB32( dst=vgpr("LocalReadAddrOrig%s"%(tc)), \ src=vgpr("LocalReadAddr%s"%(tc)), \ comment="set LocalReadAddr to LocalReadAddrOrig")) return module def lwaInitAddressesForDTLTailLoop(self, kernel, tP): module = Module("lwaInitAddressesForDTLTailLoop") tc = tP["tensorChar"] waveSize = kernel["WavefrontSize"] if kernel["DirectToLds%s"%tc] and kernel["NonDTLTailLoop%s"%tc]: module.addComment0("Set local write offsets for %s to be same as DTL %uB load"%(tc, kernel["GlobalReadVectorWidth%s"%tc] * tP["bpe"])) module.add(VAndB32(dst=vgpr("LocalWriteAddr%s"%tc), src0=(waveSize - 1), src1=vgpr("Serial"), comment="Serial % wavesize")) module.add(VLShiftLeftB32(dst=vgpr("LocalWriteAddr%s"%tc), shiftHex=hex(log2(kernel["GlobalReadVectorWidth%s"%tc] * tP["bpe"])), src=vgpr("LocalWriteAddr%s"%tc), comment="")) module.add(VAddU32(dst=vgpr("LocalWriteAddr%s"%tc), src0=sgpr("LocalWriteAddr%s"%tc), src1=vgpr("LocalWriteAddr%s"%tc), \ comment="" )) numLw = 0 if tP["isA"]: numLw = self.states.a.numVgprLocalWriteAddrTailLoop elif tP["isB"]: numLw = self.states.b.numVgprLocalWriteAddrTailLoop elif tc == "MXSA": numLw = self.states.mxsa.numVgprLocalWriteAddrTailLoop elif tc == "MXSB": numLw = self.states.mxsb.numVgprLocalWriteAddrTailLoop elif tP["isM"]: numLw = self.states.m.numVgprLocalWriteAddrTailLoop for i in range(1,numLw): module.add(VAddU32(dst=vgpr("LocalWriteAddr%s + %s"%(tc, i)), src0=(i * 0x10000), \ src1=vgpr("LocalWriteAddr%s"%tc), comment="" )) return module ############################################################################## # openShadowInit # Label after prefetches are launched. This is present even if ShadowInit not # used. ############################################################################## def openShadowInit(self): module = Module("openShadowInit") module.add(Label("ShadowInitStart", "")) return module ############################################################################## # closeShadowInit # Label after prefetches are launched. This is present even if ShadowInit not # used. ############################################################################## def closeShadowInit(self, kernel): module = Module("closeShadowInit") assert(self.states.doShadowInit and kernel["PrefetchGlobalRead"]) module.add(self.checkLastIter(kernel)) if kernel["SuppressNoLoadLoop"]: loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx]] strNtab = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTA0_NTB0" lastIterEnd = Label("LoopEnd%s%s"%(loopChar, strNtab), "") else: lastIterEnd = Label("PrefetchGlobalLastIterEnd", "") # This branch could potentially be very far e.g. > SIMM16 module.addComment1("after InitC, skip to end of prefetch last iter if numIter==0") if self.isPrefetchAcrossPersistentEnabled(kernel) and not kernel["SuppressNoLoadLoop"]: module.add(SCMovB32(dst=sgpr("SkPrefetchPrimed"), src=0, comment="discard primed PAP group when current slice skips NLL")) # use positive offset only long jump with self.allocTmpSgpr(3, tag="closeShadowInit_tmpSgprInfo") as tmpSgprInfo: module.addComment1(lastIterEnd.getLabelName()) module.add(self.longBranchScc1(lastIterEnd, posNeg=1, tmpSgprInfo=tmpSgprInfo)) return module ############################################################################## # Initialize C ############################################################################## def initC(self, kernel, initCIterWmma=False, waitForPGRLabel=None): module = Module("initC") skipInitCVmovLabel = None if initCIterWmma: loopIdx = 0 EndCounter = self.unrollLoopEntryEndCounter(kernel) # match openLoop's loop-entry guard loopCounter = self.loopCounter(kernel, loopIdx) loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][loopIdx]] skipInitCVmovLabel = Label(self.labels.getNameInc("skipInitCVmov"), "") module.add(SCmpLeU32( src0=loopCounter, \ src1=hex(EndCounter), \ comment="LoopCounter%s < EndCounter"%(loopChar) )) module.add(SCBranchSCC0(labelName=skipInitCVmovLabel.getLabelName(), \ comment="skip v_mov initC (WMMA initC will run in main loop)")) if self.states.lastValuMXSAB: self.vgprPool.remove(0 , self.states.lastValuMXSAB, "ValuMXSAB") module.addComment1("initC: remove ValuMXSA/B vgpr buffer [%u...%u) from pool"%(self.states.mxsa.startVgprValu, self.states.lastValuMXSAB)) self.vgprPool.remove(self.states.c.startVgprValu, self.states.c.numVgprValu, "ValuC") module.addComment1("initC: remove ValuC vgpr buffer [%u...%u) from pool"%(self.states.c.startVgprValu, self.states.c.startVgprValu+self.states.c.numVgprValu)) numAccvgprs = self.states.totalAgprs self.agprPool.remove(0, numAccvgprs, "ValuC") module.addComment1("initC: remove acc vgpr buffer [%u...%u) from pool"%(0, numAccvgprs)) self.vgprPool.remove(self.states.a.startVgprValu , self.states.lastValuAB - self.states.a.startVgprValu , "ValuAB") module.addComment1("initC: remove ValuA/B vgpr buffer [%u...%u) from pool"%(self.states.a.startVgprValu , self.states.lastValuAB)) numCVgpr = self.states.c.numVgprValu + numAccvgprs # TBD: optimize MfmaInitCVgprs case when using both VALU and ACC VGPRs (needs to init vgprs and acc separately) # Alternatively, allow MFMAInstruction to accept acc2=0 if kernel["MfmaInitCVgprs"] == True and self.states.asmCaps["HasMFMA_f8f6f4"] and self.states.maxLimitAgprs >= numCVgpr >= 32: tmpVgpr = self.vgprPool.checkOutAligned(2,2,"tmp vgpr for lds init C registers") module.add(VMovB64(dst=vgpr(tmpVgpr,2), src=0, comment="A/B=0")) chunkSize = 16 for i in range(0, min(16, numCVgpr)): copyInst = VMovB32 if i >= numAccvgprs else VAccvgprWrite regStr = vgpr("ValuC+%u"%(i-numAccvgprs)) if i >= numAccvgprs else accvgpr(i) module.add(copyInst(dst=regStr, src=0, comment="initC")) accvgprAlias = vgpr if kernel["MIArchVgpr"] else accvgpr for i in range(min(16, numCVgpr), numCVgpr, 16): if i + 16 <= numCVgpr: module.add(MFMAInstruction(instType=InstType.INST_I8, accType=InstType.INST_I32, variant=[32,32,16,1], mfma1k=False, \ acc=accvgprAlias(i,16), a=vgpr(tmpVgpr,2), b=vgpr(tmpVgpr,2), acc2=accvgprAlias(0,16), \ comment="initC: [%u, %u]"%(i, i+15))) else: for j in range(i,min(i+16, numCVgpr)): copyInst = VMovB32 if j >= numAccvgprs else VAccvgprWrite regStr = vgpr("ValuC+%u"%(j-numAccvgprs)) if j >= numAccvgprs else accvgpr(j) module.add(copyInst(dst=regStr, src=0, comment="initC")) self.vgprPool.checkIn(tmpVgpr) else: if kernel["LdsInitCVgprs"]: tmpAddr = self.vgprPool.checkOut(1,"tmp vgpr for lds init C registers") module.add(VMovB32(dst=vgpr(tmpAddr), src=self.consts.ldsOOB, comment="set out-of-bound addr")) for i in range(0, numCVgpr): copyInst = VMovB32 if i >= numAccvgprs else VAccvgprWrite regStr = vgpr("ValuC+%u"%(i-numAccvgprs)) if i >= numAccvgprs else accvgpr(i) if not kernel["LdsInitCVgprs"]: module.add(copyInst(dst=regStr, src=0, comment="initC")) else: module.add(DSLoadB32(dst=regStr, src=vgpr(tmpAddr), ds=DSModifiers(offset=0), comment="initC")) if kernel["LdsInitCVgprs"]: self.vgprPool.checkIn(tmpAddr) if skipInitCVmovLabel is not None: module.add(skipInitCVmovLabel) return module def initSumUnroll(self, kernel): return self.exclasses.biasSumUnroll.initSumUnroll(self, kernel) ############################################################################## # Calculate and apply stagger offsets and edge # Output: Sets sgpr(StaggerRowMask) ############################################################################## def declareStaggerParms(self, kernel): module = Module("declareStaggerParms") #Calculate StaggerUIter with self.allocTmpSgpr(4, tag="declareStaggerParms_tmpSgprInfo") as tmpSgprInfo: beginStaggerUIterLabel = Label(self.labels.getNameInc("beginStaggerUIter"),comment="") endStaggerUIterLabel = Label(self.labels.getNameInc("endStaggerUIter"), comment="") tmpSgpr = tmpSgprInfo.idx currentStaggerU = tmpSgpr shiftedStaggerU = tmpSgpr + 1 staggerUMask = tmpSgpr + 1 staggerUStrideShift = tmpSgpr + 2 staggerUMapping = tmpSgpr + 3 module.add(SAndB32(dst=sgpr(staggerUStrideShift), src0=sgpr("StaggerU"), src1=hex(0x1F00))) module.add(SLShiftRightB32(dst=sgpr(staggerUStrideShift), shiftHex=hex(8), src=sgpr(staggerUStrideShift))) module.add(SAndB32(dst=sgpr(staggerUMapping), src0=sgpr("StaggerU"), src1=hex(0xE000))) module.add(SAndB32(dst=sgpr("StaggerU"), src0=sgpr("StaggerU"), src1=hex(0xFF))) if kernel["PrefetchGlobalRead"] >= 3: module.add(SCmpLeU32(src0=sgpr("OrigLoopCounter"), src1=(kernel["PrefetchGlobalRead"] - 1), \ comment="loopCount <= PGR-1" )) module.add(SCMovB32(dst=sgpr("StaggerU"), src=0, comment="set StaggerU=0 if loopCount <= PGR-1")) module.add(SMovB32(dst=sgpr(currentStaggerU), src=sgpr("StaggerU"), comment="init staggerU")) module.add(beginStaggerUIterLabel) module.add(SLShiftLeftB32(dst=sgpr(shiftedStaggerU), src=sgpr(currentStaggerU), \ shiftHex=sgpr(staggerUStrideShift), comment="shift by StaggerUStride")) module.add(SCmpGeU32(src0=sgpr("OrigLoopCounter"), src1=sgpr(shiftedStaggerU), \ comment="loopCount >= current shift Count" )) module.add(SCBranchSCC1(labelName=endStaggerUIterLabel.getLabelName(), comment="jump to end")) module.add(SLShiftRightB32(dst=sgpr(currentStaggerU), src=sgpr(currentStaggerU), \ shiftHex=1, comment="step down to smaller stagger")) module.add(SBranch(labelName=beginStaggerUIterLabel.getLabelName(), comment="jump to begin")) module.add(endStaggerUIterLabel) module.add(SSubU32(dst=sgpr(staggerUMask), src0=sgpr(currentStaggerU), src1=1, comment="staggerU mask")) module.add(SCmpGeU32(src0=sgpr(currentStaggerU), src1=1, \ comment="if current staggerU >= 1" )) module.add(SCSelectB32(dst=sgpr("StaggerUIter"), src0=sgpr(staggerUMask), src1=0, comment="set Mask")) staggerInput = tmpSgpr staggerLabel = Label(self.labels.getNameInc("staggerInputEnd"), comment="") staggerMapLabelNames = [self.labels.getNameInc("StaggerUMapping") for _ in range(5)] for i in range(0, 5): label = Label(staggerMapLabelNames[i], comment="") module.add(SCmpEQU32(src0=sgpr(staggerUMapping), src1=hex(i << 13))) if i != 4: module.add(SCBranchSCC0(labelName=label.getLabelName())) else: module.add(SCBranchSCC0(labelName=staggerLabel.getLabelName())) if i == 0: module.add(SMovB32(dst=sgpr(staggerInput), src=sgpr("WorkGroup0"))) elif i == 1: module.add(SMovB32(dst=sgpr(staggerInput), src=sgpr("WorkGroup1"))) elif i == 2 and len(kernel["ProblemType"]["IndicesBatch"]) > 2: module.add(SMovB32(dst=sgpr(staggerInput), src=sgpr("WorkGroup2"))) elif i == 3: wgSerial = staggerInput tmp = tmpSgpr+1 if len(kernel["ProblemType"]["IndicesBatch"]) > 2: module.add(SMulI32(dst=sgpr(wgSerial), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"), \ comment="wgSerial = (nwg0*ngw1)*wg2 + (nwg0)*wg1 + wg0")) module.add(SMulI32(dst=sgpr(wgSerial), src0=sgpr(wgSerial), src1=sgpr("WorkGroup2"))) module.add(SMulI32(dst=sgpr(tmp), src0=sgpr("NumWorkGroups0"), src1=sgpr("WorkGroup1"))) module.add(SAddU32(dst=sgpr(wgSerial), src0=sgpr(wgSerial), src1=sgpr(tmp))) module.add(SAddU32(dst=sgpr(wgSerial), src0=sgpr(wgSerial), src1=sgpr("WorkGroup0"))) else: module.add(SMovB32(dst=sgpr(staggerInput), src=hex(-1))) module.add(SBranch(staggerLabel.getLabelName())) if i != 4: module.add(label) else: module.add(staggerLabel) module.add(SAndB32(dst=sgpr("StaggerUIter"), src0=sgpr("StaggerUIter"), \ src1=sgpr(staggerInput), \ comment="Compute actual stagger start for this tile")) module.add(SLShiftLeftB32(dst=sgpr("StaggerUIter"), src=sgpr("StaggerUIter"), \ shiftHex=sgpr(staggerUStrideShift), comment="shift by StaggerUStride")) skComponent = Component.StreamK.find(self) module.add(skComponent.declareStaggerParms(self, kernel)) return module ############################################################################## # Calculate and apply stagger offsets and edge ############################################################################## @staticmethod def _isTdmTensor(kernel, tc): # Explicit tc set so future tensor chars do not get classified by accident. # Metadata is intentionally absent: TDM metadata uses tdmMetadataGroup0 directly # instead of passing tc=="Metadata" through this A/B tensor helper. if tc in ("A", "MXSA"): return kernel.get("enableTDMA", False) if tc in ("B", "MXSB"): return kernel.get("enableTDMB", False) return False def calculateStagger(self, kernel, tP): imod = Module("calculateStagger") tc = tP["tensorChar"] isTDM = self._isTdmTensor(kernel, tc) if not isTDM: assert (kernel["BufferLoad"]) with self.allocTmpSgpr(3, tag="calculateStagger_tmpSgprInfo") as tmpSgprInfo: staggerTmp = tmpSgprInfo.idx incSparseSgpr = tmpSgprInfo.idx + 2 #--- imod.addComment1("addr += (StaggerUIter) * GlobalReadIncs%s+%u"% (tc, self.states.unrollIdx)) # Calculate the stagger byte offset imod.addModuleAsFlatItems(self.s_mul_i64_i32( sgpr(staggerTmp), sgpr(staggerTmp+1), \ sgpr("StaggerUIter"), sgpr("GlobalReadIncs%s+%u"%(tc, self.states.unrollIdx)), \ " stagger byte offset")) # Apply TDM stagger now while staggerTmp still holds StaggerUIter * GlobalReadIncs. # The Sparse and PGR>=3 paths below both reuse staggerTmp for other computations. if isTDM: self._applyStaggerTDM(imod, kernel, tP, staggerTmp) # Amount of bytes to add to get back to start. # on the llop iteration which matches StaggerUIter, this offset added instead of GlobalReadInc imod.addModuleAsFlatItems(self.s_mul_i64_i32(sgpr("WrapU%s+0"%tc), sgpr("WrapU%s+1"%tc), \ self.loopCounter(kernel, self.states.unrollIdx), sgpr("GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx)), \ "Number of bytes accessed by the unroll loop")) # TODO: put this asmCaps into rocisa SSubU64 if self.states.asmCaps["s_sub_u64"] and self.states.asmCaps["HasWMMA_V3"]: with self.allocTmpSgpr(2, 2, tag="calculateStagger_stmp") as stmp: imod.add(SMovB32(sgpr(stmp.idx), sgpr("GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx)))) imod.add(SMovB32(sgpr(stmp.idx+1), 0)) imod.add(SSubU64(dst=sgpr("WrapU%s"%tc, 2), src0=sgpr(stmp.idx, stmp.size), src1=sgpr("WrapU%s"%tc, 2), comment="increment-WrapU")) else: imod.add(SSubU32(dst=sgpr("WrapU%s+0"%tc), \ src0=sgpr("GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx)), \ src1=sgpr("WrapU%s+0"%tc), \ comment="remove one iteration")) imod.add(SSubBU32(dst=sgpr("WrapU%s+1"%tc), \ src0=0, \ src1=sgpr("WrapU%s+1"%tc), \ comment="remove one iteration")) # Apply stagger offset to A/B SRD (TDM A/B is handled separately above) if not isTDM: imod.add(self.incrementSrd(tP, sgpr(staggerTmp), sgpr(staggerTmp+1))) # Sparse metadata SRD stagger. Three mutually exclusive metadata storage modes: # - DirectToVgprSparseMetadata: metadata lives in VGPR (incrementMetadataSrd) # - enableTDMMetadata: metadata SRD is tdmMetadataGroup0 (TDM hardware path) # - regular global: metadata SRD is tP["tpsMetadata"] # DTV and TDM both plumb WrapUMetadata + GlobalReadIncsMetadata through; the regular # global path is left to upstream semantics (unchanged from prior behavior). if kernel["ProblemType"]["Sparse"] and \ (kernel["DirectToVgprSparseMetadata"] or kernel["enableTDMMetadata"]) and \ ((kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"])): imod.addComment1("SRDs += (StaggerUIter) * GlobalReadIncsMetadata") tc = "Metadata" if kernel["DirectToVgprSparseMetadata"]: incSparse = incSparseSgpr imod.add(self.calculateIncrementMetadata(kernel, incSparse)) else: incSparse = "GlobalReadIncsMetadata+%u"%(self.states.unrollIdx) imod.addModuleAsFlatItems(self.s_mul_i64_i32( \ sgpr(staggerTmp), sgpr(staggerTmp+1), \ sgpr("StaggerUIter"), sgpr(incSparse), " stagger byte offset of metadata")) # Amount of bytes to add to get back to start. # on the llop iteration which matches StaggerUIter, this offset added instead of GlobalReadInc imod.addModuleAsFlatItems(self.s_mul_i64_i32( \ sgpr("WrapU%s+0"%tc), sgpr("WrapU%s+1"%tc), \ self.loopCounter(kernel, self.states.unrollIdx), sgpr(incSparse), \ "Number of bytes accessed by the unroll loop")) imod.add(SSubU32(sgpr("WrapU%s+0"%tc), sgpr(incSparse), sgpr("WrapU%s+0"%tc), " remove one iteration")) imod.add(SSubBU32(sgpr("WrapU%s+1"%tc), 0, sgpr("WrapU%s+1"%tc), " remove one iteration")) if kernel["DirectToVgprSparseMetadata"]: imod.add(self.incrementMetadataSrd(sgpr(staggerTmp), sgpr(staggerTmp+1))) elif kernel["enableTDMMetadata"]: self._applyStaggerTDM(imod, kernel, tP, staggerTmp, labelName="SkipStaggerMetadata", commentTag="stagger metadata", tdmGroupName="tdmMetadataGroup0") else: imod.add(self.incrementSrd(tP["tpsMetadata"], sgpr(staggerTmp), sgpr(staggerTmp+1))) if tP["isB"]: isDTVAorB = (kernel["DirectToVgprA"] != kernel["DirectToVgprB"]) # only one of them is enabled if kernel["PrefetchGlobalRead"] >= 2 and isDTVAorB: # PGR2 + DTVA or B (only 1 side), need separate StaggerUIter for DTV load imod.add(SAddU32(dst=sgpr("StaggerUIterDTV"), src0=sgpr("StaggerUIter"), \ src1=(1), \ comment="Subtract (PGR-1); StaggerUIter now contains target iteration to wrap")) # Convert passed in S' to S for easy loop comparison. S=S-(PGR-1)' pf = 2 if kernel["PrefetchGlobalRead"] else 1 if kernel["PrefetchGlobalRead"] >= 3: pf = kernel["PrefetchGlobalRead"] imod.add(SAddU32(dst=sgpr(staggerTmp), src0=sgpr("StaggerUIter"), \ src1=(pf), \ comment="Subtract (PGR-1); StaggerUIter now contains target iteration to wrap")) imod.add(SCmpGtU32(src0=sgpr("StaggerUIter"), src1=hex(0), \ comment="StaggerUIter > 0 (means StaggerU > 0)" )) imod.add(SCMovB32(dst=sgpr("StaggerUIter"), src=sgpr(staggerTmp), \ comment="StaggerUIter+=PGR only if StaggerU > 0")) else: imod.add(SAddU32(dst=sgpr("StaggerUIter"), src0=sgpr("StaggerUIter"), \ src1=(pf), \ comment="Subtract (PGR-1); StaggerUIter now contains target iteration to wrap")) return imod ############################################################################## # TDM: apply stagger offset to TDM descriptor. # Wave-separated (numWaves > 1): descriptors are aliased, so only the # matching-parity waves should apply the offset (even=A, odd=B). # Non-wave-separated (numWaves == 1): A and B have independent descriptors, # all waves apply unconditionally. ############################################################################## def _applyStaggerTDM(self, imod, kernel, tP, offsetSgpr, labelName="SkipStagger", commentTag="stagger", tdmGroupName=None): # tdmGroupName=None: derive SRD from tP and wave-parity-gate the add. In wave-separated # mode A and B share tdm{tc}Group0 SGPR but hold different per-wave values; the parity # gate ensures only the owning waves update. # tdmGroupName set (e.g. "tdmMetadataGroup0" for the sparse metadata SRD): wave-uniform # add. Metadata is single-sided, but all waves run this function; non-sparse-side waves # harmlessly advance a SRD copy they never read from. if tdmGroupName is None: tc = tP["tensorChar"] tdmGroup0 = f"tdm{tc}Group0" useParityGate = kernel["NumWaves"] > 1 else: tc = "Metadata" tdmGroup0 = tdmGroupName useParityGate = False skipLabel = None if useParityGate: wavelen = kernel["WavefrontSize"] skipLabel = Label(label=f"{labelName}{tc}", comment="") with self.allocTmpSgpr(1, tag="_applyStaggerTDM_waveIdTmp") as waveIdTmp: imod.add(VReadfirstlaneB32(dst=sgpr(waveIdTmp.idx), src=vgpr("Serial"), comment="get tId")) imod.add(SLShiftRightB32(dst=sgpr(waveIdTmp.idx), shiftHex=ceil(log2(wavelen)), src=sgpr(waveIdTmp.idx), comment="waveId")) imod.add(SBitcmp1B32(src0=sgpr(waveIdTmp.idx), src1=0, comment="check wave parity")) if "A" in tc: imod.add(SCBranchSCC1(labelName=skipLabel.getLabelName(), comment="skip: odd waves handle B")) else: imod.add(SCBranchSCC0(labelName=skipLabel.getLabelName(), comment="skip: even waves handle A")) imod.add(SAddU32(dst=sgpr(f"{tdmGroup0}+2"), src0=sgpr(f"{tdmGroup0}+2"), \ src1=sgpr(offsetSgpr), comment=f"TDM addr += {commentTag} offset (lo)")) imod.add(SAddCU32(dst=sgpr(f"{tdmGroup0}+3"), src0=sgpr(f"{tdmGroup0}+3"), \ src1=sgpr(offsetSgpr+1), comment=f"TDM addr += {commentTag} offset (hi)")) if skipLabel is not None: imod.add(skipLabel) ############################################################################## # Remove stagger offset (before tail loop) # | | | | # |-- S'*I --| # |---------- W' --------|-I-| # ^ current SRD pos # ^unrollLoopStart ^tailLoopStart (in summation0 dimension) # # S = sgprStaggerUIter = S+(PGR+1)' # W = sgprWrapU # PGR = kernel["PrefetchGlobalRead"] # # S' = StaggUIter that is passed into the kernel = -PGR+1+S # S'*I is also the global read offset (from unrollLoopStart) at unroll loop exit ? # I = GlobalReadIncs # W' = W # Need to move it to tailLoopStart # To compute position where tail loop should start: # = W' - S'*I + I # = W - (S+PGR+1)*I) + I # = W - (S+PGR+1)*I + I # = W - (S+2+PGR)*I ############################################################################## def removeStaggerAB(self, kernel, tPA, tPB): imod = Module("removeStaggerAB") if kernel["PrefetchGlobalRead"] >= 3: labelRemoveSUEnd = Label(label="RemoveStaggerUEnd", comment="") imod.add(SCmpEQU32(src0=sgpr("StaggerUIter"), src1=0, comment="if StaggerUIter is 0 (means StaggerU is 0), skip remove stagger")) imod.add(SCBranchSCC1(labelName=labelRemoveSUEnd.getLabelName(), comment="skip remove stagger")) # Remove stagger A(MXSA) imod.add(self.removeStagger(kernel, tPA)) if "MX" in tPA: imod.add(self.removeStagger(kernel, tPA["MX"])) # Remove stagger B(MXSB) if "MX" in tPB: imod.add(self.removeStagger(kernel, tPB["MX"])) imod.add(self.removeStagger(kernel, tPB)) if kernel["PrefetchGlobalRead"] >= 3: imod.add(labelRemoveSUEnd) # Wave-separated TDM: stagger setup/remove still use GlobalReadIncs*; release here. imod.add(self.releaseGlobalReadIncsSgprsAfterTdmWaveSep(kernel)) return imod def removeStagger(self, kernel, tP): imod = Module("removeStagger") tc = tP["tensorChar"] isTDM = self._isTdmTensor(kernel, tc) imod.addComment(f" removeStagger {tc}") with self.allocTmpSgpr(3, tag="removeStagger_tmpSgprInfo") as tmpSgprInfo: tmp = tmpSgprInfo.idx tmpIncSparse = tmpSgprInfo.idx + 2 # might be able to refactor this to eliminate signed math pf = 3 if kernel["PrefetchGlobalRead"] else 2 if kernel["PrefetchGlobalRead"] >= 3: pf = kernel["PrefetchGlobalRead"] + 1 imod.add(SSubI32(dst=sgpr(tmp), src0=pf, \ src1=sgpr("StaggerUIter"))) imod.addModuleAsFlatItems(self.s_mul_i64_i32_u32(sgpr(tmp), sgpr(tmp+1), \ sgpr(tmp), sgpr("GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx)), \ "start offset S in bytes")) # TODO: put this asmCaps into rocisa SSubU64 if self.states.asmCaps["s_sub_u64"] and self.states.asmCaps["HasWMMA_V3"]: imod.add(SSubU64(dst=sgpr(tmp, 2), src0=sgpr(tmp, 2), src1=sgpr("WrapU%s"%(tc), 2), comment="S - WrapU")) else: imod.add(SSubU32(dst=sgpr(tmp), src0=sgpr(tmp), src1=sgpr("WrapU%s"%tc), comment="S - WrapU")) imod.add(SSubBU32(dst=sgpr(tmp+1), src0=sgpr(tmp+1), src1=sgpr("WrapU%s+1"%(tc)), comment="S - WrapU")) # Apply remove-stagger offset to A/B (TDM goes through _applyStaggerTDM) if isTDM: self._applyStaggerTDM(imod, kernel, tP, tmp, labelName="SkipRemoveStagger", commentTag="removeStagger") else: imod.add(self.incrementSrd(tP, sgpr(tmp), sgpr(tmp+1))) # Sparse metadata SRD remove-stagger. Mirrors calculateStagger's three metadata # storage modes (DTV / TDM / regular global). if kernel["ProblemType"]["Sparse"] and \ ((kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"])): tc = "Metadata" if kernel["DirectToVgprSparseMetadata"]: incSparse = tmpIncSparse imod.add(self.calculateIncrementMetadata(kernel, incSparse)) else: incSparse = "GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx) # might be able to refactor this to eliminate signed math imod.add(SSubI32(dst=sgpr(tmp), src0=pf, \ src1=sgpr("StaggerUIter"))) imod.addModuleAsFlatItems(self.s_mul_i64_i32_u32(sgpr(tmp), sgpr(tmp+1), \ sgpr(tmp), sgpr(incSparse), \ "start offset S in bytes")) imod.add(SSubU32(sgpr(tmp), sgpr(tmp), sgpr("WrapU%s"%tc), "S - WrapU")) imod.add(SSubBU32(sgpr(tmp+1), sgpr(tmp+1), sgpr("WrapU%s+1"%(tc)), "S - WrapU")) if kernel["DirectToVgprSparseMetadata"]: imod.add(self.incrementMetadataSrd(sgpr(tmp), sgpr(tmp+1))) elif kernel["enableTDMMetadata"]: self._applyStaggerTDM(imod, kernel, tP, tmp, labelName="SkipRemoveStaggerMetadata", commentTag="removeStagger metadata", tdmGroupName="tdmMetadataGroup0") else: imod.add(self.incrementSrd(tP["tpsMetadata"], sgpr(tmp), sgpr(tmp+1))) return imod ############################################################################## # Set parameters UNDEF in the macro module ############################################################################## def undefineMacroModule(self, macroModule): imod = Module() for inst in macroModule.items(): if isinstance(inst, (RegSet, ValueSet)): imod.add(ValueSet(name=inst.name, value="UNDEF", format = -1)) elif isinstance(inst, Module): imod.add(self.undefineMacroModule(inst)) return imod ############################################################################## # Check in VGPRs which is allocated for tail loop ############################################################################## def tailLoopFreeVgpr(self, vgprBase, moduleMacro): imod = Module("tailLoopFreeVgpr") if vgprBase != -1: imod.add(self.undefineMacroModule(moduleMacro)) self.vgprPool.checkIn(vgprBase) return imod ############################################################################## # Check out VGPRs which is used for tail loop # The functions will retuen either # a single pair of (imod, vgprBase) # or # a array of pairs of (imod, vgprBase)'s # User must call tailLoopFreeVgpr(vgprBase, imod) to release the resources. ############################################################################## def tailLoopAllocValuVgpr(self, kernel, tensorParametersA, tensorParametersB, tensorParametersM): imodA = Module("tailLoopAllocValuAVgpr") vgprBaseA = -1 numValuA = 0 numVgprValuPackA = 0 # T reg definition for F32XEmu self.macroAndSetF32XEmuTreg(kernel, tensorParametersA, tensorParametersB) # 1024 vgpr: avoid cross pool usage valuVgprAlignment = 8 if self.states.asmCaps["HasVgprMSB"] else 2 if self.states.a.numVgprValu > 0 and not kernel["DirectToVgprA"]: numValuA = self.states.a.numVgprValu if tensorParametersA["bpe"] < 4 and not kernel["UnrollMajorLDSA"] and not kernel["enableLDSTrA"]: if self.states.lrvwTileA > 1: numVgprValuPackA = ceil(kernel["VectorWidthA"] * tensorParametersA["bpe"] / self.states.bpr) * kernel["MIWaveTileA"] // kernel["VectorWidthA"] * kernel["InnerUnroll"] * self.states.numVgprBuffer * kernel["MIInputPerThreadA"] if self.states.packDTVA: # pack DTV case, double the number numVgprValuPackA *= 2 elif (kernel["UsePLRPack"] and self.states.numItersPLR): numVgprValuPackA //= 2 else: numVgprValuPackA = self.states.a.numVgprValuPerBlock * kernel["InnerUnroll"] * self.states.numVgprBufferPackA * (int(4/tensorParametersA["bpeDS"]) - 1) vgprBaseA = self.vgprPool.checkOutAligned(numValuA + numVgprValuPackA, 2, tag="tailLoopAllocValuVgpr_vgprBaseA") imodA.add(RegSet("v", "vgprValuA_X0_I0_BASE", vgprBaseA)) imodA.add(self.moduleVgprMacroValuA) imodA.add(self.moduleVgprMacroValuA_T) if numVgprValuPackA > 0: imodA.add(RegSet("v", "vgprValuA_X0_I0_D0_PACK", vgprBaseA + numValuA)) imodA.add(self.moduleVgprMacroValuAPack) imodB = Module("tailLoopAllocValuBVgpr") vgprBaseB = -1 numValuB = 0 numVgprValuPackB = 0 if self.states.b.numVgprValu > 0 and not kernel["DirectToVgprB"]: numValuB = self.states.b.numVgprValu if tensorParametersB["bpe"] < 4 and not kernel["UnrollMajorLDSB"] and not kernel["enableLDSTrB"]: if self.states.lrvwTileB > 1: numVgprValuPackB = ceil(kernel["VectorWidthB"] * tensorParametersB["bpe"] / self.states.bpr) * kernel["MIWaveTileB"] // kernel["VectorWidthB"] * kernel["InnerUnroll"] * self.states.numVgprBuffer * kernel["MIInputPerThreadB"] if self.states.packDTVB: # pack DTV case, double the number numVgprValuPackB *= 2 elif (kernel["UsePLRPack"] and self.states.numItersPLR): numVgprValuPackB //= 2 else: numVgprValuPackB = self.states.b.numVgprValuPerBlock * kernel["InnerUnroll"] * self.states.numVgprBufferPackB * (int(4/tensorParametersB["bpeDS"]) - 1) vgprBaseB = self.vgprPool.checkOutAligned(numValuB + numVgprValuPackB, 2, tag="tailLoopAllocValuVgpr_vgprBaseB") imodB.add(RegSet("v", "vgprValuB_X0_I0_BASE", vgprBaseB)) imodB.add(self.moduleVgprMacroValuB) imodB.add(self.moduleVgprMacroValuB_T) if numVgprValuPackB > 0: imodB.add(RegSet("v", "vgprValuB_X0_I0_D0_PACK", vgprBaseB + numValuB)) imodB.add(self.moduleVgprMacroValuBPack) imodM = Module("tailLoopAllocValuMetadataVgpr") vgprBaseM = -1 numValuM = 0 numVgprValuPackMetadata = 0 if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: numValuM = self.states.m.numVgprValu if not kernel["UnrollMajorLDSMetadata"]: if self.states.lrvwTileMetadata > 1: miWaveTile = kernel["MIWaveTileB"] if kernel["ProblemType"]["Sparse"] == 2 else kernel["MIWaveTileA"] numVgprValuPackMetadata = roundUp(kernel["VectorWidthMetadata"] * tensorParametersM["bpe"] / self.states.bpr) * miWaveTile // kernel["VectorWidthMetadata"] * kernel["InnerUnroll"] * self.states.numVgprBuffer * kernel["MIInputPerThreadMetadata"] else: numVgprValuPackMetadata = (kernel["MIInputPerThreadMetadata"]-1) * kernel["MIWaveTileMetadata"] * kernel["InnerUnroll"] * self.states.numVgprBufferPackMetadata # No PGR metadata for tail loop. vgprBaseM = self.vgprPool.checkOutAligned(max(numValuM + numVgprValuPackMetadata, self.states.m.numVgprG2LAllocated), 2, tag="tailLoopAllocValuVgpr_vgprBaseM") imodM.add(RegSet("v", "vgprValuMetadata_X0_I0_BASE", vgprBaseM)) imodM.add(self.moduleVgprMacroValuM) if numVgprValuPackMetadata > 0: imodM.add(RegSet("v", "vgprValuMetadata_X0_I0_D0_PACK", vgprBaseM + numValuM)) imodM.add(self.moduleVgprMacroValuMPack) imodMisc = Module("tailLoopAllocValuMiscVgpr") vgprBaseMisc = -1 numVgprPackTemp = 0 if ((tensorParametersA["bpe"] < 4 and not kernel["UnrollMajorLDSA"]) \ or (tensorParametersB["bpe"] < 4 and not kernel["UnrollMajorLDSB"]) \ or (kernel["ProblemType"]["Sparse"] and not kernel["UnrollMajorLDSMetadata"] and (kernel["MIInputPerThreadMetadata"] == 4))) \ and (kernel["ProblemType"]["DataType"].isInt8() or kernel["ProblemType"]["DataType"].is8bitFloat()) or \ (self.states.asmCaps["HasSWMMAC_gfx1250"] and kernel["ProblemType"]["Sparse"] and not kernel["UnrollMajorLDSMetadata"]): numVgprPackTemp = 1 numVgprCvtTemp = 0 if self.states.a.startVgprValuCvtTemp != -1 or self.states.b.startVgprValuCvtTemp != -1: numVgprCvtTemp = 2 if numVgprPackTemp + numVgprCvtTemp > 0: vgprBaseMisc = self.vgprPool.checkOut(numVgprPackTemp + numVgprCvtTemp, tag="tailLoopAllocValuVgpr_vgprBaseMisc") if numVgprPackTemp > 0: imodMisc.add(RegSet("v", "vgprPackTemp", vgprBaseMisc)) if numVgprCvtTemp > 0: imodMisc.add(RegSet("v", "vgprCvtTemp", vgprBaseMisc + numVgprPackTemp)) return ([vgprBaseA, imodA],[vgprBaseB, imodB],[vgprBaseM, imodM],[vgprBaseMisc, imodMisc]) def tailLoopAllocG2LVgpr(self, kernel): imod = Module("tailLoopAllocG2LVgpr") vgprBase = -1 numG2LA = 0 numG2LB = 0 numG2LMXSA = 0 numG2LMXSB = 0 numG2LMetadata = 0 if not kernel["DirectToVgprA"]: if ("ULSGRODoubleG2L" in kernel) and kernel["ULSGRODoubleG2L"] == 1: numG2LA = self.states.a.numVgprG2LTailloopAllocated*2 else: numG2LA = self.states.a.numVgprG2LTailloopAllocated if not kernel["DirectToVgprB"]: if ("ULSGRODoubleG2L" in kernel) and kernel["ULSGRODoubleG2L"] == 1: numG2LB = self.states.b.numVgprG2LTailloopAllocated*2 else: numG2LB = self.states.b.numVgprG2LTailloopAllocated if kernel["ProblemType"]["MXBlockA"] and (not kernel["DirectToVgprMXSA"]): if ("ULSGRODoubleG2L" in kernel) and kernel["ULSGRODoubleG2L"] == 1: numG2LMXSA = self.states.mxsa.numVgprG2LTailloopAllocated*2 else: numG2LMXSA = self.states.mxsa.numVgprG2LTailloopAllocated if kernel["ProblemType"]["MXBlockB"] and (not kernel["DirectToVgprMXSB"]): if ("ULSGRODoubleG2L" in kernel) and kernel["ULSGRODoubleG2L"] == 1: numG2LMXSB = self.states.mxsb.numVgprG2LTailloopAllocated*2 else: numG2LMXSB = self.states.mxsb.numVgprG2LTailloopAllocated if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: numG2LMetadata = self.states.m.numVgprG2LAllocated numG2LA = ((numG2LA+1)//2)*2 numG2LB = ((numG2LB+1)//2)*2 numG2LMXSA = ((numG2LMXSA+1)//2)*2 numG2LMXSB = ((numG2LMXSB+1)//2)*2 numG2LMetadata = ((numG2LMetadata+1)//2)*2 if numG2LA + numG2LB + numG2LMXSA + numG2LMXSB + numG2LMetadata > 0: vgprBase = self.vgprPool.checkOutAligned(numG2LA + numG2LB + numG2LMXSA + numG2LMXSB + numG2LMetadata, 2, tag="tailLoopAllocG2LVgpr_vgprBase") imod.addComment0("Check out VGPR (numG2LA,numG2LB,numG2LMXSA,numG2LMXSB,numG2LMetadata) = (%d,%d,%d,%d,%d)"%(numG2LA,numG2LB,numG2LMXSA,numG2LMXSB,numG2LMetadata)) if numG2LA > 0: imod.add(RegSet("v", "vgprG2LA_BASE", vgprBase)) if kernel["DirectToLdsA"] and kernel["NonDTLTailLoopA"]: imod.add(RegSet("v", "vgprG2LA", "vgprG2LA_BASE", 0)) else: imod.add(self.moduleVgprMacroG2LA) if numG2LB > 0: imod.add(RegSet("v", "vgprG2LB_BASE", vgprBase + numG2LA)) if kernel["DirectToLdsB"] and kernel["NonDTLTailLoopB"]: imod.add(RegSet("v", "vgprG2LB", "vgprG2LB_BASE", 0)) else: imod.add(self.moduleVgprMacroG2LB) if numG2LMXSA > 0: imod.add(RegSet("v", "vgprG2LMXSA_BASE", vgprBase + numG2LA + numG2LB)) if kernel["DirectToLdsMXSA"] and kernel["NonDTLTailLoopMXSA"]: imod.add(RegSet("v", "vgprG2LMXSA", "vgprG2LMXSA_BASE", 0)) else: imod.add(self.moduleVgprMacroG2LMXSA) if numG2LMXSB > 0: imod.add(RegSet("v", "vgprG2LMXSB_BASE", vgprBase + numG2LA + numG2LB + numG2LMXSA)) if kernel["DirectToLdsMXSB"] and kernel["NonDTLTailLoopMXSB"]: imod.add(RegSet("v", "vgprG2LMXSB", "vgprG2LMXSB_BASE", 0)) else: imod.add(self.moduleVgprMacroG2LMXSB) if numG2LMetadata > 0: imod.add(RegSet("v", "vgprG2LMetadata", vgprBase + numG2LA + numG2LB + numG2LMXSA + numG2LMXSB)) return imod, vgprBase def tailLoopAllocDTLLWVgpr(self, kernel): imod = Module("tailLoopAllocDTLLWVgpr") vgprBase = -1 numLWA = self.states.a.numVgprLocalWriteAddrTailLoop numLWB = self.states.b.numVgprLocalWriteAddrTailLoop numLWMXSA = self.states.mxsa.numVgprLocalWriteAddrTailLoop if kernel["ProblemType"]["MXBlockA"] else 0 numLWMXSB = self.states.mxsb.numVgprLocalWriteAddrTailLoop if kernel["ProblemType"]["MXBlockB"] else 0 numLWM = self.states.m.numVgprLocalWriteAddrTailLoop if numLWA + numLWB + numLWMXSA + numLWMXSB + numLWM > 0: vgprBase = self.vgprPool.checkOutAligned(numLWA + numLWB + numLWMXSA + numLWMXSB + numLWM, 2, tag="tailLoopAllocDTLLWVgpr_vgprBase") imod.addComment0("Check out VGPR (numLWA,numLWB,numLWMXSA,numLWMXSB,numLWM) = (%d,%d,%d,%d,%d)"%(numLWA,numLWB,numLWMXSA,numLWMXSB,numLWM)) if numLWA > 0 and (kernel["DirectToLdsA"] and kernel["NonDTLTailLoopA"]): imod.add(RegSet("v", "vgprLocalWriteAddrA", vgprBase)) if numLWMXSA > 0 and (kernel["DirectToLdsMXSA"] and kernel["NonDTLTailLoopMXSA"]): imod.add(RegSet("v", "vgprLocalWriteAddrMXSA", vgprBase + numLWA)) if numLWB > 0 and (kernel["DirectToLdsB"] and kernel["NonDTLTailLoopB"]): imod.add(RegSet("v", "vgprLocalWriteAddrB", vgprBase + numLWA + numLWMXSA)) if numLWMXSB > 0 and (kernel["DirectToLdsMXSB"] and kernel["NonDTLTailLoopMXSB"]): imod.add(RegSet("v", "vgprLocalWriteAddrMXSB", vgprBase + numLWA + numLWMXSA + numLWB)) if numLWM > 0 and (kernel["DirectToLdsMetadata"] and kernel["NonDTLTailLoopMetadata"]): imod.add(RegSet("v", "vgprLocalWriteAddrMetadata", vgprBase + numLWA + numLWMXSA + numLWB + numLWMXSB)) return imod, vgprBase def tailLoopAllocDTVVgpr(self, kernel, tensorParametersA, tensorParametersB): imodA = Module("tailLoopAllocDTVVgprA") vgprBaseA = -1 if kernel["DirectToVgprA"]: if self.states.packDTVA or self.states.convDTVA: numValuA = self.states.a.numVgprValu numVgprValuPackA = 0 if tensorParametersA["bpe"] < 4 and not kernel["UnrollMajorLDSA"]: if self.states.lrvwTileA > 1: numVgprValuPackA = ceil(kernel["VectorWidthA"] * tensorParametersA["bpe"] / self.states.bpr) * kernel["MIWaveTileA"] // kernel["VectorWidthA"] * kernel["InnerUnroll"] * self.states.numVgprBuffer * kernel["MIInputPerThreadA"] if self.states.packDTVA: # pack DTV case, double the number numVgprValuPackA *= 2 else: numVgprValuPackA = self.states.a.numVgprValuPerBlock * kernel["InnerUnroll"] * self.states.numVgprBufferPackA * (int(4/tensorParametersA["bpeDS"]) - 1) vgprBaseA = self.vgprPool.checkOutAligned(numValuA + numVgprValuPackA, 2, tag="tailLoopAllocDTVVgprA_vgprBaseA") imodA.add(RegSet("v", "vgprValuA_X0_I0_BASE", vgprBaseA)) if numVgprValuPackA > 0: imodA.add(RegSet("v", "vgprValuA_X0_I0_D0_PACK", vgprBaseA + numValuA)) if self.states.packDTVA: #DTVA is a.startVgprValuPack imodA.add(RegSet("v", "vgprG2LA_BASE", vgprBaseA + numValuA)) elif self.states.convDTVA: #DTVA is a.startVgprValu imodA.add(RegSet("v", "vgprG2LA_BASE", vgprBaseA)) imodA.add(self.moduleVgprMacroG2LA) imodA.add(self.moduleVgprMacroValuA) imodA.add(self.moduleVgprMacroValuAPack) else: #DTVA is G2LA itself if kernel["ULSGRODoubleG2L"] == 1: numG2LA = self.states.a.numVgprG2LAllocated*2 else: numG2LA = self.states.a.numVgprG2LAllocated vgprBaseA = self.vgprPool.checkOutAligned(numG2LA, 2, tag="tailLoopAllocDTVVgpr_vgprBaseA") imodA.add(RegSet("v", "vgprG2LA_BASE", vgprBaseA)) imodA.add(self.moduleVgprMacroG2LA) imodB = Module("tailLoopAllocDTVVgprA") vgprBaseB = -1 if kernel["DirectToVgprB"]: if self.states.packDTVB or self.states.convDTVB: numValuB = self.states.b.numVgprValu numVgprValuPackB = 0 if tensorParametersB["bpe"] < 4 and not kernel["UnrollMajorLDSB"]: if self.states.lrvwTileB > 1: numVgprValuPackB = ceil(kernel["VectorWidthB"] * tensorParametersB["bpe"] / self.states.bpr) * kernel["MIWaveTileB"] // kernel["VectorWidthB"] * kernel["InnerUnroll"] * self.states.numVgprBuffer * kernel["MIInputPerThreadB"] if self.states.packDTVB: # pack DTV case, double the number numVgprValuPackB *= 2 else: numVgprValuPackB = self.states.b.numVgprValuPerBlock * kernel["InnerUnroll"] * self.states.numVgprBufferPackB * (int(4/tensorParametersB["bpeDS"]) - 1) vgprBaseB = self.vgprPool.checkOutAligned(numValuB + numVgprValuPackB, 2, tag="tailLoopAllocDTVVgpr_vgprBaseB") imodB.add(RegSet("v", "vgprValuB_X0_I0_BASE", vgprBaseB)) if numVgprValuPackB > 0: imodB.add(RegSet("v", "vgprValuB_X0_I0_D0_PACK", vgprBaseB + numValuB)) if self.states.packDTVB: #DTVB is b.startVgprValuPack imodB.add(RegSet("v", "vgprG2LB_BASE", vgprBaseB + numValuB)) elif self.states.convDTVB: #DTVB is b.startVgprValu imodB.add(RegSet("v", "vgprG2LB_BASE", vgprBaseB)) imodB.add(self.moduleVgprMacroG2LB) imodB.add(self.moduleVgprMacroValuB) imodB.add(self.moduleVgprMacroValuBPack) else: #DTVB is G2LB itself if kernel["ULSGRODoubleG2L"] == 1: numG2LB = self.states.b.numVgprG2LAllocated*2 else: numG2LB = self.states.b.numVgprG2LAllocated vgprBaseB = self.vgprPool.checkOutAligned(numG2LB, 2, tag="tailLoopAllocDTVVgpr_vgprBaseB") imodB.add(RegSet("v", "vgprG2LB_BASE", vgprBaseB)) imodB.add(self.moduleVgprMacroG2LB) imodM = Module("tailLoopAllocDTVVgprA") vgprBaseM = -1 if kernel["ProblemType"]["Sparse"]: if kernel["DirectToVgprSparseMetadata"]: numValuM = self.states.m.numVgprValu vgprBaseM = self.vgprPool.checkOut(numValuM, tag="tailLoopAllocDTVVgpr_vgprBaseM") imodM.add(RegSet("v", "vgprValuMetadata_X0_I0_BASE", vgprBaseM)) imodM.add(self.moduleVgprMacroValuM) return ([vgprBaseA, imodA],[vgprBaseB, imodB],[vgprBaseM, imodM]) ############################################################################## # Using wider load instructions to improve the GR efficiency in tail loop. # If loading size is smaller than a dword(32bit), it will return 0 instead. # Need to use buffer_load_d16 to load the data which is out of boundary. ############################################################################## def tailLoopGlobalRead(self, kernel, tPA, tPB, doA, doB): imod = Module("tailLoopGlobalRead") glvwWorkaround = 8 * kernel["ProblemType"]["DataType"].numRegisters() dataTypeA = kernel["ProblemType"]["MacDataTypeA"] if tPA["glvw"] < glvwWorkaround else \ kernel["ProblemType"]["DataTypeA"] dataTypeB = kernel["ProblemType"]["MacDataTypeB"] if tPB["glvw"] < glvwWorkaround else \ kernel["ProblemType"]["DataTypeB"] numElementsPerLoadA = -1 numElementsPerLoadB = -1 if dataTypeA.isHalf() or dataTypeA.isBFloat16(): if tPA["glvw"] > 1 and kernel["AssertSummationElementMultiple"] % 2 == 0: numElementsPerLoadA = 2 if dataTypeB.isHalf() or dataTypeB.isBFloat16(): if tPB["glvw"] > 1 and kernel["AssertSummationElementMultiple"] % 2 == 0: numElementsPerLoadB = 2 if numElementsPerLoadA == 2: doA = False if numElementsPerLoadB == 2: doB = False tagList = ["AddressA", "AddressB", "WrapUA", "WrapUB", "StaggerU", "WGM", \ "StaggerUIter", "GlobalReadIncsA", "GlobalReadIncsB", \ "StridesA", "StridesB", "ShadowLimitA", "ShadowLimitB"] if kernel["ProblemType"]["MXBlockA"]: tagList += ["GlobalReadIncsMXSA", "ShadowLimitMXSA"] if kernel["ProblemType"]["MXBlockB"]: tagList += ["GlobalReadIncsMXSB", "ShadowLimitMXSB"] lastRegTag = None spool = self.sgprPool.getPool() imod.addComment1("release sgprs that will not be used") for i in range(0, self.sgprPool.size()): regTag = spool[i].tag if regTag != lastRegTag: lastRegTag = regTag if (lastRegTag not in self.states.nonPostLoopSgpr) and \ (spool[i].status == RegisterPool.Status.InUse) and \ (lastRegTag in tagList): imod.add(self.undefineSgpr(regTag)) loadALabel = Label(label="LoadA", comment="") loadBLabel = Label(label="LoadB", comment="") mergeALabel = Label(label="MergeA", comment="") mergeBLabel = Label(label="MergeB", comment="") tailGlobalLoadEndLabel = Label(label="TailGlobalLoadEnd", comment="") checkOtherLoadALabel = Label(label="CheckOtherLoadA", comment="") checkAddrALabel = Label(label="CheckAddrA", comment="") checkOtherLoadBLabel = Label(label="CheckOtherLoadB", comment="") checkAddrBLabel = Label(label="CheckAddrB", comment="") lspA = kernel[tPA["lsp"]] lscA = kernel[tPA["lsc"]] lspB = kernel[tPB["lsp"]] lscB = kernel[tPB["lsc"]] nlcA = kernel["NumLoadsCoalescedA"] nlcB = kernel["NumLoadsCoalescedB"] nlpA = kernel["NumLoadsPerpendicularA"] nlpB = kernel["NumLoadsPerpendicularB"] numElementsPer4BytesA = int(4 / tPA["bpeGR"]) numElementsPer4BytesB = int(4 / tPB["bpeGR"]) maxNumOOBElementsA = numElementsPer4BytesA - 1 maxNumOOBElementsB = numElementsPer4BytesB - 1 numDwordA = int(tPA["glvw"] * tPA["bpeGR"]) >> 2 numDwordB = int(tPB["glvw"] * tPB["bpeGR"]) >> 2 numDwordA = 1 if numDwordA == 0 else numDwordA numDwordB = 1 if numDwordB == 0 else numDwordB numTmpVgpr = maxNumOOBElementsA * numDwordA + maxNumOOBElementsB * numDwordB # 2 fo 16b numSingleSgpr = 10 numPairSgpr = 2 singSgprList = [] # store single sgpr pairSgprList = [] # store aligned2 sgpr for num in range(numSingleSgpr): tmpSgpr = self.sgprPool.checkOut(1, tag="tailLoopGlobalRead_Single_tmpSgpr", preventOverflow=False) singSgprList.append(tmpSgpr) for num in range(numPairSgpr): tmpSgpr = self.sgprPool.checkOutAligned(2, 2, tag="tailLoopGlobalRead_Pair_tmpSgpr", preventOverflow=False) pairSgprList.append(tmpSgpr) # If load bytes of each thread is a multiple of 4-byte, data can be fully loaded by using # buffer_load_dword, if not, load single element by buffer_load_short_d16 additionally. It's # affected by WaveSeparateGlobalRead, MT, DU. # This function finds the last element locates in which global read tile (load range of each # global read instruction. # sReloadFlag: 0 is no need to reload. # sLoadTileIdx: The last element locates in which load tile. # sLoadNum: How many elements need to be reloaded. # sValidBytes: How many valid bytes can be loaded in tail loop. def generateFindTheLastElementLocation(tc): do = doA if tc == "A" else doB strSize = "SizeI" if tc == "A" else "SizeJ" strMacroTile = "MacroTile0" if tc == "A" else "MacroTile1" strWG = "WorkGroup0" if tc == "A" else "WorkGroup1" tP = tPA if tc == "A" else tPB nlp = nlpA if tc == "A" else nlpB nlc = nlcA if tc == "A" else nlcB lsp = lspA if tc == "A" else lspB lsc = lscA if tc == "A" else lscB strWSGR = "WaveSeparateGlobalReadA" if tc == "A" else "WaveSeparateGlobalReadB" maxNumOOBElements = maxNumOOBElementsA if tc == "A" else maxNumOOBElementsB sReloadFlag = singSgprList[2] if tc == "A" else singSgprList[3] sTmp0 = singSgprList[0] if tc == "A" else singSgprList[1] sTmp1 = singSgprList[4] if tc == "A" else singSgprList[5] sTmp2 = singSgprList[6] if tc == "A" else singSgprList[7] sTmp3 = singSgprList[8] if tc == "A" else singSgprList[9] sx2Tmp0 = pairSgprList[0] sx2Tmp1 = pairSgprList[1] sLoadTileIdx = sTmp0 sLoadNum = sTmp1 sValidBytes = sTmp2 loopIdx = self.states.unrollIdx if do: imod.addComment("Calculate %s %% %s"%(strSize, strMacroTile)) imod.add(SMulI32(dst=sgpr(sTmp0), src0=sgpr(strWG), src1=kernel[strMacroTile], \ comment="Calculate the remaining dimension along I/J direction.")) imod.add(SSubU32(dst=sgpr(sTmp0), src0=sgpr(strSize), src1=sgpr(sTmp0), \ comment="Calculate the remaining dimension along I/J direction.")) imod.add(SMulI32(dst=sgpr(sTmp0), src0=sgpr(sTmp0), src1=int(tP["bpeGR"]), \ comment="In bytes")) imod.add(SAndB32(dst=sgpr(sTmp1), src0=sgpr("SizeL"), src1=(kernel["DepthU"] - 1), \ comment="Calculate the remaining dimension along L direction.")) imod.add(SLShiftRightB32(dst=sgpr(sx2Tmp1), shiftHex=hex(log2(lsc)), \ src=sgpr(sTmp1), comment="Divided by lsc(%s)"%(lsc))) imod.add(self.s_mul_u64_u32(sgpr(sValidBytes), sgpr(sReloadFlag), sgpr(sTmp0), sgpr(sTmp1), \ comment="Calculate total number of valid elements.")) imod.add(SCmpGtU32(src0=sgpr(sReloadFlag), src1=0)) imod.add(SCMovB32(dst=sgpr(sValidBytes), src=hex(0xFFFFFFFF), \ comment="If valid elements > max(U32), set the value to max")) if (kernel[strWSGR] == 0): tmpSgpr = sTmp0 else: tmpSgpr = sTmp3 imod.add(SSubU32(dst=sgpr(sTmp1), src0=sgpr(strSize), src1=1, \ comment="sLoadTileIdx starts from 0")) imod.addComment("Calculate %s - 1 %% %s"%(strSize, strMacroTile)) imod.add(scalarStaticDivideAndRemainder(tmpSgpr, tmpSgpr, sTmp1, \ kernel[strMacroTile], \ ContinuousRegister(sx2Tmp0, 2), 1)) if (kernel[strWSGR] > 0): imod.addComment("Calculate %s %% (nlp * lsp)"%(tmpSgpr)) imod.add(scalarStaticDivideAndRemainder(sTmp0, sTmp0, tmpSgpr, \ (nlp * lsp), \ ContinuousRegister(sx2Tmp0, 2), 1)) imod.add(SLShiftRightB32(dst=sgpr(sTmp0), shiftHex=hex(log2(lsp)), \ src=sgpr(sTmp0), comment="Divide lsp to get the load tile index")) imod.add(SMulI32(dst=sgpr(sTmp0), src0=sgpr(sTmp0), src1=nlc, comment="Multiply nlc")) imod.add(SAddI32(dst=sgpr(sLoadTileIdx), src0=sgpr(sTmp0), src1=sgpr(sx2Tmp1), \ comment="")) imod.add(scalarStaticDivideAndRemainder(sx2Tmp0, sTmp1, "SizesSum+%u"%loopIdx, \ kernel["DepthU"], ContinuousRegister(sx2Tmp0, 2), \ 2)) imod.add(SAndB32(dst=sgpr(sLoadNum), src0=sgpr(sTmp1), src1=(tP["glvw"] - 1), \ comment="sLoadNum = (SizesSum+%u"%(loopIdx)+" mod DU) & glvw")) imod.add(SAndB32(dst=sgpr(sReloadFlag), src0=sgpr(sLoadNum), src1=hex(maxNumOOBElements), \ comment="")) # The end of generateFindTheLastElementLocation # Generate code for jumping branch to correct load tile (sLoadTileIdx) and load correct number # of elements by checking sLoadNum. # Extract the global load instraction and merge instruction of each load tile and load number # from globalReadGuardK. # What "LOAD" do is buffer_load_short_d16 v0, vAddr, sgprSrd, soffset. # What "MERGE" do is s_waitcnt vmcnt(0) # v_or_b32 vData, vData, v0 # Flow: GRA -> GRB -> MergeA -> MergeB. def generateLoadAndMergeSingleElement(): def generateJumpBranch(idx): def generateDetailsForBehavior(tP, tmpVgpr, behavior, jumpLabel, tileSgpr, kSgpr): self.param = TailOptParams() tc = tP["tensorChar"] bpe = tP["bpeGR"] numThreadsPerp = kernel[tP["lsp"]] numThreadsCoal = int(kernel["NumThreads"] / numThreadsPerp) nlp = kernel["NumLoadsPerpendicular%c"%(tc)] nlc = kernel["NumLoadsCoalesced%c"%(tc)] numTiles = nlp * nlc numElementsPer4Bytes = int(4 / bpe) def generateFunction(idx, bevavior, jumpLabel, tileSgpr, kSgpr): tileLabelsList = [] kLabelsList = [] if (idx >= 0): labelStr = bevavior + "_"+tc+str(idx) labelTmp = Label(label=labelStr, comment="") tileLabelsList.append(labelTmp) if (idx != 0): imod.add(SCmpEQU32(src0=sgpr(tileSgpr), src1=idx, comment="")) imod.add(SCBranchSCC1(labelName=labelTmp.getLabelName(), comment="")) generateFunction(idx-1, bevavior, jumpLabel, tileSgpr, kSgpr) imod.add(labelTmp) for i in range(tP["glvw"], 0, -1): if i % numElementsPer4Bytes != 0: labelStrK = labelStr+"_K"+str(i) labelTmp = Label(label=labelStrK, comment="") kLabelsList.append(labelTmp) if idx == (numTiles - 1): finalLoop = 1 else: finalLoop = 0 self.param.idx = idx self.param.jumpLabel = jumpLabel self.param.tmpVgpr = tmpVgpr self.param.kLabelsList = kLabelsList self.param.behavior = behavior self.param.kSgpr = kSgpr self.param.finalLoop = finalLoop imod.add(self.globalReadGuardK(kernel, tP, 1, self.param)) generateFunction(numTiles - 1, behavior, jumpLabel, tileSgpr, kSgpr) # The End of generateDetailsForBehavior tc = ["A", "B"] do = [doA, doB] tP = [tPA, tPB] loadLabel = [loadALabel, loadBLabel] sReloadFlag = [sReloadFlagA, sReloadFlagB] sLoadTileIdx = [sLoadTileIdxA, sLoadTileIdxB] sLoadNum = [sLoadNumA, sLoadNumB] tmpVgprStartIdx = [0, (maxNumOOBElementsA * numDwordA)] mergeLabel = [mergeALabel, mergeBLabel] numElementsPerLoad = [numElementsPerLoadA, numElementsPerLoadB] if (idx >= 0 and idx < len(tc)): if tc[idx] == "A": jumpLabel = loadBLabel if doA and doB else mergeALabel else: jumpLabel = mergeALabel if doA and doB else mergeBLabel if do[idx]: imod.addComment1("load single element for %s"%(tc[idx])) imod.add(loadLabel[idx]) imod.add(SCmpEQU32(src0=sgpr(sReloadFlag[idx]), src1=0, \ comment="Valid loading size per thread is multiples of 4 bytes")) imod.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), comment="Skip loading "+tc[idx])) generateDetailsForBehavior(tP[idx], tmpVgpr + tmpVgprStartIdx[idx], "LOAD", jumpLabel, sLoadTileIdx[idx], sLoadNum[idx]) generateJumpBranch(idx+1) idx = idx ^ 1 if do[idx]: if tc[idx] == "A": jumpLabel = mergeBLabel if doA and doB else checkOtherLoadALabel else: jumpLabel = checkOtherLoadALabel if doA and doB else checkOtherLoadBLabel imod.addComment1("merge single element for %s"%(tc[idx])) imod.add(mergeLabel[idx]) imod.add(SCmpEQU32(src0=sgpr(sReloadFlag[idx]), src1=0, \ comment="Valid loading size per thread is multiples of 4 bytes")) imod.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), comment="Skip mergeing "+tc[idx])) generateDetailsForBehavior(tP[idx], tmpVgpr + tmpVgprStartIdx[idx], "MERGE", \ jumpLabel, sLoadTileIdx[idx], sLoadNum[idx]) generateJumpBranch(0) # The end of generateJumpBranch # The end of generateFindTheLastElementLocation # Generate tht jump branch for checking other load tiles. def jumpLabel(tP, tmpSgpr, jumpLabel): tc = tP["tensorChar"] nlp = kernel["NumLoadsPerpendicular%c"%(tc)] nlc = kernel["NumLoadsCoalesced%c"%(tc)] numTiles = nlp * nlc def func(idx, tmpSgpr): tileLabelsList = [] if (idx >= 0): labelStr = tc+str(idx) labelTmp = Label(label=labelStr, comment="") tileLabelsList.append(labelTmp) if (idx != 0): imod.add(SCmpEQU32(src0=sgpr(tmpSgpr), src1=idx, comment="")) imod.add(SCBranchSCC1(labelName=labelTmp.getLabelName(), comment="")) func(idx-1, tmpSgpr) imod.add(labelTmp) if tc == "A": addrbase = self.globalread_gpr_record.a.addrVgpr[idx] offset = self.globalread_gpr_record.a.offset[idx] else: addrbase = self.globalread_gpr_record.b.addrVgpr[idx] offset = self.globalread_gpr_record.b.offset[idx] if kernel["_UseSgprForGRO"]: imod.add(VAddU32(dst=vgpr(tmpVgpr), src0=vgpr(addrbase), src1=offset)) else: imod.add(VMovB32(dst=vgpr(tmpVgpr), src=vgpr(addrbase))) if idx != (numTiles - 1): imod.add(SBranch(labelName=jumpLabel.getLabelName(), comment="")) func(numTiles - 1, tmpSgpr) # The end of jumpLabel. # Even we have reloaded the load tile which the last element locates in, there might be some # other tiles access the last element./ # So we check the upper load next step, and do the check recursively untill no threads in the # load tile accesses the last element. def generateReLoadLoop(tc): do = doA if tc == "A" else doB nlp = nlpA if tc == "A" else nlpB nlc = nlcA if tc == "A" else nlcB checkOtherLoadLabel = checkOtherLoadALabel if tc == "A" else checkOtherLoadBLabel sLoadTileIdx = sLoadTileIdxA if tc == "A" else sLoadTileIdxB tP = tPA if tc == "A" else tPB checkAddrLabel = checkAddrALabel if tc == "A" else checkAddrBLabel sValidBytes = sValidBytesA if tc == "A" else sValidBytesB loadLabel = loadALabel if tc == "A" else loadBLabel imod.add(checkOtherLoadLabel) # At the beginning of the loop, jump to the next step if: # no need to reload (sReloadFlagA/B == 0) or # sLoadCnt == total number of global load subtiles. # Increase sLoadCnt nlc in every loop. # For generating for loop B: # LoopA/B share the same sgpr to record the load cnt. Reset it at the beginning of loop B. # Force sReloadFlagB to be 0 to skip reload B when reload A. Reset it at the beginning of # loop B. if do: if tc == "A": imod.add(SCmpEQU32(src0=sgpr(sReloadFlagA), src1=0, comment="Noneed to load single element fo A?")) if doB: imod.add(SCBranchSCC1(labelName=checkOtherLoadBLabel.getLabelName(), comment="")) else: imod.add(SCBranchSCC1(labelName=tailGlobalLoadEndLabel.getLabelName(), comment="")) imod.add(SAddU32(dst=sgpr(sLoadCnt), src0=sgpr(sLoadCnt), src1=nlc)) imod.add(SCmpEQU32(src0=sgpr(sLoadCnt), src1=(nlp * nlc), comment="Have reloaded all subtiles?")) if doB: imod.add(SCMovB32(dst=sgpr(sLoadCnt), src=0, comment="Reset loop count")) imod.add(SCBranchSCC1(labelName=checkOtherLoadBLabel.getLabelName(), comment="")) else: imod.add(SCBranchSCC1(labelName=tailGlobalLoadEndLabel.getLabelName(), comment="")) else: if doA: imod.add(SMovB32(dst=sgpr(sReloadFlagA), src=0, comment="Force to skip reload A")) imod.add(SCmpEQU32(src0=sgpr(sLoadCnt), src1=0, comment="Loop start?")) imod.add(SCMovB32(dst=sgpr(sReloadFlagB), src=sgpr(sBackupSkipLoadB), comment="Restore sReloadFlagB for B")) imod.add(SCmpEQU32(src0=sgpr(sReloadFlagB), src1=0, comment="Noneed to load single element for B?")) imod.add(SCBranchSCC1(labelName=tailGlobalLoadEndLabel.getLabelName(), comment="")) imod.add(SAddU32(dst=sgpr(sLoadCnt), src0=sgpr(sLoadCnt), src1=nlc)) imod.add(SCmpEQU32(src0=sgpr(sLoadCnt), src1=(nlp * nlc), comment="Have reloaded all subtiles?")) imod.add(SCBranchSCC1(labelName=tailGlobalLoadEndLabel.getLabelName(), comment="")) imod.add(SSubI32(dst=sgpr(sLoadTileIdx), src0=sgpr(sLoadTileIdx), src1=nlc, \ comment="Check the upper subtile")) imod.add(SCmpLtI32(src0=sgpr(sLoadTileIdx), src1=0, comment="")) imod.add(SCSelectB32(dst=sgpr(sCmpLoadStartAddrStatusx2), src0=(nlp * nlc), src1=0, \ comment="Back to the last subtile")) imod.add(SAddI32(dst=sgpr(sLoadTileIdx), src0= sgpr(sLoadTileIdx), src1=sgpr(sCmpLoadStartAddrStatusx2), \ comment="If currently reload the first subtile, \ check the last subtile next.")) jumpLabel(tP, sLoadTileIdx, checkAddrLabel) imod.add(checkAddrLabel) imod.add(VSubU32(dst=vgpr(tmpVgpr), src0=vgpr(tmpVgpr), src1=int(self.states.srdShiftLeft[tc] * tP["bpeGR"]), comment="sub prepad")) loadRangePerThread = int(tP["glvw"] * tP["bpeGR"] - 1) imod.add(VAddU32(dst=vgpr(tmpVgpr+1), src0=vgpr(tmpVgpr), src1=loadRangePerThread, \ comment="Calculate load range per thread")) imod.add(VCmpLtI32(dst=sgpr(sCmpLoadStartAddrStatusx2, 2), src0=vgpr(tmpVgpr), \ src1=sgpr(sValidBytes), \ comment="If loading start address < total valid bytes?")) imod.add(VCmpGEI32(dst=sgpr(sCmpLoadEndAddrStatusx2, 2), src0=vgpr(tmpVgpr+1), \ src1=sgpr(sValidBytes), \ comment="If loading end address >= total valid bytes?")) imod.add(SAndB32(dst=sgpr(sCmpLoadStartAddrStatusx2), \ src0=sgpr(sCmpLoadStartAddrStatusx2), \ src1=sgpr(sCmpLoadEndAddrStatusx2), \ comment="Find threads which access the last element")) imod.add(SAndB32(dst=sgpr(sCmpLoadStartAddrStatusx2+1), \ src0=sgpr(sCmpLoadStartAddrStatusx2+1), \ src1=sgpr(sCmpLoadEndAddrStatusx2+1), \ comment="Find thread that access the last element")) imod.add(SAddU32(dst=sgpr(sCmpLoadStartAddrStatusx2), \ src0=sgpr(sCmpLoadStartAddrStatusx2), \ src1=sgpr(sCmpLoadStartAddrStatusx2+1), \ comment="Find thread that access the last element")) imod.add(SCmpLgU32(src0=sgpr(sCmpLoadStartAddrStatusx2), src1=0, \ comment="Have threads access the last element?")) if tc == "A": if doB: imod.add(SCMovB32(dst=sgpr(sReloadFlagB), src=0, comment="Skip reload B temporarily")) imod.add(SCSelectB32(dst=sgpr(sLoadCnt), src0=sgpr(sLoadCnt), src1=0, \ comment="Reset loop count if needed")) imod.add(SCBranchSCC1(labelName=loadLabel.getLabelName(), comment="Reload "+tc)) # The End of generateReLoadLoop # These sgprs must be the same as used in generateFindTheLastElementLocation sLoadTileIdxA = singSgprList[0] sLoadTileIdxB = singSgprList[1] sReloadFlagA = singSgprList[2] sReloadFlagB = singSgprList[3] sLoadNumA = singSgprList[4] sLoadNumB = singSgprList[5] sValidBytesA = singSgprList[6] sValidBytesB = singSgprList[7] sLoadCnt = singSgprList[8] sBackupSkipLoadB = singSgprList[9] sCmpLoadStartAddrStatusx2 = pairSgprList[0] sCmpLoadEndAddrStatusx2 = pairSgprList[1] imod.addComment1("find the last element location for a") generateFindTheLastElementLocation("A") imod.addComment1("find the last element location for b") generateFindTheLastElementLocation("B") if doA or doB: tmpVgpr = self.vgprPool.checkOut(numTmpVgpr, tag="tailLoopGlobalRead_tmpVgpr") imod.add(SMovB32(dst=sgpr(sLoadCnt), src=0, comment="Set loop count = 0")) if doA and doB: imod.add(SMovB32(dst=sgpr(sBackupSkipLoadB), src=sgpr(sReloadFlagB), \ comment="Backup and will be restored in label_CheckB_OOB")) generateLoadAndMergeSingleElement() imod.addComment1("reload loop for a: check if there's other load range need to be reloaded") generateReLoadLoop("A") imod.addComment1("reload loop for b: check if there's other load range need to be reloaded") generateReLoadLoop("B") imod.addComment1("global read for tail done") imod.add(tailGlobalLoadEndLabel) if doA or doB: self.vgprPool.checkIn(tmpVgpr) for s in singSgprList: self.sgprPool.checkIn(s) for s in pairSgprList: self.sgprPool.checkIn(s) return imod ############################################################################## # Emit code to compute loop iterations for GSU. # See same function in KernelWriterSource.py for background explanation # This function is used to compute number of loop iters and also # for computing the global read increment for GSU case. # For multiple summation, the number of loop iterations needs to be reset # for each iteration so replicate the code in addr inc and at open of unroll loop # tmpSgpr is allocation of at least 3 tmpSgpr # Output: SGPR(destName) contains the number of unroll iterations for # this workgroup. ############################################################################## def calculateLoopNumIterGsu(self, kernel, destName, tmpSgprRes: ContinuousRegister): module = Module("calculateLoopNumIterGsu") loopCounter = sgpr(destName) quotient = destName remainder = "GSUSumIdx+1" # numIterPerWgRemainder dividend = destName tmpVgpr = self.vgprPool.checkOut(2, tag="calculateLoopNumIterGsu_tmpVgpr") tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2) module.add(SAndB32(dst=sgpr(remainder), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(scalarUInt32DivideAndRemainder(quotient, dividend, remainder, remainder, tmpVgprRes, wavewidth=kernel["WavefrontSize"])) self.vgprPool.checkIn(tmpVgpr) # if gsuSumIdx < numIterPerWgRemainder module.add(SAddU32(dst=sgpr(tmpSgprRes.idx), src0=1, src1=loopCounter, comment="tmp<-numIterMyWg+1")) module.add(SCmpLtU32(src0=sgpr("GSUSumIdx"), src1=sgpr("GSUSumIdx+1"), comment="gsuSumIdx < numIterPerWgRemainder")) module.add(SCMovB32(dst=loopCounter, src=sgpr(tmpSgprRes.idx), comment="numIterMyWg++ if needed")) return module def calculateLoopNumIterOffsetGsu(self, kernel, destName, tmpSgprRes: ContinuousRegister): module = Module("calculateLoopNumIterOffsetGsu") loopCounter = sgpr(destName) quotient = destName remainder = "GSUSumIdx+1" # numIterPerWgRemainder dividend = destName tmpVgpr = self.vgprPool.checkOut(2, tag="calculateLoopNumIterOffsetGsu_tmpVgpr") tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2) module.add(SAndB32(dst=sgpr(remainder), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(scalarUInt32DivideAndRemainder(quotient, dividend, remainder, remainder, tmpVgprRes, wavewidth=kernel["WavefrontSize"])) self.vgprPool.checkIn(tmpVgpr) # calculate offset number of loop iterations for each wg stmp = tmpSgprRes.idx module.add(SMulI32(dst=sgpr(stmp+1), src0=loopCounter, src1=sgpr("GSUSumIdx"), comment="quotient*GSUSumIdx")) module.add(SAddU32(dst=sgpr(stmp+0), src0=1, src1=loopCounter, comment="quotient+1")) module.add(SAddU32(dst=sgpr(stmp+1), src0=sgpr(stmp+1), src1=sgpr("GSUSumIdx+1"), comment="quotient*GSUSumIdx+remainder")) module.add(SMulI32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr("GSUSumIdx"), comment="(quotient+1)*GSUSumIdx")) # if gsuSumIdx < numIterPerWgRemainder module.add(SCmpLtU32(src0=sgpr("GSUSumIdx"), src1=sgpr("GSUSumIdx+1"), comment="gsuSumIdx < numIterPerWgRemainder")) module.add(SCSelectB32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr(stmp+1), comment="(quotient+1)*GSUSumIdx if needed")) return module ############################################################################## # Calculate Loop Num Iter # loopIdx is the index of the loop (used for contractions with multiple summations) # 0 is outermost; self.states.unrollIdx is the unroll index. # -1 is tail loop (used only for the unroll loop) ############################################################################## def calculateLoopNumIter(self, kernel, tPA, tPB, loopIdx, tailloopInNll=False, NLLindex=0): module = Module("calculateLoopNumIter") tailLoop = loopIdx < 0 if tailLoop: loopIdx = self.states.unrollIdx loopDim = kernel["ProblemType"]["IndicesSummation"][loopIdx] loopChar = self.states.indexChars[loopDim] ######################################## # Tail Loop if tailLoop: with self.allocTmpSgpr(4, tag="calculateLoopNumIter_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx loopCounterName = self.loopCounterName(kernel, loopIdx) module.addSpaceLine() if kernel["SuppressNoLoadLoop"]: if kernel["HalfPLR"]: # In HalfPLR case, TDM will be disabled in the last unroll loop iterations. Also LDS buffer # is not aligned for TDM & local read address. So we need to re-enable TDM & align LDS buffer SkipHalfPLRAdjustLabel = Label("Skip_HPLR_Adjust", "") module.add(SCmpEQU32(src0=sgpr("OrigLoopCounter"), src1=0, comment="skip if main loop was not executed")) module.add(SCBranchSCC1(labelName=SkipHalfPLRAdjustLabel.getLabelName(), comment="")) module.addComment0("HalfPLR: re-enable TDM & align LDS buffer") if kernel["NumWaves"] > 1: module.add(SMovB32(dst=sgpr("tdmAGroup0+0"), src=1, comment="")) module.add(self.tdmSwapLdsOffset(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.add(SMovB32(dst=sgpr("tdmMXSAGroup0+0"), src=1, comment="")) module.add(self.tdmSwapLdsOffset(kernel, tPA["MX"])) else: module.add(SMovB32(dst=sgpr("tdmAGroup0+0"), src=1, comment="")) module.add(self.tdmSwapLdsOffset(kernel, tPA)) module.add(SMovB32(dst=sgpr("tdmBGroup0+0"), src=1, comment="")) module.add(self.tdmSwapLdsOffset(kernel, tPB)) if kernel["ProblemType"]["MXBlockA"]: module.add(SMovB32(dst=sgpr("tdmMXSAGroup0+0"), src=1, comment="")) module.add(self.tdmSwapLdsOffset(kernel, tPA["MX"])) if kernel["ProblemType"]["MXBlockB"]: module.add(SMovB32(dst=sgpr("tdmMXSBGroup0+0"), src=1, comment="")) module.add(self.tdmSwapLdsOffset(kernel, tPB["MX"])) # Undo HPLR last-body dangling +=split (matching -= lives in next body). # Leak only happens when >= 2 unrolled bodies executed (LC_init > 1), since # the first body's end-of-body +=split is undone by the next body's incCode. if kernel["TDMSplit"] and not kernel["ProblemType"]["Sparse"] \ and kernel["PrefetchGlobalRead"] >= 2: SkipUndoLabel = Label("Skip_TDMSplit_Undo", "") module.add(SCmpLeU32(src0=sgpr("OrigLoopCounter"), src1=1, comment="skip TDMSplit undo if only 1 main iter ran (no cross-body leak)")) module.add(SCBranchSCC1(labelName=SkipUndoLabel.getLabelName(), comment="")) if kernel["NumWaves"] > 1: with self.allocTmpSgpr(2, tag="tdmSplitTailUndo") as undoTmp: gIncIdx = undoTmp.idx scratchIdx = undoTmp.idx + 1 ldsIncOp = self._tdmSplitMultiWaveInc(module, kernel, gIncIdx, scratchIdx, tPA, tPB) module.add(SSubU32(sgpr("tdmAGroup0+1"), sgpr("tdmAGroup0+1"), ldsIncOp, "TDMSplit tail: undo dangling +=ldsSplit")) module.add(SSubU32(sgpr("tdmAGroup0+2"), sgpr("tdmAGroup0+2"), sgpr(gIncIdx), "TDMSplit tail: undo dangling +=globalSplit")) else: module.add(SSubU32(sgpr("tdmAGroup0+1"), sgpr("tdmAGroup0+1"), sgpr("tdmALdsSplitIncs"), "TDMSplit tail: undo dangling +=ldsSplit (A)")) module.add(SSubU32(sgpr("tdmAGroup0+2"), sgpr("tdmAGroup0+2"), sgpr("tdmAGlobalSplitIncs"), "TDMSplit tail: undo dangling +=globalSplit (A)")) module.add(SSubU32(sgpr("tdmBGroup0+1"), sgpr("tdmBGroup0+1"), sgpr("tdmBLdsSplitIncs"), "TDMSplit tail: undo dangling +=ldsSplit (B)")) module.add(SSubU32(sgpr("tdmBGroup0+2"), sgpr("tdmBGroup0+2"), sgpr("tdmBGlobalSplitIncs"), "TDMSplit tail: undo dangling +=globalSplit (B)")) module.add(SkipUndoLabel) module.add((SWaitCnt(dscnt=0, comment="HalfPLR: prevent dummy ds_load from polluting LDS buffer"))) module.add((SBarrier())) module.add(SkipHalfPLRAdjustLabel) else: # If the tail loop is suppressed, then final iterations will have moved the Srd base forward # (and also moved back the srd shadow limit) and slammed Limit to 0, so need to 'undo' # those increments - see setTailSrd assert(kernel["PrefetchGlobalRead"] == 1) #if >1 would need a multiply here module.add(SCmpEQU32(src0=sgpr("OrigLoopCounter"), src1=0, comment="completely skipped unroll loop?")) module.add(SCSelectB32(dst=sgpr(tmpSgpr+0), src0=0, src1=sgpr("GlobalReadIncsA"), comment="force to 0?")) module.add(SCSelectB32(dst=sgpr(tmpSgpr+1), src0=0, src1=sgpr("GlobalReadIncsB"), comment="force to 0?")) module.add(self.setTailSrd(tPA, sgpr(tmpSgpr+0))) module.addSpaceLine() module.add(self.setTailSrd(tPB, sgpr(tmpSgpr+1))) module.addSpaceLine() if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.add(SCSelectB32(dst=sgpr(tmpSgpr+0), src0=0, src1=sgpr("GlobalReadIncsMetadata"), comment="force to 0?")) tP = tPB if kernel["ProblemType"]["Sparse"] == 2 else tPA module.add(self.setTailSrd(tP, sgpr(tmpSgpr+0))) module.addSpaceLine() # LOCAL_SPLITU * min(sizeL % LOCAL_DEPTHU, DEPTHU / LOCAL_SPLITU) module.addComment("numIter%s = LOCAL_SPLITU * min(size%s %% LOCAL_DEPTHU, DEPTHU / LOCAL_SPLITU)" \ % (self.states.unrollChar, self.states.unrollChar)) # size % DepthU module.add(scalarStaticDivideAndRemainder(tmpSgpr, loopCounterName, \ "SizesSum+%u"%loopIdx, kernel["DepthU"], ContinuousRegister(tmpSgpr+2, 2), 2)) loopCounter = sgpr(loopCounterName) if tailloopInNll and NLLindex == 0: # tailLoop in NLL case if kernel["StreamK"]: # StreamK + TailLoopINNLL case # skip TailLoopINNLL if StreamK WG not processing final iteration # Check if tile finished sIpt = self.acquireStreamKConstSgpr(kernel, "ItersPerTile") if self.isStreamKConstantsToVgprEnabled(kernel): module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(self.states.skConstVgprs["ItersPerTile"]))) module.add(SCmpLtU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr(sIpt), comment="Check if WG processes final iteration of tile")) self.releaseStreamKConstSgpr(sIpt) module.add(SCMovB32(dst=loopCounter, src=0, comment="This WG not completing tile")) module.add(SCmpEQU32(src0=loopCounter, src1=0, comment="numIter%s == 0"%loopChar)) EndOfTailLoopInNLLLabel = Label("TailLoopInNLLEnd%s"%(loopChar), "" ) module.add(SCBranchSCC1(labelName=EndOfTailLoopInNLLLabel.getLabelName(), comment="skip TailLoopInNLL code")) maxUnit = self.states.tailloopInNllmaxUnit if maxUnit > 1: module.add(SAndB32(dst=sgpr(tmpSgpr+2), src0=loopCounter, src1=maxUnit-1, \ comment="if summation is not multiple of %s, skip TailLoopInNLL"%maxUnit)) module.add(SCBranchSCC1(labelName=EndOfTailLoopInNLLLabel.getLabelName(), comment="skip TailLoopInNLL code and use TailLoop")) if kernel["GlobalSplitU"] != 0: # skip tailloopInNll code if GSU>1 module.add(SAndB32(dst=sgpr(tmpSgpr+2), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SCmpGtU32(src0=sgpr(tmpSgpr+2), src1=1, comment="GSU > 1 ?")) module.add(SCBranchSCC1(labelName=EndOfTailLoopInNLLLabel.getLabelName(), comment="skip TailLoopInNLL code and use TailLoop")) if kernel["LocalSplitU"] > 1: # we cannot set loopCounter zero and skip tail loop because we need all waves to do global read. # in order to check the k index, we have to keep the offset to check boundary. # | lsu0 | lsu1 | lsu2 | lsu3| # |----d----| # | o0 |o1| keep this offset of each lsu. # For example, 'o1' is the offset of lsu1. '-o0' is the offset of lsu0. This offset can be negative. dividend = tmpSgpr+2 tmpVgpr = self.vgprPool.checkOutAligned(2, 2, tag="calculateLoopNumIter_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) wave_id = self.vgprPool.checkOut(1, tag="calculateLoopNumIter_wave_id") numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] module.add(SMovB32(dst=sgpr(dividend), src=hex(kernel["DepthU"]//kernel["LocalSplitU"]), comment="DEPTHU / LOCAL_SPLITU" )) module.add(vectorStaticDivide(wave_id, "Serial", kernel["WavefrontSize"], tmpVgprRes)) module.add(vectorStaticDivide(wave_id, wave_id, numWaves, tmpVgprRes, comment="LSU offset: Get LSU wave_id")) #module.add(VAddU32(vgpr(wave_id), vgpr(wave_id), 1, "add 1")) module.add(VMulLOU32(vgpr(wave_id), sgpr(dividend), vgpr(wave_id), comment="multiply by (DEPTHU / LOCAL_SPLITU)")) module.add(SNop(waitState=0, comment="Wait to read lane")) module.add(VReadfirstlaneB32(dst=sgpr("LSUTailLoopOffset"), src=vgpr(wave_id), comment="Update Alpha")) module.add(SSubI32( dst=sgpr("LSUTailLoopOffset"), \ src0=loopCounter, \ src1=sgpr("LSUTailLoopOffset"), \ comment="lsu offset" )) module.add(SMinU32(dst=loopCounter, src0=sgpr(dividend), src1=loopCounter, comment="" )) self.vgprPool.checkIn(wave_id) self.vgprPool.checkIn(tmpVgpr) if not tailloopInNll: # skip tailLoopNumIter in tailloopInNll (already generated above) skComponent = Component.StreamK.find(self) module.add(skComponent.tailLoopNumIter(self, kernel, loopCounter)) gsuComponent = Component.GSU.find(self) module.add(gsuComponent.tailLoopNumIter(self, kernel, loopCounter)) # if tail numIter == 0 skip altogether skipTailLoopLabel = Label.getFormatting("SkipTailLoop%s"%(loopChar) ) module.add(SCmpEQU32(src0=loopCounter, src1=0, comment="numIter%s == 0"%loopChar )) module.add(SMovB32(dst=sgpr("OrigLoopCounter"), src=0, comment="repurpose to count each localRead increment")) module.add(SCBranchSCC1(labelName=skipTailLoopLabel, \ comment="skip to end of tail loop b/c numIter==0")) ######################################## # Unrolled Loop elif loopIdx == self.states.unrollIdx: loopCounterName = self.loopCounterName(kernel, loopIdx) loopCounter = sgpr(loopCounterName) if not self.do["PreLoop"]: module.add(ValueEndif()) with self.allocTmpSgpr(3, tag="calculateLoopNumIter_tmpSgprInfo") as tmpSgprInfo: skComponent = Component.StreamK.find(self) module.add(skComponent.calculateLoopNumIter(self, kernel, loopCounterName, loopIdx, tmpSgprInfo)) gsuComponent = Component.GSU.find(self) module.add(gsuComponent.calculateLoopNumIter(self, kernel, loopCounterName, tmpSgprInfo)) module.add(SMovB32(dst=sgpr("OrigLoopCounter"), \ src=loopCounter, \ comment="copy loop counter")) else: # other summation, not unroll loop #printExit("no assembly support for 2+ dimensional summation") module.addComment1("%sother summation, numIter%s = size%s" \ % (self.indent, loopChar, loopChar)) loopCounter = self.loopCounter(kernel, loopIdx) module.add(SMovB32(dst=loopCounter, \ src=sgpr("SizesSum+%u"%loopIdx), \ comment="init loop counter")) return module ############################################################################## # Calculate Metadata offset ############################################################################## def calculateIncrementMetadata(self, kernel, sgprOut): module = Module("calculateIncrementMetadata") gsuComponent = Component.GSU.find(self) module.add(gsuComponent.calculateIncrementMetadata(self, kernel, sgprOut)) return module def getOneStr(self, mxType, bits): value = "" if bits == 32: if mxType.toName() == "E8": value = "0x7f7f7f7f" elif mxType.toName() == "Float8": value = "0x38383838" elif mxType.toName() == "E5M3": value = "0x78787878" else: raise Exception(f"unsupport tc {mxType.toName()}") elif bits == 64: if mxType.toName() == "E8": value = "0x7f7f7f7f7f7f7f7f" elif mxType.toName() == "Float8": value = "0x3838383838383838" elif mxType.toName() == "E5M3": value = "0x7878787878787878" else: raise Exception(f"unsupport tc {mxType.toName()}") else: raise Exception(f"unsupport bits {bits}") return value ############################################################################## # Unrolled-loop entry threshold ############################################################################## def unrollLoopEntryEndCounter(self, kernel): """Threshold T for the unrolled loop: the main (global-load) unrolled loop body -- and therefore the InitCIterWmma cloned iter0 that zeroes C -- runs iff LoopCounter > T.""" pgr = kernel["PrefetchGlobalRead"] if pgr == 1: return 0 if kernel["SuppressNoLoadLoop"] else 1 if pgr == 2: return (0 if kernel["HalfPLR"] else 1) if kernel["SuppressNoLoadLoop"] else 2 if pgr >= 3: # openLoop early-exits to NoGlobalLoadLoop for LoopCounter <= PGR. return pgr return 0 ############################################################################## # Open Loop ############################################################################## def openLoop(self, kernel, tPA, tPB, loopIdx, noLabelGen=False, beginLabelOnly=False, nta=0, ntb=0): strNta = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTA%s"%nta strNtb = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTB%s"%ntb module = Module("openLoop") if bool(kernel["ProblemType"]["MXBlockA"]) ^ bool(kernel["ProblemType"]["MXBlockB"]): block = max(kernel["ProblemType"]["MXBlockA"], kernel["ProblemType"]["MXBlockB"]) mxType = kernel["ProblemType"]["DataTypeMXSA"] if kernel["ProblemType"]["MXBlockA"] else kernel["ProblemType"]["DataTypeMXSB"] if block == 32: module.add(VMovB32(dst=vgpr("ValuMXSDummy"), src=self.getOneStr(mxType, 32))) elif block == 16: module.add(VMovB64(dst=vgpr("ValuMXSDummy",2), src=self.getOneStr(mxType, 64))) else: raise Exception(f"unsupport mx block {block}") # TODO - rewrite this function to simplify control-flow between tail-loop / unroll loop tailLoop = loopIdx < 0 if tailLoop: loopIdx = self.states.unrollIdx self.states.inTailLoop = True loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][loopIdx]] if not tailLoop and not noLabelGen and not beginLabelOnly: module.add(Label("openLoop%s%s%s"%(loopChar, strNta, strNtb), "")) bStrNta = "" if tailLoop else strNta bStrNtb = "" if tailLoop else strNtb loopLabelBegin = Label("%sLoopBegin%s%s%s"%("Tail" if tailLoop else "", loopChar, bStrNta, bStrNtb), "", alignment=16 ) loopLabelEnd = Label("%sLoopEnd%s%s%s"%("Tail" if tailLoop else "", loopChar, bStrNta, bStrNtb), "" ) if beginLabelOnly: # generate only beginLabel, then, return module.add(loopLabelBegin) return module # is numIter at least 1? otherwise skip to end # PGL needs a skip-check here if not bufferload # If kernel["SuppressNoLoadLoop"] we don't have a special loop for the 'last iter' loopCounter = self.loopCounter(kernel, loopIdx) if tailLoop: # begin loop if not noLabelGen: module.add(loopLabelBegin) else: # not tailloop: if loopIdx == self.states.unrollIdx: # 1 loop check is necessary only when AssertSummationElementMultiple % (DepthU * 2) != 0 if kernel["PrefetchGlobalRead"] >= 2 and kernel["AssertSummationElementMultiple"] % (kernel["DepthU"] * 2) != 0 and not kernel["SuppressNoLoadLoop"]: module.add(SCmpEQU32( src0=loopCounter, \ src1=hex(1), \ comment="LoopCounter%s == 1 (PGR>=2, not Suppress: single-loop -> toPGR1)"%(loopChar) )) toPGR1 = Label.getFormatting(self.labels.getName("toPGR1")) module.add(SCBranchSCC1(labelName=toPGR1, comment="PGR=2 but only 1 loop, toPGR1")) if kernel["PrefetchGlobalRead"] >= 3: # early exit 1 (2<=loopCounter<=PGR-1) to second NGLL (no need GR Inc) endCounter = kernel["PrefetchGlobalRead"]-1 module.add(SCmpLeU32( src0=loopCounter, \ src1=hex(endCounter), \ comment="LoopCounter%s <= EndCounter"%(loopChar) )) remainPgr = endCounter - 1 jumpLabel = Label("NoGlobalLoadLoop_%d"%remainPgr, "") module.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), \ comment="do not enter Loop%s"%loopChar )) # early exit 2 (loopCounter==PGR) to first NGLL (need GR Inc) endCounter = self.unrollLoopEntryEndCounter(kernel) # == PGR module.add(SCmpLeU32( src0=loopCounter, \ src1=hex(endCounter), \ comment="LoopCounter%s <= EndCounter"%(loopChar) )) remainPgr = endCounter - 1 jumpLabel = Label("NoGlobalLoadLoop_%d"%remainPgr, "") module.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), \ comment="do not enter Loop%s"%loopChar )) else: endCounter = self.unrollLoopEntryEndCounter(kernel) module.add(SCmpLeU32( src0=loopCounter, \ src1=hex(endCounter), \ comment="LoopCounter%s < EndCounter"%(loopChar) )) jumpLabel = loopLabelEnd if kernel["PrefetchGlobalRead"]==2 and (not kernel["SuppressNoLoadLoop"]) and kernel["ExpandPointerSwap"]: # PGR=2 and EPS and no SuppressNoLoadLoop case, need to jump to EvenExit jumpLabel = Label("LoopEnd%s_evenexit%s%s"%(loopChar, strNta, strNtb), "" ) module.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), \ comment="do not enter Loop%s"%loopChar )) if kernel["ExpertSchedulingMode"] > 0: enableESMInstr = SSetRegIMM32B32(dst=HWRegContainer(reg="26", value=[0,2]), src=int(kernel["ExpertSchedulingMode"]), comment="enable expert scheduling mode") if kernel["ESMRuntimeGate"]: skipEnableESMLabel = Label("skipEnableESM", "") module.add(VCmpGEI32(dst=VCC(), src0=vgpr(self.states.esmRuntimeFlagVgpr), src1=1, comment="check if ESM supported at runtime")) module.add(SCBranchVCCZ(labelName=skipEnableESMLabel.getLabelName(), comment="skip s_setreg if not supported")) module.add(enableESMInstr) module.add(skipEnableESMLabel) self.vgprPool.checkIn(self.states.esmRuntimeFlagVgpr) else: module.add(enableESMInstr) if not noLabelGen: module.add(loopLabelBegin) if loopIdx != self.states.unrollIdx: # reset LRO since these may have changed due to odd-iter exit ? if kernel["PrefetchGlobalRead"]: module.addComment0("openLoop - reset LRO for possible odd-iter exit") module.add(self.localReadResetOffsets(kernel, tPA)) if kernel["ProblemType"]["MXBlockA"]: module.add(self.localReadResetOffsets(kernel, tPA["MX"])) if kernel["ProblemType"]["MXBlockB"]: module.add(self.localReadResetOffsets(kernel, tPB["MX"])) module.add(self.localReadResetOffsets(kernel, tPB)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] module.add(self.localReadResetOffsets(kernel, tPM)) return module ############################################################################## # Close Loop # finalLoop : final unroll loop ############################################################################## def closeLoop(self, kernel, tPA, tPB, loopIdx, finalLoop, emitEndLabelOnly=False, oddLabel=False, skipCondJumpCounter=-1, NLLindexLast=False, nta=0, ntb=0): strNta = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTA%s"%nta strNtb = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTB%s"%ntb module = Module("closeLoop") if emitEndLabelOnly: loopIdx = self.states.unrollIdx loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][loopIdx]] module.add(Label("SkipTailLoop%s"%(loopChar), "")) return module tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] finalJump = SCBranchSCC0 jumpNeeded = True tailLoop = loopIdx < 0 tailloopInNll = loopIdx == -2 # use -2 for tailloopInNll needTailEndCode = tailLoop and finalLoop and ((not tailloopInNll) or kernel["NoTailLoop"]) if tailLoop: loopIdx = self.states.unrollIdx loopChar = self.states.indexChars[kernel["ProblemType"]["IndicesSummation"][loopIdx]] loopLabelBegin = Label("TailLoopBegin%s"%(loopChar), "" ) loopLabelEnd = Label("TailLoopEnd%s"%(loopChar), "" ) loopCounter = self.loopCounter(kernel, loopIdx) numReadsIterCoalescedA = self.states.numReadsIterCoalescedA numReadsIterCoalescedB = self.states.numReadsIterCoalescedB numReadsIterCoalesced = max(numReadsIterCoalescedA, numReadsIterCoalescedB) unrollInc = 1 KinInnerUnroll = kernel["InnerUnroll"] # dot2: force InnerUnroll = 1 in tail loop if kernel["UseDotInstruction"]: unrollInc *= kernel["NumDotElements"] * kernel["NumWaveSplitK"] else: if kernel["EnableMatrixInstruction"]: unrollInc *= kernel["MatrixInstK"] * numReadsIterCoalesced KinInnerUnroll *= kernel["MatrixInstK"] if kernel["AssertSummationElementMultiple"] % KinInnerUnroll == 0: unrollInc *= kernel["InnerUnroll"] skipCondJump = False if skipCondJumpCounter >= 0 and (skipCondJumpCounter%numReadsIterCoalesced < numReadsIterCoalesced - 1): # skip conditional jump when numReadsIterCoalesced > 1 and skipCondJumpCounter is not the last in numReadsIterCoalesced # to support numReadsIterCoalesced > 1, MatrixInstK * numReadsIterCoalesced needs to be executed # e.g.) MatrixInstK=4, numReadsIterCoalesced=2 # skipCondJumpCounter==0 case: execute K=0,2,4,6 (here, no exit(means skip cond jump) to execute odd K(1,3,5,7)) # skipCondJumpCounter==1 case: execute K=1,3,5,7 (here, all K=0-7 are done. check condition and jump if tail loop is done) # skipCondJump=True is not to exit after skipCondJumpCounter==0. skipCondJump = True jumpNeeded = not skipCondJump if not skipCondJump: module.addComment1("closeLoop loop%s finalLoop=%d tailLoop=%d" % (loopChar, finalLoop, tailLoop)) module.add(SSubI32( dst=loopCounter, \ src0=loopCounter, \ src1=hex(unrollInc), \ comment="dec counter%s (tailLoop)"%(loopChar) )) # Track # LDS reads? module.add(SAddU32( dst=sgpr("OrigLoopCounter"), \ src0=sgpr("OrigLoopCounter"), \ src1=hex(unrollInc), comment="inc counter%s"%(loopChar) )) endCounter = 0 if kernel["LocalSplitU"] > 1: module.add(SSubI32( dst=sgpr("LSUTailLoopOffset"), \ src0=sgpr("LSUTailLoopOffset"), \ src1=hex(unrollInc), \ comment="LSU offset dec counter%s (tailLoop)"%(loopChar) )) tmp = self.sgprPool.checkOut(1, tag="closeLoop_tmpSgpr") module.add(SMinI32(dst=sgpr(tmp), src0=loopCounter, src1=sgpr("LSUTailLoopOffset"), comment="check lsu offset too")) module.add(SCmpLeI32( src0=sgpr(tmp), \ src1=hex(endCounter), \ comment="counter%s<=%d"%(loopChar,endCounter) )) self.sgprPool.checkIn(tmp) else: module.add(SCmpLeI32( src0=loopCounter, \ src1=hex(endCounter), \ comment="counter%s<=%d"%(loopChar,endCounter) )) else: # not tailloop loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][loopIdx]] loopLabelBegin = Label("LoopBegin%s%s%s"%(loopChar, strNta, strNtb), "" ) loopLabelEnd = Label("LoopEnd%s%s%s"%(loopChar, strNta, strNtb), "" ) loopLabelEndOddExit = Label("LoopEnd%s_oddexit%s%s"%(loopChar, strNta, strNtb), "unroll loop odditer exit" ) loopLabelEndEvenExit = Label("LoopEnd%s_evenexit%s%s"%(loopChar, strNta, strNtb), "unroll loop eveniter exit" ) loopCounter = self.loopCounter(kernel, loopIdx) module.addComment1("closeLoop loop%s finalLoop=%d tailLoop=%d" % (loopChar, finalLoop, tailLoop)) if kernel["enableTDMA"] and kernel["enableTDMB"] and not kernel["PrefetchGlobalRead"]: module.add(SWaitCnt(dscnt=0, comment="TDM PGR=0: wait all ds_reads before TDM overwrite")) module.add(SBarrier(comment="TDM PGR=0: signal+wait done reading LDS")) # If PrefetchGlobalRead=1 the loads in the loop prefetch next macro-tile # For the final trip through the unroll loop we need to ensure those loads stay in bounds. # One technique is to create a copy of the unroll loop with all loads removed. # However buffer load doesn't need this loop copy since we OOB loads can be suppressed by buffer limit hardware # So can do one more iteration (endCounter==0) in the main unroll loop, and adjust the pointer # increments appropriately. # Also sum idx other than unroll always compare against 0 (there is no PGR to account for) if kernel["PrefetchGlobalRead"] >= 1 and not kernel["SuppressNoLoadLoop"] and loopIdx == self.states.unrollIdx: endCounter = kernel["PrefetchGlobalRead"] else: endCounter = 0 if kernel["AssertSummationElementMultiple"] % (kernel["DepthU"] * 2) == 0 and endCounter > 0 and \ kernel["PrefetchGlobalRead"] <= 2: # if AssertSummationElementMultiple is multiple of DepthU*2, loop exit is necessary only once in 2 Loop iterations # In endCounter % 2 == 1 case, exit at lc % 2 == 0 (= oddLabel). It means no exit if not oddLabel # In endCounter % 2 == 0 case, exit at lc % 2 == 1 (= not oddLabel). It means no exit if oddLabel # No exit case, no code is necessary except for final Loop # decrement by 2 if PGR=2 and StaggerU is 0, else 1 if kernel["PrefetchGlobalRead"]==2: with self.allocTmpSgpr(2, tag="closeLoop_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx module.add(SCmpEQU32(src0=sgpr("StaggerU"), src1=0)) module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=hex(2), src1=hex(1))) decCode = SSubU32(dst=loopCounter, src0=loopCounter, \ src1=sgpr(tmpSgpr), \ comment="dec counter%s"%(loopChar) ) else: decCode = SSubU32(dst=loopCounter, src0=loopCounter, \ src1=1, \ comment="dec counter%s"%(loopChar) ) condCode = SCmpEQI32(src0=loopCounter, \ src1=hex(endCounter), \ comment="counter%s==%d"%(loopChar,endCounter) ) noExit = False if endCounter%2 != 0: if not oddLabel: noExit = True else: if oddLabel: noExit = True if noExit: # No exit. No dec code if decValue is 2 if decValue == 2: decCode = "" condCode = "" if finalLoop: # No exit and finalLoop case, use s_branch (no condition) finalJump = SBranch else: jumpNeeded =False if decCode: module.add(decCode) if condCode: module.add(condCode) else: module.add(SSubU32( dst=loopCounter, src0=loopCounter, \ src1=1, \ comment="dec counter%s"%(loopChar) )) module.add(SCmpEQI32( src0=loopCounter, \ src1=hex(endCounter), \ comment="counter%s==%d"%(loopChar,endCounter) )) jumpLabel = loopLabelEnd if not tailLoop and not kernel["SuppressNoLoadLoop"] and kernel["ExpandPointerSwap"]: # in this case, odd or/and even code is generated and use odd/even exit to avoid skipping odd/even code # (end label is generated after odd/even code) jumpLabel = loopLabelEndOddExit if oddLabel else loopLabelEndEvenExit if not finalLoop: if jumpNeeded: # just an exit check, else fall through to the next loop copy module.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), comment="exit Loop%s"%loopChar )) else: #finalLoop: if jumpNeeded: module.add(finalJump(labelName=loopLabelBegin.getLabelName(), comment="restart Loop%s"%(loopChar))) if not tailLoop and loopIdx == self.states.unrollIdx: oddIterPreCode = Module() oddIterCode = Module() evenIterPreCode = Module() evenIterCode = Module() if not kernel["SuppressNoLoadLoop"] and kernel["ExpandPointerSwap"]: oddIterPreCode.add(loopLabelEndOddExit) # In this case we kept the 'no-load' loop which has LDS offsets assuming first bank of LDS # if we exit the main loop at an odd iter - need to swap LDS read pointers # so the ds_reads read from the 'high' buffer of LDS oddIterPreCode.addComment1("Select high bank of LDS") # Generate local read address code only if DirectToVgpr is not enabled if not kernel["DirectToVgprA"] and not kernel["StoreSwapAddr"]: if kernel["ExpertSchedulingMode"] > 0: oddIterCode.add(SWaitAlu(vm_vsrc=0, comment="wait for local read to vgpr complete")) oddIterCode.add(self.localReadSwapOffsets(kernel, False, tPA)) if kernel["ProblemType"]["MXBlockA"]: oddIterCode.add(self.localReadSwapOffsets(kernel, False, tPA["MX"])) # Generate local read address code only if DirectToVgpr is not enabled if not kernel["DirectToVgprB"] and not kernel["StoreSwapAddr"]: if kernel["ExpertSchedulingMode"] > 0: oddIterCode.add(SWaitAlu(vm_vsrc=0, comment="wait for local read to vgpr complete")) oddIterCode.add(self.localReadSwapOffsets(kernel, False, tPB)) if kernel["ProblemType"]["MXBlockB"]: oddIterCode.add(self.localReadSwapOffsets(kernel, False, tPB["MX"])) if kernel["ProblemType"]["Sparse"]: if kernel["DirectToVgprSparseMetadata"]: oddIterCode.add(SWaitCnt(vlcnt=0, comment="wait for global read before moving metadata to target vgpr")) for i in range(0, self.states.m.numVgprValuPerBlock): oddIterCode.add(VMovB32(vgpr("ValuMetadata+%u"%i), vgpr("ValuMetadata+%u+%u"%(self.states.m.numVgprValuPerBlock, i)), \ comment="copy ValuMetadata blk1 to blk0")) elif not kernel["StoreSwapAddr"]: oddIterCode.add(self.localReadSwapOffsets(kernel, False, tPM)) evenIterPreCode.add(loopLabelEndEvenExit) # generate even code here (so far, for PrefetchGlobalRead>=2 only) if kernel["PrefetchGlobalRead"]>=2: if not kernel["DirectToVgprA"] or not kernel["DirectToVgprB"]: if kernel["ExpertSchedulingMode"] > 0: evenIterCode.add(SWaitAlu(vm_vsrc=0, comment="wait for local read to vgpr complete")) # Generate local write address code only for PrefetchGlobalRead>=2 #TODO: TDM if not kernel["DirectToLdsA"]: evenIterCode.add(self.localWriteSwapOffsets(kernel, False, tPA)) if ("MX" in tPA) and (not kernel["DirectToLdsMXSA"]): evenIterCode.add(self.localWriteSwapOffsets(kernel, False, tPA["MX"])) if ("MX" in tPB) and (not kernel["DirectToLdsMXSB"]): evenIterCode.add(self.localWriteSwapOffsets(kernel, False, tPB["MX"])) if not kernel["DirectToLdsB"]: evenIterCode.add(self.localWriteSwapOffsets(kernel, False, tPB)) evenIterCode.add(self.localWriteSwapOffsets(kernel, True, tPA)) if "MX" in tPA: evenIterCode.add(self.localWriteSwapOffsets(kernel, True, tPA["MX"])) if "MX" in tPB: evenIterCode.add(self.localWriteSwapOffsets(kernel, True, tPB["MX"])) evenIterCode.add(self.localWriteSwapOffsets(kernel, True, tPB)) #swap local write memory token self.states.ldsWriteTokenIdx = \ self.states.memTokenLdsBuffer1 if self.states.ldsWriteTokenIdx == self.states.memTokenLdsBuffer0 else self.states.memTokenLdsBuffer0 # generate even, odd exit code # not oddLabel case, order is even -> odd firstPreCode = evenIterPreCode firstCode = evenIterCode secondPreCode = oddIterPreCode secondCode = oddIterCode if oddLabel: # oddLabel case, swap the order (odd -> even) firstPreCode, secondPreCode = secondPreCode, firstPreCode firstCode, secondCode = secondCode, firstCode module.add(firstPreCode) module.add(firstCode) # if secondCode exist, add jump to skip secondCode if secondCode.count(): module.add(SBranch(labelName=loopLabelEnd.getLabelName(), \ comment="exit unroll loop%s (and skip second exit code)"%(loopChar))) module.add(secondPreCode) module.add(secondCode) if (not tailLoop) or needTailEndCode: module.add(loopLabelEnd) elif tailloopInNll: module.add(SBranch(labelName=loopLabelEnd.getLabelName(), \ comment="End of tailloopInNll. Jump to the end of TailLoop")) if NLLindexLast: # end of tailloopInNll label (NLLindex==NLLnum-1 only) EndOfTailLoopInNLLLabel = Label("TailLoopInNLLEnd%s"%(loopChar), "" ) module.add(EndOfTailLoopInNLLLabel) if needTailEndCode: if len(kernel["ProblemType"]["IndicesSummation"]) > 1 or kernel["StreamK"]: # recover the 'damage' done to LRO: # if LRA is backed-up before (wlr case), we simply restore the addr (sub inc*loop doesn't work) tPList = [] if self.oriLraA != None: module.add(VMovB32(dst=vgpr("LocalReadAddrA"), src=vgpr(self.oriLraA), comment="restore LRA")) self.vgprPool.checkIn(self.oriLraA) self.oriLraA = None else: tPList.append(tPA) if self.oriLraB != None: module.add(VMovB32(dst=vgpr("LocalReadAddrB"), src=vgpr(self.oriLraB), comment="restore LRA")) self.vgprPool.checkIn(self.oriLraB) self.oriLraB = None else: tPList.append(tPB) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: if self.oriLraM != None: module.add(VMovB32(dst=vgpr("LocalReadAddrMetadata"), src=vgpr(self.oriLraM), comment="restore LRA")) self.vgprPool.checkIn(self.oriLraM) self.oriLraM= None else: tPList.append(tPM) for tP in tPList: tc = tP["tensorChar"] LdsPad = int(kernel["LdsPad%s" % tc] * tP["bpeDS"]) if kernel["LdsBlockSizePerPad%s"%tc] == 0 else 0 inc = kernel["LocalSplitU"] * (kernel["MacroTile%s"%tc] + LdsPad) * tP["bpeDS"] # aligned with localReadInc if kernel["EnableMatrixInstruction"]: if kernel["UnrollMajorLDS%s" % tP["tensorChar"]]: inc = kernel["LocalSplitU"] * tP["bpeDS"] # No need to *= K, because LoopCounter is increased by K each time # inc *= kernel["MatrixInstK"] if not (tP["isA"] and kernel["DirectToVgprA"] or tP["isB"] and kernel["DirectToVgprB"]): # no local read code if DirectToVgpr is enabled with self.allocTmpSgpr(1, tag="closeLoop_tmpSgprInfo2") as tmpSgprInfo: stmp = tmpSgprInfo.idx module.add(SMovB32(dst=sgpr(stmp), src=int(inc), comment="tailloop lds offset")) module.add(SMulI32(dst=sgpr(stmp), src0=sgpr("OrigLoopCounter"), src1=sgpr(stmp), comment="scale by mul")) module.add(VSubU32(dst=vgpr("LocalReadAddr%s"%tc), src0=vgpr("LocalReadAddr%s"%tc), src1=sgpr(stmp), comment="remove lro damage")) # if LWA is backed-up before, we simply restore the addr if self.oriLwaA != None: module.add(VMovB32(dst=vgpr("LocalWriteAddrA"), src=vgpr(self.oriLwaA), comment="restore LWA")) module.add(VMovB32(dst=vgpr("LocalWriteAddrB"), src=vgpr(self.oriLwaB), comment="restore LWA")) self.vgprPool.checkIn(self.oriLwaA) self.vgprPool.checkIn(self.oriLwaB) self.oriLwaA = None self.oriLwaB = None if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: module.add(VMovB32(dst=vgpr("LocalWriteAddrMetadata"), src=vgpr(self.oriLwaM), comment="restore LWA")) self.vgprPool.checkIn(self.oriLwaM) self.oriLwaM = None # restore all threads if needTailEndCode and kernel["LocalSplitU"] > 1: sgprCnt = self.states.laneSGPRCount waveSize = kernel["WavefrontSize"] module.addComment1("restore full exec mask") with self.allocTmpSgpr(sgprCnt, tag="closeLoop_tmpSgprInfo3") as tmpSgprInfo: fullExec = tmpSgprInfo.idx activeMask = "0xFFFFFFFF" if (waveSize == 32) else "0xFFFFFFFFFFFFFFFF" SMovBX = SMovB32 if (waveSize == 32) else SMovB64 SOrSaveExecBX = SOrSaveExecB32 if (waveSize == 32) else SOrSaveExecB64 module.add(SMovBX(dst=sgpr(fullExec,sgprCnt), src=activeMask, comment="restore all threads active")) module.add(SOrSaveExecBX (dst=sgpr(fullExec,sgprCnt), src=sgpr(fullExec,sgprCnt), comment="full mask -> exec" )) return module ############################################################################## # End Summation ############################################################################## def endSummation(self, kernel, tPA, tPB, noSkipLoad = True, label = None, isOptNLL = False): module = Module("endSummation") module.add(Label((self.labels.getUniqueNamePrefix("Summation_End") if label is None else label), "")) if kernel["StorePriorityOpt"]: module.add(SSetPrior(prior=0, comment="optimization store")) if isOptNLL: vbegin = self.states.startVgprMisc vsize = self.states.lastVgprForReads - vbegin if self.states.a.startVgprLocalReadAddr > self.states.lastVgprForReads: vsize = self.states.a.startVgprLocalReadAddr + self.states.a.numVgprLocalReadAddr - vbegin if self.states.b.startVgprLocalReadAddr > self.states.lastVgprForReads: vsize = self.states.b.startVgprLocalReadAddr + self.states.b.numVgprLocalReadAddr - vbegin if self.states.m.startVgprLocalReadAddr > self.states.lastVgprForReads: vsize = self.states.m.startVgprLocalReadAddr + self.states.m.numVgprLocalReadAddr - vbegin self.vgprPool.add(vbegin, vsize, "endSummation") module.addComment0("endSummation: add vgpr [%u...%u) to pool" % \ (vbegin, vbegin+vsize)) # Write bias A, B data to LDS if kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["UseBias"] and (kernel["ProblemType"]["BiasSrc"] == "A" or kernel["ProblemType"]["BiasSrc"] == "B"): tP = tPA if kernel["ProblemType"]["BiasSrc"] == "A" else tPB module.add(self.exclasses.biasSumUnroll.storeSumLDS(self, kernel, tP)) vbegin = self.states.bias.startVgprValu vsize = self.states.bias.numVgprValu self.vgprPool.add(vbegin, vsize, "free vgpr of biasSumUnroll") module.addComment0("endSummation: add vgpr [%u...%u) to pool" % \ (vbegin, vbegin+vsize)) keptSgprs = [] # FP32 to FP8 SR without v_prng_b32 needs RNDSeed sgpr preserved if kernel["ProblemType"]["DestDataType"].is8bitFloat() and kernel["ProblemType"]["StochasticRounding"] and not self.states.asmCaps["v_prng_b32"]: keptSgprs.append("RNDSeed") lastRegTag=None spool = self.sgprPool.getPool() for i in range(0, self.sgprPool.size()): regTag = spool[i].tag if regTag != lastRegTag: lastRegTag = regTag if lastRegTag in keptSgprs: continue if (lastRegTag not in self.states.nonPostLoopSgpr) and (spool[i].status == RegisterPool.Status.InUse): if label == "Summation_End_OptNLL": self.undefineSgpr(regTag) else: module.add(self.undefineSgpr(regTag)) if self.db["InitVgpr"] & 0x2: module.add(self.vgprPool.initTmps(self.consts.initVgprValue,start=0, stop=100)) if 0: # FIXME: Can remove? for i in range(0,16+1): module.add(VMovB32(dst=vgpr(21), src=vgpr(21), comment="hack tmp in pool")) # this doesn't seem to do anything if self.db["InitSgpr"] & 0x2: module.add(self.sgprPool.initTmps(self.consts.initSgprValue)) if self.db["ConservativeWaitCnt"] & 0x10: module.add(SBarrier(comment="debug")) module.add(SWaitCnt(dscnt=0, vlcnt=0, vscnt=0, comment="")) if kernel["SuppressNoLoadLoop"]: module.add(SWaitCnt(dscnt=0, vlcnt=0, vscnt=0, comment="wait for all summation activity")) # Load bias data from LDS and write to global ######################################## # Load kernel args needed by global write batch module.addComment0("load store sgprs") # Define sgprs for kernel args runActivation = True if ((kernel["ProblemType"]["ActivationType"] != 'none') \ and kernel["ActivationFused"]) else False def fixPreloadOffset(offset, sgpxIdxVec, numStoreSgprToLoad): item = None startVgprName = sgpxIdxVec[0] if kernel["ProblemType"]["UseScaleAB"] == "Scalar": if (kernel["ProblemType"]["DataTypeA"].numRegisters() > kernel["ProblemType"]["MacDataTypeA"].numRegisters()) and (kernel["ProblemType"]["DataTypeB"].numRegisters() > kernel["ProblemType"]["MacDataTypeB"].numRegisters()): self.argLoader.setOffset(offset + ((self.states.rpga * self.states.bpr) * 2)) elif kernel["ProblemType"]["DataTypeA"].numRegisters() > kernel["ProblemType"]["MacDataTypeA"].numRegisters(): assert sgpxIdxVec[0] == self.sgprs["AddressScaleB"] self.argLoader.setOffset(offset + (self.states.rpga * self.states.bpr)) elif kernel["ProblemType"]["DataTypeB"].numRegisters() > kernel["ProblemType"]["MacDataTypeB"].numRegisters(): assert sgpxIdxVec[0] == self.sgprs["AddressScaleA"] item = self.argLoader.loadKernArg(self.sgprs["AddressScaleA"], "KernArgAddress", dword=2) if len(sgpxIdxVec) > 1: startVgprName = sgpxIdxVec[1] numStoreSgprToLoad -= self.states.rpga self.argLoader.setOffset(offset + ((self.states.rpga * self.states.bpr) * 2)) return (item, startVgprName, numStoreSgprToLoad) if self.states.numStoreSgprToLoad: sgpxIdxVec = self.defineMultiSgprIndex(self.states.numStoreSgprNames, self.states.numStoreSgprNameSizes, align=4) for name in self.states.numStoreSgprNames: module.add(RegSet("s", "sgpr"+name, self.sgprs[name])) if noSkipLoad and kernel["GlobalSplitU"] != 0: gsuLabel = Label(label=self.labels.getNameInc("GSU"), comment="") with self.allocTmpSgpr(1, tag="endSummation_tmpSgprGSU") as tmpSgprGSU: module.add(SAndB32(dst=sgpr(tmpSgprGSU.idx), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SCmpEQU32(src0=sgpr(tmpSgprGSU.idx), src1=1, comment="GSU == 1 ?")) if (kernel["_GlobalAccumulation"] != 'MultipleBufferSingleKernel'): module.add(SCBranchSCC0(labelName=gsuLabel.getLabelName(), comment="branch if GSU != 1")) if kernel["ProblemType"]["SupportUserArgs"]: extReadEpilogueLabel = Label(label=self.labels.getNameInc("LoadExternalEpilogueStruct"), comment="") extReadEpilogueLabelEnd = Label(label=self.labels.getNameInc("LoadExternalEpilogueStructEnd"), comment="") module.addComment0("Check if custom structure pointer is null") module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=2, comment="ArgType == 2 ?")) module.add(SCBranchSCC1(labelName=extReadEpilogueLabel.getLabelName(), comment="branch if ArgType == 2")) argOffset = self.argLoader.getOffset() # Backup offset numStoreSgprToLoad = self.states.numStoreSgprToLoad (item, startVgprName, numStoreSgprToLoad) = fixPreloadOffset(argOffset, sgpxIdxVec, numStoreSgprToLoad) if item: module.add(item) loadModule = module.addModuleAsFlatItems(self.argLoader.loadAllKernArg(startVgprName, "KernArgAddress", numStoreSgprToLoad)) self.states.numStoreSgprInst = countSMemLoad(loadModule) self.argLoader.setOffset(argOffset) # Restore offset module.add(SBranch(extReadEpilogueLabelEnd.getLabelName())) module.add(extReadEpilogueLabel) extArgOffset = self.externalArgLoader.getOffset() backupExtArgOffset = extArgOffset loadList = [[-1, 0, extArgOffset]] extArgOffset += self.states.userArgsInfo.scaleASize if (kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (not self.states.preloadScaleA)) or kernel["ProblemType"]["UseScaleAB"] == "Vector": if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["AddressScaleA"] loadList[-1][1] += self.states.userArgsInfo.scaleASize else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.scaleBSize if (kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (not self.states.preloadScaleB)) or kernel["ProblemType"]["UseScaleAB"] == "Vector": if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["AddressScaleB"] loadList[-1][1] += self.states.userArgsInfo.scaleBSize else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.scaleCSize + self.states.userArgsInfo.scaleDSize if kernel["ProblemType"]["UseScaleCD"]: if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["AddressScaleC"] loadList[-1][1] += self.states.userArgsInfo.scaleCSize + self.states.userArgsInfo.scaleDSize else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.scaleAlphaVecSize if kernel["ProblemType"]["UseScaleAlphaVec"]: if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["AddressScaleAlphaVec"] loadList[-1][1] += self.states.userArgsInfo.scaleAlphaVecSize else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.biasSize if self.states.numSgprAddressBias: biasLoadSize = (self.states.numSgprAddressBias + self.states.BiasType + self.states.BiasStride) * 4 if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["AddressBias"] loadList[-1][1] += biasLoadSize if biasLoadSize < self.states.userArgsInfo.biasSize: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.factorDimSize if self.states.FactorDim == 3: if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["FactorDim"] loadList[-1][1] += self.states.userArgsInfo.factorDimSize else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.eSize if kernel["ProblemType"]["UseE"]: if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs["AddressE"] loadList[-1][1] += self.states.userArgsInfo.eSize else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. extArgOffset += self.states.userArgsInfo.activationSize if runActivation: needActTypeArg = 1 if self.states.numActivationTypeArgSize else 0 actNames = kernel["ProblemType"]["ActivationType"].getAdditionalArgStringList() actLoadSize = (len(actNames) * self.states.numActivationArgSize + needActTypeArg) * 4 if (actLoadSize == self.states.userArgsInfo.activationSize) or len(actNames) > 0: if loadList[-1][0] == -1: loadList[-1][0] = self.sgprs[actNames[0]] loadList[-1][1] += actLoadSize else: loadList.append(["ActivationType", actLoadSize]) # Need to start a new loadAllKernArg cause no AdditionalArgStringList is needed loadList.append([-1, 0, extArgOffset - (needActTypeArg * 4)]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. else: loadList.append([-1, 0, extArgOffset]) # Need to start a new loadAllKernArg cause the argument is not consecutively anymore. # Start reading arguments loadModuleExt = Module("Count Inst") for loadInfo in loadList: if loadInfo[0] == -1: continue dwordLen = loadInfo[1] // 4 self.externalArgLoader.setOffset(loadInfo[2]) loadModuleExt.addModuleAsFlatItems(module.addModuleAsFlatItems(self.externalArgLoader.loadAllKernArg(loadInfo[0], "KernArgAddress", dwordLen))) self.states.numStoreSgprInstExt = countSMemLoad(loadModuleExt) self.externalArgLoader.setOffset(backupExtArgOffset) module.add(extReadEpilogueLabelEnd) else: argOffset = self.argLoader.getOffset() # Backup offset numStoreSgprToLoad = self.states.numStoreSgprToLoad (item, startVgprName, numStoreSgprToLoad) = fixPreloadOffset(argOffset, sgpxIdxVec, numStoreSgprToLoad) if item: module.add(item) loadModule = module.addModuleAsFlatItems(self.argLoader.loadAllKernArg(startVgprName, "KernArgAddress", numStoreSgprToLoad)) self.states.numStoreSgprInst = countSMemLoad(loadModule) self.argLoader.setOffset(argOffset) # Restore offset if noSkipLoad and kernel["GlobalSplitU"] != 0: module.add(gsuLabel) ######################################## # Load kernel args needed by global write batch # Define sgprs for kernel args if self.states.numStoreSgprToLoad2: module.addComment0("load store sgprs2") for name in self.states.numStoreSgprNames2: module.add(RegSet("s", "sgpr"+name, self.sgprs[name])) argOffset = self.argLoader.getOffset() # Backup offset if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (kernel["ProblemType"]["DataTypeA"].numRegisters() > kernel["ProblemType"]["MacDataTypeA"].numRegisters() \ or kernel["ProblemType"]["DataTypeB"].numRegisters() > kernel["ProblemType"]["MacDataTypeB"].numRegisters()): self.argLoader.setOffset(argOffset + (self.states.numStoreSgprToLoad)*4 + (self.states.rpga * self.states.bpr)) # Restore offset else: self.argLoader.setOffset(argOffset + (self.states.numStoreSgprToLoad)*4) # Restore offset numStoreSgprToLoad = self.states.numStoreSgprToLoad2 if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (kernel["ProblemType"]["DataTypeA"].numRegisters() > kernel["ProblemType"]["MacDataTypeA"].numRegisters() \ or kernel["ProblemType"]["DataTypeB"].numRegisters() > kernel["ProblemType"]["MacDataTypeB"].numRegisters()): argOffsettmp = (argOffset + (self.states.numStoreSgprToLoad)*4 + (self.states.rpga * self.states.bpr)) # Restore offset else: argOffsettmp = (argOffset + (self.states.numStoreSgprToLoad)*4) # Restore offset extReadEpilogueLabeltmp = Label(label=self.labels.getNameInc("LoadExternalEpilogueStruct"), comment="") module.addComment0("Check if custom structure pointer is null") if kernel["ProblemType"]["SupportUserArgs"]: module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=2, comment="ArgType == 2 ?")) module.add(SCBranchSCC1(labelName=extReadEpilogueLabeltmp.getLabelName())) module.add(self.argLoader.loadKernArg("AddressTD", "KernArgAddress", sgprOffset=hex(argOffsettmp), dword=2)) module.add(self.argLoader.loadKernArg("Synchronizer", "KernArgAddress", sgprOffset=hex(argOffsettmp + (2)*4), dword=2)) module.add(extReadEpilogueLabeltmp) self.argLoader.setOffset(argOffset) # Restore offset # define the rest sgprs if (not self.states.doShadowInit) and kernel["BufferStore"]: self.defineSgpr("SrdD", 4, 4) self.defineSgpr("SrdC", 4, 4) module.add(RegSet("s", "sgprSrdC", self.sgprs["SrdC"])) module.add(RegSet("s", "sgprSrdD", self.sgprs["SrdD"])) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): self.defineSgpr("SrdScaleA", 4, 4) self.defineSgpr("SrdScaleB", 4, 4) module.add(RegSet("s", "sgprSrdScaleA", self.sgprs["SrdScaleA"])) module.add(RegSet("s", "sgprSrdScaleB", self.sgprs["SrdScaleB"])) if kernel["ProblemType"]["UseScaleAlphaVec"]: self.defineSgpr("SrdScaleAlphaVec", 4, 4) module.add(RegSet("s", "sgprSrdScaleAlphaVec", self.sgprs["SrdScaleAlphaVec"])) if self.states.useBias != DataDirection.NONE: self.defineSgpr("SrdBias", 4, 4) module.add(RegSet("s", "sgprSrdBias", self.sgprs["SrdBias"])) if ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1)): self.defineSgpr("SrdTD", 4, 4) module.add(RegSet("s", "sgprSrdTD", self.sgprs["SrdTD"])) self.defineSgpr("GSUSync", 1) module.add(RegSet("s", "sgprGSUSync", self.sgprs["GSUSync"])) if self.states.useGateResidual: self.defineSgpr("SrdGate", 4, 4) module.add(RegSet("s", "sgprSrdGate", self.sgprs["SrdGate"])) self.defineSgpr("GateNullOne", 1) module.add(RegSet("s", "sgprGateNullOne", self.sgprs["GateNullOne"])) if kernel["ProblemType"]["UseE"]: self.defineSgpr("SrdE", 4, 4) module.add(RegSet("s", "sgprSrdE", self.sgprs["SrdE"])) for idx in range(0, kernel["ProblemType"]["NumIndicesC"]): i = idx idxChar= self.states.indexChars[idx] if i == 0 and not kernel["ProblemType"]["UseInitialStridesCD"]: module.add(ValueSet("constStrideE%s"%idxChar, 1)) else: if not kernel["ProblemType"]["UseInitialStridesCD"]: i = i-1 module.add(RegSet("s", "sgprStrideE%s"%idxChar, \ "sgprStridesE", i)) if kernel["ProblemType"]["UseScaleCD"]: assert kernel["ProblemType"]["ComputeDataType"].isSingle() self.defineSgpr("ScaleD", 2, 2) module.add(RegSet("s", "sgprScaleD", self.sgprs["ScaleD"])) if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1: self.defineSgpr("SrdSync", 4, 4) module.add(RegSet("s", "sgprSrdSync", self.sgprs["SrdSync"])) # Load kernel args end ######################################## # copy accumulated C from agpr to vgpr if kernel["EnableMatrixInstruction"]: #TODO avoid s_nop if its possible #instCycles = kernel["MatrixInstM"] // 2 # 32x32 is 64 cycles, 16x16 is 32 cycles, 4x4 is 8 cycles #module.add(SNop(waitState=instCycles)) module.addComment1("Mapping of Acc register -> C Vgpr register") # For subtile kernels with mixed agpr/vgpr accumulators the spilled # D-tile values live in arch vgprs allocated from the pool (not at # ValuC+N). Determine their base vgpr so mapAcctoArchRegs can address # them correctly. spilledVgprBase = None if kernel.get("UseSubtileImpl"): # For subtile kernels, D-tile accumulators that overflow the accvgpr # pool are placed in arch vgprs allocated from the vgpr pool. # mapAcctoArchRegs needs to know the base address of those vgprs so it # can emit correct moves instead of referencing "ValuC+N" (which points # to the wrong location in the subtile allocation scheme). for vtile in self.states.d.tileInfo.vgprTiles: if vtile.regList.is_vgpr: spilledVgprBase = vtile.regList.indices[0] break self.codes.accVgprRead = mapAcctoArchRegs(kernel, self.states.maxLimitAgprs, write=False, spilledVgprBase=spilledVgprBase) if (kernel["StreamK"] > 0 and kernel["StreamKAtomic"] == 0) or \ ((kernel["GlobalSplitU"] == -1 or kernel["GlobalSplitU"] > 0) and (kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1)): self.codes.accVgprWrite = mapAcctoArchRegs(kernel, self.states.maxLimitAgprs, write=True, spilledVgprBase=spilledVgprBase) # same spilledVgprBase if kernel["MIArchVgpr"]: module.addComment1("Multiply MI out register with Alpha -> C Vgpr register") self.codes.mulAlphaMultipleBuffer = moveMIoutToArch(kernel, self.states.startVgprAlphaTmp) self.codes.mulAlphaOther = mulMIoutAlphaToArch(kernel, self.states.startVgprAlphaTmp) return module def mfmaIter_waitCount(self, kernel): if self.states.version in [(9,4,2), (9,5,0)]: dataType = kernel["ProblemType"]["DataType"] miM = kernel["MatrixInstM"] miN = kernel["MatrixInstN"] if dataType.isSingle() or dataType.isHalf() or dataType.isBFloat16(): if miM == 4 and miN == 4: return 2 return 0 ############################################################################## # src A,B str for MFMA ############################################################################## def generateSrcStrForMFMAshiftK(self, kernel, tP, innerUnroll, vregSetIdx, vgprPerInput, m, u, iui, idxAB, lc=0, bk=None): return self.generateSrcStrForMFMA(kernel, tP, innerUnroll, vregSetIdx, vgprPerInput, m, u, iui, idxAB, lc=lc, bk=bk, shiftK=True) def generateSrcStrForMFMA(self, kernel, tP, innerUnroll, vregSetIdx, vgprPerInput, m, u, iui, idxAB, lc=0, bk=None, shiftK=False): tc = tP["tensorChar"] if tc == "A": statesTc = self.states.a numIterPerCoalescedRead = self.states.numIterPerCoalescedReadA numReadsIterCoalesced = self.states.numReadsIterCoalescedA elif tc == "MXSA": statesTc = self.states.mxsa numIterPerCoalescedRead = self.states.numIterPerCoalescedReadMXSA numReadsIterCoalesced = self.states.numReadsIterCoalescedMXSA elif tc == "MXSB": statesTc = self.states.mxsb numIterPerCoalescedRead = self.states.numIterPerCoalescedReadMXSB numReadsIterCoalesced = self.states.numReadsIterCoalescedMXSB elif tc == "B": statesTc = self.states.b numIterPerCoalescedRead = self.states.numIterPerCoalescedReadB numReadsIterCoalesced = self.states.numReadsIterCoalescedB else: raise Exception(f"unsupport tc %s{tc}") numVgprValuPerBlock = int(kernel["MIWaveTile%s"%tc] * ceil(kernel["MIInputPerThread%s"%tc] * tP["bpe"] / self.states.bpr)) # calculate vgprBufferA_new ( or B) and offset for DirectToVgpr. Use u instead of m (number of local prefetch buffer does not matter) m_or_u = u if kernel["DirectToVgpr%s"%tc] else m vgprBuffer_new = (m_or_u//numIterPerCoalescedRead)*numIterPerCoalescedRead vgprBuffer_new_offset = m_or_u%numIterPerCoalescedRead*innerUnroll*vgprPerInput # DirectToVgpr + pack special case # offset vgprBuffer_new packDTV = self.states.packDTVA if (tc in ("A", "MXSA")) else self.states.packDTVB convDTV = self.states.convDTVA if (tc in ("A", "MXSA")) else self.states.convDTVB if packDTV or convDTV: # DTV + pack case, offset bufferIdx for local read packing instructions numBi = kernel["LoopIters"] vgprBuffer_new += vregSetIdx * numBi iui_new = (iui//numReadsIterCoalesced)*numReadsIterCoalesced iui_new_offset = iui%numReadsIterCoalesced*vgprPerInput ab_new = idxAB*vgprPerInput*numReadsIterCoalesced abStr = "Valu%s_X%u_I%u+%u+%u+%u" % (tc, vgprBuffer_new, iui_new, ab_new, vgprBuffer_new_offset, iui_new_offset) if statesTc.useDirect32XEmulationThis or statesTc.useDirect32XEmulationNext: lrvwTile = self.states.lrvwTileA if tP["isA"] else self.states.lrvwTileB component = Component.LocalRead.find(self) # tranpose case, we do ShiftK -> cnv -> MFMA # Vreg can be different between ShiftK and MFMA # ShiftK: src of cvt # MFMA: dst of cvt (set dst=True) # X32Emu case, need to add all index to get correct tranpose index idx = ab_new + vgprBuffer_new_offset + iui_new_offset if bk != None: idx += bk bk = None if shiftK: abStr = component.getVgprStrForEmuMfma(self, kernel, tc, vgprBuffer_new, iui_new, idx, lrvwTile, u) else: abStr = component.getVgprStrForEmuMfma(self, kernel, tc, vgprBuffer_new, iui_new, idx, lrvwTile, u, dst=True) # need to add dummy for after string conversion abStr += "+0+0" if kernel["DirectToVgpr%s"%tc] and not (packDTV or convDTV): # overwrite aStr/bStr for DirectToVgpr (except for pack DTV case) numVgprPerBlock = statesTc.numVgprG2LAllocated numVgprPerBlock //= 2 # DTV case, buffer is doubled. //2 to calculate single size ab_new += vregSetIdx * numVgprPerBlock + ( vgprBuffer_new * innerUnroll) * numVgprValuPerBlock abStr = "G2L%s+%u+%u" % (tc, ab_new, vgprBuffer_new_offset) if bk != None: if kernel["UseDirect32XEmulation"] and (int(bk) % 8) < 4: # T0 set is allocated with half the VGPRs (ri // 2). Map bk to the # correct T0 offset: groups of 8 elements share 4 T0 slots each. adjustedBk = (int(bk) // 8) * 4 + (int(bk) % 4) abStr += "+%u"%(adjustedBk) else: abStr += "+%u"%(bk) # halfPLR halfPLR = (tP["isA"] or tP["isB"]) and kernel["HalfPLR%c"%tc] if halfPLR: halfPLRGroups = self.getHalfPLRGroups(kernel, lc, u) bkVal = 0 if bk == None else bk abStr = self.getHalfPLRValuStr(halfPLRGroups, ab_new + vgprBuffer_new_offset + iui_new_offset + bkVal, tc) return abStr ############################################################################## # MXS TileSpan scale-select # # With the TileSpan optimization, N MXS scale ds_loads collapse to N/2 loads: within # each group of 2*VW scale blocks, the lower half-wave of the single ds_load holds the # first VW blocks and the upper half-wave holds the partner VW blocks (LRA produces this # layout for both MIWaveGroup==1 and >1). The consuming WMMA reads each block directly # from that one loaded register via the gfx1250 matrix_{a,b}_scale:N select. This maps a # logical scale block index # to the register that actually holds it plus the scale-select (0 = lower half-wave, # 1 = partner half-wave). ############################################################################## def mxsUsesScaleSel(self, kernel): # The scale-select path is only valid where localReadMX (gfx1250 WMMA_V3 + # InMemorySwizzle) produces the half-wave partner scale layout. return (self.states.asmCaps.get("HasWMMA_V3", False) and kernel.get("MXScaleFormat") == "InMemorySwizzle") def mxsTileSpanScaleSel(self, kernel, tP, idxAB): tc = tP["tensorChar"] if "MXS" not in tc or not self.mxsUsesScaleSel(kernel): return idxAB, 0 component = Component.LocalRead.find(self) info = component.getMxsTileSpanInfo(kernel, tc, tP["tile01Idx"], self.states.asmCaps) if info is None: return idxAB, 0 vectorWidth = info["vectorWidth"] groupSize = 2 * vectorWidth group = idxAB // groupSize within = idxAB % groupSize scaleSel = within // vectorWidth # 0 = lower half-wave, 1 = partner half-wave regInHalf = within % vectorWidth # Only the lower half-wave of each group is loaded (LocalRead packs the loaded groups # contiguously and drops the partner-half vgprs), so map the logical block to the # compacted register: group g owns vectorWidth registers, not 2*vectorWidth. The partner # block (scaleSel==1) shares the same loaded register and is selected by matrix_*_scale. mappedIdx = group * vectorWidth + regInHalf return mappedIdx, scaleSel ############################################################################## # MAC Iteration ############################################################################## def macIter(self, kernel, tPA, tPB, bufferIdx, iuiCount, useMacro, isTail=False): imod = Module("macIter_X%u_I%u"%(bufferIdx, iuiCount)) # if kernel["ProblemType"]["MacDataTypeA"].isHalf(): # # imod.addText(".align32 8, 0xbf800001", "align v_pk_fma") # Align v_pk_fma instructions used in MAC_ blocks # imod.addText(".align32 8, 0xbf800001\n") # Align v_pk_fma instructions used in MAC_ blocks # dot2: add shiftK module to prevent read out of K bound shiftK = Module("shiftK") inputType = kernel["ProblemType"]["DataType"] numRegistersIn = inputType.numRegisters() loopCounterName = self.loopCounterName(kernel, self.states.unrollIdx) numInput = kernel["NumDotElements"] if kernel["UseDotInstruction"] else 1 kPerIter = numInput * kernel["NumWaveSplitK"] vgprPerInput = 1 shiftPerElement = int(numRegistersIn * 32) s_nop = 0 # alloc vgpr kReg = None abReg = None dummy = -1 if isTail and (kernel["AssertSummationElementMultiple"] % kPerIter != 0): kReg = self.vgprPool.checkOut(1, tag="macIter_kReg") # remainder loopCntSgpr = loopCounterName with self.allocTmpSgpr(1, tag="macIter_tmpSgprInfo") as tmpSgprInfo: shiftK.add(vectorStaticRemainder(dummy, kReg, "Serial", kernel["NumWaveSplitK"], None, tmpSgprInfo)) numTmpSgpr = 4 if (vgprPerInput > 2) else 3 with self.allocTmpSgpr(numTmpSgpr, alignment=1, tag="macIter_tmpSgprInfo2") as tmpSgprInfo: if tmpSgprInfo.idx % 2 == 0: tmpSgprX2 = tmpSgprInfo.idx tmpSgprX1 = tmpSgprInfo.idx+2 else: tmpSgprX2 = tmpSgprInfo.idx+1 tmpSgprX1 = tmpSgprInfo.idx # replace 0 for differnet thread shiftK.add(vectorStaticMultiply(vgpr(kReg), vgpr(kReg), numInput, tmpSgprInfo)) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) for a in range(0, kernel["ThreadTileA"]): for iui in range(0, iuiCount): aStr = vgpr("ValuA_X%u_I%u+%u" % (bufferIdx, iui, a*vgprPerInput), 1) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=0, src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="set 0 if K_idx >= sizeL")) for b in range(0, kernel["ThreadTileB"]): for iui in range(0, iuiCount): bStr = vgpr("ValuB_X%u_I%u+%u" % (bufferIdx, iui, b*vgprPerInput), 1) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=0, src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="set 0 if K_idx >= sizeL")) # replace 0 for some thread if kernel["AssertSummationElementMultiple"] % numInput != 0: abReg = self.vgprPool.checkOutAligned(vgprPerInput, vgprPerInput, tag="macIter_abReg") shiftK.add(VSubU32(dst=vgpr(kReg), src0=sgpr(loopCntSgpr), src1=vgpr(kReg), comment="get distance between size and k index")) shiftK.add(VCmpLtI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=vgpr(kReg), src1=numInput, comment="set partial 0 if distance less than numInput")) shiftK.add(SAndB32(dst=sgpr(tmpSgprX1), src0=sgpr(loopCntSgpr), src1=numInput-1, comment="get inputs for edge thread")) shiftK.add(SSubU32(dst=sgpr(tmpSgprX1), src0=numInput, src1=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) shiftK.add(SLShiftLeftB32(dst=sgpr(tmpSgprX1), shiftHex=log2(shiftPerElement), src=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) # currently vgprPerInput must be 1 VShiftLeft = VLShiftLeftB32 for a in range(0, kernel["ThreadTileA"]): for iui in range(0, iuiCount): aStr = vgpr("ValuA_X%u_I%u+%u" % (bufferIdx, iui, a*vgprPerInput), vgprPerInput) shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInput), shiftHex=sgpr(tmpSgprX1), src=aStr, comment="")) for bk in range(0, vgprPerInput): aStr = vgpr("ValuA_X%u_I%u+%u+%u" % (bufferIdx, iui, a*vgprPerInput, bk), 1) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=vgpr(abReg+bk), src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) for b in range(0, kernel["ThreadTileB"]): for iui in range(0, iuiCount): bStr = vgpr("ValuB_X%u_I%u+%u" % (bufferIdx, iui, b*vgprPerInput), vgprPerInput) shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInput), shiftHex=sgpr(tmpSgprX1), src=bStr, comment="")) for bk in range(0, vgprPerInput): bStr = vgpr("ValuB_X%u_I%u+%u+%u" % (bufferIdx, iui, b*vgprPerInput, bk), 1) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=vgpr(abReg+bk), src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) if kernel["LocalSplitU"] > 1: self.sgprPool.checkIn(loopCntSgpr) s_nop = 2 if s_nop != 0: imod.add(SNop(waitState=(s_nop - 1), comment="")) tmpIU = kernel["InnerUnroll"] if kernel["InnerUnroll"] > 1 and iuiCount==1: # This it tail-loop case where we just want one IUI, instr = "MAC_%ux%u_X%u_OneIUI" % (kernel["ThreadTile0"],kernel["ThreadTile1"], bufferIdx) tmpIU = 1 else: if not useMacro: printExit("MAC doesn't support useMacro=False") instr = "MAC_%ux%u_X%u" % (kernel["ThreadTile0"],kernel["ThreadTile1"], bufferIdx) # release register if kReg is not None: self.vgprPool.checkIn(kReg) if abReg is not None: self.vgprPool.checkIn(abReg) if self.do["MAC"]: imod.add(shiftK) if self.states.asmCaps["HasVgprMSB"]: component = Component.MAC.find(self) if not component: printExit("Assembly doesn't support datatype %s" % kernel["ProblemType"]["DataType"]) innerModule = component(self, tPA, tPB, bufferIdx, tmpIU) imod.add(innerModule) else: imod.add(MacroInstruction(name=instr, args=[])) imod.addSpaceLine() return imod ############################################################################## # ACC Vgpr R/W Function ############################################################################## def accVgprReadWriteFunction(self, kernel, idx, read=True): if not kernel["MIArchVgpr"]: if idx >= self.states.maxLimitAgprs: return VMovB32 else: return VAccvgprReadB32 if read else VAccvgprWriteB32 else: return VMovB32 def accVgprReadWriteIndex(self, kernel, idx, sz=1): if not kernel["MIArchVgpr"]: if idx >= self.states.maxLimitAgprs: return vgpr(idx - self.states.maxLimitAgprs, sz) else: return accvgpr(idx, sz) else: return vgpr("ValuC+%u"%idx, sz) ############################################################################## # MFMA Iteration ############################################################################## def mfmaIter(self, kernel, tPA, tPB, u, innerUnroll, vregSetIdx, unrollLoopIdx = 0, unrollIdx = 0, tail = False, firstIter = False, postShiftK = Module()): imod = Module("mi") shiftK = Module("shiftK") m = (u) % (self.states.numVgprBuffer) # local to use for MACs def dataTypeToMfmaInstTypePair(dataTypeA: DataType, dataTypeB: DataType, sourceSwap: bool) -> Tuple[InstType, InstType]: miInTypeStrA = dataTypeA.toNameAbbrev() miInTypeStrB = dataTypeB.toNameAbbrev() miInTypeStr = miInTypeStrA + "_" + miInTypeStrB if miInTypeStrA != miInTypeStrB else miInTypeStrA miInInstType = dataTypeNameAbbrevToInstType(miInTypeStr, sourceSwap) # v_mfma_[...xK] miOutInstType = dataTypeNameAbbrevToInstType(dataTypeA.MIOutputTypeNameAbbrev()) # v_mfma_.. return miInInstType, miOutInstType def dataTypeNameAbbrevToInstType(abbrev: str, sourceSwap: bool = False) -> InstType: if abbrev == 'f64': return InstType.INST_F64 elif abbrev == 'f32': return InstType.INST_F32 elif abbrev == 'f16': return InstType.INST_F16 elif abbrev == 'i32': return InstType.INST_I32 elif abbrev == 'i8': return InstType.INST_I8 elif abbrev == 'bf16': return InstType.INST_BF16 elif abbrev == 'xf32': return InstType.INST_XF32 elif abbrev == 'fp8': return InstType.INST_F8 elif abbrev == 'bf8': return InstType.INST_BF8 elif (abbrev == 'fp8_bf8' and sourceSwap == False) or \ (abbrev == 'bf8_fp8' and sourceSwap == True): return InstType.INST_F8_BF8 elif (abbrev == 'bf8_fp8' and sourceSwap == False) or \ (abbrev == 'fp8_bf8' and sourceSwap == True): return InstType.INST_BF8_F8 elif abbrev == 'fp6': return InstType.INST_F6 elif abbrev == 'bf6': return InstType.INST_BF6 elif (abbrev == 'fp6_bf6' and sourceSwap == False) or \ (abbrev == 'bf6_fp6' and sourceSwap == True): return InstType.INST_F6_B6 elif (abbrev == 'bf6_fp6' and sourceSwap == False) or \ (abbrev == 'fp6_bf6' and sourceSwap == True): return InstType.INST_B6_F6 elif abbrev == 'fp4': return InstType.INST_F4 elif (abbrev == 'fp8_fp4' and sourceSwap == False) or \ (abbrev == 'fp4_fp8' and sourceSwap == True): return InstType.INST_F8_F4 elif (abbrev == 'fp4_fp8' and sourceSwap == False) or \ (abbrev == 'fp8_fp4' and sourceSwap == True): return InstType.INST_F4_F8 elif (abbrev == 'fp6_fp4' and sourceSwap == False) or \ (abbrev == 'fp4_fp6' and sourceSwap == True): return InstType.INST_F6_F4 elif (abbrev == 'fp4_fp6' and sourceSwap == False) or \ (abbrev == 'fp6_fp4' and sourceSwap == True): return InstType.INST_F4_F6 elif (abbrev == 'fp8_fp6' and sourceSwap == False) or \ (abbrev == 'fp6_fp8' and sourceSwap == True): return InstType.INST_F8_F6 elif (abbrev == 'fp6_fp8' and sourceSwap == False) or \ (abbrev == 'fp8_fp6' and sourceSwap == True): return InstType.INST_F6_F8 elif (abbrev == 'fp8_bf6' and sourceSwap == False) or \ (abbrev == 'bf6_fp8' and sourceSwap == True): return InstType.INST_F8_B6 elif (abbrev == 'bf6_fp8' and sourceSwap == False) or \ (abbrev == 'fp8_bf6' and sourceSwap == True): return InstType.INST_B6_F8 elif (abbrev == 'bf8_fp4' and sourceSwap == False) or \ (abbrev == 'fp4_bf8' and sourceSwap == True): return InstType.INST_B8_F4 elif (abbrev == 'fp4_bf8' and sourceSwap == False) or \ (abbrev == 'bf8_fp4' and sourceSwap == True): return InstType.INST_F4_B8 elif (abbrev == 'bf6_fp4' and sourceSwap == False) or \ (abbrev == 'fp4_bf6' and sourceSwap == True): return InstType.INST_B6_F4 elif (abbrev == 'fp4_bf6' and sourceSwap == False) or \ (abbrev == 'bf6_fp4' and sourceSwap == True): return InstType.INST_F4_B6 elif (abbrev == 'bf8_fp6' and sourceSwap == False) or \ (abbrev == 'fp6_bf8' and sourceSwap == True): return InstType.INST_B8_F6 elif (abbrev == 'fp6_bf8' and sourceSwap == False) or \ (abbrev == 'bf8_fp6' and sourceSwap == True): return InstType.INST_F6_B8 elif (abbrev == 'bf8_bf6' and sourceSwap == False) or \ (abbrev == 'bf6_bf8' and sourceSwap == True): return InstType.INST_B8_B6 elif (abbrev == 'bf6_bf8' and sourceSwap == False) or \ (abbrev == 'bf8_bf6' and sourceSwap == True): return InstType.INST_B6_B8 elif abbrev == 'e8': return InstType.INST_E8 elif abbrev == 'e5m3': return InstType.INST_E5M3 else: assert("Unsupported data type.") return InstType.INST_NOTYPE isgfx950 = kernel["ISA"][:2] == (9, 5) miInputTypeA = kernel["ProblemType"]["F32XdlMathOp"] if kernel["EnableF32XdlMathOp"] else kernel["ProblemType"]["MacDataTypeA"] miInputTypeB = kernel["ProblemType"]["F32XdlMathOp"] if kernel["EnableF32XdlMathOp"] else kernel["ProblemType"]["MacDataTypeB"] # calculate constant is_mfma = self.states.asmCaps["HasMFMA"] is_wmma_v1 = self.states.asmCaps["HasWMMA_V1"] is_wmma_v2 = self.states.asmCaps["HasWMMA_V2"] is_wmma_v3 = self.states.asmCaps["HasWMMA_V3"] numRegistersInA = miInputTypeA.numRegisters() numRegistersInB = miInputTypeB.numRegisters() # TODO: Consider correctness of tail-loop numRegistersIn = max(numRegistersInA, numRegistersInB) numRegistersInMXSA = kernel["ProblemType"]["DataTypeMXSA"].numRegisters() if kernel["ProblemType"]["MXBlockA"] else 0 numRegistersInMXSB = kernel["ProblemType"]["DataTypeMXSB"].numRegisters() if kernel["ProblemType"]["MXBlockB"] else 0 numRegistersOut = kernel["MIRegPerOut"] loopCounterName = self.loopCounterName(kernel, self.states.unrollIdx) accs_per_wave = kernel["MatrixInstM"] * kernel["MatrixInstN"] * kernel["MatrixInstB"] \ // self.states.kernel["WavefrontSize"] * numRegistersOut dividerFortidInK = kernel["MatrixInstN"] * kernel["MatrixInstB"] matrixInstT = min(kernel["MatrixInstM"], kernel["MatrixInstN"]) numTileInInstA = kernel["MatrixInstM"] // matrixInstT numTileInInstB = kernel["MatrixInstN"] // matrixInstT numMIInputA = kernel["MIInputPerThreadA"] numMIInputB = kernel["MIInputPerThreadB"] numMIInputMXSA = kernel["MIInputPerThreadMXSA"] if kernel["ProblemType"]["MXBlockA"] else 0 numMIInputMXSB = kernel["MIInputPerThreadMXSB"] if kernel["ProblemType"]["MXBlockB"] else 0 numMIInputM = kernel["MIInputPerThreadMetadata"] if kernel["ProblemType"]["Sparse"] else 0 numMIInput = max(numMIInputA, numMIInputB) numMIInUnroll = max(numMIInputA//numTileInInstA, numMIInputB//numTileInInstB) miInInstType, miOutInstType = dataTypeToMfmaInstTypePair(miInputTypeA, miInputTypeB, kernel["SourceSwap"]) neg_flag = True if ((not is_mfma) and (miInInstType == InstType.INST_I8)) else False miInInstType = InstType.INST_U8 if ((not is_mfma) and miInInstType == InstType.INST_I8) else miInInstType miOutInstType = miOutInstType if (is_mfma or kernel["ProblemType"]["Sparse"]) else dataTypeNameAbbrevToInstType(kernel["ProblemType"]["ComputeDataType"].toNameAbbrev()) miInScaleAInstType = dataTypeNameAbbrevToInstType(kernel["ProblemType"]["DataTypeMXSA"].toNameAbbrev()) miInScaleBInstType = dataTypeNameAbbrevToInstType(kernel["ProblemType"]["DataTypeMXSB"].toNameAbbrev()) numReadsIterCoalescedA = self.states.numReadsIterCoalescedA numReadsIterCoalescedB = self.states.numReadsIterCoalescedB numReadsIterCoalesced = max(numReadsIterCoalescedA, numReadsIterCoalescedB) tPM = (tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"]) if kernel["ProblemType"]["Sparse"] else None vgprPerInputA = int(numMIInputA * numRegistersInA) vgprPerInputB = int(numMIInputB * numRegistersInB) vgprPerInputMXSA = ceil(numMIInputMXSA * numRegistersInMXSA) vgprPerInputMXSB = ceil(numMIInputMXSB * numRegistersInMXSB) vgprPerInputM = int(ceil(numMIInputM / self.states.bpr)) if kernel["ProblemType"]["Sparse"] else 0 vgprPerInput = max(vgprPerInputA, vgprPerInputB) vgprPerInUnrollA = vgprPerInputA // numTileInInstA vgprPerInUnrollB = vgprPerInputB // numTileInInstB vgprPerInUnroll = max(vgprPerInUnrollA, vgprPerInUnrollB) shiftPerElement = int(numRegistersIn * 32) s_nop = 0 accumRegType = "acc" if not kernel["MIArchVgpr"] else "v" mfma_1k = True if kernel["MFMA_BF16_1K"] else False accStoreCIdx = 0 # alloc vgpr kReg_first = None tmpVgpr2 = None kReg = None abReg = None mxReg = None tmpVgpr = None dummy = -1 if (numRegistersIn < 1) and ((kernel["UnrollMajorLDSA"] == False) or (kernel["UnrollMajorLDSB"] == False)): s_nop = 2 if kernel["ConvertAfterDS"] and (numRegistersIn < 1) and ((tPA["bpe"] > tPA["bpeDS"]) or (tPB["bpe"] > tPB["bpeDS"])): s_nop = 2 # here we remap index to where it read for wider local read # ex. if we read 2 iteration at a time, # original : _ds_load_b64 valuA_X0_I0 # read 2 iter: _ds_load_b128 valuA_X0_I0 (we read valuA_X0_I0 and valuA_X1_I0) # instead of using valuA_X1_I0, we use valuA_X0_I0+2 as mfma input if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: vgprBufferM_new = (m//self.states.numIterPerCoalescedReadMetadata)*self.states.numIterPerCoalescedReadMetadata vgprBufferM_new_offset = m%self.states.numIterPerCoalescedReadMetadata*kernel["InnerUnroll"]*vgprPerInputM # handle multiple K element in MFMA instruction # MIK=1 case, we still need this code for Coalesced case if tail and (kernel["MatrixInstK"] > 1 or numReadsIterCoalescedA > 1 or numReadsIterCoalescedB > 1) and \ (not kernel["enableTDMA"] and not kernel["enableTDMB"]): if not is_wmma_v1: #mfma or wmma_v2 # ToDo: Avoid using extra kReg_first for wmma_v2 case kReg_first = self.vgprPool.checkOut(1,"kReg_first") # the first vgpr of remainder kReg = self.vgprPool.checkOut(1,"kReg") # remainder if kernel["LocalSplitU"] > 1: loopCntSgpr = self.sgprPool.checkOut(1, tag="mfmaIter_loopCntSgpr") else: loopCntSgpr = loopCounterName with self.allocTmpSgpr(1, tag="mfmaIter_tmpSgprInfo") as tmpSgprInfo: shiftK.add(vectorStaticRemainder(dummy, kReg_first, "Serial", self.states.kernel["WavefrontSize"], tmpVgpr, tmpSgprInfo)) shiftK.add(vectorStaticDivide(kReg_first, kReg_first, dividerFortidInK, tmpVgpr)) numTmpSgpr = 4 if (vgprPerInUnroll > 2 and is_wmma_v2) else 3 # ToDo: check if this is correct for sparse mfma with self.allocTmpSgpr(numTmpSgpr, alignment=1, tag="mfmaIter_tmpSgprInfo2") as tmpSgprInfo: if tmpSgprInfo.idx % 2 == 0: tmpSgprX2 = tmpSgprInfo.idx tmpSgprX1 = tmpSgprInfo.idx+2 if numTmpSgpr == 4: tmpSgprX3 = tmpSgprInfo.idx+3 else: tmpSgprX2 = tmpSgprInfo.idx+1 tmpSgprX1 = tmpSgprInfo.idx if numTmpSgpr == 4: tmpSgprX3 = tmpSgprInfo.idx+3 def findSparseOffset(isA:bool): blocksPerTGroupSMFMA = 1 elementsPerBlockSMFMA = 1 blockOffsetSMFMA =1 if kernel["ProblemType"]["Sparse"] != 0: if kernel["MIInputPerThread"] * kernel["ProblemType"]["DataTypeB"].numBytes() > 16: # double K isSparseTrack = (kernel["ProblemType"]["Sparse"] == 1 and isA) or (kernel["ProblemType"]["Sparse"] == 2 and not isA) # gfx950 sparse track only has one block for each thread group. # TODO adjust this value for other arch. blocksPerTGroupSMFMA = 1 if isSparseTrack else 2 if self.states.asmCaps["HasSWMMAC_gfx1250"]: blocksPerTGroupSMFMA = 2 if blocksPerTGroupSMFMA > 1: threadGroups = kernel["MatrixInstK"] // kernel["MIInputPerThread"] elementsPerBlockSMFMA = kernel["MIInputPerThread%s" % ("A" if isA else "B")] // blocksPerTGroupSMFMA blockStride = elementsPerBlockSMFMA * threadGroups blockOffsetSMFMA = (blockStride - elementsPerBlockSMFMA) * (2 if isSparseTrack else 1) return blocksPerTGroupSMFMA, elementsPerBlockSMFMA, blockOffsetSMFMA blocksPerTGroupSMFMAA, elementsPerBlockSMFMAA, blockOffsetSMFMAA = findSparseOffset(True) blocksPerTGroupSMFMAB, elementsPerBlockSMFMAB, blockOffsetSMFMAB = findSparseOffset(False) # replace 0 for differnet thread if kernel["ProblemType"]["Sparse"] == 1 and numMIInUnroll//8 >= 1: vgprPerSet0Group = 1 elif vgprPerInUnrollA <= 2: shiftK.add(vectorStaticMultiply(vgpr(kReg_first), vgpr(kReg_first), numMIInUnroll * numReadsIterCoalesced, tmpSgprInfo)) kStepForCoalesced = (u%numReadsIterCoalesced) * numMIInUnroll if kStepForCoalesced > 0: shiftK.add(VAddU32(vgpr(kReg_first), hex(kStepForCoalesced), vgpr(kReg_first), "k += (u%%numReadsIterCoalesced) * numMIInUnroll")) if kernel["LocalSplitU"] > 1: shiftK.add(SMinI32(dst=sgpr(loopCntSgpr), src0=sgpr(loopCounterName), src1=sgpr("LSUTailLoopOffset"), comment="check lsu bound")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=vgpr(kReg_first), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) #ToDo: check value of laneSGPRCount vgprPerSet0Group = vgprPerInUnrollA elif is_wmma_v2 and vgprPerInUnrollA > 2: vgprPerSet0Group = 4 else: if is_wmma_v3 and kernel["UseF32XEmulation"]: vgprPerSet0Group = 1 elif kernel["ProblemType"]["DataType"].is6bitFloat(): vgprPerSet0Group = 3 else: vgprPerSet0Group = 2 if not tPA["isSwizzled"] and tPA["bpe"] == 0.75 and kernel["enableLDSTrA"]: for a in range(0, kernel["MIWaveTileA"]): for ti in range(0, numTileInInstA): for iui in range(0, innerUnroll): aStr = self.generateSrcStrForMFMA(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, unrollLoopIdx) shiftK.add(_shiftLrElements(aStr, vgprPerInUnrollA, a)) numSet0GroupA = vgprPerInUnrollA // vgprPerSet0Group for group in range(0, numSet0GroupA): if numSet0GroupA > 1 or (is_wmma_v2 and vgprPerInUnrollA > 2): if group == 0: multiplyBy = numMIInUnroll if kernel["ProblemType"]["Sparse"] == 1 and self.states.asmCaps["HasSWMMAC_gfx1250"]: multiplyBy = numMIInUnroll // blocksPerTGroupSMFMAA elif kernel["ProblemType"]["Sparse"] and not self.states.asmCaps["HasSWMMAC_gfx1250"]: multiplyBy = numMIInUnroll // blocksPerTGroupSMFMAA elif is_wmma_v3 and kernel["UseF32XEmulation"]: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPA["bpe"] * 8) multiplyBy = vgprLayout[-1] // vgprLayout[2] elif is_wmma_v3: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPA["bpe"] * 8) multiplyBy = vgprLayout[-1] elif vgprPerInUnrollA == 8 and not is_wmma_v2: multiplyBy = numMIInUnroll // 2 shiftK.add(vectorStaticMultiply(vgpr(kReg_first), vgpr(kReg_first), multiplyBy, tmpSgprInfo)) shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), 0, "")) elif is_wmma_v3 and kernel["UseF32XEmulation"]: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPA["bpe"] * 8) miVectorWidth = vgprLayout[-1] # Absolute K offset for interleaved layout: # groups within each half are contiguous, halves are miVectorWidth apart. elemPerHalf = miVectorWidth // vgprLayout[2] kOffsetA = (group // elemPerHalf) * miVectorWidth + (group % elemPerHalf) shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), kOffsetA, "K index for group %d" % group)) elif blocksPerTGroupSMFMAA == 2 and (group * vgprPerSet0Group) == (elementsPerBlockSMFMAA * numRegistersIn): if kernel["ProblemType"]["Sparse"]: kIncA = blockOffsetSMFMAA + (numMIInUnroll//numSet0GroupA) * (max(group - 1, 0) if not self.states.asmCaps["HasSWMMAC_gfx1250"] else 1) else: kIncA = blockOffsetSMFMAA + (numMIInUnroll//numSet0GroupA) shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncA, "add part of K (block offset + group offset)")) else: kIncA = numMIInUnroll // numSet0GroupA if kernel["ProblemType"]["Sparse"]: kIncA = kIncA elif is_wmma_v3: bpe = tPA["bpe"] vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPA["bpe"] * 8) miVectorWidth = vgprLayout[-1] curElemIdx = group * vgprPerSet0Group * self.states.bpr // bpe if group and curElemIdx % miVectorWidth == 0: kIncA += miVectorWidth elif self.states.asmCaps["HasMFMA_f8f6f4"]: if group == 2 and vgprPerInUnrollA == 8 and not is_wmma_v3: #special layout for F8 kIncA = 56 if kernel["MatrixInstK"] == 128 else 24 if kernel["UseF32XEmulation"]: kIncA = 14 if kernel["MatrixInstK"] == 32 else 6 shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncA, "add part of K")) if kernel["LocalSplitU"] > 1: shiftK.add(SMinI32(dst=sgpr(loopCntSgpr), src0=sgpr(loopCounterName), src1=sgpr("LSUTailLoopOffset"), comment="check lsu bound")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) if not tPA["isSwizzled"]: for se in range(0, vgprPerSet0Group): for a in range(0, kernel["MIWaveTileA"]): for ti in range(0, numTileInInstA): for iui in range(0, innerUnroll): bk = se + group * vgprPerSet0Group + ti * vgprPerInUnrollA aStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, lc=unrollLoopIdx, bk=bk), 1) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=0, src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="set 0 if K_idx >= sizeL")) # separate code for gfx950 mx if group == 0 and kernel["ProblemType"]["MXBlockA"] and isgfx950: for mxsa in range(0, kernel["MIWaveTileMXSA"]): for iui in range(0, innerUnroll): for bk in range(0, vgprPerInputMXSA): mxsaStr_base = self.generateSrcStrForMFMAshiftK(kernel, tPA["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSA, m, u, iui, mxsa, bk=bk) mxsaStr = vgpr(mxsaStr_base, 1) shiftK.add(VCndMaskB32(dst=mxsaStr, src0=mxsaStr, src1=0, src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) if kernel["ProblemType"]["Sparse"] == 2 and numMIInUnroll//8 >= 1: shiftK.add(vectorStaticRemainder(dummy, kReg_first, "Serial", kernel["WavefrontSize"], tmpVgpr, tmpSgprInfo)) shiftK.add(vectorStaticDivide(kReg_first, kReg_first, dividerFortidInK, tmpVgpr)) vgprPerSet0Group = 1 elif vgprPerInUnrollB <= 2: vgprPerSet0Group = vgprPerInUnrollB elif is_wmma_v2 and vgprPerInUnrollB > 2: shiftK.add(vectorStaticRemainder(dummy, kReg_first, "Serial", kernel["WavefrontSize"], tmpVgpr, tmpSgprInfo)) shiftK.add(vectorStaticDivide(kReg_first, kReg_first, dividerFortidInK, tmpVgpr)) vgprPerSet0Group = 4 else: shiftK.add(vectorStaticRemainder(dummy, kReg_first, "Serial", kernel["WavefrontSize"], tmpVgpr, tmpSgprInfo)) shiftK.add(vectorStaticDivide(kReg_first, kReg_first, dividerFortidInK, tmpVgpr)) if is_wmma_v3 and kernel["UseF32XEmulation"]: vgprPerSet0Group = 1 elif kernel["ProblemType"]["DataType"].is6bitFloat(): vgprPerSet0Group = 3 else: vgprPerSet0Group = 2 if tPB["bpe"] == 0.75 and kernel["enableLDSTrB"]: for b in range(0, kernel["MIWaveTileB"]): for ti in range(0, numTileInInstB): for iui in range(0, innerUnroll): bStr = self.generateSrcStrForMFMA(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, b, unrollLoopIdx) shiftK.add(_shiftLrElements(bStr, vgprPerInUnrollB, b)) numSet0GroupB = vgprPerInUnrollB//vgprPerSet0Group for group in range(0, numSet0GroupB): if numSet0GroupB > 1 or (is_wmma_v2 and vgprPerInUnrollB > 2): if group == 0: multiplyBy = numMIInUnroll if kernel["ProblemType"]["Sparse"] == 2 and self.states.asmCaps["HasSWMMAC_gfx1250"]: multiplyBy = numMIInUnroll // blocksPerTGroupSMFMAB elif kernel["ProblemType"]["Sparse"] and not self.states.asmCaps["HasSWMMAC_gfx1250"]: multiplyBy = numMIInUnroll // blocksPerTGroupSMFMAB elif is_wmma_v3 and kernel["UseF32XEmulation"]: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPB["bpe"] * 8) multiplyBy = vgprLayout[-1] // vgprLayout[2] elif is_wmma_v3: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPB["bpe"] * 8) multiplyBy = vgprLayout[-1] elif vgprPerInUnrollB == 8 and not is_wmma_v2: multiplyBy = numMIInUnroll//2 shiftK.add(vectorStaticMultiply(vgpr(kReg_first), vgpr(kReg_first), multiplyBy, tmpSgprInfo)) shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), 0, "")) elif is_wmma_v3 and kernel["UseF32XEmulation"]: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPB["bpe"] * 8) miVectorWidth = vgprLayout[-1] elemPerHalf = miVectorWidth // vgprLayout[2] kOffsetB = (group // elemPerHalf) * miVectorWidth + (group % elemPerHalf) shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), kOffsetB, "K index for group %d" % group)) elif blocksPerTGroupSMFMAB == 2 and (group * vgprPerSet0Group) == (elementsPerBlockSMFMAB * numRegistersIn): if kernel["ProblemType"]["Sparse"]: kIncB = blockOffsetSMFMAB + (numMIInUnroll//numSet0GroupB) * (max(group - 1, 0) if not self.states.asmCaps["HasSWMMAC_gfx1250"] else 1) else: kIncB = blockOffsetSMFMAB + (numMIInUnroll//numSet0GroupB) shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncB, "add part of K (block offset + group offset)")) else: kIncB = numMIInUnroll//numSet0GroupB if kernel["ProblemType"]["Sparse"]: kIncB = kIncB elif is_wmma_v3: bpe = tPB["bpe"] vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPB["bpe"] * 8) miVectorWidth = vgprLayout[-1] curElemIdx = group * vgprPerSet0Group * self.states.bpr // bpe if group and curElemIdx % miVectorWidth == 0: kIncB += miVectorWidth elif self.states.asmCaps["HasMFMA_f8f6f4"]: if group == 2 and vgprPerInUnrollB == 8 and not is_wmma_v3: kIncB = 56 if kernel["MatrixInstK"] == 128 else 24 if kernel["UseF32XEmulation"]: kIncB = 14 if kernel["MatrixInstK"] == 32 else 6 shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncB, "add part of K")) # replace 0 for differnet thread if kernel["LocalSplitU"] > 1: shiftK.add(SMinI32(dst=sgpr(loopCntSgpr), src0=sgpr(loopCounterName), src1=sgpr("LSUTailLoopOffset"), comment="check lsu bound")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) for se in range(0, vgprPerSet0Group): for b in range(0, kernel["MIWaveTileB"]): for ti in range(0, numTileInInstB): for iui in range(0, innerUnroll): bk = se + group * vgprPerSet0Group + ti * vgprPerInUnrollB bStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, b, lc=unrollLoopIdx, bk=bk), 1) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=0, src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="set 0 if K_idx >= sizeL")) # separate code for gfx950 mx if group == 0 and kernel["ProblemType"]["MXBlockB"] and isgfx950: for mxsb in range(0, kernel["MIWaveTileMXSB"]): for iui in range(0, innerUnroll): for bk in range(0, vgprPerInputMXSB): mxsbStr_base = self.generateSrcStrForMFMAshiftK(kernel, tPB["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSB, m, u, iui, mxsb, bk=bk) mxsbStr = vgpr(mxsbStr_base, 1) shiftK.add(VCndMaskB32(dst=mxsbStr, src0=mxsbStr, src1=0, src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) # replace elements with 0 for same thread, this conducting shift and mask between numElementsPerRead # Skip when numRegistersIn >= 1 on wmma_v3 with XF32 emulation: each element fills a full register, # so Phase 1 per-element masking (vgprPerSet0Group=1) already handles all cases. if numMIInUnroll > 1 and kernel["AssertSummationElementMultiple"] < 32 and tPA["bpe"] != 0.75 and tPB["bpe"] != 0.75 \ and not (is_wmma_v3 and kernel["UseF32XEmulation"]): alignment = vgprPerInUnroll if is_wmma_v2 else (2 if vgprPerInUnroll > 1 else 1) abReg = self.vgprPool.checkOutAligned(vgprPerInUnroll, alignment, tag="mfmaIter_abReg") if (vgprPerInUnroll < 4 and is_mfma) or is_wmma_v2: shiftK.add(VSubU32(dst=vgpr(kReg), src0=sgpr(loopCntSgpr), src1=vgpr(kReg_first), comment="get distance between size and k index")) shiftK.add(VCmpLtI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=vgpr(kReg), src1=numMIInUnroll, comment="set partial 0 if distance less than input per thread")) TailLoop_SkipZeroOutMask = Label((self.labels.getUniqueNamePrefix("TailLoop_SkipZeroOutMask")), comment="") if is_wmma_v3: shiftK.add(SAndB32(dst=sgpr(tmpSgprX1), src0=sgpr(loopCntSgpr), src1=int((64 // (numRegistersIn * 32))-1), comment="get inputs for edge thread")) shiftK.add(SSubU32(dst=sgpr(tmpSgprX1), src0=int((64 // (numRegistersIn * 32))), src1=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) shiftK.add(SLShiftLeftB32(dst=sgpr(tmpSgprX1), shiftHex=log2(shiftPerElement), src=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) else: shiftK.add(SAndB32(dst=sgpr(tmpSgprX1), src0=sgpr("SizeL"), src1=8-1, comment="if summation is multiple of 8, skip masking")) shiftK.add(SCmpEQU32(src0=sgpr(tmpSgprX1), src1=0)) shiftK.add(SCBranchSCC1(labelName=TailLoop_SkipZeroOutMask.getLabelName(), comment="skip mask")) shiftK.add(SAndB32(dst=sgpr(tmpSgprX1), src0=sgpr(loopCntSgpr), src1=numMIInUnroll-1, comment="get inputs for edge thread")) shiftK.add(SSubU32(dst=sgpr(tmpSgprX1), src0=numMIInUnroll, src1=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) shiftK.add(SLShiftLeftB32(dst=sgpr(tmpSgprX1), shiftHex=log2(shiftPerElement), src=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) if vgprPerInUnroll == 1: VShiftLeft = VLShiftLeftB32 elif vgprPerInUnroll >= 2: # for 2, 4, 8 VShiftLeft = VLShiftLeftB64 if vgprPerInUnroll == 4 and is_wmma_v2: tmpVgpr2 = self.vgprPool.checkOutAligned(2, 2, tag="mfmaIter_tmpVgpr2") for a in range(0, kernel["MIWaveTileA"]): for ti in range(0, numTileInInstA): for iui in range(0, innerUnroll): aStr_base = self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, lc=unrollLoopIdx, bk=ti*vgprPerInputA) aStr = vgpr(aStr_base, min(2, vgprPerInUnrollA)) if is_wmma_v2: if vgprPerInUnrollA == 4: a_64_shift = Label(label=self.labels.getNameInc("a_64_Shift"), comment="") a_32_shift = Label(label=self.labels.getNameInc("a_32_Shift"), comment="") a_common = Label(label=self.labels.getNameInc("a_shift_end"), comment="") if kernel["UseF32XEmulation"]: aStr1 = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInUnrollA, m, u, iui, a+1), min(2, vgprPerInUnrollA)) aStr2 = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInUnrollA, m, u, iui, a+2), min(2, vgprPerInUnrollA)) else: aStr1 = vgpr(aStr_base + "+1", min(2, vgprPerInUnrollA)) aStr2 = vgpr(aStr_base + "+2", min(2, vgprPerInUnrollA)) shiftK.add(SMovB32(dst=sgpr(tmpSgprX3), src=sgpr(tmpSgprX1), comment="sgpr used for minic shift 128 bit")) shiftK.add(SCmpGeI32(src0=sgpr(tmpSgprX3), src1=64, comment="check offset > 63")) shiftK.add(SCBranchSCC1(labelName=a_64_shift.getLabelName(), comment="jump when positive")) shiftK.add(SCmpGeI32(src0=sgpr(tmpSgprX3), src1=32, comment="check offset > 32")) shiftK.add(SCBranchSCC1(labelName=a_32_shift.getLabelName(), comment="jump when positive")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=aStr, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg), src=vgpr(tmpVgpr2),comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+1), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=aStr1, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+2), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=aStr2, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+3), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(SBranch(a_common.getLabelName())) shiftK.add(a_32_shift) shiftK.add(SSubU32(dst=sgpr(tmpSgprX3), src0=sgpr(tmpSgprX3), src1=32, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg), src=0, comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=aStr, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+1), src=vgpr(tmpVgpr2),comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+2), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=aStr1, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+3), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(SAddU32(dst=sgpr(tmpSgprX3), src0=sgpr(tmpSgprX3), src1=32, comment="")) shiftK.add(SBranch(a_common.getLabelName())) shiftK.add(a_64_shift) shiftK.add(SSubU32(dst=sgpr(tmpSgprX3), src0=sgpr(tmpSgprX3), src1=64, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg), src=0, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+1), src=0, comment="")) shiftK.add(VShiftLeft(dst=vgpr(abReg+2, 2), shiftHex=sgpr(tmpSgprX3), src=aStr, comment="")) shiftK.add(SAddU32(dst=sgpr(tmpSgprX3), src0=sgpr(tmpSgprX3), src1=64, comment="")) shiftK.add(a_common) elif vgprPerInUnrollA == 2: shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInUnrollA), shiftHex=sgpr(tmpSgprX1), src=aStr, comment="")) elif vgprPerInUnrollA > 1: assert False, f"Invalid vgprPerInUnrollA: {vgprPerInUnrollA}" for ir in range(0, vgprPerInUnrollA): bk = ir + ti * vgprPerInUnrollA aStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, lc=unrollLoopIdx, bk=bk), 1) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=vgpr(abReg+ir), src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) else: # mfma or wmma_v3 if kernel["ProblemType"]["Sparse"]: if vgprPerInUnrollA == 2: aStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, lc=unrollLoopIdx, bk=ti*vgprPerInUnroll), 2) shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInUnrollA), shiftHex=sgpr(tmpSgprX1), src=aStr, comment="")) elif vgprPerInUnroll <= 2: shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInUnrollA), shiftHex=sgpr(tmpSgprX1), src=aStr, comment="")) if vgprPerInUnrollA >= 4 and vgprPerInUnrollA % 2 == 0: for ivgpr in range(0, vgprPerInUnrollA, 2): bk = ivgpr + ti * vgprPerInUnrollA aStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, lc=unrollLoopIdx, bk=bk), 2) shiftK.add(VShiftLeft(dst=vgpr(abReg+ivgpr, 2), shiftHex=sgpr(tmpSgprX1), src=aStr, comment="")) for ir in range(0, vgprPerInUnrollA): bk = ir + ti * vgprPerInUnrollA aStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, a, lc=unrollLoopIdx, bk=bk), 1) elemIdx = ir * self.states.bpr // tPA["bpe"] if is_wmma_v3: # may check 64 bit vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPA["bpe"] * 8) mivw = vgprLayout[-1] if vgprPerInUnrollA >= 2: kIncA = int((64 // (numRegistersIn * 32))) if elemIdx == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), kIncA, "add part of K: 64 bit groupd")) elif elemIdx % mivw == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncA+mivw, "add part of K: 64 bit groupd")) elif ir % 2 == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncA, "add part of K: 64 bit groupd")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=vgpr(abReg+ir), src2=sgpr(tmpSgprX2), comment="")) else: if vgprPerInUnrollA >= 4: if ir == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), numMIInUnroll//numSet0GroupA, "add part of K")) elif ir % 2 == 0: kIncA = numMIInUnroll//numSet0GroupA if ir == 4 and self.states.asmCaps["HasMFMA_f8f6f4"]: kIncA = 56 if kernel["MatrixInstK"] == 128 else 24 if kernel["UseF32XEmulation"]: kIncA = 14 if kernel["MatrixInstK"] == 32 else 6 shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncA, "add part of K")) # replace 0 for differnet thread if kernel["LocalSplitU"] > 1: shiftK.add(SMinI32(dst=sgpr(loopCntSgpr), src0=sgpr(loopCounterName), src1=sgpr("LSUTailLoopOffset"), comment="check lsu bound")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2,2), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=vgpr(abReg+ir), src2=sgpr(tmpSgprX2,2), comment="")) for b in range(0, kernel["MIWaveTileB"]): for ti in range(0, numTileInInstB): for iui in range(0, innerUnroll): bStr_base = self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, b, lc=unrollLoopIdx, bk=ti*vgprPerInUnrollB) bStr = vgpr(bStr_base, min(2, vgprPerInUnrollB)) if is_wmma_v2: if vgprPerInUnrollB == 4: b_64_shift = Label(label=self.labels.getNameInc("b_64_Shift"), comment="") b_32_shift = Label(label=self.labels.getNameInc("b_32_Shift"), comment="") b_common = Label(label=self.labels.getNameInc("b_shift_end"), comment="") if kernel["UseF32XEmulation"]: bStr1 = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInUnrollB, m, u, iui, b+1), min(2, vgprPerInUnrollB)) bStr2 = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInUnrollB, m, u, iui, b+2), min(2, vgprPerInUnrollB)) else: bStr1 = vgpr(bStr_base + "+1", min(2, vgprPerInUnrollB)) bStr2 = vgpr(bStr_base + "+2", min(2, vgprPerInUnrollB)) shiftK.add(SMovB32(dst=sgpr(tmpSgprX3), src=sgpr(tmpSgprX1), comment="sgpr used for minic shift 128 bit")) shiftK.add(SCmpGeI32(src0=sgpr(tmpSgprX3), src1=64, comment="check offset >63")) shiftK.add(SCBranchSCC1(labelName=b_64_shift.getLabelName(), comment="jump when positive")) shiftK.add(SCmpGeI32(src0=sgpr(tmpSgprX3), src1=32, comment="check offset >32")) shiftK.add(SCBranchSCC1(labelName=b_32_shift.getLabelName(), comment="jump when positive")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=bStr, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg), src=vgpr(tmpVgpr2),comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+1), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=bStr1, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+2), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=bStr2, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+3), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(SBranch(b_common.getLabelName())) shiftK.add(b_32_shift) shiftK.add(SSubU32(dst=sgpr(tmpSgprX3), src0=sgpr(tmpSgprX3), src1=32, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg), src=0, comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=bStr, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+1), src=vgpr(tmpVgpr2),comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+2), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(VShiftLeft(dst=vgpr(tmpVgpr2, 2), shiftHex=sgpr(tmpSgprX3), src=bStr1, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+3), src=vgpr(tmpVgpr2+1),comment="")) shiftK.add(SBranch(b_common.getLabelName())) shiftK.add(b_64_shift) shiftK.add(SSubU32(dst=sgpr(tmpSgprX3), src0=sgpr(tmpSgprX3), src1=64, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg), src=0, comment="")) shiftK.add(VMovB32(dst=vgpr(abReg+1), src=0, comment="")) shiftK.add(VShiftLeft(dst=vgpr(abReg+2,2), shiftHex=sgpr(tmpSgprX3), src=bStr, comment="")) shiftK.add(b_common) else: shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInUnrollB), shiftHex=sgpr(tmpSgprX1), src=bStr, comment="")) for ir in range(0, vgprPerInUnrollB): bk = ir + ti * vgprPerInUnrollB bStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, b, lc=unrollLoopIdx, bk=bk), 1) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=vgpr(abReg+ir), src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) else: # mfma or wmma_v3 if kernel["ProblemType"]["Sparse"]: if vgprPerInUnrollB == 2: aStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, a, lc=unrollLoopIdx, bk=ti*vgprPerInUnrollB), 2) shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInUnrollB), shiftHex=sgpr(tmpSgprX1), src=aStr, comment="")) elif vgprPerInUnrollB <= 2: shiftK.add(VShiftLeft(dst=vgpr(abReg, vgprPerInUnrollB), shiftHex=sgpr(tmpSgprX1), src=bStr, comment="")) if vgprPerInUnrollB >= 4 and vgprPerInUnrollB % 2 == 0: for ivgpr in range(0, vgprPerInUnrollB, 2): bk = ivgpr + ti * vgprPerInUnrollB bStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, b, lc=unrollLoopIdx, bk=bk), 2) shiftK.add(VShiftLeft(dst=vgpr(abReg+ivgpr, 2), shiftHex=sgpr(tmpSgprX1), src=bStr, comment="")) for ir in range(0, vgprPerInUnrollB): bk = ir + ti * vgprPerInUnrollB bStr = vgpr(self.generateSrcStrForMFMAshiftK(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, b, lc=unrollLoopIdx, bk=bk), 1) elemIdx = ir * self.states.bpr // tPB["bpe"] if is_wmma_v3: vgprLayout = wmmaV3InputVgprLayout(kernel["MatrixInstruction"], tPB["bpe"] * 8) mivw = vgprLayout[-1] if vgprPerInUnrollB >= 2: kIncB = int((64 // (numRegistersIn * 32))) if elemIdx == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), kIncB, "add part of K: 64 bits group")) elif elemIdx % mivw == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncB+mivw, "add part of K: 64 bit groupd")) elif ir % 2 == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncB, "add part of K: 64 bits group")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=vgpr(abReg+ir), src2=sgpr(tmpSgprX2), comment="")) else: if vgprPerInUnrollB >= 4: if ir == 0: shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg_first), numMIInUnroll//numSet0GroupB, "add part of K")) elif ir % 2 == 0: kIncB = numMIInUnroll//numSet0GroupB if ir == 4: # when vgprPerInUnrollB == 8 kIncB = 56 if kernel["MatrixInstK"] == 128 else 24 if kernel["UseF32XEmulation"]: kIncB = 14 if kernel["MatrixInstK"] == 32 else 6 shiftK.add(VAddU32(vgpr(kReg), vgpr(kReg), kIncB, "add part of K")) # replace 0 for differnet thread if kernel["LocalSplitU"] > 1: shiftK.add(SMinI32(dst=sgpr(loopCntSgpr), src0=sgpr(loopCounterName), src1=sgpr("LSUTailLoopOffset"), comment="check lsu bound")) shiftK.add(VCmpGEI32(dst=sgpr(tmpSgprX2,2), src0=vgpr(kReg), src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=vgpr(abReg+ir), src2=sgpr(tmpSgprX2,2), comment="")) shiftK.add(TailLoop_SkipZeroOutMask) if vgprPerInUnrollB == 4 and is_wmma_v2: if tmpVgpr2 is not None: self.vgprPool.checkIn(tmpVgpr2) if (kernel["ProblemType"]["MXBlockA"] or kernel["ProblemType"]["MXBlockB"]) and not isgfx950: mxK = kernel["MatrixInstK"] mxBlock = max(kernel["ProblemType"]["MXBlockA"], kernel["ProblemType"]["MXBlockB"]) vgprPerInUnroll = max(vgprPerInputMXSA, vgprPerInputMXSB) mxReg = self.vgprPool.checkOutAligned(vgprPerInUnroll, vgprPerInUnroll, "mxReg") shiftK.add(SAndB32(dst=sgpr(tmpSgprX1), src0=sgpr(loopCntSgpr), src1=int(mxK-1), comment="get inputs for edge thread")) shiftK.add(SLShiftRightB32(dst=sgpr(tmpSgprX1), src=sgpr(tmpSgprX1), shiftHex=log2(mxBlock), comment="calculate 64bit groups index")) shiftK.add(SSubU32(dst=sgpr(tmpSgprX1), src0=(mxK//mxBlock), src1=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) shiftK.add(SLShiftLeftB32(dst=sgpr(tmpSgprX1), shiftHex=log2(8), src=sgpr(tmpSgprX1), comment="use shift to fill 0 for outside element")) if vgprPerInUnroll == 1: VShiftLeft = VLShiftLeftB32 VShiftRight = VLShiftRightB32 elif vgprPerInUnroll == 2: VShiftLeft = VLShiftLeftB64 VShiftRight = VLShiftRightB64 else: raise Exception(f"unsupport vgprPerInUnroll {vgprPerInUnroll}") if kernel["ProblemType"]["MXBlockA"] and not isgfx950: for mxsa in range(0, kernel["MIWaveTileMXSA"]): for iui in range(0, innerUnroll): mxsaStr_base = self.generateSrcStrForMFMA(kernel, tPA["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSA, m, u, iui, mxsa) mxsaStr = vgpr(mxsaStr_base, vgprPerInputMXSA) shiftK.add(VShiftLeft(dst=vgpr(mxReg, vgprPerInUnroll), shiftHex=sgpr(tmpSgprX1), src=mxsaStr, comment="")) shiftK.add(VShiftRight(dst=vgpr(mxReg, vgprPerInUnroll), shiftHex=sgpr(tmpSgprX1), src=vgpr(mxReg, vgprPerInUnroll), comment="")) shiftK.add(VCmpGTI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=mxK, src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) for bk in range(0, vgprPerInUnroll): mxsaStr_base = self.generateSrcStrForMFMA(kernel, tPA["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSA, m, u, iui, mxsa, bk=bk) mxsaStr = vgpr(mxsaStr_base, 1) shiftK.add(VCndMaskB32(dst=mxsaStr, src0=mxsaStr, src1=vgpr(mxReg+bk), src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) if kernel["ProblemType"]["MXBlockB"] and not isgfx950: for mxsb in range(0, kernel["MIWaveTileMXSB"]): for iui in range(0, innerUnroll): mxsbStr_base = self.generateSrcStrForMFMA(kernel, tPB["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSB, m, u, iui, mxsb) mxsbStr = vgpr(mxsbStr_base, vgprPerInputMXSB) shiftK.add(VShiftLeft(dst=vgpr(mxReg, vgprPerInUnroll), shiftHex=sgpr(tmpSgprX1), src=mxsbStr, comment="")) shiftK.add(VShiftRight(dst=vgpr(mxReg, vgprPerInUnroll), shiftHex=sgpr(tmpSgprX1), src=vgpr(mxReg, vgprPerInUnroll), comment="")) shiftK.add(VCmpGTI32(dst=sgpr(tmpSgprX2, self.states.laneSGPRCount), src0=mxK, src1=sgpr(loopCntSgpr), comment="check K index >= Size L")) for bk in range(0, vgprPerInUnroll): mxsbStr_base = self.generateSrcStrForMFMA(kernel, tPB["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSB, m, u, iui, mxsb, bk=bk) mxsbStr = vgpr(mxsbStr_base, 1) shiftK.add(VCndMaskB32(dst=mxsbStr, src0=mxsbStr, src1=vgpr(mxReg+bk), src2=sgpr(tmpSgprX2, self.states.laneSGPRCount), comment="")) if kernel["LocalSplitU"] > 1: self.sgprPool.checkIn(loopCntSgpr) else: #wmma iui = 0 abReg = self.vgprPool.checkOutAligned(2, 2, tag="wmmaIter_abReg") with self.allocTmpSgpr(3, tag="wmmaIter_tmpSgprInfo") as tmpSgprInfo: sgprShift = tmpSgprInfo.idx sgpr64bIdx = sgprShift + 1 sgprMask = sgprShift + 2 shiftK.add(SSubI32(dst=sgpr(sgprShift), src0=sgpr(loopCounterName), src1=1, comment="calculate 64bit groups index")) shiftK.add(SLShiftRightB32(dst=sgpr(sgpr64bIdx), src=sgpr(sgprShift), shiftHex=log2(64 // (numRegistersIn * 32)), comment="calculate 64bit groups index")) shiftK.add(SAndB32(dst=sgpr(sgprShift), src0=sgpr(sgprShift), src1=int((64 // (numRegistersIn * 32))-1), comment="calculate shift value")) shiftK.add(SSubI32(dst=sgpr(sgprShift), src0=int((64 // (numRegistersIn * 32))-1), src1=sgpr(sgprShift), comment="calculate shift value")) shiftK.add(SLShiftLeftB32(dst=sgpr(sgprShift), shiftHex=log2(numRegistersIn * 32), src=sgpr(sgprShift), comment="calculate shift value")) for it in range(int((kernel["MatrixInstK"] * numRegistersIn) // 2)): # handle 64 bit per iteration shiftK.add(VCmpEQI32(dst=sgpr(sgprMask), src0=sgpr(sgpr64bIdx), src1=it, comment='handle this 64bit group: part 1')) for a in range(0, kernel["MIWaveTileA"]): aStr = vgpr("ValuA_X%u_I%u+%u+%u" % (m, iui, a*vgprPerInUnroll, it*2), 2) shiftK.add(VLShiftLeftB64(dst=vgpr(abReg,2), shiftHex=sgpr(sgprShift), src=aStr, comment=f"shfit for ValuA[{it*2}:{it*2+1}]")) for bk in range(2): aStr = vgpr("ValuA_X%u_I%u+%u+%u+%u" % (m, iui, a*vgprPerInUnroll, it*2, bk)) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=vgpr(abReg+bk), src2=sgpr(sgprMask), comment="shift if in this 64b group")) for b in range(0, kernel["MIWaveTileB"]): bStr = vgpr("ValuB_X%u_I%u+%u+%u" % (m, iui, b*vgprPerInUnroll, it*2), 2) shiftK.add(VLShiftLeftB64(dst=vgpr(abReg,2), shiftHex=sgpr(sgprShift), src=bStr, comment=f"shfit for ValuB[{it*2}:{it*2+1}]")) for bk in range(2): bStr = vgpr("ValuB_X%u_I%u+%u+%u+%u" % (m, iui, b*vgprPerInUnroll, it*2, bk)) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=vgpr(abReg+bk), src2=sgpr(sgprMask), comment="shift if in this 64b group")) if it > 0: shiftK.add(VCmpLtI32(dst=sgpr(sgprMask), src0=sgpr(sgpr64bIdx), src1=it, comment='handle this 64bit group: part 2')) for a in range(0, kernel["MIWaveTileA"]): for bk in range(2): aStr = vgpr("ValuA_X%u_I%u+%u+%u+%u" % (m, iui, a*vgprPerInUnroll, it*2, bk)) shiftK.add(VCndMaskB32(dst=aStr, src0=aStr, src1=0, src2=sgpr(sgprMask), comment="shift if in this 64b group")) for b in range(0, kernel["MIWaveTileB"]): for bk in range(2): bStr = vgpr("ValuB_X%u_I%u+%u+%u+%u" % (m, iui, b*vgprPerInUnroll, it*2, bk)) shiftK.add(VCndMaskB32(dst=bStr, src0=bStr, src1=0, src2=sgpr(sgprMask), comment="shift if in this 64b group")) s_nop = 2 if s_nop != 0: imod.add(SNop(waitState=(s_nop - 1), comment="")) prevAccIdx = -1 for iui in range(0, innerUnroll): if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: iuiM_new = (iui//self.states.numReadsIterCoalescedMetadata)*self.states.numReadsIterCoalescedMetadata iuiM_new_offset = iui%self.states.numReadsIterCoalescedMetadata*vgprPerInputM zgemmVaddSrcCheck = [[], [], []] # to avoid generating redundant v_add outer = 1 loopSwap = False # complex case, swap inner loop and outer loop so that idxA comes outer # this is to re-use same tmp vgpr to nagate ai or ar if (kernel["ProblemType"]["DataType"].isComplex() and tPB["tile01Idx"]) or kernel["HalfPLRA"]: outer = 0 loopSwap = True inner = 1 - outer # inner is the opposite of outer shiftedIndicesA = set() shiftedIndicesB = set() idxOuter_start = 0 idxInner_start = 0 idxOuter_stop = kernel["MIWaveTile"][outer] idxInner_stop = kernel["MIWaveTile"][inner] numSubTiles = kernel["numSubTiles"] if numSubTiles > 1 and not self.states.inTailLoop: # iter (idxOuter_start, idxOuter_stop) (idxInner_start, idxInner_stop) MFMA # 0 (0,4) (0,4) MFMA(A0,B0) # 1 (0,4) (4,8) MFMA(A1,B0) # 2 (4,8) (0,4) MFMA(A0,B1) # 3 (4,8) (4,8) MFMA(A1,B1) outerBy2=(kernel["MIWaveTile"][outer]//numSubTiles) innerBy2=(kernel["MIWaveTile"][inner]//numSubTiles) outerMod2=(kernel["MIWaveTile"][outer]%numSubTiles) innerMod2=(kernel["MIWaveTile"][inner]%numSubTiles) idxHalfO = u//numSubTiles idxHalfI = u % numSubTiles idxOuter_start = (outerBy2 + outerMod2)*idxHalfO idxInner_start = (innerBy2 + innerMod2)*idxHalfI idxOuter_stop = kernel["MIWaveTile"][outer] - (1-idxHalfO)* outerBy2 idxInner_stop = kernel["MIWaveTile"][inner] - (1-idxHalfI)* innerBy2 for idxOuter in range(idxOuter_start, idxOuter_stop): for idxInner in range(idxInner_start, idxInner_stop): idx0 = idxInner idx1 = idxOuter if loopSwap: idx0, idx1 = idx1, idx0 accIdx = idx1 * kernel["MIWaveTile"][0] + idx0 accStart = accIdx * accs_per_wave accEnd = accStart + accs_per_wave - 1 idxA = idx0 if tPB["tile01Idx"] else idx1 idxB = idx1 if tPB["tile01Idx"] else idx0 aStr_base = self.generateSrcStrForMFMA(kernel, tPA, innerUnroll, vregSetIdx, vgprPerInputA, m, u, iui, idxA, unrollLoopIdx) bStr_base = self.generateSrcStrForMFMA(kernel, tPB, innerUnroll, vregSetIdx, vgprPerInputB, m, u, iui, idxB, unrollLoopIdx) aStr = vgpr(aStr_base, vgprPerInputA) bStr = vgpr(bStr_base, vgprPerInputB) if not tail and tPA["bpe"] == 0.75 and kernel["enableLDSTrA"]: if idxInner not in shiftedIndicesA: imod.add(_shiftLrElements(aStr_base, vgprPerInputA, idxInner)) shiftedIndicesA.add(idxInner) if not tail and tPB["bpe"] == 0.75 and kernel["enableLDSTrB"]: if idxOuter not in shiftedIndicesB: imod.add(_shiftLrElements(bStr_base, vgprPerInputB, idxOuter)) shiftedIndicesB.add(idxOuter) Str0 = aStr if tPB["tile01Idx"] else bStr Str1 = bStr if tPB["tile01Idx"] else aStr mxsaScaleSel = 0 mxsbScaleSel = 0 if kernel["ProblemType"]["MXBlockA"]: mxsaIdx, mxsaScaleSel = self.mxsTileSpanScaleSel(kernel, tPA["MX"], idxA) mxsaStr_base = self.generateSrcStrForMFMA(kernel, tPA["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSA, m, u, iui, mxsaIdx) mxsaStr = vgpr(mxsaStr_base, vgprPerInputMXSA) else: mxsaStr = vgpr("ValuMXSDummy") if kernel["ProblemType"]["MXBlockB"] == 32 else vgpr("ValuMXSDummy",2) if kernel["ProblemType"]["MXBlockB"]: mxsbIdx, mxsbScaleSel = self.mxsTileSpanScaleSel(kernel, tPB["MX"], idxB) mxsbStr_base = self.generateSrcStrForMFMA(kernel, tPB["MX"], innerUnroll, vregSetIdx, vgprPerInputMXSB, m, u, iui, mxsbIdx) mxsbStr = vgpr(mxsbStr_base, vgprPerInputMXSB) else: mxsbStr = vgpr("ValuMXSDummy") if kernel["ProblemType"]["MXBlockA"] == 32 else vgpr("ValuMXSDummy",2) StrMX0 = mxsaStr if tPB["tile01Idx"] else mxsbStr StrMX1 = mxsbStr if tPB["tile01Idx"] else mxsaStr # matrix_a_scale/matrix_b_scale correspond to the mxsa/mxsb operand *positions* # (which follow the same tile01Idx swap as StrMX0/StrMX1). scaleSelMX0 = mxsaScaleSel if tPB["tile01Idx"] else mxsbScaleSel scaleSelMX1 = mxsbScaleSel if tPB["tile01Idx"] else mxsaScaleSel if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: idxM = idxB if kernel["ProblemType"]["Sparse"] == 2 else idxA m_new = idxM * (self.states.numReadsIterCoalescedMetadata * vgprPerInputM if not kernel["enableLDSTrMetadata"] else ceil(self.states.numReadsIterCoalescedMetadata * vgprPerInputM / 2) * 2) mStr = "ValuMetadata_X%u_I%u+%u+%u+%u" % (vgprBufferM_new, iuiM_new, m_new, vgprBufferM_new_offset, iuiM_new_offset) mStr = vgpr(mStr, vgprPerInputM) if kernel["ProblemType"]["DataType"].isComplex(): # override because complex mul is emulated by 4 mfma insts # TODO: adopt component system miInInstType = miOutInstType #"f32" for SingleComplex, "f64" for DoubleComplex ccA = kernel["ProblemType"]["ComplexConjugateA"] ccB = kernel["ProblemType"]["ComplexConjugateB"] ccVgprs = [None]*3 # three terms that can be negated: [real1, imag0, imag1] ccInsts = [None]*3 accImOffset = accVgprImagNumOffset(kernel) accStartSrcImg = accStart+accImOffset accEndSrcImg = accStartSrcImg + accs_per_wave - 1 # vgpr A,B setting. In complex case, numRegistersIn does not match. Use numRegistersOut instead ar_base = aStr_base ai_base = ar_base + "+%u"%numRegistersOut ar = vgpr(ar_base, numRegistersOut) ai = vgpr(ai_base, numRegistersOut) br_base = bStr_base bi_base = br_base + "+%u"%numRegistersOut br = vgpr(br_base, numRegistersOut) bi = vgpr(bi_base, numRegistersOut) minus_ar = ar.getMinus() minus_ai = ai.getMinus() if miOutInstType == InstType.INST_F32: VAddX = VAddF32 elif miOutInstType == InstType.INST_F64: VAddX = VAddF64 else: printExit("Unsupported v_add type %s"%miOutInstType) offsetVgpr = [0,0,0] forceGenerate = True # always generate negate since ccVgprs are checked in/out each iteration if ccA == ccB: arrayIndex = 0 ccVgprs[arrayIndex] = self.vgprPool.checkOutAligned(numRegistersOut, numRegistersOut, "negate r1") # generate negate code only when same code is not generated (avoid generating same (redundant) code again if forceGenerate or (ai not in zgemmVaddSrcCheck[arrayIndex]): ccInsts[arrayIndex] = VAddX(dst=vgpr(ccVgprs[arrayIndex] + offsetVgpr[arrayIndex], numRegistersOut), src0=minus_ai, src1=0, comment="Ai=-Ai") zgemmVaddSrcCheck[arrayIndex].append(ai) if ccA: arrayIndex = 1 ccVgprs[arrayIndex] = self.vgprPool.checkOutAligned(numRegistersOut, numRegistersOut, "negate i0") # generate negate code only when same code is not generated (avoid generating same (redundant) code again if forceGenerate or (ai not in zgemmVaddSrcCheck[arrayIndex]): ccInsts[arrayIndex] = VAddX(dst=vgpr(ccVgprs[arrayIndex] + offsetVgpr[arrayIndex], numRegistersOut), src0=minus_ai, src1=0, comment="Ai=-Ai") zgemmVaddSrcCheck[arrayIndex].append(ai) if ccB: arrayIndex = 2 ccVgprs[arrayIndex] = self.vgprPool.checkOutAligned(numRegistersOut, numRegistersOut, "negate i1") # generate negate code only when same code is not generated (avoid generating same (redundant) code again if forceGenerate or (ar not in zgemmVaddSrcCheck[arrayIndex]): ccInsts[arrayIndex] = VAddX(dst=vgpr(ccVgprs[arrayIndex] + offsetVgpr[arrayIndex], numRegistersOut), src0=minus_ar, src1=0, comment="Ar=-Ar") zgemmVaddSrcCheck[arrayIndex].append(ar) (src0, src1) = (br, ar) if kernel["SourceSwap"] else (ar, br) for inst in ccInsts: if inst is not None: imod.add(inst) variant = [kernel["MIBlock"][0], kernel["MIBlock"][1], kernel["MatrixInstK"], kernel["MatrixInstB"]] imod.add(MFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=False, \ acc=self.accVgprReadWriteIndex(kernel, accStart, (accEnd-accStart+1)), a=src0, b=src1, acc2=self.accVgprReadWriteIndex(kernel, accStart, (accEnd-accStart+1)), \ comment="Cr += Ar*Br")) (src0, src1) = (bi, (vgpr(ccVgprs[0] + offsetVgpr[0], numRegistersOut) if ccVgprs[0] else ai)) if kernel["SourceSwap"] else ((vgpr(ccVgprs[0] + offsetVgpr[0], numRegistersOut) if ccVgprs[0] else ai), bi) imod.add(MFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=False, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), a=src0, b=src1, acc2=self.accVgprReadWriteIndex(kernel, accStart, (accEnd-accStart+1)), \ comment="Cr += %sAi*Bi"%("-" if ccVgprs[0] else ""))) (src0, src1) = (br, (vgpr(ccVgprs[1] + offsetVgpr[1], numRegistersOut) if ccVgprs[1] else ai)) if kernel["SourceSwap"] else ((vgpr(ccVgprs[1] + offsetVgpr[1], numRegistersOut) if ccVgprs[1] else ai), br) imod.add(MFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=False, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accImOffset), (accEnd-accStart+1)), a=src0, b=src1, acc2=self.accVgprReadWriteIndex(kernel, accStartSrcImg, (accEndSrcImg-accStartSrcImg+1)), \ comment="Ci += %sAi*Br"%("-" if ccVgprs[1] else ""))) (src0, src1) = (bi, (vgpr(ccVgprs[2] + offsetVgpr[2], numRegistersOut) if ccVgprs[2] else ar)) if kernel["SourceSwap"] else ((vgpr(ccVgprs[2] + offsetVgpr[2], numRegistersOut) if ccVgprs[2] else ar), bi) imod.add(MFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=False, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accImOffset+accStoreCIdx), (accEnd-accStart+1)), a=src0, b=src1, acc2=self.accVgprReadWriteIndex(kernel, accStartSrcImg, (accEndSrcImg-accStartSrcImg+1)), \ comment="Ci += %sAr*Bi"%("-" if ccVgprs[2] else ""))) for v in ccVgprs: if v is not None: self.vgprPool.checkIn(v) else: if kernel["SourceSwap"]: src0 = Str1 src1 = Str0 srcMX0 = StrMX1 srcMX1 = StrMX0 scaleSelSrc0 = scaleSelMX1 scaleSelSrc1 = scaleSelMX0 miInScale0InstType = miInScaleBInstType miInScale1InstType = miInScaleAInstType else: src0 = Str0 src1 = Str1 srcMX0 = StrMX0 srcMX1 = StrMX1 scaleSelSrc0 = scaleSelMX0 scaleSelSrc1 = scaleSelMX1 miInScale0InstType = miInScaleAInstType miInScale1InstType = miInScaleBInstType variant = [kernel["MIBlock"][0], kernel["MIBlock"][1], kernel["MatrixInstK"], kernel["MatrixInstB"]] waits = self.mfmaIter_waitCount(kernel) if waits > 0 and prevAccIdx == accIdx: imod.add(SNop(waits - 1, "Wait for C")) if(kernel["ProblemType"]["Sparse"]): if kernel["DirectToVgprSparseMetadata"]: miWaveTile = kernel["MIWaveTileB"] if kernel["ProblemType"]["Sparse"] == 2 else kernel["MIWaveTileA"] idx = idx1 if kernel["ProblemType"]["Sparse"] == 2 else idx0 accInStart = miWaveTile * kernel["LoopIters"] * unrollLoopIdx + idx * kernel["LoopIters"] + unrollIdx imod.add(SMFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0, b=src1, metadata=vgpr("ValuMetadata+%u"%(accInStart)), neg=neg_flag, \ comment="left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) else: imod.add(SMFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0, b=src1, metadata=mStr, neg=neg_flag, \ comment="left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) else: acc2_args = {"acc2": self.accVgprReadWriteIndex(kernel, accStart, (accEnd-accStart+1))} if kernel["UseF32XEmulation"]: abOffsetStr = "+" + str(vgprPerInputA // 2) if kernel["SourceSwap"]: src1_0 = vgpr(aStr_base[:-4], vgprPerInputA / 2) src1_1 = vgpr(aStr_base[:-4] + abOffsetStr, vgprPerInputA / 2) src0_0 = vgpr(bStr_base[:-4], vgprPerInputB / 2) src0_1 = vgpr(bStr_base[:-4] + abOffsetStr, vgprPerInputB / 2) else: src1_0 = vgpr(bStr_base[:-4], vgprPerInputB / 2) src1_1 = vgpr(bStr_base[:-4] + abOffsetStr, vgprPerInputB / 2) src0_0 = vgpr(aStr_base[:-4], vgprPerInputA / 2) src0_1 = vgpr(aStr_base[:-4] + abOffsetStr, vgprPerInputA / 2) imod.add(MFMAInstruction(instType=InstType.INST_BF16, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0_0, b=src1_0, **acc2_args, neg=neg_flag,\ comment="src0_h*src1_h, left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) if kernel["SourceSwap"]: imod.add(MFMAInstruction(instType=InstType.INST_BF16, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0_0, b=src1_1, **acc2_args, neg=neg_flag,\ comment="src0_h*src1_l, left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) imod.add(MFMAInstruction(instType=InstType.INST_BF16, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0_1, b=src1_0, **acc2_args, neg=neg_flag,\ comment="src0_l*src1_h, left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) else: imod.add(MFMAInstruction(instType=InstType.INST_BF16, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0_1, b=src1_0, **acc2_args, neg=neg_flag,\ comment="src0_l*src1_h, left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) imod.add(MFMAInstruction(instType=InstType.INST_BF16, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0_0, b=src1_1, **acc2_args, neg=neg_flag,\ comment="src0_h*src1_l, left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) elif kernel["ProblemType"]["MXBlockA"] or kernel["ProblemType"]["MXBlockB"]: block = max(kernel["ProblemType"]["MXBlockA"], kernel["ProblemType"]["MXBlockB"]) imod.add(MXMFMAInstruction(instType=miInInstType, accType=miOutInstType, \ mxScaleAType=miInScale0InstType, mxScaleBType=miInScale1InstType, variant=variant, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0, b=src1, **acc2_args, \ mxsa=srcMX0, mxsb=srcMX1, block=block, \ mxScaleASel=scaleSelSrc0, mxScaleBSel=scaleSelSrc1, \ comment="left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) else: imod.add(MFMAInstruction(instType=miInInstType, accType=miOutInstType, variant=variant, mfma1k=mfma_1k, \ acc=self.accVgprReadWriteIndex(kernel, (accStart+accStoreCIdx), (accEnd-accStart+1)), \ a=src0, b=src1, **acc2_args, neg=neg_flag,\ comment="left value = %s[%u+%u:%u+%u]" % (accumRegType, accStart, accStoreCIdx, accEnd, accStoreCIdx))) prevAccIdx = accIdx if kernel["ExpertSchedulingMode"] > 0: imod.add(SWaitAlu(va_vdst=0, comment="wait for the current iter's writes to complete")) # release register if kReg_first is not None: self.vgprPool.checkIn(kReg_first) if kReg is not None: self.vgprPool.checkIn(kReg) if abReg is not None: self.vgprPool.checkIn(abReg) if mxReg is not None: self.vgprPool.checkIn(mxReg) if tmpVgpr is not None: self.vgprPool.checkIn(tmpVgpr) mfmaMod = Module("mfmaCode") if self.do["MAC"]: shiftK.add(postShiftK) mfmaMod.add(shiftK) mfmaMod.add(imod) return mfmaMod ############################################################################## # At Least 1 Unroll # prefetch means this is in the prefetch code, either before unroll loop # or in the PAP code. # isOptNLL : this is for the store-interleaved NLL optimization ############################################################################## def openSumAtLeastUnroll(self, kernel, prefetch, isOptNLL, isNGLL=False, NLLindex=0, NLLnum=1, tailloopInNll=False): isLongBranch = False if kernel["EnableMatrixInstruction"] and kernel["ProblemType"]["ActivationType"] in ['all', 'hipblaslt_all']: acclen = getAccToArchLen(kernel) # Just a rough calculation if acclen > 100 or (kernel["GlobalSplitUAlgorithm"] == 'MultipleBufferSingleKernel'): isLongBranch = True module = Module("openSumAtLeastUnroll") if prefetch: if not isOptNLL: module.add(self.checkLastIter(kernel)) if kernel["StorePriorityOpt"]: module.add(SSetPrior(prior=0, comment="optimization store")) if self.states.doShadowInit: shadowName = Label.getFormatting("ShadowInitStart") module.add(SCBranchSCC1(labelName=shadowName, \ comment="skip to ShadowInitStart iter b/c numIter==0")) else: # This branch could potentially be very far e.g. > SIMM16 module.addComment1("after InitC, skip to end of prefetch last iter if numIter==0") # use positive offset only long jump if kernel["SuppressNoLoadLoop"]: loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx]] strNtab = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTA0_NTB0" lastIterEnd = Label("LoopEnd%s%s"%(loopChar, strNtab), "") module.add(SCBranchSCC1(labelName=lastIterEnd.getLabelName(), \ comment="skip to unrollLoop end loop%s iter b/c numIter==0" % loopChar)) else: lastIterEnd = Label("PrefetchGlobalLastIterEnd", "") if self.isPrefetchAcrossPersistentEnabled(kernel): module.add(SCMovB32(dst=sgpr("SkPrefetchPrimed"), src=0, comment="discard primed PAP group when current slice skips NLL")) # use positive offset only long jump with self.allocTmpSgpr(3, tag="openSumAtLeastUnroll_tmpSgprInfo") as tmpSgprInfo: module.add(self.longBranchScc1(lastIterEnd, posNeg=1, tmpSgprInfo=tmpSgprInfo)) else: if isOptNLL and NLLindex==0: skipOptNLL = Label("OptNLL_End", "") with self.allocTmpSgpr(4, tag="openSumAtLeastUnroll_tmpSgprInfo2") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx placeHolder="skipOptNLL_placeholder" if self.states.FactorDim == 3 else None module.add(self.checkIsBetaZero(kernel, tmpSgprInfo, skipOptNLL, isLongBranch=isLongBranch, placeHolder=placeHolder, posNeg=1)) # check alpha if self.do["ApplyAlpha"]: # (The new hgemm (h,h,h,h,s,s) is included in ComputeType=Single) if kernel["ProblemType"]["ComputeDataType"].isHalf(): if kernel["ProblemType"]["HighPrecisionAccumulate"] and kernel["StreamK"]: module.add(SCmpEQU32(src0=sgpr("Alpha"), src1=1.0, comment="Alpha == 1.0 ?")) else: # for (h,h,h,h,h,h) no HPA, module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(0x3c003c00), comment="Packed alpha==1.0")) module.add(SCmpEQU32(src0=sgpr("Alpha"), src1=sgpr(tmpSgpr), comment="alpha == 1.0?")) # Shouldn't go here. Currently, DataType=B->ComputeDataType=S # (bf-gemm is included in ComputeType=Single) elif kernel["ProblemType"]["ComputeDataType"].isBFloat16(): module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(0x3f803f80), comment="Packed alpha==1.0")) module.add(SCmpEQU32(src0=sgpr("Alpha"), src1=sgpr(tmpSgpr), comment="alpha == 1.0?")) elif kernel["ProblemType"]["ComputeDataType"].isInt32(): module.add(SCmpEQU32(src0=sgpr("Alpha"), src1=1, comment="Alpha == 1.0 ?")) # This covers sgemm, bfgemm + HPA (b,b,b,b,s,s), and also hgemm (h,h,h,h,s,s) elif kernel["ProblemType"]["ComputeDataType"].isSingle(): module.add(SCmpEQU32(src0=sgpr("Alpha"), src1=1.0, comment="Alpha == 1.0 ?")) elif kernel["ProblemType"]["ComputeDataType"].isDouble(): module.add(SMovB32(dst=sgpr(tmpSgpr+0), src=hex(0x00000000), comment="Low part of double 1.0")) module.add(SMovB32(dst=sgpr(tmpSgpr+1), src=hex(0x3ff00000), comment="High part of double 1.0")) module.add(SCmpEQU64(src0=sgpr("Alpha",2), src1=sgpr(tmpSgpr,2), comment="Alpha == 1.0 ?")) elif kernel["ProblemType"]["ComputeDataType"].isSingleComplex(): module.add(SMovB32(dst=sgpr(tmpSgpr+0), src=1.0, comment="Real part of 1.0")) module.add(SMovB32(dst=sgpr(tmpSgpr+1), src=0.0, comment="Imaginary part of 1.0")) module.add(SCmpEQU64(src0=sgpr("Alpha",2), src1=sgpr(tmpSgpr,2), comment="Alpha == 1.0 ?")) elif kernel["ProblemType"]["ComputeDataType"].isDoubleComplex(): module.add(SMovB32(dst=sgpr(tmpSgpr+0), src=hex(0x00000000), comment="lsb of real part of 1.0")) module.add(SMovB32(dst=sgpr(tmpSgpr+1), src=hex(0x3ff00000), comment="msb of real part of 1.0")) module.add(SCmpEQU64(src0=sgpr("Alpha",2), src1=sgpr(tmpSgpr,2), comment="Alpha.real == 1.0 ?")) if placeHolder == None: if isLongBranch: module.add(self.longBranchScc0(skipOptNLL, posNeg=1, tmpSgprInfo=tmpSgprInfo, comment="branch if alpha != 1")) else: module.add(SCBranchSCC0(labelName=skipOptNLL.getLabelName(), comment="branch if alpha != 1")) else: skipOptNLLModule = Module("skipOptNLL_placeholder") skipOptNLLModule.addComment1("branch if alpha.real != 1") module.add(skipOptNLLModule) module.add(SMovB32(dst=sgpr(tmpSgpr+0), src=hex(0x00000000), comment="lsb of imag part of 0.0")) module.add(SMovB32(dst=sgpr(tmpSgpr+1), src=hex(0x00000000), comment="msb of imag part of 0.0")) module.add(SCmpEQU64(src0=sgpr("Alpha+2",2), src1=sgpr(tmpSgpr,2), comment="Alpha.imag == 0.0 ?")) if placeHolder == None: if isLongBranch: module.add(self.longBranchScc0(skipOptNLL, posNeg=1, tmpSgprInfo=tmpSgprInfo, comment="branch if alpha != 1")) else: module.add(SCBranchSCC0(labelName=skipOptNLL.getLabelName(), comment="branch if alpha != 1")) else: skipOptNLLModule = Module("skipOptNLL_placeholder") skipOptNLLModule.addComment1("branch if alpha != 1") module.add(skipOptNLLModule) module.addSpaceLine() placeHolder = "skipOptNLL_scc1_placeholder" if self.states.FactorDim == 3 else None module.add(self.checkIsEdge(kernel, tmpSgprInfo, skipOptNLL, kernel["MacroTile0"], isLongBranch=isLongBranch, placeHolder=placeHolder)) module.add(self.checkIsEdge(kernel, tmpSgprInfo, skipOptNLL, kernel["MacroTile1"], isSize1=True, isLongBranch=isLongBranch, placeHolder=placeHolder)) module.addSpaceLine() # Check tail loop required: # Skip tail loop check if noTailLoop is true (except for tailloopInNll) if not kernel["NoTailLoop"] or self.states.tailloopInNll: loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx]] module.add(scalarStaticDivideAndRemainder(tmpSgpr, tmpSgpr+1, "SizesSum+%u"%self.states.unrollIdx, \ kernel["DepthU"], ContinuousRegister(tmpSgpr+2, 2), 2)) module.add(SCmpEQU32(src0=sgpr(tmpSgpr+1), src1=0, comment="numIter%s == 0"%loopChar )) if placeHolder == None: if isLongBranch: module.add(self.longBranchScc0(skipOptNLL, posNeg=1, tmpSgprInfo=tmpSgprInfo, comment="skip if tail loop required")) else: module.add(SCBranchSCC0(labelName=skipOptNLL.getLabelName(), comment="skip if tail loop required")) else: skipOptNLLModule = Module("skipOptNLL_placeholder") skipOptNLLModule.addComment1("skip if tail loop required") module.add(skipOptNLLModule) # save the vgprPool for generating the normal path. # so we can avoid explicit vgpr check-in/out self.savedVgprPool = deepcopy(self.vgprPool) self.savedSgprPool = deepcopy(self.sgprPool) # comment out the following codes that attempt to reduce vgpr consumption # however, the kernel vgpr count is governed by peak vgpr consumption so saving # a few here shouldn't affect kernel's overall vgpr consumption. # the following code is for reference and will be removed in the future """ added = [] # track registers added to pool if kernel["PrefetchGlobalRead"]: if not kernel["DirectToLdsA"]: added.append(self.vgprPool.addRange(self.states.a.startVgprG2L, \ self.states.a.startVgprG2L+self.states.a.numVgprG2L-1, "startOptNLL")) added.append(self.vgprPool.addRange(self.states.a.startVgprLocalWriteAddr, \ self.states.a.startVgprLocalWriteAddr, "startOptNLL")) if not kernel["DirectToLdsB"]: added.append(self.vgprPool.addRange(self.states.b.startVgprG2L, \ self.states.b.startVgprG2L+self.states.b.numVgprG2L-1, "startOptNLL")) added.append(self.vgprPool.addRange(self.states.b.startVgprLocalWriteAddr, \ self.states.b.startVgprLocalWriteAddr, "startOptNLL")) if kernel["BufferLoad"]: added.append(self.vgprPool.addRange(self.startVgprGlobalReadOffsetA, \ self.startVgprGlobalReadOffsetB, "startOptNLL")) else: added.append(self.vgprPool.addRange(self.startVgprGlobalReadAddressesA, \ self.startVgprGlobalReadAddressesB, "startOptNLL")) module.addComment1("reclaim VGPRS: " + ", ".join(added)) """ if (not isNGLL) and NLLnum == 2: OptOrOrd = "Opt" if isOptNLL else "Ord" if tailloopInNll: # tailloopInNll case, add "_TI_" to avoid duplicated label OptOrOrd += "_TI_" loopLabel2ndNLL = Label("%sNLL_second"%(OptOrOrd), "second %s NoLoadLoop entry"%OptOrOrd ) # NLL + double buffer (NLLnum==2) case (PGR1/2), we need to generate 2 NLL (first and second buffer) if NLLindex == 0: # first NLL, jump to second code if OrigLoopCounter is odd module.add(SBitcmp1B32(src0=sgpr("OrigLoopCounter"), src1=0, comment="test if OrigLoopCounter is Odd ?")) module.add(SCBranchSCC1(labelName=loopLabel2ndNLL.getLabelName(), comment="jump to second NoLoadLoop" )) else: # NLLindex==1 module.add(loopLabel2ndNLL) return module ############################################################################## def closeSumAtLeastUnroll(self, kernel, tPA, tPB, prefetch, isOptNLL, isNGLL, isNotLast=False, tailloopInNll=False, remainPgr=0): module = Module("closeSumAtLeastUnroll") if not prefetch: if isNGLL: if kernel["ProblemType"]["Sparse"] and kernel["PrefetchGlobalRead"] >= 2 and kernel["DirectToVgprSparseMetadata"]: for i in range(0, self.states.m.numVgprValuPerBlock): module.add(VMovB32(vgpr("ValuMetadata+%u"%i), vgpr("ValuMetadata+%u+%u"%(self.states.m.numVgprValuPerBlock, i)), \ comment="copy ValuMetadata blk1 to blk0")) loopIdx = self.states.unrollIdx loopChar = self.states.indexChars[kernel["ProblemType"]["IndicesSummation"][loopIdx]] loopCounter = self.loopCounter(kernel, loopIdx) if (kernel["DirectToVgprA"] or kernel["DirectToVgprB"]) or (kernel["PrefetchGlobalRead"] >= 3 and remainPgr >= 2): # PGR2 + DTVA/B case, we need to decrement loop counter after NGLL # This needs to be before toPGR1 label to avoid decrementing loop counter in loopCounter==1 case # Also, we need to decrement loop counter for PGR>=3 and remainPgr>=2 for early exit in NGLL module.add(SSubU32(dst=loopCounter, src0=loopCounter, src1=1, comment="dec counter%s"%(loopChar))) toPGR1 = Label(self.labels.getName("toPGR1"), "") if kernel["PrefetchGlobalRead"]-1 > remainPgr >= 2: # generate early exit in NGLL (no need for the first NGLL(remainPgr==PGR-1)) module.add(SCmpEQU32( src0=loopCounter, \ src1=hex(1), \ comment="LoopCounter%s == 1"%(loopChar) )) module.add(SCBranchSCC1(labelName=toPGR1.getLabelName(), comment="PGR>=3 in NGLL. Only 1 loop, toPGR1")) elif remainPgr <= 1: if isNotLast: module.add(SBranch(labelName=toPGR1.getLabelName(), comment="Branch to toPGR1")) else: module.add(toPGR1) else: suffix = "OptNLL" if isOptNLL else "OrdNLL" toPGR1end = Label(self.labels.getName("toPGR1end_%s"%suffix), "") if isNotLast: module.add(SBranch(labelName=toPGR1end.getLabelName(), comment="Branch to toPGR1end")) else: if not tailloopInNll: # generate toPGR1end label except for tailloopInNll case module.add(toPGR1end) if isOptNLL: endSumLabel = "Summation_End_OptNLL" module.addComment0("Stores for OptNLL") module.add(self.endSummation(kernel, tPA, tPB, False, endSumLabel, isOptNLL)) # dot2: WaveSplitK reduction for opt NLL if kernel["NumWaveSplitK"] > 1: module.add(self.waveSplitKReduction(kernel)) # perhaps could work with LSU>1 by adding other indices here, but not tested assert (kernel["LocalSplitU"] == 1) module.add(self.notLocalSplitUGlobalWriteIndices(kernel)) # add stores for opt NLL (fullVws, elements, fullVws_1, elements_1) = self.notLocalFullTileElements(kernel) alpha = False beta = False storeModule, _ = self.globalWriteElements(kernel, tPA, tPB, [fullVws[0]], [fullVws_1[0]], [elements[0]], [elements_1[0]], True, applyAlpha=alpha, betas=[beta], edge=False) module.add(storeModule) self.cleanupGlobalWrite(kernel) module.addSpaceLine() module.add(self.functionEnd(kernel, addLabel=False)) module.add(Label("OptNLL_End", "")) else: if not tailloopInNll: # generate PrefetchGlobalLastIterEnd label except for tailloopInNll case module.add(Label("PrefetchGlobalLastIterEnd", "")) # swap back vgpr pool if any def swapBackGprPool(gprPool, savedGprPool, isNotLast): if savedGprPool != None and (not isNotLast): # in case pool size in current path is larger than pool size in main path # and it will miss allocate vgpr since allocating vgpr is based on pool size in main path newSize = gprPool.size() savedGprPool.appendPool(newSize) return savedGprPool, None # restore vgprPool before alternate path return gprPool, savedGprPool self.vgprPool, self.savedVgprPool = swapBackGprPool(self.vgprPool, self.savedVgprPool, isNotLast) self.sgprPool, self.savedSgprPool = swapBackGprPool(self.sgprPool, self.savedSgprPool, isNotLast) return module ############################################################################## # incLower must be constant or SGPR unsigned value def incrementSrd(self, tP, incLower, incUpper): imod = Module("incrementSrd") tc = tP["tensorChar"] use64bShadowLimit = self.states.use64bShadowLimitMX if tc in ["MXSA", "MXSB"] else self.states.use64bShadowLimit imod.add(SAddU32(dst=sgpr("Srd%s+0"%(tc)), \ src0=sgpr("Srd%s+0"%(tc)), \ src1=incLower, \ comment="gra SRD += inc(lower)" )) imod.add(SAddCU32(dst=sgpr("Srd%s+1"%(tc)), \ src0=sgpr("Srd%s+1"%(tc)), \ src1=incUpper, \ comment="gra SRD += inc(upper)" )) # also have to move the boundary since we change the base # so less buffers to the edge: if use64bShadowLimit: imod.add(SSubU32(dst=sgpr("ShadowLimit%s+0"%tc), \ src0=sgpr("ShadowLimit%s+0"%tc), \ src1=incLower, \ comment="limit -= inc)")) imod.add(SSubBU32(dst=sgpr("ShadowLimit%s+1"%tc), \ src0=sgpr("ShadowLimit%s+1"%tc), \ src1=incUpper, \ comment="limit -= inc)" )) imod.add(SCmpEQU32(src0=sgpr("ShadowLimit%s+1"%tc), src1=0, comment="are we within 2^32?")) if self.states.staggerUCode: # staggerU case, need to restore BufferLimit when ShadowLimit goes to negative value imod.add(SCSelectB32(dst=sgpr("Srd%s+2"%tc), src0=sgpr("ShadowLimit%s+0"%tc), src1="BufferLimit", comment="Move shadow to real if we are within 2^32")) else: imod.add(SCMovB32(dst=sgpr("Srd%s+2"%tc), src=sgpr("ShadowLimit%s+0"%tc), comment="Move shadow to real if we are within 2^32")) imod.add(self.shiftSrd(tc)) else: imod.add(SSubU32(dst=sgpr("Srd%s+2"%(tc)), \ src0=sgpr("Srd%s+2"%(tc)), \ src1=incLower, \ comment="limit -= inc)" )) return imod def incrementMetadataSrd(self, incSparseLower, incSparseUpper): imod = Module("incrementMetadataSrd") imod.add(SAddU32(sgpr("SrdMetadata+0"), \ sgpr("SrdMetadata+0"), \ incSparseLower, \ "gra SRD += incSparse(lower)")) imod.add(SAddCU32(sgpr("SrdMetadata+1"), \ sgpr("SrdMetadata+1"), \ incSparseUpper, \ "gra SRD += incSparse(uppper)" )) # also have to move the boundary since we change the base # so less buffers to the edge: if self.states.use64bShadowLimit: imod.add(SSubU32(sgpr("ShadowLimitMetadata+0"), \ sgpr("ShadowLimitMetadata+0"), \ incSparseLower, \ "limit -= incSparse(lower)")) imod.add(SSubBU32(sgpr("ShadowLimitMetadata+1"), \ sgpr("ShadowLimitMetadata+1"), \ incSparseUpper, \ "limit -= incSparse(uppper)" )) imod.add(SCmpEQU32(sgpr("ShadowLimitMetadata+1"), 0, "are we within 2^32?")) if 1: # self.states.staggerUCode: # staggerU case, need to restore BufferLimit when ShadowLimit goes to negative value imod.add(SCSelectB32(dst=sgpr("SrdMetadata+2"), src0=sgpr("ShadowLimitMetadata+0"), src1="BufferLimit", comment="Move shadow to real if we are within 2^32")) else: imod.add(SCMovB32(sgpr("SrdMetadata+2"), sgpr("ShadowLimitMetadata+0"), "Move shadow to real if we are within 2^32")) else: imod.addInst(SSubU32(sgpr("SrdMetadata+2"), \ sgpr("SrdMetadata+2"), \ incSparseLower, \ "limit -= inc)" )) return imod ############################################################################## # incLower must be constant or SGPR unsigned value def setTailSrd(self, tP, incLower): # In SuppressNoLoadLoop, the final loop iteration moves the SRD base forward # and the ShadowLimit backwards by one extra 'click' of GlobalReadIncs[AB]. # Note the ShadowLimit may become negative - for example edge tiles where the # increment is > tile width. # The SuppressNoLoadLoop mode also forces the SRD limit to 0 on the final iteration. # The code here undoes the final click step by moving the base backwards and the # limit forwards (reading from the ShadowLimit). # It only works if use64bShadowLimit is enabled (since this enables use of the ShadowLimit) tc = tP["tensorChar"] module = Module("setTailSrd") incUpper = 0 module.add(SSubU32(dst=sgpr("Srd%s+0"%(tc)), \ src0=sgpr("Srd%s+0"%(tc)), src1=incLower, \ comment="gra SRD -= inc(lower)" )) module.add(SSubBU32(dst=sgpr("Srd%s+1"%(tc)), \ src0=sgpr("Srd%s+1"%(tc)), src1=incUpper, \ comment="gra SRD -= inc(upper)" )) # using Shadow limit here which only works with 64-bit PBC: use64bShadowLimit = self.states.use64bShadowLimitMX if tc in ["MXSA", "MXSB"] else self.states.use64bShadowLimit assert(use64bShadowLimit) module.add(SAddU32(dst=sgpr("ShadowLimit%s+0"%tc), \ src0=sgpr("ShadowLimit%s+0"%tc), src1=incLower, \ comment="limit -= inc)")) module.add(SAddCU32(dst=sgpr("ShadowLimit%s+1"%tc), \ src0=sgpr("ShadowLimit%s+1"%tc), src1=incUpper, \ comment="limit -= inc)" )) module.add(SCmpEQU32(src0=sgpr("ShadowLimit%s+1"%tc), src1=0, comment="are we within 2^32?")) module.add(SCMovB32(dst=sgpr("Srd%s+2"%tc), src=sgpr("ShadowLimit%s+0"%tc), comment="Move shadow to real if we are within 2^32")) return module ############################################################################## # Global Read: Increment A/B # loopIdx is summation idx: # self.states.unrollIdx, or an idx from 0..NumIndicesSummation # prefetchIndex is >0 (1...PrefetchGlobalRead) if this increment follows a # global prefetch or 0 otherwise ############################################################################## def globalReadIncrement(self, kernel, imod, loopIdx, tP, prefetchIndex): if not self.do["GlobalInc"]: return "" tc = tP["tensorChar"] loopChar = self.states.indexChars[ \ kernel["ProblemType"]["IndicesSummation"][loopIdx]] imod.addComment1("global read inc %s loop%s"%(tc,loopChar)) if kernel["BufferLoad"]: # TODO - does this handle N-dim tensors correctly? #if tP["isB"]: # module.add(SMovB32(dst=sgpr("OffsetB"), src=sgpr("SrdB+0"), comment="hack to save")) if loopIdx == self.states.unrollIdx and self.states.staggerUCode: # add a wrap increment, if needed: with self.allocTmpSgpr(4, tag="globalReadIncrement_tmpSgprInfo") as tmpSgprInfo: incLower = tmpSgprInfo.idx incUpper = incLower + 1 tmpS = incLower + 2 tmpIncSparse = incLower + 3 suStr = "StaggerUIter" tcOther = tc.replace("A", "B") if ("A" in tc) else tc.replace("B", "A") if kernel["PrefetchGlobalRead"] >= 2 and (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc] and (not kernel["DirectToVgpr%s"%tcOther]): suStr += "DTV" if prefetchIndex: imod.add(SAddU32(dst=sgpr(tmpS), src0=self.loopCounter(kernel, self.states.unrollIdx), src1=prefetchIndex, comment="remove pf(%u)"%prefetchIndex)) imod.add(SCmpEQU32(src0=sgpr(suStr), src1=sgpr(tmpS), comment="Is this wrapIter? (pf)")) else: imod.add(SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=sgpr(suStr), comment="Is this the wrapIter?")) imod.add(SCSelectB32(dst=sgpr(incLower), src0=sgpr("WrapU%s+0"%tc), src1=sgpr("GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx)), \ comment="incLower <- ?")) imod.add(SCSelectB32(dst=sgpr(incUpper), src0=sgpr("WrapU%s+1"%tc), src1=0, comment="incUpper <- ?")) imod.addModuleAsFlatItems(self.incrementSrd(tP, sgpr(incLower), sgpr(incUpper))) if kernel["ProblemType"]["Sparse"]: if (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) : tc = "Metadata" if kernel["DirectToVgprSparseMetadata"]: incSparse = tmpIncSparse imod.addModuleAsFlatItems(self.calculateIncrementMetadata(kernel, incSparse)) else: incSparse = "GlobalReadIncs%s+%u"%(tc,self.states.unrollIdx) if prefetchIndex: imod.add(SCmpEQU32(src0=sgpr("StaggerUIter"), src1=sgpr(tmpS), comment="Is this wrapIter? (pf)")) else: imod.add(SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=sgpr("StaggerUIter"), comment="Is this the wrapIter?")) imod.add(SCSelectB32(dst=sgpr(incLower), src0=sgpr("WrapU%s+0"%tc), src1=sgpr(incSparse), \ comment="incLower <- ?")) imod.add(SCSelectB32(dst=sgpr(incUpper), src0=sgpr("WrapU%s+1"%tc), src1=0, comment="incUpper <- ?")) if kernel["DirectToVgprSparseMetadata"]: imod.addModuleAsFlatItems(self.incrementMetadataSrd(sgpr(incLower), sgpr(incUpper))) else: imod.addModuleAsFlatItems(self.incrementSrd(tP["tpsMetadata"], sgpr(incLower), sgpr(incUpper))) else: grIncTcs = ('A', 'B', 'MXSA', 'MXSB') if loopIdx != self.states.unrollIdx or (tc in grIncTcs and kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx] in kernel["ProblemType"]["MirrorDims%s"%tc]): with self.allocTmpSgpr(1, tag="globalReadIncrement_tmpSgprInfo2") as tmpSgprInfo: incUpper = tmpSgprInfo.idx # GRO may be negative for other summation if stride-other < stride-unroll or if mirror dim. imod.add(SAShiftRightI32(dst=sgpr(incUpper), shiftHex=31, src=sgpr("GlobalReadIncs%s+%u"%(tc,loopIdx)), comment="sign-extend")) imod.addModuleAsFlatItems(self.incrementSrd(tP, sgpr("GlobalReadIncs%s+%u"%(tc,loopIdx)), sgpr(incUpper))) else: incUpper = 0 # GRO is positive for loop unroll srcGRInc = sgpr("GlobalReadIncs%s+%u"%(tc,loopIdx)) if tc in grIncTcs and self.useConstSgprGlobalReadIncsForTc(tc): srcGRInc = "GlobalReadIncs%s"%tc imod.addModuleAsFlatItems(self.incrementSrd(tP, srcGRInc, hex(incUpper))) if kernel["ProblemType"]["Sparse"]: if (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) : tc = "Metadata" imod.addComment1("global read inc metadata loop%s"%(loopChar)) if kernel["DirectToVgprSparseMetadata"]: with self.allocTmpSgpr(1, tag="globalReadIncrement_tmpSgprInfo3") as tmpSgprInfo: incSparse = tmpSgprInfo.idx imod.addModuleAsFlatItems(self.calculateIncrementMetadata(kernel, incSparse)) imod.addModuleAsFlatItems(self.incrementMetadataSrd(sgpr(incSparse), 0)) else: incSparse = "GlobalReadIncs%s+%u"%(tc,loopIdx) if loopIdx != self.states.unrollIdx or (kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx] in kernel["ProblemType"]["MirrorDims%s"%tc]): with self.allocTmpSgpr(1, tag="globalReadIncrement_tmpSgprInfo4") as tmpSgprInfo: incUpper = tmpSgprInfo.idx # GRO may be negative for other summation if stride-other < stride-unroll or if mirror dim. imod.add(SAShiftRightI32(dst=sgpr(incUpper), shiftHex=31, src=sgpr(incSparse), comment="sign-extend")) imod.addModuleAsFlatItems(self.incrementMetadataSrd(sgpr(incSparse), sgpr(incUpper))) else: incUpper = 0 # GRO is positive for loop unroll imod.addModuleAsFlatItems(self.incrementSrd(tP["tpsMetadata"], sgpr(incSparse), hex(incUpper))) else: graIdx = 0 for _ in range(0, tP["nrp"]): for _ in range(0, tP["nrpv"]): for _ in range(0, tP["nrc"]): for _ in range(0, tP["nrcv"]//tP["nrcvpi"]): if self.states.globalReadIncsUseVgpr: imod.add(VAddCOU32( \ dst=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ src1=vgpr("GlobalReadIncs%s+%u+0"%(tP["tensorChar"], 2*loopIdx)), \ comment="gra += inc%s%s (lower)"%(tP["tensorChar"], loopChar))) imod.add(VAddCCOU32( \ dst=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ src1=vgpr("GlobalReadIncs%s+%u+1"%(tP["tensorChar"], 2*loopIdx)), \ src2=VCC(), \ comment="gra += inc%s%s (upper)"%(tP["tensorChar"], loopChar))) else: imod.add(VAddCOU32( \ dst=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ src1=sgpr("GlobalReadIncs%s+%u"%(tP["tensorChar"], loopIdx)), \ comment="gra += inc%s%s (lower)"%(tP["tensorChar"], loopChar))) imod.add(VAddCCOU32( \ dst=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ src1=0, src2=VCC(), \ comment="gra += inc%s%s (upper)"%(tP["tensorChar"], loopChar))) graIdx += self.states.rpga def globalReadIncrementAB(self, kernel, tPA, tPB, loopIdx, prefetchIndex): imod = Module("globalReadIncrementAB") tdmA: bool = kernel["enableTDMA"] tdmB: bool = kernel["enableTDMB"] tdmMetadata: bool = kernel["enableTDMMetadata"] # TODO: check correctness of non-MI kernels. numWaves: int = kernel["NumWaves"] #TDM Wave Separated if tdmA and tdmB and tPA and tPB and numWaves > 1: incCodeA = imod.add(Module("globalReadIncrementA")) if kernel["HalfPLR"] and prefetchIndex == 0: incCodeA.add(self.graIncrementMask(kernel, tPA, tPB)) incCodeA.add(self.tdmIncrementABWaveSperated(kernel, tPA, tPB, loopIdx, prefetchIndex)) # makeSchedule expects both increment modules to exist. For non-MX wave-separated # TDM, B is intentionally empty because the A module advances the aliased A/B SRD. incCodeB = imod.add(Module("globalReadIncrementB")) if "MX" in tPA and "MX" in tPB: incCodeB.add(self.tdmIncrementABWaveSperated(kernel, tPA["MX"], tPB["MX"], loopIdx, prefetchIndex)) return imod incCodeA = imod.add(Module("globalReadIncrementA")) if tPA != None: if not tdmA: self.globalReadIncrement(kernel, incCodeA, loopIdx, tPA, prefetchIndex) else: if kernel["HalfPLR"] and prefetchIndex == 0: incCodeA.add(self.graIncrementMask(kernel, tPA, tPB)) incCodeA.add(self.tdmIncrementAB(kernel, tPA, loopIdx, prefetchIndex)) if "MX" in tPA and not tdmA: self.globalReadIncrement(kernel, incCodeA, loopIdx, tPA["MX"], prefetchIndex) incCodeB = imod.add(Module("globalReadIncrementB")) if tPB != None: if not tdmB: self.globalReadIncrement(kernel, incCodeB, loopIdx, tPB, prefetchIndex) else: incCodeB.add(self.tdmIncrementAB(kernel, tPB, loopIdx, prefetchIndex)) if "MX" in tPB and not tdmB: self.globalReadIncrement(kernel, incCodeB, loopIdx, tPB["MX"], prefetchIndex) return imod ############################################################################## # Global Read: # globalReadTrueGuardK is called for loads in the tail loop # Must ensure each load is in bounds - either using buffer bounds # or exec-mask checks. ############################################################################## def globalReadGuardK(self, kernel, tP, doTailOpt = 0, optParams = None): module = Module("globalReadGuardK") tc = tP["tensorChar"] problemType = self.states.kernel["ProblemType"] numElementsPer4Bytes = int(4 / tP["bpeGR"]) isTr = (tc == "A" or tc == "B") and kernel["enableGLTr%s"%tc] # For example, it replaces buffer_load_b128 with 4 buffer_load_b32 if HasPartialOOB is false. tailGRVW = kernel["GlobalReadVectorWidth%s"%tc] if self.states.asmCaps["HasPartialOOB"] else 1 ######################################## # Calculate Max Addr ######################################## if isTr: # DirectToVgpr case, we need to calculate max address module.addComment1("Max read address offset for GLTr%s"%tc) maxGroVgpr = self.vgprPool.checkOut(1, tag="globalReadGuardK_maxGroVgpr") tmpVgpr = self.vgprPool.checkOutAligned(2, 2, tag="globalReadGuardK_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) tmp = self.vgprPool.checkOut(1, tag="globalReadGuardK_tmp") tmp2 = self.vgprPool.checkOut(1, tag="globalReadGuardK_tmp2") WvG_M = kernel["MIWaveGroup"][0] numKr = kernel["MatrixInstK"] // tP["glvw"] module.addComment0("calc last tile offset") module.add(vectorStaticDivide(maxGroVgpr, "Serial", kernel["WavefrontSize"], tmpVgprRes)) if tP["isA"]: module.add(VAndB32(dst=vgpr(maxGroVgpr), src0=hex(WvG_M-1), src1=vgpr(maxGroVgpr), comment="GLTr%s: wave_id (along_M) mod MIWG[0]"%tc)) module.add(VMulU32U24(dst=vgpr(maxGroVgpr), src0=numKr, src1=vgpr(maxGroVgpr), comment="GLTr%s: wave_id (along_M) *= numKr"%tc)) elif tP["isB"]: # NB: # Calc of w_id is: /= MIWG[0], not %= MIWG[1] module.add(VLShiftRightB32(dst=vgpr(maxGroVgpr), shiftHex=log2(WvG_M), src=vgpr(maxGroVgpr), comment="GLTr%s: wave_id (along_N) /= MIWG[0]"%tc)) module.add(VMulU32U24(dst=vgpr(maxGroVgpr), src0=numKr, src1=vgpr(maxGroVgpr), comment="GLTr%s: wave_id (along_N) *= numKr"%tc)) module.add(VBfeU32(dst=vgpr(tmp2), src0=vgpr("Serial"), src1=int(tP["bpeGR"])+1, src2=1, comment="GLTr%s: offset for the right half of the tile"%(tc))) module.add(VAddU32(dst=vgpr(maxGroVgpr), src0=vgpr(tmp2), src1=vgpr(maxGroVgpr), comment="GLTr%s: wave_id += offset for the right half of the tile"%(tc))) with self.allocTmpSgpr(1, tag="globalReadGuardK_tmpSgprInfo") as tmpSgprInfo: if tP["glvw"] > 1: if tP["tlu"]: module.add(vectorStaticMultiply(vgpr(maxGroVgpr), vgpr(maxGroVgpr), tP["glvw"], tmpSgprInfo, comment="GLTr%s: tile * glvw(%u)"%(tc, tP["glvw"]))) else: module.add(vectorStaticMultiply(vgpr(maxGroVgpr), vgpr(maxGroVgpr), tP["glvw"], tmpSgprInfo, comment="GLTr%s: unroll * glvw(%u)"%(tc, tP["glvw"]))) strideIdx = tP["lsc"] if tP["tlu"] else tP["lsp"] stride = kernel[strideIdx] module.add(VAddU32(dst=vgpr(maxGroVgpr), src0=vgpr(maxGroVgpr), src1=stride*(tP["nrt"]-1))) module.addComment0("calc last unroll offset") module.add(VMovB32(dst=vgpr(tmp), src=sgpr("SizesSum+%u"%self.states.unrollIdx))) with self.allocTmpSgpr(1, tag="globalReadGuardK_tmpSgprInfo2") as tmpSgprInfo: module.add(vectorStaticRemainder(tmp, tmp2, tmp, kernel["DepthU"], tmpVgprRes, tmpSgprInfo)) module.addComment0("final offset") module.add(VSubU32(dst=vgpr(tmp2), src0=vgpr(tmp2), src1=1, comment="GLTr%s: unroll idx - 1"%(tc))) bfArgs = (maxGroVgpr, maxGroVgpr, tmp2, tmp) module.add(MacroInstruction(name="GLOBAL_OFFSET_%s"%tc, args=bfArgs)) module.add(vectorMultiplyBpe(maxGroVgpr, maxGroVgpr, tP["bpeGR"])) module.addSpaceLine() self.vgprPool.checkIn(tmp) self.vgprPool.checkIn(tmp2) self.vgprPool.checkIn(tmpVgpr) elif not kernel["BufferLoad"]: with self.allocTmpSgpr(2, tag="globalReadGuardK_tmpSgprInfo3") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx maxAddrSgpr = tmpSgpr module.addComment0("flat addressing - max read address = size[n] * stride[n-1]") dim = len(tP["ia"])-1 # dim sizeIdx = tP["ia"][dim] sizeIdxIsSum = sizeIdx in kernel["ProblemType"]["IndicesSummation"] if sizeIdxIsSum: sizeIdx -= kernel["ProblemType"]["NumIndicesC"] # TODO-multiply by largest stride module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(maxAddrSgpr+0), sgpr(maxAddrSgpr+1), \ sgpr("Sizes%s+%u"%("Sum" if sizeIdxIsSum else "Free", sizeIdx)), \ sgpr("Stride%s%s"%(tc, self.states.indexChars[tP['ia'][-1]])), \ comment="64b tensor%s size in elements"%tc)) module.add(scalarMultiply64Bpe(maxAddrSgpr, maxAddrSgpr, tP["bpeGR"], tmpSgpr, comment="<- tensor%s size in bytes"%tc)) module.add(SAddU32( dst=sgpr(maxAddrSgpr+0), \ src0=sgpr(self.sgprs["AddressA"] if tP["isA"] else self.sgprs["AddressB"]), \ src1=sgpr(maxAddrSgpr+0), \ comment="prepend address lower")) module.add(SAddCU32( dst=sgpr(maxAddrSgpr+1), \ src0=sgpr((self.sgprs["AddressA"] if tP["isA"] else self.sgprs["AddressB"])+1), \ src1=sgpr(maxAddrSgpr+1), \ comment="prepend address upper")) # sgpr->vgpr maxAddrVgpr = self.vgprPool.checkOutAligned(2, 2, "maxAddrVgpr") module.add(VMovB32(dst=vgpr(maxAddrVgpr+0), src=sgpr(maxAddrSgpr+0), comment="sgpr->vgpr")) module.add(VMovB32(dst=vgpr(maxAddrVgpr+1), src=sgpr(maxAddrSgpr+1), comment="sgpr->vgpr")) # full exec mask fullExec = tmpSgpr sgprCnt = self.states.laneSGPRCount waveSize = kernel["WavefrontSize"] activeMask = "0xFFFFFFFF" if (waveSize == 32) else "0xFFFFFFFFFFFFFFFF" SMovBX = SMovB32 if (waveSize == 32) else SMovB64 module.add(SMovBX(dst=sgpr(fullExec,sgprCnt), src=activeMask, comment="to restore all threads active")) bpeVgpr = self.vgprPool.checkOut(1, "bpeVgpr") module.add(VMovB32(dst=vgpr(bpeVgpr), src=int(tP["bpeGR"]), comment="bpeGR")) # can remove this? zeroVgpr = self.vgprPool.checkOut(1,"zeroVgpr") module.add(VMovB32(dst=vgpr(zeroVgpr), src=0, comment="zero")) def globalReadGuardKBody(tP, optParams = None): if optParams != None: jumpLabel = optParams.jumpLabel idx = optParams.idx tmpVgpr = optParams.tmpVgpr kLabelsList = optParams.kLabelsList behavior = optParams.behavior kSgpr = optParams.kSgpr vgprSlot = optParams.vgprSlot periodParam = optParams.periodParam loadNum = optParams.loadNum preDirectToLdsLoads = optParams.preDirectToLdsLoads finalLoop = optParams.finalLoop else: jumpLabel = None idx = 0 tmpVgpr = None kLabelsList = [] behavior = "" kSgpr = None vgprSlot = [] periodParam = [] loadNum = 0 preDirectToLdsLoads = 0 finalLoop = 0 tc = tP["tensorChar"] self.vgprs.globalReadRegisters[tc] = [] tcDataType = "" if tc == "Metadata" else tc graIdx = 0 g2lIdx = 0 loadWidth = tP["globalReadInstruction"].totalWidth # FIXME: Don't know why for grvw == 1, need further investigate glvwWorkaround = 8 * kernel["ProblemType"]["MacDataType%s"%tcDataType if (tcDataType=='A' or tcDataType=='B') else "DataType"].numRegisters() dataType = kernel["ProblemType"]["DataType%s"%tcDataType] if ("MXS" not in tcDataType) and tP["glvw"] < glvwWorkaround: dataType = kernel["ProblemType"]["MacDataType%s"%tcDataType if (tcDataType=='A' or tcDataType=='B') else "DataType"] isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal( self.states.asmCaps, tP["NonTemporal"], TemporalHint(tP["TemporalHint"]), _nonVolatile(kernel, tc)) isLds = True if (kernel["DirectToLds%s"%tc] and not kernel["NonDTLTailLoop%s"%tc]) else False isTr = (tc == "A" or tc == "B") and kernel["enableGLTr%s"%tc] is16b = dataType.isHalf() or dataType.isBFloat16() directToLdsLoads = 0 if doTailOpt == 2 and behavior == "LOAD": directToLdsLoads += preDirectToLdsLoads prevLdsOffset = 0 instOffset = 0 loopCnt = -1 vgprIdx = 0 loadCnt = 0 perpStart = periodParam[0] if doTailOpt == 2 else 0 perpEnd = periodParam[1] if doTailOpt == 2 else tP["nrp"] sPerpStart = periodParam[2] if doTailOpt == 2 else 0 sPerpEnd = periodParam[3] if doTailOpt == 2 else tP["nrpv"] paraStart = periodParam[4] if doTailOpt == 2 else 0 paraEnd = periodParam[5] if doTailOpt == 2 else tP["nrc"] sParaStart = periodParam[6] if doTailOpt == 2 else 0 sParaEnd = periodParam[7] if doTailOpt == 2 else (tP["nrcv"]//tP["nrcvpi"]) rStart = periodParam[8] if doTailOpt == 2 else 0 rEnd = periodParam[9] if doTailOpt == 2 else 0 requiredSgprs = min(self.sgprPool.size(), self.states.regCaps["MaxSgpr"]) maxVgprs, occupancy = self.getMaxRegsForOccupancy(kernel["NumThreads"], self.vgprPool.size(), requiredSgprs, \ self.getLdsSize(kernel), self.agprPool.size(), self.states.doubleVgpr) tmpVgprPool = [] packInst = Module() maxNumTempToUse = max(maxVgprs - self.vgprPool.size(), 1) for perp in range(perpStart, perpEnd): for sPerp in range(sPerpStart, sPerpEnd): for para in range(paraStart, paraEnd): for sPara in range(sParaStart, sParaEnd): i = sPara + (tP["nrcv"] // tP["nrcvpi"]) * (para + tP["nrc"] * (sPerp + tP["nrpv"] * perp)) loopCnt = sPara + para * sParaEnd + sPerp * paraEnd * sParaEnd + perp * sPerpEnd * paraEnd * sParaEnd graIdx = i * self.states.rpgo if kernel["BufferLoad"] else i * self.states.rpga g2lIdx = i * loadWidth * tP["bpeRatio"] if loadWidth == 3: g2lIdx = i * (loadWidth + 1) * tP["bpeRatio"] if loadWidth == 6: g2lIdx = i * (loadWidth + 2) * tP["bpeRatio"] if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc] and kernel["ConvertAfterDS"] and not isTr: # DTV + ConvertAfterDS case, if bpe > bpeGR, we need to shift g2lIdx for conversion if tP["bpe"] > tP["bpeGR"]: g2lIdx *= tP["bpe"] // tP["bpeGR"] destVgprHi = None tmpVgprIdx = tmpVgpr dataIsByte = False packInt8Code = None eccOffset = 0 instOffsetInc = 0 # increment value for instOffset. Need to apply after r loop r = 0 numLoadVectorComp = int(loadWidth*self.states.bpr//tP["bpeGR"]) if (dataType.isDouble() or dataType.isSingle())\ and (kernel["BufferLoad"] or tP["tlu"]): # BufferLoad: SRD limit handles OOB, use same wide instruction as main loop. # Flat + TLU=True: vector spans tile (not K) so all elements share the # same K index. A single exec-mask check guards the whole vector; graShift # keeps the tile-edge in bounds. numLoadVectorComp = numLoadVectorComp // tailGRVW if isLds or isTr: numLoadVectorComp = 1 int8TempVgpr = numLoadVectorComp - 1 # for each component in vector while r < numLoadVectorComp: if doTailOpt == 1 and i == idx and ((r+1) % numElementsPer4Bytes != 0): if kLabelsList != None: module.add(kLabelsList.pop()) module.add(SCmpGeU32(src0=sgpr(kSgpr), src1=(r + 1), comment="")) module.add(SCBranchSCC0(labelName=jumpLabel.getLabelName(), comment="")) numElementsPerLoad = 1 if kernel["ConvertAfterDS"]: dataType = kernel["ProblemType"]["DataType%s"%tcDataType] if dataType.isInt8() or dataType.is8bitFloat() or tP["isM"]: # TODO-Int8, Check this: # if tP["glvw"]>1 and kernel["AssertSummationElementMultiple"] % 2 == 0: # # Pack two FP16 values into a single load dword x2 # numElementsPerLoad = 2 # elif self.states.archCaps["HasEccHalf"]: # destVgprHi = self.vgprPool.checkOut(1, 'destVgprHi') if isTr: numElementsPerLoad = tailGRVW if (not tP["isM"]) and isLds: if not kernel["NonDTLTailLoop%s"%tc]: numElementsPerLoad = tailGRVW dataIsByte = False else: dataIsByte = True # Check out 3 regs once , for component 1,2,3 (r = 1,2,3) if doTailOpt == 1: if r == 1: packInt8Code = Module() destVgprHi = tmpVgprIdx else: if doTailOpt == 2: if r == 1: packInt8Code = Module() if r != 0: destVgprHi = vgprSlot[vgprIdx] if r >= rStart and r < rEnd and (r % 4 != 0): vgprIdx = vgprIdx + 1 elif r == 1: packInt8Code = Module() destVgprHi = self.vgprPool.checkOut( int8TempVgpr , tag="globalReadGuardK_destVgprHi1") regIdx = r // 4 if (tP["localWriteInstruction"].blockWidth <= 0.5) and (r%2 == 0) and not tP["isM"]: numVgprG2L = self.states.a.numVgprG2L if tc == 'A' else self.states.b.numVgprG2L eccBpe = tP["bpeDS"] if kernel["ConvertAfterDS"] else max(tP["bpeGR"], tP["bpe"]) eccOffset = _getEccOffset(tP["globalReadInstruction"].totalWidth, bpr=self.states.bpr, bpe=eccBpe, \ glvw=tP["glvw"], idx=loopCnt, numVgprG2L=numVgprG2L) elif dataType.isHalf() or dataType.isBFloat16(): if isTr: numElementsPerLoad = tailGRVW elif tP["glvw"]>1 and kernel["AssertSummationElementMultiple"] % 2 == 0 and not isLds: # Pack two FP16 values into a single load dword x2 numElementsPerLoad = 2 elif isLds and not kernel["NonDTLTailLoop%s"%tc]: numElementsPerLoad = tailGRVW elif self.states.archCaps["HasEccHalf"]: # In some cards, loading half types into register will zero out # the other half. Therefore we need to load into a separate register # then pack 2 registers into one if doTailOpt == 1: destVgprHi = tmpVgprIdx else: if (tP["localWriteInstruction"].blockWidth == 0.5) and (r%2 == 0): numVgprG2L = self.states.a.numVgprG2L if tc == 'A' else self.states.b.numVgprG2L eccBpe = tP["bpeDS"] if kernel["ConvertAfterDS"] else max(tP["bpeGR"], tP["bpe"]) eccOffset = _getEccOffset(tP["globalReadInstruction"].totalWidth, bpr=self.states.bpr, bpe=eccBpe, \ glvw=tP["glvw"], idx=loopCnt, numVgprG2L=numVgprG2L) else: if doTailOpt == 2: destVgprHi = vgprSlot[vgprIdx] if r >= rStart and r < rEnd and (r % 2 != 0): vgprIdx = vgprIdx + 1 else: #doTailOpt == 0 case if r % 2 == 1: tmpVgprPool.append(self.vgprPool.checkOut(1, tag="globalReadGuardK_destVgprHi2")) destVgprHi = tmpVgprPool[-1] regIdx = r // 2 elif dataType.isFloat4(): numElementsPerLoad = 8 regIdx = r // 8 elif dataType.is6bitFloat(): numElementsPerLoad = 16 regIdx = r // 4 elif dataType.isInt8x4() or dataType.isSingle(): if kernel["BufferLoad"] or tP["tlu"]: numElementsPerLoad = tailGRVW regIdx = r elif dataType.isDouble(): if kernel["BufferLoad"] or tP["tlu"]: numElementsPerLoad = tailGRVW regIdx = r*2 elif dataType.isSingleComplex(): regIdx = r*2 elif dataType.isDoubleComplex() : regIdx = r*4 else: printWarning("DataType unsupported") if doTailOpt == 0: module.addComment0("g2l=%u, load component %u"%(g2lIdx, r)) offset = 0 hi8 = 0 hi16 = 0 bpl = numElementsPerLoad*(tP["bpeGR"] if not tP["isM"] else tP["bpe"]) if kernel["BufferLoad"] or isTr: # Use buffer limit to stay in-bounds - the limit was set to edge when SRD initialized # and each increment of SRD base in the unroll loop does a corresponding decrement # of the srd limit - so base+limit stays constant and also points at maximum # element that should be accessed. if kernel["_UseSgprForGRO"]: offsetVgpr = "GlobalReadOffset%s+0"%(tc) else: offsetVgpr = "GlobalReadOffset%s+%u"%(tc, graIdx) # Vgpr for GRO if not kernel["_UseSgprForGRO"]: soffset = "0" # instruction offset with Sgpr for GRO elif isLds and kernel["UseInstOffsetForGRO"]: soffset = sgpr("ScalarGlobalReadOffset%s+%u"%(tc, graIdx)) # Sgpr for GRO else: soffset = "0" if graIdx == 0 else sgpr("ScalarGlobalReadOffset%s+%u"%(tc, graIdx-1)) unrollMirrorWithSoffset = kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx] in problemType["MirrorDims%s"%tc] and soffset != "0" # ScalarGlobalReadOffset should be negative value with unroll mirroring. # However, buffer_load uses soffset as uint value, so GRO - SGRO, SGRO = 0 if unrollMirrorWithSoffset: codeMod = Module("mirrorIdx%u"%loopCnt) codeMod.add(VSubU32(dst=vgpr(offsetVgpr), src0=vgpr(offsetVgpr), src1=soffset, comment="mirror unroll: GRO=GRO-SGRO, soffset=0")) module.add(codeMod) soffset_prev = soffset soffset = "0" if isLds: # need to increment ldsInc only once per each loopCnt # this is pre count up, so increment it at r == 0 if r == 0: ldsInc = int((self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * tailGRVW * tP["bpeGR"]) else: ldsInc = 0 if kernel["LdsBlockSizePerPad%s"%tc] != 0: ldsInc += int((ldsInc // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeGR"]) else: padInterval = (self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * self.states.bpr ldsInc += int((ldsInc // padInterval) * kernel["LdsPad%s"%tc] * tP["bpeGR"]) if kernel["UseInstOffsetForGRO"]: # buffer_load only support 12 bit instruction offset # we have to increase m0 if offset is larger thant 12 bits if instOffset >= self.buff_load_inst_offset_max: inc = int((instOffset // self.buff_load_inst_offset_max) * self.buff_load_inst_offset_max) module.add(SAddU32(dst=mgpr(0), src0=mgpr(0), src1=inc, comment="Move LDS write address to next base" )) instOffset -= inc elif directToLdsLoads != 0 and ldsInc > 0: if tP["nrc"] > 1: # another address conversion for DirectToLds + NumLoadsCoalesced > 1 divisorName = tP["lvc"] divisor = kernel[divisorName] # DirectToLds + NumLoadsCoalesced>1 case, need to adjust m0 increment value to store values to correct location in LDS wSize = max(self.states.kernel["WavefrontSize"], divisor) lscaOffset = int(para * wSize * tP["bpeGR"] * tP["glvw"]) ldsOffset = int(ldsInc * tP["nrc"] * (sPerp + tP["nrpv"] * perp) + lscaOffset) ldsInc = ldsOffset - prevLdsOffset prevLdsOffset = ldsOffset if (doTailOpt == 0) or (doTailOpt == 2 and behavior == "LOAD") or\ (doTailOpt == 1 and i == idx and ((r+1) % numElementsPer4Bytes != 0) and r != 0): module.add(SAddU32(dst=mgpr(0), src0=mgpr(0), src1=int(ldsInc), comment="Move LDS write address to next line" )) destVgpr=0 self.vgprs.globalReadRegisters[tc].append(0) else: g2lIdxM = i * max(loadWidth * tP["bpeRatio"], 1) destVgpr="G2L%s+%u+%u"%(tc, g2lIdx + tP["shiftGR"] if not tP["isM"] else g2lIdxM, regIdx+eccOffset) self.vgprs.globalReadRegisters[tc].append( (g2lIdx + tP["shiftGR"] if not tP["isM"] else g2lIdxM) + regIdx+eccOffset) offset = int(r * tP["bpeGR"] + instOffset) comment = "load one buffer value" if (dataType.isHalf() or dataType.isBFloat16()) and not tP["isM"]: if numElementsPerLoad==2: # Pack two FP16 values into a single load dword x2 r += 1 # skip next element since we loaded 2X here comment = "load packed 2X half buffer value" elif not isLds: hi16=loopCnt%2 if tP["glvw"]==1 else r%2 comment="load one buffer value" if ((dataType.isInt8() or dataType.is8bitFloat()) \ and not tP["isM"]) or (tP["isM"] and destVgprHi != None): # TODO-Int8, Check this: #if numElementsPerLoad==2: # Pack two FP16 values into a single load dword x2 # r += 1 # skip next element since we loaded 2X here # comment = "load packed 2X half buffer value" if not isLds: hi8 = (loopCnt%4) %2 if tP["glvw"]==1 else (r%4) %2 hi16 = False if tP["glvw"]==1 else (r%4)//2 comment="load one buffer value" if (dataType.isFloat4()) and not tP["isM"]: if numElementsPerLoad==8: # Pack 8 FP4 elements into a single load dword r += numElementsPerLoad-1 # skip next (numElementsPerLoad-1) element since we loaded dword here comment = "Load 8 elements for Float4 in single VGPR." if dataType.is6bitFloat() and not tP["isM"]: r += numElementsPerLoad-1 comment = f"Load {numElementsPerLoad} elements for 6 bits" useBuffer = not isTr # if hi8=1 or hi16=1 (component 1,2,3 for int8) or (component 1 for half), use the temp destVgprHi # but only when hi16=1 we use the _d16_hi version instruction, see the below visualized int8 comment if doTailOpt == 1: loadVgpr = destVgprHi else: loadVgpr = destVgprHi if ((hi16 or hi8) and destVgprHi != None) else destVgpr self.vgprs.globalReadRegisters[tc][-1] = destVgprHi if ((hi16 or hi8) and destVgprHi != None) else self.vgprs.globalReadRegisters[tc][-1] if ((kernel["ProblemType"]["DataType%s"%tcDataType].isInt8() or kernel["ProblemType"]["DataType%s"%tcDataType].is8bitFloat() or tP["isM"]) \ and (not self.states.archCaps["HasEccHalf"]) and useBuffer) \ or (tP["isM"] and self.states.asmCaps["HasSWMMAC_gfx1250"]): module.add(VMovB32(dst=vgpr(loadVgpr), src=0, comment="set to zero to avoid unexpected value")) if isTr: module.add(VMinI32(dst=vgpr(offsetVgpr), src0=vgpr(maxGroVgpr), src1=vgpr(offsetVgpr), comment="truncated load: clamp GRO to legal range")) if (kernel["ProblemType"]["DataType%s"%tcDataType].isHalf() or kernel["ProblemType"]["DataType%s"%tcDataType].isBFloat16()) and not tP["isM"] and self.states.asmCaps["HasWMMA_V3"]: module.add(VMovB32(dst=vgpr(loadVgpr), src=0, comment="set to zero to avoid unexpected value")) if doTailOpt == 1: if behavior == "LOAD" and i == idx: if (numElementsPerLoad == 2 and r % numElementsPer4Bytes != 0) or \ (numElementsPerLoad != 2 and ((r + 1) % numElementsPer4Bytes != 0)): module.addComment0("g2l=%u, load component %u"%(g2lIdx, r)) module.add(self.chooseGlobalRead(useBuffer, \ bpl, destVgpr=loadVgpr, \ addr0=vgpr(offsetVgpr), addr1=sgpr("Srd%s"%tc, 2 if isTr else 4), \ soffset=soffset, offset=offset, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ tr=isTr, hi16=hi16, \ comment=comment, scope=scope, th=th, nv=nv)) tmpVgprIdx += 1 if (numElementsPerLoad == 2 and r == (numLoadVectorComp - 1)) or \ (numElementsPerLoad != 2 and (r + 1) == (numLoadVectorComp - 1)): module.add(SBranch(labelName=jumpLabel.getLabelName(), comment="")) else: if (doTailOpt == 0) or (doTailOpt == 2 and behavior == "LOAD" and r >= rStart and r < rEnd): module.add(self.chooseGlobalRead(useBuffer, \ bpl, destVgpr=loadVgpr, \ addr0=vgpr(offsetVgpr), addr1=sgpr("Srd%s"%tc, 2 if isTr else 4), \ soffset=soffset, offset=offset, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ tr=isTr, hi16=hi16, \ comment=comment, scope=scope, th=th, nv=nv)) loadCnt = loadCnt + 1 if unrollMirrorWithSoffset: codeMod = Module("mirrorIdx%u"%loopCnt) codeMod.add(VAddU32(dst=vgpr(offsetVgpr), src0=vgpr(offsetVgpr), src1=soffset_prev, comment="mirror unroll: restore GRO=GRO+SGRO")) module.add(codeMod) if isLds and kernel["UseInstOffsetForGRO"]: instOffsetInc += ldsInc else: # Not buffer load, ie 'global' load # mask if current address if in bounds module.add(VCmpXLtU64(dst=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u"%(tP["tensorChar"], graIdx),2), \ src1=vgpr(maxAddrVgpr,2), \ comment="addr < maxAddr")) if kernel["ProblemType"]["DataType%s"%tcDataType].isHalf() or kernel["ProblemType"]["DataType%s"%tcDataType].isBFloat16(): hi16 = loopCnt%2 if tP["glvw"]==1 else r%2 else: hi16 = 0 destVgpr="G2L%s+%u+%u"%(tc, g2lIdx, regIdx) module.add(self.chooseGlobalRead(False, \ bpl, destVgpr=destVgprHi if (hi16 and destVgprHi != None) else destVgpr, \ addr0=vgpr("GlobalReadAddr%s+%u"%(tc,graIdx),2), addr1="", \ soffset=0, offset=0, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ hi16=hi16, \ comment="load %u flat value(s)"%(numElementsPerLoad), scope=scope, th=th, nv=nv)) # restore full exec mask SOrSaveExecBX = SOrSaveExecB32 if self.states.kernel["WavefrontSize"] == 32 else SOrSaveExecB64 module.add(SOrSaveExecBX(dst=VCC(), src=sgpr(fullExec,self.states.laneSGPRCount), comment="all threads active")) # increment address by numElementsPerLoad elements module.add(VAddCOU32( dst=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+0"%(tP["tensorChar"], graIdx)), \ src1=bpl, comment="gra += %u element(s) (lower)"%(numElementsPerLoad))) module.add(VAddCCOU32( dst=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ dst1=VCC(), \ src0=vgpr("GlobalReadAddr%s+%u+1"%(tP["tensorChar"], graIdx)), \ src1=vgpr(zeroVgpr), \ src2=VCC(), comment="gra += 1 (upper)")) # int8 byte: # |--------|--------|--------|---V0---|, r = 0, hi8=0, hi16=0, load d16 # |--------|--------|--------|---V1---|, r = 1, hi8=1, hi16=0, load d16 # |--------|---V2---|--------|--------|, r = 2, hi8=0, hi16=1, load d16_hi # |--------|---V3---|--------|--------|, r = 3, hi8=1, hi16=1, load d16_hi # V1, V3 -> shift left 8 bits, or 4 regs (pack) # DestV0|=(V1 << 8), DestV0|= V2, DestV0|=(V3 << 8) # Int8 (byte) if doTailOpt == 2 and r >= rStart and r < rEnd: loadNum = loadNum - 1 if dataIsByte and (destVgprHi != None): if doTailOpt == 1: if behavior == "MERGE" and i == idx and ((r + 1) % numElementsPer4Bytes != 0): # hi8 -> r = 1,3 # hi16 -> r = 2,3 module.add(SWaitCnt(vlcnt=0, comment="")) if hi8: module.add(VLShiftLeftB32(dst=vgpr(destVgprHi), shiftHex=hex(0x8), src=vgpr(destVgprHi), comment="shift left to higher 8 bits")) module.add(VOrB32(dst=vgpr(destVgpr), src0=vgpr(destVgpr), src1=vgpr(destVgprHi), comment="pack a sub 8-bit with dest")) tmpVgprIdx += 1 if (r + 1) == (numLoadVectorComp - 1): module.add(SBranch(labelName=jumpLabel.getLabelName(), comment="")) else: if (doTailOpt == 0) or (doTailOpt == 2 and behavior == "MERGE" and r >= rStart and r < rEnd): # hi8 -> r = 1,3 # hi16 -> r = 2,3 if hi8 or hi16: # r = 1,2,3, vlcnt needed for one packing if doTailOpt == 0: packInt8Code.add(SWaitCnt(vlcnt=(int8TempVgpr-r), comment="")) else: packInt8Code.add(SWaitCnt(vlcnt=(loadNum), comment="")) if hi8: # r = 1,3, shift needed packInt8Code.add(VLShiftLeftB32(dst=vgpr(destVgprHi), shiftHex=hex(0x8), src=vgpr(destVgprHi), comment="shift left to higher 8 bits")) if hi8 or hi16: # r = 1,2,3, packing packInt8Code.add(VOrB32(dst=vgpr(destVgpr), src0=vgpr(destVgpr), src1=vgpr(destVgprHi), comment="pack a sub 8-bit with dest")) destVgprHi += 1 # Half elif destVgprHi != None: if doTailOpt == 1: if behavior == "MERGE" and i == idx: if (numElementsPerLoad == 2 and r % numElementsPer4Bytes != 0) or \ (numElementsPerLoad != 2 and ((r + 1) % numElementsPer4Bytes != 0)): module.add(SWaitCnt(vlcnt=0, comment="")) # 8b checks are needed for mixed 16b/8b precisions if kernel["ProblemType"]["DataType%s"%tcDataType].is8bitFloat(): module.add(VLShiftRightB32(dst=vgpr(destVgprHi), shiftHex=hex(8), src=vgpr(destVgprHi), comment="shift right to lower 8 bits")) module.add(VOrB32(dst=vgpr(destVgpr), src0=vgpr(destVgpr), src1=vgpr(destVgprHi), comment="HasEccHalf: pack")) if kernel["ProblemType"]["DataType%s"%tcDataType].is8bitFloat() and (g2lIdx % 2 == 1): module.add(VLShiftLeftB32(dst=vgpr(destVgpr), shiftHex=hex(16), src=vgpr(destVgpr), comment="shift left to higher 16 bits")) tmpVgprIdx += 1 if (numElementsPerLoad == 2 and r == (numLoadVectorComp - 1)) or\ (numElementsPerLoad != 2 and (r + 1) == (numLoadVectorComp - 1)): module.add(SBranch(labelName=jumpLabel.getLabelName(), comment="")) else: if r % 2 == 1: # 8b checks are needed for mixed 16b/8b precisions if kernel["ProblemType"]["DataType%s"%tcDataType].is8bitFloat(): packInst.add(VLShiftRightB32(dst=vgpr(destVgprHi), shiftHex=hex(8), src=vgpr(destVgprHi), comment="shift right to lower 8 bits")) packInst.add(VOrB32(dst=vgpr(destVgpr), src0=vgpr(destVgpr), src1=vgpr(destVgprHi), comment="HasEccHalf: pack")) if kernel["ProblemType"]["DataType%s"%tcDataType].is8bitFloat() and (g2lIdx % 2 == 1): packInst.add(VLShiftLeftB32(dst=vgpr(destVgpr), shiftHex=hex(16), src=vgpr(destVgpr), comment="shift left to higher 16 bits")) vlcntVal = 0 if doTailOpt == 0 else loadNum if (doTailOpt == 0 and len(tmpVgprPool) == maxNumTempToUse) or (doTailOpt == 2 and behavior == "MERGE" and r >= rStart and r < rEnd): module.add(SWaitCnt(vlcnt=(vlcntVal), comment="")) module.add(packInst) packInst = Module() else: if doTailOpt == 1 and i == idx and behavior == "MERGE" and\ ((numElementsPerLoad == 2 and r == (numLoadVectorComp - 1)) or\ (numElementsPerLoad != 2 and (r + 1) == (numLoadVectorComp - 1))): module.add(SBranch(labelName=jumpLabel.getLabelName(), comment="")) # For half (bf16). Note: for int8, we will checkin after loading all components if (destVgprHi != None) and (not dataIsByte): if len(tmpVgprPool) == maxNumTempToUse and doTailOpt == 0: for v in tmpVgprPool: self.vgprPool.checkIn(v) tmpVgprPool = [] destVgprHi = None r += 1 # next component (for half, byte) # end R loop instOffset += instOffsetInc # add increment value for instOffset. Need to apply after r loop # increment once per r loop (at the end) directToLdsLoads+=1 # for int8: # we do the 3 packs, and checking the 3 extra vgprs after loading all components if dataIsByte and int8TempVgpr and not isLds: assert packInt8Code != None and destVgprHi != None module.add(packInt8Code) if doTailOpt == 0: self.vgprPool.checkIn(destVgprHi - int8TempVgpr) destVgprHi = None # Handles doTailOpt==0 packing after all previous buffer loads have completed if is16b and doTailOpt == 0 and tmpVgprPool: module.add(SWaitCnt(vlcnt=0, comment="Wait for previous GR to finish")) module.add(packInst) for v in tmpVgprPool: self.vgprPool.checkIn(v) tmpVgprPool = [] if kernel["ProblemType"]["Sparse"]: if kernel["DirectToVgprSparseMetadata"]: miWaveTile = kernel["MIWaveTileB"] if (tP["is_sparse"] and tP["isB"]) else kernel["MIWaveTileA"] if (tP["is_sparse"] and tP["isA"]) else 0 for wtIdx in range(0, miWaveTile): offsetVgpr= "GlobalReadOffsetMetadata+%u"%wtIdx for unrollIdx in range(0, kernel["LoopIters"]): bpl = kernel["MIInputPerThread"]//8 # bytes per load: 1 byte for fp16,bf16, 2 bytes for int8 constOffset = unrollIdx * kernel["MatrixInstK"] // 8 for byteIdx in range(0, bpl): constOffset += byteIdx if byteIdx == 0: # For PGR=2 read metadata into the 3rd blk of vpgrValuMetadata if kernel["PrefetchGlobalRead"] >= 2: offsetBlk = self.states.m.numVgprValuPerBlock * 2 else: offsetBlk = 0 destVgprLow="ValuMetadata+%u+%u"%(offsetBlk, (wtIdx*kernel["LoopIters"]+unrollIdx)) destVgpr=destVgprLow else: destVgprHi = self.vgprPool.checkOut(1, tag="globalReadGuardK_destVgprHi3") destVgpr=destVgprHi module.add(self.chooseGlobalRead(kernel["BufferLoad"], \ 1, \ destVgpr=destVgpr, \ addr0=vgpr(offsetVgpr), addr1=sgpr("SrdMetadata",4), \ soffset=0, offset=constOffset, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ hi16=0, \ comment="G -> Reg ValuMetadata", scope=scope, th=th, nv=nv)) if bpl == 2: #pack 2bytes module.add(SWaitCnt(vlcnt=0)) module.add(VLShiftLeftB32(dst=vgpr(destVgprHi), shiftHex="0x8", src=vgpr(destVgprHi), comment="shift left to higher 8 bits")) module.add(VOrB32(dst=vgpr(destVgprLow), src0=vgpr(destVgprLow), src1=vgpr(destVgprHi), comment="pack 2 bytes")) self.vgprPool.checkIn(destVgprHi) destVgprHi = None return loadCnt, self.vgprs.globalReadRegisters[tc], directToLdsLoads loadCnt, vgprList, directToLdsLoads = globalReadGuardKBody(tP, optParams) if doTailOpt == 0 or \ (doTailOpt == 2 and optParams.finalLoop == 1 and optParams.behavior == "MERGE"): if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: if tP["is_sparse"]: globalReadGuardKBody(tP["tpsMetadata"]) if self.db["ConservativeWaitCnt"] & 0x1: module.add(SBarrier(comment="debug")) module.add(SWaitCnt(dscnt=0, vlcnt=0, vscnt=0, comment="")) module.add(SBarrier(comment="debug")) if isTr: self.vgprPool.checkIn(maxGroVgpr) # BufferLoad=0 VGPRs are local to each call — always free them if not isTr and not kernel["BufferLoad"]: self.vgprPool.checkIn(maxAddrVgpr) self.vgprPool.checkIn(bpeVgpr) self.vgprPool.checkIn(zeroVgpr) if doTailOpt == 2: return module, loadCnt, vgprList, directToLdsLoads else: return module ############################################################################## # Redorder the wait instructions to reduce overall waiting time in tail loop. # It needs more vgprs to store data which is needed to be merged into a dword. ############################################################################## def doTailLoopOpt(self, kernel, tP): module = Module("doTailLoop") tc = tP["tensorChar"] orinrp = tP["nrp"] orinrpv = tP["nrpv"] orinrc = tP["nrc"] orinrcv_div_nrcvpi = tP["nrcv"] // tP["nrcvpi"] loadWidth = tP["globalReadInstruction"].totalWidth oriNumLoadVectorComp = (int(loadWidth*self.states.bpr//tP["bpeGR"])) numElementsPerLoad = 1 glvwWorkaround = 8 * kernel["ProblemType"]["MacDataType%s"%tc].numRegisters() dataType = kernel["ProblemType"]["MacDataType%s"%tc] if tP["glvw"] < glvwWorkaround \ else kernel["ProblemType"]["DataType%s"%tc] if dataType.isHalf() or dataType.isBFloat16(): if tP["glvw"] > 1 and kernel["AssertSummationElementMultiple"] % 2 == 0: # Pack two FP16 values into a single load dword x2 numElementsPerLoad = 2 oriNumLoadVectorComp = oriNumLoadVectorComp // 2 totalVgprNum = (tP["nrp"] * tP["nrpv"] * tP["nrc"] * (tP["nrcv"] // tP["nrcvpi"])) * (oriNumLoadVectorComp) if kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["UseBias"] \ and (kernel["ProblemType"]["BiasSrc"] == "A" or kernel["ProblemType"]["BiasSrc"] == "B"): totalVgprNum += 1 numLoadVectorComp = oriNumLoadVectorComp nrp = orinrp nrpv = orinrpv nrc = orinrc nrcv_div_nrcvpi = orinrcv_div_nrcvpi currentSize = self.vgprPool.size() VgprSlot = [] VgprSlotBk = [] maxVgpr = 50 while (len(VgprSlot) < totalVgprNum and len(VgprSlot) <= maxVgpr): tempVgpr = self.vgprPool.checkOut(1, tag="doTailLoopOpt_tempVgpr") if tempVgpr >= currentSize: self.vgprPool.checkIn(tempVgpr) break if kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["UseBias"] \ and (kernel["ProblemType"]["BiasSrc"] == "A" or kernel["ProblemType"]["BiasSrc"] == "B"): if tempVgpr != self.states.bias.startVgprValu: VgprSlot.append(tempVgpr) else: VgprSlot.append(tempVgpr) VgprSlotBk.append(tempVgpr) loopNum = 1 finalVgprNum = totalVgprNum while (finalVgprNum > len(VgprSlot) and totalVgprNum > 0): if nrp > 1: nrp -= 1 elif nrpv > 1: nrpv -= 1 elif nrc > 1: nrc -= 1 elif nrcv_div_nrcvpi > 1: nrcv_div_nrcvpi -= 1 else: numLoadVectorComp = numLoadVectorComp - 1 finalVgprNum = (nrp * nrpv * nrc * nrcv_div_nrcvpi) * (numLoadVectorComp) nrpLoopNum = ceil(tP["nrp"] / nrp) nrpvLoopNum = ceil(tP["nrpv"] / nrpv) nrcLoopNum = ceil(tP["nrc"] / nrc) nrcv_div_nrcvpiLoopNum = ceil(tP["nrcv"] // tP["nrcvpi"] / nrcv_div_nrcvpi) loadVectorCompLoopNum = int(ceil(oriNumLoadVectorComp / numLoadVectorComp)) totalLoopNum = nrpLoopNum * nrpvLoopNum * nrcLoopNum * nrcv_div_nrcvpiLoopNum * loadVectorCompLoopNum i = 0 globalReadRegisters = [] directToLdsLoads = 0 finalLoop = 0 self.optParamsLoad = TailOptParams() self.optParamsMerge = TailOptParams() for nrp_idx in range(nrpLoopNum): for nrpv_idx in range(nrpvLoopNum): for nrc_idx in range(nrcLoopNum): for nrcv_div_nrcvpi_idx in range(nrcv_div_nrcvpiLoopNum): for loadVectorComp_idx in range(loadVectorCompLoopNum): firstLoop = 0 if numElementsPerLoad == 2: periodParam = [nrp_idx * nrp, min(orinrp, (nrp_idx + 1) * nrp), \ nrpv_idx * nrpv, min(orinrpv, (nrpv_idx + 1) * nrpv), \ nrc_idx * nrc, min(orinrc, (nrc_idx + 1) * nrc), \ nrcv_div_nrcvpi_idx * nrcv_div_nrcvpi, \ min(orinrcv_div_nrcvpi, (nrcv_div_nrcvpi_idx + 1) * nrcv_div_nrcvpi), \ loadVectorComp_idx * numLoadVectorComp * 2, \ min(oriNumLoadVectorComp * 2, (loadVectorComp_idx + 1) * numLoadVectorComp * 2)] else: periodParam = [nrp_idx * nrp, min(orinrp, (nrp_idx + 1) * nrp), \ nrpv_idx * nrpv, min(orinrpv, (nrpv_idx + 1) * nrpv), \ nrc_idx * nrc, min(orinrc, (nrc_idx + 1) * nrc), \ nrcv_div_nrcvpi_idx * nrcv_div_nrcvpi, \ min(orinrcv_div_nrcvpi, (nrcv_div_nrcvpi_idx + 1) * nrcv_div_nrcvpi), \ loadVectorComp_idx * numLoadVectorComp, \ min(oriNumLoadVectorComp, (loadVectorComp_idx + 1) * numLoadVectorComp)] if i == 0: firstLoop = 1 self.optParamsLoad.behavior = "LOAD" self.optParamsLoad.vgprSlot = VgprSlot self.optParamsLoad.periodParam = periodParam self.optParamsLoad.directToLdsLoads = directToLdsLoads self.optParamsLoad.firstLoop = firstLoop self.optParamsLoad.finalLoop = finalLoop imod, loadCnt, vgprList, directToLdsLoads = self.globalReadDo(kernel, 2, tP, -1, 0, 2, self.optParamsLoad) module.add(imod) globalReadRegisters = globalReadRegisters + vgprList if i == totalLoopNum - 1: finalLoop = 1 self.optParamsMerge.behavior = "MERGE" self.optParamsMerge.vgprSlot = VgprSlot self.optParamsMerge.periodParam = periodParam self.optParamsMerge.loadNum = loadCnt self.optParamsMerge.directToLdsLoads = directToLdsLoads self.optParamsMerge.firstLoop = firstLoop self.optParamsMerge.finalLoop = finalLoop imod, loadCnt, vgprList, directToLdsLoads = self.globalReadDo(kernel, 2, tP, -1, 0, 2, self.optParamsMerge) module.add(imod) i += 1 self.vgprs.globalReadRegisters[tc] = globalReadRegisters while VgprSlotBk: tempVgpr = VgprSlotBk.pop(0) self.vgprPool.checkIn(tempVgpr) return module ############################################################################## # DirectToLds M0 update: Do It A/B ############################################################################## def directToLdsM0Update(self, kernel, mode, tP, skipWait = False): tc = tP["tensorChar"] tail = mode == 2 imod = Module("directToLdsM0Update%s_%u"%(tc,mode)) DtldsModule = imod.add(Module("dtls_offset%s"%tP["tensorChar"])) if not self.do["GlobalRead%s"%tP["tensorChar"]]: return imod if kernel["DirectToLds%s"%tc]: # DirectToLds only enabled for TLU=1 cases, where the registers are directly copied into LDS # for cases both A&B are DTLS, updating m0 for each GlobalRead requires instruction schedule # along with global reads assert (kernel["LocalWriteUseSgpr%s"%tc]) if self.states.IncLdsBufSwitch: DtldsModule.add(SAddU32(dst=mgpr(0), src0=sgpr("LocalWriteAddr%s"%tc), \ src1=sgpr("LDSBufferWriteInc"), comment="m0 <- LDS write address (base + inc)")) elif self.states.useCommonSgprSwap: DtldsModule.add(SAddU32(dst=mgpr(0), src0=sgpr("LocalWriteAddr%s"%tc), \ src1=sgpr("SwapCommon"), comment="m0 <- LDS write address (base + inc)")) elif kernel["ExpandPointerSwap"]: DtldsModule.add(SAddU32(dst=mgpr(0), src0=sgpr("LocalWriteAddr%s"%tc), \ src1=tP["localWriteSwapByteOffset"], comment="m0 <- LDS write address")) else: DtldsModule.add(SMovB32(dst=mgpr(0), src=sgpr("LocalWriteAddr%s"%tc), comment="m0 <- LDS write address")) # PrefetchGlobalRead>=2 + (oneBufferScheduling or tailLoop) case, generate local read wait for DirectToLds # Early exit + Tailloop case, barrier sync might not be executed when reaching Tailloop. # Need sync before running GR for Tailloop. if kernel["PrefetchGlobalRead"]>=2 and (self.states.oneBufferScheduling or tail) and not skipWait: # do not generate local read wait for PGR=2 DtldsModule.addComment0("before DirectToLds load, ensure prior ds_reads have finished") DtldsModule.add(SWaitCnt(dscnt=0, comment="")) _barrier = SBarrier() DtldsModule.add(_barrier) return imod ############################################################################## # Global Read: Do It A/B ############################################################################## def globalReadDo(self, kernel, mode, tP, unrollLoopIdx=-1, g2lBufIdx=0, \ doTailOpt = 0, optParams = None, tPM = None): tc = tP["tensorChar"] problemType = self.states.kernel["ProblemType"] numWaves: int = kernel["NumWaves"] imod = StructuredModule("globalReadDo%s_%u"%(tc,mode)) if not self.do["GlobalRead%s"%tP["tensorChar"]]: return imod loopIdx = self.states.unrollIdx # TODO - does this handle multiple summation indices? if kernel["SuppressNoLoadLoop"]: if mode==1 and tP["isA"]: if kernel["HalfPLR"]: # only support TDM if kernel["PrefetchGlobalRead"] > 0: imod.header.addComment0("HalfPLR: disable TDM in last #PGR loops") imod.header.add(SCmpLeI32( src0=self.loopCounter(kernel, loopIdx), \ src1=kernel["PrefetchGlobalRead"], \ comment="%s"%"is this the last iters")) if kernel["NumWaves"] > 1: imod.header.add(SCMovB32(dst=sgpr("tdmAGroup0+0"), src=0, comment="")) if kernel["ProblemType"]["MXBlockA"]: imod.header.add(SCMovB32(dst=sgpr("tdmMXSAGroup0+0"), src=0, comment="")) else: imod.header.add(SCMovB32(dst=sgpr("tdmAGroup0+0"), src=0, comment="")) imod.header.add(SCMovB32(dst=sgpr("tdmBGroup0+0"), src=0, comment="")) if kernel["ProblemType"]["MXBlockA"]: imod.header.add(SCMovB32(dst=sgpr("tdmMXSAGroup0+0"), src=0, comment="")) if kernel["ProblemType"]["MXBlockB"]: imod.header.add(SCMovB32(dst=sgpr("tdmMXSBGroup0+0"), src=0, comment="")) else: imod.header.add(SCmpLeI32( src0=self.loopCounter(kernel, loopIdx), \ src1=kernel["PrefetchGlobalRead"], \ comment="is this one of the last %u iteration(s)"%kernel["PrefetchGlobalRead"])) imod.header.add(SCMovB32( dst=sgpr("SrdA+2"), src=0, comment="Set limit to 0 for last PGR iteration(s)")) imod.header.add(SCMovB32( dst=sgpr("SrdB+2"), src=0, comment="Set limit to 0 for last iteration")) if tc == "A" and kernel["enableTDMA"]: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) useSplitTokens = bool(kernel["TDMSplit"]) and not kernel["ProblemType"]["Sparse"] tdmParity = self.states.ldsTensorTokenIdx if useSplitTokens: comp.setMemToken([self.states.memTokenLdsSplit[tdmParity][0]]) else: comp.setMemToken([self.states.ldsTensorTokenIdx]) if self.states.inTailLoop and not kernel["1LDSBuffer"] and kernel["StreamK"]: ldsAddrSgprName = comp.getLdsAddrSgprName("tdmAGroup0") if self.isPrefetchAcrossPersistentEnabled(kernel): with self.allocTmpSgpr(1) as tmpSgprRes: tailBankSgpr = tmpSgprRes.idx imod.middle.add(self.papTdmSelectTailLdsBank(kernel, tailBankSgpr)) imod.middle.add(self.papTdmSetTailLdsBank(kernel, ldsAddrSgprName, tailBankSgpr)) else: clearMask = ~kernel["LdsOffsetA_Blk"] & 0xFFFFFFFF imod.middle.add(SAndB32(dst=sgpr(ldsAddrSgprName), src0=sgpr(ldsAddrSgprName), src1=hex(clearMask), comment="Reset TDM LDS swap bit for tail loop")) # WS mode: this single shared load also serves B (odd waves) via # A's aliased SGPRs, so pass iter operands when either tile is iterate. isIterA = kernel.get("_TDMIterateModeA", False) if self.isTdmWaveSeparated(kernel): isIterA = isIterA or kernel.get("_TDMIterateModeB", False) tdmAGroup2 = "tdmAGroup2" if isIterA else None tdmAGroup3 = "tdmAGroup3" if isIterA else None imod.middle.add(comp.issueLoad("tdmAGroup0", "tdmAGroup1", tdmAGroup2, tdmAGroup3)) # TODO: Embed metadata TDM issueLoad here (mirrors non-TDM pattern where globalReadBody(tP["tpsMetadata"]) # is called after globalReadBody(tP) for the sparse tensor). This would allow _splitTdmLoad in SIA.py # to extract and defer both A and metadata TDM loads together, eliminating the separate globalReadMetadata # module and the special-case handling in noSchedGlobalRead. # if kernel["enableTDMMetadata"] and tP["is_sparse"]: # imod.middle.add(comp.issueLoad("tdmMetadataGroup0", "tdmMetadataGroup1", None, None)) if kernel["TDMSplit"] and not kernel["ProblemType"]["Sparse"]: if numWaves > 1: # Multi-wave: recompute the LDS split boundary and global split increment # transiently per use (see _tdmSplitMultiWaveInc). with self.allocTmpSgpr(2, tag="tdmSplitInc") as incTmp: gIncIdx = incTmp.idx scratchIdx = incTmp.idx + 1 ldsIncOp = self._tdmSplitMultiWaveInc(imod.middle, kernel, gIncIdx, scratchIdx, self.tPA, self.tPB) imod.middle.add(SAddU32(sgpr(f"tdm{tc}Group0+1"), sgpr(f"tdm{tc}Group0+1"), ldsIncOp)) imod.middle.add(SAddU32(sgpr(f"tdm{tc}Group0+2"), sgpr(f"tdm{tc}Group0+2"), sgpr(gIncIdx))) imod.middle.add(SAddCU32(sgpr(f"tdm{tc}Group0+3"), sgpr(f"tdm{tc}Group0+3"), 0, f"tdm{tc} split carry")) else: ldsIncSgprName = f"tdm{tc}LdsSplitIncs" globalIncSgprName = f"tdm{tc}GlobalSplitIncs" imod.middle.add(SAddU32(sgpr(f"tdm{tc}Group0+1"), sgpr(f"tdm{tc}Group0+1"), sgpr(ldsIncSgprName))) imod.middle.add(SAddU32(sgpr(f"tdm{tc}Group0+2"), sgpr(f"tdm{tc}Group0+2"), sgpr(globalIncSgprName))) imod.middle.add(SAddCU32(sgpr(f"tdm{tc}Group0+3"), sgpr(f"tdm{tc}Group0+3"), 0, f"tdm{tc} split carry")) if numWaves > 1: # Recompute the two per-half dim1 boundaries from the live descriptor. # halfRows is per-wave: even waves load A (MacroTile0//2), odd load B # (MacroTile1//2). group1 = f"tdm{tc}Group1" halfRowsA = kernel["MacroTile0"] // 2 halfRowsB = kernel["MacroTile1"] // 2 with self.allocTmpSgpr(3, tag="tdmSplitDim1Recompute") as tmpSgprRes: h0 = tmpSgprRes.idx h1 = tmpSgprRes.idx + 1 hr = tmpSgprRes.idx + 2 if halfRowsA == halfRowsB: imod.middle.add(SMovB32(sgpr(hr), halfRowsA, "halfRows")) else: with self.allocTmpSgpr(1, tag="tdmSplitDim1Parity") as waveIdTmp: self._emitTdmWaveParitySCC(imod.middle, kernel, waveIdTmp.idx) imod.middle.add(SCSelectB32(sgpr(hr), halfRowsB, halfRowsA, "halfRows = parity ? B : A")) imod.middle.add(SLShiftRightB32(sgpr(h0), hex(16), sgpr(f"{group1}+2"), "H0 = dim1 lo")) imod.middle.add(SLShiftLeftB32(sgpr(h1), hex(16), sgpr(f"{group1}+3"), "H0 hi << 16")) imod.middle.add(SOrB32(sgpr(h0), sgpr(h0), sgpr(h1), "H0 = full dim1")) imod.middle.add(SSubU32(sgpr(h1), sgpr(h0), sgpr(hr), "H1 = H0 - halfRows")) imod.middle.add(SCSelectB32(sgpr(h1), 0, sgpr(h1), "clamp H1 to 0")) imod.middle.add(comp.setTensorDim1(group1, h1, self)) comp.setMemToken([self.states.memTokenLdsSplit[tdmParity][1]]) imod.middle.add(comp.issueLoad("tdmAGroup0", "tdmAGroup1", tdmAGroup2, tdmAGroup3)) imod.middle.add(comp.setTensorDim1(group1, h0, self)) else: comp.setMemToken([self.states.memTokenLdsSplit[tdmParity][1]]) imod.middle.add(comp.issueLoad("tdmAGroup0", "tdmAGroup1", tdmAGroup2, tdmAGroup3)) return imod if tc == "MXSA" and kernel["enableTDMA"]: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) comp.setMemToken([self.states.ldsTensorTokenIdx]) if kernel["ProblemType"]["MXBlockA"]: if self.states.inTailLoop and not kernel["1LDSBuffer"] and kernel["StreamK"]: ldsAddrSgprName = comp.getLdsAddrSgprName("tdmMXSAGroup0") if self.isPrefetchAcrossPersistentEnabled(kernel): with self.allocTmpSgpr(1) as tmpSgprRes: tailBankSgpr = tmpSgprRes.idx imod.middle.add(self.papTdmSelectTailLdsBank(kernel, tailBankSgpr)) imod.middle.add(self.papTdmSetTailLdsBank(kernel, ldsAddrSgprName, tailBankSgpr)) else: clearMask = ~kernel["LdsOffsetA_Blk"] & 0xFFFFFFFF imod.middle.add(SAndB32(dst=sgpr(ldsAddrSgprName), src0=sgpr(ldsAddrSgprName), src1=hex(clearMask), comment="Reset TDM LDS swap bit for tail loop")) imod.middle.add(comp.issueLoad("tdmMXSAGroup0", "tdmMXSAGroup1", None, None)) return imod if tc == "Metadata" and kernel["enableTDMMetadata"]: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) comp.setMemToken([self.states.ldsTensorTokenIdx]) imod.add(comp.issueLoad("tdmMetadataGroup0", "tdmMetadataGroup1", None, None)) return imod if tc == "B" and kernel["enableTDMB"]: #TODO: TDM refactor, wave separated TDM only issues 1 tensor load if numWaves == 1: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) useSplitTokens = bool(kernel["TDMSplit"]) and not kernel["ProblemType"]["Sparse"] tdmParity = self.states.ldsTensorTokenIdx if useSplitTokens: comp.setMemToken([self.states.memTokenLdsSplit[tdmParity][0]]) else: comp.setMemToken([self.states.ldsTensorTokenIdx]) if self.states.inTailLoop and not kernel["1LDSBuffer"] and kernel["StreamK"]: ldsAddrSgprName = comp.getLdsAddrSgprName("tdmBGroup0") if self.isPrefetchAcrossPersistentEnabled(kernel): with self.allocTmpSgpr(1) as tmpSgprRes: tailBankSgpr = tmpSgprRes.idx imod.middle.add(self.papTdmSelectTailLdsBank(kernel, tailBankSgpr)) imod.middle.add(self.papTdmSetTailLdsBank(kernel, ldsAddrSgprName, tailBankSgpr)) else: clearMask = ~kernel["LdsOffsetA_Blk"] & 0xFFFFFFFF imod.middle.add(SAndB32(dst=sgpr(ldsAddrSgprName), src0=sgpr(ldsAddrSgprName), src1=hex(clearMask), comment="Reset TDM LDS swap bit for tail loop")) isIterB = kernel.get("_TDMIterateModeB", False) tdmBGroup2 = "tdmBGroup2" if isIterB else None tdmBGroup3 = "tdmBGroup3" if isIterB else None imod.middle.add(comp.issueLoad("tdmBGroup0", "tdmBGroup1", tdmBGroup2, tdmBGroup3)) # TODO: Embed metadata TDM issueLoad here (mirrors non-TDM pattern where globalReadBody(tP["tpsMetadata"]) # is called after globalReadBody(tP) for the sparse tensor). This would allow _splitTdmLoad in SIA.py # to extract and defer both B and metadata TDM loads together, eliminating the separate globalReadMetadata # module and the special-case handling in noSchedGlobalRead. # if kernel["enableTDMMetadata"] and tP["is_sparse"]: # imod.middle.add(comp.issueLoad("tdmMetadataGroup0", "tdmMetadataGroup1", None, None)) if kernel["TDMSplit"] and not kernel["ProblemType"]["Sparse"]: ldsIncSgprName = f"tdm{tc}LdsSplitIncs" globalIncSgprName = f"tdm{tc}GlobalSplitIncs" imod.middle.add(SAddU32(sgpr(f"tdm{tc}Group0+1"), sgpr(f"tdm{tc}Group0+1"), sgpr(ldsIncSgprName))) imod.middle.add(SAddU32(sgpr(f"tdm{tc}Group0+2"), sgpr(f"tdm{tc}Group0+2"), sgpr(globalIncSgprName))) imod.middle.add(SAddCU32(sgpr(f"tdm{tc}Group0+3"), sgpr(f"tdm{tc}Group0+3"), 0, f"tdm{tc} split carry")) comp.setMemToken([self.states.memTokenLdsSplit[tdmParity][1]]) imod.middle.add(comp.issueLoad("tdmBGroup0", "tdmBGroup1", tdmBGroup2, tdmBGroup3)) return imod if tc == "MXSB" and kernel["enableTDMB"]: #TODO: TDM refactor, wave separated TDM only issues 1 tensor load if kernel["NumWaves"] == 1: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) comp.setMemToken([self.states.ldsTensorTokenIdx]) if self.states.inTailLoop and not kernel["1LDSBuffer"] and kernel["StreamK"]: ldsAddrSgprName = comp.getLdsAddrSgprName("tdmMXSBGroup0") if self.isPrefetchAcrossPersistentEnabled(kernel): with self.allocTmpSgpr(1) as tmpSgprRes: tailBankSgpr = tmpSgprRes.idx imod.middle.add(self.papTdmSelectTailLdsBank(kernel, tailBankSgpr)) imod.middle.add(self.papTdmSetTailLdsBank(kernel, ldsAddrSgprName, tailBankSgpr)) else: clearMask = ~kernel["LdsOffsetA_Blk"] & 0xFFFFFFFF imod.middle.add(SAndB32(dst=sgpr(ldsAddrSgprName), src0=sgpr(ldsAddrSgprName), src1=hex(clearMask), comment="Reset TDM LDS swap bit for tail loop")) imod.middle.add(comp.issueLoad("tdmMXSBGroup0", "tdmMXSBGroup1", None, None)) return imod # sizeK % LOCAL_DEPTHU guardK = (mode==2) # set the first tc for below wait code for DirectToLds tc1st = 'A' # if DirectToVgpr is enabled and swapGlobalRead is true, change the first to B if self.isSwapGlobalReadOrderForDtvOrDtl(kernel, prefetch1=(mode==0)): tc1st = 'B' if tc == tc1st and (kernel["DirectToLdsA"] or kernel["DirectToLdsB"]) \ and not kernel["PrefetchGlobalRead"]>=2: if doTailOpt == 0 or \ (doTailOpt == 2 and optParams.behavior == "LOAD" and optParams.firstLoop == 1): # generate local read wait for DirectToLds except for # PrefetchGlobalRead=2 (for PGR=2, generate wait after m0 value setting) imod.header.addComment0("before DirectToLds load, ensure prior ds_reads have finished") if (kernel["DirectToVgprA"] or kernel["DirectToVgprB"]) and not guardK and mode != 3: # no need to generate sync here if DirectToVgpr is enabled and not tail loop imod.header.add(SWaitCnt(dscnt=0, comment="wait for LDS read/write")) else: imod.header.add(self._syncThreads(kernel)) if guardK: if doTailOpt == 0: imod.middle.add(self.globalReadGuardK(kernel, tP)) return imod else: module, loadCnt, vgprList, directToLdsLoads = \ self.globalReadGuardK(kernel, tP, doTailOpt, optParams) imod.middle.add(module) return imod, loadCnt, vgprList, directToLdsLoads # else not-guardK below: record = [False, False] # record global read gpr info for A/B record[0] = True if tc == "A" else False record[1] = True if tc == "B" else False def globalReadBody(tP): tc = tP["tensorChar"] isAB = tc in ("A", "B") self.vgprs.globalReadRegisters[tc] = [] graIdx = 0 g2lIdx = 0 loadWidth = tP["globalReadInstruction"].totalWidth # load width in elements? bpe = tP["bpeGR"] if isAB else tP["bpe"] bpl = bpe * tP["glvw"] # bytes per load isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal( self.states.asmCaps, tP["NonTemporal"], TemporalHint(tP["TemporalHint"]), _nonVolatile(kernel, tc)) isLds = True if kernel["DirectToLds%s"%tc] else False isTr = (tc == "A" or tc == "B") and kernel["enableGLTr%s"%tc] directToLdsLoads = 0 instOffset = 0 prevLdsOffset = 0 if g2lBufIdx >= 1: # G2L vgpr base string. DirectToVgpr or swapAB case. Need to toggle destination vreg set destVgprPrefix = "G2L%s%u"%(tc, g2lBufIdx + 1) else: destVgprPrefix = "G2L%s"%(tc) loopCnt = -1 for perp in range(0, tP["nrp"]): for sPerp in range(0, tP["nrpv"]): for para in range(0, tP["nrc"]): for sPara in range(0, tP["nrcv"]//tP["nrcvpi"]): if tc != "Metadata" and kernel["DirectToVgpr%s"%tc] and kernel["reorderGRInstForDTV%s"%tc]: i = perp + tP["nrp"] * (sPerp + tP["nrpv"] * (para + sPara * tP["nrc"])) else: i = sPara + (tP["nrcv"]//tP["nrcvpi"]) * (para + tP["nrc"] * (sPerp + tP["nrpv"] * perp)) loopCnt += 1 graIdx = i * self.states.rpgo if kernel["BufferLoad"] else i * self.states.rpga g2lIdx = i * loadWidth * tP["bpeRatio"] if loadWidth != 3 else i * (loadWidth + 1) * tP["bpeRatio"] #TODO: remove this if upcoming compiler changes getting merged if loadWidth == 3: g2lIdx = i * (loadWidth + 1) * tP["bpeRatio"] if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc] and kernel["ConvertAfterDS"]: # DTV + ConvertAfterDS case, if bpe > bpeGR, we need to shift g2lIdx for conversion if tP["bpe"] > tP["bpeGR"]: g2lIdx *= tP["bpe"] // tP["bpeGR"] # Each load may contains a small bundle of instructions, package them together in loadModule: loadModule = Module("load%u"%loopCnt) imod.middle.add(loadModule) if (self.states.archCaps["HasEccHalf"] or not self.states.asmCaps["HasWMMA_V1"]) and isAB: numVgprG2L = self.states.a.numVgprG2L if tc == 'A' else self.states.b.numVgprG2L if tc =='B' else self.states.m.numVgprG2L eccBpe = tP["bpeDS"] if kernel["ConvertAfterDS"] else max(tP["bpeGR"], tP["bpe"]) eccOffset = _getEccOffset(loadWidth, bpr=self.states.bpr, bpe=eccBpe, \ glvw=tP["glvw"], idx=i, numVgprG2L=numVgprG2L) else: eccOffset = 0 if kernel["BufferLoad"]: if kernel["_UseSgprForGRO"]: offsetVgpr= "GlobalReadOffset%s+0"%(tc) else: offsetVgpr= "GlobalReadOffset%s+%u"%(tc, graIdx) # vgpr for GRO if not kernel["_UseSgprForGRO"]: soffset = "0" # instruction offset with Sgpr for GRO elif kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"]: soffset = sgpr("ScalarGlobalReadOffset%s+%u"%(tc, graIdx)) # Sgpr for GRO else: soffset = "0" if graIdx == 0 else sgpr("ScalarGlobalReadOffset%s+%u"%(tc, graIdx-1)) unrollMirrorWithSoffset = kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx] in problemType["MirrorDims%s"%tc] and soffset != "0" # ScalarGlobalReadOffset should be negative value with unroll mirroring. # However, buffer_load uses soffset as uint value, so GRO - SGRO, SGRO = 0 if unrollMirrorWithSoffset: codeMod = Module("mirrorIdx%u"%loopCnt) codeMod.add(VSubU32(dst=vgpr(offsetVgpr), src0=vgpr(offsetVgpr), src1=soffset, comment="mirror unroll: GRO=GRO-SGRO, soffset=0")) loadModule.add(codeMod) soffset_prev = soffset soffset = "0" if kernel["DirectToLds%s"%tc]: # use bpe with GlobalReadVectorWidth ldsInc = int((self.states.kernel["WavefrontSize"] * kernel["GlobalReadVectorWidth%s"%tc] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"] * kernel["GlobalReadVectorWidth%s"%tc]) * tP["bpeGR"]) if kernel["LdsBlockSizePerPad%s"%tc] != 0: ldsInc += int((ldsInc // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeGR"]) else: padInterval = (self.states.kernel["WavefrontSize"] if kernel["WaveSeparateGlobalRead%s"%tc] else kernel["NumThreads"]) * self.states.bpr ldsInc += int((ldsInc // padInterval) * kernel["LdsPad%s"%tc] * tP["bpeGR"]) if kernel["UseInstOffsetForGRO"]: # buffer_load only support 12 bit instruction offset # we have to increase m0 if offset is larger thant 12 bits if instOffset >= self.buff_load_inst_offset_max: inc = (instOffset // self.buff_load_inst_offset_max) * self.buff_load_inst_offset_max loadModule.add(SAddU32(dst=mgpr(0), src0=mgpr(0), src1=inc, comment="Move LDS write address to next base" )) instOffset -= inc elif directToLdsLoads != 0: # m0 offset conversion (only for UseInstOffsetForGRO == 0) # in tP["glvw"] == 1 and tP["nrc"] > 1 case, only m0 offset conversion is necessary. row and column index conversion is not necessary. if tP["nrc"] > 1: # another address conversion for DirectToLds + NumLoadsCoalesced > 1 divisorName = tP["lvc"] divisor = kernel[divisorName] # DirectToLds + NumLoadsCoalesced>1 case, need to adjust m0 increment value to store values to correct location in LDS wSize = max(self.states.kernel["WavefrontSize"], divisor) lscaOffset = int(para * wSize * tP["bpeGR"] * tP["glvw"]) ldsOffset = int(ldsInc * tP["nrc"] * (sPerp + tP["nrpv"] * perp) + lscaOffset) ldsInc = ldsOffset - prevLdsOffset prevLdsOffset = ldsOffset loadModule.add(SAddU32(dst=mgpr(0), src0=mgpr(0), src1=int(ldsInc), comment="Move LDS write address to next line" )) directToLdsLoads+=1 destVgpr=0 self.vgprs.globalReadRegisters[tc].append(0) else: g2lIdxM = i * max(loadWidth * tP["bpeRatio"], 1) destVgpr = destVgprPrefix + "+%u"%((g2lIdx+eccOffset+tP["shiftGR"]) if isAB else g2lIdxM) self.vgprs.globalReadRegisters[tc].append(g2lIdx+eccOffset+tP["shiftGR"] if isAB else g2lIdxM) if tP["isM"]: assert(graIdx <= self.states.m.numVgprG2LAllocated) # TODO: is it possible to load only hi16 when no in tail? (need to check INT8 too) datatype = kernel["ProblemType"]["DataType%s"%tc] if kernel["ConvertAfterDS"] else kernel["ProblemType"]["MacDataType%s"%tc if (tc=='A' or tc=='B') else "DataType"] isHigh16Bits = (datatype.isHalf() or datatype.isBFloat16()) and loopCnt%2==1 if isAB else False if tc == "A" and record[0] == True: self.globalread_gpr_record.a.addrVgpr.append(offsetVgpr) self.globalread_gpr_record.a.offset.append(soffset) elif tc == "B" and record[1] == True: self.globalread_gpr_record.b.addrVgpr.append(offsetVgpr) self.globalread_gpr_record.b.offset.append(soffset) useBuffer = not isTr loadModule.add( self.chooseGlobalRead(useBuffer, \ bpl, destVgpr=destVgpr, \ addr0=vgpr(offsetVgpr), addr1=sgpr("Srd%s"%tc, 2 if isTr else 4), \ soffset=soffset, offset=instOffset, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ tr=isTr, hi16=isHigh16Bits , \ comment="G -> Reg %u_%u_%u_%u"%(para, sPara, perp, sPerp), scope=scope, th=th, nv=nv)) if unrollMirrorWithSoffset: codeMod = Module("mirrorIdx%u"%loopCnt) codeMod.add(VAddU32(dst=vgpr(offsetVgpr), src0=vgpr(offsetVgpr), src1=soffset_prev, comment="mirror unroll: restore GRO=GRO+SGRO")) loadModule.add(codeMod) if kernel["DirectToLds%s"%tc] and kernel["UseInstOffsetForGRO"]: instOffset += ldsInc #print "IM=", type(imod.instList[-1]), imod.instList[-1], else: # not buffer load offsetVgpr = "GlobalReadAddr%s+%u"%(tc,graIdx) if tc == "A" and record[0] == True: self.globalread_gpr_record.a.addrVgpr.append(offsetVgpr) self.globalread_gpr_record.a.offset.append("0") elif tc == "B" and record[1] == True: self.globalread_gpr_record.b.addrVgpr.append(offsetVgpr) self.globalread_gpr_record.b.offset.append("0") destVgpr = destVgprPrefix + "+%u"%(g2lIdx + tP["shiftGR"]) loadModule.add( self.chooseGlobalRead(False, \ bpl, \ destVgpr=destVgpr, \ addr0=vgpr(offsetVgpr,2), addr1="", \ soffset=0, offset=0, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ hi16=(kernel["ProblemType"]["MacDataType%s"%tc if not tP["isM"] else "DataType"].isHalf() or kernel["ProblemType"]["MacDataType%s"%tc if not tP["isM"] else "DataType"].isBFloat16()) and loopCnt%2==1, \ comment="G -> Reg %u_%u_%u_%u"%(para, sPara, perp, sPerp ))) if kernel["ProblemType"]["Sparse"] and kernel["DirectToVgprSparseMetadata"]: if tP["is_sparse"]: if kernel["PrefetchGlobalRead"] == 1 and unrollLoopIdx % 2 == 0: offsetBlk = self.states.m.numVgprValuPerBlock elif kernel["PrefetchGlobalRead"] >= 2: offsetBlk = self.states.m.numVgprValuPerBlock * 2 else: offsetBlk = 0 miWaveTile = kernel["MIWaveTileA"] if tP["isA"] else kernel["MIWaveTileB"] for wtIdx in range(0, miWaveTile): offsetVgpr= "GlobalReadOffsetMetadata+%u"%wtIdx for unrollIdx in range(0, kernel["LoopIters"]): bpl = kernel["MIInputPerThread"]//8 # bytes per load: 1 byte for fp16,bf16, 2 bytes for int8 constOffset = unrollIdx * kernel["MatrixInstK"] // 8 codeMod = Module("load metadata%u"%loopCnt) imod.middle.add(codeMod) codeMod.add( self.chooseGlobalRead(kernel["BufferLoad"], \ bpl, \ destVgpr="ValuMetadata+%u+%u"%(offsetBlk, (wtIdx*kernel["LoopIters"]+unrollIdx)), \ addr0=vgpr(offsetVgpr), addr1=sgpr("SrdMetadata",4), \ soffset=0, offset=constOffset, \ glc=isGlc, slc=isSlc, nt=isNT, lds=isLds, \ hi16=0, \ comment="G -> Reg ValuMetadata", scope=scope, th=th, nv=nv)) if tc == "A" and record[0] == True: self.globalread_gpr_record.a.addrVgpr = [] self.globalread_gpr_record.a.offset = [] elif tc == "B" and record[1] == True: self.globalread_gpr_record.b.addrVgpr = [] self.globalread_gpr_record.b.offset = [] globalReadBody(tP) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: # Workaround, two cases to put metadata GR to non-sparse side # 1. Sparse B + DTLA0_DTLB1_DTLM0 # 2. Sparse A + DTLA1_DTLB1_DTLM0 swapMetadataGlobalReadTc = False if kernel["ProblemType"]["Sparse"] == 2 and kernel["DirectToLdsB"] and not kernel["DirectToLdsA"] and not kernel["DirectToLdsMetadata"]: swapMetadataGlobalReadTc = True if kernel["ProblemType"]["Sparse"] == 1 and kernel["DirectToLdsB"] and kernel["DirectToLdsA"] and not kernel["DirectToLdsMetadata"]: swapMetadataGlobalReadTc = True if tP["is_sparse"] ^ swapMetadataGlobalReadTc: if kernel["DirectToLdsMetadata"]: imod.middle.add(SMovB32(dst=mgpr(0), src=sgpr("LocalWriteAddrMetadata"),\ comment="DTLM: m0 <- metadata LDS write base")) globalReadBody(tPM if tPM else tP["tpsMetadata"]) if self.db["ConservativeWaitCnt"] & 0x1: imod.footer.add(SBarrier(comment="debug")) imod.footer.add(SWaitCnt(dscnt=0, vlcnt=0, vscnt=0, comment="conservative wait")) imod.footer.add(SBarrier(comment="debug")) #module.add(self.getCmpAssert(self.asmAssert.lt, vgpr("Serial"), 64)) # examine second wavefront return imod def tdmSwapLdsOffset(self, kernel, tP) -> Module: tc: str = tP["tensorChar"] strippedTc: str = tc[-1] oppositeStrippedTc = "A" if strippedTc == "B" else "B" needSwap: bool = not kernel["1LDSBuffer"] if not needSwap: return Module("TDM LDS swap (Empty)") module: Module = Module("TDM LDS swap") storeSwapAddr = kernel["StoreSwapAddr"] module.addComment(f"TDM LDS swap(aligned pow2: {not storeSwapAddr})") comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) ldsAddrSgprName: str = comp.getLdsAddrSgprName(f"tdm{tc}Group0") if not kernel["StoreSwapAddr"]: swapMask: int = kernel[f"LdsOffsetA_Blk"] module.add(SXorB32(sgpr(ldsAddrSgprName), sgpr(ldsAddrSgprName), hex(swapMask))) else: assert kernel["NumWaves"] > 1 with self.allocTmpSgpr(1, tag="tdmSwapLdsOffset_tmpSgprRes") as tmpSgprRes: tmpSgprIdx: int = tmpSgprRes.idx module.add(SCmpLtU32(sgpr(ldsAddrSgprName), kernel["LdsOffsetA_Blk"], "Check if < blk offset")) module.add(SMovB32(sgpr(tmpSgprIdx), -kernel["LdsOffsetA_Blk"], "Init as -blk")) module.add(SCSelectB32(sgpr(tmpSgprIdx), kernel["LdsOffsetA_Blk"], sgpr(tmpSgprIdx), "<: +blk, >=: -blk")) module.add(SAddI32(sgpr(ldsAddrSgprName), sgpr(ldsAddrSgprName), sgpr(tmpSgprIdx), "Do swap")) return module ############################################################################## # Local Write: Swap Offsets A/B ############################################################################## def localWriteSwapOffsets(self, kernel, internalPointerSwap, tP, prefetch=False, skipMetaSwap=False): tc = tP["tensorChar"] # TDM has its own swap path (tdmSwapLdsOffset); skip here for that case. if not self.do["LocalWrite%s"%tc] or (kernel["enableTDMA"] and kernel["enableTDMB"]): return Module("localWriteSwapOffsets (No local write%s)"%tc) needSwap = False if kernel["1LDSBuffer"] else True doMetadataCheck = kernel["ProblemType"]["Sparse"] and \ ((kernel["ProblemType"]["Sparse"] ==2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"])) needMetaSwap = needSwap and doMetadataCheck # swap not needed if DirectToVgpr is enabled (do not use DTVA/B for metaData. Change needSwap after setting needMetaSwap) if kernel["DirectToVgpr%s"%tc]: needSwap = False if doMetadataCheck: if kernel["DirectToVgprSparseMetadata"]: needMetaSwap = (kernel["PrefetchGlobalRead"] >= 2 and kernel["ExpandPointerSwap"]) if not (needSwap or needMetaSwap): return Module("localWriteSwapOffsets (Empty)") module = Module("localWriteSwapOffsets") def localWriteSwapXOR(tc, src0Val, numLwa): if kernel["LocalWriteUseSgpr%s"%tc]: module.add(SXorB32( dst=sgpr("LocalWriteAddr%s"%tc), \ src0=src0Val, \ src1=sgpr("LocalWriteAddr%s"%tc), \ comment="swap Red Blk SGPR")) else: module.add(VXorB32( dst=vgpr("LocalWriteAddr%s"%tc), \ src0=src0Val, \ src1=vgpr("LocalWriteAddr%s"%tc), \ comment="swap Red Blk")) for i in range(1,numLwa): module.add(VAddU32( dst=vgpr("LocalWriteAddr%s+%u"%(tc,i)), \ src0=(i * self.states.regCaps["maxLDSConstOffset"]), \ src1=vgpr("LocalWriteAddr%s"%tc), \ comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024))) def localWriteAddRound(tc): is1st = tc == "A" # so far, A is always first # LDSBufferWriteInc is common for A and B. Add this only for the first one (tc=="A") if is1st: module.add(SAddU32( dst=sgpr("LDSBufferWriteInc"), \ src0="LdsOneBlockSize", \ src1=sgpr("LDSBufferWriteInc"), \ comment="add LDS block size to incSgpr")) module.add(SCmpEQU32(src0=sgpr("LDSBufferWriteInc"), src1="LdsBlockEndSize", comment="LDSBufferWriteInc == End ?")) module.add(SCMovB32(dst=sgpr("LDSBufferWriteInc"), src=0, comment="LDSBufferWriteInc loop back to 0")) def localWriteSwapCommon(tc): is1st = tc == "A" # so far, A is always first if is1st: module.add(SXorB32( dst=sgpr("SwapCommon"), \ src0="LdsOneBlockSize", \ src1=sgpr("SwapCommon"), \ comment="xor LDS block size")) def getSrc0Val(tc): src0Val = None if kernel["StoreSwapAddr"]: if kernel["LocalWriteUseSgpr%s"%tc]: src0Val = sgpr("Swap%s"%tc) else: src0Val = vgpr("LocalWriteSwapAddr%s"%tc) else: # Using inlined constants src0Val = hex(kernel["LdsOffsetA_Blk"]) return src0Val if needSwap: #fixme-iui need to use wrapping increment for double or triple buffering: if internalPointerSwap and not kernel["StoreSwapAddr"]: tP["localWriteSwapByteOffset"] = 0 if tP["localWriteSwapByteOffset"] else kernel["LdsOffsetA_Blk"] module.addComment1("(EPS=1) local write swap internal offset -> %u" % tP["localWriteSwapByteOffset"]) elif self.states.IncLdsBufSwitch: # 3 or more LDS block case, we do not use xor. Instead, use add and max check for round back # (numLDSBlk>=3 is for DTL (and LocalWriteUseSgpr) only) localWriteAddRound(tc) elif self.states.useCommonSgprSwap: # commonSwap case, need only 1 swap # (generate at "A" only) localWriteSwapCommon(tc) else: src0Val = getSrc0Val(tc) numLwa = self.states.a.numVgprLocalWriteAddr if tP["isA"] else self.states.b.numVgprLocalWriteAddr localWriteSwapXOR(tc, src0Val, numLwa) # This used to control where to store the metadata if needMetaSwap and not skipMetaSwap: if kernel["DirectToVgprSparseMetadata"]: tP["metadataWriteSwapByteOffset"] = 0 if tP["metadataWriteSwapByteOffset"] else self.states.m.numVgprValuPerBlock module.addComment1("metadata write swap offset -> %u" % tP["metadataWriteSwapByteOffset"]) else: tc = "Metadata" tPM = tP["tpsMetadata"] if internalPointerSwap and not kernel["StoreSwapAddr"]: tPM["localWriteSwapByteOffset"] = 0 if tPM["localWriteSwapByteOffset"] else kernel["LdsOffsetA_Blk"] module.addComment1("(EPS=1) local write swap internal offset -> %u" % tPM["localWriteSwapByteOffset"]) else: src0Val = getSrc0Val(tc) numLwa = self.states.m.numVgprLocalWriteAddr localWriteSwapXOR(tc, src0Val, numLwa) return module ############################################################################## # Local Write: Reset Offsets A/B # used for global-read + tail-loop to reset to writing in red ############################################################################## def localWriteResetOffsets(self, kernel, internalPointerSwap, tP): tc = tP["tensorChar"] if not self.do["LocalWrite%s"%tc]: return Module("localWriteResetOffsets (no local write%s)"%tc) needReset = not (kernel["1LDSBuffer"]) doMetadataCheck = kernel["ProblemType"]["Sparse"] and \ ((kernel["ProblemType"]["Sparse"] ==2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"])) needMetaReset = needReset and doMetadataCheck # reset not needed if DirectToVgpr is enabled (do not use DTVA/B for metaData. Change needReset after setting needMetaReset) if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc]: needReset = False if doMetadataCheck: if kernel["DirectToVgprSparseMetadata"]: needMetaReset = (kernel["PrefetchGlobalRead"] >= 2 and kernel["ExpandPointerSwap"]) if not (needReset or needMetaReset): # no local write code if DirectToVgpr is enabled return Module("localWriteResetOffsets (Empty)") module = Module("localWriteResetOffsets") if needReset: resetMask = hex(kernel["LdsOffsetA_Blk"]-1 | self.consts.ldsOOB) useSwapAddr = (internalPointerSwap or (kernel["StoreSwapAddr"] and not self.states.useCommonSgprSwap)) if useSwapAddr: if internalPointerSwap: tP["localWriteSwapByteOffset"] = 0 else: if kernel["LocalWriteUseSgpr%s"%tc]: tmpsgpr = self.sgprPool.checkOut(1, tag="localWriteResetOffsets_tmpsgpr") module.add(SXorB32( dst=sgpr(tmpsgpr), \ src0=sgpr("Swap%s"%tc), \ src1=sgpr("LocalWriteAddr%s"%tP["tensorChar"]), \ comment="Get other lds buffer offset value")) module.add(SMinU32(dst=sgpr("LocalWriteAddr%s"%tP["tensorChar"]), \ src0=sgpr("LocalWriteAddr%s"%tP["tensorChar"]), \ src1=sgpr(tmpsgpr), comment="Set LWA to first buffer offset" )) self.sgprPool.checkIn(tmpsgpr) else: tmpvgpr = self.vgprPool.checkOut(1, tag="localWriteResetOffsets_tmpvgpr") module.add(VXorB32( dst=vgpr(tmpvgpr), \ src0=vgpr("LocalWriteSwapAddr%s"%tc), \ src1=vgpr("LocalWriteAddr%s"%tc), \ comment="Get other lds buffer offset")) module.add(VMinI32(dst=vgpr("LocalWriteAddr%s"%tc), \ src0=vgpr("LocalWriteAddr%s"%tc), \ src1=vgpr(tmpvgpr), comment="Set LWA to first buffer offset")) self.vgprPool.checkIn(tmpvgpr) elif self.states.IncLdsBufSwitch: # IncLdsBufSwitch case, round back to 0 # (numLDSBlk>=3 is for DTL (and LocalWriteUseSgpr) only) module.add(SMovB32( dst=sgpr("LDSBufferWriteInc"), \ src=0, \ comment="reset incSgpr")) elif self.states.useCommonSgprSwap: # commonSwap case, back to 0 # (generate at "A" only) if tc == "A": module.add(SMovB32( dst=sgpr("SwapCommon"), \ src=0, \ comment="reset swapCommon")) else: if kernel["LocalWriteUseSgpr%s"%tc]: module.add(SAndB32( dst=sgpr("LocalWriteAddr%s"%tP["tensorChar"]), \ src0=resetMask, \ src1=sgpr("LocalWriteAddr%s"%tP["tensorChar"]), \ comment="reset to Red")) else: module.add(VAndB32( dst=vgpr("LocalWriteAddr%s"%tc), \ src0=resetMask, \ src1=vgpr("LocalWriteAddr%s"%tc), \ comment="reset to Red")) if needMetaReset: if kernel["DirectToVgprSparseMetadata"]: tP["metadataWriteSwapByteOffset"] = 0 module.addComment1("reset metadata write offset to %u" % tP["metadataWriteSwapByteOffset"]) else: tPM = tP["tpsMetadata"] resetMask = hex(kernel["LdsOffsetA_Blk"]-1 | self.consts.ldsOOB) if internalPointerSwap: tPM["localWriteSwapByteOffset"] = 0 elif kernel["StoreSwapAddr"]: if kernel["LocalWriteUseSgpr%s"%tPM["tensorChar"]]: tmpsgpr = self.sgprPool.checkOut(1, tag="localWriteResetOffsets_StoreSwapAddr_tmpsgpr") module.add(SXorB32( dst=sgpr(tmpsgpr), \ src0=sgpr("Swap%s"%tPM["tensorChar"]), \ src1=sgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ comment="Get other lds buffer offset value")) module.add(SMinU32(dst=sgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ src0=sgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ src1=sgpr(tmpsgpr), comment="Set LWA to first buffer offset" )) self.sgprPool.checkIn(tmpsgpr) else: tmpvgpr = self.vgprPool.checkOut(1, tag="localWriteResetOffsets_StoreSwapAddr_tmpvgpr") module.add(VXorB32( dst=vgpr(tmpvgpr), \ src0=vgpr("LocalWriteSwapAddr%s"%tPM["tensorChar"]), \ src1=vgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ comment="Get other lds buffer offset")) module.add(VMinI32(dst=vgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ src0=vgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ src1=vgpr(tmpvgpr), comment="Set LWA to first buffer offset")) self.vgprPool.checkIn(tmpvgpr) else: if kernel["LocalWriteUseSgpr%s"%tPM["tensorChar"]]: module.add(SAndB32( dst=sgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ src0=resetMask, \ src1=sgpr("LocalWriteAddr%s"%tPM["tensorChar"]), \ comment="reset to Red")) else: module.add(VAndB32( dst=vgpr("LocalWriteAddr%s+%u"%(tPM["tensorChar"], 0)), \ src0=resetMask, \ src1=vgpr("LocalWriteAddr%s+%u"%(tPM["tensorChar"], 0)), \ comment="reset to Red")) numLwa = 0 if tP["isA"]: numLwa = self.states.a.numVgprLocalWriteAddr elif tP["isMXSA"]: numLwa = self.states.mxsa.numVgprLocalWriteAddr elif tP["isMXSB"]: numLwa = self.states.mxsb.numVgprLocalWriteAddr elif tP["isB"]: numLwa = self.states.b.numVgprLocalWriteAddr elif tP["isMXSA"]: numLwa = self.states.mxsa.numVgprLocalWriteAddr elif tP["isMXSB"]: numLwa = self.states.mxsb.numVgprLocalWriteAddr elif tP["isM"]: numLwa = self.states.m.numVgprLocalWriteAddr else: raise Exception(f"unsupport tc %s{tc}") for i in range(1, numLwa): module.add(VAddU32( dst=vgpr("LocalWriteAddr%s+%u"%(tc, i)), \ src0=(i * self.states.regCaps["maxLDSConstOffset"]), \ src1=vgpr("LocalWriteAddr%s"%tc), \ comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024))) return module ############################################################################## # Calculate offset to use for LDS write # Intro: # Each WI has a 2D tile index (coal, perp). # - Code above computes global mem address by scaling one dim by the # lda and adding the other. # - Here we compute a linear LDS offset by scaling one dim by the MT # dim and adding the other. # Result is we map a tile from global memory into LDS. Consecutive LDS # locations contain elements from different summation 'rows' - therefore # loading a row of LDS will feed computations for different C tile indices. # Notes: # Total load insts is nrc * nrp which load the macro-tile. # Par and coalesced are ~synonyms referring to same dimension # Either nrpv or nrvc must be 1 - can't have vectors in both dimensions. # Thus either sPerp or sPara is 0. # Inputs: # perp : index of the load in perp dimension (0...nrp) # par : index of the load in the para dim (0...nrc) # sPerp : component index of the perp vector (0...nrpv) # sPara : component index of the par vector (0...nrcv) # Outputs: # offsetBytes : Offset in bytes for the _ds_store instruction # i : i-th instruction # comment : Comment with the text version of the formula ############################################################################# def calculateLdsWriteOffset(self, perp, para, sPerp, sPara, kernel, tP): tc = tP["tensorChar"] mask = 0 #print "tc ", tc, " perp ", perp, " para ", para, " sPerp ", sPerp, " sPara ", sPara lscaOffset = para * kernel[tP["lsc"]] perp_masked = perp perp_rem = 0 lspaOffset = perp_masked * kernel[tP["lsp"]] rem = 0 # Add component offset to interleave from different regs # and compute mysterious "i" assert(sPerp==0 or sPara==0) if tP["tlu"] != kernel["UnrollMajorLDS%s" % tP["tensorChar"]]: lspaOffset += sPerp & mask lscaOffset += sPara rem = (sPerp & ~mask) i = sPara + (tP["nrcv"]//tP["nrcvpi"]) * (para + tP["nrc"] * (sPerp + tP["nrpv"] * perp_masked)) #print "nrcv ", tP["nrcv"], " nrcvpi ", tP["nrcvpi"], " nrc ", tP["nrc"], " nrpv ", tP["nrpv"] else: lscaOffset += sPara lspaOffset += sPerp rem = 0 i = sPara + (tP["nrcv"]//tP["nrcvpi"]) * (para * tP["glvw"] + tP["nrc"] * (sPerp + tP["glvw"] * tP["nrpv"] * perp )) #if not tP["tlu"]: # tmp = sPara # sPara = sPerp # sPerp = tmp # print("0lspaOffset", lspaOffset) # print("0lscaOffset", lscaOffset) LdsPad = kernel["LdsPad%s"%tc] if kernel["LdsBlockSizePerPad%s"%tc] == 0 else 0 lds_stride = (kernel["_DepthU%s"%tc] + LdsPad) if kernel["UnrollMajorLDS%s" % tP["tensorChar"]] \ else (kernel[tP["mt"]] + LdsPad) if tP["tlu"] != kernel["UnrollMajorLDS%s" % tP["tensorChar"]]: lspaOffset *= lds_stride lspaOffset += rem + perp_rem else: lscaOffset *= lds_stride lscaOffset += rem # print("1lspaOffset", lspaOffset) # print("1lscaOffset", lscaOffset) #if tP["tlu"]: # lspaOffset *= tP["glvw"] # lscaOffset *= tP["glvw"] # print("2lspaOffset", lspaOffset) # print("2lscaOffset", lscaOffset) if kernel["UseGeneralizedNLCOne%s"%tc]: offsetElements = perp * kernel["GlobalReadVectorWidth%s"%tc] * kernel["WavefrontSize"] * kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] else: offsetElements = (lspaOffset + lscaOffset) # print("offsetElements", offsetElements) offsetBytes = int(offsetElements*tP["bpeDS"]) if kernel["LdsBlockSizePerPad%s"%tc] != 0 and kernel["LdsPad%s"%tc] != 0: offsetBytes = int(offsetBytes + (offsetBytes // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"]) offsetBytes += tP["localWriteSwapByteOffset"] #print("offsetBytes", offsetBytes) #print "offset", offset comment = "lwo%s_%u_%u_%u_%u = (%s%d*%s)" \ % (tP["tensorChar"], \ para, sPara, perp, sPerp, \ (("%u + "%sPara) if tP["wtc"] else ""), \ para, tP["lsc"] ) if not tP["tlu"]: comment += "*(MT%s+PAD)" % (tP["tileChar"]) comment += " + (%d*%s)" % (perp, tP["lsp"]) if tP["tlu"]: comment += "(*MT%s+PAD)" % (tP["tileChar"]) comment += " = %u" % (offsetBytes) return (offsetBytes, i, comment) def recalcLocalWriteAddresses(self, kernel, tP): tc = tP["tensorChar"] module = Module("recalcLocalWriteAddresses") module.addComment1("recalculate LocalWriteAddr{}".format(tc)) lwvw = getattr(self, "localWriteWidth{}".format(tc)) newInstIdx = self.selectMemoryInstruction("LocalWrite", lwvw, \ False, \ tP["localWrite2Coalesced"], tP["localWrite2Perpendicular"], [tP["localWriteStrideTile"], tP["localWriteStrideUnroll"]] ) tP["localWriteInstruction"] = self.memoryInstructions["LocalWrite"][newInstIdx] loopComponent = Component.PersistentLoop.find(self) module.add(loopComponent.recalcLocalWriteAddresses(self, kernel, tc)) # local write tile assignments module.add(self.lwaTileAssignment(kernel, tP)) # local write unroll assignments module.add(self.lwaUnrollAssignment(kernel, tP)) # local write local write first offsets module.add(self.lwaFirstOffset(kernel, tP)) # global read tile assignment module.add(self.graTileAssignment(kernel, tP)) # global read tile offsets module.add(self.graTileOffsets(kernel, tP)) # global read unroll offsets module.add(self.graUnrollOffsets(kernel, tP)) # still needed for vgpr resource management # intentionally not emitting code self.graFinalOffsets(kernel, tP) if kernel["ProblemType"]["Sparse"] and kernel["DirectToVgprSparseMetadata"]: if (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]): graMetadataFinalOffsets(kernel, tP) return module def recalcLocalReadAddressesAB(self, kernel, tPA, tPB): imod = Module() if self.states.inTailLoop: # it do 1 iteration each loop in tail loop, and is no use to wider local read next iteration. # In 1 block MI, it remap localReadAddr in order to let each thread wider local read continuous k # this decrease performance since it require more loop to handle continuous k in each thread. # reCalculating localread address because we disable wider local read in tail loop if not kernel["UseDotInstruction"] and ((self.states.numReadsIterCoalescedA > 1 or self.states.numReadsIterCoalescedB > 1)): loopComponent = Component.PersistentLoop.find(self) imod.add(loopComponent.recalcLocalReadAddressesAB(self, kernel)) self.states.numReadsIterCoalescedA = 1 self.states.numReadsIterCoalescedB = 1 if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: self.states.numReadsIterCoalescedMetadata = 1 tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] imod.add(self.lraTileAssignment(kernel, tPA, tPB)) imod.add(self.lraFinalOffset(kernel, tPA)) imod.add(self.lraDeclareAddresses(kernel, tPA)) imod.add(self.lraFinalOffset(kernel, tPB)) imod.add(self.lraDeclareAddresses(kernel, tPB)) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: imod.add(self.lraFinalOffset(kernel, tPM)) imod.add(self.lraDeclareAddresses(kernel, tPM)) localRead2Perpendicular = False instructions = self.memoryInstructions if kernel["UnrollMajorLDSA"]: localReadWidth = (kernel["MIInputPerThreadA"] * tPA["bpeDS"]) // self.states.bpr localReadInstructionIdxA = \ self.selectMemoryInstruction("LocalRead", localReadWidth, \ False, \ tPA["localRead2Coalesced"], localRead2Perpendicular, [tPB["localReadStrideCoalesced"]] ) tPA["localReadInstruction"] = instructions["LocalRead"][localReadInstructionIdxA] if kernel["UnrollMajorLDSB"]: localReadWidth = (kernel["MIInputPerThreadB"] * tPB["bpeDS"]) // self.states.bpr localReadInstructionIdxB = \ self.selectMemoryInstruction("LocalRead", localReadWidth, \ False, \ tPB["localRead2Coalesced"], localRead2Perpendicular, [tPB["localReadStrideCoalesced"]] ) tPB["localReadInstruction"] = instructions["LocalRead"][localReadInstructionIdxB] if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: localReadWidth = tPM["bpeDS"] / self.states.bpr if kernel["UnrollMajorLDSMetadata"]: localReadWidth = (kernel["MIInputPerThreadMetadata"] * tPM["bpeDS"]) // self.states.bpr localReadInstructionIdxM = \ self.selectMemoryInstruction("LocalRead", localReadWidth, \ False, \ tPM["localRead2Coalesced"], localRead2Perpendicular, [ tPM["localReadStrideCoalesced"]] ) tPM["localReadInstruction"] = instructions["LocalRead"][ \ localReadInstructionIdxM] # dot2: remove wider local read in tail loop, need to recalculate lra offset elif kernel["UseDotInstruction"] and (kernel["InnerUnroll"] > 1): self.states.numReadsIterCoalescedA = 1 self.states.numReadsIterCoalescedB = 1 self.states.lrvwUnrollA //= kernel["InnerUnroll"] self.states.lrvwUnrollB //= kernel["InnerUnroll"] localRead2Perpendicular = False instructions = self.memoryInstructions if kernel["UnrollMajorLDSA"]: localReadWidth = (kernel["NumDotElements"] * tPA["bpeDS"]) // self.states.bpr localReadInstructionIdxA = \ self.selectMemoryInstruction("LocalRead", localReadWidth, \ False, \ tPA["localRead2Coalesced"], localRead2Perpendicular, [tPB["localReadStrideCoalesced"]] ) tPA["localReadInstruction"] = instructions["LocalRead"][localReadInstructionIdxA] if kernel["UnrollMajorLDSB"]: localReadWidth = (kernel["NumDotElements"] * tPB["bpeDS"]) // self.states.bpr localReadInstructionIdxB = \ self.selectMemoryInstruction("LocalRead", localReadWidth, \ False, \ tPB["localRead2Coalesced"], localRead2Perpendicular, [tPB["localReadStrideCoalesced"]] ) tPB["localReadInstruction"] = instructions["LocalRead"][localReadInstructionIdxB] imod.add(self.lraTileAssignment(kernel, tPA, tPB)) imod.add(self.lraFinalOffset(kernel, tPA)) imod.add(self.lraDeclareAddresses(kernel, tPA)) imod.add(self.lraFinalOffset(kernel, tPB)) imod.add(self.lraDeclareAddresses(kernel, tPB)) return imod ############################################################################## # Local Write in Prefetch Pass (PreLoop): Do It A/B ############################################################################## def preLoopLocalWriteDo(self, kernel, tPA, tPB): imod = Module() LWDoMod = imod.add(Module()) LWDoA = self.localWriteDo(kernel, tPA) if self.do["LocalWrite%s"%tPA["tensorChar"]] else Module() if ("MX" in tPA) and self.do["LocalWrite%s"%tPA["MX"]["tensorChar"]]: LWDoMXSA = self.localWriteDo(kernel, tPA["MX"]) else: LWDoMXSA = Module() if ("MX" in tPB) and self.do["LocalWrite%s"%tPB["MX"]["tensorChar"]]: LWDoMXSB = self.localWriteDo(kernel, tPB["MX"]) else: LWDoMXSB = Module() LWDoB = self.localWriteDo(kernel, tPB) if self.do["LocalWrite%s"%tPB["tensorChar"]] else Module() LWDoMod.addComment1("local write a") LWDoMod.add(LWDoA) if ("MX" in tPA): LWDoMod.addComment1("local write mxsa") LWDoMod.add(LWDoMXSA) if ("MX" in tPB): LWDoMod.addComment1("local write mxsb") LWDoMod.add(LWDoMXSB) LWDoMod.addComment1("local write b") LWDoMod.add(LWDoB) return imod def preLoopLocalWriteDoMX(self, kernel, tPA, tPB): imod = Module() LWDoMod = imod.add(Module()) if ("MX" in tPA) and self.do["LocalWrite%s"%tPA["MX"]["tensorChar"]]: LWDoMXSA = self.localWriteDo(kernel, tPA["MX"]) else: LWDoMXSA = Module() if ("MX" in tPB) and self.do["LocalWrite%s"%tPB["MX"]["tensorChar"]]: LWDoMXSB = self.localWriteDo(kernel, tPB["MX"]) else: LWDoMXSB = Module() if ("MX" in tPA): LWDoMod.addComment1("local write mxsa") LWDoMod.add(LWDoMXSA) if ("MX" in tPB): LWDoMod.addComment1("local write mxsb") LWDoMod.add(LWDoMXSB) return imod ############################################################################## # Local Write: Do It A/B ############################################################################## def localWriteDo(self, kernel, tP, swapAB=0): ti = rocIsa.getInstance() tc = tP["tensorChar"] imod = Module() isBpeInputLarger = True if tP["bpeGR"] > tP["bpe"] else False def localWriteBody(tP): tc = tP["tensorChar"] instruction = tP["localWriteInstruction"] numBlocks = instruction.numBlocks numOffsets = instruction.numOffsets blockWidth = instruction.blockWidth #offsetMultiplier = instruction.offsetMultiplier g2lIdx = 0 tmpLocalWriteAddr = -1 isAB = tc in ('A', 'B') numVgprG2L = 0 if tc == 'A': numVgprG2L = self.states.a.numVgprG2L elif tc == 'MXSA': numVgprG2L = self.states.mxsa.numVgprG2L elif tc == 'B': numVgprG2L = self.states.b.numVgprG2L elif tc == 'MXSB': numVgprG2L = self.states.mxsb.numVgprG2L else: numVgprG2L = self.states.m.numVgprG2L # using _ds_store_b8: need one more vgpr space to do lshr tmpVgprOffset = (numVgprG2L / 2) if (blockWidth == 0.25) else 0 # if transposing, positions of sPerp and sPara are transposed instructionCnt = 0 Hcvt2BMap = {} g2lIdxDict = {} regTmpVgprBlock = None if swapAB == 1: destVgprPrefix = "G2L%s2"%(tc) else: destVgprPrefix = "G2L%s"%(tc) for perp in range(0, tP["nrp"]): localWriteCode = imod.add(Module("LocalWrite%u perp=%d"%(instructionCnt,perp))) lwa = "LocalWriteAddr%s"%tc # default for para in range(0, tP["nrc"]): if para>=1: localWriteCode = imod.add(Module("LocalWrite%u perp=%d para=%d"%(instructionCnt,perp,para))) for s in range(0, max(tP["nwcv"],tP["nwpv"])//tP["nwcvpi"]): localWriteCVTCode = Module() sPerp = 0 sPara = 0 needToSplitMetadata = False metadataScalar = 1.0 if tP["tlu"] != kernel["UnrollMajorLDS%s" % tP["tensorChar"]]: if tP["wtc"]: sPerp = s else: if tP["wtc"]: sPara = s if tP["isM"]: needToSplitMetadata = True metadataScalar = max(1, self.states.bpr // int(tP["glvw"] * tP["bpeGR"])) #print("perp:{}/{} para:{}/{} sPerp:{} sPara:{}".format(perp,tP["nrp"],para,tP["nrc"],sPerp,sPara)) (offset, i, comment) = self.calculateLdsWriteOffset(perp, para, sPerp, sPara, kernel, tP) # Need refactor, the pattern < glvw in fp8 is not the same as the original. # Thus the offset calculation here does not match global read. if tP["bpeGR"] == 4 and tP["bpeDS"] == 2: # float to fp16/bf16, keep idx in the lower idx. g2lIdx = int((i * blockWidth) * metadataScalar) * 2 elif tP["glvw"] <= 2: # In Metadata use metadataScalar to group # of i g2lIdx = (i * blockWidth) * metadataScalar if isBpeInputLarger: g2lIdx *= (tP["bpeGR"]// tP["bpeDS"]) g2lIdx = int(g2lIdx) else: g2lIdx = int((i * blockWidth) * metadataScalar) if isBpeInputLarger: g2lIdx *= int(tP["bpeGR"]// tP["bpeDS"]) graIdx = i * self.states.rpgo if kernel["BufferLoad"] else i * self.states.rpga if tP["isM"]: if not needToSplitMetadata: g2lIdx = graIdx * ceil(blockWidth) #TODO: remove this if upcoming compiler changes getting merged if tP["globalReadInstruction"].blockWidth == 3: g2lIdx = graIdx * 4 # If g2lIdx is already in the dict and blockWidth < 1, the data may # be packed into one register. instHi = 0 if g2lIdx in g2lIdxDict: g2lIdxDict[g2lIdx] += 1 else: g2lIdxDict[g2lIdx] = 0 instHi = g2lIdxDict[g2lIdx] if self.states.archCaps["HasEccHalf"] or not self.states.asmCaps["HasWMMA_V1"]: eccinstHi = instHi # FIXME: Workaround, unique pattern in 8bit + glvw == 2... if tP["bpeDS"] == tP["bpeGR"] and (tP["globalReadInstruction"].totalWidth) == 0.5 and (blockWidth == 0.25) and isAB: eccinstHi = i // 2 eccBpe = tP["bpeDS"] if kernel["ConvertAfterDS"] else max(tP["bpeGR"], tP["bpe"]) eccOffset = _getEccOffset(tP["globalReadInstruction"].totalWidth, bpr=self.states.bpr, bpe=eccBpe, \ glvw=tP["glvw"], idx=eccinstHi, numVgprG2L=numVgprG2L) else: eccOffset = 0 # TODO- INT8: check uDu if (blockWidth == 0.25) and ((s % 4) == 0) and (isAB or needToSplitMetadata): src = destVgprPrefix + "+%u" % (g2lIdx + eccOffset) dst = destVgprPrefix + "+%u+%u" % (tmpVgprOffset, g2lIdx) if tP["bpeDS"] != tP["bpeGR"]: if kernel["ProblemType"]["DataType%s"%tc].isHalf(): if tP["glvw"] > 1: dst = destVgprPrefix + "+%u+%u" % (tmpVgprOffset, g2lIdx // 2) localWriteCVTCode.add(VPackF16toB32(dst=vgpr(dst), src0=vgpr(src), src1=vgpr(destVgprPrefix + "+%u" % (g2lIdx+1)), \ vop3=VOP3PModifiers(op_sel=[1,1,0]), comment="Pack with neighbor")) localWriteCVTCode.add(VPackF16toB32(dst=vgpr(src), src0=vgpr(src), src1=vgpr(destVgprPrefix + "+%u" % (g2lIdx+1)), \ vop3=VOP3PModifiers(op_sel=[0,0,0]), comment="Pack with neighbor")) else: printExit("Unsupported combination DataType%s (%s) -> MacDataType (%s)"%(tc, kernel["ProblemType"]["DataType%s"%tc].toChar(), kernel["ProblemType"]["MacDataType%s"%tc if (tc=='A' or tc=='B') else "DataType"].toChar())) elif tP["glvw"] > 1: localWriteCVTCode.add(VMovB32(dst=vgpr(dst), src=vgpr(src), comment="another VGPR storing lshr 8-bit value")) localWriteCVTCode.add(VLShiftRightB32(dst=vgpr(dst), shiftHex=hex(8), src=vgpr(dst), comment="G2L Vpgr >> 8")) if kernel["ExpertSchedulingMode"] > 0: localWriteCVTCode.add(SWaitAlu(va_vdst=0, comment="wait for writes to complete")) #comment = "Reg -> L %u_%u_%u_%u"%(para, sPara, perp, sPerp) isHigh16Bits = False isCvtHighBits = False datatype = kernel["ProblemType"]["DataType%s"%tc] if kernel["ConvertAfterDS"] else kernel["ProblemType"]["MacDataType%s"%tc if (tc=='A' or tc=='B') else "DataType"] if (datatype.isHalf() or datatype.isBFloat16()) and isAB: if s%2==1: isHigh16Bits = True if (blockWidth == 0.5) and (instHi % 2 == 1): isHigh16Bits = True if kernel["ProblemType"]["DataType%s"%tc].isAnyFloat8(): if g2lIdx%2 == 1: isCvtHighBits = True # | hi16 | hi16 | | | # | hi8 | | hi8 | | ############################################# # VGPR: |---w4---|---w3---|---w2---|---w1---| -> b8_d16: get w1 / _b8_d16_hi: get w3 # LSHR: |--------|---w4---|--------|---w2---| -> b8_d16: get w2 / _b8_d16_hi: get w4 elif datatype.isInt8() or datatype.is8bitFloat() or (not isAB): isHigh16Bits = (s % 4) > 1 # 2,3 # TODO # if tP["glvw"]==1 and instructionCnt%2==1: # isHigh16Bits = True paramList = [] numsOfRegister = [] globalBlockWidth = tP["globalReadInstruction"].totalWidth # print("tc = ", tc, ", numBlocks = ", numBlocks) # print("regs: ", self.vgprs.globalReadRegisters[tc]) # print("destVgprPrefix = ", destVgprPrefix, ", blockWidth = ", blockWidth) for _ in range(0, numBlocks): # FIXME: In the future all registers should pass from global read instead of recalculate them if globalBlockWidth == 1 and blockWidth == 0.5 and tP["glvw"] == 1: if isHigh16Bits: paramList.append(vgpr(destVgprPrefix + "+%u"%(g2lIdx + 1))) else: paramList.append(vgpr(destVgprPrefix + "+%u"%(g2lIdx))) numsOfRegister.append(blockWidth) elif globalBlockWidth == blockWidth and tP["glvw"] == 1 and i < len(self.vgprs.globalReadRegisters[tc]): # print("destVgprPrefix = ", destVgprPrefix) # print("tc = ", tc, ", i", i) # print(self.vgprs.globalReadRegisters[tc][i]) # print("blockWidth = ", blockWidth) # print(vgpr(destVgprPrefix + "+%u"%(self.vgprs.globalReadRegisters[tc][i]), blockWidth)) # print("i = ", i) # print("_ = ", _) paramList.append(vgpr(destVgprPrefix + "+%u"%(self.vgprs.globalReadRegisters[tc][i]), blockWidth)) # print("DONE") elif blockWidth == 1: paramList.append(vgpr(destVgprPrefix + "+%u"%(g2lIdx))) numsOfRegister.append(1) elif blockWidth == 0.25 and ((s % 2) == 1): # Int8, s = 1 or 3 (high8Bits) if tP["bpeDS"] != tP["bpeGR"] and tmpVgprOffset != 0: paramList.append(vgpr(destVgprPrefix + "+%u+%u"%(tmpVgprOffset, g2lIdx // 2))) else: paramList.append(vgpr(destVgprPrefix + "+%u+%u"%(tmpVgprOffset, g2lIdx))) numsOfRegister.append(1) else: paramList.append(vgpr(destVgprPrefix + "+%u"%(g2lIdx + eccOffset), blockWidth)) numsOfRegister.append(blockWidth) if self.db["ForceInputValue%s"%tc]: localWriteCVTCode.add(VMovB32(dst=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), src=self.db["ForceValue%s"], comment="ForceInputValue")) if (kernel["ProblemType"]["MacDataType%s"%tc if (tc=='A' or tc=='B') else "DataType"].isBFloat16() and kernel["ProblemType"]["DataType%s"%tc].isHalf()) and isAB: numIters = 1 if blockWidth <= 1 else blockWidth vgprTmp = self.vgprPool.checkOut(2, tag="localWriteDo_vgprTmp") for iter in range(0, numIters): f16Tobf16Idx = g2lIdx + iter if f16Tobf16Idx in Hcvt2BMap: Hcvt2BMap[f16Tobf16Idx] += 2 else: Hcvt2BMap[f16Tobf16Idx] = 0 f16Tobf16Idx += Hcvt2BMap[f16Tobf16Idx] localWriteCVTCode.add(ECvtF16toF32(dst=vgpr(vgprTmp), src=vgpr(destVgprPrefix + "+%u"%(f16Tobf16Idx)), sel=HighBitSel.LOW)) localWriteCVTCode.add(ECvtF16toF32(dst=vgpr(vgprTmp+1), src=vgpr(destVgprPrefix + "+%u"%(f16Tobf16Idx)), sel=HighBitSel.HIGH)) localWriteCVTCode.add(VPackF16toB32(dst=vgpr(destVgprPrefix + "+%u"%(f16Tobf16Idx)), src0=vgpr(vgprTmp), src1=vgpr(vgprTmp+1), vop3=VOP3PModifiers(op_sel=[1,1,0]))) self.vgprPool.checkIn(vgprTmp) for oIdx in range(0, numOffsets): paramList.append(offset) #print "offset", offset # Need cvt if tP["bpeDS"] != tP["bpeGR"]: assert numBlocks == 1 if (kernel["ProblemType"]["DataType%s"%tc].isSingle() and kernel["ProblemType"]["MacDataType%s"%tc if (tc=='A' or tc=='B') else "DataType"].isHalf()): newBlockWidth = (tP["bpeGR"] / tP["bpe"]) * blockWidth if newBlockWidth == 1: new_src = deepcopy(paramList[0]) if isHigh16Bits: new_src.regName.addOffset(1) localWriteCVTCode.add(ECvtF32toF16(dst=paramList[0], src=new_src, sel=HighBitSel.HIGH if isHigh16Bits else HighBitSel.LOW, comment="convert C to fp16")) else: for vi in range(0, int(newBlockWidth)): localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi)), sel=HighBitSel.HIGH if vi % 2 else HighBitSel.LOW, comment="convert C to fp16")) elif (kernel["ProblemType"]["DataType%s"%tc].isHalf() and kernel["ProblemType"]["MacDataType%s"%tc if (tc=='A' or tc=='B') else "DataType"].is8bitFloat()): #HH_F8/B8/F8B8/B8F8_ toF8 = False if ((tc == "A" and kernel["ProblemType"]["MacDataTypeA"].isAnyFloat8()) or (tc == "B" and kernel["ProblemType"]["MacDataTypeB"].isAnyFloat8())): toF8 = True newBlockWidth = (tP["bpeGR"] / tP["bpe"]) * blockWidth if newBlockWidth == 0.5: # We cannot use optimized f16->fp8 instruction here as we have only one half element to be converted # and we only have a pack instruction which can convert 2 half elements to 2 f8 element if kernel["ProblemType"]["StochasticRounding"]: if self.states.asmCaps["v_prng_b32"]: vgprTmp = self.vgprPool.checkOutAligned(2, 2, tag="localWriteDo_vgprTmp1") else: vgprTmp = self.vgprPool.checkOutAligned(4, 2, tag="localWriteDo_vgprTmp2") else: vgprTmp = self.vgprPool.checkOutAligned(1, 2, tag="localWriteDo_vgprTmp3") sel = 1 if isHigh16Bits else 0 localWriteCVTCode.add(ECvtF16toF32(dst=vgpr(vgprTmp), src=paramList[0], sel=HighBitSel.HIGH if isHigh16Bits else HighBitSel.LOW, comment="convert to F32")) # ScaleA/B if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and kernel["ProblemType"]["DataType%s"%tc].numRegisters() > kernel["ProblemType"]["MacDataType%s"%tc].numRegisters(): localWriteCVTCode.add(VMulF32(dst=vgpr(vgprTmp), src0=vgpr(vgprTmp), src1=sgpr("Scale%s"%tc), comment="Input *= scale %s"%tc)) if kernel["ProblemType"]["StochasticRounding"]: vRand = vgprTmp+1 #seed if self.states.asmCaps["v_prng_b32"]: localWriteCVTCode.add(VPrngB32(dst=vgpr(vRand),src=vgpr(vgprTmp),comment="Psudo Random Number Generator")) else: vTemp0 = vgprTmp+2 vTemp1 = vgprTmp+3 if self.states.asmCaps["HasVgprMSB"]: localWriteCVTCode.add(PseudoRandomGeneratorModule(vRand, vgprTmp, vTemp0, vTemp1)) else: localWriteCVTCode.add(MacroInstruction(name="PRND_GENERATOR", args=[vRand, vgprTmp, vTemp0, vTemp1])) if (toF8): localWriteCVTCode.add(VCvtSRF32toFP8(dst=paramList[0], src0=vgpr(vgprTmp), src1=vgpr(vRand), sels=[sel], comment="Convert to FP8")) else: localWriteCVTCode.add(VCvtSRF32toBF8(dst=paramList[0], src0=vgpr(vgprTmp), src1=vgpr(vRand), sels=[sel], comment="Convert to BF8")) else: if (toF8): localWriteCVTCode.add(VCvtPkF32toFP8(dst=paramList[0], src0=vgpr(vgprTmp), src1=vgpr(vgprTmp), vop3=VOP3PModifiers(op_sel=[0,0,sel]), comment="Convert to FP8")) else: localWriteCVTCode.add(VCvtPkF32toBF8(dst=paramList[0], src0=vgpr(vgprTmp), src1=vgpr(vgprTmp), vop3=VOP3PModifiers(op_sel=[0,0,sel]), comment="Convert to BF8")) self.vgprPool.checkIn(vgprTmp) else: if kernel["ProblemType"]["StochasticRounding"]: if self.states.asmCaps["v_prng_b32"] : vgprTmp = self.vgprPool.checkOutAligned(3, 2, tag="localWriteDo_vgprTmp4") else: vgprTmp = self.vgprPool.checkOutAligned(5, 2, tag="localWriteDo_vgprTmp5") else: vgprTmp = self.vgprPool.checkOutAligned(2, 2, tag="localWriteDo_vgprTmp6") vgprTmp2 = vgprTmp + 1 for vi in range(0, int(newBlockWidth)): sel = 1 if vi %2 == 1 else 0 if self.states.asmCaps["Hascvtfp8_f16"] and not kernel["ProblemType"]["StochasticRounding"] and not kernel["ProblemType"]["UseScaleAB"] == "Scalar": if (toF8): localWriteCVTCode.add(VCvtScalePkF16toFP8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[0,0,sel]), comment="convert F16 to F8")) else: localWriteCVTCode.add(VCvtScalePkF16toBF8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[0,0,sel]), comment="convert F16 to BF8")) else: localWriteCVTCode.add(ECvtF16toF32(dst=vgpr(vgprTmp), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi)), sel=HighBitSel.LOW, comment="convert to F32")) localWriteCVTCode.add(ECvtF16toF32(dst=vgpr(vgprTmp2), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi)), sel=HighBitSel.HIGH, comment="convert to F32")) if kernel["ProblemType"]["StochasticRounding"]: # ScaleA/B, sgpr upper is dummy. if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and kernel["ProblemType"]["DataType%s"%tc].numRegisters() > kernel["ProblemType"]["MacDataType%s"%tc].numRegisters(): localWriteCVTCode.add(VMulPKF32S(dst=vgpr(vgprTmp, 2), src0=vgpr(vgprTmp, 2), src1=sgpr("Scale%s"%tc, 2), vop3=VOP3PModifiers(op_sel_hi=[1,0,1]), comment="Input *= scale %s"%tc)) vRand = vgprTmp+2 if self.states.asmCaps["v_prng_b32"]: localWriteCVTCode.add(VPrngB32(dst=vgpr(vRand),src=vgpr(vgprTmp),comment="Psudo Random Number Generator")) else: vTemp0 = vgprTmp+3 vTemp1 = vgprTmp+4 if self.states.asmCaps["HasVgprMSB"]: localWriteCVTCode.add(PseudoRandomGeneratorModule(vRand, vgprTmp, vTemp0, vTemp1)) else: localWriteCVTCode.add(MacroInstruction(name="PRND_GENERATOR", args=[vRand, vgprTmp, vTemp0, vTemp1])) if (toF8): localWriteCVTCode.add(VCvtSRF32toFP8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(vgprTmp), src1=vgpr(vRand), sels=[0,sel], comment="Convert to FP8")) if self.states.asmCaps["v_prng_b32"]: localWriteCVTCode.add(VPrngB32(dst=vgpr(vRand),src=vgpr(vgprTmp2),comment="Psudo Random Number Generator")) else: if self.states.asmCaps["HasVgprMSB"]: localWriteCVTCode.add(PseudoRandomGeneratorModule(vRand, vgprTmp2, vTemp0, vTemp1)) else: localWriteCVTCode.add(MacroInstruction(name="PRND_GENERATOR", args=[vRand, vgprTmp2, vTemp0, vTemp1])) localWriteCVTCode.add(VCvtSRF32toFP8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(vgprTmp2), src1=vgpr(vRand), sels=[1,sel], comment="Convert to FP8")) else: localWriteCVTCode.add(VCvtSRF32toBF8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(vgprTmp), src1=vgpr(vRand), sels=[0,sel], comment="Convert to BF8")) if self.states.asmCaps["v_prng_b32"]: localWriteCVTCode.add(VPrngB32(dst=vgpr(vRand),src=vgpr(vgprTmp2),comment="Psudo Random Number Generator")) else: if self.states.asmCaps["HasVgprMSB"]: localWriteCVTCode.add(PseudoRandomGeneratorModule(vRand, vgprTmp2, vTemp0, vTemp1)) else: localWriteCVTCode.add(MacroInstruction(name="PRND_GENERATOR", args=[vRand, vgprTmp2, vTemp0, vTemp1])) localWriteCVTCode.add(VCvtSRF32toBF8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(vgprTmp2), src1=vgpr(vRand), sels=[1,sel], comment="Convert to BF8")) else: # ScaleA/B, sgpr upper is dummy. if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and kernel["ProblemType"]["DataType%s"%tc].numRegisters() > kernel["ProblemType"]["MacDataType%s"%tc].numRegisters(): localWriteCVTCode.add(VMulPKF32S(dst=vgpr(vgprTmp, 2), src0=vgpr(vgprTmp, 2), src1=sgpr("Scale%s"%tc, 2), vop3=VOP3PModifiers(op_sel_hi=[1,0,1]), comment="Input *= scale %s"%tc)) if (toF8): localWriteCVTCode.add(VCvtPkF32toFP8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(vgprTmp), src1=vgpr(vgprTmp2), vop3=VOP3PModifiers(op_sel=[0,0,sel]), comment="Convert to FP8")) else: localWriteCVTCode.add(VCvtPkF32toBF8(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(vgprTmp), src1=vgpr(vgprTmp2), vop3=VOP3PModifiers(op_sel=[0,0,sel]), comment="Convert to BF8")) self.vgprPool.checkIn(vgprTmp) elif (kernel["ProblemType"]["DataType%s"%tc].isAnyFloat8() and kernel["ProblemType"]["MacDataType%s"%tc].isHalf()): noSDWA = ti.getArchCaps()["NoSDWA"] newBlockWidth = tP["globalReadInstruction"].blockWidth if newBlockWidth == 0.25: new_src = deepcopy(paramList[0]) if tP["glvw"] == 1: vgprTmp = self.vgprPool.checkOut(1, tag="localWriteDo_vgprTmp7") new_src.regName.addOffset(tP["shiftGR"]) if self.states.asmCaps["Hascvtf16_fp8_sf32"]: sel = [0,1,1,0] if isHigh16Bits else [0,0,0,0] localWriteCVTCode.add(VCvtScaleFP8toF16(dst=paramList[0], src=new_src, scale=0x3f800000, vop3=VOP3PModifiers(op_sel=sel), comment="A convert fp8 to f16")) else: sel = 1 if isHigh16Bits else 0 if noSDWA: localWriteCVTCode.add(VCvtFP8toF32(dst=vgpr(vgprTmp), src=new_src, vop3=VOP3PModifiers(byte_sel=[sel]), comment="convert C to fp32")) else: src_sel = SelectBit.BYTE_2 if isHigh16Bits else SelectBit.BYTE_0 localWriteCVTCode.add(VCvtFP8toF32(dst=vgpr(vgprTmp), src=new_src, sdwa=SDWAModifiers(src0_sel=src_sel), comment="convert C to fp32")) localWriteCVTCode.add(ECvtF32toF16(dst=paramList[0], src=vgpr(vgprTmp), sel=HighBitSel.HIGH if isHigh16Bits else HighBitSel.LOW, comment="convert C to fp16")) self.vgprPool.checkIn(vgprTmp) else: if isCvtHighBits and isHigh16Bits: src_sel = SelectBit.BYTE_3 elif isHigh16Bits and (not isCvtHighBits): new_src.regName.addOffset(tP["shiftGR"]) src_sel = SelectBit.BYTE_1 elif (not isHigh16Bits) and isCvtHighBits: src_sel = SelectBit.BYTE_2 else: new_src.regName.addOffset(tP["shiftGR"]) src_sel = SelectBit.BYTE_0 dst_sel = SelectBit.WORD_1 if isHigh16Bits else SelectBit.WORD_0 if self.states.asmCaps["Hascvtf16_fp8_sf32"]: if isHigh16Bits: if isCvtHighBits: sel = [1,1,1,0] else: sel = [0,1,1,0] else: if isCvtHighBits: sel = [0,1,0,0] else: sel = [0,0,0,0] if new_src == paramList[0]: if self.states.asmCaps["Hascvtf16_fp8_sf32"]: localWriteCVTCode.add(VCvtScaleFP8toF16(dst=paramList[0], src=new_src , scale=0x3f800000, vop3=VOP3PModifiers(op_sel=sel), comment="B convert fp8 to f16")) else: if src_sel == SelectBit.BYTE_0 or src_sel == SelectBit.BYTE_2: if regTmpVgprBlock == None: regTmpVgprBlock = self.vgprPool.checkOutAligned(2, 2) src_sel2 = SelectBit.WORD_0 if src_sel == SelectBit.BYTE_0 else SelectBit.WORD_1 localWriteCVTCode.add(ECvtPkFP8toF32(dst=vgpr(regTmpVgprBlock, 2), src=new_src, sel=HighBitSel.HIGH if src_sel2 == SelectBit.WORD_1 else HighBitSel.LOW, comment="convert C to fp32")) localWriteCVTCode.add(ECvtF32toF16(dst=paramList[0], src=vgpr(regTmpVgprBlock), sel=HighBitSel.HIGH if isHigh16Bits else HighBitSel.LOW, comment="convert C to fp16")) else: localWriteCVTCode.add(ECvtF32toF16(dst=paramList[0], src=vgpr(regTmpVgprBlock + 1), sel=HighBitSel.HIGH if isHigh16Bits else HighBitSel.LOW, comment="convert C to fp16")) else: vgprTmp = self.vgprPool.checkOut(1, tag="localWriteDo_vgprTmp8") if self.states.asmCaps["Hascvtf16_fp8_sf32"]: localWriteCVTCode.add(VCvtScaleFP8toF16(dst=paramList[0], src=new_src, scale=0x3f800000, vop3=VOP3PModifiers(op_sel=sel), comment="C convert fp8 to f16")) else: if noSDWA: localWriteCVTCode.add(VCvtFP8toF32(dst=vgpr(vgprTmp), src=new_src, vop3=VOP3PModifiers(op_sel=[1 if isHigh16Bits else 0]), comment="convert C to fp32")) else: localWriteCVTCode.add(VCvtFP8toF32(dst=vgpr(vgprTmp), src=new_src, sdwa=SDWAModifiers(src0_sel=src_sel), comment="convert C to fp32")) localWriteCVTCode.add(ECvtF32toF16(dst=paramList[0], src=vgpr(vgprTmp), sel=HighBitSel.HIGH if isHigh16Bits else HighBitSel.LOW, comment="convert C to fp16")) self.vgprPool.checkIn(vgprTmp) elif newBlockWidth == 0.5: vgprTmp = self.vgprPool.checkOutAligned(2, 2, tag="localWriteDo_vgprTmp9") src_sel = SelectBit.WORD_1 if isCvtHighBits else SelectBit.WORD_0 modNum = max(1, int(newBlockWidth / blockWidth)) if self.states.asmCaps["Hascvtf16_fp8_sf32"]: sel = [1,0,0,0] if isCvtHighBits else [0,0,0,0] localWriteCVTCode.add(VCvtScalePkFP8toF16(dst=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), src=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=sel), comment="D convert fp8 to f16")) else: if (not isHigh16Bits) and (g2lIdx % modNum == 0): localWriteCVTCode.add(ECvtPkFP8toF32(dst=vgpr(vgprTmp, 2), src=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), sel=HighBitSel.HIGH if src_sel == SelectBit.WORD_1 else HighBitSel.LOW, comment="convert to F32")) localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), src=vgpr(vgprTmp), sel=HighBitSel.LOW, comment="Convert to FP16")) if (newBlockWidth <= blockWidth): localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), src=vgpr(vgprTmp+1), sel=HighBitSel.HIGH, comment="Convert to FP16")) elif (newBlockWidth > blockWidth): localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u"%(g2lIdx)), src=vgpr(vgprTmp+1), sel=HighBitSel.HIGH, comment="Convert to FP16")) self.vgprPool.checkIn(vgprTmp) else: modNum = max(1, int(newBlockWidth / blockWidth)) vgprTmp = self.vgprPool.checkOutAligned(2, 2, tag="localWriteDo_vgprTmp10") if (not isHigh16Bits) and (newBlockWidth <= blockWidth): for vi in range(0, int(newBlockWidth)): if self.states.asmCaps["Hascvtf16_fp8_sf32"]: localWriteCVTCode.add(VCvtScaleFP8toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx+tP["shiftGR"], vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[0,0,0,0]), comment="E convert fp8 to f16")) localWriteCVTCode.add(VCvtScaleFP8toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx+tP["shiftGR"], vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[1,0,1,0]), comment="E convert fp8 to f16")) localWriteCVTCode.add(VCvtScaleFP8toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2 + 1)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx+tP["shiftGR"], vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[0,1,0,0]), comment="E convert fp8 to f16")) localWriteCVTCode.add(VCvtScaleFP8toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2 + 1)), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx+tP["shiftGR"], vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[1,1,1,0]), comment="E convert fp8 to f16")) else: localWriteCVTCode.add(ECvtPkFP8toF32(dst=vgpr(vgprTmp, 2), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx+tP["shiftGR"], vi)), sel=HighBitSel.LOW, comment="convert to F32")) localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2)), src=vgpr(vgprTmp), sel=HighBitSel.LOW, comment="Convert to FP16")) localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2)), src=vgpr(vgprTmp+1), sel=HighBitSel.HIGH, comment="Convert to FP16")) localWriteCVTCode.add(ECvtPkFP8toF32(dst=vgpr(vgprTmp, 2), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx+tP["shiftGR"], vi)), sel=HighBitSel.HIGH, comment="convert to F32")) localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2 + 1)), src=vgpr(vgprTmp), sel=HighBitSel.LOW, comment="Convert to FP16")) localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi * 2 + 1)), src=vgpr(vgprTmp+1), sel=HighBitSel.HIGH, comment="Convert to FP16")) else: if (not isHigh16Bits): idxMod = g2lIdx % modNum g2lIdxTmp = g2lIdx - idxMod vi = idxMod // 2 interOffset = 0 if idxMod % 2 == 0 else 1 if self.states.asmCaps["Hascvtf16_fp8_sf32"]: sel = 0 if idxMod % 2 == 0 else 1 localWriteCVTCode.add(VCvtScalePkFP8toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdxTmp, vi * 2 + interOffset)),\ src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdxTmp+tP["shiftGR"], vi)), scale=0x3f800000,\ vop3=VOP3PModifiers(op_sel=[0,0,0,sel]), comment="F convert fp8 to f16")) else: localWriteCVTCode.add(ECvtPkFP8toF32(dst=vgpr(vgprTmp, 2), src=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdxTmp+tP["shiftGR"], vi)), sel=HighBitSel.HIGH if idxMod % 2 else HighBitSel.LOW, comment="convert to F32")) localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdxTmp, vi * 2 + interOffset)), src=vgpr(vgprTmp), sel=HighBitSel.LOW, comment="Convert to FP16")) if blockWidth != 0.5: localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdxTmp, vi * 2 + interOffset)), src=vgpr(vgprTmp+1), sel=HighBitSel.HIGH, comment="Convert to FP16")) elif blockWidth == 0.5: localWriteCVTCode.add(ECvtF32toF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdxTmp, vi * 2 + interOffset)), src=vgpr(vgprTmp+1), sel=HighBitSel.HIGH, comment="Convert to FP16")) self.vgprPool.checkIn(vgprTmp) elif (kernel["ProblemType"]["DataType%s"%tc].isSingle() and kernel["ProblemType"]["DataType"].isBFloat16()): newBlockWidth = (tP["bpeGR"] / tP["bpe"]) * blockWidth if tP["glvw"] == 1: # use the same VGPR localWriteCVTCode.add(VCvtPkF32toBF16(dst=paramList[0], src0=paramList[0], src1=paramList[0], comment="convert one fp32 to bf16 to the same VGPR")) elif newBlockWidth == 1: new_src = deepcopy(paramList[0]) if isHigh16Bits: new_src.regName.addOffset(1) localWriteCVTCode.add(VCvtPkF32toBF16(dst=paramList[0], src0=new_src, src1=new_src, comment="convert one fp32 to bf16 idx=%d"%g2lIdx)) else: for vi in range(0, int(newBlockWidth)): if vi%2 == 0: localWriteCVTCode.add(VCvtPkF32toBF16(dst=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi//2)), src0=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi)), src1=vgpr(destVgprPrefix + "+%u+%u"%(g2lIdx, vi+1)), comment="convert/pack fp32 to bf16")) else: printExit("Unsupported combination DataType%s (%s) -> DataType (%s)"%(tc, kernel["ProblemType"]["DataType%s"%tc].toChar(), kernel["ProblemType"]["MacDataType%s"%tc].toChar())) LocalWriteX = tP["localWriteInstruction"].getInst(isHigh16Bits) localWriteMemToken = [self.states.ldsWriteTokenIdx] if len(localWriteMemToken) == 1: memTokenComment = "sync LDS%u"%(localWriteMemToken[0]) else: memTokenComment = "sync LDS %s"%(localWriteMemToken) commentWithMemToken = "%s %s"%(comment, memTokenComment) if numBlocks == 1: addrIdx = paramList[1] // 65536 olwa = "LocalWriteAddr%s+%u"%(tc, addrIdx) paramList[1] -= addrIdx * 65536 ds = DSModifiers(na=1, offset=paramList[1]) writeInst = LocalWriteX(dstAddr=vgpr(olwa), src=paramList[0], ds=ds, comment=commentWithMemToken) else: ds = DSModifiers(na=2, offset0=paramList[2], offset1=paramList[3]) writeInst = LocalWriteX(dstAddr=vgpr(lwa), src0=paramList[0], src1=paramList[1], ds=ds, comment=commentWithMemToken) # Attach LDS memory token to local write instructions so downstream # StinkyTofu passes can track local write dependencies. writeInst.setMemToken(MemTokenData(localWriteMemToken)) if self.do["LocalWriteCVT"]: localWriteCode.add(localWriteCVTCode) if self.do["LocalWrite%s"%tc]: localWriteCode.add(writeInst) instructionCnt += 1 if blockWidth < 8 else 2 if regTmpVgprBlock != None: self.vgprPool.checkIn(regTmpVgprBlock) if tmpLocalWriteAddr != -1: self.vgprPool.checkIn(tmpLocalWriteAddr) #if vgprFp32NanInfFlag != None: #self.vgprPool.checkIn(vgprFp32NanInfFlag) if kernel["ProblemType"]["Sparse"] and kernel["DirectToVgprSparseMetadata"]: miWaveTile = kernel["MIWaveTileB"] if (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) else kernel["MIWaveTileA"] if (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) else 0 if kernel["PrefetchGlobalRead"] >= 2 and miWaveTile > 0: #vpgr to store metadata offsetBlk = tP["metadataWriteSwapByteOffset"] instructionCnt = -1 for wtIdx in range(0, miWaveTile): for unrollIdx in range(0, kernel["LoopIters"]): instructionCnt +=1 localWriteCode = imod.add(Module("MetadataWrite%u "%(instructionCnt))) localWriteCode.add(VMovB32( \ vgpr("ValuMetadata+%u+%u"%(offsetBlk, (wtIdx * kernel["LoopIters"]+unrollIdx))), \ vgpr("ValuMetadata+%u+%u"%(self.states.m.numVgprValuPerBlock * 2, (wtIdx*kernel["LoopIters"]+unrollIdx))), \ comment="copy ValuMetadata from blk2")) if 0 and tP["isB"]: # post-lds-write localWriteCode.add(SWaitCnt(dscnt=0, vlcnt=0, vscnt=0, comment="")) localWriteCode.add(SBarrier(comment="dump LDS")) localWriteCode.add(self.getCmpAssert(self.asmAssert.ne, sgpr("WorkGroup0"),1)) # Enable local write if not DTL or using nonDTL loads in tail loop if (not kernel["DirectToLds%s"%tc]) or \ ((tc in ("A","B","MXSA","MXSB") and kernel["NonDTLTailLoop%s"%tc]) and self.states.inTailLoop): # Skip local write if DTVA or DTVB if not ((tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc]): localWriteBody(tP) if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"] and \ (not kernel["DirectToLdsMetadata"] or (kernel["NonDTLTailLoopMetadata"] and self.states.inTailLoop)): if tP["is_sparse"]: localWriteBody(tP["tpsMetadata"]) return imod ############################################################################## # Local Read: Swap Offsets A/B # internalPointerSwap: swap internally tracked offsets - rather than # emit specific instructions to do the pointer swap ############################################################################## def localReadSwapOffsets(self, kernel, internalPointerSwap, tP): tc=tP["tensorChar"] if (not self.do["LocalRead%s"%tc]): return Module("localReadSwapOffsets (no local read)") if kernel["1LDSBuffer"] or ((tc in ("A", "B", "MXSA", "MXSB")) and kernel["DirectToVgpr%s"%tc]): # no local read code if DirectToVgpr is enabled return Module("localReadSwapOffsets (Empty)") module = Module("localReadSwapOffsets") numLra = 0 if tc == "A": numLra = self.states.a.numVgprLocalReadAddr elif tc == "B": numLra = self.states.b.numVgprLocalReadAddr elif tc == "Metadata": numLra = self.states.m.numVgprLocalReadAddr elif tc == "MXSA": numLra = self.states.mxsa.numVgprLocalReadAddr elif tc == "MXSB": numLra = self.states.mxsb.numVgprLocalReadAddr else: raise Exception(f"unsupport tc %s{tc}") if self.states.IncLdsBufSwitch: # IncLdsBufSwitch case, we do not use xor. Instead, use add and max check for round back # (numLDSBlk>=3 is for DTL (and LocalWriteUseSgpr) only) is1st = tc == "A" # so far, A is always first # LDSBufferReadInc is common for A and B. Add this only for the first one (tc=="A") if is1st: module.add(SAddU32( dst=sgpr("LDSBufferReadInc"), \ src0="LdsOneBlockSize", \ src1=sgpr("LDSBufferReadInc"), \ comment="add LDS block size to incSgpr")) module.add(SCmpEQU32( src0=sgpr("LDSBufferReadInc"), \ src1="LdsBlockEndSize", \ comment="LDSBufferReadInc == End ?")) module.add(SCMovB32( dst=sgpr("LDSBufferReadInc"), \ src=0, comment="LDSBufferReadInc loop back to 0")) module.add(VAddU32( dst=vgpr("LocalReadAddr%s"%(tc)), \ src0=sgpr("LDSBufferReadInc"), \ src1=vgpr("LocalReadAddrOrig%s"%(tc)), \ comment="LocalReadAddr = Inc + Orig")) elif internalPointerSwap or kernel["StoreSwapAddr"]: if not kernel["StoreSwapAddr"]: tP["localReadSwapByteOffset"] = 0 if tP["localReadSwapByteOffset"] else kernel["LdsOffsetA_Blk"] module.addComment1("local read swap internal offset -> %u" % tP["localReadSwapByteOffset"]) else: module.add(VXorB32( dst=vgpr("LocalReadAddr%s"%tc), \ src0=vgpr("LocalReadSwapAddr%s"%tc), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="swap Red Blk")) else: module.add(VXorB32( dst=vgpr("LocalReadAddr%s"%tc), \ src0=hex(kernel["LdsOffsetA_Blk"]), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="swap Red Blk")) for i in range(1,numLra): module.add(VAddU32( dst=vgpr("LocalReadAddr%s+%u"%(tc,i)), \ src0=(i * self.states.regCaps["maxLDSConstOffset"]), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024))) return module ############################################################################## # Local Read: Reset Offsets A/B # x % n == n & (n-1) for n power of 2 # tP[localReadOffset] maintains running count of offsets # This is called from the tail loop to reset read offsets? ############################################################################## def localReadResetOffsets(self, kernel, tP): tc = tP["tensorChar"] if not self.do["LocalRead%s"%tc]: return Module("localReadResetOffsets (no local read)") # no local read code if DirectToVgpr is enabled if kernel["1LDSBuffer"] or ((tP["isA"] or tP["isB"] or tP["isMXSA"] or tP["isMXSB"]) and kernel["DirectToVgpr%s"%tc]): return Module("localReadResetOffsets (Empty)") module = Module("localReadResetOffsets") if tP["localReadInstruction"].numOffsets == 1: tP["localReadSwapByteOffset"] = 0 module.addComment1("localReadResetOffsets") tP["localReadOffset"] = 0 module.addComment0("handled internally") numLra = 0 if tP["isA"]: numLra = self.states.a.numVgprLocalReadAddr elif tP["isB"]: numLra = self.states.b.numVgprLocalReadAddr elif tP["isMXSA"]: numLra = self.states.mxsa.numVgprLocalReadAddr elif tP["isMXSB"]: numLra = self.states.mxsb.numVgprLocalReadAddr elif tP["isM"]: numLra = self.states.m.numVgprLocalReadAddr else: raise Exception(f"unsupport tc %s{tc}") if self.states.IncLdsBufSwitch: # 3 or more LDS block case, round back to 0 and set LocalReadAddrOrig to LocalReadAddr # (numLDSBlk>=3 is for DTL (and LocalWriteUseSgpr) only) module.add(SMovB32( dst=sgpr("LDSBufferReadInc"), \ src=0, \ comment="reset incSgpr")) module.add(VMovB32( dst=vgpr("LocalReadAddr%s"%(tc)), \ src=vgpr("LocalReadAddrOrig%s"%(tc)), \ comment="set LocalReadAddrOrig to LocalReadAddr")) elif kernel["StoreSwapAddr"]: # Reset offset, by picking smaller of the two tmpvgpr = self.vgprPool.checkOut(1, tag="localReadResetOffsets_tmpvgpr") # contains other offsets module.add(VXorB32( dst=vgpr(tmpvgpr), \ src0=vgpr("LocalReadSwapAddr%s"%tc), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="Get other lds buffer offset value")) module.add(VMinI32(dst=vgpr("LocalReadAddr%s"%tc), src0=vgpr("LocalReadAddr%s"%tc), src1=vgpr(tmpvgpr), comment="Set LRA to first buffer offset")) self.vgprPool.checkIn(tmpvgpr) else: module.add(VAndB32( dst=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], 0)), \ src0=hex(kernel["LdsOffsetA_Blk"]-1), \ src1=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], 0)), \ comment="reset Red,Blk -> Red")) numLra = 0 if tP["isA"]: numLra = self.states.a.numVgprLocalReadAddr elif tP["isB"]: numLra = self.states.b.numVgprLocalReadAddr elif tP["isM"]: numLra = self.states.m.numVgprLocalReadAddr for i in range(1,numLra): module.add(VAddU32( dst=vgpr("LocalReadAddr%s+%u"%(tc,i)), \ src0=(i * self.states.regCaps["maxLDSConstOffset"]), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024))) return module ############################################################################## # Local Read: Init Pointers A/B ############################################################################## def localReadInitPointers(self, kernel, tPA, tP): tc=tP["tensorChar"] if (not self.do["LocalRead%s"%tc]) or ((tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc]): # no local read code if DirectToVgpr is enabled return Module("localReadInitPointers (Empty)") module = Module("localReadInitPointers") if tPA["localReadInstruction"].numOffsets == 1: module.addComment1("localReadInitPointers") tP["localReadOffset"] = 0 else: numLra = 0 if tP["isA"]: numLra = self.states.a.numVgprLocalReadAddr elif tP["isB"]: numLra = self.states.b.numVgprLocalReadAddr elif tP["isMXSA"]: numLra = self.states.mxsa.numVgprLocalReadAddr elif tP["isMXSB"]: numLra = self.states.mxsb.numVgprLocalReadAddr elif tP["isM"]: numLra = self.states.m.numVgprLocalReadAddr module.add(VAndB32( dst=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], 0)), \ src0=hex(kernel["LdsOffset%s_Blk"%tP["tensorChar"]]-1), \ src1=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], 0)), \ comment="init Red,Blk -> Red")) for i in range(1, numLra): module.add(VAddU32( dst=vgpr("LocalReadAddr%s+%u"%(tc,i)), \ src0=(i * self.states.regCaps["maxLDSConstOffset"]), \ src1=vgpr("LocalReadAddr%s"%tc), \ comment="Final Offset Plus %uK"%((i * self.states.regCaps["maxLDSConstOffset"]) / 1024))) return module ############################################################################## # Local Read: Increment A/B ############################################################################## def localReadInc(self, kernel, iui, tP): tc = tP["tensorChar"] if (not self.do["LocalRead%s" % tc]) or ((tc in ("A", "MXSA", "B", "MXSB")) and kernel["DirectToVgpr%s"%tc]): # no local read code if DirectToVgpr is enabled return Module("localReadInc (Empty)") module = Module("localReadInc") offsetInc = 0 LdsPad = kernel["LdsPad%s"%tc] if kernel["LdsBlockSizePerPad%s"%tc] == 0 else 0 if kernel["EnableMatrixInstruction"]: # TODO: remove ceil matrixInstK = kernel["MatrixInstK"] if tc == "MXSA": matrixInstK = ceil(kernel["MatrixInstK"] / kernel["ProblemType"]["MXBlockA"]) elif tc == "MXSB": matrixInstK = ceil(kernel["MatrixInstK"] / kernel["ProblemType"]["MXBlockB"]) if self.states.inTailLoop: if kernel["UseDotInstruction"]: # dot2 inc = int(self.states.lrvwUnrollA * kernel["NumWaveSplitK"] * tP["bpeDS"]) comment = "(LocalReadVectorWidth*NumWaveSplitK*bpeDS)" else: inc = int((kernel["MacroTile%s" % tP["tensorChar"]] + LdsPad) * tP["bpeDS"]) comment = " ((MT+PAD)*bpeDS)" if kernel["EnableMatrixInstruction"]: if kernel["UnrollMajorLDS%s" % tc]: if tc in ("MXSA", "MXSB"): # Tail-loop K-step between MFMA-K sub-iterations for MX scales, # gated by MXScaleFormat: # - Swizzled (HostPreSwizzle/InMemorySwizzle): MT * mxUnit * bpeDS, # scaled by matrixInstK (M-blocks interleaved on K). # - NoSwizzle (canonical): mxUnit * bpeDS; mxUnit already encodes # the per-K-scale stride and is not multiplied by matrixInstK. subTc = tc[3] mxUnit = kernel["MatrixInstK"] // kernel["ProblemType"]["MXBlock%s" % subTc] mxScaleFormat = kernel.get("MXScaleFormat", "NoSwizzle") isMxSwizzled = mxScaleFormat in ("InMemorySwizzle", "HostPreSwizzle") if isMxSwizzled: inc = kernel["MacroTile%s"%tP["tensorChar"]] * tP["bpeDS"] * max(self.states.numReadsIterCoalescedMXSA,self.states.numReadsIterCoalescedMXSB) comment = " (bpeDS)" inc *= matrixInstK else: inc = mxUnit * tP["bpeDS"] * max(self.states.numReadsIterCoalescedMXSA, self.states.numReadsIterCoalescedMXSB) comment = " (mxUnit*bpeDS)" else: inc = tP["bpeDS"] * max(self.states.numReadsIterCoalescedA, self.states.numReadsIterCoalescedB) comment = " (bpeDS)" inc *= matrixInstK else: inc *= matrixInstK if kernel["ProblemType"]["Sparse"]: if (kernel["ProblemType"]["Sparse"] == 2 and tc == "B") or (kernel["ProblemType"]["Sparse"] == 1 and tc == "A"): inc //= 2 elif tc == "Metadata": inc //= 8 padd = 0 # Apply additional padding if needed when cumulative incs add upt to ldsblocksize if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0): ldsBSPad = kernel["LdsBlockSizePerPad%s"%tc] if ( (iui+1) * inc) % kernel["LdsBlockSizePerPad%s"%tc] == 0: totalIncForNextLR = (iui+1) * inc totalIncForCurrLR = iui * inc # Calculate number padding needed for inc used in next LR, and subtract # any padding added previously to avoid double counting. padd = (totalIncForNextLR // ldsBSPad - totalIncForCurrLR // ldsBSPad ) \ * kernel["LdsPad%s"%tc] * tP["bpeDS"] module.addComment0("Adding additional %u pad since cumulative inc has reached %u"\ %(padd, kernel["LdsBlockSizePerPad%s"%tc])) inc = int(inc) with self.allocTmpSgpr(1, tag="localReadInc_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx module.add(SMovB32(dst=sgpr(tmpSgpr), src=int(inc + padd), comment="inc")) numLra = 0 if tP["isA"]: numLra = self.states.a.numVgprLocalReadAddr elif tP["isB"]: numLra = self.states.b.numVgprLocalReadAddr elif tP["isMXSA"]: numLra = self.states.mxsa.numVgprLocalReadAddr elif tP["isMXSB"]: numLra = self.states.mxsb.numVgprLocalReadAddr elif tP["isM"]: numLra = self.states.m.numVgprLocalReadAddr for i in range(numLra): module.add(VAddCOU32( dst=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], i)), \ dst1=VCC(), \ src0=sgpr(tmpSgpr), \ src1=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], i)), \ comment="lr%s += %u%s"%(tP["tensorChar"], inc, comment) )) else: if tP["localReadInstruction"].numOffsets == 1: if kernel["EnableMatrixInstruction"]: if "MXS" in tc: subTc = tc[3] mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"] # K-step between MFMA-K sub-iterations for MX scales: # - Swizzled (HostPreSwizzle/InMemorySwizzle): # MT * mxUnit (M-blocks interleaved on K) # - NoSwizzle (canonical), LDS layout follows UnrollMajorLDS: # UMLDS=1 (K-major LDS): mxUnit (K-scales contiguous per M) # UMLDS=0 (M-major LDS): (MT + LdsPad) * mxUnit (step over M-row, mxUnit K-scales) mxScaleFormat = kernel.get("MXScaleFormat", "NoSwizzle") isMxSwizzled = mxScaleFormat in ("InMemorySwizzle", "HostPreSwizzle") if isMxSwizzled: offsetInc = kernel["MacroTile%s"%tP["tensorChar"]] * mxUnit elif kernel["UnrollMajorLDS%s" % tP["tensorChar"]]: offsetInc = mxUnit else: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * mxUnit elif kernel["UnrollMajorLDS%s" % tP["tensorChar"]]: if tc in ("MXSA", "MXSB"): offsetInc = matrixInstK * max(self.states.numReadsIterCoalescedMXSA, self.states.numReadsIterCoalescedMXSB) else: offsetInc = matrixInstK * max(self.states.numReadsIterCoalescedA, self.states.numReadsIterCoalescedB) if kernel["ProblemType"]["Sparse"]: if (kernel["ProblemType"]["Sparse"] == 2 and tc == "B") or (kernel["ProblemType"]["Sparse"] == 1 and tc == "A"): offsetInc //= 2 elif tc == "Metadata": offsetInc //= 8 else: if tc == "A": sparseA = kernel["ProblemType"]["Sparse"] == 1 lrvw = kernel["LocalReadVectorWidth%s"%tc] wlr = max(lrvw//kernel["MIInputPerThreadA"], 1) if kernel["ProblemType"]["Sparse"] and lrvw < kernel["MIInputPerThreadA"]: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * (kernel["MatrixInstK"]) if sparseA: offsetInc //= 2 elif (self.states.localReadDoCntA)%(wlr): offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * kernel["MIInputPerThreadA"] else: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * ((matrixInstK * wlr) - (kernel["MIInputPerThreadA"] * (wlr - 1))) if sparseA: offsetInc //= 2 elif tc in ("MXSA", "MXSB"): offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * (matrixInstK) elif tc == "Metadata": lrvw = kernel["LocalReadVectorWidth%s"%tc] // 4 wlr = max(lrvw//kernel["MIInputPerThreadMetadata"], 1) if lrvw < kernel["MIInputPerThreadMetadata"]: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * kernel["MatrixInstK"] // 8 elif (self.states.localReadDoCntMetadata)%(wlr): offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * kernel["MIInputPerThreadMetadata"] else: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * (kernel["MatrixInstK"] * wlr - (kernel["MIInputPerThreadMetadata"] * (wlr - 1))) offsetInc //= 8 elif tc == "B": sparseB = kernel["ProblemType"]["Sparse"] == 2 lrvw = kernel["LocalReadVectorWidth%s"%tc] wlr = max(lrvw//kernel["MIInputPerThreadB"], 1) if kernel["ProblemType"]["Sparse"] and lrvw < kernel["MIInputPerThreadB"]: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * (kernel["MatrixInstK"]) if sparseB: offsetInc //= 2 elif (self.states.localReadDoCntB)%(wlr): offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * kernel["MIInputPerThreadB"] else: offsetInc = (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) * (kernel["MatrixInstK"] * wlr - (kernel["MIInputPerThreadB"] * (wlr - 1))) if sparseB: offsetInc //= 2 else: raise Exception(f"unsupport tc %s{tc}") else: # dot2 offsetInc = self.states.lrvwUnrollA * kernel["NumWaveSplitK"] if kernel["UseDotInstruction"] else (kernel["MacroTile%s"%tP["tensorChar"]] + LdsPad) tP["localReadOffset"] += offsetInc module.addComment0("N/A, lro->%d" % tP["localReadOffset"]) module.addComment0("localReadDoCntA %d localReadDoCntMXSA %d localReadDoCntB %d localReadDoCntMXSB %d localReadDoCntM %d" \ % (self.states.localReadDoCntA, self.states.localReadDoCntMXSA, self.states.localReadDoCntB, self.states.localReadDoCntMXSB, self.states.localReadDoCntMetadata)) else: inc = (kernel["MacroTile%s" % tP["tensorChar"]] + LdsPad) numLra = 0 if tP["isA"]: numLra = self.states.a.numVgprLocalReadAddr elif tP["isB"]: numLra = self.states.b.numVgprLocalReadAddr elif tP["isMXSA"]: numLra = self.states.mxsa.numVgprLocalReadAddr elif tP["isMXSB"]: numLra = self.states.mxsb.numVgprLocalReadAddr elif tP["isM"]: numLra = self.states.m.numVgprLocalReadAddr for i in range(numLra): module.add(VAddCOU32( dst=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], i)), \ dst1=VCC(), \ src0=hex(inc), \ src1=vgpr("LocalReadAddr%s+%u"%(tP["tensorChar"], i)), \ comment="lr%s += %u ((MT+Pad)*bpe)"%(tP["tensorChar"], inc) )) return module ############################################################################## # Local Read: Do It A/B # iui = Inner Unroll Idx # uIdx - Unroll Idx # epsi = expand pointer swap index. Only used for PAP ############################################################################## def localReadDo(self, kernel, bufferIdx, iui, epsi, tP): if not self.do["LocalRead%s" % tP["tensorChar"]]: imod = Module("LocalReadDo%s_I%s" % (tP["tensorChar"], iui)) pack = Module("pack%s_I%s" % (tP["tensorChar"], iui)) return imod, pack component = Component.LocalRead.find(self) if component: return component(self, kernel, bufferIdx, iui, epsi, tP) ############################################################################## # Save the local read pointers, for example when creating a duplicated # optimized path (like optNLL) ############################################################################## def saveLocalPointers(self, kernel, tPA, tPB): tPA["savedLocalReadOffset"] = tPA["localReadOffset"] if kernel["ProblemType"]["MXBlockA"]: tPA["MX"]["savedLocalReadOffset"] = tPA["MX"]["localReadOffset"] if kernel["ProblemType"]["MXBlockB"]: tPB["MX"]["savedLocalReadOffset"] = tPB["MX"]["localReadOffset"] tPB["savedLocalReadOffset"] = tPB["localReadOffset"] tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: tPM["savedLocalReadOffset"] = tPM["localReadOffset"] self.states.savedLocalReadDoCntA = self.states.localReadDoCntA if kernel["ProblemType"]["MXBlockA"]: self.states.savedLocalReadDoCntMXSA = self.states.localReadDoCntMXSA if kernel["ProblemType"]["MXBlockB"]: self.states.savedLocalReadDoCntMXSB = self.states.localReadDoCntMXSB self.states.savedLocalReadDoCntB = self.states.localReadDoCntB self.states.savedLocalReadDoCntMetadata = self.states.localReadDoCntMetadata if kernel["ExpandPointerSwap"]: tPA["savedLocalWriteSwapByteOffset"] = tPA["localWriteSwapByteOffset"] if kernel["ProblemType"]["MXBlockA"]: tPA["MX"]["savedLocalWriteSwapByteOffset"] = tPA["MX"]["localWriteSwapByteOffset"] if kernel["ProblemType"]["MXBlockB"]: tPB["MX"]["savedLocalWriteSwapByteOffset"] = tPB["MX"]["localWriteSwapByteOffset"] tPB["savedLocalWriteSwapByteOffset"] = tPB["localWriteSwapByteOffset"] if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: tPM["savedLocalWriteSwapByteOffset"] = tPM["localWriteSwapByteOffset"] ############################################################################## # Restore the saved local read pointers # Must be paired with an earlier call to savePointers ############################################################################## def restoreLocalPointers(self, kernel, tPA, tPB): tPA["localReadOffset"] = tPA["savedLocalReadOffset"] if kernel["ProblemType"]["MXBlockA"]: tPA["MX"]["localReadOffset"] = tPA["MX"]["savedLocalReadOffset"] if kernel["ProblemType"]["MXBlockB"]: tPB["MX"]["localReadOffset"] = tPB["MX"]["savedLocalReadOffset"] tPB["localReadOffset"] = tPB["savedLocalReadOffset"] tPM = tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"] if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: tPM["localReadOffset"] = tPM["savedLocalReadOffset"] self.states.localReadDoCntA = self.states.savedLocalReadDoCntA if kernel["ProblemType"]["MXBlockA"]: self.states.localReadDoCntMXSA = self.states.savedLocalReadDoCntMXSA if kernel["ProblemType"]["MXBlockB"]: self.states.localReadDoCntMXSB = self.states.savedLocalReadDoCntMXSB self.states.localReadDoCntB = self.states.savedLocalReadDoCntB self.states.localReadDoCntMetadata = self.states.savedLocalReadDoCntMetadata if kernel["ExpandPointerSwap"]: tPA["localWriteSwapByteOffset"] = tPA["savedLocalWriteSwapByteOffset"] if kernel["ProblemType"]["MXBlockA"]: tPA["MX"]["localWriteSwapByteOffset"] = tPA["MX"]["savedLocalWriteSwapByteOffset"] if kernel["ProblemType"]["MXBlockB"]: tPB["MX"]["localWriteSwapByteOffset"] = tPB["MX"]["savedLocalWriteSwapByteOffset"] tPB["localWriteSwapByteOffset"] = tPB["savedLocalWriteSwapByteOffset"] if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: tPM["localWriteSwapByteOffset"] = tPM["savedLocalWriteSwapByteOffset"] ############################################################################## # Shift Vector Components d0,1 ############################################################################## def shiftVectorComponents(self, kernel, tP): component = Component.ShiftVectorComponents.find(self) if component: return component(self, kernel, tP) ############################################################################## # computeStoreSrd # Add tile assignment fields to store srd # This is based on WG not the WI/TT assignment ############################################################################## def computeGSUFlatOffset(self, kernel, module, addrDVgpr): """Add per-GSU partition offset to a flat address VGPR pair for MultipleBuffer/MBSK. Mirrors computeStoreSrdStartCommon (which updates SrdD for buffer stores). No-op when GSU==1 or _GlobalAccumulation is not MultipleBuffer/MBSK.""" if kernel["_GlobalAccumulation"] not in ("MultipleBuffer", "MultipleBufferSingleKernel"): return gsuFlatLabel = Label(label=self.labels.getNameInc("GSU_FlatAddr"), comment="") with self.allocTmpSgpr(1) as tmpSgprGSU: module.add(SAndB32(dst=sgpr(tmpSgprGSU.idx), src0=sgpr("GSU"), src1=hex(0x3FFF), comment="Restore GSU")) module.add(SCmpEQU32(src0=sgpr(tmpSgprGSU.idx), src1=1, comment="GSU == 1 ?")) module.add(SCBranchSCC1(labelName=gsuFlatLabel.getLabelName(), comment="skip GSU offset if GSU == 1")) numDim = kernel["ProblemType"]["NumIndicesC"] with self.allocTmpSgpr(4, alignment=1) as tmpSgprInfo: if tmpSgprInfo.idx % 2 == 0: tmpSgprX2 = tmpSgprInfo.idx tmpSgpr0, tmpSgpr1, tmpSgpr2, tmpSgpr3 = (tmpSgprInfo.idx+i for i in range(4)) else: tmpSgprX2 = tmpSgprInfo.idx+1 tmpSgpr0, tmpSgpr1, tmpSgpr2, tmpSgpr3 = tmpSgprInfo.idx+1, tmpSgprInfo.idx+2, tmpSgprInfo.idx+0, tmpSgprInfo.idx+3 module.addComment("Flat store: GSU offset for per-GSU partition") module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tmpSgpr0), sgpr(tmpSgpr1), sgpr("SizesFree+0"), sgpr("GSUSumIdx"), comment="Free0")) for i in range(1, numDim): module.add(SSubU32(dst=sgpr(tmpSgpr2), src0=sgpr("SizesFree+%u"%i), src1=1, comment="Free%u" % i)) module.add(SMulI32(dst=sgpr(tmpSgpr2), src0=sgpr(tmpSgpr2), src1=sgpr("GSUSumIdx"), comment="Free%u" % i)) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tmpSgpr2), sgpr(tmpSgpr3), sgpr(tmpSgpr2), sgpr("StrideC%s"%self.states.indexChars[i]), comment="Free%u" % i)) module.add(SAddU32(dst=sgpr(tmpSgpr0), src0=sgpr(tmpSgpr0), src1=sgpr(tmpSgpr2), comment="Free%u" % i)) module.add(SAddCU32(dst=sgpr(tmpSgpr1), src0=sgpr(tmpSgpr1), src1=sgpr(tmpSgpr3), comment="Free%u" % i)) module.add(SLShiftLeftB64(dst=sgpr(tmpSgprX2,2), src=sgpr(tmpSgprX2,2), shiftHex=log2(self.states.bpeCinternal), comment="scale by bpe")) tmpVgprGSU = self.vgprPool.checkOut(2) module.add(VMovB32(dst=vgpr(tmpVgprGSU+0), src=sgpr(tmpSgprX2), comment="sgpr -> vgpr")) module.add(VMovB32(dst=vgpr(tmpVgprGSU+1), src=sgpr(tmpSgpr1), comment="sgpr -> vgpr")) module.add(VAddCOU32(dst=vgpr(addrDVgpr+0), dst1=VCC(), src0=vgpr(tmpVgprGSU+0), src1=vgpr(addrDVgpr+0), comment="add lo GSU offset to flat addr")) module.add(VAddCCOU32(dst=vgpr(addrDVgpr+1), dst1=VCC(), src0=vgpr(tmpVgprGSU+1), src1=vgpr(addrDVgpr+1), src2=VCC(), comment="add hi GSU offset to flat addr")) self.vgprPool.checkIn(tmpVgprGSU) module.add(gsuFlatLabel) def computeStoreSrdStart(self, kernel, srdTcList: list, sgprBpeList = [], useSize: list = [], noMultipleBuffer = False): module = Module("computeStoreSrdStart") if useSize: assert len(srdTcList) == len(useSize) else: useSize = [False for _ in srdTcList] # For subtile StreamK kernels (StreamK==3, no atomic), the SGPR pool is exhausted # after endSummation. Temporarily expose SrdWS (s60-s63) as Available scratch so # that allocTmpSgpr calls within this function (and the SK component call below) # can borrow those slots. Restore SrdWS as InUse at the end. srdWsAvailableCtx = ( self.states.streamK.borrowsSrdWsInEpilogue and kernel.get("StreamKAtomic", 1) == 0 and "SrdWS" in self.sgprs and "SrdWS" not in self.states.freeSgprVarPool ) if srdWsAvailableCtx: self.addSgprVarToPool("SrdWS") # Keep tmp SGPR usage lean for the common path. with self.allocTmpSgpr(3, tag="computeStoreSrdStart_tmpSgprInfo") as tmpSgprInfo: tmpS0 = tmpSgprInfo.idx tmpS1 = tmpS0+1 wgMT1 = tmpS0+2 assert kernel["BufferStore"] module.addSpaceLine() module.add(SMulI32(dst=sgpr(wgMT1), src0="MT1", src1=sgpr("WorkGroup1"), comment="<- wg1*MT1")) # Overall strategy is to set the SRD to the top-left of the macro-tile. # TT offsets are from this base (and include the column) # In non-packed mode: # higher-order tensor dims are static since this kernel operates within # the 2D Tensor formed by Index0 and Indexa. # Index0 and Index1 vary for each work-item (aka 'dynamic') so roll these into the VGPR # In packed mode: # Higher-order dimensions may be packed into coord0 / coord1 - see rowstart calculation below # Walk through addressing components (each tensor index) in C # For static dims add to SrdC / SrdD to compute a new base. # For dynamic (based on TT assignment) - save in coutRowPtrD in computeStoreVgprs, # which saves the TT assignment for each WI scaled by StrideC0 # TODO - future opportunities for store vgpr and other optimization # - coutRowPtrD and tid1 are strongly related - can we merge or remove one of these? # Packed follows same philosophy but may have more vector components indices = list(range(0, kernel["ProblemType"]["NumIndicesC"])) numDim = len(indices) #addrSrcSgpr = "Address" # use "Address" only for the first iteration addrSrcSgpr = "Srd" # Since SrdC/D are initialized with AddressC/D for non-General Batched GEMM case. _gateMultiBpeShiftSgpr = None _gl = kernel["ProblemType"].get("GateResidualDataTypeList", []) if "Gate" in srdTcList and _gl and len(_gl) > 1: _gateMultiBpeShiftSgpr = self.sgprPool.checkOut(1, "gateBpeShift") _gateShiftEnd = Label(self.labels.getNameInc("Gate_SrdBpe_End"), "") _fallbackBpe = max(1, self.states.bpeGate) for _gDtype in _gl: _gNext = Label(self.labels.getNameInc("Gate_SrdBpe_Next_%s" % _gDtype.toNameAbbrev()), "") module.add(self.getSCMPKInstruction("LGU32", "GateType", _gDtype.value, comment="Gate SRD bpe: GateType != %u" % _gDtype.value)) module.add(SCBranchSCC1(_gNext.getLabelName(), "try next gate dtype")) _dtypeBpe = max(1, int(self.states.bpr * _gDtype.numRegisters())) module.add(SMovB32(dst=sgpr(_gateMultiBpeShiftSgpr), src=log2(_dtypeBpe), comment="Gate SRD log2(bpe)=%u for %s" % (log2(_dtypeBpe), _gDtype.toNameAbbrev()))) module.add(SBranch(labelName=_gateShiftEnd.getLabelName(), comment="done")) module.add(_gNext) module.add(SMovB32(dst=sgpr(_gateMultiBpeShiftSgpr), src=log2(_fallbackBpe), comment="Gate SRD log2(bpe) fallback")) module.add(_gateShiftEnd) for i in range(1, numDim): if i == kernel["ProblemType"]["Index0"]: # Used if the output is transposed? addToSrd = False elif i == kernel["ProblemType"]["Index1"] and len(kernel["PackedC1IndicesX"]) == 1: coord = sgpr(wgMT1) addToSrd = True elif i != kernel["ProblemType"]["Index0"] and i != kernel["ProblemType"]["Index1"] and not isPackedIndex(kernel, i): # group index, this is higher-order Tensor dimension, just add to SRD base: isStridedBuffer = kernel["ProblemType"]["StridedBatched"] or kernel["_GlobalAccumulation"] coord = sgpr("WorkGroup2") if isStridedBuffer else None addToSrd = True if isStridedBuffer else False else: # could be packed higher-order index, just ignore coord = None addToSrd = False if not sgprBpeList: sgprBpeList = [""] * len(srdTcList) assert len(srdTcList) == len(sgprBpeList) if addToSrd: for mat, sgprBpe, us in zip(srdTcList, sgprBpeList, useSize): generalBatchedGemmLoad = Label(label="GeneralBatchedGemmLoad"+mat, comment="Computing the Batch Matrix's base address for General Batched GEMM") generalBatchedGemmLoad_End = Label(label="GeneralBatchedGemmLoad"+mat+"_End", comment="End of label GeneralBatchedGemmLoad"+mat) argTypeChecks = Label(label="ArgTypeChecks"+mat, comment="Checks for ArgType to General Batched or non-General Batched") stridedBatchedGemmLoad = Label(label="StridedBatchedGemmLoad"+mat, comment="Computing the Batch Matrix's base address for Strided Batched GEMM") bpe = self.states.bpeCinternal if mat =="Bias" else (self.states.bpeE if mat == "E" else self.states.bpeCexternal) if mat == "Gate": bpe = self.states.bpeGate bpe = int(self.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters()) if kernel["_GlobalAccumulation"] == 'MultipleBuffer' and mat =="C" else bpe if mat == "Gate" and _gateMultiBpeShiftSgpr is not None: bpe = sgpr(_gateMultiBpeShiftSgpr) else: bpe = sgpr(sgprBpe) if sgprBpe else log2(bpe) # sgprBpe cannot be 0 if(kernel["GlobalSplitU"] != 0): module.add(SAndB32(dst=sgpr(tmpS1), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) # These are constant across all workitems, just add to the SRD: if us: if i == 0: # Skip cause stride = 1 if use size instead continue if i > 1: strideC0 = "Size%s"%(INDEX_CHARS[0]) strideC1 = "Size%s"%(INDEX_CHARS[1]) module.add(SMulI32(dst=sgpr(tmpS0), src0=sgpr(strideC0), src1=sgpr(strideC1))) for x in range(2, i - 1): strideC = "Size%s"%(INDEX_CHARS[x]) module.add(SMulI32(dst=sgpr(tmpS0), src0=sgpr(tmpS0), src1=sgpr(strideC))) if(i == 2 and (mat == "C" or mat == "D")): gsuComp = Component.GSU.find(self) module.add(gsuComp.routeToGeneralBatchedOrStridedBatched(stridedBatchedGemmLoad, argTypeChecks, generalBatchedGemmLoad, mat, kernel, tmpS1)) skComp = Component.StreamK.find(self) module.add(skComp.routeToGeneralBatchedOrStridedBatched(stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel)) module.add(stridedBatchedGemmLoad) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tmpS0), sgpr(tmpS1), coord, sgpr(tmpS0), comment="Scale%s %s by Stride"%(mat, coord))) else: strideC = "Size%s"%(INDEX_CHARS[i-1]) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tmpS0), sgpr(tmpS1), coord, sgpr(strideC), comment="Scale%s %s by Stride"%(mat, coord))) else: if mat == "Gate": strideC = "GateStride+%u" % (i - 1) else: strideC = "Stride%s%s"%(mat, self.states.indexChars[i]) if(i == 2 and (mat == "C" or mat == "D")): gsuComp = Component.GSU.find(self) module.add(gsuComp.routeToGeneralBatchedOrStridedBatched(stridedBatchedGemmLoad, argTypeChecks, generalBatchedGemmLoad, mat, kernel, tmpS1)) skComp = Component.StreamK.find(self) module.add(skComp.routeToGeneralBatchedOrStridedBatched(stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel)) module.add(stridedBatchedGemmLoad) module.addModuleAsFlatItems(self.s_mul_u64_u32(sgpr(tmpS0), sgpr(tmpS1), coord, sgpr(strideC), comment="Scale%s %s by Stride"%(mat, coord))) module.add(SLShiftLeftB64(dst=sgpr(tmpS0,2), src=sgpr(tmpS0,2), shiftHex=bpe, comment="scale by bpe")) module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("%s%s+0"%(addrSrcSgpr, mat)), src1=sgpr(tmpS0), comment="add lo to SRD")) module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("%s%s+1"%(addrSrcSgpr, mat)), src1=sgpr(tmpS1), comment="add hi to SRD")) if(i == 2 and (mat == "C" or mat == "D")): module.add(SBranch(labelName=generalBatchedGemmLoad_End.getLabelName())) module.add(generalBatchedGemmLoad) module.add(SMulI32(dst=sgpr(tmpS0), src0=8, src1=coord, comment="Compute stride in bytes into Pointer Array")) module.add(SAddU32(dst=sgpr(tmpS0), src0=sgpr(tmpS0), src1=sgpr("Address%s+0"%mat), comment="Offsetting to the location [Lower half of address]")) module.add(SAddCU32(dst=sgpr(tmpS1), src0=sgpr("Address%s+1"%mat), src1=0, comment="Offsetting to the location [Higher half of address]")) module.add(SLoadB64(dst=sgpr(tmpS0, 2), base=sgpr(tmpS0, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array")) module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("Srd%s+0"%mat), src1=sgpr(tmpS0), comment="Offsetting within the Batch Matrix [Lower half of address]")) module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Offsetting within the Batch Matrix [Higher half of address]")) module.add(generalBatchedGemmLoad_End) module.addSpaceLine() addrSrcSgpr = "Srd" # update src Sgpr for the second or later iterations if _gateMultiBpeShiftSgpr is not None: self.sgprPool.checkIn(_gateMultiBpeShiftSgpr) if noMultipleBuffer: if srdWsAvailableCtx: self.removeSgprVarFromPool("SrdWS") return module gsuComponent = Component.GSU.find(self) module.add(gsuComponent.computeStoreSrdStart(self, kernel)) skComponent = Component.StreamK.find(self) module.add(skComponent.computeStoreSrdStart(self, kernel)) for cdir in (0,1): indices = kernel["PackedC%uIndicesX"%cdir] packedSizes = "PackedSize%u"%cdir if len(indices) > 1: for i,idx in enumerate(indices[1:]): if i==0: module.add(SMulI32(dst=sgpr(packedSizes), src0=self.sizeRef(indices[0]), \ src1=self.sizeRef(idx), comment="first packed size")) else: module.add(SMulI32(dst=sgpr(packedSizes), src0=sgpr(packedSizes), \ src1=self.sizeRef(idx), comment="first packed size")) if srdWsAvailableCtx: self.removeSgprVarFromPool("SrdWS") return module ############################################################################## # computeStoreVgprs # Compute workitem/TT offsets in VGPRS # and coord0/coord1 ############################################################################## def computeStoreVgprs(self, kernel, divisor=None, tid0Scale=None, tid1Scale=None): module = Module("computeStoreVgprs") module.addComment0("computeStoreVgprs") component = Component.ComputeStoreVgprs.find(self) if component: if kernel["EnableMatrixInstruction"]: module.add(component(self, kernel)) else: module.add(component(self, kernel, divisor, tid0Scale, tid1Scale)) return module ############################################################################## # globalWriteWorkGroupInit: ############################################################################## def SrdTDInit(self, kernel): module = Module("SrdTDInit") tmpspgr0 = self.sgprPool.checkOut(1, tag="SrdTDInit_tmpspgr0", preventOverflow=False) tmpspgr1 = self.sgprPool.checkOutAligned(2, 4, tag="SrdTDInit_tmpspgr1", preventOverflow=False) tmpspgr2 = self.sgprPool.checkOutAligned(2, 4, tag="SrdTDInit_tmpspgr2", preventOverflow=False) module.addComment0("calculate SrdTD address") module.add(SMovB32(dst=sgpr("SrdTD+2"), src="BufferOOB")) module.add(SMovB32(dst=sgpr("SrdTD+3"), src="Srd127_96", comment="Set bits 127_96 in post-loop SRD")) module.add(SMulI32(dst=sgpr(tmpspgr0), src0="MT1", src1=sgpr("WorkGroup1"), comment="")) module.add(SMulHIU32(dst=sgpr(tmpspgr1+1), src0=sgpr(tmpspgr0), src1=sgpr("StrideD1J"), comment="")) module.add(SMulI32(dst=sgpr(tmpspgr1+0), src0=sgpr(tmpspgr0), src1=sgpr("StrideD1J"), comment="")) bpe = int(self.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters()) # self.states.bpeCinternal module.add(SLShiftLeftB64(dst=sgpr(tmpspgr1,2), src=sgpr(tmpspgr1,2), shiftHex=log2(bpe), comment="scale by bpe")) SrdTDGeneralBatched = Label(label="SrdTDInit_GeneralBatched", comment="") SrdTDGeneralBatched_End = Label(label="SrdTDInit_GeneralBatched_End", comment="") if kernel["ProblemType"]["SupportUserArgs"]: module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=3, comment="ArgType == 3 for General Batched GEMM")) module.add(SCBranchSCC1(labelName=SrdTDGeneralBatched.getLabelName(), comment="Initializing General Batched GEMM SrdTD differently")) module.add(SAddU32(dst=sgpr("SrdTD+0"), src0=sgpr("AddressTD+0"), src1=sgpr(tmpspgr1+0), comment="add lo to SRTD" )) module.add(SAddCU32(dst=sgpr("SrdTD+1"), src0=sgpr("AddressTD+1"), src1=sgpr(tmpspgr1+1), comment="add hi to SRTD" )) module.add(SMulHIU32(dst=sgpr(tmpspgr1+1), src0=sgpr("WorkGroup2"), src1=sgpr("StrideDK"), comment="")) module.add(SMulI32(dst=sgpr(tmpspgr1+0), src0=sgpr("WorkGroup2"), src1=sgpr("StrideDK"), comment="")) module.add(SLShiftLeftB64(dst=sgpr(tmpspgr1,2), src=sgpr(tmpspgr1,2), shiftHex=log2(bpe), comment="scale by bpe")) module.add(SBranch(labelName=SrdTDGeneralBatched_End.getLabelName())) module.add(SrdTDGeneralBatched) module.add(SMulI32(dst=sgpr(tmpspgr2+0), src0=8, src1=sgpr("WorkGroup2"), comment="Compute stride in bytes into Pointer Array")) module.add(SAddU32(dst=sgpr(tmpspgr2+0), src0=sgpr(tmpspgr2+0), src1=sgpr("AddressTD+0"), comment="Offsetting to the location [Lower half of address]")) module.add(SAddCU32(dst=sgpr(tmpspgr2+1), src0=sgpr("AddressTD+1"), src1=0, comment="Offsetting to the location [Higher half of address]")) module.add(SLoadB64(dst=sgpr("SrdTD", 2), base=sgpr(tmpspgr2, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array")) module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) module.add(SrdTDGeneralBatched_End) module.add(SAddU32(dst=sgpr("SrdTD+0"), src0=sgpr("SrdTD+0"), src1=sgpr(tmpspgr1+0), comment="add lo to SRTD" )) module.add(SAddCU32(dst=sgpr("SrdTD+1"), src0=sgpr("SrdTD+1"), src1=sgpr(tmpspgr1+1), comment="add hi to SRTD" )) self.sgprPool.checkIn(tmpspgr0) self.sgprPool.checkIn(tmpspgr1) self.sgprPool.checkIn(tmpspgr2) module.add(self.shiftSrd("TD")) return module def globalWriteWorkGroupInit(self, kernel): module = Module("globalWriteWorkGroupInit") if kernel["BufferStore"]: module.add(self.allocPostLoopSrd("D", kernel)) module.add(self.allocPostLoopSrd("C", kernel)) self.sgprBpeList = ["GSULog2BpeC", "GSULog2BpeD"] if kernel["GlobalSplitU"] != 0 else [] # Set BPE based on reduction algorithm if self.states.streamK.emitsWorkspaceReductionBpe and not kernel["StreamKForceDPOnly"]: sgprLog2BpeC = self.sgprPool.checkOut(1, tag="globalWriteWorkGroupInit_sgprLog2BpeC", preventOverflow=False) sgprLog2BpeD = self.sgprPool.checkOut(1, tag="globalWriteWorkGroupInit_sgprLog2BpeD", preventOverflow=False) module.add(SMovB32(dst=sgpr(sgprLog2BpeC), src=log2(self.states.bpeCexternalGSU1))) module.add(SMovB32(dst=sgpr(sgprLog2BpeD), src=log2(self.states.bpeCexternalGSU1))) bpeDoneLabel = Label(label=self.labels.getNameInc("BPEDone"), comment="") module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer")) module.add(SCBranchSCC0(labelName=bpeDoneLabel.getLabelName(), comment="If synchronizer, use regular output BPE")) sSkt = self.acquireStreamKConstSgpr(kernel, "skTiles") if self.isStreamKConstantsToVgprEnabled(kernel): module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(self.states.skConstVgprs["skTiles"]))) module.add(SCmpEQU32(src0=sgpr(sSkt), src1=1, comment="split == 1 ?")) self.releaseStreamKConstSgpr(sSkt) module.add(SCBranchSCC1(labelName=bpeDoneLabel.getLabelName(), comment="If split == 1, use reguler output BPE")) # BPE for parallel reduction module.add(SMovB32(dst=sgpr(sgprLog2BpeC), src=log2(int(self.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters())))) module.add(SMovB32(dst=sgpr(sgprLog2BpeD), src=log2(self.states.bpeCinternal))) module.add(bpeDoneLabel) self.sgprBpeList = [sgprLog2BpeC, sgprLog2BpeD] module.add(self.computeStoreSrdStart(kernel, ["C", "D"], sgprBpeList=self.sgprBpeList)) if kernel["GlobalSplitU"] != 0: module.add(self.undefineSgpr("GSULog2BpeC")) if kernel["StreamK"] == 0: module.add(self.undefineSgpr("AddressC")) if self.states.streamK.emitsWorkspaceReductionBpe and not kernel["StreamKForceDPOnly"]: if not kernel["StoreRemapVectorWidth"]: self.sgprPool.checkIn(sgprLog2BpeD) self.sgprPool.checkIn(sgprLog2BpeC) return module ############################################################################## # LocalSplitU: Global Write Indices ############################################################################## def localSplitUGlobalWriteIndices(self, kernel): module = Module("localSplitUGlobalWriteIndices") # Add LSU Offset back packedC1 = kernel["PackedC1IndicesX"] strideC1 = "StrideC%s" % (self.states.indexChars[packedC1[0]]) strideD1 = "StrideD%s" % (self.states.indexChars[packedC1[0]]) wave_id = self.vgprPool.checkOut(1, "tmpWaveID") tmpVgpr1 = self.vgprPool.checkOutAligned(2,2,"tmpVgpr1") tmpVgpr1Res = ContinuousRegister(tmpVgpr1, 2) module.add(vectorStaticDivide(wave_id, "Serial", kernel["WavefrontSize"], tmpVgpr1Res)) numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] module.add(vectorStaticDivide(wave_id, wave_id, numWaves, tmpVgpr1Res)) with self.allocTmpSgpr(1, tag="localSplitUGlobalWriteIndices_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx if self.LSUValidOffset0 > 0: module.add(SMovB32(dst=sgpr(tmpSgpr), \ src=hex(self.LSUValidOffset0), comment="a valid offset")) module.add(VMulLOU32(dst=vgpr(tmpVgpr1), src0=vgpr(wave_id), src1=sgpr(tmpSgpr), comment="wave LSU offset")) module.add(VAddU32(dst=vgpr(self.vgprs.coord0), src0=vgpr(tmpVgpr1), src1=vgpr(self.vgprs.coord0), comment="coord0 += LSU offset0")) else: module.addComment0("valid offset coord0 is zero.") if self.LSUValidOffset1 > 0: module.add(SMovB32(dst=sgpr(tmpSgpr), \ src=hex(self.LSUValidOffset1), comment="a valid offset")) module.add(VMulLOU32(dst=vgpr(tmpVgpr1), src0=vgpr(wave_id), src1=sgpr(tmpSgpr), comment="wave LSU offset")) module.add(VAddU32(dst=vgpr(self.vgprs.coord1), src0=vgpr(tmpVgpr1), src1=vgpr(self.vgprs.coord1), comment="coord1 += LSU offset1")) module.add(VAddU32(dst=vgpr(self.vgprs.coord1InMT), src0=vgpr(tmpVgpr1), src1=vgpr(self.vgprs.coord1InMT), comment="coord1InMT += LSU offset1")) # this code is from CouputeStoreVgprs. coord 1 : offset part # cinRowPtr/coutRowPtrD are only allocated for BufferStore (SRD-based row offsets) if kernel["BufferStore"]: packedC1 = kernel["PackedC1IndicesX"] strideC1 = "StrideC%s" % (self.states.indexChars[packedC1[0]]) strideD1 = "StrideD%s" % (self.states.indexChars[packedC1[0]]) module.add(VMulLOU32(dst=vgpr(self.vgprs.cinRowPtr), src0=vgpr(self.vgprs.coord1InMT), src1=sgpr(strideC1), comment=" offset 1")) module.add(VMulLOU32(dst=vgpr(self.vgprs.coutRowPtrD), src0=vgpr(self.vgprs.coord1InMT), src1=sgpr(strideD1), comment=" offset 1")) if kernel["ProblemType"]["UseE"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0): module.add(VMovB32(dst=vgpr(self.vgprs.coutRowPtrE), src=vgpr(self.vgprs.coord1InMT), comment=" save offset 1 for E")) if self.vgprs.coutRowPtrGate != -1: module.add(VMovB32(dst=vgpr(self.vgprs.coutRowPtrGate), src=vgpr(self.vgprs.coord1InMT), comment=" save offset 1 for Gate")) if self.vgprs.coutRowPtrBias != -1: index = packedC1[0] - 1 strideW1 = "Size%s" % "I" if index == 0 else ("J" if index == 1 else (self.states.indexChars[index])) module.add(VMulLOU32(dst=vgpr(self.vgprs.coutRowPtrBias), src0=vgpr(self.vgprs.coord1InMT), src1=sgpr(strideW1), comment=" offset 1")) else: module.addComment0("valid offset coord1 is zero.") self.vgprPool.checkIn(tmpVgpr1) self.vgprPool.checkIn(wave_id) self.vgprPool.checkIn(self.lsuCoordOffset) self.vgprPool.checkIn(self.vgprs.coord0InMT) self.vgprPool.checkIn(self.vgprs.coord1InMT) if kernel["BufferStore"]: #print "----AddressC-LocalSplitU" #print self.vgprPool.state() self.vgprs.addrE = -1 self.vgprs.addrD = -1 self.vgprs.addrC = -1 self.vgprs.addrBias = -1 self.vgprs.addrScaleAVec = -1 self.vgprs.addrScaleBVec = -1 self.vgprs.addrScaleAlphaVec = -1 else: self.vgprs.addrE = -1 self.vgprs.addrBias = -1 self.vgprs.addrScaleAVec = -1 self.vgprs.addrScaleBVec = -1 self.vgprs.addrScaleAlphaVec = -1 self.vgprs.addrD = self.vgprPool.checkOut(2, tag="LSU_addrD") module.add(VMovB32( dst=vgpr(self.vgprs.addrD+0), \ src=sgpr("AddressD+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( dst=vgpr(self.vgprs.addrD+1), \ src=sgpr("AddressD+1"), \ comment="sgpr -> vgpr")) self.computeGSUFlatOffset(kernel, module, self.vgprs.addrD) self.vgprs.addrC = self.vgprPool.checkOut(2, tag="LSU_addrC") module.add(VMovB32( dst=vgpr(self.vgprs.addrC+0), \ src=sgpr("AddressC+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( dst=vgpr(self.vgprs.addrC+1), \ src=sgpr("AddressC+1"), \ comment="sgpr -> vgpr")) if kernel["GlobalSplitU"] != 0: gsuLabel = Label(label=self.labels.getNameInc("GSU"), comment="") with self.allocTmpSgpr(1, tag="localSplitUGlobalWriteIndices_tmpSgprGSU") as tmpSgprGSU: module.add(SAndB32(dst=sgpr(tmpSgprGSU.idx), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SCmpEQU32(src0=sgpr(tmpSgprGSU.idx), src1=1, comment="GSU == 1 ?")) module.add(SCBranchSCC0(labelName=gsuLabel.getLabelName(), comment="branch if GSU != 1")) if kernel["ProblemType"]["UseE"]: self.vgprs.addrE = self.vgprPool.checkOut(2, 'addrE') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrE+0), \ src=sgpr("AddressE+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrE+1), \ src=sgpr("AddressE+1"), \ comment="sgpr -> vgpr")) if self.states.useBias == DataDirection.READ: self.vgprs.addrBias = self.vgprPool.checkOut(2, 'addrBias') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrBias+0), \ src=sgpr("AddressBias+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrBias+1), \ src=sgpr("AddressBias+1"), \ comment="sgpr -> vgpr")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): self.vgprs.addrScaleAVec = self.vgprPool.checkOut(2, 'addrScaleAVec') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAVec+0), \ src=sgpr("AddressScaleA+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAVec+1), \ src=sgpr("AddressScaleA+1"), \ comment="sgpr -> vgpr")) self.vgprs.addrScaleBVec = self.vgprPool.checkOut(2, 'addrScaleVVec') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleBVec+0), \ src=sgpr("AddressScaleB+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleBVec+1), \ src=sgpr("AddressScaleB+1"), \ comment="sgpr -> vgpr")) if kernel["ProblemType"]["UseScaleAlphaVec"]: self.vgprs.addrScaleAlphaVec = self.vgprPool.checkOut(2, 'addrScaleAlphaVec') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAlphaVec+0), \ src=sgpr("AddressScaleAlphaVec+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAlphaVec+1), \ src=sgpr("AddressScaleAlphaVec+1"), \ comment="sgpr -> vgpr")) if kernel["GlobalSplitU"] != 0: module.add(gsuLabel) return module ############################################################################## def shiftSrd(self, tc) -> Module: return self._shiftSrdImpl(sgprLo=sgpr("Srd%s+1"%tc), sgprHi=sgpr("Srd%s+2"%tc)) def shiftSrdByIdx(self, srdIdx) -> Module: return self._shiftSrdImpl(sgprLo=sgpr(srdIdx+1), sgprHi=sgpr(srdIdx+2)) def _shiftSrdImpl(self, sgprLo, sgprHi) -> Module: module = Module("shiftSrd") if self.states.version[:2] == (12, 5): with self.allocTmpSgpr(1, tag="shiftSrd_tmpSgprRes") as stmpRes: # gfx125x encodes low record-count bits in Srd+1, so this must run # after the final Srd+2 limit value is known. module.addComment("Shift num records for gfx125x") module.add(SAndB32(sgpr(stmpRes.idx), sgprHi, 0x7F)) module.add(SLShiftLeftB32(sgpr(stmpRes.idx), 25, sgpr(stmpRes.idx))) module.add(SAndB32(sgprLo, sgprLo, 0x1FFFFFF)) module.add(SOrB32(sgprLo, sgprLo, sgpr(stmpRes.idx))) module.add(SLShiftRightB32(sgprHi, 7, sgprHi)) return module def allocPostLoopSrd(self, ch: str, kernel): module = Module("allocPostLoopSrd") GeneralBatchedGemmSrdInitiation = Label(label="GeneralBatchedGemmSrdInitiation"+ch, comment="Handling General Batched GEMM SRD initialization") GeneralBatchedGemmSrdInitiation_End = Label(label="GeneralBatchedGemmSrdInitiation"+ch+"_End", comment="End of handling General Batched GEMM SRD initialization") ArgTypeCheckLabel = Label(label="ArgTypeCheck"+ch, comment="Check if ArgType is for General Batched GEMM for "+ch) RegularSrdInitialization = Label(label="RegularSrdInitialization"+ch, comment="Regular SRD initialization for non-General Batched GEMM for "+ch) gsuComponent = Component.GSU.find(self) module.add(gsuComponent.initializeSrd(self, ArgTypeCheckLabel, GeneralBatchedGemmSrdInitiation_End, kernel, ch)) if kernel["ProblemType"]["SupportUserArgs"]: module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=3, comment="ArgType == 3 for General Batched GEMM")) module.add(SCBranchSCC0(labelName=RegularSrdInitialization.getLabelName())) # Check for StreamK Kernel when ArgType == 3 (General Batched GEMM) # AddressFlags == 0, then parallel reduction in StreamK and SrdC/D needs to be initialized to workspace pointer (AddressC/D) # AddressFlags != 0, then not parallel reduction in StreamK and SrdC/D should be initialized to batch matrix address from pointer array (AddressC/D) skComponent = Component.StreamK.find(self) module.add(skComponent.initializeSrdAddressFlagsCheck(GeneralBatchedGemmSrdInitiation)) module.add(RegularSrdInitialization) module.add(SMovB64(dst=sgpr("Srd%s+0"%ch, 2), src=sgpr("Address%s+0"%ch, 2), comment="init SRD base address" )) module.add(SBranch(labelName=GeneralBatchedGemmSrdInitiation_End.getLabelName())) module.add(GeneralBatchedGemmSrdInitiation) module.add(SMovB64(dst=sgpr("Srd%s+0"%ch, 2), src=0, comment="init SRD to 0" )) module.add(GeneralBatchedGemmSrdInitiation_End) module.add(SMovB32(dst=sgpr("Srd%s+2"%ch), src="BufferOOB")) module.add(SMovB32(dst=sgpr("Srd%s+3"%ch), src="Srd127_96", comment="Set bits 127_96 in post-loop SRD")) module.add(self.shiftSrd(ch)) module.addSpaceLine() return module def allocPostLoopSrdSuppressRaw(self, ch: str, chAddress: str, labelStr: str, sgprLength, evenAlignedAddr: bool = True) -> Module: module = Module("allocPostLoopSrdSuppress") label = Label("%sAddrValid"%labelStr, "") label2 = Label("%sAddrValid_End"%labelStr, "") # Set SRD base pointer. evenAlignedAddr=False uses 32-bit ops instead of # s_mov_b64/s_cmp_eq_u64, which require an even-aligned AddressX pair. if evenAlignedAddr: module.add(SMovB64(dst=sgpr("Srd%s+0"%ch, 2), src=sgpr("Address%s+0"%chAddress, 2), comment="init SRD base address" )) module.add(SMovB32(dst=sgpr("Srd%s+3"%ch), src="Srd127_96", comment="Set bits 127_96 in post-loop SRD")) module.add(BranchIfNotZero("Address%s"%chAddress, DataType('int64').toEnum(), label)) else: # 32-bit base init (odd-aligned safe): copy lo/hi halves independently. module.add(SMovB32(dst=sgpr("Srd%s+0"%ch), src=sgpr("Address%s+0"%chAddress), comment="init SRD base address (lo, unaligned-safe)" )) module.add(SMovB32(dst=sgpr("Srd%s+1"%ch), src=sgpr("Address%s+1"%chAddress), comment="init SRD base address (hi, unaligned-safe)" )) module.add(SMovB32(dst=sgpr("Srd%s+3"%ch), src="Srd127_96", comment="Set bits 127_96 in post-loop SRD")) # Address != 0 (64-bit) without s_cmp_eq_u64: branch to valid if either half is nonzero. module.add(SCmpEQU32(src0=sgpr("Address%s+0"%chAddress), src1=0, comment="Address%s lo == 0 ?"%chAddress)) module.add(SCBranchSCC0(label.getLabelName(), "branch if Address%s lo != 0"%chAddress)) module.add(SCmpEQU32(src0=sgpr("Address%s+1"%chAddress), src1=0, comment="Address%s hi == 0 ?"%chAddress)) module.add(SCBranchSCC0(label.getLabelName(), "branch if Address%s hi != 0"%chAddress)) module.add(SMovB32(dst=sgpr("Srd%s+2"%ch), src=0)) module.add(SBranch(label2.getLabelName())) module.add(label) module.add(SMovB32(dst=sgpr("Srd%s+2"%ch), src=sgprLength)) module.add(label2) module.addSpaceLine() return module def allocPostLoopSrdSuppress(self, ch: str, labelStr: str, sgprLength, evenAlignedAddr: bool = True) -> Module: return self.allocPostLoopSrdSuppressRaw(ch, ch, labelStr, sgprLength, evenAlignedAddr=evenAlignedAddr) ############################################################################## # Not LocalSplitU: Global Write Indices ############################################################################## def notLocalSplitUGlobalWriteIndices(self, kernel): #print "GlobalWriteIndices" if not self.do["PostLoop"]: return "" module = Module("notLocalSplitUGlobalWriteIndices") if kernel["EnableMatrixInstruction"]: module.add(self.computeStoreVgprs(kernel)) else: module.add(self.computeStoreVgprs(kernel, divisor = kernel["SubGroup0"], tid0Scale = kernel["VectorWidthA"], tid1Scale = kernel["VectorWidthB"])) if kernel["BufferStore"]: #print "----AddressC-nonLSU-----" #print self.vgprPool.state() self.vgprs.addrE = -1 self.vgprs.addrD = -1 self.vgprs.addrC = -1 self.vgprs.addrBias = -1 self.vgprs.addrScaleAVec = -1 self.vgprs.addrScaleBVec = -1 self.vgprs.addrScaleAlphaVec = -1 else: self.vgprs.addrE = -1 self.vgprs.addrBias = -1 self.vgprs.addrScaleAVec = -1 self.vgprs.addrScaleBVec = -1 self.vgprs.addrScaleAlphaVec = -1 self.vgprs.addrD = self.vgprPool.checkOut(2, 'addrD') module.add(VMovB32( dst=vgpr(self.vgprs.addrD+0), \ src=sgpr("AddressD+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( dst=vgpr(self.vgprs.addrD+1), \ src=sgpr("AddressD+1"), \ comment="sgpr -> vgpr")) self.computeGSUFlatOffset(kernel, module, self.vgprs.addrD) self.vgprs.addrC = self.vgprPool.checkOut(2, 'addrC') module.add(VMovB32( dst=vgpr(self.vgprs.addrC+0), \ src=sgpr("AddressC+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( dst=vgpr(self.vgprs.addrC+1), \ src=sgpr("AddressC+1"), \ comment="sgpr -> vgpr")) if kernel["GlobalSplitU"] != 0: gsuLabel = Label(label=self.labels.getNameInc("GSU"), comment="") with self.allocTmpSgpr(1, tag="notLocalSplitUGlobalWriteIndices_tmpSgprGSU") as tmpSgprGSU: module.add(SAndB32(dst=sgpr(tmpSgprGSU.idx), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SCmpEQU32(src0=sgpr(tmpSgprGSU.idx), src1=1, comment="GSU == 1 ?")) module.add(SCBranchSCC0(labelName=gsuLabel.getLabelName(), comment="branch if GSU != 1")) if kernel["ProblemType"]["UseE"]: self.vgprs.addrE = self.vgprPool.checkOut(2, 'addrE') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrE+0), \ src=sgpr("AddressE+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrE+1), \ src=sgpr("AddressE+1"), \ comment="sgpr -> vgpr")) if self.states.useBias == DataDirection.READ: self.vgprs.addrBias = self.vgprPool.checkOut(2, 'addrBias') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrBias+0), \ src=sgpr("AddressBias+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrBias+1), \ src=sgpr("AddressBias+1"), \ comment="sgpr -> vgpr")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): self.vgprs.addrScaleAVec = self.vgprPool.checkOut(2, 'addrScaleAVec') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAVec+0), \ src=sgpr("AddressScaleA+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAVec+1), \ src=sgpr("AddressScaleA+1"), \ comment="sgpr -> vgpr")) self.vgprs.addrScaleBVec = self.vgprPool.checkOut(2, 'addrScaleBVec') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleBVec+0), \ src=sgpr("AddressScaleB+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleBVec+1), \ src=sgpr("AddressScaleB+1"), \ comment="sgpr -> vgpr")) if kernel["ProblemType"]["UseScaleAlphaVec"]: self.vgprs.addrScaleAlphaVec = self.vgprPool.checkOut(2, 'addrScaleAlphaVec') module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAlphaVec+0), \ src=sgpr("AddressScaleAlphaVec+0"), \ comment="sgpr -> vgpr")) module.add(VMovB32( \ dst=vgpr(self.vgprs.addrScaleAlphaVec+1), \ src=sgpr("AddressScaleAlphaVec+1"), \ comment="sgpr -> vgpr")) if kernel["GlobalSplitU"] != 0: module.add(gsuLabel) return module ############################################################################## # Release any resources used by the global write def cleanupGlobalWrite(self, kernel): self.vgprPool.checkIn(self.vgprs.coord0) self.vgprPool.checkIn(self.vgprs.coord1) if kernel["StoreRemapVectorWidth"]: self.vgprPool.checkIn(self.vgprs.storeRemapLW) self.vgprPool.checkIn(self.vgprs.storeRemapLR) 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) if self.vgprs.coutRowPtrE != -1: self.vgprPool.checkIn(self.vgprs.coutRowPtrE) self.vgprs.coutRowPtrE = -1 if self.vgprs.coutRowPtrBias != -1: self.vgprPool.checkIn(self.vgprs.coutRowPtrBias) self.vgprs.coutRowPtrBias = -1 if self.vgprs.coutRowPtrGate != -1: self.vgprPool.checkIn(self.vgprs.coutRowPtrGate) self.vgprs.coutRowPtrGate = -1 if not kernel["BufferStore"]: self.vgprPool.checkIn(self.vgprs.addrD) self.vgprPool.checkIn(self.vgprs.addrC) if self.vgprs.addrE != -1: self.vgprPool.checkIn(self.vgprs.addrE) self.vgprs.addrE = -1 if self.states.useBias == DataDirection.READ: self.vgprPool.checkIn(self.vgprs.addrBias) self.vgprs.addrBias = -1 if self.vgprs.addrScaleAVec != -1: self.vgprPool.checkIn(self.vgprs.addrScaleAVec) self.vgprs.addrScaleAVec = -1 if self.vgprs.addrScaleBVec != -1: self.vgprPool.checkIn(self.vgprs.addrScaleBVec) self.vgprs.addrScaleBVec = -1 if self.vgprs.addrScaleAlphaVec != -1: self.vgprPool.checkIn(self.vgprs.addrScaleAlphaVec) self.vgprs.addrScaleAlphaVec = -1 ############################################################################## # Return max global write vector width, in elements def maxGwvw(self, kernel): atomic = (kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1) and (kernel["_GlobalAccumulation"] != 'MultipleBuffer' and kernel["_GlobalAccumulation"] != 'MultipleBufferSingleKernel') if kernel["BufferStore"]: if atomic: return self.getVectorAtomicWidth(kernel) else: return 1000 # no limit else: if atomic: return 1 # flat vector atomic is not tested else: return 1000 # no limit def getVectorAtomicWidth(self, kernel): if kernel["ProblemType"]["DataType"].isHalf() and (not kernel["_GlobalAccumulation"]): return 2 return 1 ############################################################################## # Partition thread-tile into writeElements for store code # This function creates the writeElement mapping for full tiles # (ie non-edge cases) ############################################################################## def notLocalFullTileElements(self, kernel): component = Component.NotLocalFullTileElements.find(self) if component: return component(self, kernel) ############################################################################## # Store Remap: Local Write ############################################################################## def storeRemapAddLocalWrite(self, kernel, ss, addrCalc, srcVgpr): """ Add localWrite for the element with addrCalc and srcVgpr. """ module = Module("storeRemapAddLocalWrite srcVgpr %s"%str(srcVgpr)) bps = self.states.bpeCexternal * ss.cfg.gwvw rpv = self.states.bpeCexternal * ss.cfg.gwvw / self.states.bpr addr0 = vgpr(self.vgprs.storeRemapLW) offset = addrCalc.coordOffset0 * self.states.bpeCexternal ds = DSModifiers(offset=offset) numRegs = int(max(1, rpv)) # Attach LDS mem token (same bucket as bias/SAV loads + GSU ds_store in this # epilogue store BB) so the rocisa MemTokenConsistencyCheck pass doesn't abort # on a BB that mixes tokened and un-tokened LDS ops when SRVW != 0. module.add(dsStore(bps, dstAddr=addr0, src=vgpr(srcVgpr, numRegs), \ ds=ds, comment="storeRemap lw", memToken=MemTokenData([self.states.memTokenLdsBuffer0]))) return module ############################################################################## # Store Remap: Local Read and Global Write ############################################################################## def storeRemapAddStore(self, kernel, tmpVgpr, tmpS01, edge, StoreRemapLastBatch): module = Module("storeRemapAddStore") module.add(SWaitCnt(dscnt=0, comment="wait for LDS write")) numStoreInst = 0 #Data exchange between different waves #Make sure LDS writes are finished of all waves if kernel["MIWaveGroup"][0] > 1: # FIXME: Indent? module.add(SBarrier(comment="wait all lds write finished")) module.addSpaceLine() gwvw = kernel["StoreRemapVectorWidth"] nElements = kernel["MacroTile0"]*kernel["MatrixInstN"]//kernel["MIWaveGroup"][0]//self.states.kernel["WavefrontSize"] bpe = self.states.bpeCexternal bps = bpe * gwvw rpe = self.states.bpeCexternal / self.states.bpr rpv = rpe * gwvw storeRegs = self.vgprs.storeRemapAS src = vgpr(self.vgprs.storeRemapLR) for rIdx, i in enumerate(range(0, nElements, gwvw)): offset = self.storeRemapLrOffset * bpe * (i//gwvw) ds = DSModifiers(offset=offset) dst = vgpr(storeRegs[rIdx], rpv) # Match the storeRemap lw token so the whole epilogue store BB is consistent # for the rocisa MemTokenConsistencyCheck pass (avoids abort when SRVW != 0). module.add(dsLoad(bps, dst=dst, src=src, ds=ds, comment="storeRemap lr", memToken=MemTokenData([self.states.memTokenLdsBuffer0]))) module.addSpaceLine() # Global Write addr1 = sgpr("SrdD", 4) packedD1 = kernel["PackedC1IndicesX"] strideD1 = "StrideD%s" % (self.states.indexChars[packedD1[0]]) vTmp = self.vgprPool.checkOut(1, "SR Store temp addr0") addr0 = vgpr(vTmp) isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal( self.states.asmCaps, kernel["NonTemporalD"], _temporalHint(kernel, "D"), _nonVolatile(kernel, "D")) if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel": isGlc, isSlc, isNT, scope, th, nv = forceCoherentNonTemporal( self.states.asmCaps, isNT, _temporalHint(kernel, "D"), _nonVolatile(kernel, "D"), scope) if not edge: for rIdx, i in enumerate(range(0, nElements, gwvw)): if i == 0: module.add(VMovB32(dst=addr0, src=vgpr(self.vgprs.storeRemapOffsetCoord1), comment="coord1")) else: currentStep = i//gwvw module.add(VAddU32(dst=addr0, src0=vgpr(self.vgprs.storeRemapOffsetCoord1), src1=(self.storeRemapNCPL * currentStep), comment="coord1 += nColPerLoad")) module.add(VMulLOU32(dst=addr0, src0=addr0, src1=sgpr(strideD1), comment="coord1 offset = coord1 * StrideD")) module.add(VAddLShiftLeftU32(dst=addr0, src0=addr0, src1=vgpr(self.vgprs.storeRemapCoord0), shiftHex=hex(log2(bpe)), comment="global write D address")) dscnt = min((nElements-i)//gwvw - 1, 15) module.add(SWaitCnt(dscnt=dscnt, comment="wait for LDS read")) numStoreInst += 1 module.add(self.chooseGlobalWrite(True, bps, storeRegs[rIdx], rpv, addr0, addr1, 0, glc=isGlc, slc=isSlc, nt=isNT, comment="store D StoreRemapVectorWidth", scope=scope, th=th, nv=nv)) else: tmpS23 = tmpS01+self.states.laneSGPRCount coord0 = tmpVgpr coord1 = coord0+1 srvw = kernel["StoreRemapVectorWidth"] edgeVw = min(kernel["AssertFree0ElementMultiple"],kernel["StoreRemapVectorWidth"]) bps = self.states.bpeCexternal * edgeVw rpv = self.states.bpeCexternal / self.states.bpr * edgeVw for rIdx, i in enumerate(range(0, nElements, srvw)): for vi in range (0, srvw, edgeVw): if vi == 0: dscnt = min((nElements-i)//srvw - 1, 15) module.add(SWaitCnt(dscnt=dscnt, comment="wait for LDS read")) sizeBoundary = [0,0] sizeBoundary[0] = \ sgpr("PackedSize0") if len(kernel["PackedC0IndicesX"]) > 1 \ else self.sizeRef(kernel["ProblemType"]["Index0"]) sizeBoundary[1] = \ sgpr("PackedSize1") if len(kernel["PackedC1IndicesX"]) > 1 \ else self.sizeRef(kernel["ProblemType"]["Index1"]) currentStep = i//srvw # calculate global coordination module.add(VAddU32(dst=vgpr(coord1), src0=vgpr(self.vgprs.storeRemapCoord1), src1=(self.storeRemapNCPL * currentStep), comment="coord1 += nColPerLoad")) module.add(VAddU32(dst=vgpr(coord0), src0=vgpr(self.vgprs.storeRemapCoord0), src1=vi, comment="coord0 += element index of load vector")) module.add(VAddU32(dst=addr0, src0=vgpr(self.vgprs.storeRemapOffsetCoord1), src1=(self.storeRemapNCPL * currentStep), \ comment="offset coord1 += nColPerLoad")) module.add(VCmpLtU32(dst=sgpr(tmpS01,self.states.laneSGPRCount), src0=vgpr(coord0), src1=sizeBoundary[0], comment="coord0 < size0" )) module.add(VCmpLtU32(dst=sgpr(tmpS23,self.states.laneSGPRCount), src0=vgpr(coord1), src1=sizeBoundary[1], comment="coord1 < size1" )) SAndBX = SAndB32 if self.states.kernel["WavefrontSize"] == 32 else SAndB64 module.add(SAndBX( dst=sgpr(tmpS23,self.states.laneSGPRCount), src0=sgpr(tmpS01,self.states.laneSGPRCount), src1=sgpr(tmpS23,self.states.laneSGPRCount), comment="in0 && in1" )) module.add(VMulLOU32(dst=addr0, src0=addr0, src1=sgpr(strideD1), comment="coord1 element offset = coord1 * StrideD")) module.add(VAddLShiftLeftU32(dst=addr0, src0=addr0, src1=vgpr(coord0), shiftHex=hex(log2(bpe)), comment="scale to BPE")) module.add(VCndMaskB32(dst=addr0, src0=-1, src1=addr0, src2=sgpr(tmpS23,self.states.laneSGPRCount), comment="clip if OOB. offset" )) sumIdx = storeRegs[rIdx] + int(vi*rpe) numStoreInst += 1 if bps == 2: module.add(self.chooseGlobalWrite(True, bpe, sumIdx, rpe, addr0, addr1, 0, glc=isGlc, slc=isSlc, nt=isNT, hi16=vi%2, comment="store D StoreRemapVectorWidth", scope=scope, th=th, nv=nv)) else: module.add(self.chooseGlobalWrite(True, bps, sumIdx, rpv, addr0, addr1, 0, glc=isGlc, slc=isSlc, nt=isNT, comment="store D StoreRemapVectorWidth", scope=scope, th=th, nv=nv)) if bps == 1: module.add(VAShiftRightI32(dst=vgpr("ValuC+%u"%sumIdx), shiftHex=8, src=vgpr("ValuC+%u"%sumIdx), comment=" shift 1 byte" )) module.addSpaceLine() self.vgprPool.checkIn(vTmp) #Data exchange between different waves #Make sure LDS reads are finished of all waves if kernel["MIWaveGroup"][0] > 1: module.add(SBarrier(comment="wait all lds read finished")) return module, numStoreInst ############################################################################## # Store remap compute vgprs: ############################################################################## def storeRemapComputeStoreVgprs(self, kernel): module = Module("storeRemapComputeStoreVgprs") module.addComment0("Store Remap Local Write address") with self.allocTmpSgpr(2, tag="storeRemapComputeStoreVgprs_tmpSgprInfo") as tmpSgprInfo: tmpS0 = tmpSgprInfo.idx wgMT1 = tmpS0 + 1 wg0="WorkGroup0" wg1="WorkGroup1" tid0 = self.vgprPool.checkOut(1, "SR coord0") tid1 = self.vgprPool.checkOut(1, "SR coord1") coord1Offset = self.vgprPool.checkOut(1, "SR coord1 offset") storeRemapLW = self.vgprPool.checkOut(1, "SR local write") storeRemapLR = self.vgprPool.checkOut(1, "SR local read") tmpV0 = self.vgprPool.checkOut(5, tag="storeRemapComputeStoreVgprs_tmpV0") waveCoord0 = tmpV1 = tmpV0+1 ldsStride = tmpV0+2 coord0 = tmpV0+3 waveCoord1 = tmpV0+4 gwvw = kernel["StoreRemapVectorWidth"] ldsPad = max(kernel["StoreRemapVectorWidth"],kernel["MIOutputVectorWidth"]) #calculate local write Address: v[vgprLocalWriteAddrC] module.add(vectorStaticDivideAndRemainder(tid1, tid0, "Serial", self.states.kernel["WavefrontSize"]*kernel["MIWaveGroup"][0], \ ContinuousRegister(tmpV0, 2))) module.add(VMulLOU32(dst=vgpr(waveCoord1), src0=hex(kernel["MatrixInstN"]), src1=vgpr(tid1), comment="coord1 offset of LDS for each Wave")) module.add(VAndB32(dst=vgpr(tid1), src0=hex(kernel["MatrixInstN"]-1), src1=vgpr("Serial"), comment="coord1 offset of LDS for each thread")) module.add(VAddU32(dst=vgpr(tid1), src0=vgpr(waveCoord1), src1=vgpr(tid1), comment="coord1 offset in MacroTile")) module.add(VMovB32(dst=vgpr(ldsStride), src=hex(kernel["MacroTile0"]+ldsPad), \ comment="lds stride = MT0 + PAD")) module.add(VMulLOU32(dst=vgpr(tmpV0), src0=vgpr(tid1), src1=vgpr(ldsStride), \ comment="lds coord1 offset = Col-id* lds stride")) module.add(vectorStaticDivideAndRemainder(waveCoord0, tid0, tid0, self.states.kernel["WavefrontSize"], ContinuousRegister(tmpV0, 2))) module.add(VLShiftRightB32(dst=vgpr(coord0), shiftHex=hex(log2(kernel["MatrixInstN"])), src=vgpr(tid0), \ comment="tid / matrixInstN")) if kernel["MIOutputVectorWidth"] > 1: module.add(VLShiftLeftB32(dst=vgpr(coord0), shiftHex=hex(log2(kernel["MIOutputVectorWidth"])), src=vgpr(coord0), \ comment="lds coord0 offset *= 4 (each thread hold 4 element)")) module.add(VMadU32U24(dst=vgpr(coord0), src0=(kernel["MatrixInstM"]*kernel["MatrixInstBM"]), src1=vgpr(waveCoord0), src2=vgpr(coord0), \ comment="coord0 += waveCoord0 * wave M shape(blockM*MiM)")) if kernel["StreamK"] == 3 and not kernel["StreamKForceDPOnly"]: module.add(VAddLShiftLeftU32( dst=vgpr(storeRemapLW), \ src0=vgpr(tmpV0), \ src1=vgpr(coord0), \ shiftHex=sgpr(self.sgprBpeList[1]), \ comment="local write C address")) else: module.add(VAddLShiftLeftU32( dst=vgpr(storeRemapLW), \ src0=vgpr(tmpV0), \ src1=vgpr(coord0), \ shiftHex=sgpr("GSULog2BpeD"), \ comment="local write C address")) module.addSpaceLine() # calculate local read address : v[vgprLocalReadAddrC] module.addComment0("Store Remap Local Read address") module.add(vectorStaticDivideAndRemainder(tid1, tid0, "Serial", self.states.kernel["WavefrontSize"], \ ContinuousRegister(tmpV0, 2))) module.add(VMulLOU32(dst=vgpr(waveCoord1), src0=hex(kernel["MatrixInstN"]//kernel["MIWaveGroup"][0]), src1=vgpr(tid1), comment="coord1 offset of LDS for each Wave")) nThreadPerCol = kernel["MacroTile0"] // gwvw nColPerLoad = self.states.kernel["WavefrontSize"] // nThreadPerCol self.storeRemapLrOffset = (kernel["MacroTile0"]+ldsPad) * nColPerLoad self.storeRemapNCPL = nColPerLoad module.add(VLShiftRightB32(dst=vgpr(tmpV1),\ shiftHex=hex(log2(nThreadPerCol)), src=vgpr(tid0), \ comment="tid / nThreadPerCol")) module.add(VAddU32(dst=vgpr(coord1Offset), src0=vgpr(waveCoord1), src1=vgpr(tmpV1), comment="coord1 offset in MacroTile")) module.add(VMulLOU32(dst=vgpr(tmpV0), src0=vgpr(coord1Offset), src1=vgpr(ldsStride), \ comment="lds coord1 offset = Col-id* lds stride")) module.add(VAndB32(dst=vgpr(coord0), src0=hex(nThreadPerCol-1), src1=vgpr(tid0), comment="coord0 offset of LDS for each thread")) module.add(VLShiftLeftB32(dst=vgpr(coord0), shiftHex=hex(log2(gwvw)), src=vgpr(coord0), \ comment="lds coord0 offset *= gwvw (each thread hold gwvw element)")) if kernel["StreamK"] == 3 and not kernel["StreamKForceDPOnly"]: module.add(VAddLShiftLeftU32( dst=vgpr(storeRemapLR), \ src0=vgpr(tmpV0), \ src1=vgpr(coord0), \ shiftHex=sgpr(self.sgprBpeList[1]), \ comment="local read C address")) else: module.add(VAddLShiftLeftU32( dst=vgpr(storeRemapLR), \ src0=vgpr(tmpV0), \ src1=vgpr(coord0), \ shiftHex=sgpr("GSULog2BpeD"), \ comment="local read C address")) module.addSpaceLine() # calculate global write coord0 and coord1 module.addComment0("Store Remap global write coord0 and coord1") module.add(vectorStaticDivideAndRemainder(tid1, tid0, "Serial", self.states.kernel["WavefrontSize"]*kernel["MIWaveGroup"][0], \ ContinuousRegister(tmpV0, 2))) ColsPerBlockShape = kernel["MatrixInstN"] * kernel["MatrixInstBN"] module.add(VMulLOU32(dst=vgpr(waveCoord1), src0=hex(ColsPerBlockShape), src1=vgpr(tid1), comment="coord1 offset of global memory for each Wave")) module.add(vectorStaticDivideAndRemainder(tid1, tid0, tid0, self.states.kernel["WavefrontSize"], \ ContinuousRegister(tmpV0, 2))) module.add(VMadU32U24(dst=vgpr(waveCoord1), src0=(kernel["MatrixInstN"]//kernel["MIWaveGroup"][0]), src1=vgpr(tid1), src2=vgpr(waveCoord1), \ comment="waveCoord1 += waveCoord0 * MiN / WaveGroupM")) module.add(VLShiftRightB32(dst=vgpr(tmpV1), shiftHex=hex(log2(nThreadPerCol)), src=vgpr(tid0), \ comment="tid / nThreadPerCol")) module.add(VAddU32(dst=vgpr(coord1Offset), src0=vgpr(waveCoord1), src1=vgpr(tmpV1), comment="coord1 offset in MacroTile")) module.add(SMulI32(dst=sgpr(tmpS0), src0=hex(kernel["MacroTile0"]), src1=sgpr(wg0), comment="%s = wg0*MT0"%sgpr(tmpS0))) module.add(VAddCOU32(dst=vgpr(tid0), dst1=VCC(), src0=sgpr(tmpS0), src1=vgpr(coord0), comment="coord0 = coord0 + wg0 * MT0")) module.add(SMulI32(dst=sgpr(wgMT1), src0="MT1", src1=sgpr(wg1), comment="<- wg1*MT1")) module.add(VAddCOU32(dst=vgpr(tid1), dst1=VCC(), src0=sgpr(wgMT1), src1=vgpr(coord1Offset), comment="coord1 = tid1*VW + wg1*MT1")) module.addSpaceLine() module.add(self._syncThreads(kernel, "StoreRemap Start")) self.vgprs.storeRemapLW = storeRemapLW #local write self.vgprs.storeRemapLR = storeRemapLR #local read self.vgprs.storeRemapCoord0 = tid0 #global coord0 self.vgprs.storeRemapCoord1 = tid1 #global coord1 self.vgprs.storeRemapOffsetCoord1 = coord1Offset #offset coord1 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")) if kernel["StreamK"] == 3 and not kernel["StreamKForceDPOnly"]: self.sgprPool.checkIn(self.sgprBpeList[1]) self.sgprPool.checkIn(self.sgprBpeList[0]) return module ############################################################################## # Not LocalSplitU: Global Write # Determine write batching pattern # element() specifies TT 'coordinate' to write # vectorWidths specifies width of vector to store # TODO - why does this use VectorWidth to control store width? Could be GlobalWriteVectorWidth? # # Function creates one mapping for full tiles and one for edge tiles, # then calls globalWriteElements to generate the code for the new tiles. ############################################################################## def notLocalSplitUGlobalWrite(self, kernel, tPA, tPB): if not self.do["PostLoop"]: return Module("notLocalSplitUGlobalWrite"), None (fullVws, elements, fullVws_1, elements_1) = self.notLocalFullTileElements(kernel) # print("len(elements)= ", len(elements_1)) noGSUBranch = (kernel["GlobalSplitU"] == 0 and (not self.states.streamK.requiresWorkspaceReductionStorePath or kernel["StreamKForceDPOnly"])) module = Module("notLocalSplitUGlobalWrite") storeModule, deferredGSU0 = self.globalWriteElements(kernel, tPA, tPB, fullVws, fullVws_1, elements, elements_1, noGSUBranch=noGSUBranch) module.add(storeModule) self.cleanupGlobalWrite(kernel) return module, deferredGSU0 ############################################################################## # LocalSplitU: Global Write ############################################################################## def localSplitUGlobalWrite(self, kernel, tPA, tPB): if not self.do["PostLoop"]: return "" elements_0 = [[] for y in range(2)] elements_1 = [[] for y in range(2)] elements_f0 = [[] for y in range(2)] elements_f1 = [[] for y in range(2)] # TODO: LocalSplitU has special global write elements, we will not use all StoreVectorWidths (fullVws, elementss_0, fullVws_1, elementss_1) = self.notLocalFullTileElements(kernel) (fullVw, elements_0[False], fullVw_1, elements_1[False]) = (fullVws[0], elementss_0[0], fullVws_1[0], elementss_1[0]) (edgeVw, elements_0[True], edgeVw_1, elements_1[True] ) = (fullVws[-1], elementss_0[-1], fullVws_1[-1], elementss_1[-1]) edgeScaled_0 = len(elements_0[True]) // len(elements_1[False]) edgeScaled_1 = len(elements_1[True]) // len(elements_1[False]) noEgScaled_0 = len(elements_0[False]) // len(elements_1[False]) for i in range(0, len(elements_1[False])): element = elements_1[False][i] if element in self.LSUelementsPerLSUWave: elements_f1[False].append(element) for j in range(0, edgeScaled_0): # in general, edge will affect vc0 dimension. element = elements_0[True][i*edgeScaled_0+j] elements_f0[True].append(element) for j in range(0, edgeScaled_1): # in general, edge will affect vc0 dimension. element = elements_1[True][i*edgeScaled_1+j] elements_f1[True].append(element) for j in range(0, noEgScaled_0): # in general, edge will affect vc0 dimension. element = elements_0[False][i*noEgScaled_0+j] elements_f0[False].append(element) vectorWidths = [fullVw, edgeVw] vectorWidths_1 = [fullVw_1, edgeVw_1] noGSUBranch = (kernel["GlobalSplitU"] == 0 and (not self.states.streamK.requiresWorkspaceReductionStorePath or kernel["StreamKForceDPOnly"])) module = Module("localSplitUGlobalWrite") storeModule, _ = self.globalWriteElements(kernel, tPA, tPB, vectorWidths, vectorWidths_1, elements_f0, elements_f1, noGSUBranch=noGSUBranch) module.add(storeModule) self.cleanupGlobalWrite(kernel) self.vgprPool.checkIn(self.accVgprLdsReduction) return module ############################################################################## # generateActivationModules # Generate activation modules for different vector widths ############################################################################## def generateActivationModules(self, kernel, activation, activationLabelList, activationEnumStrList, activationSetPCStruct, tmpVgpr, actPCGwvwVgpr, actTempSgpr): module = Module("ActivationModules") assert activationEnumStrList and activationSetPCStruct for key, activationLabelModules in activationLabelList.items(): gwvw = key actModules = Module(getActFuncModuleName(gwvw, \ activationSetPCStruct.vgprActCopy, tmpVgpr.idx, actTempSgpr)) for index, activationLabelModule in enumerate(activationLabelModules): actModule = Module(activationLabelModule.getLabelName()) actModule.isCallable = True actModule.callableName = activationLabelModule.getLabelName() actModule.add(activationLabelModule) activationTypeStr = activationEnumStrList[index] vgprIdx = activationSetPCStruct.vgprActCopy if self.insertActivationAfterPacked(kernel, activationTypeStr): actModule.appendModule(self.getActivationDestDataType(kernel, activation, \ activationTypeStr, gwvw, vgprIdx, vgprIdx, (tmpVgpr.idx + actPCGwvwVgpr), \ actTempSgpr)) else: actModule.appendModule(self.getActivationActivationComputeType(kernel, activation, \ activationTypeStr, gwvw, vgprIdx, vgprIdx, (tmpVgpr.idx + actPCGwvwVgpr), \ actTempSgpr)) actModule.add(SSetPCB64(src=sgpr(activationSetPCStruct.sgprOffsetBack,2))) actModules.add(actModule) module.add(actModules) return module ############################################################################## # generateBetaModules # Generate beta modules with edge detection and vector width handling ############################################################################## def generateBetaModules(self, kernel, tPA, tPB, activation, applyAlpha, betas, edge, atomic, vectorWidths, elements, activationLabelList, tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, writeLabels, endLabel, vectorDataTypes, factorDims, globalWriteMode, hasMultipleGlobalWriteModes): betaModules = Module("Betas") # Base deferral condition — per-factorDim bias check is applied below. # ScaleAlphaVec has similar LDS barriers that block deferral unconditionally. allowDeferBase = ( kernel.get("UseSubtileImpl") ) currentInstLength = 0 for betaIdx in reversed(range(len(betas))): beta = betas[betaIdx] if beta and kernel["_GlobalAccumulation"] == "SingleBuffer" and (kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1): continue betaModule = Module("Beta_%u"%betaIdx) for fdIdx in reversed(range(len(factorDims))): factorDim = factorDims[fdIdx] # Create StoreVectorWidth branches generatedVectorWidths = set() for vectorWidth, element in zip(reversed(vectorWidths), reversed(elements)): # Avoid duplicated vectorWidth if vectorWidth not in generatedVectorWidths: generatedVectorWidths.add(vectorWidth) else: continue # edge module # B0 FD0 edge paths are barrier-free — safe to defer. # FD0 edge paths are safe to defer (edge check is workgroup-uniform, # all waves take the same path so all waves hit the same barrier). isMultipleBuffer = kernel["_GlobalAccumulation"] in ("MultipleBufferSingleKernel", "MultipleBuffer") deferEdge = ( edge and allowDeferBase and ( factorDim == 0 # FD0: safe to defer (workgroup-uniform edge check) or (not isMultipleBuffer and self.states.useBias == DataDirection.NONE) # no bias: safe ) ) edgeModule = Module("Edge_B%u_FD%u_VW%u" % (beta, factorDim, vectorWidth)) if deferEdge: # Generate Edge store into a deferred module edgeDeferredModule = Module("Edge_B%u_FD%u_VW%u_DeferredBlock" % (beta, factorDim, vectorWidth)) # Use ThenDeferredReturn as endLabel so the batch jumps back to inline directly currentInstLength, activationTypeStr = \ self.globalWriteElementBatch(kernel, tPA, tPB, activation, applyAlpha, beta, edge, atomic, vectorWidth, element, activationLabelList, tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, edgeDeferredModule, writeLabels[beta][factorDim][vectorWidth]["ThenDeferred"], writeLabels[beta][factorDim][vectorWidth]["ThenDeferredReturn"], currentInstLength, betaIdx, fdIdx, vectorDataTypes, factorDims) if not hasattr(self.states, 'deferredEdgeModules'): self.states.deferredEdgeModules = [] self.states.deferredEdgeModules.append(edgeDeferredModule) # Inline stub: keep "Then" label, jump to deferred, return + jump to GW_End edgeModule.add(writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["Then"]) with self.allocTmpSgpr(3, tag="generateBetaModules_tmpSgprInfo") as tmpSgprInfo: posLabel = self.labels.getNameInc("ThenDeferredDir") edgeModule.add(SLongBranch(writeLabels[beta][factorDim][vectorWidth]["ThenDeferred"], tmpSgprInfo, posLabel, comment="edge store (deferred)")) edgeModule.addComment0("=" * 60) edgeModule.addComment0(" Edge store B%u FD%u VW%u deferred to after persistent loop" % (beta, factorDim, vectorWidth)) edgeModule.addComment0(" (would have been inline here in non-deferred version)") edgeModule.addComment0("=" * 60) edgeModule.add(writeLabels[beta][factorDim][vectorWidth]["ThenDeferredReturn"]) with self.allocTmpSgpr(2, alignment=2, tag="generateBetaModules_tmpPair") as tmpPair: with self.allocTmpSgpr(1, tag="generateBetaModules_tmpOff") as tmpOff: posLabel = self.labels.getNameInc("ThenDeferredReturnDir") edgeModule.add(SLongBranch(endLabel, tmpPair, tmpOff, posLabel, comment="jump to end")) else: currentInstLength, activationTypeStr = \ self.globalWriteElementBatch(kernel, tPA, tPB, activation, applyAlpha, beta, edge, atomic, vectorWidth, element, activationLabelList, tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, edgeModule, writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["Then"], endLabel, currentInstLength, betaIdx, fdIdx, vectorDataTypes, factorDims) # Edge conditions and branches if edge == True: # Else label elseLabelPos = 0 if vectorWidth == min(vectorWidths) else -1 edgeModule.add(writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["Else"], pos=elseLabelPos) currentInstLength += 1 # This check is dedicated to adaptive kernel if kernel["AdaptiveGemm"] == 1 and vectorWidth != min(vectorWidths): # If module, checking Size0 % vectorWidth > 0 isLongBranch = True if currentInstLength >= self.states.asmCaps["ShortBranchMaxLength"] else False with self.allocTmpSgpr(4, tag="generateBetaModules_tmpSgprInfo2") as tmpSgprInfo: checkIsEdge = edgeModule.add(self.checkIsEdge(kernel, tmpSgprInfo, \ writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["Else"], vectorWidth, isLongBranch=isLongBranch), pos=0) currentInstLength += countInstruction(checkIsEdge) # Longest vectorWidth should have non edge module if vectorWidth == max(vectorWidths): # NonEdgeEnd label edgeModule.add(writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["NonEdgeEnd"], pos=0) currentInstLength += 1 # Non edge module # Keep B0 FD0 NonEdge inline (optimized store path with permute). # Defer other NonEdge paths when no bias and no MultipleBuffer. deferNonEdge = ( allowDeferBase and not (not beta and factorDim == 0) # B0 FD0: keep inline and not isMultipleBuffer and self.states.useBias == DataDirection.NONE ) if deferNonEdge: nonEdgeDeferredModule = Module("NonEdge_B%u_FD%u_VW%u_DeferredBlock" % (beta, factorDim, vectorWidth)) currentInstLength, activationTypeStr = \ self.globalWriteElementBatch(kernel, tPA, tPB, activation, applyAlpha, beta, False, atomic, vectorWidth, element, activationLabelList, tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, nonEdgeDeferredModule, writeLabels[beta][factorDim][vectorWidth]["NonEdgeDeferred"], writeLabels[beta][factorDim][vectorWidth]["NonEdgeDeferredReturn"], currentInstLength, betaIdx, fdIdx, vectorDataTypes, factorDims) if not hasattr(self.states, 'deferredEdgeModules'): self.states.deferredEdgeModules = [] self.states.deferredEdgeModules.append(nonEdgeDeferredModule) nonEdgeModule = Module("Non_Edge_B%u_FD%u_VW%u" % (beta, factorDim, vectorWidth)) nonEdgeModule.add(writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["NonEdge"]) with self.allocTmpSgpr(3, tag="generateBetaModules_tmpSgprInfo3") as tmpSgprInfo: posLabel = self.labels.getNameInc("NonEdgeDeferredDir") nonEdgeModule.add(SLongBranch(writeLabels[beta][factorDim][vectorWidth]["NonEdgeDeferred"], tmpSgprInfo, posLabel, comment="beta NonEdge store (deferred)")) nonEdgeModule.addComment0("=" * 60) nonEdgeModule.addComment0(" NonEdge store B%u FD%u VW%u deferred to after persistent loop" % (beta, factorDim, vectorWidth)) nonEdgeModule.addComment0(" (would have been inline here in non-deferred version)") nonEdgeModule.addComment0("=" * 60) nonEdgeModule.add(writeLabels[beta][factorDim][vectorWidth]["NonEdgeDeferredReturn"]) with self.allocTmpSgpr(2, alignment=2, tag="generateBetaModules_tmpPair2") as tmpPair: with self.allocTmpSgpr(1, tag="generateBetaModules_tmpOff2") as tmpOff: posLabel = self.labels.getNameInc("NonEdgeDeferredReturnDir") nonEdgeModule.add(SLongBranch(endLabel, tmpPair, tmpOff, posLabel, comment="jump to end")) else: nonEdgeModule = Module("Non_Edge_B%u_FD%u_VW%u" % (beta, factorDim, vectorWidth)) currentInstLength, activationTypeStr = \ self.globalWriteElementBatch(kernel, tPA, tPB, activation, applyAlpha, beta, False, atomic, vectorWidth, element, activationLabelList, tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, nonEdgeModule, writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["NonEdge"], endLabel, currentInstLength, betaIdx, fdIdx, vectorDataTypes, factorDims) edgeModule.add(nonEdgeModule, pos=0) # NOTE: isEdgeTarget of normal and adaptive kernels are different # Normal kernel: to Then/Else label, followed by edge store # Adaptive kernel: to NonEdgeEnd label, followed by Size0 % vectorWidth check isEdgeTarget = writeLabels[beta][factorDim][vectorWidth][globalWriteMode] # When UseSubtileImpl is active (and not a multi-buffer GSU accumulation), # use subtile-aligned edge check: remainder must be a multiple of the # subtile block size (32 for M, 16 for N) rather than requiring a full tile. useSubtileEdgeCheck = ( kernel.get("UseSubtileImpl") and kernel["_GlobalAccumulation"] not in ("MultipleBufferSingleKernel", "MultipleBuffer") ) # If module, checking Size1 % MT1 > 0 (or subtile alignment for N) # Force long branch when Edge code is deferred isLongBranch = True if currentInstLength >= 16384 else False with self.allocTmpSgpr(4, tag="generateBetaModules_tmpSgprInfo4") as tmpSgprInfo: if useSubtileEdgeCheck: checkIsEdge = edgeModule.add(self.checkIsEdgeSubtile(kernel, tmpSgprInfo, \ isEdgeTarget["Then"] if kernel["AdaptiveGemm"] == 0 else isEdgeTarget["NonEdgeEnd"], \ isSize1=True, isLongBranch=isLongBranch), pos=0) else: checkIsEdge = edgeModule.add(self.checkIsEdge(kernel, tmpSgprInfo, \ isEdgeTarget["Then"] if kernel["AdaptiveGemm"] == 0 else isEdgeTarget["NonEdgeEnd"], \ kernel["MacroTile1"], isSize1=True, isLongBranch=isLongBranch), pos=0) currentInstLength += countInstruction(checkIsEdge) # If module, checking Size0 % MT0 > 0 (or subtile alignment for M) isLongBranch = True if currentInstLength >= 16384 else False with self.allocTmpSgpr(4, tag="generateBetaModules_tmpSgprInfo5") as tmpSgprInfo: if useSubtileEdgeCheck: checkIsEdge = edgeModule.add(self.checkIsEdgeSubtile(kernel, tmpSgprInfo, \ isEdgeTarget["Else"] if kernel["AdaptiveGemm"] == 0 else isEdgeTarget["NonEdgeEnd"], \ isSize1=False, isLongBranch=isLongBranch), pos=0) else: checkIsEdge = edgeModule.add(self.checkIsEdge(kernel, tmpSgprInfo, \ isEdgeTarget["Else"] if kernel["AdaptiveGemm"] == 0 else isEdgeTarget["NonEdgeEnd"], \ kernel["MacroTile0"], isLongBranch=isLongBranch), pos=0) currentInstLength += countInstruction(checkIsEdge) betaModule.add(edgeModule, pos=0) # FactorDim label betaModule.add(writeLabels[beta][factorDim][globalWriteMode]["Label"], pos=0) currentInstLength += 1 # If module, checking factorDim is zero if len(factorDims) == 2: isLongBranch = True if currentInstLength >= self.states.asmCaps["ShortBranchMaxLength"] else False with self.allocTmpSgpr(3, tag="generateBetaModules_tmpSgprInfo6") as tmpSgprInfo: checkIsFactorDimZero = betaModule.add(self.checkIsFactorDimZero(kernel, tmpSgprInfo, \ writeLabels[beta][factorDims[1]][globalWriteMode]["Label"], isLongBranch=isLongBranch), pos=0) currentInstLength += countInstruction(checkIsFactorDimZero) # Beta label betaModule.add(writeLabels[beta][globalWriteMode]["Label"], pos=0) currentInstLength += 1 betaModules.add(betaModule, pos=0) return betaModules, currentInstLength ############################################################################## # checkBetaBranchExceeds # Check if branch exceeds and add beta zero check if needed # Returns a module to be added to the main module ############################################################################## def checkBetaBranchExceeds(self, kernel, betas, betaModules, writeLabels, globalWriteMode): module = Module("BetaBranchCheck") # Check if branch exceeds if False in betas and True in betas: isBetaLongBranch = False count = 0 count, found = findInstCount(betaModules, writeLabels[1][globalWriteMode]["Label"], count) if found: if count >= self.states.asmCaps["ShortBranchMaxLength"]: isBetaLongBranch = True with self.allocTmpSgpr(3 if isBetaLongBranch else 1, tag="checkBetaBranchExceeds_tmpSgprInfo") as tmpSgprInfo: module.add(self.checkIsBetaZero(kernel, tmpSgprInfo, writeLabels[1][globalWriteMode]["Label"], isBetaLongBranch, posNeg=1)) return module ############################################################################## # checkIsBetaZero # tmpSgpr is one temp sgpr # betaLabel is label to branch to if beta != 0 ############################################################################## def checkIsBetaZero(self, kernel, tmpSgprInfo, betaLabel, isLongBranch=False, placeHolder=None, posNeg: int=0): module = Module("checkIsBetaZero label %s"%betaLabel) assert(isinstance(betaLabel, Label)) betaLabelName = betaLabel.getLabelName() if kernel["ProblemType"]["UseBeta"]: if self.states.bpeCinternal <= self.states.bpr: # 1 register to check for Beta==0 #module.add(SCmpKEQU32(src=sgpr("Beta"), simm16=0, comment="Beta == 0")) module.add(self.getSCMPKInstruction("EQU32", "Beta", 0, comment="Beta == 0")) else: # multiple registers to check for Beta==0 module.add(SMovB32(dst=sgpr(tmpSgprInfo.idx), src=sgpr("Beta+0"), comment="tmp = Beta[0]")) for i in range(1, self.states.bpeCinternal//self.states.bpr): module.add(SOrB32(dst=sgpr(tmpSgprInfo.idx), src0=sgpr("Beta+%u"%i), src1=sgpr(tmpSgprInfo.idx), comment="tmp |= Beta[%u] " % i)) #module.add(SCmpKEQU32(src=sgpr(tmpSgprInfo.idx), simm16=0, comment="Beta == 0")) module.add(self.getSCMPKInstruction("EQU32", tmpSgprInfo.idx, 0, comment="Beta == 0")) if placeHolder == None: if isLongBranch: module.add(self.longBranchScc0(betaLabel, posNeg, tmpSgprInfo)) else: module.add(SCBranchSCC0(labelName=betaLabelName, comment="Branch if Beta is not zero")) else: placeHolderModule = Module(placeHolder) module.add(placeHolderModule) module.addSpaceLine() return module ############################################################################## # checkIsEdgeSubtile # Used when UseSubtileImpl is active. Checks whether the wave's M/N rows # are subtile-aligned so the NonEdge paired-store path can be used. # # Non-last workgroups always take the NonEdge path (their tile is full). # For the last workgroup in each dimension, we check that the partial # remainder is subtile-aligned: # isSize1=False: (SizeI % MT0) % blockSizeM == 0 → NonEdge (else → edge) # blockSizeM = 16 for fp32 dest, 32 for 16-bit dest # isSize1=True : (SizeJ % MT1) % 16 == 0 → NonEdge (else → edge) # # tmpSgpr must have at least 4 free SGPRs (same as checkIsEdge). # isEdgeTarget is the label to branch to when the tile IS an edge. ############################################################################## def checkIsEdgeSubtile(self, kernel, tmpSgprInfo, isEdgeTarget, isSize1=False, isLongBranch=False): assert(isinstance(isEdgeTarget, Label)) isEdgeTargetLabel = isEdgeTarget.getLabelName() module = Module("checkIsEdgeSubtile") dim = "N (isSize1)" if isSize1 else "M" module.addComment1("Edge/NonEdge store path check (%s): subtile-aligned remainder -> NonEdge paired store; unaligned -> Edge scalar store" % dim) tmpS0 = tmpSgprInfo.idx tmpS1 = tmpS0 + 1 tmpS23 = tmpS1 + 1 sizeBoundary = [0, 0] sizeBoundary[0] = \ sgpr("PackedSize0") if len(kernel["PackedC0IndicesX"]) > 1 \ else self.sizeRef(kernel["ProblemType"]["Index0"]) sizeBoundary[1] = \ sgpr("PackedSize1") if len(kernel["PackedC1IndicesX"]) > 1 \ else self.sizeRef(kernel["ProblemType"]["Index1"]) if not isSize1: divisor = kernel["MacroTile0"] destBpe = int(kernel["ProblemType"]["DestDataType"].numBytes()) if self.states.storeAlign8 else 1 alignSize = 16 // destBpe # storeAlign8: dwordx4 store width (16B) / destBpe; else: 16 wgSgpr = "WorkGroup0" nwgSgpr = "NumWorkGroups0" # tmpS0 = SizeI % MT0 (the trailing-row count for the last WG) module.add(scalarStaticDivideAndRemainder(tmpS1, tmpS0, sizeBoundary[0], divisor, ContinuousRegister(tmpS23, 2), 2)) # tmpS1 = nwg0 - 1 module.add(SAddU32(dst=sgpr(tmpS1), src0=hex(-1), src1=sgpr(nwgSgpr))) # SCC = 1 if this is the last WG in dim 0 module.add(SCmpGeU32(src0=sgpr(wgSgpr), src1=sgpr(tmpS1), comment="wg0 >= nwg0-1 ?")) else: divisor = kernel["MacroTile1"] alignSize = 1 if self.states.storeAlign8 else kernel["MatrixInstN"] wgSgpr = "WorkGroup1" nwgSgpr = "NumWorkGroups1" # tmpS0 = SizeJ % MT1 module.add(scalarStaticDivideAndRemainder(tmpS1, tmpS0, sizeBoundary[1], divisor, ContinuousRegister(tmpS23, 2), 2)) # tmpS1 = nwg1 - 1 module.add(SAddU32(dst=sgpr(tmpS1), src0=hex(-1), src1=sgpr(nwgSgpr))) # SCC = 1 if this is the last WG in dim 1 module.add(SCmpGeU32(src0=sgpr(wgSgpr), src1=sgpr(tmpS1), comment="wg1 >= nwg1-1 ?")) # myRem = last WG ? (SizeX % divisor) : 0 # Non-last WGs always take NonEdge (full tile), so myRem = 0 keeps them out of the edge branch. module.add(SCSelectB32(dst=sgpr(tmpS0), src0=sgpr(tmpS0), src1=0, comment="myRem = last WG ? rem : 0")) # Check alignment: myRem % alignSize != 0 → edge. # alignSize is a compile-time value. For power-of-2 use AND; for non-power-of-2 # (e.g., waveGroupM=48 from MIWT3), enumerate the valid multiples and branch-chain. if alignSize & (alignSize - 1) == 0: # Power of 2: use AND for fast modulo module.add(SAndB32(dst=sgpr(tmpS0), src0=sgpr(tmpS0), src1=alignSize - 1, comment="myRem %% %d (subtile alignment check)" % alignSize)) module.add(self.getSCMPKInstruction("GTU32", tmpS0, 0, comment="not subtile-aligned → edge")) if isLongBranch: module.add(self.longBranchScc1(isEdgeTarget, posNeg=1, tmpSgprInfo=tmpSgprInfo, comment="jump to edge if not subtile-aligned")) else: module.add(SCBranchSCC1(labelName=isEdgeTargetLabel, comment="jump to edge if not subtile-aligned")) else: # Non-power-of-2: enumerate valid aligned multiples {0, alignSize, 2*alignSize, ...}. # myRem is in [0, divisor-1]; divisor = alignSize * numWaves. At most numWaves values to check. # Strategy: branch to Edge if NOT any valid multiple. # We emit: for each valid k: if myRem==k*alignSize goto NonEdge # fall-through → Edge if not isSize1: numWaves = kernel["MIWaveGroup"][0] else: numWaves = kernel["MIWaveGroup"][1] nonEdgeLabel = Label(self.labels.getNameInc("subtile_nonedge_aligned"), "myRem is a valid multiple of alignSize=%d" % alignSize) for k in range(numWaves): multiple = k * alignSize module.add(self.getSCMPKInstruction("EQU32", tmpS0, multiple, comment="myRem == %d (aligned multiple k=%d)?" % (multiple, k))) module.add(SCBranchSCC1(labelName=nonEdgeLabel.getLabelName(), comment="aligned → NonEdge")) # Not any valid multiple → Edge if isLongBranch: module.add(self.longBranchScc0(isEdgeTarget, posNeg=1, tmpSgprInfo=tmpSgprInfo, comment="not subtile-aligned (alignSize=%d) → edge" % alignSize)) else: module.add(SBranch(labelName=isEdgeTargetLabel, comment="not subtile-aligned (alignSize=%d) → edge" % alignSize)) module.add(nonEdgeLabel) return module ############################################################################## # _emitSubtileGuards # Compute both M and N OOB guard SGPRs for the UseSubtileImpl NonEdge path. # Results are stored in self.states.subtileM32ValidBlocksSgpr and # self.states.subtileN16ValidBlocksSgpr for use by _emitSubtileOobGuard. # # waveId (serial >> 6) is read once and used for both dimensions: # waveIdM = waveId & (numWavesM - 1) lower bits (M is innermost) # waveIdN = waveId >> log2(numWavesM) upper bits (N is outermost) # # M algorithm: # validM = SizeI - WG0 * MT0 # waveBase = waveIdM * waveGroupM # remainder = max(validM - waveBase, 0) # numValidMBlocks = min(ceil(remainder / mBlockSize), MIWaveTile[0]) # mBlockSize = 32 for 16-bit dest, 16 for fp32 dest. # # N algorithm: # validN = SizeJ - WG1 * MT1 # waveBaseN = waveIdN * waveGroupN (skipped if numWavesN == 1) # validN_wave = max(validN - waveBaseN, 0) # clamped = min(validN_wave, waveGroupN) # numValid16NBlocks = clamped >> 4 ############################################################################## def _emitSubtileGuards(self, kernel, edgeModule): numWavesM = kernel["MIWaveGroup"][0] numWavesN = kernel["MIWaveGroup"][1] log2numWavesM = int(log(numWavesM, 2)) mBlockSize = 16 if kernel["ProblemType"]["DestDataType"].isSingle() else 32 mBlockShift = int(log(mBlockSize, 2)) waveGroupM = kernel["MIWaveTile"][0] * kernel["MatrixInstM"] waveGroupN = kernel["MIWaveTile"][1] * kernel["MatrixInstN"] mt0, mt1 = kernel["MacroTile0"], kernel["MacroTile1"] # Use pre-allocated permanent guard SGPRs (allocated at start of post-loop). assert self.states.subtileM32ValidBlocksSgpr is not None, \ "SubtileMGuard must be pre-allocated before _emitSubtileGuards" tmpM = self.sgprPool.checkOut(1, "subtileWaveIdM", preventOverflow=False) tmpN = self.sgprPool.checkOut(1, "subtileWaveIdN", preventOverflow=False) edgeModule.addComment1("UseSubtileImpl NonEdge guards: numValidD1Steps (MatrixInstM=%d) and numValid16NBlocks" % kernel["MatrixInstM"]) # Read waveId once; extract M (lower bits) and N (upper bits) before AND destroys waveId. waveSize = kernel["WavefrontSize"] log2WaveSize = int(log(waveSize, 2)) edgeModule.add(VReadfirstlaneB32(dst=sgpr(tmpM), src=vgpr("Serial"), comment="lane 0 serial of this wave")) edgeModule.add(SLShiftRightB32(dst=sgpr(tmpM), src=sgpr(tmpM), shiftHex=log2WaveSize, comment="waveId = serial >> %d" % log2WaveSize)) if numWavesN > 1: edgeModule.add(SLShiftRightB32(dst=sgpr(tmpN), src=sgpr(tmpM), shiftHex=log2numWavesM, comment="waveIdN = waveId >> log2(numWavesM=%d)" % numWavesM)) # MIWaveGroup[0] is always a power of 2, so AND is correct for modulo. edgeModule.add(SAndB32(dst=sgpr(tmpM), src0=sgpr(tmpM), src1=numWavesM - 1, comment="waveIdM = waveId & (numWavesM-1=%d)" % (numWavesM - 1))) # --- M guard --- # Each d1 step in the C-load batch corresponds to MatrixInstM rows. # We compute numValidD1Steps = min(ceil(max(validM-waveBase,0)/MatrixInstM), MIWaveTile[0]). # The guard check is (numValidD1Steps > d1): true iff this wave's d1-th block has valid rows. miM = kernel["MatrixInstM"] miMShift = int(log(miM, 2)) edgeModule.addComment0("M-guard: numValidD1Steps = min(ceil(max(validM-waveBase,0)/%d), MIWaveTile[0]=%d)" % (miM, kernel["MIWaveTile"][0])) edgeModule.add(SMulI32(dst=sgpr("SubtileMGuard"), src0=sgpr("WorkGroup0"), src1=mt0, comment="WG0 * MT0")) edgeModule.add(SMulI32(dst=sgpr(tmpM), src0=sgpr(tmpM), src1=waveGroupM, comment="waveBase = waveIdM * waveGroupM(%d)" % waveGroupM)) edgeModule.add(SAddU32(dst=sgpr(tmpM), src0=sgpr(tmpM), src1=sgpr("SubtileMGuard"), comment="totalMOffset = waveBase + WG0*MT0")) edgeModule.add(SSubU32(dst=sgpr("SubtileMGuard"), src0=sgpr("SizeI"), src1=sgpr(tmpM), comment="validM_wave = SizeI - totalMOffset; SCC=1 if OOB")) edgeModule.add(SCSelectB32(dst=sgpr("SubtileMGuard"), src0=0, src1=sgpr("SubtileMGuard"), comment="remainder = 0 if OOB")) # Precompute clamped validM_wave for per-store exec mask (avoids recomputing per store) edgeModule.add(SMinU32(dst=sgpr(tmpM), src0=sgpr("SubtileMGuard"), src1=waveGroupM, comment="validM_wave = min(remainder, waveGroupM=%d) (precomputed for exec mask)" % waveGroupM)) edgeModule.add(SAddU32(dst=sgpr("SubtileMGuard"), src0=sgpr("SubtileMGuard"), src1=miM - 1, comment="ceil: remainder + (%d-1)" % miM)) edgeModule.add(SLShiftRightB32(dst=sgpr("SubtileMGuard"), src=sgpr("SubtileMGuard"), shiftHex=miMShift, comment="numValidD1Steps = ceil(remainder / %d)" % miM)) # Clamp: guard comparison is (numValidD1Steps > d1); d1 < MIWaveTile[0] always. edgeModule.add(SMinU32(dst=sgpr("SubtileMGuard"), src0=sgpr("SubtileMGuard"), src1=kernel["MIWaveTile"][0], comment="clamp to MIWaveTile[0]=%d" % kernel["MIWaveTile"][0])) # Keep tmpM alive — it holds clamped validM_wave for use in per-store exec mask. self.states.subtileTotalMOffsetSgpr = tmpM # --- N guard --- edgeModule.addComment0("N-guard: numValid16NBlocks = min(max(validN-waveBaseN,0), waveGroupN=%d) >> 4" % waveGroupN) edgeModule.add(SMulI32(dst=sgpr("SubtileNGuard"), src0=sgpr("WorkGroup1"), src1=mt1, comment="WG1 * MT1")) edgeModule.add(SSubU32(dst=sgpr("SubtileNGuard"), src0=self.sizeRef(kernel["ProblemType"]["Index1"]), src1=sgpr("SubtileNGuard"), comment="validN = SizeJ - WG1*MT1")) if numWavesN > 1: edgeModule.add(SMulI32(dst=sgpr(tmpN), src0=sgpr(tmpN), src1=waveGroupN, comment="waveBaseN = waveIdN * waveGroupN(%d)" % waveGroupN)) edgeModule.add(SSubU32(dst=sgpr("SubtileNGuard"), src0=sgpr("SubtileNGuard"), src1=sgpr(tmpN), comment="validN - waveBaseN; SCC=1 if OOB")) edgeModule.add(SCSelectB32(dst=sgpr("SubtileNGuard"), src0=0, src1=sgpr("SubtileNGuard"), comment="validN_wave = 0 if OOB")) # clamped = min(validN_wave, waveGroupN); SCC=1 on borrow → keep validN_wave, else waveGroupN. edgeModule.add(SSubU32(dst=sgpr(tmpN), src0=sgpr("SubtileNGuard"), src1=waveGroupN, comment="validN_wave - waveGroupN; SCC=1 if validN_wave < waveGroupN")) edgeModule.add(SCSelectB32(dst=sgpr("SubtileNGuard"), src0=sgpr("SubtileNGuard"), src1=waveGroupN, comment="min(validN_wave, waveGroupN)")) if self.states.storeAlign8: # Keep SubtileNGuard = clamped (not shifted). At use sites: # numBlocks = SubtileNGuard >> 4, partialN = SubtileNGuard & 0xF edgeModule.addComment0("SubtileStoreAlign8: SubtileNGuard stores clamped (not >>4) for arbitrary N mask") else: edgeModule.add(SLShiftRightB32(dst=sgpr("SubtileNGuard"), src=sgpr("SubtileNGuard"), shiftHex=4, comment="numValid16NBlocks = clamped >> 4")) self.sgprPool.checkIn(tmpN) self.states.subtileMBlockSize = mBlockSize ############################################################################## # checkIsEdge # tmpSgpr must have at least 4 free SGPR # isEdgeTarget is the branch target if Size % divisor > 0 ############################################################################## def checkIsEdge(self, kernel, tmpSgprInfo, isEdgeTarget, divisor, isSize1=False, isLongBranch=False, placeHolder=None): assert(isinstance(isEdgeTarget, Label)) isEdgeTargetLabel = isEdgeTarget.getLabelName() module = Module("checkIsEdge") dim = "N (isSize1)" if isSize1 else "M" module.addComment1("Edge/NonEdge store path check (%s): Size %% %d > 0 -> Edge store; else -> NonEdge store" % (dim, divisor)) tmpS0 = tmpSgprInfo.idx tmpS1 = tmpS0 + 1 tmpS23 = tmpS1 + 1 wg0="WorkGroup0" wg1="WorkGroup1" sizeBoundary = [0,0] sizeBoundary[0] = \ sgpr("PackedSize0") if len(kernel["PackedC0IndicesX"]) > 1 \ else self.sizeRef(kernel["ProblemType"]["Index0"]) sizeBoundary[1] = \ sgpr("PackedSize1") if len(kernel["PackedC1IndicesX"]) > 1 \ else self.sizeRef(kernel["ProblemType"]["Index1"]) if isSize1 == False: # check edge0 ### # s23 = rem = Size0 % divisor #-- module.add(scalarStaticDivideAndRemainder(tmpS1, tmpS0, sizeBoundary[0], divisor, \ ContinuousRegister(tmpS23, 2), 2)) # s23 = nwg0-1 module.add(SAddU32(dst=sgpr(tmpS1), src0=hex(-1), src1=sgpr("NumWorkGroups0"))) module.add(SCmpGeU32(src0=sgpr(wg0), src1=sgpr(tmpS1), comment="wg0 >= nwg0-1 ?")) else: # check edge1 ### # TODO-packed - this only needs to change to handle packing into C1 index # change would be similar to above - multiply by product of packed sizes in C1 # -- module.add(scalarStaticDivideAndRemainder(tmpS1, tmpS0, sizeBoundary[1], divisor, \ ContinuousRegister(tmpS23, 2), 2)) # s23 = nwg1-1 module.add(SAddU32(dst=sgpr(tmpS1), src0=hex(-1), src1=sgpr("NumWorkGroups1"))) module.add(SCmpGeU32(src0=sgpr(wg1), src1=sgpr(tmpS1), comment="wg1 >= nwg1-1")) # s01 now = myRem = wg < nwg-1 ? divisor : rem module.add(SCSelectB32(dst=sgpr(tmpS0), src0=sgpr(tmpS0), src1=0, comment="set rem")) # if myRem > 0 goto target if self.do["EdgeWrite"]: #module.add(SCmpKGtU32(src=sgpr(tmpS0), simm16=0, comment="rem > 0")) module.add(self.getSCMPKInstruction("GTU32", tmpS0, 0, comment="rem > 0")) if self.db["ForceEdgeStores"]: module.add(SCmpEQU32(src0=sgpr(tmpS0), src1=sgpr(tmpS0), comment="ForceEdgeStores!")) if placeHolder == None: if isLongBranch: module.add(self.longBranchScc1(isEdgeTarget, posNeg=1, tmpSgprInfo=tmpSgprInfo, comment="jump if edges required")) else: module.add(SCBranchSCC1(labelName=isEdgeTargetLabel, comment="jump if edges required")) else: placeHolderModule = Module(placeHolder) placeHolderModule.addComment1("jump if edges required") module.add(placeHolderModule) return module ############################################################################## # checkIsFactorDimZero # tmpSgpr is one temp sgpr # factorDimLabel is label to branch to if factorDim != 0 ############################################################################## def checkIsFactorDimZero(self, kernel, tmpSgprInfo, factorDimLabel, isLongBranch=False, posNeg: int=0): module = Module("checkIsFactorDimZero label %s"%factorDimLabel) assert(isinstance(factorDimLabel, Label)) factorDimLabelName = factorDimLabel.getLabelName() if kernel["ProblemType"]["UseBias"] or kernel["ProblemType"]["UseScaleAlphaVec"]: if self.states.bpeCinternal <= self.states.bpr: # 1 register to check for Beta==0 module.add(self.getSCMPKInstruction("EQU32", "FactorDim", 0, comment="FactorDim == 0")) else: # multiple registers to check for Beta==0 module.add(SMovB32(dst=sgpr(tmpSgprInfo.idx), src=sgpr("FactorDim+0"), comment="tmp = FactorDim[0]")) for i in range(1, self.states.bpeCinternal//self.states.bpr): module.add(SOrB32(dst=sgpr(tmpSgprInfo.idx), src0=sgpr("FactorDim+%u"%i), src1=sgpr(tmpSgprInfo.idx), comment="tmp |= FactorDim[%u] " % i)) module.add(self.getSCMPKInstruction("EQU32", tmpSgprInfo.idx, 0, comment="FactorDim == 0")) if isLongBranch: module.add(self.longBranchScc0(factorDimLabel, posNeg, tmpSgprInfo)) else: module.add(SCBranchSCC0(labelName=factorDimLabelName, comment="Branch if FactorDim is not zero")) module.addSpaceLine() return module ############################################################################## # Global Write Elements ############################################################################## class BF16CVTVgprStruct(NamedTuple): # class for bf16 vgprs vgprBf16Temp: int = -1 # rounding bias constant (0x7fff) for standard bf16 pack path vgprBf16Mask: int = -1 # mask constant (0xffff0000) for extracting bf16 bits from f32 vgprFp32Nan: int = -1 # NaN sentinel used during bf16 saturation/rounding vgprBf16Inc: int = -1 # increment constant for bf16 rounding (standard pack path) # UseSubtileImpl paired dwordx4 store extras (+0..+3 above are reused as pack/perm staging): vgprPermAddr: int = -1 # per-lane ds_bpermute byte address (partner_lane*4); constant for the whole batch vgprLaneGroupDelta: int = -1 # per-lane lane_group*8: M-row byte offset added to addrDVgpr for the dwordx4 store vgprAddrScratch: int = -1 # per-store scratch: holds (addrDVgpr scaled + lane_group*8) without modifying addrDVgpr class FP8CVTVgprStruct(NamedTuple): vgprFp8NanInf: int = -1 vgprFp8Temp: int = -1 vgprFp8Min: int = -1 vgprFp8Max: int = -1 class BF8CVTVgprStruct(NamedTuple): vgprBF8NanInf: int = -1 vgprBF8Temp: int = -1 vgprBF8Min: int = -1 vgprBF8Max: int = -1 class I8CVTVgprStruct(NamedTuple): vgprI8Temp0: int = -1 vgprI8Temp1: int = -1 vgprI8Mask0: int = -1 vgprI8Mask1: int = -1 class ActivationSetPCStruct(NamedTuple): sgprOffsetActivation: int = -1 sgprOffsetBack: int = -1 vgprActCopy: int = -1 calleeLabelsByGwvw: Optional[Mapping[int, Tuple[str, ...]]] = None def globalWriteElements(self, kernel, tPA, tPB, vectorWidths_2, vectorWidths_1, elements_2, elements_1, noGSUBranch=False, applyAlpha=True, # defaults to generating *=alpha codes betas=None, # if left unspecified, then let global parameter decide edge=True # defaults to using edge write ): if not self.do["PostLoop"]: return Module("GlobalWriteElements (Empty)"), Module("DeferredGSU0 (Empty)") module = Module("GlobalWriteElements") deferredGSU0 = Module("DeferredGSU0") module.addComment2("Global Write Elements") if kernel["ProblemType"]["OutputAmaxD"]: module.add(VMovB32(dst=vgpr("AmaxOut"), src="0")) if self.states.numStoreSgprToLoad or self.states.numStoreSgprToLoad2: # Wait for kernel args module.add(SWaitCnt(kmcnt=0, comment="wait for %u bytes of kern args."%((self.states.numStoreSgprToLoad+self.states.numStoreSgprToLoad2) * 4))) if not noGSUBranch: labelMB = Label(self.labels.getNameInc("GW_MB"), comment="Global Write MB") labelMBSK = Label(self.labels.getNameInc("GW_MBSK"), comment="Global Write MBSK") skBackup = kernel["StreamK"] gsuBackup = kernel["GlobalSplitU"] gsuAccumBackup = kernel["_GlobalAccumulation"] bpeCexternalBackup = self.states.bpeCexternal afcBackup = kernel["ActivationFuncCall"] useBiasBackup = self.states.useBias betasBackup = betas assert len(vectorWidths_2) == len(elements_2) # Only full and edge vectorWidths/elements are used in normal kernel if kernel["AdaptiveGemm"] == 0 and len(vectorWidths_2) > 1: vectorWidths_2 = [vectorWidths_2[0], vectorWidths_2[-1]] elements_2 = [elements_2[0], elements_2[-1]] assert len(vectorWidths_1) == len(elements_1) # Only full and edge vectorWidths/elements are used in normal kernel if kernel["AdaptiveGemm"] == 0 and len(vectorWidths_1) > 1: vectorWidths_1 = [vectorWidths_1[0], vectorWidths_1[-1]] elements_1 = [elements_1[0], elements_1[-1]] # GSU0 always sets useBias=NONE (no bias in workspace writes), so no barrier issue. deferGSU0 = ( not noGSUBranch and kernel.get("UseSubtileImpl") and kernel.get("StreamK", 0) > 0 ) gsu0DeferredLabel = None gsu0ReturnLabel = None gsuLimit = 1 if noGSUBranch or self.debugConfig.splitGSU else 2 if gsuLimit > 1: gsuLabel = Label(label=self.labels.getNameInc("GSU"), comment="") if kernel["StreamK"]: if deferGSU0: gsu0DeferredLabel = Label(label=self.labels.getNameInc("GW_B0_Deferred"), comment="") gsu0ReturnLabel = Label(label=self.labels.getNameInc("GW_B0_Deferred_Return"), comment="") # Keep original GSU check unchanged — falls through to GSU0, branches to gsuLabel for GSU1 module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer")) module.add(SCBranchSCC0(labelName=gsuLabel.getLabelName(), comment="Branch to stream-k store code")) sSkt = self.acquireStreamKConstSgpr(kernel, "skTiles") if self.isStreamKConstantsToVgprEnabled(kernel): module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(self.states.skConstVgprs["skTiles"]))) module.add(SCmpEQU32(src0=sgpr(sSkt), src1=1, comment="split == 1 ?")) self.releaseStreamKConstSgpr(sSkt) # TODO May need long branch?? module.add(SCBranchSCC1(labelName=gsuLabel.getLabelName(), comment="branch if split == 1")) else: with self.allocTmpSgpr(1, tag="globalWriteElements_tmpSgprGSU") as tmpSgprGSU: module.add(SAndB32(dst=sgpr(tmpSgprGSU.idx), src0=sgpr("GSU"), src1=self.gsuMaskHex(kernel), comment="Restore GSU")) module.add(SCmpEQU32(src0=sgpr(tmpSgprGSU.idx), src1=1, comment="GSU == 1 ?")) if (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel" or kernel["AdaptiveGemmGSUA"] == 1): module.add(self.longBranchScc1(label=gsuLabel, posNeg=1, comment="long branch if GSU == 1")) else: module.add(SCBranchSCC1(labelName=gsuLabel.getLabelName(), comment="branch if GSU == 1")) gsuLimitRange = range(0, gsuLimit) # generate GSU1 and GSUM label for gsuLimitIdx in gsuLimitRange: # Redirect GSU0 output to deferred module, keeping original label as stub if deferGSU0 and gsuLimitIdx == 0: # Keep original GW_B0 label inline as a stub with jump to deferred gsu0InlineLabel = Label(label=self.labels.getNameInc("GW_B0"), comment="") module.add(gsu0InlineLabel) with self.allocTmpSgpr(2, alignment=2, tag="globalWriteElements_tmpPair") as tmpPair: with self.allocTmpSgpr(1, tag="globalWriteElements_tmpOff") as tmpOff: module.add(SLongBranchPositive(gsu0DeferredLabel, tmpPair, tmpOff, comment="GSU0 reduction (deferred)")) module.addComment0("=" * 60) module.addComment0(" GSU0 reduction block deferred to after persistent loop") module.addComment0(" (would have been inline here in non-deferred version)") module.addComment0("=" * 60) module.add(gsu0ReturnLabel) # Redirect code generation to deferred module savedModule = module module = Module("GSU0_DeferredBlock") module.add(gsu0DeferredLabel) if gsuLimit > 1: betas = betasBackup if gsuLimitIdx == 0: self.states.bpeCexternal = self.states.bpeCinternal if (kernel["_GlobalAccumulation"] != 'MultipleBufferSingleKernel'): self.states.useBias = self.states.useBias if self.states.useBias == DataDirection.WRITE else DataDirection.NONE if self.states.useBias == DataDirection.WRITE and kernel["ProblemType"]["BiasSrc"] == "D": self.states.useBias = DataDirection.NONE if (kernel["_GlobalAccumulation"] != 'MultipleBufferSingleKernel'): kernel["LdsOffsetBias"] = kernel["LdsOffsetBiasGSU"] kernel["ActivationFuncCall"] = False kernel["GlobalSplitU"] = 2 kernel["_GlobalAccumulation"] = kernel["_GlobalAccumulation"] if kernel["StreamK"]: kernel["StreamK"] = 0 kernel["_GlobalAccumulation"] = "MultipleBuffer" vectorWidths = vectorWidths_2 elements = elements_2 else: module.add(gsuLabel) self.states.bpeCexternal = self.states.bpeCexternalGSU1 self.states.useBias = useBiasBackup kernel["LdsOffsetBias"] = kernel["LdsOffsetBiasNonGSU"] kernel["ActivationFuncCall"] = afcBackup kernel["GlobalSplitU"] = 1 kernel["_GlobalAccumulation"] = None kernel["StreamK"] = skBackup if kernel["StreamK"]: kernel["GlobalSplitU"] = gsuBackup kernel["_GlobalAccumulation"] = gsuAccumBackup vectorWidths = vectorWidths_1 elements = elements_1 else: if (kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1) and (kernel["_GlobalAccumulation"] != 'MultipleBufferSingleKernel'): self.states.useBias = self.states.useBias if self.states.useBias == DataDirection.WRITE else DataDirection.NONE kernel["LdsOffsetBias"] = kernel["LdsOffsetBiasGSU"] kernel["ActivationFuncCall"] = False vectorWidths = vectorWidths_2 elements = elements_2 else: kernel["LdsOffsetBias"] = kernel["LdsOffsetBiasNonGSU"] vectorWidths = vectorWidths_1 elements = elements_1 ''' Post process for loop ''' ssslist = [] useSize = [] if gsuLimit > 1 and gsuLimitIdx > 0: if kernel["ProblemType"]["UseScaleAB"]: if not self.states.preloadScaleA: module.add(self.setSgprToInUseState("AddressScaleA")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): module.add(self.setSgprToInUseState("SrdScaleA")) if not self.states.preloadScaleB: module.add(self.setSgprToInUseState("AddressScaleB")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): module.add(self.setSgprToInUseState("SrdScaleB")) if kernel["ProblemType"]["UseScaleAlphaVec"]: module.add(self.setSgprToInUseState("AddressScaleAlphaVec")) module.add(self.setSgprToInUseState("SrdScaleAlphaVec")) isSingleKernel = ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel') or kernel["StreamK"] > 0 # Issue read scale A/B value for later use if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and \ isSingleKernel and \ ((kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["MacDataTypeA"].numRegisters()) or \ (kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["MacDataTypeB"].numRegisters())): assert(kernel["ProblemType"]["ComputeDataType"].isSingle()) sgprScaleA = self.sgprPool.checkOut(1, tag="globalWriteElements_sgprScaleA", preventOverflow=False) sgprScaleB = self.sgprPool.checkOut(1, tag="globalWriteElements_sgprScaleB", preventOverflow=False) for i,name in enumerate(['A','B']): if kernel["ProblemType"]["DataType%s"%name].numRegisters() <= kernel["ProblemType"]["MacDataType%s"%name].numRegisters(): sgprScale = sgprScaleA if name == 'A' else sgprScaleB module.add(SMovB32(dst=sgpr(sgprScale), src=1.0 , comment="init as 1" )) label = Label(self.labels.getNameInc("Scale%sValid"%name), "") module.add(BranchIfZero("AddressScale%s"%name, DataType('int64').toEnum(), -1, int(kernel["WavefrontSize"]/32), label, kernel["WavefrontSize"])) # load scale data module.add(SLoadB32(dst=sgpr(sgprScale), base=sgpr("AddressScale%s"%name,2), soffset=0, comment="load scale%s"%name)) module.add(label) # Issue read scale C/D value for later use if kernel["ProblemType"]["UseScaleCD"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0): module.add(SMovB32(dst=sgpr("ScaleD"), src=1.0 , comment="init as 1" )) module.add(SMovB32(dst=sgpr("ScaleD+1"), src=1.0 , comment="init as 1" )) label = Label(self.labels.getNameInc("ScaleDValid"), "") module.add(BranchIfZero("AddressScaleD", DataType('int64').toEnum(), -1, int(kernel["WavefrontSize"]/32), label, kernel["WavefrontSize"])) # load scale data module.add(SLoadB32(dst=sgpr("ScaleD"), base=sgpr("AddressScaleD",2), soffset=0, comment="load scaleD")) module.add(label) sgprScaleC = self.sgprPool.checkOut(1, tag="globalWriteElements_sgprScaleC", preventOverflow=False) module.add(SMovB32(dst=sgpr(sgprScaleC), src=1.0 , comment="init as 1" )) label = Label(self.labels.getNameInc("ScaleCValid"), "") module.add(BranchIfZero("AddressScaleC", DataType('int64').toEnum(), -1, int(kernel["WavefrontSize"]/32), label, kernel["WavefrontSize"])) # load scale data module.add(SLoadB32(dst=sgpr(sgprScaleC), base=sgpr("AddressScaleC",2), soffset=0, comment="load scaleC")) module.add(label) factorDims = [0] if self.states.FactorDim == 3: factorDims.append(1) elif self.states.FactorDim == 2: factorDims = [1] vectorDataTypes = VectorDataTypes() if (kernel["ProblemType"]["UseScaleAlphaVec"]) and isSingleKernel: labelStr = self.labels.getNameInc("ScaleAlphaVec") # Init ScaleAlphaVec Srd (buffer load only; flat load uses pre-computed 64-bit address) if kernel["BufferLoad"]: if self.states.FactorDim == 3: with self.allocTmpSgpr(1,1, tag="globalWriteElements_tmpSgprRes") as tmpSgprRes: tmpSgpr = tmpSgprRes.idx module.add(self.getSCMPKInstruction("EQU32", "FactorDim", 0, comment="FactorDim == 0")) module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=sgpr("SizeI"), src1=sgpr("SizeJ"))) module.add(self.allocPostLoopSrdSuppress("ScaleAlphaVec", labelStr, sgprLength=sgpr(tmpSgpr))) elif self.states.FactorDim == 2: module.add(self.allocPostLoopSrdSuppress("ScaleAlphaVec", labelStr, sgprLength=sgpr("SizeJ"))) else: module.add(self.allocPostLoopSrdSuppress("ScaleAlphaVec", labelStr, sgprLength=sgpr("SizeI"))) module.add(SMulI32(dst=sgpr("SrdScaleAlphaVec+2"), src0=hex(self.states.bpeCinternal), src1=sgpr("SrdScaleAlphaVec+2"), comment="ScaleAlphaVec scaled by BPE"))# scaled by BPE for d in range(len(factorDims)): vectorDataTypes.scaleAlpha(d).dataType = kernel["ProblemType"]["ComputeDataType"] # Add ScaleABVec support here # Issue read scale A/B vector value for later use if ((kernel["ProblemType"]["UseScaleAB"] == "Vector")) and isSingleKernel: labelStrA = self.labels.getNameInc("ScaleAVec") labelStrB = self.labels.getNameInc("ScaleBVec") # Init ScaleABVec Srd (buffer load only; flat load uses pre-computed 64-bit address) if kernel["BufferLoad"]: module.add(self.allocPostLoopSrdSuppress("ScaleA", labelStrA, sgprLength=sgpr("SizeI"))) module.add(self.allocPostLoopSrdSuppress("ScaleB", labelStrB, sgprLength=sgpr("SizeJ"))) module.add(SMulI32(dst=sgpr("SrdScaleA+2"), src0=hex(self.states.bpeCinternal), src1=sgpr("SrdScaleA+2"), comment="ScaleAVec scaled by BPE"))# scaled by BPE module.add(SMulI32(dst=sgpr("SrdScaleB+2"), src0=hex(self.states.bpeCinternal), src1=sgpr("SrdScaleB+2"), comment="ScaleBVec scaled by BPE"))# scaled by BPE module.add(self.shiftSrd("ScaleA")) module.add(self.shiftSrd("ScaleB")) vectorDataTypes.scaleA.dataType = kernel["ProblemType"]["ComputeDataType"] vectorDataTypes.scaleB.dataType = kernel["ProblemType"]["ComputeDataType"] factorDim0Label = Label(self.labels.getNameInc("Load_FactorDim_0"), "") factorDim1Label = Label(self.labels.getNameInc("Load_FactorDim_1"), "") # Add bias lds isLdsLoaded = False if self.states.useBias == DataDirection.READ and isSingleKernel: # Init bias Srd labelStr = self.labels.getNameInc("Bias") with self.allocTmpSgpr(1,1, tag="globalWriteElements_tmpSgprRes2") as tmpSgprRes: set_bs_label = Label(self.labels.getNameInc("Dont_Set_BiasStride"), "") tmpSgpr = tmpSgprRes.idx module.add(SAddU32(dst=sgpr(tmpSgpr), src0=sgpr("WorkGroup2"), src1=hex(1))) module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr("BiasStride"), src1=sgpr(tmpSgpr), comment="stride * (wg+1)")) module.add(SCmpEQU32(sgpr(tmpSgpr), 0, comment="bias stride = 0?")) if self.states.FactorDim == 3: module.add(SCBranchSCC0(set_bs_label.getLabelName())) module.add(self.getSCMPKInstruction("EQU32", "FactorDim", 0, comment="FactorDim == 0")) module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=sgpr("SizeI"), src1=sgpr("SizeJ"))) module.add(set_bs_label) elif self.states.FactorDim == 2: module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=sgpr("SizeJ"), src1=sgpr(tmpSgpr))) else: module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=sgpr("SizeI"), src1=sgpr(tmpSgpr))) # module.add(SMovB32(dst=sgpr(tmpSgpr), src=128, comment="set bias stride")) # Init bias Srd (buffer load only; flat load uses the pre-computed 64-bit address) if kernel["BufferLoad"]: module.add(self.allocPostLoopSrdSuppress("Bias", labelStr, sgprLength=sgpr(tmpSgpr))) loadBiasEndLabel = Label(self.labels.getNameInc("Load_Bias_End"), "") if self.states.FactorDim == 3: module.add(factorDim0Label) module.add(self.getSCMPKInstruction("LGU32", "FactorDim", 0, comment="FactorDim != 0")) module.add(SCBranchSCC1(factorDim1Label.getLabelName(), "Branch if true")) for d in range(len(factorDims)): # Calculate max vgpr for bias read vectorDataTypes.bias(d).dataType = kernel["ProblemType"]["BiasDataTypeList"][0] totalTmpVgpr = self.getNumOfTempVgprs(vectorDataTypes, kernel, 1, factorDims[d]) tmpVgpr = self.vgprPool.checkOutAligned(totalTmpVgpr, 2, "store tmps") tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=4) if d == 1: module.add(factorDim1Label) multiBiasTypeLabel = [] for i in kernel["ProblemType"]["BiasDataTypeList"]: name = self.labels.getNameInc("Load_Bias%s_%u"%(i.toNameAbbrev(), factorDims[d])) multiBiasTypeLabel.append(Label(name, "")) multiBiasTypeLabel.append(loadBiasEndLabel) offsetVgpr = self.vgprPool.checkOut(1, tag="globalWriteElements_offsetVgpr") if len(kernel["ProblemType"]["BiasDataTypeList"]) == 1: vectorDataTypes.bias(d).dataType = kernel["ProblemType"]["BiasDataTypeList"][0] module.add(self.readVectorToLDS(vectorDataTypes, kernel, 1, offsetVgpr, tmpVgprRes, factorDims[d])) if len(factorDims) == 2: if d == 0: module.add(SBranch(labelName=loadBiasEndLabel.getLabelName(), comment="Branch to load bias end")) else: module.add(loadBiasEndLabel) else: for i, label in enumerate(multiBiasTypeLabel[1:]): typeValue = kernel["ProblemType"]["BiasDataTypeList"][i].value module.add(multiBiasTypeLabel[i]) #module.add(SCmpKLGU32(sgpr("BiasType"), typeValue, "BiasType != %u"%typeValue)) module.add(self.getSCMPKInstruction("LGU32", "BiasType", typeValue, comment="BiasType != %u"%typeValue)) module.add(SCBranchSCC1(label.getLabelName(), "Branch if true")) vectorDataTypes.bias(d).dataType = kernel["ProblemType"]["BiasDataTypeList"][i] module.add(self.readVectorToLDS(vectorDataTypes, kernel, 1, offsetVgpr, tmpVgprRes, factorDims[d])) module.add(SBranch(labelName=loadBiasEndLabel.getLabelName(), comment="Branch to load bias end")) if d == len(factorDims) -1: module.add(loadBiasEndLabel) isLdsLoaded = True self.vgprPool.checkIn(offsetVgpr) self.vgprPool.checkIn(tmpVgpr) elif self.states.useBias == DataDirection.WRITE: labelStr = self.labels.getNameInc("Bias") if kernel["ProblemType"]["BiasSrc"] == "A" or kernel["ProblemType"]["BiasSrc"] == "B": # Calculate max vgpr for bias write A, B tP = tPA if kernel["ProblemType"]["BiasSrc"] == "A" else tPB tile01 = tP["tile01Idx"] maxKId = self.states.lraTileProperties[tile01].maxKId biasMaxVgpr = kernel["VectorWidthA"] * kernel["ProblemType"]["ComputeDataType"].numRegisters() * maxKId maxAlign = max(1, (kernel["VectorWidthA"] + 1) // 2 * 2) tmpVgpr = self.vgprPool.checkOutAligned(biasMaxVgpr, maxAlign, "store tmps") tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=biasMaxVgpr) # Skip bias store skipGlobalStoreLabel = Label(self.labels.getNameInc("SkipBiasStore"), comment="") wgIdx = 1 if tile01 == 0 else 0 #module.add(SCmpKLGU32(sgpr("WorkGroup%d"%wgIdx), 0, "WorkGroup%d != 0"%wgIdx)) module.add(self.getSCMPKInstruction("LGU32", "WorkGroup%d"%wgIdx, 0, comment="WorkGroup%d != 0"%wgIdx)) module.add(SCBranchSCC1(skipGlobalStoreLabel.getLabelName(), "Branch if true")) if (kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1) and kernel["_GlobalAccumulation"] == "MultipleBuffer": sourceAddress = "D" else: sourceAddress = "Bias" numRecordsStr = "SizeI" if kernel["ProblemType"]["BiasSrc"] == "A" else "SizeJ" # Init bias Srd (buffer store only; flat store uses the pre-computed 64-bit address) if kernel["BufferStore"]: module.add(self.allocPostLoopSrdSuppressRaw("Bias", sourceAddress, labelStr, sgprLength=sgpr(numRecordsStr))) if sourceAddress == "D": module.add(self.undefineSgpr("AddressD")) multiBiasTypeLabel = [] for i in kernel["ProblemType"]["BiasDataTypeList"]: name = self.labels.getNameInc("Write_Bias%s"%i.toNameAbbrev()) multiBiasTypeLabel.append(Label(name, "")) writeBiasEndLabel = Label(self.labels.getNameInc("Write_Bias_End"), "") multiBiasTypeLabel.append(writeBiasEndLabel) # Get gwvw ''' gwvw is set to max(mt // kernel["NumThreads"], kernel["VectorWidthA"]) instead of kernel["VectorWidthA"] is that we don't want batch exists. If VW is set to 1, MT=512, and flat work group = 256. We will have to set gwvw to 2 to store all the bias data. ''' tile01 = tP["tile01Idx"] mt = kernel["MacroTile%u" % tile01] gwvw = min(max(mt // kernel["NumThreads"], kernel["VectorWidthA"]), tP["glvw"]) offsetVgpr = self.vgprPool.checkOut(gwvw, tag="globalWriteElements_Bias_offsetVgpr") with self.allocTmpSgpr(5, 2, tag="globalWriteElements_tmpSgprRes3") as tmpSgprRes: if kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1: module.add(self.writeBiasToGlobal(kernel["ProblemType"]["ComputeDataType"], kernel, tP, gwvw, offsetVgpr, tmpSgprRes, tmpVgprRes)) elif len(kernel["ProblemType"]["BiasDataTypeList"]) == 1: module.add(self.writeBiasToGlobal(kernel["ProblemType"]["BiasDataTypeList"][0], kernel, tP, gwvw, offsetVgpr, tmpSgprRes, tmpVgprRes)) else: for i, label in enumerate(multiBiasTypeLabel[1:]): typeValue = kernel["ProblemType"]["BiasDataTypeList"][i].value module.add(multiBiasTypeLabel[i]) #module.add(SCmpKLGU32(sgpr("BiasType"), typeValue, "BiasType != %u"%typeValue)) module.add(self.getSCMPKInstruction("LGU32", "BiasType", typeValue, comment="BiasType != %u"%typeValue)) module.add(SCBranchSCC1(label.getLabelName(), "Branch if true")) module.add(self.writeBiasToGlobal(kernel["ProblemType"]["BiasDataTypeList"][i], kernel, tP, gwvw, offsetVgpr, tmpSgprRes, tmpVgprRes)) module.add(SBranch(labelName=writeBiasEndLabel.getLabelName(), comment="Branch to write bias end")) module.add(writeBiasEndLabel) self.vgprPool.checkIn(offsetVgpr) self.vgprPool.checkIn(tmpVgpr) module.add(skipGlobalStoreLabel) else: # Init bias Srd (buffer store only; flat store uses the pre-computed 64-bit address) if kernel["BufferStore"]: module.add(self.allocPostLoopSrdSuppress("Bias", labelStr, "BufferOOB")) module.add(self.shiftSrd("Bias")) ssslist.append("Bias") useSize.append(True) if vectorDataTypes.isValid() and (not isLdsLoaded): if self.states.FactorDim == 3: module.add(factorDim0Label) module.add(self.getSCMPKInstruction("LGU32", "FactorDim", 0, comment="FactorDim != 0")) module.add(SCBranchSCC1(factorDim1Label.getLabelName(), "Branch if true")) labelDimEnd = Label(self.labels.getNameInc("MultiDimEnd"), "") for d in range(len(factorDims)): totalTmpVgpr = self.getNumOfTempVgprs(vectorDataTypes, kernel, 1, factorDims[d]) tmpVgpr = self.vgprPool.checkOutAligned(totalTmpVgpr, 2, "store tmps") tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=4) offsetVgpr = self.vgprPool.checkOut(1, tag="globalWriteElements_offsetVgpr2") if d == 1: module.add(factorDim1Label) module.add(self.readVectorToLDS(vectorDataTypes, kernel, 1, offsetVgpr, tmpVgprRes, factorDims[d])) if self.states.FactorDim == 3 and d == 0: module.add(SBranch(labelName=labelDimEnd.getLabelName(), comment="Branch to load end")) self.vgprPool.checkIn(offsetVgpr) self.vgprPool.checkIn(tmpVgpr) if len(factorDims) > 1: module.add(labelDimEnd) # Undefine LDS load related sgprs if gsuLimit > 1 and gsuLimitIdx == 0: if kernel["ProblemType"]["UseScaleAB"]: if not self.states.preloadScaleA: module.add(self.setSgprToFreeState("AddressScaleA")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): module.add(self.setSgprToFreeState("SrdScaleA")) if not self.states.preloadScaleB: module.add(self.setSgprToFreeState("AddressScaleB")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): module.add(self.setSgprToFreeState("SrdScaleB")) if kernel["ProblemType"]["UseScaleAlphaVec"]: module.add(self.setSgprToFreeState("AddressScaleAlphaVec")) module.add(self.setSgprToFreeState("SrdScaleAlphaVec")) else: if kernel["ProblemType"]["UseScaleAB"]: if not self.states.preloadScaleA: module.add(self.undefineSgpr("AddressScaleA")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): module.add(self.undefineSgpr("SrdScaleA")) if not self.states.preloadScaleB: module.add(self.undefineSgpr("AddressScaleB")) if (kernel["ProblemType"]["UseScaleAB"] == "Vector"): module.add(self.undefineSgpr("SrdScaleB")) if kernel["ProblemType"]["UseScaleAlphaVec"]: module.add(self.undefineSgpr("AddressScaleAlphaVec")) module.add(self.undefineSgpr("SrdScaleAlphaVec")) if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0) or kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel') and \ ((kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["MacDataTypeA"].numRegisters()) or \ (kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["MacDataTypeB"].numRegisters())): assert(kernel["ProblemType"]["ComputeDataType"].isSingle()) newAlphaVgpr = self.vgprPool.checkOut(1, tag="globalWriteElements_newAlphaVgpr") module.add(VMovB32(dst=vgpr(newAlphaVgpr), src=sgpr("Alpha"))) module.add(SWaitCnt(kmcnt=0, comment="wait for scaleAB load")) if kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["MacDataTypeA"].numRegisters(): module.add(VMulF32(dst=vgpr(newAlphaVgpr), src0=vgpr(newAlphaVgpr), src1=sgpr(sgprScaleA))) if kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["MacDataTypeB"].numRegisters(): module.add(VMulF32(dst=vgpr(newAlphaVgpr), src0=vgpr(newAlphaVgpr), src1=sgpr(sgprScaleB))) module.add(SNop(waitState=0, comment="1 wait states")) if kernel["StreamK"] > 0: oldAlpha = self.sgprPool.checkOut(1, tag="globalWriteElements_oldAlpha", preventOverflow=False) module.add(SMovB32(dst=sgpr(oldAlpha), src=sgpr("Alpha"), comment="Save alpha value")) module.add(VReadfirstlaneB32(dst=sgpr("Alpha"), src=vgpr(newAlphaVgpr), comment="Update Alpha")) self.vgprPool.checkIn(newAlphaVgpr) self.sgprPool.checkIn(sgprScaleA) self.sgprPool.checkIn(sgprScaleB) # Update beta with ScaleC (only when Beta is actually used) if kernel["ProblemType"]["UseBeta"] and kernel["ProblemType"]["UseScaleCD"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0): assert(kernel["ProblemType"]["ComputeDataType"].isSingle()) newBetaVgpr = self.vgprPool.checkOut(1, tag="globalWriteElements_newBetaVgpr") module.add(VMovB32(dst=vgpr(newBetaVgpr), src=sgpr("Beta"))) if (not ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel')) or kernel["StreamK"] > 0: module.add(SWaitCnt(kmcnt=0, comment="wait for scaleC load")) module.add(VMulF32(dst=vgpr(newBetaVgpr), src0=vgpr(newBetaVgpr), src1=sgpr(sgprScaleC))) module.add(SNop(waitState=0, comment="1 wait states")) module.add(VReadfirstlaneB32(dst=sgpr("Beta"), src=vgpr(newBetaVgpr), comment="Update Beta")) self.vgprPool.checkIn(newBetaVgpr) self.sgprPool.checkIn(sgprScaleC) # Copy scaleD for PK calculations module.add(SMovB32(dst=sgpr("ScaleD+1"), src=sgpr("ScaleD"))) if kernel["ProblemType"]["UseE"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0): # Update E offset1 strideE1 = "StrideE%s" % (self.states.indexChars[kernel["PackedC1IndicesX"][0]]) module.add(VMulLOU32(dst=vgpr(self.vgprs.coutRowPtrE), src0=vgpr(self.vgprs.coutRowPtrE), src1=sgpr(strideE1), comment=" offset 1")) labelEStr = self.labels.getNameInc("E") # Init E Srd (buffer store only; flat store uses the pre-computed 64-bit address) if kernel["BufferStore"]: module.add(self.allocPostLoopSrdSuppress("E", labelEStr, "BufferOOB")) module.add(self.shiftSrd("E")) ssslist.append("E") useSize.append(False) if self.states.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"): if self.vgprs.coutRowPtrGate != -1: module.add(VMulLOU32(dst=vgpr(self.vgprs.coutRowPtrGate), src0=vgpr(self.vgprs.coutRowPtrGate), src1=sgpr("GateStride+0"), comment="finalize coutRowPtrGate = row * GateStride+0")) labelGateStr = self.labels.getNameInc("Gate") gateAddrEven = (self.sgprs.get("AddressGate", 0) % 2 == 0) module.add(self.allocPostLoopSrdSuppress("Gate", labelGateStr, "BufferOOB", evenAlignedAddr=gateAddrEven)) module.add(self.shiftSrd("Gate")) ssslist.append("Gate") useSize.append(False) if ssslist: module.add(self.computeStoreSrdStart(kernel, ssslist, useSize=useSize, noMultipleBuffer=True)) # s = gate null ? 1.0 : 0.0. FMA does acc = (gate+s)*acc + gate. if self.states.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"): module.add(self.getSCMPKInstruction("EQU32", "SrdGate+2", 0, comment="gate null? (SrdGate num_records==0)")) module.add(SCSelectB32(dst=sgpr("GateNullOne"), src0=hex(0x3f800000), src1=0, comment="GateNullOne = (gate null) ? 1.0 : 0.0")) ''' Post process for loop end ''' atomic = (kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1) and (kernel["_GlobalAccumulation"] != 'MultipleBuffer' and kernel["_GlobalAccumulation"] != 'MultipleBufferSingleKernel') activation = self.exclasses.activation # write possibilities and labels # if beta/fd combo not specified fall back to global param definition if betas is None: hasBeta = kernel["ProblemType"]["UseBeta"] and \ (kernel["_GlobalAccumulation"] != 'MultipleBuffer' or \ kernel["GlobalSplitUAlgorithm"] == 'MultipleBufferSingleKernel' or \ kernel["GlobalSplitU"] == 0) betas = [False, True] if hasBeta else [False] # Avoid wrong edge arg, making edge true for many vectorWidths if edge == False and len(vectorWidths) > 1: edge = True if factorDims is None: factorDims = [0, 1] if self.states.FactorDim == 3 else [1] if self.states.FactorDim == 2 else [0] writeLabels = {} for beta in betas: writeLabels[beta] = {} writeLabels[beta]["Label"] = Label(self.labels.getNameInc("GW_B%u" % (beta) ), "") for factorDim in factorDims: writeLabels[beta][factorDim] = {} writeLabels[beta][factorDim]["Label"] = Label(self.labels.getNameInc("GW_B%u_FD%u" % (beta, factorDim) ), "") generatedVectorWidths = set() for vectorWidthIdx, vectorWidth in enumerate(reversed(vectorWidths)): # Avoid duplicated vectorWidth if vectorWidth not in generatedVectorWidths: generatedVectorWidths.add(vectorWidth) else: continue writeLabels[beta][factorDim][vectorWidth] = {} writeLabels[beta][factorDim][vectorWidth]["NonEdge"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_NonEdge" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["NonEdgeEnd"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_NonEdgeEnd" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["Then"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_Then" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["Else"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_Else" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["ThenDeferred"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_Then_Deferred" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["ThenDeferredReturn"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_Then_Deferred_Return" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["ElseDeferred"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_Else_Deferred" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["ElseDeferredReturn"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_Else_Deferred_Return" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["NonEdgeDeferred"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_NonEdge_Deferred" % (beta, factorDim, vectorWidth) ), "") writeLabels[beta][factorDim][vectorWidth]["NonEdgeDeferredReturn"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_NonEdge_Deferred_Return" % (beta, factorDim, vectorWidth) ), "") endLabel = Label(self.labels.getNameInc("GW_End"), "") # Layout """ if B1 goto label_B1 if BD1 goto label_B0_FD1 label_B0_FD0: edge writes goto label_End label_B0_FD1: edge writes goto label_End label_B1: if BD1 goto label_B1_FD1 label_B1_FD0: edge writes goto label_End label_B1_FD1: edge writes goto label_End """ ######################################## # Vgprs maxAlign = 2 if kernel["BufferStore"]: numTmpVgpr = 2 if len(kernel["PackedC0IndicesX"]) > 1: numTmpVgpr += 1 else: numTmpVgpr = 2 + 3 # GLOBAL_OFFSET_C needs 3, plus 2 tmps? # BufferOOB if (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or kernel["BufferStore"]: numTmpVgpr += 1 # dot2: for WaveSplitK reduction if kernel["NumWaveSplitK"] > 1: numTmpVgpr += 1 # Get max vgpr and sgpr for activation actPCGwvwVgpr = 0 actPCMaxTempSgpr = 0 actTempSgpr = 0 actExportType = ActivationType.Export.GRADONLY if kernel["ProblemType"]["Gradient"] else ActivationType.Export.NORMAL if kernel["ActivationFuncCall"] or \ ((((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0) and kernel["ActivationFused"]) and \ (kernel["ProblemType"]["ActivationType"] != 'none')): maxVw = max(vectorWidths) # Here is where activation creates cache if cache is enabled usage = activation.getAllGprUsage(kernel["ProblemType"]["ActivationComputeDataType"], kernel["ProblemType"]["ActivationType"], exportType=actExportType) actPCMaxTempVgpr = 0 for _, gprs in usage.items(): actPCMaxTempVgpr = max(actPCMaxTempVgpr, gprs["vgpr"]) actPCMaxTempSgpr = max(actPCMaxTempSgpr, gprs["sgpr"]) actPCGwvwVgpr = int(ceil(maxVw * kernel["ProblemType"]["ActivationComputeDataType"].numRegisters())) numTmpVgpr = max(numTmpVgpr, actPCMaxTempVgpr + actPCGwvwVgpr) if kernel["ProblemType"]["UseE"] and (not kernel["ProblemType"]["Gradient"]): maxVw = max(vectorWidths) gwvwVgpr = int(ceil(maxVw * kernel["ProblemType"]["ActivationComputeDataType"].numRegisters())) if kernel["ActivationFuncCall"]: gwvwVgpr += actPCMaxTempVgpr + actPCGwvwVgpr numTmpVgpr = max(numTmpVgpr, gwvwVgpr) tmpVgpr = self.vgprPool.checkOutAligned(numTmpVgpr, maxAlign, "store tmps") tmpVgpr = ContinuousRegister(idx=tmpVgpr, size=numTmpVgpr) cvtVgprStruct = None cvtVgpr = None is16bitHPA = (kernel["ProblemType"]["DestDataType"].isBFloat16() or kernel["ProblemType"]["DestDataType"].isHalf()) and \ kernel["ProblemType"]["HighPrecisionAccumulate"] if is16bitHPA: # For UseSubtileImpl, allocate 7 vgprs with 2-alignment (64-bit aligned) so # that the first 4 (reused as pack scratch for the paired 16bit store) satisfy # the buffer_store_dwordx4 alignment requirement. Any pool vgpr skipped for # alignment becomes a hole that subsequent element-address checkouts fill, so # pool.size() (= startVgprValu) stays within budget for large macro-tiles. # If a very large tile causes accvgpr staging to exceed 256 vgprs despite the # alignment overhead, reduce the batch via NumElementsPerBatchStore. # +0..+3: scratch for pack output + ds_bpermute + v_permlane32_swap # (reuses vgprBf16Temp/Mask/Nan/Inc slots; constants written at batch # start are overwritten by the packed 16bit values before the store) # +4: vgprPermAddr — ds_permute partner-lane byte address # +5: vgprLaneGroupDelta — lane_group*8, pre-computed once per batch # +6: vgprAddrScratch — per-store adjusted D address; avoids modifying addrDVgpr numCvtVgprs = 7 if kernel.get("UseSubtileImpl") else 4 cvtAlign = 2 if kernel.get("UseSubtileImpl") else 1 cvtVgpr = self.vgprPool.checkOutAligned(numCvtVgprs, cvtAlign, tag="globalWriteElements_cvtVgpr") cvtVgprStruct = self.BF16CVTVgprStruct(vgprBf16Temp=cvtVgpr, vgprBf16Mask=(cvtVgpr+1), \ vgprFp32Nan=(cvtVgpr+2), vgprBf16Inc=(cvtVgpr+3), \ vgprPermAddr=(cvtVgpr+4) if kernel.get("UseSubtileImpl") else -1, \ vgprLaneGroupDelta=(cvtVgpr+5) if kernel.get("UseSubtileImpl") else -1, \ vgprAddrScratch=(cvtVgpr+6) if kernel.get("UseSubtileImpl") else -1) elif kernel["ProblemType"]["DestDataType"].isAnyFloat8() and kernel["ProblemType"]["HighPrecisionAccumulate"]: cvtVgpr = self.vgprPool.checkOut(4, tag="globalWriteElements_cvtVgpr2") cvtVgprStruct = self.FP8CVTVgprStruct(vgprFp8Temp=cvtVgpr, vgprFp8NanInf=(cvtVgpr+1), \ vgprFp8Min=(cvtVgpr+2), vgprFp8Max=(cvtVgpr+3)) elif kernel["ProblemType"]["DestDataType"].isAnyBFloat8(): cvtVgpr = self.vgprPool.checkOut(4, tag="globalWriteElements_cvtVgpr3") cvtVgprStruct = self.BF8CVTVgprStruct(vgprBF8Temp=cvtVgpr, vgprBF8NanInf=(cvtVgpr+1), \ vgprBF8Min=(cvtVgpr+2), vgprBF8Max=(cvtVgpr+3)) elif kernel["ProblemType"]["DestDataType"].isInt8(): cvtVgpr = self.vgprPool.checkOut(4, tag="globalWriteElements_cvtVgpr4") cvtVgprStruct = self.I8CVTVgprStruct(vgprI8Temp0=cvtVgpr, vgprI8Temp1=(cvtVgpr+1), \ vgprI8Mask0=(cvtVgpr+2), vgprI8Mask1=(cvtVgpr+3)) activationSetPCStruct = None activationLabelList = None activationEnumStrList = None toActModuleList = None isInsertActFunctionCallAddrCalc = True if kernel["ActivationFuncCall"]: sgprOffsetActivation = self.sgprPool.checkOutAligned(2, 2, tag="globalWriteElements_sgprOffsetActivation", preventOverflow=False) sgprOffsetBack = self.sgprPool.checkOutAligned(2, 2, tag="globalWriteElements_sgprOffsetBack", preventOverflow=False) activationCDataType = kernel["ProblemType"]["ActivationComputeDataType"] activationLabelList = {} toActModuleList = {} supportedBy = ActivationType.SupportedBy.ALL if kernel["ProblemType"]["ActivationType"] == 'all' else ActivationType.SupportedBy.HIPBLASLT activationEnumStrList = ActivationType.getEnumStrList(activationCDataType, supportedBy, exportType=actExportType) for gwvw in vectorWidths: if gwvw in activationLabelList: continue activationLabelList[gwvw] = [] toActModuleList[gwvw] = [] for enumStr in activationEnumStrList: name = self.labels.getNameInc("Activation_%s_VW%u"% (enumStr.capitalize(), gwvw)) activationLabelList[gwvw].append(Label(name, "")) toActModuleList[gwvw].append(Label("To_%s"% (name), "")) calleeLabelsByGwvw = { gwvw: tuple(label.getLabelName() for label in labels) for gwvw, labels in activationLabelList.items() } activationSetPCStruct = self.ActivationSetPCStruct(sgprOffsetActivation=sgprOffsetActivation, \ sgprOffsetBack=sgprOffsetBack, vgprActCopy=tmpVgpr.idx, \ calleeLabelsByGwvw=calleeLabelsByGwvw) # Add branch here if all elements are identical if vectorWidths.count(vectorWidths[0]) == len(vectorWidths): isInsertActFunctionCallAddrCalc = False module.add(self.insertActFunctionCallAddrCalc(activationSetPCStruct.sgprOffsetActivation, \ vectorWidths[0], toActModuleList, activationEnumStrList, activationLabelList)) ######################################## # Sgprs # allocate tmps for the store header (before the batch implementations) # branch B1 or B0 skPartialsLabel = Label(label=self.labels.getNameInc("SK_Partials"), comment="") skComponent = Component.StreamK.find(self) module.add(skComponent.storeBranches(self, kernel, skPartialsLabel, vectorWidths_1, elements_1, tmpVgpr.idx, cvtVgprStruct)) # support dynamic MBSK/MB selection by checking synchronizer after bias write if kernel["AdaptiveGemmGSUA"] == 1: if skBackup == 0: if gsuLimit > 1 and gsuLimitIdx == 0: if (kernel["_GlobalAccumulation"] == "MultipleBuffer" or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"): module.add(SCmpEQU64(src0=sgpr("Synchronizer", 2), src1=hex(0), comment="Check for synchronizer")) module.add(self.longBranchScc1(label=labelMB, posNeg=1, comment="long branch if Synchronizer is null")) gsuComponent = Component.GSU.find(self) if kernel["GlobalSplitU"] > 1 and kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel": module.add(self.SrdTDInit(kernel)) reductionStartLabel = Label(self.labels.getNameInc("Reduction_Start"), comment="Reduction start") reductionEndLabel = Label(self.labels.getNameInc("Reduction_End"), comment="Reduction end") module.add(reductionStartLabel) self.defineSgpr("GSUStartWGIdx", 1) module.add(RegSet("s", "sgprGSUStartWGIdx", self.sgprs["GSUStartWGIdx"])) module.add(gsuComponent.reductionBranches(self, kernel, tPB, vectorWidths_1, elements_1, tmpVgpr, cvtVgprStruct, \ vectorDataTypes, factorDims, reductionEndLabel, endLabel)) module.add(self.undefineSgpr("GSUStartWGIdx")) module.add(reductionEndLabel) module.addSpaceLine() # Generate write labels for different global write modes globalWriteModes = [] hasMultipleGlobalWriteModes = False if gsuLimitIdx == 0 and kernel["GlobalSplitU"] > 0: if kernel["AdaptiveGemmGSUA"] == 0: if kernel["GlobalSplitUAlgorithm"] == "MultipleBuffer": globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] elif kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel": globalWriteModes = ["OptNLL_MBSK"] if noGSUBranch else ["MBSK"] elif kernel["GlobalSplitUAlgorithm"] == "SingleBuffer": globalWriteModes = ["OptNLL_SB"] if noGSUBranch else ["SB"] else: if kernel["GlobalSplitUAlgorithm"] == "MultipleBuffer": # StreamK and dot2 cannot be enabled with MBSK globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] if skBackup or kernel["UseDotInstruction"] else ["MBSK", "MB"] hasMultipleGlobalWriteModes = False if noGSUBranch else False if skBackup or kernel["UseDotInstruction"] else True if kernel["NumElementsPerBatchStore"] == 1: # too many store at MultipleBufferSingleKernel direct reject globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] hasMultipleGlobalWriteModes = False if kernel["ProblemType"]["UseScaleCD"]: # MultipleBufferSingleKernel not support UseScaleCD yet globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] hasMultipleGlobalWriteModes = False if kernel["ProblemType"]["UseE"]: # MultipleBufferSingleKernel not support UseE yet globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] hasMultipleGlobalWriteModes = False if kernel["ProblemType"]["BiasSrc"] != "D": # MultipleBufferSingleKernel not support BiasSrc not D yet globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] hasMultipleGlobalWriteModes = False if kernel["ProblemType"]["DataType"].isDouble(): # MultipleBufferSingleKernel not support double yet globalWriteModes = ["OptNLL_MB"] if noGSUBranch else ["MB"] hasMultipleGlobalWriteModes = False elif kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel": globalWriteModes = ["OptNLL_MBSK"] if noGSUBranch else ["MBSK", "MB"] hasMultipleGlobalWriteModes = False if noGSUBranch else True elif kernel["GlobalSplitUAlgorithm"] == "SingleBuffer": globalWriteModes = ["OptNLL_SB"] if noGSUBranch else ["SB"] hasMultipleGlobalWriteModes = False else: globalWriteModes = ["GSU1"] hasMultipleGlobalWriteModes = False for globalWriteMode in globalWriteModes: for beta in betas: writeLabels[beta][globalWriteMode] = {} writeLabels[beta][globalWriteMode]["Label"] = Label(self.labels.getNameInc("GW_B%u_%s" % (beta, globalWriteMode) ), "") for factorDim in factorDims: writeLabels[beta][factorDim][globalWriteMode]= {} writeLabels[beta][factorDim][globalWriteMode]["Label"] = Label(self.labels.getNameInc("GW_B%u_FD%u_%s" % (beta, factorDim, globalWriteMode) ), "") for vectorWidthIdx, vectorWidth in enumerate(reversed(vectorWidths)): writeLabels[beta][factorDim][vectorWidth][globalWriteMode] = {} writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["NonEdge"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_%s_NonEdge" % (beta, factorDim, vectorWidth, globalWriteMode) ), "") writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["NonEdgeEnd"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_%s_NonEdgeEnd" % (beta, factorDim, vectorWidth, globalWriteMode) ), "") writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["Then"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_%s_Then" % (beta, factorDim, vectorWidth, globalWriteMode) ), "") writeLabels[beta][factorDim][vectorWidth][globalWriteMode]["Else"] = Label(self.labels.getNameInc("GW_B%u_FD%u_VW%u_%s_Else" % (beta, factorDim, vectorWidth, globalWriteMode) ), "") # Generate beta modules for each global write mode activationLabelListBackup = activationLabelList for globalWriteMode in globalWriteModes: if kernel["AdaptiveGemmGSUA"] == 1: if not noGSUBranch: if globalWriteMode == "MBSK": module.add(labelMBSK) kernel["_GlobalAccumulation"] = "MultipleBufferSingleKernel" elif globalWriteMode == "MB": module.add(labelMB) kernel["_GlobalAccumulation"] = "MultipleBuffer" if hasMultipleGlobalWriteModes: self.states.useBias = DataDirection.NONE kernel["ActivationFuncCall"] = False activationLabelList = {} betaModules, currentInstLength = self.generateBetaModules( kernel, tPA, tPB, activation, applyAlpha, betas, edge, atomic, vectorWidths, elements, activationLabelList, tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, writeLabels, endLabel, vectorDataTypes, factorDims, globalWriteMode, hasMultipleGlobalWriteModes) # Check if branch exceeds and add beta zero check betaBranchCheckModule = self.checkBetaBranchExceeds(kernel, betas, betaModules, writeLabels, globalWriteMode) # Append the beta modules to the main module module.appendModule(betaBranchCheckModule) module.appendModule(betaModules) if kernel["AdaptiveGemmGSUA"] == 1: kernel["ActivationFuncCall"] = afcBackup kernel["_GlobalAccumulation"] =gsuAccumBackup self.states.useBias = useBiasBackup activationLabelList = activationLabelListBackup # Generate activation modules if activationLabelList: activationModules = self.generateActivationModules( kernel, activation, activationLabelList, activationEnumStrList, activationSetPCStruct, tmpVgpr, actPCGwvwVgpr, actTempSgpr) # Defer activation blocks to end of kernel when other blocks are deferred # (called via s_setpc/s_swappc, position-independent). if kernel.get("UseSubtileImpl"): self.states.deferredActivationModules = activationModules else: module.appendModule(activationModules) self.sgprPool.checkIn(activationSetPCStruct.sgprOffsetActivation) self.sgprPool.checkIn(activationSetPCStruct.sgprOffsetBack) module.add(skComponent.writePartials(self, kernel, skPartialsLabel, vectorWidths_1, elements_1, tmpVgpr.idx, cvtVgprStruct, endLabel)) # End label module.add(endLabel) if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and kernel["StreamK"] > 0 and \ ((kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["MacDataTypeA"].numRegisters()) or \ (kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["MacDataTypeB"].numRegisters())): assert(kernel["ProblemType"]["ComputeDataType"].isSingle()) module.add(SMovB32(dst=sgpr("Alpha"), src=sgpr(oldAlpha), comment="Restore alpha value")) self.sgprPool.checkIn(oldAlpha) if self.states.FactorDim == 3 or hasMultipleGlobalWriteModes: self.updateBranchPlaceHolder(module, ["end_placeholder"], [endLabel.label], ["SBranch"]) self.vgprPool.checkIn(tmpVgpr.idx) if cvtVgpr is not None: self.vgprPool.checkIn(cvtVgpr) if gsuLimit > 1 and gsuLimitIdx == 0: if deferGSU0: # GSU0 store code is done. Append it to deferredGSU0 (placed after persistent loop), # then restore `module` to savedModule (the inline stub region) so subsequent code # (e.g. the SLongBranchPositive to KernelEnd) lands inline, not in the deferred block. # The deferred block falls through to GW_End -> KernelEnd -> s_endpgm directly, # so no explicit return branch back to inline is needed. deferredGSU0.appendModule(module) module = savedModule with self.allocTmpSgpr(3, tag="globalWriteElements_tmpSgprInfo") as tmpSgprInfo: module.add(SLongBranchPositive(Label("KernelEnd", ""), tmpSgprInfo, comment="GSU0 done, skip to end")) else: with self.allocTmpSgpr(3, tag="globalWriteElements_tmpSgprInfo2") as tmpSgprInfo: module.add(SLongBranchPositive(Label("KernelEnd", ""), tmpSgprInfo)) kernel["GlobalSplitU"] = gsuBackup kernel["_GlobalAccumulation"] = gsuAccumBackup self.states.bpeCexternal = bpeCexternalBackup return module, deferredGSU0 def getMBSKGSUTotal(self, kernel): if kernel["MbskPrefetchMethod"]: return kernel["NumMbskPrefetchElements"] + 1 GSUtotal = 16 if (kernel["MIWaveTile"][0] * kernel["MIWaveTile"][1]) * (kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1]) > 8: GSUtotal = int(GSUtotal/int((kernel["MIWaveTile"][0] * kernel["MIWaveTile"][1]) * (kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1])/8)) GSUtotal = max(2,GSUtotal) return GSUtotal def setOccupancy(self, kernel): # Use VGPR up to next occupancy threshold: # Account for additional temp sgprs that will be required for code gen, up to physical limits. +5 approximates upper end of required temp space for GSUSynccodegenOpt requiredSgprs = min(self.sgprPool.size() + 5, self.states.regCaps["MaxSgpr"]) maxVgprs, occupancy = self.getMaxRegsForOccupancy(kernel["NumThreads"], self.vgprPool.size(), requiredSgprs, \ 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, element, actPCMaxTempSgpr, \ gwvw, maxVgprs, ss): # Get estimated numVgprAvailable # Use block-aware counting to avoid overestimating with fragmented pools # print("Max vgprs =", maxVgprs, self.vgprPool.size(), self.vgprPool.availableBlock(ss.numVgprsPerElement, ss.align)) numVgprAvailable = self.vgprPool.availableBlockMaxVgpr(maxVgprs, ss.numVgprsPerElement, ss.align) # Grow the register pool if needed - we need enough regs for at least one element # Unfortunate since this means the write logic is setting the VGPR requirement # for the entire kernel but at least we have a functional kernel. # Before growing the pool, see if we can shrink the write vector width instead? # TODO : the vgprSerial is needed for-ever and if we grow here will split the # range of the tmps. Maybe want to move vgprSerial to first vgpr? # TODO: Minimum elems for StoreRemap # TODO: Which of DataType or DestDataType is in a better sense? 0114: Check Using DestDataType + HSS minElements = 2 if (kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16()) else 1 if numVgprAvailable < minElements*ss.numVgprsPerElement: self.vgprPool.resetOccupancyLimit() print2("info: growing pool += %d * %d for GlobalWrite\n" \ % (minElements,ss.numVgprsPerElement)) self.vgprPool.growPool(0, minElements, ss.numVgprsPerElement, \ "grow-pool for GlobalWrite") maxVgprsN, occupancy = self.setOccupancy(kernel) if maxVgprs != maxVgprsN: #print("refineOccupancy maxVgprs, new", maxVgprsN, "old", maxVgprs) return self.refineOccupancy(kernel, atomic, element, actPCMaxTempSgpr, gwvw, maxVgprsN, ss) # After growing, try block-aware count again; fall back to simple count # if the pool is too fragmented for aligned blocks numVgprAvailable = self.vgprPool.availableBlockMaxVgpr(maxVgprs, ss.numVgprsPerElement, ss.align) if numVgprAvailable == 0: numVgprAvailable = self.vgprPool.available() # print("NumVgprAvailable", numVgprAvailable) if ss.numVgprsPerElement: numElementsPerBatch = numVgprAvailable // ss.numVgprsPerElement else: # numVgprsPerElement==0: accvgprs are pre-staged (e.g. UseSubtileImpl) so no pool # vgprs are needed per element. Default to the full element list; use # NumElementsPerBatchStore to cap the batch size (e.g. for very large macro-tiles # where a smaller batch reduces register pressure or improves store pipelining). numElementsPerBatch = len(element) # Cap batch size to align on MIWaveTile[0] (M-tile) boundaries. # The acc-to-VGPR mapping interleaves M and N tiles, so a batch that # partially covers an N-column still touches the full acc range of that # column. Aligning to MIWaveTile[0] ensures batches break on N-column # boundaries, avoiding accesses beyond the ValuC range. if kernel.get("UseSubtileImpl") and kernel.get("EnableMatrixInstruction"): miwt0 = kernel["MIWaveTile"][0] totalElems = kernel["MIWaveTile"][0] * kernel["MIWaveTile"][1] if numElementsPerBatch >= totalElems: numElementsPerBatch = totalElems elif miwt0 > 1 and numElementsPerBatch >= miwt0: numElementsPerBatch = (numElementsPerBatch // miwt0) * miwt0 assert(self.states.c.numVgprValu % gwvw == 0) # sanity check numElementsPerBatch = numElementsPerBatch if not kernel["NumElementsPerBatchStore"] else min(kernel["NumElementsPerBatchStore"],numElementsPerBatch) # print("NumElementsPerBatch=", numElementsPerBatch, "LimitedBySgprs=", ss.cfg.numElementsPerBatchLimitedBySgprs, \ # "WARNING" if ss.cfg.numElementsPerBatchLimitedBySgprs < numElementsPerBatch else "okay") if ss.cfg.numElementsPerBatchLimitedBySgprs < numElementsPerBatch: numElementsPerBatch = ss.cfg.numElementsPerBatchLimitedBySgprs # TODO: Which of DataType or DestDataType is in a better sense? 0114: Check Using DestDataType + HSS destType = kernel["ProblemType"]["DestDataType"] srcType = kernel["ProblemType"]["DataType"] subtileImplDest16b = kernel.get("UseSubtileImpl") and (destType.isHalf() or destType.isBFloat16()) if (srcType.isHalf() or srcType.isBFloat16() or subtileImplDest16b): # only do an even number of halves - since these share hi/lo pieces of some registers? if numElementsPerBatch > 1: numElementsPerBatch = int(numElementsPerBatch/2)*2 # UseSubtileImpl paired-store: batch must be aligned to MIWaveTile[0] # (the number of M-tiles per N-column) so that batch boundaries don't # split sba=0/sba=1 pairs within an N-column. if kernel.get("UseSubtileImpl") and kernel["MIWaveTile"][0] > 1: miwt0 = kernel["MIWaveTile"][0] if numElementsPerBatch >= miwt0: numElementsPerBatch = (numElementsPerBatch // miwt0) * miwt0 # dot2: no this constraint elif not kernel["EnableMatrixInstruction"] and not kernel["UseDotInstruction"]: # The globalWriteBatch routine below can't handle odd elements per batch # and 0 elements per batch is illegal. # so if we don't have *GPR resources to handle a larger batch then need # to mark overflowedResources rather than generate a kernel that won't work. # It might be possible to fix globalWriteBatch to handle this case but these # are likely to be low-performing so likely not worth optimizing. print2("WARNING: half requires at least two elements per batch") self.states.overflowedResources = 3 assert numElementsPerBatch > 0, "numElementsPerBatch=0 for %s"%self.states.kernelName #numElementsPerBatch=min(2,numElementsPerBatch) # hack to control number of batches if atomic and (ss.optSingleColVgpr or ss.optSharedColVgpr): # hack to avoid re-using address vgpr across rows # atomics need to perform several memory operations # if the batch spans multiple rows, need multiple address vgpr # which is not currently supported in the two opt*ColVgpr modes firstRow = [e for e in element if e[0]==0 and e[2]==0] numElementsPerBatch=min(len(firstRow),numElementsPerBatch) # 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 element if e[0]==0 and e[2]==0] # format for element = (tt1, tt0, vc1, vc0) # find the largest factor and smaller than numElementPerBatch nBatchesPerRow = 1 for d in range(1, len(firstRow)+1): largestFactor = len(firstRow)//d if len(firstRow)%d == 0 and largestFactor <= numElementsPerBatch: numElementsPerBatch = largestFactor nBatchesPerRow = d break # if no atomics and no edge, then write whole vectors #if not atomic and not edge: # numVectorsPerBatch = numElementsPerBatch / kernel["GlobalWriteVectorWidth"] # #print " NumVectorsPerBatch", numVectorsPerBatch # numElementsPerBatch = numVectorsPerBatch * kernel["GlobalWriteVectorWidth"] numBatches = max(1, ceilDivide(len(element),numElementsPerBatch)) # Grow pool if needed # Get true numVgprAvailable numVgprAvailable = self.vgprPool.availableBlock(ss.numVgprsPerElement, ss.align) 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, element, actPCMaxTempSgpr, 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, tag="refineOccupancy_numSgprs", preventOverflow=False)) maxVgprsN, occupancy = self.setOccupancy(kernel) if maxVgprs != maxVgprsN: #print("refineOccupancy maxVgprs, new", maxVgprsN, "old", maxVgprs) return self.refineOccupancy(kernel, atomic, element, actPCMaxTempSgpr, gwvw, maxVgprsN, ss) return numElementsPerBatch, nBatchesPerRow, numBatches, numSgprs ############################################################################## # globalWriteElementBatch : ############################################################################## def globalWriteElementBatch(self, kernel, tPA, tPB, activation, \ applyAlpha, beta, edge, atomic, \ vectorWidth, element, activationLabelList, \ tmpVgpr, cvtVgprStruct, activationSetPCStruct, activationEnumStrList, \ actPCMaxTempSgpr, isInsertActFunctionCallAddrCalc, toActModuleList, \ edgeModule, writeLabel, endLabel, \ currentInstLength, \ betaIdx, fdIdx, vectorDataTypes, factorDims, hasMultipleGlobalWriteModes=False): factorDim = factorDims[fdIdx] edgeModule.add(writeLabel) # CompactLoopStore: allocate two dedicated, named sgprs for the CLS loop so # it no longer hijacks WorkGroup2 / ArgType (which are still live in some # store paths -- e.g. batched-GEMM batch offsets read WorkGroup2, and # ArgType is read during store-arg load). Checked out once here so the same # registers persist across all batches (and activation branches) of this # store, and freed at the end of the function. # CLSm0Base : M0 base accumulator (src VGPR index offset) # CLSLoopCounter : CLS loop iteration countdown if kernel["CompactLoopStore"]: edgeModule.add(self.defineSgpr("CLSm0Base", 1)) edgeModule.add(self.defineSgpr("CLSLoopCounter", 1)) # for storeRemap edge case, non-beta still can enable vector stores gwvw = vectorWidth #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, tag="globalWriteElements_sgprList1", preventOverflow=False)) sgprList.append(self.sgprPool.checkOut(1, tag="globalWriteElements_sgprList2", preventOverflow=False)) sgprList.append(self.sgprPool.checkOut(1, tag="globalWriteElements_sgprList3", preventOverflow=False)) sgprList.append(self.sgprPool.checkOutAligned(2, 2, tag="globalWriteElements_sgprList4", preventOverflow=False)) sgprList.append(self.sgprPool.checkOutAligned(2, 2, tag="globalWriteElements_sgprList5", preventOverflow=False)) sgprList.append(self.sgprPool.checkOutAligned(4, 4, tag="globalWriteElements_sgprList6", 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, tag="globalWriteElements_actPCMaxTempSgpr", 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, tag="globalWriteElements_tmpVgprDynamic"), size=tmpVgprDynamicSize) maxVgprs, occupancy = self.setOccupancy(kernel) ss = StoreState(self, kernel, gwvw, edge, beta, atomic, element, vectorDataTypes, dim=factorDim) actPCMaxTempSgpr_ = None if activationLabelList and isInsertActFunctionCallAddrCalc: assert activationSetPCStruct, activationEnumStrList and activationLabelList and toActModuleList actPCMaxTempSgpr_ = actPCMaxTempSgpr numElementsPerBatch, nBatchesPerRow, numBatches, numSgprs = self.refineOccupancy(kernel, atomic, element, actPCMaxTempSgpr_, 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, \ betaIdx)) if self.db["PrintStoreRegisterDb"]: print("edge", edge, "NumBatches", numBatches, "NumElementsPerBatch", numElementsPerBatch, "numVgprsPerElement", ss.numVgprsPerElement, "len(element)", len(element)) print ("numSgprs=", numSgprs, "sgprPool.size()=", self.sgprPool.size(), "numTempSgprPerBatch=", ss.cfg.numTempSgprPerBatch, "numMaskSgprPerBatch=", ss.cfg.numMaskSgprPerBatch, "numMaskSgprPerElement=", ss.cfg.numMaskSgprPerElement) print(self.sgprPool.state()) edgeModule.addComment1("edge=%d, allocate %u sgpr. perBatchTmpS=%u perBatchMaskS=%u perElementMaskS=%u elementsPerBatch=%u" % (edge, numSgprs, ss.cfg.numTempSgprPerBatch, ss.cfg.numMaskSgprPerBatch, ss.cfg.numMaskSgprPerElement, numElementsPerBatch)) #edgeModule.addComment("storeStats, %d, %d, %d"% (edge, numSgprs, numElementsPerBatch)) # so if we don't have *GPR resources to handle a larger batch then need # to mark overflowedResources rather than generate a kernel that won't work. # UseSubtileImpl NonEdge guard: compute numValidMBlocks / numValidNBlocks so # stores can skip OOB wave groups. Active for any NonEdge UseSubtileImpl path # that is not multi-buffer GSU accumulation. isSubtileNonEdge = ( not edge and kernel.get("UseSubtileImpl") and kernel["_GlobalAccumulation"] not in ("MultipleBufferSingleKernel", "MultipleBuffer") ) if isSubtileNonEdge: self._emitSubtileGuards(kernel, edgeModule) else: # Don't clear permanent guard SGPRs — they persist across batches if "SubtileMGuard" not in self.sgprs: self.states.subtileM32ValidBlocksSgpr = None self.states.subtileN16ValidBlocksSgpr = None self.states.subtileTotalMOffsetSgpr = None self.states.subtileMBlockSize = 0 # Activation actLoopEndLabel, actLoopLabelModules, actLoopEnumStrList = self.initActivationLoop(kernel, beta) actLoopModuleList = [] actLoopModuleCodeLength = [] with self.allocTmpSgpr(numSgprs, 2, tag="globalWriteElementBatch_tmpSgprRes") as tmpSgprRes: for index, activationLabelModule in enumerate(actLoopLabelModules): actLoopModule = Module("Activation Loop %s"%index) activationTypeStr = actLoopEnumStrList[index] if activationLabelModule: actLoopModule.add(activationLabelModule) tmpSgpr = tmpSgprRes.idx actTempSgpr = tmpSgpr # Get sgpr start address, should always be the same elementSgprs = tmpSgpr + ss.cfg.numTempSgprPerBatch codeAccVgprRead = deepcopy(self.codes.accVgprRead) if self.states.serializedStore else None mulAlpha = self.codes.mulAlphaMultipleBuffer if (kernel["_GlobalAccumulation"] == 'MultipleBuffer' or kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel') else self.codes.mulAlphaOther codeMulAlpha = deepcopy(mulAlpha) if self.states.serializedStore else None self.alphaBeforeLoadC = False if kernel["MIArchVgpr"] and applyAlpha and not kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel': # CompactLoopStore: emit alpha via Pattern 1 -- keep codeAccVgprRead so the # "copy MI out reg" stays a standalone per-element move, and force # codeMulAlpha=None to suppress the fused Pattern 3 (v_mul_f32 copy+alpha). # The CLS loop later steps that standalone copy across MI slices via M0, # which the fused form cannot do. Non-CLS keeps Pattern 3 verbatim. if kernel["CompactLoopStore"]: codeMulAlpha = None else: codeAccVgprRead = None #Only apply when 2 wave optimization features are enabled if (kernel["StorePriorityOpt"] or kernel["StoreSyncOpt"]) and beta: self.alphaBeforeLoadC = True #When LSU>1, don't use the VGPRs from the endSum. if (kernel["LocalSplitU"] > 1): codeMulAlpha = None else: codeMulAlpha = None biasLocalBarrierInit = False # If LSU, the VGPRs are from LSU reduction. # We need a variable to read from correct VGPR index when numBatches > 1. ss.lsuStartVgprOffset = 0 # CompactLoopStore: only emit ONE CLS body's worth of batches # (= batchesPerCLSBody, computed from SourceSwap / outerTT1 layout in # GlobalWriteBatchWriter.computeBatchesPerCLSBody). The CLS loop tail # emitted by GlobalWriteBatch.emit() at batchIdx == batchesPerCLSBody-1 # branches the loop back at runtime to re-execute the body CLS-iter-count # times via M0 indirection. Emitting additional batches beyond the CLS body # is unnecessary and would only duplicate store code that is intended to be # re-executed via the runtime loop. # For non-CLS, numBatchesCLS == numBatches so behaviour is unchanged. if kernel["CompactLoopStore"]: numBatchesCLS = GlobalWriteBatchWriter.computeBatchesPerCLSBody(kernel, numBatches) else: numBatchesCLS = numBatches # CompactLoopStore: pre-compute per-batch "next_rowInc" (rowInc from # this batch's last elt to the NEXT batch's first elt). Passed as # `inter_iter_rowInc` to globalWriteBatch so the per-elt look-ahead # scan in GlobalWriteBatchWriter falls through to this value when no # further intra-batch emit is found (= last emitting elt of the batch # needs the cross-batch advance). 0 for non-CLS / last batch / atomic. next_rowInc_per_batch = [0] * numBatches if kernel["CompactLoopStore"] and not atomic: # Use the single authoritative coordOffset1 formula in # getStoreElementsInfoForBatch (returns coord1 per element for the # whole list) instead of a duplicated helper. _, _coord1All = ss.getStoreElementsInfoForBatch(kernel, element) for _bIdx in range(numBatches): _elStopIdx = min((_bIdx + 1) * numElementsPerBatch, len(element)) if _elStopIdx < len(element): next_rowInc_per_batch[_bIdx] = _coord1All[_elStopIdx] - _coord1All[_elStopIdx - 1] for batchIdx in range(0, numBatchesCLS): elementStartIdx = batchIdx * numElementsPerBatch elementStopIdx = min( elementStartIdx + numElementsPerBatch, len(element) ) elementsThisBatch = element[elementStartIdx:elementStopIdx] #print("BATCH[%u/%u]: element[%u:%u] VGPRs=%u" % (batchIdx, numBatches, elementStartIdx, elementStopIdx,ss.numVgprsPerElement )) # elementVgprs can be large and should be perfectly tuned to the number of available # VGPRS. We do not want to accidentally overflow and grow the pool here: if kernel["StoreRemapVectorWidth"]: #Indication if this batch is last batch for this column block shape self.StoreRemapLastBatch = 1 if (batchIdx+1) % nBatchesPerRow == 0 else 0 # Look ahead through next_rowInc_per_batch for the first non-zero # transition: SRVW path skips batches whose `addrCalc.rowInc == 0`, # so the chain primer at the current batch primes for the NEXT # firing consumer (= first batch with a real row transition). # Cumulative skip is implicit since skipped transitions are 0 and # contribute nothing to the sum. _next_firing_rowInc = 0 for _k in range(batchIdx, numBatches): if next_rowInc_per_batch[_k] != 0: _next_firing_rowInc = next_rowInc_per_batch[_k] break _direct_next_rowInc = next_rowInc_per_batch[batchIdx] if batchIdx < numBatches else 0 actLoopModule.add(self.globalWriteBatch(kernel, tPA, tPB, activation, ss, batchIdx, \ applyAlpha, beta, edge, atomic, gwvw, atomicW, \ elementsThisBatch, self.vgprs.addrE, self.vgprs.addrD, self.vgprs.addrC, self.vgprs.addrBias, \ self.vgprs.addrScaleAVec, self.vgprs.addrScaleBVec, self.vgprs.addrScaleAlphaVec, \ biasLocalBarrierInit, tmpVgpr, tmpVgprDynamic, cvtVgprStruct, activationSetPCStruct, \ activationTypeStr, elementSgprs, tmpSgpr, codeAccVgprRead, codeMulAlpha, factorDim, numBatches, \ _next_firing_rowInc, _direct_next_rowInc)) biasLocalBarrierInit = True ss.resetState() actLoopModuleList.append(actLoopModule) actLoopModuleCodeLength.append(countInstruction(actLoopModule)) ################# # Free after final vgpr calculation # Only free locally-allocated guard SGPRs, not permanent ones (SubtileMGuard). if self.states.subtileTotalMOffsetSgpr is not None: self.sgprPool.checkIn(self.states.subtileTotalMOffsetSgpr) self.states.subtileTotalMOffsetSgpr = None if self.states.subtileM32ValidBlocksSgpr is not None and "SubtileMGuard" not in self.sgprs: self.sgprPool.checkIn(self.states.subtileM32ValidBlocksSgpr) self.sgprPool.checkIn(self.states.subtileN16ValidBlocksSgpr) self.states.subtileM32ValidBlocksSgpr = None self.states.subtileN16ValidBlocksSgpr = None self.states.subtileMBlockSize = 0 if tmpVgprDynamic: self.vgprPool.checkIn(tmpVgprDynamic.idx) if len(actLoopLabelModules) > 1: actInstCounter = 0 # Add activation branch for index, actLoopLabelModule in enumerate(actLoopLabelModules): enumIndex = ActivationType.getEnumIndex(actLoopEnumStrList[index]) #edgeModule.add(SCmpKEQU32(sgpr("ActivationType"), enumIndex, "activationType == %u"%enumIndex)) edgeModule.add(self.getSCMPKInstruction("EQU32", "ActivationType", enumIndex, comment="activationType == %u"%enumIndex)) if actInstCounter >= self.states.asmCaps["ShortBranchMaxLength"]: edgeModule.add(self.longBranchScc1(actLoopLabelModule, posNeg=1, comment="Branch if true")) else: edgeModule.add(SCBranchSCC1(actLoopLabelModule.getLabelName(), "Branch if true")) actInstCounter += actLoopModuleCodeLength[index] # Add jump to activation end for index, _ in enumerate(actLoopLabelModules): actLoopModule = actLoopModuleList[index] if (index < (len(actLoopLabelModules) - 1)): if actInstCounter >= self.states.asmCaps["ShortBranchMaxLength"]: with self.allocTmpSgpr(3, tag="globalWriteElementBatch_tmpSgprInfo") as tmpSgprInfo: actLoopModule.add(SLongBranchPositive(actLoopEndLabel, tmpSgprInfo)) else: actLoopModule.add(SBranch(labelName=actLoopEndLabel.getLabelName())) actInstCounter -= actLoopModuleCodeLength[index] # Append to edgeModule for actLoopModule in actLoopModuleList: edgeModule.appendModule(actLoopModule) # Add actLoopEndLabel if needed if len(actLoopLabelModules) > 1: edgeModule.add(actLoopEndLabel) # CompactLoopStore: release the dedicated CLS sgprs now that all batches / # activation branches of this store have been emitted. if kernel["CompactLoopStore"]: edgeModule.add(self.undefineSgpr("CLSLoopCounter")) edgeModule.add(self.undefineSgpr("CLSm0Base")) if len(factorDims) == 1: isDeferredReturn = "Deferred" in endLabel.getLabelName() if currentInstLength >= self.states.asmCaps["ShortBranchMaxLength"] or isDeferredReturn: posLabel = self.labels.getNameInc("DeferredReturnDir") with self.allocTmpSgpr(3, tag="globalWriteElementBatch_tmpSgprInfo2") as tmpSgprInfo: edgeModule.add(SLongBranch(endLabel, tmpSgprInfo, posLabel, comment="jump to end")) else: edgeModule.add(SBranch(labelName=endLabel.getLabelName(), comment="jump to end")) else: end_placeholder = Module("end_placeholder") edgeModule.add(end_placeholder) currentInstLength += countInstruction(edgeModule) del ss return currentInstLength, activationTypeStr ############################################################################## # chooseGlobalRead : # create the load instruction for requested vector width and other parms # return an Inst class # # bpl = bytes per load op # tr = transpose read (for matrix instruction) ############################################################################## def chooseGlobalRead(self, useBuffer, bpl, destVgpr, \ addr0, addr1, soffset, offset, \ glc=False, slc=False, nt=False, lds=False, \ tr=False, hi16=0, comment="load C", \ scope=CacheScope.SCOPE_NONE, th=TemporalHint.TH_NONE, nv=NonVolatile.NV_NONE): # rpv = regs per vector rpv = bpl/4.0 if useBuffer: rv = Module("Global Read") mubuf = MUBUFModifiers(offen=True, offset12=int(offset), glc=glc, slc=slc, nt=nt, lds=lds, scope=scope, th=th, nv=nv) # Nested buffer load implementation function for easy branching for soffset def bufferLoadImpl(soffset): nonlocal rv factor = max(1, 4//bpl) dst = None if lds else vgpr(destVgpr, rpv*factor) if bpl < 4: if self.do["EmulatedECCBufferLoad"]: rv.add(VMovB32(dst, 0)) if bpl==1 and hi16: rv.add(BufferLoadD16HIU8(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==1 and not hi16: rv.add(BufferLoadD16U8(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==2 and hi16: rv.add(BufferLoadD16HIB16(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==2 and not hi16: rv.add(BufferLoadD16B16(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==4: rv.add(BufferLoadB32(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==8: rv.add(BufferLoadB64(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==12: # Cannot easily adopt to this due to alignment requirement rv.add(BufferLoadB96(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) # split into dwordx2 and dword loads. Second load offset is 8 bytes. # rv = Module("emulated _buffer_load_b192") # dst = None if lds else vgpr(destVgpr, 2) # rv.add(BufferLoadB64(dst=dst, vaddr=addr0, saddr=addr1, \ # soffset=soffset, mubuf=mubuf, comment=comment)) # mubuf2 = MUBUFModifiers(offen=True, offset12=offset+8, glc=glc, slc=slc, nt=nt, lds=lds) # if isinstance(destVgpr, str): # dst2 = f"{destVgpr}+{2}" # elif isinstance(destVgpr, int): # dst2 = destVgpr + 2 # dst = None if lds else vgpr(dst2, 1) # rv.add(BufferLoadB32(dst=dst, vaddr=addr0, saddr=addr1, \ # soffset=soffset, mubuf=mubuf2, comment=comment)) return rv elif bpl==16: rv.add(BufferLoadB128(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) return rv elif bpl==24: # split into dwordx4 and dwordx2 loads. Second load offset is 16 bytes. rv = Module("emulated _buffer_load_b192") dst = None if lds else vgpr(destVgpr, 4) rv.add(BufferLoadB128(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) mubuf2 = MUBUFModifiers(offen=True, offset12=offset+16, glc=glc, slc=slc, nt=nt, lds=lds, scope=scope, th=th, nv=nv) dst = None if lds else vgpr(_vgprOffset(destVgpr, 4), 2) rv.add(BufferLoadB64(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf2, comment=comment)) return rv elif bpl==32: # split into two dwordx4 loads. Second load offset is +0.5 bpl rv = Module("emulated _buffer_load_b256") dst = None if lds else vgpr(destVgpr, rpv//2) rv.add(BufferLoadB128(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf, comment=comment)) mubuf2 = MUBUFModifiers(offen=True, offset12=int(offset + bpl/2), glc=glc, slc=slc, nt=nt, lds=lds, scope=scope, th=th, nv=nv) dst = None if lds else vgpr(_vgprOffset(destVgpr, int(rpv//2)), rpv//2) rv.add(BufferLoadB128(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf2, comment=comment)) return rv elif bpl==64 and not lds: rv = Module("emulated _buffer_load_b512") quarter_rpv = int(rpv//4) for i in range(4): mubuf_n = MUBUFModifiers(offen=True, offset12=int(offset + i*bpl/4), glc=glc, slc=slc, nt=nt, lds=lds, scope=scope, th=th, nv=nv) dst = vgpr(_vgprOffset(destVgpr, i*quarter_rpv), quarter_rpv) rv.add(BufferLoadB128(dst=dst, vaddr=addr0, saddr=addr1, \ soffset=soffset, mubuf=mubuf_n, comment=comment)) else: assert 0, "%s\nchooseGlobalRead: bad bpl %u"%(self.states.kernelName,bpl) # buffer_load offset field is 12-bit. # if offset >= 4096, use soffset instead if offset >= 4096: if soffset in (0, "0"): mubuf = MUBUFModifiers(offen=True, offset12=0, glc=glc, slc=slc, nt=nt, lds=lds, scope=scope, th=th, nv=nv) with self.allocTmpSgpr(1, tag="chooseGlobalRead_tmpSgprInfo") as tmpSgprInfo: soffset = sgpr(tmpSgprInfo.idx) rv.add(SMovB32(dst=soffset, src=offset, comment="large offset")) bufferLoadImpl(soffset) else: assert 0, "offset too large and soffset set" else: bufferLoadImpl(soffset) return rv elif tr: # global_load offset field is 24-bit. if offset >= 1<<24: raise RuntimeError("chooseGlobalRead: offset too large for global_load %u"%(offset)) modifier = GLOBALModifiers(offset=offset) if bpl==8: return GlobalLoadTR8B64(dst=vgpr(destVgpr, rpv), vaddr=addr0, saddr=addr1, modifier=modifier, comment=comment) elif bpl==16: return GlobalLoadTR16B128(dst=vgpr(destVgpr, rpv), vaddr=addr0, saddr=addr1, modifier=modifier, comment=comment) else: modifier = GLOBALModifiers(offset=int(offset), glc=glc, slc=slc, dlc=False, scope=CacheScope.SCOPE_NONE, lds=lds, isStore=False) saddr_off = vgpr("off", 1, False, False, True) if bpl==1 and hi16: return GlobalLoadD16HIU8(dst=vgpr(destVgpr, rpv*4), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==1 and not hi16: return GlobalLoadD16U8(dst=vgpr(destVgpr, rpv*4), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==2 and hi16: return GlobalLoadD16HIB16(dst=vgpr(destVgpr, rpv*2), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==2 and not hi16: return GlobalLoadD16B16(dst=vgpr(destVgpr, rpv*2), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==4: return GlobalLoadB32(dst=vgpr(destVgpr, rpv), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==8: return GlobalLoadB64(dst=vgpr(destVgpr, rpv), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==12: return GlobalLoadB96(dst=vgpr(destVgpr, rpv), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==16: return GlobalLoadB128(dst=vgpr(destVgpr, rpv), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment) elif bpl==24: # emulate global_load_b192 as global_load_b128 + global_load_b64 rv = Module("emulated _global_load_b192") rv.add(GlobalLoadB128(dst=vgpr(destVgpr, 4), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment)) modifier2 = GLOBALModifiers(offset=int(offset+16), glc=glc, slc=slc, dlc=False, scope=CacheScope.SCOPE_NONE, lds=lds, isStore=False) rv.add(GlobalLoadB64(dst=vgpr(_vgprOffset(destVgpr, 4), 2), vaddr=addr0, saddr=saddr_off, modifier=modifier2, comment=comment)) return rv elif bpl >= 32 and bpl % 16 == 0: # split into multiple dwordx4 loads using GLOBALModifiers offset rv = Module("emulated flat load %u bytes" % bpl) rounds = int(bpl // 16) shiftRpv = rpv // rounds rv.add(GlobalLoadB128(dst=vgpr(destVgpr, shiftRpv), vaddr=addr0, saddr=saddr_off, modifier=modifier, comment=comment)) for i in range(1, rounds): modifierN = GLOBALModifiers(offset=int(offset+16*i), glc=glc, slc=slc, dlc=False, scope=CacheScope.SCOPE_NONE, lds=lds, isStore=False) vgprOff = _vgprOffset(destVgpr, int(shiftRpv * i)) rv.add(GlobalLoadB128(dst=vgpr(vgprOff, shiftRpv), vaddr=addr0, saddr=saddr_off, modifier=modifierN, comment=comment)) return rv else: assert 0, "chooseGlobalRead: bad bpl" ############################################################################## def chooseGlobalWrite(self, useBuffer, bps, srcVgpr, rpv, \ addr0, addr1, offset, soffset=0, \ glc=False, slc=False, nt=False, dlc=False, \ hi16=0, comment="store", \ scope=CacheScope.SCOPE_NONE, th=TemporalHint.TH_NONE, nv=NonVolatile.NV_NONE): """ create the store instruction for requested vector width and other parms rpv = regs per vector """ module = Module("chooseGlobalWrite %s -> %s (%s)"%(srcVgpr, addr0, addr1)) def bufferStoreImpl(tmpSgpr, mubuf): if bps==1 and hi16: module.add(BufferStoreD16HIB16(src=vgpr(srcVgpr, rpv*4), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps==1 and not hi16: module.add(BufferStoreB8(src=vgpr(srcVgpr, rpv*4), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps==2 and hi16: module.add(BufferStoreD16HIB16(src=vgpr(srcVgpr, rpv*2), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps==2 and not hi16: module.add(BufferStoreB16(src=vgpr(srcVgpr, rpv*2), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps==4: module.add(BufferStoreB32(src=vgpr(srcVgpr, rpv), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps==8: module.add(BufferStoreB64(src=vgpr(srcVgpr, rpv), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps==16: module.add(BufferStoreB128(src=vgpr(srcVgpr, rpv), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) elif bps >= 32 and bps % 32 == 0: # split into several dwordx4 loads. Offset the next by +0.5 bps rounds = bps // 16 shiftByte = bps // rounds shiftRpv = rpv // rounds module.add(BufferStoreB128(src=vgpr(srcVgpr, shiftRpv), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf, comment=comment)) for i in range(1, rounds): offset2 = offset+shiftByte*i if offset2 >= 4096: module.add(SMovB32(dst=tmpSgpr, src=offset2, comment="large offset")) offset2 = 0 mubuf2 = MUBUFModifiers(offen=True, offset12=offset2, glc=glc, slc=slc, dlc=dlc, nt=nt, isStore=True, scope=scope, th=th, nv=nv) vgprOff = _vgprOffset(srcVgpr, int(shiftRpv * i)) module.add(BufferStoreB128(src=vgpr(vgprOff, shiftRpv), vaddr=addr0, \ saddr=addr1, soffset=tmpSgpr, mubuf=mubuf2, comment=comment)) else: assert 0, "bad bps" if useBuffer: mubuf = MUBUFModifiers(offen=True, offset12=offset, glc=glc, slc=slc, dlc=dlc, nt=nt, isStore=True, scope=scope, th=th, nv=nv) if soffset != 0: assert offset < 4096, "sgpr offset provided with large const offset" # buffer_load offset field is 12-bit. # if offset >= 4096, use soffset instead maxShift = max(bps - 16, 0) #if bps = 32 or bps = 64 if (offset + maxShift) >= 4096: with self.allocTmpSgpr(1, tag="chooseGlobalWrite_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = sgpr(tmpSgprInfo.idx) if offset >= 4096: module.add(SMovB32(dst=tmpSgpr, src=offset, comment="large offset")) mubuf = MUBUFModifiers(offen=True, offset12=0, glc=glc, slc=slc, dlc=dlc, nt=nt, isStore=True, scope=scope, th=th, nv=nv) bufferStoreImpl(tmpSgpr, mubuf) else: bufferStoreImpl(soffset, mubuf) else: modifier = GLOBALModifiers(offset=0, glc=glc, slc=slc, dlc=False, scope=CacheScope.SCOPE_NONE, lds=False, isStore=True) saddr_off = vgpr("off", 1, False, False, True) if bps==1 and hi16: module.add(GlobalStoreD16HIB16(vaddr=addr0, src=vgpr(srcVgpr, rpv*4), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps==1 and not hi16: module.add(GlobalStoreB8(vaddr=addr0, src=vgpr(srcVgpr, rpv*4), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps==2 and hi16: module.add(GlobalStoreD16HIB16(vaddr=addr0, src=vgpr(srcVgpr*2), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps==2 and not hi16: module.add(GlobalStoreB16(vaddr=addr0, src=vgpr(srcVgpr, rpv*2), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps==4: module.add(GlobalStoreB32(vaddr=addr0, src=vgpr(srcVgpr, rpv), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps==8: module.add(GlobalStoreB64(vaddr=addr0, src=vgpr(srcVgpr, rpv), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps==16: module.add(GlobalStoreB128(vaddr=addr0, src=vgpr(srcVgpr, rpv), saddr=saddr_off, modifier=modifier, comment=comment)) elif bps >= 32 and bps % 16 == 0: # split into multiple dwordx4 stores using GLOBALModifiers offset rounds = int(bps // 16) shiftByte = 16 shiftRpv = rpv // rounds module.add(GlobalStoreB128(vaddr=addr0, src=vgpr(srcVgpr, shiftRpv), saddr=saddr_off, modifier=modifier, comment=comment)) for i in range(1, rounds): modifierN = GLOBALModifiers(offset=shiftByte*i, glc=glc, slc=slc, dlc=False, scope=CacheScope.SCOPE_NONE, lds=False, isStore=True) vgprOff = _vgprOffset(srcVgpr, int(shiftRpv * i)) module.add(GlobalStoreB128(vaddr=addr0, src=vgpr(vgprOff, shiftRpv), saddr=saddr_off, modifier=modifierN, comment=comment)) else: assert 0, "bad bps" return module def addVecGlobalLoad(self, dataType, kernel, vecVgpr, addr0, addr1, offset, gwvw, comment=""): """ Add vec for the element with addrCalc, elementIdx, and vecVgpr. vecVgpr is one or more vgpr :temp vGPR ( = gwvw * numbytes // 4 + 1 if cvt is needed) """ # Add vec here module = Module(comment) bps = dataType.numBytes() * gwvw useBuffer = kernel["BufferLoad"] if not (dataType.isHalf() or dataType.isBFloat16() or dataType.isInt32() or dataType.isSingle() or dataType.isDouble() or dataType.isSingleComplex()): printExit("Unsupported %s type %s."%(comment, str(dataType))) module.add(self.chooseGlobalRead(useBuffer, bps, vecVgpr, \ addr0, addr1, soffset=0, offset=offset, hi16=0, comment=comment)) return module ############################################################################## def addScaleVecLoad(self, kernel, ss, name: str, srdName: str, addrScaleVecVgpr, scaleVecVgpr, gwvw, scaleVecOffset, addVecPostfix = True): """ Add scaleAlphaVec for the element with addrCalc, elementIdx, and scaleVecVgpr. scaleVecVgpr is one or more vgpr :temp vGPR ( = gwvw * numbytes // 4 + 1 if cvt is needed) """ module = Module("addScale%sVec"%srdName) if kernel["ProblemType"]["UseScale%s"%name] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["GlobalSplitUAlgorithm"] == 'MultipleBufferSingleKernel' or kernel["StreamK"] > 0): bps = kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvw if kernel["BufferLoad"]: addr0 = vgpr(addrScaleVecVgpr) addr1 = sgpr("SrdScale%sVec"%srdName, 4) if addVecPostfix else sgpr("SrdScale%s"%srdName, 4) else: addr0 = vgpr(addrScaleVecVgpr,2) addr1 = "" useBuffer = kernel["BufferLoad"] computeDataType = kernel["ProblemType"]["ComputeDataType"] if not (computeDataType.isHalf() or computeDataType.isBFloat16() or computeDataType.isInt32() or computeDataType.isSingle() or computeDataType.isDouble() or computeDataType.isSingleComplex()): printExit("Unsupported scale%sVec type %s."%(srdName, str(computeDataType))) module.add(self.chooseGlobalRead(useBuffer, bps, scaleVecVgpr, \ addr0, addr1, soffset=0, offset=scaleVecOffset, hi16=0, comment="load scale%sVec"%srdName)) return module def addLdsLoad(self, dataType, dstVgpr, srcAddrVgpr, dsOffset, gwvw, comment=""): module = Module(comment) src = vgpr(srcAddrVgpr) ds = DSModifiers(offset=dsOffset) bpl = dataType.numBytes() * gwvw memToken = MemTokenData([self.states.memTokenLdsBuffer0]) if bpl <= 16: numRegs = max(1, bpl // 4) dst = vgpr(dstVgpr, numRegs) if numRegs > 1 else vgpr(dstVgpr) module.add(dsLoad(bpl, dst=dst, src=src, ds=ds, comment=comment, memToken=memToken)) elif bpl==32: inst = DSLoadB128(dst=vgpr(dstVgpr, 4), src=src, ds=ds, comment=comment) inst.setMemToken(memToken) module.add(inst) ds = DSModifiers(offset=int(dsOffset+bpl/2)) inst = DSLoadB128(dst=vgpr(dstVgpr+4, 4), src=src, ds=ds, comment=comment) inst.setMemToken(memToken) module.add(inst) elif bpl==64: for chunk in range(4): inst = DSLoadB128(dst=vgpr(dstVgpr + chunk*4, 4), src=src, ds=DSModifiers(offset=int(dsOffset + chunk*16)), comment=comment) inst.setMemToken(memToken) module.add(inst) else: assert 0, "bad bpl" return module def addBiasLoad(self, dataType, kernel, gwvw, addrCalc, biasVgpr, factorDim, isLocal=False): if isLocal and (self.states.useBias == DataDirection.READ): return self.addLdsLoad(dataType, biasVgpr, addrCalc.addrBiasVgpr, addrCalc.biasOffset[factorDim], gwvw, comment="Load Bias") if self.states.useBias == DataDirection.READ: if kernel["BufferLoad"]: addr0 = vgpr(addrCalc.addrBiasVgpr) addr1 = sgpr("SrdBias", 4) else: addr0 = vgpr(addrCalc.addrBiasVgpr,2) addr1 = "" else: return Module("Empty load") return self.addVecGlobalLoad(dataType, kernel, biasVgpr, addr0, addr1, addrCalc.biasOffset[factorDim], gwvw, comment="Load Bias") ############################################################################## def addStore(self, kernel, ss, tc: str, addrCalc, sumIdx, tmpS01, edge, elementIdx=None, batchIdx=None, wsOffset=0, overrideAfterPrimerRows=0, comment="addStore"): """ Add stores for the element with addrCalc and sumIdx. tmpS01 is a single :temp sGPR CompactLoopStore: `elementIdx` is used so that elt-0 (rowInc=0) also emits incrementToNextRow as a chain seed -- a no-op s_add + s_lshl primer that seeds s[stmp] for elt-1's real advance. Default 1 keeps the legacy rowInc != 0 gate for non-CLS callers. `overrideAfterPrimerRows` is forwarded to incrementToNextRow's CLS delayed AFTER-primer so the primer encodes the NEXT EMITTING elt's rowInc (look-ahead). 0 means no override. """ module = Module("addStore sumIdx %s"%(str(sumIdx))) if self.do["GlobalWrite"]: # perform vector stores here, so no VI indexing. # if GWVW > Vw, might need to support loops to # implement wider stores isGlc = False isSlc = False isNT = False scope = CacheScope.SCOPE_NONE th = TemporalHint.TH_NONE nv = NonVolatile.NV_NONE isDlc = False if tc == 'D': isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal( self.states.asmCaps, kernel["NonTemporalD"], _temporalHint(kernel, "D"), _nonVolatile(kernel, "D")) bps = self.states.bpeCexternal * ss.cfg.gwvw rpv = self.states.bpeCexternal * ss.cfg.gwvw / self.states.bpr if kernel["BufferStore"]: addr0 = vgpr(addrCalc.addrDVgpr) addr1 = sgpr("SrdD", 4) else: addr0 = vgpr(addrCalc.addrDVgpr,2) addr1 = "" # CompactLoopStore: drop rowInc gate so elt-0 (rowInc=0) also emits a # chain seed -- no-op s_add (s[stmp]=0 from preamble) + s_lshl that # primes s[stmp]=StrideD< Module: packdata = Component.PackData.find(self) gwriter = Component.GlobalWriteComponents.find(self) return gwriter(kernel, tPA, tPB, activation, ss, \ batchIdx, applyAlpha, beta, edge, atomic, gwvw, atomicW, \ batchElements, addrE, addrD, addrC, addrBias, \ addrScaleAVec, addrScaleBVec, addrScaleAlphaVec, biasLocalBarrierInit, \ tmpVgpr, tmpVgprDynamic, cvtVgprStruct, activationSetPCStruct, activationTypeStr, \ batchElementSgprs, tmpSgpr, codeAccVgprRead, codeMulAlpha, packdata, self, factorDim, \ self.assembler.version, numBatches, inter_iter_rowInc, direct_next_rowInc) ############################################################################## def openPrefetchGlobalRead2orMore(self, kernel, idxPgr): imod = Module() # early exit check # jump to idxPgr=1 label (wait 0) loopCounter = self.loopCounter(kernel, self.states.unrollIdx) imod.add(SCmpEQU32(src0=loopCounter, src1=hex(idxPgr), comment="PGR=%d but only %d loop"%(kernel["PrefetchGlobalRead"], idxPgr))) skipPGRn = Label(self.labels.getName("skipPGR%d_%d"%(kernel["PrefetchGlobalRead"], 1)), "") imod.add(SCBranchSCC1(labelName=skipPGRn.getLabelName(), comment="PGR=%d but only %d loop"%(kernel["PrefetchGlobalRead"], idxPgr))) return imod def closePrefetchGlobalRead2orMore(self, kernel, tensorParametersA, tensorParametersB, idxPgr): imod = Module() PGR = kernel["PrefetchGlobalRead"] if idxPgr == 0: # first, generate branch to the last idxPgr (= PGR) skipPGRn = Label(self.labels.getName("skipPGR%d_%d"%(PGR, PGR)), "") imod.add(SBranch(labelName=skipPGRn.getLabelName(), comment="jump to PGR=%d label"%(PGR))) elif idxPgr == 1 or idxPgr == kernel["PrefetchGlobalRead"]: skipPGRn = Label(self.labels.getName("skipPGR%d_%d"%(PGR, idxPgr)), "") imod.add(skipPGRn) if (kernel["DirectToLdsA"] and kernel["DirectToLdsB"]): # early exit case (idxPgr < PGR), need to wait for all prefetch here (vmcnt=0). # non early exit case, we wait only 1 set. Skip wait for (PGR-1) set numWait = idxPgr - 1 if idxPgr == kernel["PrefetchGlobalRead"] else 0 comment = "wait for global reads with lds" if idxPgr == 1: # add more comment for early exit comment += " (for early exit)" imod.add(self._wait(kernel, tensorParametersA, tensorParametersB, numWait, -1, -1, comment)) return imod ######################################## # Read vector to LDS ######################################## def calculateVectorGlobalOffset(self, kernel, offsetVgpr, tmpSgpr, dim): module = Module("") def calculateVectorGlobalOffsetCommon(s): module.add(SMulI32(dst=sgpr(s), src0=kernel["MacroTile%d"%dim], src1=sgpr("WorkGroup%d"%dim), comment="wgp%d * MT%d"%(dim, dim))) module.add(VAddU32(dst=vgpr(offsetVgpr), src0=sgpr(s), src1=vgpr("Serial"), comment=f"coord {dim} = wgp{dim} * MT{dim} + thread offset")) if tmpSgpr: calculateVectorGlobalOffsetCommon(tmpSgpr) else: with self.allocTmpSgpr(1, tag="calculateVectorGlobalOffset_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx calculateVectorGlobalOffsetCommon(tmpSgpr) return module def calculateVectorGlobalStride(self, kernel, offsetInVgpr, offsetOutVgpr, tmpSgpr, dim, strideName:str): module = Module("") def calculateVectorGlobalStrideCommon(s): module.add(SMulI32(dst=sgpr(s), src0=sgpr(strideName), src1=sgpr("WorkGroup2"), comment="Stride * WG")) module.add(VAddU32(dst=vgpr(offsetOutVgpr), src0=sgpr(s), src1=vgpr(offsetInVgpr), comment=f"coord {dim} = wgp{dim} * MT{dim} + thread offset + Stride * WG")) if tmpSgpr: calculateVectorGlobalStrideCommon(tmpSgpr) else: with self.allocTmpSgpr(1, tag="calculateVectorGlobalStride_tmpSgprInfo") as tmpSgprInfo: tmpSgpr = tmpSgprInfo.idx calculateVectorGlobalStrideCommon(tmpSgpr) return module def getGlobalShiftOffset(self, kernel, dataType, gwvw): numVgprs = int(ceil(dataType.numRegisters() * gwvw)) reg = dataType.numRegisters() if dataType.numRegisters() >= kernel["ProblemType"]["ComputeDataType"].numRegisters() \ else kernel["ProblemType"]["ComputeDataType"].numRegisters() return gwvw * reg - numVgprs def getEpilogueGlobalLoadTurn(self, kernel, gwvw, dim): """Macro-tile coverage turns for epilogue vector global loads and LDS staging.""" divisor = kernel["SubGroup0"] * kernel["SubGroup1"] turn = ceil(kernel["MacroTile%d"%dim] / (divisor * gwvw)) return turn, divisor def getEpilogueGlobalLoadStrideBpe(self, kernel, gwvw, dim, bpe=None): _, divisor = self.getEpilogueGlobalLoadTurn(kernel, gwvw, dim) if bpe is None: bpe = kernel["ProblemType"]["ComputeDataType"].numBytes() return (divisor * gwvw) * bpe def getTurn(self, kernel, gwvw, dim): """Epilogue vector turn count (global load + LDS staging). GW batch slot spacing uses coordOffset.""" return self.getEpilogueGlobalLoadTurn(kernel, gwvw, dim) def addVectorGlobalLoad(self, kernel, srdName: str, offsetVgpr, shiftOffset, dataType, bpe, gwvw, tmpVgpr1Res: ContinuousRegister, dstOffset, dim): module = Module("") tmpVgpr1 = tmpVgpr1Res.idx + dstOffset turn, divisor = self.getEpilogueGlobalLoadTurn(kernel, gwvw, dim) if kernel["BufferLoad"]: addr0 = vgpr(offsetVgpr) addr1 = sgpr("Srd%s"%srdName, 4) else: addr0 = vgpr(offsetVgpr, 2) addr1 = "" offset = self.getEpilogueGlobalLoadStrideBpe(kernel, gwvw, dim, bpe) for i in range(turn): if i != 0: module.add(VAddU32(dst=vgpr(offsetVgpr), src0=offset, src1=vgpr(offsetVgpr), comment="add subgroup offset")) module.add(self.addVecGlobalLoad(dataType, kernel, tmpVgpr1 + shiftOffset, addr0, addr1, 0, gwvw, comment="Load %s"%srdName)) # TODO: Will this work if gwvw > 1? tmpVgpr1 += 1 return module def addVectorLocalStore(self, kernel, addressStr: str, offsetVgpr, shiftOffset, dataType, gwvw, tmpVgpr1Res: ContinuousRegister, srcOffset, subGroupOffset, dim, setToOne=False, comment=""): module = Module("") tmpVgpr1 = tmpVgpr1Res.idx + srcOffset turn, divisor = self.getEpilogueGlobalLoadTurn(kernel, gwvw, dim) offset = self.getEpilogueGlobalLoadStrideBpe(kernel, gwvw, dim) if setToOne: if kernel["BufferLoad"]: # SRD+2 holds num records; compare to 0 to check if address is valid module.add(VCmpGtU32(dst=sgpr("Address%s"%addressStr, self.states.laneSGPRCount), src0=sgpr("Srd%s+2"%addressStr), src1=0, comment=" == 0 ?")) else: # For flat addressing, check if the base pointer is non-null module.add(VCmpNeU64(dst=sgpr("Address%s"%addressStr, self.states.laneSGPRCount), src0=sgpr("Address%s"%addressStr, 2), src1=0, comment="address != 0 ?")) # Set maskConst to 1.0 or 1 if kernel["ProblemType"]["ComputeDataType"].isSingle(): maskConst = 1.0 elif kernel["ProblemType"]["ComputeDataType"].isInt32(): maskConst = 1 turnOffset = 0 for i in range(turn): if i != 0: turnOffset += offset bps = kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvw ds = DSModifiers(offset=(subGroupOffset[0] + turnOffset)) dst = vgpr(offsetVgpr) for vi in range(gwvw): # Does not support hi/lo yet shiftOffset2 = shiftOffset + int(vi * dataType.numRegisters()) if kernel["ProblemType"]["ComputeDataType"].isSingle(): if dataType.isHalf(): module.add(ECvtF16toF32(dst=vgpr(tmpVgpr1 + vi + i * gwvw), src=vgpr(tmpVgpr1 + shiftOffset2+ i * gwvw), sel=HighBitSel.HIGH if vi % 2 else HighBitSel.LOW, comment="convert to FP32")) elif dataType.isBFloat16(): module.add(VCvtBF16toFP32(dst=vgpr(tmpVgpr1 + vi + i * gwvw), src=vgpr(tmpVgpr1 + shiftOffset2+ i * gwvw), vgprMask=None, vi=0)) elif dataType.isInt32(): module.add(VCvtI32toF32(dst=vgpr(tmpVgpr1 + vi + i * gwvw), src=vgpr(tmpVgpr1 + shiftOffset2+ i * gwvw), comment="convert I32 to FP32 (bias)")) elif dataType == kernel["ProblemType"]["ComputeDataType"]: pass # Same, no need to convert else: printExit("[Compute fp32] Unrecognized data type %s."%str(dataType)) elif kernel["ProblemType"]["ComputeDataType"].isInt32(): if dataType == kernel["ProblemType"]["ComputeDataType"]: pass # Same, no need to convert else: printExit("[Compute int32] Unrecognized data type %s."%str(dataType)) else: printExit("Does not support ComputeDataType == %s"%str(kernel["ProblemType"]["ComputeDataType"])) if setToOne: module.add(VCndMaskB32( dst=vgpr(tmpVgpr1 + i * gwvw), \ src1=vgpr(tmpVgpr1 + i * gwvw), \ src0=maskConst, \ src2=sgpr("Address%s"%addressStr, self.states.laneSGPRCount), \ comment="1. mul 1 if 0")) numRegs = max(1, bps // 4) module.add(dsStore(bps, dstAddr=dst, src=vgpr(tmpVgpr1 + i * gwvw * numRegs, numRegs), ds=ds, comment=comment)) return module def getNumOfTempVgprs(self, vectorDataTypes: VectorDataTypes, kernel, gwvw, dim): biasDataType = vectorDataTypes.bias(dim).dataType scaleADataType = vectorDataTypes.scaleA.dataType scaleBDataType = vectorDataTypes.scaleB.dataType scaleAlphaDataType = vectorDataTypes.scaleAlpha(dim).dataType # Calculate nums of vgpr for store data totalReg = 0 regPerVec = gwvw * kernel["ProblemType"]["ComputeDataType"].numRegisters() if biasDataType: vectorDataTypes.bias(dim).dstVgpr = totalReg vectorDataTypes.bias(dim).turn = self.getTurn(kernel, gwvw, dim)[0] totalReg = totalReg + (self.getTurn(kernel, gwvw, dim)[0] * regPerVec) if scaleAlphaDataType: vectorDataTypes.scaleAlpha(dim).dstVgpr = totalReg vectorDataTypes.scaleAlpha(dim).turn = self.getTurn(kernel, gwvw, dim)[0] totalReg = totalReg + (self.getTurn(kernel, gwvw, dim)[0] * regPerVec) if scaleADataType: vectorDataTypes.scaleA.dstVgpr = totalReg vectorDataTypes.scaleA.turn = self.getTurn(kernel, gwvw, 0)[0] totalReg = totalReg + (self.getTurn(kernel, gwvw, 0)[0] * regPerVec) if scaleBDataType: vectorDataTypes.scaleB.dstVgpr = totalReg vectorDataTypes.scaleB.turn = self.getTurn(kernel, gwvw, 1)[0] totalReg = totalReg + (self.getTurn(kernel, gwvw, 1)[0] * regPerVec) # Check how many additional sgpr is needed for global read tmpVgprNum = 0 offsetVgprStart = totalReg # Only vector without stride input can add to dimKey dimKey = {} if biasDataType: vectorDataTypes.bias(dim).offsetVgpr = offsetVgprStart tmpVgprNum = tmpVgprNum + 1 if scaleAlphaDataType: if (scaleAlphaDataType, 1) in dimKey: vectorDataTypes.scaleAlpha(dim).offsetVgpr = dimKey[(scaleAlphaDataType, dim)] else: vectorDataTypes.scaleAlpha(dim).offsetVgpr = offsetVgprStart + tmpVgprNum tmpVgprNum = tmpVgprNum + 1 if scaleADataType: if (scaleADataType, 0) in dimKey: vectorDataTypes.scaleA.offsetVgpr = dimKey[(scaleADataType, dim)] else: vectorDataTypes.scaleA.offsetVgpr = offsetVgprStart + tmpVgprNum tmpVgprNum = tmpVgprNum + 1 if scaleBDataType: if (scaleBDataType, 1) in dimKey: vectorDataTypes.scaleB.offsetVgpr = dimKey[(scaleBDataType, dim)] else: vectorDataTypes.scaleB.offsetVgpr = offsetVgprStart + tmpVgprNum tmpVgprNum = tmpVgprNum + 1 return totalReg + tmpVgprNum def readVectorToLDS(self, vectorDataTypes: VectorDataTypes, kernel, gwvw, offsetVgpr, tmpVgpr1Res: ContinuousRegister, dim, tmpSgpr = None): assert gwvw == 1 # Params biasDataType = vectorDataTypes.bias(dim).dataType scaleADataType = vectorDataTypes.scaleA.dataType scaleBDataType = vectorDataTypes.scaleB.dataType scaleAlphaDataType = vectorDataTypes.scaleAlpha(dim).dataType biasBpe = int(self.states.bpr * biasDataType.numRegisters()) if biasDataType else 0 scaleABpe = int(self.states.bpr * scaleADataType.numRegisters()) if scaleADataType else 0 scaleBBpe = int(self.states.bpr * scaleBDataType.numRegisters()) if scaleBDataType else 0 scaleAlphaBpe = int(self.states.bpr * scaleAlphaDataType.numRegisters()) if scaleAlphaDataType else 0 biasDstVgpr = vectorDataTypes.bias(dim).dstVgpr scaleADstVgpr = vectorDataTypes.scaleA.dstVgpr scaleBDstVgpr = vectorDataTypes.scaleB.dstVgpr scaleAlphaDstVgpr = vectorDataTypes.scaleAlpha(dim).dstVgpr biasOffsetVgpr = vectorDataTypes.bias(dim).offsetVgpr + tmpVgpr1Res.idx scaleAOffsetVgpr = vectorDataTypes.scaleA.offsetVgpr + tmpVgpr1Res.idx scaleBOffsetVgpr = vectorDataTypes.scaleB.offsetVgpr + tmpVgpr1Res.idx scaleAlphaOffsetVgpr = vectorDataTypes.scaleAlpha(dim).offsetVgpr + tmpVgpr1Res.idx module = Module("ReadVecToLds") module.addComment2("Read vector to LDS") # Calculate global offset- macro tile X part ## Common codes module.addModuleAsFlatItems(self.calculateVectorGlobalOffset(kernel, offsetVgpr, tmpSgpr, dim)) ## Scale for each component offsetIsInit = {} if biasDataType: # Recalculate bias length (SrdBias only valid for BufferLoad=1) if kernel["BufferLoad"]: module.add(SMulI32(dst=sgpr("SrdBias+2"), src0=hex(biasBpe), src1=sgpr("SrdBias+2"), comment="scaled by BPE")) module.add(self.shiftSrd("Bias")) if biasOffsetVgpr not in offsetIsInit: offsetIsInit[biasOffsetVgpr] = 1 module.addModuleAsFlatItems(self.calculateVectorGlobalStride(kernel, offsetVgpr, biasOffsetVgpr, tmpSgpr, dim, "BiasStride")) module.add(VLShiftLeftB32(dst=vgpr(biasOffsetVgpr), \ shiftHex=hex(log2(biasBpe)), \ src=vgpr(biasOffsetVgpr), \ comment="Global bias address scaled by BPE")) offsetSequences = [] if dim == 0: offsetSequences.append([scaleAlphaDataType, scaleAlphaOffsetVgpr, scaleAlphaBpe, "scaleAlpha", dim]) offsetSequences.append([scaleADataType, scaleAOffsetVgpr, scaleABpe, "scaleA", 0]) offsetSequences.append([scaleBDataType, scaleBOffsetVgpr, scaleBBpe, "scaleB", 1]) else: offsetSequences.append([scaleAlphaDataType, scaleAlphaOffsetVgpr, scaleAlphaBpe, "scaleAlpha", dim]) offsetSequences.append([scaleBDataType, scaleBOffsetVgpr, scaleBBpe, "scaleB", 1]) offsetSequences.append([scaleADataType, scaleAOffsetVgpr, scaleABpe, "scaleA", 0]) for index, offsetSequence in enumerate(offsetSequences): if offsetSequence[4] != dim: dimAnother = 1 if dim == 0 else 0 module.addModuleAsFlatItems(self.calculateVectorGlobalOffset(kernel, offsetVgpr, tmpSgpr, dimAnother)) if offsetSequence[0] and (offsetSequence[1] not in offsetIsInit): offsetIsInit[offsetSequence[1]] = 1 module.add(VLShiftLeftB32(dst=vgpr(offsetSequence[1]), \ shiftHex=hex(log2(offsetSequence[2])), \ src=vgpr(offsetVgpr), \ comment="Global %s address scaled by BPE"%offsetSequence[3])) # global load globalLoadsModule = Module("Global Loads") if biasDataType: biasShiftOffset = self.getGlobalShiftOffset(kernel, biasDataType, gwvw) globalLoadsModule.addModuleAsFlatItems(self.addVectorGlobalLoad(kernel, "Bias", biasOffsetVgpr, biasShiftOffset, biasDataType, biasBpe, gwvw, tmpVgpr1Res, biasDstVgpr, dim)) if scaleAlphaDataType: scaleAlphaShiftOffset = self.getGlobalShiftOffset(kernel, scaleAlphaDataType, gwvw) globalLoadsModule.addModuleAsFlatItems(self.addVectorGlobalLoad(kernel, "ScaleAlphaVec", scaleAlphaOffsetVgpr, scaleAlphaShiftOffset, scaleAlphaDataType, scaleAlphaBpe, gwvw, tmpVgpr1Res, scaleAlphaDstVgpr, dim)) if scaleADataType: scaleAShiftOffset = self.getGlobalShiftOffset(kernel, scaleADataType, gwvw) globalLoadsModule.addModuleAsFlatItems(self.addVectorGlobalLoad(kernel, "ScaleA", scaleAOffsetVgpr, scaleAShiftOffset, scaleADataType, scaleABpe, gwvw, tmpVgpr1Res, scaleADstVgpr, 0)) if scaleBDataType: scaleBShiftOffset = self.getGlobalShiftOffset(kernel, scaleBDataType, gwvw) globalLoadsModule.addModuleAsFlatItems(self.addVectorGlobalLoad(kernel, "ScaleB", scaleBOffsetVgpr, scaleBShiftOffset, scaleBDataType, scaleBBpe, gwvw, tmpVgpr1Res, scaleBDstVgpr, 1)) # Count global loads (covers both buffer loads and flat loads) vlcnt = 0 for item in globalLoadsModule.items(): if isinstance(item, GlobalReadInstruction): vlcnt = vlcnt + 1 module.add(globalLoadsModule) assert vlcnt > 0 # Local write # In local write, all vector shares the same offsetVgpr since the internal data types are all the same. module.add(VLShiftLeftB32(dst=vgpr(offsetVgpr), \ shiftHex=hex(log2(self.states.bpeCinternal)), \ src=vgpr("Serial"), \ comment="Local address scaled by BPE")) if kernel["LdsOffsetBias"] != 0: module.add(VAddU32(dst=vgpr(offsetVgpr), \ src0=(kernel["LdsOffsetBias"]), \ src1=vgpr(offsetVgpr), \ comment="add lds offset")) # Get all local stores storeModules = Module("Store") subGroupOffset = [0] if biasDataType: vectorDataTypes.bias(dim).ldsOffset = subGroupOffset[0] storeModules.add(self.addVectorLocalStore(kernel, "Bias", offsetVgpr, biasShiftOffset, biasDataType, gwvw, tmpVgpr1Res, biasDstVgpr, subGroupOffset, dim, comment="store bias")) subGroupOffset[0] += kernel["NumThreads"] * int(kernel["ProblemType"]["ComputeDataType"].numBytes()) * vectorDataTypes.bias(dim).turn if scaleAlphaDataType: vectorDataTypes.scaleAlpha(dim).ldsOffset = subGroupOffset[0] storeModules.add(self.addVectorLocalStore(kernel, "ScaleAlphaVec", offsetVgpr, scaleAlphaShiftOffset, scaleAlphaDataType, gwvw, tmpVgpr1Res, scaleAlphaDstVgpr, subGroupOffset, dim, setToOne=True, comment="store scaleAlpha")) subGroupOffset[0] += kernel["NumThreads"] * int(kernel["ProblemType"]["ComputeDataType"].numBytes()) * vectorDataTypes.scaleAlpha(dim).turn if scaleADataType: vectorDataTypes.scaleA.ldsOffset = subGroupOffset[0] storeModules.add(self.addVectorLocalStore(kernel, "ScaleA", offsetVgpr, scaleAShiftOffset, scaleADataType, gwvw, tmpVgpr1Res, scaleADstVgpr, subGroupOffset, 0, setToOne=True, comment="store scaleA")) subGroupOffset[0] += kernel["NumThreads"] * int(kernel["ProblemType"]["ComputeDataType"].numBytes()) * vectorDataTypes.scaleA.turn if scaleBDataType: vectorDataTypes.scaleB.ldsOffset = subGroupOffset[0] storeModules.add(self.addVectorLocalStore(kernel, "ScaleB", offsetVgpr, scaleBShiftOffset, scaleBDataType, gwvw, tmpVgpr1Res, scaleBDstVgpr, subGroupOffset, 1, setToOne=True, comment="store scaleB")) subGroupOffset[0] += kernel["NumThreads"] * int(kernel["ProblemType"]["ComputeDataType"].numBytes()) * vectorDataTypes.scaleB.turn # We move s_barrier before local load. Add barrier here to avoid race condition if lds offset starts from 0 if kernel["LdsOffsetBias"] == 0: module.add(SBarrier(comment="wait for all global loads.")) # rearrange them and add waitcnt for storeModule in storeModules.items(): isAdded = False if isinstance(storeModule, Module): for item in storeModule.items(): if isinstance(item, DSStoreInstruction): item.setMemToken(MemTokenData([self.states.memTokenLdsBuffer0])) if (not isAdded) and isinstance(item, (VCvtInstruction, DSStoreInstruction, VCndMaskB32, VLShiftLeftB32, VAndB32)): vlcnt = vlcnt - 1 module.add(SWaitCnt(vlcnt=(vlcnt), comment="wait for global load")) module.add(item) isAdded = True else: module.add(item) # restore after ds_store if isinstance(item, DSStoreInstruction): isAdded = False else: if isinstance(storeModule, DSStoreInstruction): storeModule.setMemToken(MemTokenData([self.states.memTokenLdsBuffer0])) module.add(storeModule) # Emit the epilogue vector-LDS drain barrier for multi-DU kernels only # (matches the GlobalWriteBatch bias/SAV barrier gating). Single-DU MX # kernels do not need it. if kernel.get("UseSubtileImpl") and isSubtileMultiDU(kernel): module.add(SWaitCnt(dscnt=0, comment="drain epilogue vector LDS writes")) module.add(SBarrier(comment="sync waves after epilogue vector LDS writes")) return module ''' Read reduction results from LDS In edge case the gwvw will be set to 1. ''' def writeBiasToBlobalLdsRead(self, kernel, offsetVgpr, gwvw, maxKId, outVgpr): module = Module("WriteBiasToGlobalLdsRead") module.add(VLShiftLeftB32(dst=vgpr(offsetVgpr), \ shiftHex=hex(log2(self.states.bpeCinternal)), \ src=vgpr(offsetVgpr), \ comment="Local bias address scaled by BPE")) if kernel["LdsOffsetBias"] != 0: module.add(VAddU32(dst=vgpr(offsetVgpr), \ src0=(kernel["LdsOffsetBias"]), \ src1=vgpr(offsetVgpr), \ comment="add bias lds offset")) srcAddr = vgpr(offsetVgpr) outVgprN = outVgpr if maxKId == 1: bps = kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvw numRegs = max(1, bps // 4) ds = DSModifiers(offset=0) module.add(dsLoad(bps, dst=vgpr(outVgprN, numRegs), src=srcAddr, ds=ds, comment="load bias")) outVgprN += numRegs else: dsOffset = 0 for _ in range(0, gwvw): idx = 0 while idx < maxKId: gwvwK = 2 if (idx + 1 < maxKId * gwvw) else 1 bps = kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvwK numRegs = max(1, bps // 4) ds = DSModifiers(offset=dsOffset) module.add(dsLoad(bps, dst=vgpr(outVgprN, numRegs), src=srcAddr, ds=ds, comment="load bias")) outVgprN += gwvwK idx += gwvwK dsOffset += bps return module, outVgprN ''' This is a tiny version of globalWriteElements for Bias reduction. Every MT saves MTx1 bias data to global. When GSU = 1, the address is set to the bias pointer. When GSU > 1 in multiple buffer mode, the address is set to the work space pointer. Wider DS load is always enabled. Wider global store only enables when freeElementMultiple % gwvw == 0 since each thread only stores 1, 2 elements. ''' def writeBiasToGlobal(self, biasDataType, kernel, tP, gwvw, offsetVgpr, tmpSgprRes, tmpVgpr1Res: ContinuousRegister): tile01 = tP["tile01Idx"] mt = kernel["MacroTile%u" % tile01] maxKId = self.states.lraTileProperties[tile01].maxKId assert tmpSgprRes.size >= 1 assert tmpVgpr1Res.size >= kernel["VectorWidthA"] * kernel["ProblemType"]["ComputeDataType"].numRegisters() * maxKId assert gwvw % 2 == 0 or gwvw == 1 # Params biasBpe = int(self.states.bpr * biasDataType.numRegisters()) module = Module("WriteBiasToGlobal") module.addComment2("Write Bias to Global") module.add(SBarrier(comment="wait for bias lds store.")) # Recalculate bias length gsuComponent = Component.GSU.find(self) module.add(gsuComponent.writeBiasToGlobal(self, kernel, biasDataType, tP, tmpSgprRes, biasBpe)) # Num records (SrdBias only valid for BufferStore=1; flat store uses pre-computed 64-bit address) if kernel["BufferStore"]: module.add(SMulI32(dst=sgpr("SrdBias+2"), src0=hex(biasBpe), src1=sgpr("SrdBias+2"), comment="scaled by BPE")) module.add(self.shiftSrd("Bias")) # Local read # remaining size % VW module.add(vectorStaticMultiply(vgpr(offsetVgpr), vgpr("Serial"), maxKId, tmpSgprRes, \ "offset = serial * maxKId")) module.add(vectorStaticMultiply(vgpr(offsetVgpr), vgpr(offsetVgpr), gwvw, tmpSgprRes, \ "apply VectorWidth: offset = bnOffset * vw(%u)" % gwvw)) enableEdge = False serialOffsetVgpr = tmpVgpr1Res.idx serialOffsetVgpr2 = tmpVgpr1Res.idx + 1 # ShiftPtr if not (kernel["BufferLoad"] and kernel["GuaranteeNoPartialA"]) and (kernel["ProblemType"]["BiasSrc"] == "A"): enableEdge = True if not (kernel["BufferLoad"] and kernel["GuaranteeNoPartialB"]) and (kernel["ProblemType"]["BiasSrc"] == "B"): enableEdge = True if kernel["EdgeType"] == "ShiftPtr" and enableEdge: jumpLabel = Label(self.labels.getNameInc("ShiftPtrSkip"), comment="Skip shift ptr") jumpLabelEnd = Label(self.labels.getNameInc("ShiftPtrEnd"), comment="Skip shift ptr end") assert tmpSgprRes.size >= 5 # For ShiftPtr assert tP["glvw"] % gwvw == 0 # Make sure the magic trick works (serial = serial * gwvw) tmpSgpr = tmpSgprRes.idx tmpSgprShift = ContinuousRegister(idx=tmpSgprRes.idx+2, size=3) margin = tP["glvw"] if tP["rtv"] else 1 module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tP["wg"]), src1=kernel[tP["mt"]], comment="WorkGroup[01] * MT")) module.add(SSubU32(dst=sgpr(tmpSgpr), src0=self.sizeRef(tP["idx"]), src1=sgpr(tmpSgpr), \ comment="edge = Size%s - WG*MT"%(tP["tileChar"]))) module.add(SCmpGeU32(src0=sgpr(tmpSgpr), src1=kernel[tP["mt"]], comment="skip shift ptr if edge >= MT")) module.add(SCBranchSCC1(labelName=jumpLabel.getLabelName(), comment="" )) module.add(scalarStaticRemainder(tmpSgpr+1, tmpSgpr+1, tmpSgpr, tP["glvw"], tmpSgprShift, comment="remainder = edge %% glvw(%d)"%tP["glvw"])) module.add(SSubU32(dst=sgpr(tmpSgpr+1), src0=hex(tP["glvw"]), src1=sgpr(tmpSgpr+1), comment="shift = glvw(%d) - remainder"%tP["glvw"])) #module.add(SCmpKEQU32(src=sgpr(tmpSgpr+1), simm16=hex(tP["glvw"]), comment="if shift == glvw(%d)?"%tP["glvw"])) module.add(self.getSCMPKInstruction("EQU32", tmpSgpr+1, tP["glvw"], comment="if shift == glvw(%d)?"%tP["glvw"])) module.add(SCMovB32(dst=sgpr(tmpSgpr+1), src=0, comment="shift = glvw(%d) ? 0 : NOP"%(tP["glvw"]))) module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=margin, comment="edge -= margin(%u)"%(margin))) module.add(SCMovB32(dst=sgpr(tmpSgpr), src=0, comment="edge = edge < 0 ? 0 : edge")) # Saturation if gwvw != 1: ''' Edge case, the gwvw is set to 1 instead, do boundary check on every element. if gwvw > 1: for g in rangw(gwvw): if(out_of_bound(g)) offset_at_g += shit_offset ds_read vpgr_at_g, offset_at_g else: if(out_of_bound(g)) offset_at_g += shit_offset ds_read vpgr_at_g, offset_at_g ''' module.add(vectorStaticMultiply(vgpr(serialOffsetVgpr), vgpr("Serial"), gwvw, tmpSgprShift, comment="serial = serial * gwvw")) for g in range(gwvw): if g != 0: module.add(VAddU32(dst=vgpr(offsetVgpr+g), src0=vgpr(offsetVgpr+(g-1)), src1=maxKId, comment="new offset = offset + maxKId")) module.add(VMulLOU32(dst=vgpr(serialOffsetVgpr2), src0=hex(maxKId), src1=sgpr(tmpSgpr+1), comment="ds_offset = K * offset")) for g in range(gwvw): module.add(VCmpXGeU32(dst=EXEC(), src0=vgpr(serialOffsetVgpr), src1=sgpr(tmpSgpr), comment="needs shift if serial > edge")) module.add(VAddU32(dst=vgpr(offsetVgpr+g), src0=vgpr(offsetVgpr+g), src1=vgpr(serialOffsetVgpr2), comment="real offset = offset + ds_offset")) if g < gwvw - 1: module.add(VAddU32(dst=vgpr(serialOffsetVgpr), src0=1, src1=vgpr(serialOffsetVgpr), comment="inc += 1")) else: module.add(VCmpXGeU32(dst=EXEC(), src0=vgpr("Serial"), src1=sgpr(tmpSgpr), comment="needs shift if serial > edge")) module.add(VMulLOU32(dst=vgpr(serialOffsetVgpr), src0=hex(maxKId), src1=sgpr(tmpSgpr+1), comment="ds_offset = K * offset")) module.add(VAddU32(dst=vgpr(offsetVgpr), src0=vgpr(offsetVgpr), src1=vgpr(serialOffsetVgpr), comment="real offset = offset + ds_offset")) SMovBX_sp = SMovB32 if (kernel["WavefrontSize"] == 32) else SMovB64 module.add(SMovBX_sp(dst=EXEC(), src=-1, comment="reset mask")) outVgprNext = tmpVgpr1Res.idx # Shiftptr for g in range(gwvw): wb2blr, outVgprNext = self.writeBiasToBlobalLdsRead(kernel, offsetVgpr + g, 1, maxKId, outVgprNext) module.add(wb2blr) module.add(SBranch(labelName=jumpLabelEnd.getLabelName(), comment="")) module.add(jumpLabel) # Shiftptr case, but no need to shift wb2blrGwvw, _ = self.writeBiasToBlobalLdsRead(kernel, offsetVgpr, gwvw, maxKId, tmpVgpr1Res.idx) module.add(wb2blrGwvw) module.add(jumpLabelEnd) else: # Non-shiftptr case wb2blrGwvw, _ = self.writeBiasToBlobalLdsRead(kernel, offsetVgpr, gwvw, maxKId, tmpVgpr1Res.idx) module.add(wb2blrGwvw) module.add(SWaitCnt(dscnt=0, comment="wait for bias lds load")) # Sum K (if needed) ''' The vgprs are rearranged in this step. For example, gwvw = 2, k = 2, we have [v6, v7] [v8, v9] v6 = v6 + v7 v7 = v8 + v9 ''' tmpVgpr1 = tmpVgpr1Res.idx tmpVgprN = tmpVgpr1 + 1 if maxKId != 1: for gidx in range(0, gwvw): tmpVgprAccum = tmpVgpr1 + gidx if gidx != 0: module.add(VMovB32(dst=vgpr(tmpVgprAccum), src=vgpr(tmpVgprN), comment="Copy address")) tmpVgprN += 1 for idx in range(1, maxKId): if kernel["ProblemType"]["ComputeDataType"].isSingle(): module.add(VAddF32(dst=vgpr(tmpVgprAccum), src0=vgpr(tmpVgprN), src1=vgpr(tmpVgprAccum), comment="Sum K")) else: assert 0 tmpVgprN += 1 # Convert freeElementMultiple = kernel["AssertFree%dElementMultiple"%tile01] enablePack = True if (freeElementMultiple % gwvw == 0) else False tmpVgprN = tmpVgpr1 if biasDataType != kernel["ProblemType"]["ComputeDataType"]: bf16CVTVgprStruct = None bf16CVTVgpr = None if biasDataType.isBFloat16(): bf16CVTVgpr = self.vgprPool.checkOut(4, tag="writeBiasToGlobal_bf16CVTVgpr") bf16CVTVgprStruct = self.BF16CVTVgprStruct(vgprBf16Temp=bf16CVTVgpr, vgprBf16Mask=(bf16CVTVgpr+1), \ vgprFp32Nan=(bf16CVTVgpr+2), vgprBf16Inc=(bf16CVTVgpr+3)) module.add(VMovB32(vgpr(bf16CVTVgprStruct.vgprBf16Mask), "0xffff0000", comment="mask for pack two bfloat16 element to 32bit" )) module.add(VMovB32(vgpr(bf16CVTVgprStruct.vgprFp32Nan), "0x7fff0000", comment="fp32 Nan" )) module.add(VMovB32(vgpr(bf16CVTVgprStruct.vgprBf16Inc), "0x7fff", comment="rounding bias for bfloat16" )) for vi in range(gwvw): # Does not support hi/lo yet if kernel["ProblemType"]["ComputeDataType"].isSingle(): if biasDataType.isHalf(): module.add(VCvtF32toF16(dst=vgpr(tmpVgprN), src=vgpr(tmpVgprN), comment="convert to FP16")) if vi % 2 == 1 and enablePack: module.add(VPackF16toB32(dst=vgpr(tmpVgprN - 1), src0=vgpr(tmpVgprN - 1), src1=vgpr(tmpVgprN), \ comment="Pack with neighbor")) elif biasDataType.isBFloat16(): module.add(VCmpUF32(dst=sgpr(tmpSgprRes.idx,self.states.laneSGPRCount), src0=vgpr(tmpVgprN), src1=vgpr(tmpVgprN), comment="check Nan")) module.add(VBfeU32(dst=vgpr(bf16CVTVgprStruct.vgprBf16Temp), src0=vgpr(tmpVgprN), src1=16, src2=1, \ comment="Non-Nan case: store lsb of bf16" )) module.add(VAdd3U32(dst=vgpr(bf16CVTVgprStruct.vgprBf16Temp), src0=vgpr(tmpVgprN), src1=vgpr(bf16CVTVgprStruct.vgprBf16Temp), \ src2=vgpr(bf16CVTVgprStruct.vgprBf16Inc), comment="Non-Nan case: add lsb and the increment for rounding" )) module.add(VCndMaskB32(dst=vgpr(tmpVgprN), src0=vgpr(bf16CVTVgprStruct.vgprBf16Temp), \ src1=vgpr(bf16CVTVgprStruct.vgprFp32Nan), src2=sgpr(tmpSgprRes.idx,self.states.laneSGPRCount))) module.add(VLShiftRightB32(dst=vgpr(tmpVgprN), shiftHex=16, src=vgpr(tmpVgprN), comment="convert to bf16")) if vi % 2 == 1 and enablePack: module.add(VPackF16toB32(dst=vgpr(tmpVgprN - 1), src0=vgpr(tmpVgprN - 1), src1=vgpr(tmpVgprN), \ comment="Pack with neighbor")) elif biasDataType == kernel["ProblemType"]["ComputeDataType"]: pass # Same, no need to convert else: printExit("Unrecognized bias type %s."%str(biasDataType)) tmpVgprN += 1 else: printExit("Does not support ComputeDataType != float") if bf16CVTVgpr != None: self.vgprPool.checkIn(bf16CVTVgpr) # Global write # Calculate global offset- macro tile 0 part tmpSgpr = tmpSgprRes.idx module.add(SMulI32(dst=sgpr(tmpSgpr), src0=mt, src1=sgpr("WorkGroup%u" % tile01), comment="wgp * MT")) module.add(vectorStaticMultiply(vgpr(offsetVgpr), vgpr("Serial"), gwvw, tmpSgprRes, \ "apply VectorWidth: offset = serial * vw(%u)" % gwvw)) module.add(VAddU32(dst=vgpr(offsetVgpr), src0=sgpr(tmpSgpr), src1=vgpr(offsetVgpr), comment="coord = wgp * MT + thread offset")) module.add(VLShiftLeftB32(dst=vgpr(offsetVgpr), \ shiftHex=hex(log2(biasBpe)), \ src=vgpr(offsetVgpr), \ comment="Global bias address scaled by BPE")) with self.allocTmpSgpr(1, 1, tag="writeBiasToGlobal_tmpSgprRes") as tmpSgprRes: module.add(SMovB32(dst=sgpr(tmpSgprRes.idx), src=hex(mt//gwvw), comment="%d=%d//%d"%(mt//gwvw, mt, gwvw))) module.add(VCmpXLtU32(dst=EXEC(), src0=vgpr("Serial"), src1=sgpr(tmpSgprRes.idx), comment="if serial < MacroTile%d/gwvw"%tile01)) if kernel["BufferStore"]: addr0 = vgpr(offsetVgpr) addr1 = sgpr("SrdBias", 4) else: addr0 = vgpr(offsetVgpr, 2) addr1 = "" dataType = biasDataType useBuffer = kernel["BufferStore"] tmpVgprN = tmpVgpr1 if enablePack: # no partial bps = int(biasDataType.numBytes()) * gwvw rpe = int(biasDataType.numBytes()) / self.states.bpr rpv = rpe * gwvw if dataType.isHalf() or dataType.isBFloat16(): module.add(self.chooseGlobalWrite(useBuffer, bps, tmpVgprN, rpv, \ addr0, addr1, offset=0, hi16=0, comment="global store bias")) elif dataType.isInt32() or dataType.isSingle(): module.add(self.chooseGlobalWrite(useBuffer, bps, tmpVgprN, rpv, \ addr0, addr1, offset=0, comment="global store bias")) elif dataType.isDouble() or dataType.isSingleComplex() : module.add(self.chooseGlobalWrite(useBuffer, bps, tmpVgprN, rpv, \ addr0, addr1, offset=0, comment="global store bias")) else: # edge tmpVgprNStep = max(1, biasDataType.numRegisters()) globalOffset = 0 for gidx in range(0, gwvw): bps = int(biasDataType.numBytes()) rpe = int(biasDataType.numBytes()) / self.states.bpr rpv = rpe if dataType.isHalf() or dataType.isBFloat16(): module.add(self.chooseGlobalWrite(useBuffer, bps, tmpVgprN, rpv, \ addr0, addr1, offset=globalOffset, hi16=0, comment="global store bias")) elif dataType.isInt32() or dataType.isSingle(): module.add(self.chooseGlobalWrite(useBuffer, bps, tmpVgprN, rpv, \ addr0, addr1, offset=globalOffset, comment="global store bias")) elif dataType.isDouble() or dataType.isSingleComplex() : module.add(self.chooseGlobalWrite(useBuffer, bps, tmpVgprN, rpv, \ addr0, addr1, offset=globalOffset, comment="global store bias")) tmpVgprN += tmpVgprNStep globalOffset += biasBpe SMovBX = SMovB32 if (kernel["WavefrontSize"] == 32) else SMovB64 module.add(SMovBX(dst=EXEC(), src=-1, comment="Reset exec mask")) return module ######################################## # Amax related ######################################## def amax_define_load_res(self) -> Module: module = Module("AmaxD Set and Load") module.addComment0("AmaxD Set and Load") self.amaxVgprIdxVec = self.defineMultiVgprIndex(self.amaxVgprNames, self.amaxVgprSizes, align=1) for i in range(0, len(self.amaxVgprNames)): name = self.amaxVgprNames[i] idx = self.amaxVgprIdxVec[i] module.add(RegSet("v", "vgpr"+name, idx)) module.addSpaceLine() module.add(self.defineSgpr("Src", 4, 4)) module.add(self.defineSgpr("Dst", 4, 4)) module.add(self.defineSgpr("Offset", 1)) module.add(self.defineSgpr("Tmp", 6, 2)) module.add(self.defineSgpr("NumGroup", 1)) module.add(self.defineSgpr("WGIdx", 1)) module.addSpaceLine() # defineMulti (ensure they are checkout together) for SGPRs that are used to load args self.amaxSgprIdxVec = self.defineMultiSgprIndex(self.amaxSgprArgNames, self.amaxSgprArgSizes, align=4) for name in self.amaxSgprArgNames: module.add(RegSet("s", "sgpr"+name, self.sgprs[name])) module.addSpaceLine() # TODO- why we don't directly update the offset in the last argLoader ? argOffset = self.argLoader.getOffset() argOffset += (self.states.numStoreSgprToLoad) * 4 module.add(self.argLoader.loadKernArg("AddrAmaxOut", "KernArgAddress", sgprOffset=hex(argOffset), dword=4)) argOffset += 16 # advance dwordx4 module.add(self.argLoader.loadKernArg("AddressSy", "KernArgAddress", sgprOffset=hex(argOffset), dword=2)) module.add(SMulI32(sgpr("NumGroup"), sgpr("NumWorkGroups0"), sgpr("NumWorkGroups1"), "get total num_wgs")) module.add(SMulI32(sgpr("WGIdx"), sgpr("WorkGroup1"), sgpr("NumWorkGroups0"), "wgId = wg1 * numWG0")) module.add(SAddI32(sgpr("WGIdx"), sgpr("WGIdx"), sgpr("WorkGroup0"), "wgId += wg0")) module.addSpaceLine() module.add(SWaitCnt(kmcnt=0)) module.addSpaceLine() return module def amax_intra_wave_reduction(self, kernel, postfix) -> Module: wave_size = kernel["WavefrontSize"] label = Label(f"permute_{postfix}", f"permute_{postfix}") mod = Module("intra_wave_reduction") mod.addComment0("intra_wave_reduction") mod.add(SMovB32(sgpr("Tmp"), 1)) mod.add(label) mod.addSpaceLine() mod.add(VAddU32(vgpr("Tmp"), sgpr("Tmp"), vgpr("Serial"))) mod.add(VAndB32(vgpr("Tmp"), wave_size-1, vgpr("Tmp"))) mod.add(VLShiftLeftB32(vgpr("Tmp"), 0x2, vgpr("Tmp"))) mod.addSpaceLine() mod.add(DSBPermuteB32(vgpr("AmaxOutB"), vgpr("Tmp"), vgpr("AmaxOut"))) mod.add(SWaitCnt(dscnt=0)) mod.addSpaceLine() # TODO- F16 mod.add(VMaxF32(vgpr("AmaxOut"), vgpr("AmaxOut"), vgpr("AmaxOutB"))) mod.add(SLShiftLeftB32(sgpr("Tmp"), 1, sgpr("Tmp"))) mod.add(SCmpLtU32(sgpr("Tmp"), wave_size)) mod.add(SCBranchSCC1(label.getLabelName())) mod.addSpaceLine() return mod def amax_inter_wave_reduction(self, kernel) -> Module: wave_size = kernel["WavefrontSize"] numWorkItems = kernel["NumThreads"] amaxOutType = kernel["ProblemType"]["DataTypeAmaxD"] amax_lds_start = kernel["LdsBytesNoAmax"] label_wave_inter = Label("wave_inter", 'wave_inter') label_wave_upper = Label("wave_upper", 'wave_upper') label_wave_lower = Label("wave_lower", 'wave_lower') label_wave_empty = Label("wave_empty", 'wave_empty') label_wave_end = Label("wave_end", 'wave_end') mod = Module("inter_wave_reduction") mod.addComment0("inter_wave_reduction") mod.add(VLShiftRightB32(vgpr("Widx"), int(log2(wave_size)), vgpr("Serial"))) mod.add(SMovB32(sgpr("Offset"), numWorkItems // wave_size)) mod.add(label_wave_inter) mod.add(SLShiftRightB32(sgpr("Offset"), 1, sgpr("Offset"))) mod.add(SCmpEQU32(sgpr("Offset"), 0)) mod.add(SCBranchSCC1(label_wave_end.getLabelName())) mod.add(SLShiftLeftB32(sgpr("Tmp"), 1, sgpr("Offset"))) mod.add(VCmpLtU32(sgpr("Tmp+2",2), vgpr("Widx"), sgpr("Tmp"))) mod.add(VCmpGEU32(sgpr("Tmp+4",2), vgpr("Widx"), sgpr("Offset"))) mod.add(SAndB64(VCC(), sgpr("Tmp+2",2), sgpr("Tmp+4",2))) mod.add(SCBranchVCCNZ(label_wave_upper.getLabelName())) mod.add(VCmpLtU32(VCC(), vgpr("Widx"), sgpr("Offset"))) mod.add(SCBranchVCCNZ(label_wave_lower.getLabelName())) mod.add(SBranch(label_wave_empty.getLabelName())) mod.add(label_wave_upper) mod.add(VSubU32(vgpr("Tmp"), vgpr("Widx"), sgpr("Offset"))) mod.add(VLShiftLeftB32(vgpr("Tmp"), int(log2(amaxOutType.numBytes())), vgpr("Tmp"))) # TODO- select inst ds = DSModifiers(offset=amax_lds_start) mod.add(DSStoreB32(vgpr("Tmp"), vgpr("AmaxOut"), ds)) mod.add(SWaitCnt(dscnt=0)) mod.add(SBarrier()) mod.add(SBranch(label_wave_inter.getLabelName())) mod.add(label_wave_lower) mod.add(SBarrier()) mod.add(VLShiftLeftB32(vgpr("Tmp"), int(log2(amaxOutType.numBytes())), vgpr("Widx"))) # TODO- select inst mod.add(DSLoadB32(vgpr("AmaxOutB"), vgpr("Tmp"), ds)) mod.add(SWaitCnt(dscnt=0)) # TODO- F16 mod.add(VMaxF32(vgpr("AmaxOut"), vgpr("AmaxOut"), vgpr("AmaxOutB"))) mod.add(SBranch(label_wave_inter.getLabelName())) mod.add(label_wave_empty) mod.add(SBarrier()) mod.add(SBranch(label_wave_inter.getLabelName())) mod.add(label_wave_end) mod.addSpaceLine() return mod def amax_broadcast(self, kernel) -> Module: amax_lds_start = kernel["LdsBytesNoAmax"] label_lower = Label("broadcast_lower", f'broadcast_lower') label_end = Label("broadcast_end", f'broadcast_end') mod = Module("broadcast") mod.addComment0("broadcast") mod.add(VCmpEQU32(VCC(), vgpr("Widx"), 0)) mod.add(SCBranchVCCZ(label_lower.getLabelName())) # TODO- select inst ds = DSModifiers(offset=amax_lds_start) mod.add(DSStoreB32(vgpr("Widx"), vgpr("AmaxOut"), ds)) mod.add(SWaitCnt(dscnt=0)) mod.add(SBarrier()) mod.add(SBranch(label_end.getLabelName())) mod.add(label_lower) mod.add(SBarrier()) mod.add(VMovB32(vgpr("Tmp"), 0)) # TODO- select inst mod.add(DSLoadB32(vgpr("AmaxOut"), vgpr("Tmp"), ds)) mod.add(SWaitCnt(dscnt=0)) mod.add(label_end) mod.addSpaceLine() mod.addSpaceLine() return mod def amax_output_result(self, kernel) -> Module: wave_size = kernel["WavefrontSize"] amaxInType = kernel["ProblemType"]["ComputeDataType"] amaxOutType = kernel["ProblemType"]["DataTypeAmaxD"] mod = Module("output_result") mod.addComment0("output_result") label_end = Label("end", 'end') label_final_loop = Label("final_loop", 'final_loop') label_final_output = Label("final_output", 'final_output') mod.addSpaceLine() mod.add(VReadfirstlaneB32(sgpr("Tmp"), vgpr("Serial"))) mod.add(SCmpEQU32(sgpr("Tmp"), 0)) mod.add(SCBranchSCC0(label_end.getLabelName())) mod.addSpaceLine() # if self.arch.find("gfx94") != -1: mod.addSpaceLine() mod.add(SCmpEQU32(sgpr("NumGroup"), 1)) mod.add(SCBranchSCC1(label_final_output.getLabelName())) mod.add(SLShiftLeftB32(sgpr("Tmp"), int(log2(amaxInType.numBytes())), sgpr("NumGroup"))) mod.add(SMovB32(sgpr("Dst+0"), sgpr("AddressWk+0"))) mod.add(SMovB32(sgpr("Dst+1"), sgpr("AddressWk+1"))) mod.add(SMovB32(sgpr("Dst+2"), sgpr("Tmp"))) mod.add(SMovB32(sgpr("Dst+3"), "Srd127_96")) mod.add(SLShiftLeftB32(sgpr("Offset"), int(log2(amaxInType.numBytes())), sgpr("WGIdx"))) mod.add(VMovB32(vgpr("Offset"), 0)) # TODO- select inst mod.add(BufferStoreB32(vgpr("AmaxOut"), vgpr("Offset"), sgpr("Dst",4), sgpr("Offset"), MUBUFModifiers(offen=True, glc=True, slc=True))) mod.add(SWaitCnt(vlcnt=0)) mod.addSpaceLine() mod.add(SSubI32(sgpr("Tmp"), sgpr("NumGroup"), 1)) mod.add(SAtomicDec(sgpr("Tmp"), sgpr("AddressSy",2), SMEMModifiers(glc=True))) mod.add(SWaitCnt(vlcnt=0, kmcnt=0)) mod.add(SCmpEQU32(sgpr("Tmp"), 1)) mod.add(SCBranchSCC0(label_end.getLabelName())) mod.addSpaceLine() mod.add(SLShiftLeftB32(sgpr("Tmp"), int(log2(amaxInType.numBytes())), sgpr("NumGroup"))) mod.add(SMovB32(sgpr("Src+0"), sgpr("AddressWk+0"))) mod.add(SMovB32(sgpr("Src+1"), sgpr("AddressWk+1"))) mod.add(SMovB32(sgpr("Src+2"), sgpr("Tmp"))) mod.add(SMovB32(sgpr("Src+3"), "Srd127_96")) mod.addSpaceLine() mod.add(VLShiftLeftB32(vgpr("Offset"), int(log2(amaxOutType.numBytes())), vgpr("Serial"))) mod.addSpaceLine() mod.add(VMovB32(vgpr("AmaxOut"), "0")) mod.addSpaceLine() mod.add(label_final_loop) # TODO- select inst mod.add(BufferLoadB32(vgpr(f"Value"), vgpr("Offset"), sgpr("Src",4), 0, MUBUFModifiers(offen=True, glc=True, slc=True))) mod.add(SWaitCnt(vlcnt=0)) mod.addSpaceLine() # TODO- F16? mod.add(VMaxF32(vgpr("AmaxOut"), vgpr("AmaxOut"), vgpr("Value", isAbs=True))) mod.addSpaceLine() mod.add(SMovB32(sgpr("Tmp"), wave_size * int(amaxInType.numBytes()))) mod.add(VAddU32(vgpr("Offset"), vgpr("Offset"), sgpr("Tmp"))) mod.addSpaceLine() mod.add(SSubI32(sgpr("NumGroup"), sgpr("NumGroup"), wave_size)) mod.add(SCmpGtI32(sgpr("NumGroup"), 0)) mod.add(SCBranchSCC1(label_final_loop.getLabelName())) mod.addSpaceLine() mod.add(self.amax_intra_wave_reduction(kernel, "final")) mod.addSpaceLine() mod.add(label_final_output) mod.add(SMovB32(sgpr("Dst+0"), sgpr("AddrAmaxOut+0"))) mod.add(SMovB32(sgpr("Dst+1"), sgpr("AddrAmaxOut+1"))) mod.add(SMovB32(sgpr("Dst+2"), int(amaxOutType.numBytes()))) mod.add(SMovB32(sgpr("Dst+3"), "Srd127_96")) mod.addSpaceLine() mod.add(VMovB32(vgpr("Offset"), 0)) # TODO- select inst mod.add(BufferStoreB32(vgpr("AmaxOut"), vgpr("Offset"), sgpr("Dst",4), 0, MUBUFModifiers(offen=True))) mod.addSpaceLine() mod.add(label_end) mod.addSpaceLine() return mod def insertAmaxD(self, kernel): module = Module("AmaxD Output") module.addComment0("AmaxD Output") self.amaxVgprNames = ["Widx", "Offset", "Tmp", "Value"] self.amaxVgprSizes = [1, 1, 1, 1] self.amaxSgprArgNames = ["AddrAmaxOut", "AddressWk", "AddressSy"] self.amaxSgprArgSizes = [2, 2, 2] module.addSpaceLine() module.add(SBarrier()) module.add(self.amax_define_load_res()) module.add(self.amax_intra_wave_reduction(kernel, "middle")) module.add(self.amax_inter_wave_reduction(kernel)) module.add(self.amax_broadcast(kernel)) module.add(self.amax_output_result(kernel)) for i in self.amaxVgprIdxVec: self.vgprPool.checkIn(i) for i in self.amaxSgprIdxVec: self.sgprPool.checkIn(i) return module ######################################## # Activation related ######################################## def initActivationLoop(self, kernel, beta): # Create a suffix and check if the string exists activationLabelSuffix = self.labels.getNameInc( \ "%s"%("_Beta" if beta else "")) activationCDataType = kernel["ProblemType"]["ActivationComputeDataType"] activationType = kernel["ProblemType"]["ActivationType"] activationEndLabel = Label("Activation_End%s"%activationLabelSuffix, "") activationLabelModules = [] activationEnumStrList = [] if kernel["ActivationFuncCall"]: activationLabelModules.append("") activationEnumStrList.append("none") elif (((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel' or kernel["StreamK"] > 0) and kernel["ActivationFused"]) and \ (activationType != 'none'): if activationType in ['all', 'hipblaslt_all']: exportType = ActivationType.Export.GRADONLY if kernel["ProblemType"]["Gradient"] else ActivationType.Export.NORMAL supportedBy = ActivationType.SupportedBy.ALL if activationType == 'all' else ActivationType.SupportedBy.HIPBLASLT activationEnumStrList = ActivationType.getEnumStrList(activationCDataType, supportedBy, exportType=exportType) for _, enumStr in enumerate(activationEnumStrList): activationLabelModule = Label("Activation_%s%s"% (enumStr.capitalize(), activationLabelSuffix), "") activationLabelModules.append(activationLabelModule) else: activationEnumStrList.append(str(kernel["ProblemType"]["ActivationType"]).lower()) else: activationLabelModules.append("") activationEnumStrList.append("none") return activationEndLabel, activationLabelModules, activationEnumStrList def insertActFunctionCallAddrCalc(self, sgprOffset, gwvw, \ toActModuleList, activationEnumStrList, activationLabelList, \ betaIdx = -1): activationLabelModules = activationLabelList[gwvw] module = Module(getActFuncBranchModuleName()) setAddrEndLabel = Label(self.labels.getNameInc("ActivationSetPCAddrEnd"), "") toActModules = deepcopy(toActModuleList[gwvw]) for index, toActModule in enumerate(toActModules): if betaIdx >= 0: toActModule.label = self.labels.getNameInc(toActModule.label + "_beta_%u"%(betaIdx)) if index != 0: enumIndex = ActivationType.getEnumIndex(activationEnumStrList[index]) #module.add(SCmpKEQU32(sgpr("ActivationType"), enumIndex, "activationType == %u"%enumIndex)) module.add(self.getSCMPKInstruction("EQU32", "ActivationType", enumIndex, comment="activationType == %u"%enumIndex)) module.add(SCBranchSCC1(toActModule.getLabelName(), "Branch if true")) for index, activationLabelModule in enumerate(activationLabelModules): toActModule = toActModules[index] module.add(toActModule) with self.allocTmpSgpr(1, tag="insertActFunctionCallAddrCalc_tmpSgprInfo") as tmpSgprInfo: module.addModuleAsFlatItems(SGetPositivePCOffset(sgprOffset, activationLabelModule, tmpSgprInfo)) module.add(SBranch(setAddrEndLabel.getLabelName())) module.add(setAddrEndLabel) return module def insertActivationAfterPacked(self, kernel, activationTypeStr): result = False if kernel["ProblemType"]["UseScaleCD"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0): return result elif ((kernel["ProblemType"]["ActivationType"] != 'none') and \ ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["GlobalSplitUAlgorithm"] == 'MultipleBufferSingleKernel' or kernel["StreamK"] > 0) and kernel["ActivationFused"]): if kernel["ActivationFuncCall"]: return (kernel["ProblemType"]["ActivationComputeDataType"] == kernel["ProblemType"]["DestDataType"]) elif kernel["ProblemType"]["DestDataType"].isBFloat16() and (activationTypeStr == 'abs'): result = True elif kernel["ProblemType"]["DestDataType"].isHalf() and \ ((activationTypeStr == 'abs') or (activationTypeStr == 'relu')): result = True elif kernel["ProblemType"]["ActivationComputeDataType"] == kernel["ProblemType"]["DestDataType"]: result = True return result def getActivationDestDataType(self, kernel, activation, activationTypeStr: str, gwvw, \ elementSumIdxIn, elementSumIdxOut, tmpVgpr, tmpSgpr): module = Module("ActivationAfterPack") for vi in range(0, gwvw): sumIdxVIn = elementSumIdxIn + vi sumIdxVOut = elementSumIdxOut + vi if kernel["ProblemType"]["DestDataType"].isHalf() or \ kernel["ProblemType"]["DestDataType"].isBFloat16(): if kernel["ProblemType"]["HighPrecisionAccumulate"]: # Generate single f16 code if edge is detected. if ((vi + 1) == gwvw) and ((gwvw % 2) == 1): activation.setUsePK(False) # Original packed route elif vi%2 == 1: assert (gwvw % 2 == 0) else: continue vgprIn = elementSumIdxIn + vi//2 vgprOut = elementSumIdxOut + vi//2 else: if (sumIdxVIn % 2 != 0): continue vgprIn = sumIdxVIn // 2 vgprOut = sumIdxVOut // 2 elif kernel["ProblemType"]["DestDataType"].isSingle(): vgprIn = sumIdxVIn vgprOut = sumIdxVOut elif kernel["ProblemType"]["DestDataType"].isDouble(): vgprIn = sumIdxVIn * 2 vgprOut = sumIdxVOut * 2 elif kernel["ProblemType"]["DestDataType"].isInt32(): vgprIn = sumIdxVIn vgprOut = sumIdxVOut else: raise RuntimeError("Unsupported data type %s for activation vgpr index."%str(self.states.kernel["ProblemType"]["DestDataType"])) # Here we still use DestDataType cause the data is ready to be written to global actModule = activation.getModule(self.states.kernel["ProblemType"]["DestDataType"], activationTypeStr, vgprIn, vgprOut) module.add(activation.assignGpr(actModule, tmpVgpr, tmpSgpr)) activation.setUsePK(True) return module def getActivationActivationComputeType(self, kernel, activation, activationTypeStr: str, gwvw, \ elementSumIdxIn, elementSumIdxOut, tmpVgpr, tmpSgpr, satInt8=False, enableValuCPrefix=False): module = Module("ActivationBeforePack") if satInt8: activation.setSaturationForInt8(True) if enableValuCPrefix: # The register is from ValuC allocation instead of allocated by itself activation.setVgprPrefixFormat("ValuC+%u") for vi in range(0, gwvw): vgprIn = elementSumIdxIn + vi vgprOut = elementSumIdxOut + vi actModule = activation.getModule(kernel["ProblemType"]["ActivationComputeDataType"], activationTypeStr, vgprIn, vgprOut) module.add(activation.assignGpr(actModule, tmpVgpr, tmpSgpr)) activation.setSaturationForInt8(False) activation.setVgprPrefixFormat("") return module ############################################################################## # PAP tile-identity group. # # The StreamK PAP handoff maps WorkGroup*/StreamKLocal* to the next # persistent tile to compute next-tile addresses, but current NLL/tail code # resumes immediately after the prefetch. Keep tile identity borrowed. ############################################################################## def papTileIdentityNames(self, kernel): names = [ "WorkGroup0", "WorkGroup1", "WorkGroup2", ] # DP-only tiles are always full: StreamKLocalStart/End are constant # (0 / ItersPerTile) and the next-tile setup recomputes the same values, # so they need no checkpoint/restore. (DP-only PAP: skip unneeded state.) if not kernel["StreamKForceDPOnly"]: names.append("StreamKLocalStart") names.append("StreamKLocalEnd") if len(kernel["SpaceFillingAlgo"]): names.append("StreamKTileID") return names @contextmanager def allocPapTileIdentitySgprs(self, kernel): names = self.papTileIdentityNames(kernel) with self.allocTmpSgpr(len(names), alignment=1, tag="PAP tile identity") as papTileIdentitySgpr: yield {name: papTileIdentitySgpr.idx + i for i, name in enumerate(names)} def papCheckpointCurrentTileIdentity(self, kernel, prevTile): module = Module("papCheckpointCurrentTileIdentity") for name in self.papTileIdentityNames(kernel): module.add(SMovB32(dst=sgpr(prevTile[name]), src=sgpr(name), comment="checkpoint %s for PAP restore" % name)) return module def papRestoreCurrentTileIdentity(self, kernel, prevTile): module = Module("papRestoreCurrentTileIdentity") # PAP temporarily maps WorkGroup*/StreamKLocal* to the next persistent tile # so it can issue the first PGR early. Restore the current tile for the # remaining NLL/tail code; StreamKIter already points at the next chunk. for name in self.papTileIdentityNames(kernel): module.add(SMovB32(dst=sgpr(name), src=sgpr(prevTile[name]), comment="restore current %s after PAP" % name)) return module ############################################################################## # Prefetch across persistent: prefetch next tile's data during the NLL. # # Durable output from this sequence is limited to the issued first-PGR loads, # SkPrefetchPrimed, and saved PAP TDM/DTL LDS-bank bits. The setup below borrows tile # identity, descriptor, stagger, and TDM/DTL descriptor state and restores it # before current-tile code observes those registers again. ############################################################################## def prefetchAcrossPersistent(self, kernel, tensorParametersA, tensorParametersB, skipBarrier=False): module = Module("prefetchAcrossPersistent") if not self.isPrefetchAcrossPersistentEnabled(kernel): return module skipLabel = Label(self.labels.getNameInc("SK_SkipNllPAP"), "") # Parallel reduction (no synchronizer): WGs do not advance across tiles module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Parallel reduction: skip PAP")) module.add(SCBranchSCC1(labelName=skipLabel.getLabelName(), comment="")) module.add(SCmpGeU32(src0=sgpr("StreamKIter"), src1=sgpr("StreamKIterEnd"), comment="No next persistent iteration")) module.add(SCBranchSCC1(labelName=skipLabel.getLabelName(), comment="")) if not skipBarrier: module.add(SBarrier(comment="PAP: sync before next-tile prefetch")) skComponent = Component.StreamK.find(self) with self.allocPapTileIdentitySgprs(kernel) as prevTile: module.add(self.papCheckpointCurrentTileIdentity(kernel, prevTile)) module.add(skComponent.prefetchAcrossPersistentSetupNextTile(self, kernel, tensorParametersA, tensorParametersB, skipLroReset=True)) if kernel["enableTDMA"] and kernel["enableTDMB"]: module.add(self.papTdmUpdateDescriptor(kernel, tensorParametersA, tensorParametersB)) if kernel["ProblemType"]["MXBlockA"] and kernel["ProblemType"]["MXBlockB"]: module.add(self.papTdmUpdateDescriptor(kernel, tensorParametersA["MX"], tensorParametersB["MX"])) loopCounterName = self.loopCounterName(kernel, self.states.unrollIdx) # DP-only: LoopCounter is constant ItersPerTile and OrigLoopCounter is a # per-problem constant, so calculateLoopNumIter recomputes the same values # (idempotent) and PAP never runs on the last tile. Skip the 2-VGPR # checkpoint/restore. (DP-only PAP saving.) snapshotLoopCounter = not kernel["StreamKForceDPOnly"] if snapshotLoopCounter: prevLoopVgpr = self.vgprPool.checkOutAligned(2, 1, "PAP loop counters") module.add(VMovB32(dst=vgpr(prevLoopVgpr), src=sgpr(loopCounterName), comment="checkpoint LoopCounter for PAP restore")) module.add(VMovB32(dst=vgpr(prevLoopVgpr + 1), src=sgpr("OrigLoopCounter"), comment="checkpoint OrigLoopCounter for PAP restore")) module.add(self.calculateLoopNumIter(kernel, tensorParametersA, tensorParametersB, self.states.unrollIdx)) module.add(self.setupPrefetchAcrossPersistentLoads(kernel, tensorParametersA, tensorParametersB, isOptNLL=True)) if snapshotLoopCounter: module.add(VReadfirstlaneB32(dst=sgpr(loopCounterName), src=vgpr(prevLoopVgpr), comment="restore LoopCounter after PAP")) module.add(VReadfirstlaneB32(dst=sgpr("OrigLoopCounter"), src=vgpr(prevLoopVgpr + 1), comment="restore OrigLoopCounter after PAP")) self.vgprPool.checkIn(prevLoopVgpr) if kernel["enableTDMA"] and kernel["enableTDMB"]: module.add(self.papTdmSaveLdsBank(kernel)) module.add(self.papRestoreCurrentTileIdentity(kernel, prevTile)) if kernel["enableTDMA"] and kernel["enableTDMB"] and not kernel["NoTailLoop"]: module.add(self.papTdmUpdateDescriptor(kernel, tensorParametersA, tensorParametersB, preservePapBank=False)) if kernel["ProblemType"]["MXBlockA"] and kernel["ProblemType"]["MXBlockB"]: module.add(self.papTdmUpdateDescriptor(kernel, tensorParametersA["MX"], tensorParametersB["MX"], preservePapBank=False)) module.add(skipLabel) return module ############################################################################## # Function End ############################################################################## def functionEnd(self, kernel, addLabel=True): imod = Module() loopComponent = Component.PersistentLoop.find(self) imod.add(loopComponent.closePersistentLoop(self, kernel)) if addLabel: imod.add(Label("KernelEnd", "")) skComponent = Component.StreamK.find(self) imod.add(skComponent.kernelEnd(self, kernel)) # TODO- refine this part, put outside of this function if kernel["ProblemType"]["OutputAmaxD"]: imod.add(self.insertAmaxD(kernel)) imod.add(SEndpgm(comment="Kernel End")) return imod ############################################################################## # waitcnt code for DirectToVgpr ############################################################################## def getWaitcntCodeForDirectToVgpr(self, kernel, tensorParametersA, tensorParametersB, localWriteEndIter, u, isNLL=False, beforeBarrier=False, NLLlast=False, oddLast=False): module = Module("DTV wait") # generate wait for DTV # TODO: add precise wait count logic for DTVA + DTVB if kernel["DirectToVgprA"] and kernel["DirectToVgprB"]: needToWait = 0 waitComment = "global read wait for DirectToVgpr" module.add(self.getWaitcntCodeForDTVSub(kernel, tensorParametersA, tensorParametersB, needToWait, waitComment)) elif kernel["DirectToVgprA"] ^ kernel["DirectToVgprB"]: numGlobalReadA = kernel["NumLoadsPerpendicularA"] * kernel["NumLoadsCoalescedA"] numGlobalReadB = kernel["NumLoadsPerpendicularB"] * kernel["NumLoadsCoalescedB"] numReadsIterCoalesced = self.states.numReadsIterCoalescedA if kernel["DirectToVgprA"] else self.states.numReadsIterCoalescedB numReadsIterCoalesced *= kernel["NumLoadsCoalescedA"] if kernel["reorderGRInstForDTVA"] and \ kernel["NumLoadsCoalescedA"] % 2 == 0 else 1 numReadsIterCoalesced *= kernel["NumLoadsCoalescedB"] if kernel["reorderGRInstForDTVB"] and \ kernel["NumLoadsCoalescedB"] % 2 == 0 else 1 numGlobalReadNonDTV = 0 if not kernel["DirectToVgprA"]: numGlobalReadNonDTV += numGlobalReadA if not kernel["DirectToVgprB"]: numGlobalReadNonDTV += numGlobalReadB numGlobalReadDTV = numGlobalReadA + numGlobalReadB - numGlobalReadNonDTV waitComment = "global read wait for DirectToVgpr" # delay DirectToVgpr global read (from previous iteration) which is not referred yet (do not delay in beforeBarrier case) numRegsIn1set = (numGlobalReadDTV * numReadsIterCoalesced) // kernel["LoopIters"] numSet = (u + numReadsIterCoalesced) // numReadsIterCoalesced numSetMod = (u + numReadsIterCoalesced) % numReadsIterCoalesced if (not beforeBarrier) and numSetMod > 0: # if mod > 0, wait is already done by mod == 0 case and no need to wait for same set of global read return module needToWait = numGlobalReadDTV - numSet * numRegsIn1set if isNLL: # NLL case, no (non DTV) global load A, B in no load loop. Reset numGlobalReadNonDTV numGlobalReadNonDTV = 0 if kernel["PrefetchGlobalRead"] >= 2: # PGR=2 case, add numGlobalReadNonDTV for second set of prefetch needToWait += numGlobalReadNonDTV if u > 0: # count number of global read for i < u count = 0 for i in range(u): globalReadStr = ' '.join([str(x) for x in self.codes.perIterGlobalRead[i].flatitems()]) count += countGlobalRead(self.codes.perIterGlobalRead[i]) # PGR=2 case, global read is in LocalWriteCode count += countGlobalRead(self.codes.perIterLocalWrite[i][1]) needToWait += count if u == localWriteEndIter + 1 and beforeBarrier: # beforeBarrier case, reduce the amount of non-Vgpr global read needToWait -= numGlobalReadNonDTV # adjustment for oddLast # oddLast case or ScheduleIterAlg < 3 case, ignore all of above and set 0 if oddLast or kernel["_ScheduleIterAlg"] < 3: needToWait = 0 # generate waitcnt code module.add(self.getWaitcntCodeForDTVSub(kernel, tensorParametersA, tensorParametersB, needToWait, waitComment)) return module ############################################################################## # waitcnt code for PrefetchGlobalRead ############################################################################## def getWaitcntCodeForPGR(self, kernel, tensorParametersA, tensorParametersB, comment): module = Module("PGR wait") # generate wait if (kernel["DirectToVgprA"] and kernel["DirectToVgprB"]) or \ (kernel["DirectToLdsA"] and kernel["DirectToLdsB"] and (not kernel["ProblemType"]["Sparse"] or kernel["DirectToLdsMetadata"]) and kernel["PrefetchGlobalRead"] >= 2): # DTVA and DTVB, or ((DTLA + DTLB) + PGR>=2) case, wait code for prefetch is unnecessary # (because wait is for local write code in PGR>=2 case) return module count = 0 if kernel["DirectToVgprA"]: # increase vlcnt for DTVA (no need to wait for DTVA global load) count += kernel["NumLoadsPerpendicularA"] * kernel["NumLoadsCoalescedA"] if kernel["DirectToVgprB"]: # increase vlcnt for DTVB (no need to wait for DTVB global load) count += kernel["NumLoadsPerpendicularB"] * kernel["NumLoadsCoalescedB"] module.add(self.getWaitcntCodeForDTVSub(kernel, tensorParametersA, tensorParametersB, count, comment)) return module ############################################################################## # waitcnt code for DirectToVgpr Sub function ############################################################################## def getWaitcntCodeForDTVSub(self, kernel, tensorParametersA, tensorParametersB, count, waitComment): return self._wait(kernel, tensorParametersA, tensorParametersB, 0, -1, -1, waitComment, skipGlobalReadInst = count) ############################################################################## # isSwapGlobalReadOrderForDtvOrDtl ############################################################################## def isSwapGlobalReadOrderForDtvOrDtl(self, kernel, prefetch1=False): # check the conditions to swap global read order (from A, B to B, A) if kernel["DirectToVgprA"] and kernel["DirectToVgprB"]: # DTVA+DTVB case, no swap return False elif kernel["DirectToVgprA"]: # DTVA. Swap with the following conditions # - PGR=2 and prefetch1 # - PGR=1 if kernel["PrefetchGlobalRead"] >= 2 and prefetch1: return True elif kernel["PrefetchGlobalRead"] == 1: return True elif kernel["DirectToVgprB"]: # DTVB. Swap with the following conditions # - PGR=2 and not prefetch1 if kernel["PrefetchGlobalRead"] >= 2 and (not prefetch1): return True elif (not kernel["DirectToLdsA"]) and kernel["DirectToLdsB"]: # (no DTLA) + DTLB (need to put DTLB first) return True elif kernel["SwapGlobalReadOrder"] == 1: return True return False ############################################################################## # WaveSplitK Reduction ############################################################################## def waveSplitKReduction(self, kernel): module = Module("WaveSplitK Reduction") module.addComment2("WaveSplitK Reduction") # the last instruction may use ValuC VGPR (ex. no tailloop case), need 2 snop module.add(SNop(waitState=0, comment="1 wait states")) module.add(SNop(waitState=0, comment="1 wait states")) accK = 1 while accK < kernel["NumWaveSplitK"]: # if only one VGPR, need to insert one SNOP in between if (kernel["ThreadTile1"] == 1 and kernel["ThreadTile0"] == 1) and accK != 1: module.add(SNop(waitState=0, comment="1 wait states")) for block1 in range(0, kernel["ThreadTile1"]): for block0 in range(0, kernel["ThreadTile0"]): vgprIdx = block0 + block1 * kernel["ThreadTile0"] if accK <= 8: module.add(VAddF32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr("ValuC+%d"%vgprIdx), src1=vgpr("ValuC+%d"%vgprIdx), dpp=DPPModifiers(bound_ctrl=0, row_shr=accK), comment="Reduce %s elements"%(accK*2) )) else: module.add(VAddF32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr("ValuC+%d"%vgprIdx), src1=vgpr("ValuC+%d"%vgprIdx), dpp=DPPModifiers(bound_ctrl=0, row_bcast=accK-1), comment="Reduce %s elements"%(accK*2) )) accK *= 2 module.addSpaceLine() return module def simdSpecDispatch(self, kernel, numCodePath, nta=0, ntb=0): strNta = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTA%s"%nta strNtb = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTB%s"%ntb module = Module() loopLabelBegin = [] loopLabelSkipBegin = [] loopChar = self.states.indexChars[kernel["ProblemType"]["IndicesSummation"][self.states.unrollIdx]] if numCodePath == 1: module.add(Label("LoopBegin%s_0%s%s"%(loopChar, strNta, strNtb), "" )) module.add(MacroInstruction(name="MAINLOOP%s%s"%(strNta, strNtb), args=[0])) module.add(SCBranchSCC0(labelName="label_LoopBegin%s_0%s%s"%(loopChar, strNta, strNtb), comment="" )) module.add(Label("LoopEnd%s%s%s"%(loopChar, strNta, strNtb), "" )) return module module.addComment0("SIMD specialized dispatch") for l in range(numCodePath): loopLabelBegin.append(Label("LoopBegin%s_%u%s%s"%(loopChar, l, strNta, strNtb), "", alignment=16)) loopLabelSkipBegin.append(Label("LoopSkipBegin%s_%u%s%s"%(loopChar, l, strNta, strNtb), "" )) loopLabelEnd = Label("LoopEnd%s%s%s"%(loopChar, strNta, strNtb), "" ) tmpSgpr = self.sgprPool.checkOut(1, tag="simdSpecDispatch_tmpSgpr") numbits = 1 if numCodePath == 2 else 2 module.add(SGetRegB32(dst=sgpr(tmpSgpr), src="hwreg(HW_REG_HW_ID, 4, %u)"%numbits)) tmpSgpr1 = self.sgprPool.checkOutAligned(3, 2, tag="simdSpecDispatch_tmpSgpr1") sgprPC = ContinuousRegister(tmpSgpr1, 3) for l in range(numCodePath): module.addComment0("Branch to code-path %u logic"%l) module.add(SCmpEQU32(src0=sgpr(tmpSgpr), src1=hex(l), comment="" )) #module.add(SCBranchSCC1(labelName=loopLabelBegin[l].getLabelName(), comment="" )) module.add(SCLongBranchScc1(loopLabelBegin[l], sgprPC, \ loopLabelSkipBegin[l].getLabelName(), \ loopLabelBegin[l].getLabelName())) self.sgprPool.checkIn(tmpSgpr1) self.sgprPool.checkIn(tmpSgpr) for l in range(numCodePath): module.addComment0("SIMD %u code-path"%l) module.add(loopLabelBegin[l]) module.add(MacroInstruction(name="MAINLOOP%s%s"%(strNta, strNtb), args=[l])) module.add(SCBranchSCC0(labelName=loopLabelBegin[l].getLabelName(), comment="" )) tmpSgpr1 = self.sgprPool.checkOutAligned(2, 2, tag="simdSpecDispatch_tmpSgpr1_2") sgprPC = ContinuousRegister(tmpSgpr1, 3) #module.add(SBranch(labelName=loopLabelEnd.getLabelName(), comment="" )) module.add(SLongBranchPositive(loopLabelEnd, sgprPC)) self.sgprPool.checkIn(tmpSgpr1) # close unrolled loop endStr = "" module.addComment2("Unrolled Loop - End%s"%(endStr)) module.add(loopLabelEnd) return module ############################################################################## # Wrappers ############################################################################## ############################################################################## # longBranchScc0 - 32 bit offset # Conditional branch to label when SCC == 0 # Use when erroring out "invalid operand due to label > SIMM16" ############################################################################## def longBranchScc0(self, label: Label, posNeg: int=0, tmpSgprInfo: Optional[ContinuousRegister]=None, comment=""): if tmpSgprInfo: return SCLongBranchScc0(label, tmpSgprInfo, \ self.labels.getUniqueNamePrefix("NoBranch"), \ self.labels.getUniqueNamePrefix("Positive"), posNeg, comment) else: with self.allocTmpSgpr(3, tag="longBranchScc0_tmpSgprInfo") as tmpSgprInfo: return SCLongBranchScc0(label, tmpSgprInfo, \ self.labels.getUniqueNamePrefix("NoBranch"), \ self.labels.getUniqueNamePrefix("Positive"), posNeg, comment) ############################################################################## # longBranchScc1 - 32 bit offset # Conditional branch to label when SCC == 1 # Use when erroring out "invalid operand due to label > SIMM16" ############################################################################## def longBranchScc1(self, label: Label, posNeg: int=0, tmpSgprInfo: Optional[ContinuousRegister]=None, comment=""): if tmpSgprInfo: return SCLongBranchScc1(label, tmpSgprInfo, \ self.labels.getUniqueNamePrefix("NoBranch"), \ self.labels.getUniqueNamePrefix("Positive"), posNeg, comment) else: with self.allocTmpSgpr(3, tag="longBranchScc1_tmpSgprInfo") as tmpSgprInfo: return SCLongBranchScc1(label, tmpSgprInfo, \ self.labels.getUniqueNamePrefix("NoBranch"), \ self.labels.getUniqueNamePrefix("Positive"), posNeg, comment) ############################################################################## # longBranchVccnz - 32 bit offset # Conditional branch to label when VCC != 0 # Use when erroring out "invalid operand due to label > SIMM16" ############################################################################## def longBranchVccnz(self, label: Label, posNeg: int=0, tmpSgprInfo: Optional[ContinuousRegister]=None, comment=""): if tmpSgprInfo: return SCLongBranchVccnz(label, tmpSgprInfo, \ self.labels.getUniqueNamePrefix("NoBranch"), \ self.labels.getUniqueNamePrefix("Positive"), posNeg, comment) else: with self.allocTmpSgpr(3, tag="longBranchVccnz_tmpSgprInfo") as tmpSgprInfo: return SCLongBranchVccnz(label, tmpSgprInfo, \ self.labels.getUniqueNamePrefix("NoBranch"), \ self.labels.getUniqueNamePrefix("Positive"), posNeg, comment) def sMagicDivWrapper(self, dest, dividend, magicNumber, magicShift): tmpVgpr = self.vgprPool.checkOut(2, tag="sMagicDivWrapper_tmpVgpr") tmpVgprRes = ContinuousRegister(tmpVgpr, 2) module = sMagicDiv(dest, self.states.asmCaps["HasSMulHi"], dividend, magicNumber, magicShift, tmpVgprRes) self.vgprPool.checkIn(tmpVgpr) return module def s_mul_u64_u32 (self, dst0, dst1, src0, src1, tmpVgpr=None, comment=""): vtmp0 = self.vgprPool.checkOut(2, tag="s_mul_u64_u32_vtmp0") if tmpVgpr == None else tmpVgpr tmpVgprRes = ContinuousRegister(vtmp0, 2) module = SMulInt64to32(dst0, dst1, src0, src1, tmpVgprRes, bool(self.states.asmCaps["HasSMulHi"]), False, comment) if tmpVgpr == None: self.vgprPool.checkIn(vtmp0) return module def s_mul_i64_i32 (self, dst0, dst1, src0, src1, comment): vtmp0 = self.vgprPool.checkOut(2, tag="s_mul_i64_i32_vtmp0") tmpVgprRes = ContinuousRegister(vtmp0, 2) module = SMulInt64to32(dst0, dst1, src0, src1, tmpVgprRes, bool(self.states.asmCaps["HasSMulHi"]), True, comment) self.vgprPool.checkIn(vtmp0) return module def s_mul_i64_i32_u32 (self, dst0, dst1, src0, src1, comment): module = Module("S_MUL_I64_I32_U32") vtmp0 = self.vgprPool.checkOut(2, tag="s_mul_i64_i32_u32_vtmp0") negativeLabel = Label((self.labels.getUniqueNamePrefix("Negative")), comment="") multiplydoneLabel = Label((self.labels.getUniqueNamePrefix("MultiplyDone")), comment="") module.add(SCmpGeI32(src0, 0)) module.add(SCBranchSCC0(labelName=negativeLabel.getLabelName(), comment="")) tmpVgprRes = ContinuousRegister(vtmp0, 2) module.add(SMulInt64to32(dst0, dst1, src0, src1, tmpVgprRes, bool(self.states.asmCaps["HasSMulHi"]), False, comment)) module.add(SBranch(labelName=multiplydoneLabel.getLabelName(), comment="")) module.add(negativeLabel) module.add(SAbsI32(src0, src0, comment="")) module.add(SMulInt64to32(dst0, dst1, src0, src1, tmpVgprRes, bool(self.states.asmCaps["HasSMulHi"]), False, comment)) module.add(SXorB32(dst0, dst0, hex(0xFFFFFFFF), comment="")) module.add(SXorB32(dst1, dst1, hex(0xFFFFFFFF), comment="")) module.add(SAddU32(dst0, dst0, hex(0x1), comment="")) module.add(SAddCU32(dst1, dst1, 0, comment="")) module.add(multiplydoneLabel) self.vgprPool.checkIn(vtmp0) return module def getCmpAssert(self, function, val0, val1, cookie=-1): scratchVgpr = self.vgprPool.checkOut(2, tag="getCmpAssert_scratchVgpr") function(val0, val1, scratchVgpr, cookie) self.vgprPool.checkIn(scratchVgpr) def getMultipleB32Assert(self, sval, multiple2, cookie=-1): scratchVgpr = self.vgprPool.checkOut(2, tag="getMultipleB32Assert_scratchVgpr") self.asmAssert.multiple_b32(sval, multiple2, scratchVgpr, cookie) self.vgprPool.checkIn(scratchVgpr) def getVectorDiffAssert(self, v0, v1, expectedScalarDiff, cookie=-1): cmpvtmp = self.vgprPool.checkOut(1, tag="getVectorDiffAssert_cmpvtmp") vtmp = self.vgprPool.checkOut(2, tag="getVectorDiffAssert_vtmp") self.asmAssert.assert_vector_diff(v0, v1, expectedScalarDiff, cmpvtmp, vtmp, cookie) self.vgprPool.checkIn(vtmp) self.vgprPool.checkIn(cmpvtmp) #to-do, tmp solution, need to move to instruction.py def getSCMPKInstruction(self, instOP, s0, simm16, comment="") -> Module: imodscmpk = Module() Inst0 = SCmpKEQU32 Inst1 = SCmpEQU32 if instOP == "EQU32": Inst0 = SCmpKEQU32 Inst1 = SCmpEQU32 elif instOP == "LGU32": Inst0 = SCmpKLGU32 Inst1 = SCmpLgU32 elif instOP == "GEU32": Inst0 = SCmpKGeU32 Inst1 = SCmpGeU32 elif instOP == "GTU32": Inst0 = SCmpKGtU32 Inst1 = SCmpGtU32 else: assert 0,"getSCMPKInstruction failed" if self.states.asmCaps["HasSCMPK"]: imodscmpk.add(Inst0(sgpr(s0), simm16, comment=comment)) else: tmpScmp = self.sgprPool.checkOut(1, tag="getSCMPKInstruction_tmpScmp", preventOverflow=False) imodscmpk.add(SMovB32(sgpr(tmpScmp), simm16)) imodscmpk.add(Inst1(src0=sgpr(s0), src1=sgpr(tmpScmp), comment=comment)) self.sgprPool.checkIn(tmpScmp) return imodscmpk def dump(self, vgprStore): return self.dumpData.dumpVgpr(vgprStore, self.labels.getUniqueName()) def dumpSgpr(self, sgprStore): tmp = ContinuousRegister(idx=self.vgprPool.checkOut(1, tag="dumpSgpr_tmp"), size=1) module = self.dumpData.dumpSgpr(sgprStore, tmp, self.labels.getUniqueName()) self.vgprPool.checkIn(tmp.idx) return module def dumpLDS(self, kernel, startU, numU): tmp = ContinuousRegister(idx=self.vgprPool.checkOut(2, tag="dumpLDS_tmp"), size=2) module = self.dumpData.dumpLds(startU, numU, tmp, self.states.bpeAB, kernel["NumThreads"], \ self.labels.getUniqueName()) self.vgprPool.checkIn(tmp.idx) return module def globalOffset(self, kernel, tP, tc, args, comment=""): module = Module("globalOffset") ######################################## # Global Offsets ######################################## # justOffset32 means we should only write the 32-bit offset # This is used in Buffer addressing modes. # Flat addressing modes expect the GLOBAL_OFFSET to initialize a full 64-bit address # GOList = [ \ # ("C", list(range(0, kernel["ProblemType"]["NumIndicesC"])), kernel["BufferStore"], None, False), \ # ("A", kernel["ProblemType"]["IndexAssignmentsA"], kernel["BufferLoad"], tPA, False), \ # ("B", kernel["ProblemType"]["IndexAssignmentsB"], kernel["BufferLoad"], tPB, False) ] # if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]: # GOList.append(("Metadata", kernel["ProblemType"]["IndexAssignmentsMetadata"], kernel["BufferLoad"], tPM, False)) # if kernel["ProblemType"]["SwizzleTensorA"]: # GOList.append(("A", kernel["ProblemType"]["IndexAssignmentsA"], kernel["BufferLoad"], tPA, True)) # if kernel["ProblemType"]["SwizzleTensorB"]: # GOList.append(("B", kernel["ProblemType"]["IndexAssignmentsB"], kernel["BufferLoad"], tPB, True)) # for (tc, indices, justOffset32, tP, isSwizzled) in GOList: isSwizzled = False if tP == None else tP["isSwizzled"] indices = kernel["ProblemType"]["NumIndicesC"] if tc == "C" else kernel["ProblemType"]["IndexAssignments%s"%tc] justOffset32 = kernel["BufferStore"] if tc == "C" else kernel["BufferLoad"] aidx = 1 if args[0].isdigit(): tmp = int(args[0]) addrVgpr0 = vgpr(tmp) addrVgpr1 = vgpr(tmp+1) addrVgpr01 = vgpr(tmp,2) else: tmp = args[0][4:] addrVgpr0 = vgpr(tmp) addrVgpr1 = vgpr(tmp+"+1") addrVgpr01 = vgpr(tmp,2) # function name and comment suffix_tc = tc + "_SWIZZLED" if isSwizzled else tc module.addComment1("Global Offset %s"%suffix_tc) numDim = len(indices) idxChars = [] for i in indices: idxChars.append(self.states.indexChars[i]) # macro declaration calcDims = [] # dimensions which are participating in the address calc (ignores other summation) mirrorSumDims = [] macroArgs = [] offsets = {} for i in range(0, numDim): if tc == 'C': useInitialStrides = kernel["ProblemType"]["UseInitialStridesCD"] idxChar = self.states.indexChars[i] else: useInitialStrides = kernel["ProblemType"]["UseInitialStridesAB"] idxChar = self.states.indexChars[tP['ia'][i]] # tile index or unroll vgpr or summation # other summation (other than unroll) are included in the GLOBAL_OFFSET macro but not used in address calc if tc in ('A','C') and indices[i] == kernel["ProblemType"]["Index0"] \ or tc in ('B','C', "Metadata") and indices[i] == kernel["ProblemType"]["Index1"] \ or indices[i] == kernel["ProblemType"]["IndexUnroll"]: # macroArgs.append("vgprOffset%s:req" % idxChars[i]) offsets[i] = int(args[aidx]) if args[aidx].isdigit() else args[aidx][4:] aidx += 1 calcDims.append(i) elif indices[i] in kernel["ProblemType"]["IndicesSummation"]: # other summation index (not unroll) if tc in ('A', 'B', "Metadata") and indices[i] in kernel["ProblemType"]["MirrorDims%s" % tc]: mirrorSumDims.append(i) continue else: # other batch or free index if isPackedIndex(kernel, indices[i]): calcDims.append(i) # macroArgs.append("vgprOffset%s:req" % idxChars[i]) offsets[i] = int(args[aidx]) if args[aidx].isdigit() else args[aidx][4:] aidx += 1 elif not justOffset32: # buffer/justOffset32 scalars are included in SRD not the offset, so skip here calcDims.append(i) # macroArgs.append("sgprOffset%s:req" % idxChars[i]) offsets[i] = int(args[aidx]) if args[aidx].isdigit() else args[aidx][4:] aidx += 1 if args[aidx].isdigit(): tmp = int(args[aidx]) tmpVgpr0 = vgpr(tmp) tmpVgpr1 = vgpr(tmp+1) tmpVgpr2 = vgpr(tmp+2) else: tmp = args[aidx][4:] tmpVgpr0 = vgpr(tmp) tmpVgpr1 = vgpr(tmp+"+1") tmpVgpr2 = vgpr(tmp+"+2") aidx += 1 # Each index may be skipped, scaled by stride, or unscaled # If destLo is unset, no accumulation is necessary. # if the first index (i==0) is unscaled (UseInitialStrides), # it can be combined at the next update or moved at end # (if there is no next update) pendingOffset = None # offset pending for accumulation offsetIsVgpr = False # True if the source is VGPR ; False if SGPR destLo = None # true for first addr calc. In this case, we can directly write addr # rather than accumulating through a tmp writeDirectToAddr = 1 # mirror other summation indices for i in mirrorSumDims: if writeDirectToAddr: dest = addrVgpr0 needAdd = 0 # don't need add since writing address directly. writeDirectToAddr = 0 else: dest = tmpVgpr0 needAdd = 1 module.add(VSubU32(dst=dest, \ src0=sgpr("Size%s"%INDEX_CHARS[indices[i]]), \ src1=1, \ comment="mirror %s%s 1"%(tc, INDEX_CHARS[indices[i]]))) module.add(VMulLOU32(dst=dest, \ src0=dest, \ src1=self.strideRef(tc, indices[i]), \ comment="mirror %s%s 2"%(tc, INDEX_CHARS[indices[i]]))) if needAdd: writeDirectToAddr = 0 # safety net, once we write address can't directly overwrite it later destLo = addrVgpr0 destHi = addrVgpr1 srcLo = pendingOffset if pendingOffset else destLo srcHi = 0 if pendingOffset else destHi module.add(VAddCOU32(dst=destLo, \ dst1=VCC(), \ src0=srcLo, \ src1=vgprTmp0, \ comment="accumulate %s lower"%idxChar)) for i in calcDims: # should have eliminated these above idx = indices[i] isMirrorIdx = tc in ('A', 'B', "Metadata") and idx in kernel["ProblemType"]["MirrorDims%s" % tc] assert not (idx in kernel["ProblemType"]["IndicesSummation"] and idx != kernel["ProblemType"]["IndexUnroll"]) if indices[i] == kernel["ProblemType"]["Index0"] \ or indices[i] == kernel["ProblemType"]["Index1"] \ or indices[i] == kernel["ProblemType"]["IndexUnroll"]: offsetIsVgpr = True # other c index sgpr (free or batch) elif indices[i] < kernel["ProblemType"]["NumIndicesC"]: if isPackedIndex(kernel, indices[i]): offsetIsVgpr = True else: offsetIsVgpr = False else: assert(0) # no other type allowed # offset is VGPR or SGPR string to use for the offset if offsetIsVgpr: offset = vgpr(offsets[i]) else: offset = sgpr(offsets[i]) # macro.addComment0("dim%s pendingOffset=%s offset=%s offsetIsVgpr=%s" \ # % (self.states.indexChars[indices[i]], pendingOffset, offset, offsetIsVgpr)) needAdd = 0 # should be indices[i]?? if i==0 and not useInitialStrides: # slide into next address calc - can do addr = pendingOffset + nextAddrCalc pendingOffset = offset writeDirectToAddr = 0 else: # tile index or unroll vgpr if offsetIsVgpr: if writeDirectToAddr: destLo = addrVgpr0 destHi = addrVgpr1 needAdd = 0 # don't need add since writing address directly. writeDirectToAddr = 0 else: destLo = tmpVgpr0 destHi = tmpVgpr1 needAdd = 1 if isMirrorIdx: module.add(VSubI32( dst=tmpVgpr0, src0=sgpr("Size%s"%INDEX_CHARS[idx]), \ src1=offset, \ comment="mirror %s%s 1"%(tc, INDEX_CHARS[indices[i]]))) module.add(VSubI32(\ dst=tmpVgpr0, src0=tmpVgpr0, \ src1=1, \ comment="mirror %s%s 2"%(tc, INDEX_CHARS[indices[i]]))) offset = tmpVgpr0 # offset * stride if isSwizzled: # Swizzle directly uses the offset base that has already been multiplied by the stride. module.add(VMovB32(dst=destLo, src=offset)) else: module.add(VMulLOU32(dst=destLo, src0=self.strideRef(tc, indices[i]), \ src1=offset, \ comment="mul d%u lower"%i)) if not justOffset32: module.add(VMulHIU32(dst=destHi, src0=self.strideRef(tc, indices[i]), \ src1=offset, \ comment="mul d%u upper"%i)) else: # offset is SGPR: assert not isMirrorIdx if not justOffset32: # buffer mode (aka justOffset32) does scalars into SRD not offset module.add(VMovB32(dst=tmpVgpr2, src=offset, \ comment="sgprOffset -> vgprTmp+2")) # offset * stride module.add(VMulLOU32(dst=tmpVgpr0, \ src0=self.strideRef(tc, indices[i]), src1=tmpVgpr2, \ comment="other stride mul d%u lower"%i)) module.add(VMulHIU32(dst=tmpVgpr1, \ src0=self.strideRef(tc, indices[i]), src1=tmpVgpr2, \ comment="mul d%u upper"%i)) needAdd = 1 if needAdd: writeDirectToAddr = 0 # safety net, once we write address can't directly overwrite it later destLo = addrVgpr0 destHi = addrVgpr1 # addr += offset * stride (lo) : accumulate just-computed address term into addr srcLo = pendingOffset if pendingOffset else destLo srcHi = 0 if pendingOffset else destHi module.add(VAddCOU32(dst=destLo, dst1=VCC(), \ src0=srcLo, src1=tmpVgpr0, \ comment="accumulate %s lower"%idxChar)) # addr += offset * stride (hi) if not justOffset32: module.add(VAddCCOU32(dst=addrVgpr1, dst1=VCC(), \ src0=tmpVgpr1, src1=srcHi, src2=VCC(), \ comment="accumulate %s upper"%idxChar)) pendingOffset = None # pendingOffset but never got a chance to apply it, # need to just add an explicit move or add: # this can happen for small-order tensors if pendingOffset != None: destLo = addrVgpr0 if writeDirectToAddr: module.add(VMovB32(dst=destLo, src=offset, comment="setup d0 lower")) if not justOffset32: module.add(VMovB32(dst=addrVgpr1, src=0, comment="d0 upper")) else: module.add(VAddCOU32(dst=destLo, dst1=VCC(), \ src0=destLo, src1=pendingOffset, \ comment="accumulate final pendingOffset")) if tP != None and kernel["BufferLoad"] and self.states.srdShiftLeft[tc]: module.add(VAddU32(dst=addrVgpr0, \ src0=hex(self.states.srdShiftLeft[tc]), \ src1=addrVgpr0, \ comment="add prepad for pointer shift")) module.addComment1("Global Offset %s (end)"%suffix_tc) return module def shiftVgpr6bitFloat(self, tPA, tPB): module = Module("shiftVgpr6bitFloat") tPList = [] tPList.append(tPA) if tPA["glvw"] == 32 else None tPList.append(tPB) if tPB["glvw"] == 32 else None for tP in tPList: tc = tP["tensorChar"] loadWidth = tP["globalReadInstruction"].totalWidth for perp in range(0, tP["nrp"]): for sPerp in range(0, tP["nrpv"]): for para in range(0, tP["nrc"]): for sPara in range(0, tP["nrcv"]//tP["nrcvpi"]): i = sPara + (tP["nrcv"] // tP["nrcvpi"]) * (para + tP["nrc"] * (sPerp + tP["nrpv"] * perp)) g2lIdx = i * (loadWidth + 2) * tP["bpeRatio"] numLoadVectorComp = int(loadWidth * self.states.bpr // tP["bpeGR"]) r = 0 regIdx = 0 while r < numLoadVectorComp: regIdx = r // 4 idx = g2lIdx + regIdx iter = idx / 4 if iter > 0: for shift_cnt in range(0, 3): srcVgpr = f"G2L{tc}+{idx+shift_cnt}" destVgpr = f"G2L{tc}+{idx+shift_cnt-iter}" module.add(VMovB32(dst=vgpr(destVgpr), src=vgpr(srcVgpr), comment="shift vgpr for 6 bits")) r += 16 return module @staticmethod def _tdmIterTileDim1(kernel, tc, du, dtype) -> int: """N-rows-per-iter for a TDM iterate-mode descriptor: how many rows of `bytes_per_row = round(DepthU * bpe)` fit in one LDS pad block.""" lbspp = kernel["LdsBlockSizePerPad%s" % tc] bytes_per_row = int(round(du * dtype.numBytes())) if bytes_per_row == 0 or lbspp % bytes_per_row != 0: raise RuntimeError( f"TDM iterate {tc}: LdsBlockSizePerPad({lbspp}) not a multiple of " f"bytes_per_row({bytes_per_row}=round(DepthU*bpe)).") return lbspp // bytes_per_row def _emitTdmIterateInit(self, mod, kernel, tc, dtype, du, mt, perIssueLoadRowDivisor, descSgprName, strideRefName): comp = TensorDataMoverLoad.find(self) bpe = dtype.numBytes() dss = TensorDataMoverLoad.dataSizeShift(dtype) lbspp = kernel["LdsBlockSizePerPad%s" % tc] pad_bytes = int(round(kernel["LdsPad%s" % tc] * bpe)) tile_dim1 = self._tdmIterTileDim1(kernel, tc, du, dtype) rows_per_il = mt // perIssueLoadRowDivisor if tile_dim1 == 0 or rows_per_il % tile_dim1 != 0: raise RuntimeError( f"TDM iterate {tc}: rows_per_issueLoad({rows_per_il}) not divisible by tile_dim1({tile_dim1}).") iter_count = rows_per_il // tile_dim1 lds_inc = (lbspp + pad_bytes) >> dss if not (0 < iter_count <= 256): raise RuntimeError( f"TDM iterate {tc}: iter_count({iter_count}) outside HW range 1~256 (field encodes n-1).") # global_inc field = tile_dim1 * strideN * bpe >> dss. with self.allocTmpSgpr(2) as tmp: sIter, sGInc = tmp.idx, tmp.idx + 1 strRef = strideRefName() srcArg = strRef if isinstance(strRef, RegisterContainer) else sgpr(strRef) mod.add(SMulI32(sgpr(sGInc), srcArg, tile_dim1, f"TDM iter global_inc = strideN * tile_dim1({tile_dim1})")) if dtype.isFloat4(): mod.add(SLShiftRightB32(sgpr(sGInc), 1, sgpr(sGInc), "fp4 sub-byte: bytes = elements / 2")) mod.add(comp.setIterationEnabled(descSgprName(1), True)) mod.add(comp.setIterationIncrements(descSgprName(2), lds_inc, sGInc)) mod.add(SMovB32(sgpr(sIter), hex(iter_count - 1), f"TDM iter_count = rows_per_il({rows_per_il}) / tile_dim1({tile_dim1}) - 1")) mod.add(comp.setIterations(descSgprName(2), sIter)) def _emitTdmWaveParitySCC(self, module: Module, kernel: Mapping, dstTmpIdx: int, comment: str = "wave parity"): """Leave this wave's parity (bit0 of WaveId) in SCC. WaveIdx is not live inside the loop, so recompute it from Serial. Clobbers sgpr(dstTmpIdx).""" wavelen: int = kernel["WavefrontSize"] module.add(VReadfirstlaneB32(dst=sgpr(dstTmpIdx), src=vgpr("Serial"), comment="get tId")) module.add(SLShiftRightB32(dst=sgpr(dstTmpIdx), shiftHex=ceil(log2(wavelen)), src=sgpr(dstTmpIdx), comment="waveId")) module.add(SBitcmp1B32(src0=sgpr(dstTmpIdx), src1=0, comment=comment)) def tdmSplitLdsBoundary(self, kernel: Mapping, tP: Mapping) -> int: """LDS split boundary (bytes) for the second half of a TDMSplit tile. Assumes the TDMSplit && !MXS && !Sparse precondition, i.e. dim1Divisor == 2.""" tc: str = tP['tensorChar'] ti: int = tP["idx"] mt: int = kernel[f"MacroTile{ti}"] du: int = kernel["DepthU"] bpe: float = tP["bpeGR"] if not tP["isM"] else 1 dim1Divisor = 2 ldsBlockSizePerPad: int = kernel[f"LdsBlockSizePerPad{tc}"] ldsPadSize: int = int(kernel[f"LdsPad{tc}"] * bpe) half = round(mt * du * bpe // dim1Divisor) extraPadSize = half // ldsBlockSizePerPad * ldsPadSize if ldsBlockSizePerPad != 0 and ldsPadSize != 0 else 0 return half + extraPadSize def tdmSplitGlobalInc(self, kernel: Mapping, tP: Mapping): """Return (strideRef, const) for the TDMSplit global split increment (stride * mt*bpe//2). strideRef mirrors the descriptor init: strideRef(tc, ti) for unrolled-major, else strideRef(tc, 3). const is a compile-time integer. Assumes the TDMSplit && !MXS && !Sparse precondition (dim1Divisor == 2).""" tc: str = tP["tensorChar"] ti: int = tP["idx"] unrolledMajor = not tP["tlu"] mt: int = kernel[f"MacroTile{ti}"] bpe = tP["bpeGR"] if not tP["isM"] else 1 strideRef = self.strideRef(tc, ti) if unrolledMajor else self.strideRef(tc, 3) return strideRef, round(mt * bpe) // 2 def _tdmSplitMultiWaveInc(self, module: Module, kernel: Mapping, gIncIdx: int, scratchIdx: int, tPA=None, tPB=None): """Emit the per-wave-parity TDMSplit increments for the multi-wave (aliased AB) path. Even waves load A, odd waves load B, so both the LDS split boundary and the global split increment are per-tensor compile-time values. Recompute them transiently (0 persistent SGPRs): the global increment is written to sgpr(gIncIdx); the LDS increment operand is returned (an immediate int when A == B, otherwise sgpr(scratchIdx)). Clobbers SCC; scratchIdx is reused as scratch, so the caller must consume the returned operand before overwriting it. """ tPA = self.tPA if tPA is None else tPA tPB = self.tPB if tPB is None else tPB ldsBoundA = self.tdmSplitLdsBoundary(kernel, tPA) ldsBoundB = self.tdmSplitLdsBoundary(kernel, tPB) strideRefA, constA = self.tdmSplitGlobalInc(kernel, tPA) strideRefB, constB = self.tdmSplitGlobalInc(kernel, tPB) # The split advances along the tile's free dimension, whose stride is never the # unit-stride (constStride*) index, so strideRef is always a live SGPR here. assert isinstance(strideRefA, RegisterContainer) and isinstance(strideRefB, RegisterContainer), \ "TDMSplit multi-wave split stride must be a live SGPR, not a constStride symbol" strideArgA = strideRefA strideArgB = strideRefB with self.allocTmpSgpr(1, tag="tdmSplitParity") as waveIdTmp: self._emitTdmWaveParitySCC(module, kernel, waveIdTmp.idx, "wave parity (A=even/B=odd)") def selectByParity(dstIdx: int, valB: int, valA: int, comment: str): # dst = SCC ? valB : valA. SALU allows only one literal operand per # instruction, so materialize A first (s_mov leaves SCC intact) and # s_cselect against the B literal. s_mov/s_cselect/s_mul do not touch SCC, # so the parity result stays live across the whole sequence. module.add(SMovB32(sgpr(dstIdx), valA, comment)) module.add(SCSelectB32(sgpr(dstIdx), valB, sgpr(dstIdx), comment)) # All parity selects run while SCC holds the parity result. module.add(SCSelectB32(sgpr(gIncIdx), strideArgB, strideArgA, "split stride = parity ? B : A")) if constA == constB: module.add(SMulI32(sgpr(gIncIdx), sgpr(gIncIdx), constA, f"globalSplitInc = stride * {constA}")) else: selectByParity(scratchIdx, constB, constA, "globalSplit const = parity ? B : A") module.add(SMulI32(sgpr(gIncIdx), sgpr(gIncIdx), sgpr(scratchIdx), "globalSplitInc = stride * const")) if ldsBoundA == ldsBoundB: return ldsBoundA selectByParity(scratchIdx, ldsBoundB, ldsBoundA, "ldsSplit = parity ? B : A") return sgpr(scratchIdx) def initTDMDescriptor(self, kernel: Mapping, tP: Mapping) -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tc: str = tP['tensorChar'] tlu: int = tP["tlu"] unrolledMajor = not tlu ti: int = tP["idx"] tileChar: str = tP["tileChar"] enableCluster = clusterEnabled(kernel["ClusterDim"]) mod = Module(f"Init TDM Descriptor {tc}") def descSgprName(idx: int) -> str: assert idx < 4 return f"tdm{tc}Group{idx}" def strideRefName() -> str: return self.strideRef(tc, ti) if unrolledMajor else self.strideRef(tc, 3) def sizeRefName(idx: int) -> str: idxChar= INDEX_CHARS[idx] return f"Size{idxChar}" def maskSgprName(tc: str) -> str: if tc.startswith("MXS"): string = tc.removeprefix("MXS") else: string = tc return f"MulticastMask{string}" dtype: DataType = kernel["ProblemType"][f"DataType{tc}"] mt: int = kernel[f"MacroTile{ti}"] du: int = kernel["DepthU"] if (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) or (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"] or tP["isM"]): du = du // 2 if tP["isM"]: du = du // 4 sizeTile0, sizeTile1 = (du, mt) if unrolledMajor else (mt, du) bpe: float = tP["bpeGR"] if not tP["isM"] else 1 #TODO: temp hack numWaves: int = kernel["NumWaves"] wavelen: int = kernel["WavefrontSize"] wgM, wgN = kernel["MIWaveGroup"] numWavesThisAxis: int = wgM if ti == 0 else wgN ldsConstOffset: int = kernel[f"LdsOffset{tc}"] ldsBlockSizePerPad: int = kernel[f"LdsBlockSizePerPad{tc}"] ldsPadSize: int = int(kernel[f"LdsPad{tc}"] * bpe) dim1Divisor = 2 if (kernel["TDMSplit"] and not ("MXS" in tc) and not kernel["ProblemType"]["Sparse"]) else 1 isSparseTrack: bool = (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) or (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) isMetadata: bool = tP["isM"] isMetadataML1: bool = isMetadata and kernel["ProblemType"]["Sparse"] and kernel["ProblemType"]["MetadataLayout"] isTdmIter = (not isMetadata and kernel.get("_TDMIterateMode%s" % tc, False)) group2Name = descSgprName(2) if isTdmIter else None # Group 3 SGPRs alias Group 2; the alias-side zero-init in initOperands # would be redundant. Pass None to skip it. mod.add(comp.initOperands(descSgprName(0), descSgprName(1), group2Name, None)) mod.add(comp.setDataType(dtype, descSgprName(1), isMetadata)) mod.add(comp.setGlobalAddr(descSgprName(0), f"Address{tc}")) if kernel["Multicast"] and enableCluster: mod.add(comp.setMulticastMask(descSgprName(1), maskSgprName(tc), self)) with self.allocTmpSgpr(2, tag="initTDMDescriptor_tmpSgprRes") as tmpSgprRes: waveOffsetSgprIdx: int = tmpSgprRes.idx tmpPadSgprIdx: int = tmpSgprRes.idx + 1 dataBytes = mt // numWaves * du * int(bpe * 4) // (4 * dim1Divisor) mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr("WaveIdx"), dataBytes, f"woffset = WaveIdx * (mt // numWaves * du * bpe // dim1Divisor)")) if ldsBlockSizePerPad != 0 and ldsPadSize != 0: mod.add(SLShiftRightB32(sgpr(tmpPadSgprIdx), int(log2(ldsBlockSizePerPad)), sgpr(waveOffsetSgprIdx), \ f"numPadBlocks = woffset >> log2({ldsBlockSizePerPad=})")) mod.add(SMulI32(sgpr(tmpPadSgprIdx), sgpr(tmpPadSgprIdx), ldsPadSize, \ f"padBytes = numPadBlocks * ({ldsPadSize=})")) mod.add(SAddU32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), sgpr(tmpPadSgprIdx), \ "woffset += padBytes")) mod.add(SAddU32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), ldsConstOffset, "ldsOffset = woffset + ldsConstOffset")) mod.add(comp.setLdsAddr(descSgprName(0), sgpr(waveOffsetSgprIdx))) #TODO: refactor, currently special handling for FP4 along K-dim sizeShifter: int = 1 if dtype.isFloat4() else 0 is6bit: bool = dtype.is6bitFloat() mod.add(comp.setIterationEnabled(descSgprName(1), False)) if isTdmIter: # Iterate-mode supplies the pad via the LDS write stride; disable pad_interval. mod.add(comp.setPadding(descSgprName(1), 0, 0)) else: mod.add(comp.setPadding(descSgprName(1), ldsBlockSizePerPad, ldsPadSize)) dim0Idx, dim1Idx = (3, ti) if unrolledMajor else (ti, 3) if isMetadataML1: mod.add(comp.setTensorDim0(descSgprName(1), sizeRefName(ti), self, sizeShifter)) mod.add(comp.setTensorDim1(descSgprName(1), sizeRefName(3), self, sizeShifter, False, isSparseTrack=isSparseTrack, isMetadata=isMetadata)) mod.add(comp.setTensorTile0(descSgprName(1), sizeTile0, self, sizeShifter)) mod.add(comp.setTensorTile1(descSgprName(1), sizeTile1 // numWaves, self)) else: # isSparseTrack/isMetadata apply to the K dimension (index 3). # For unrolledMajor (TN): K is dim0. For tlu (NN): K is dim1. dim0IsK = unrolledMajor if is6bit: # F6: 0.75 bytes/element. dim0 in bytes = elements * 3 / 4. # NOTE: F6 currently dense-only (TN GEMM). Composition with isSparseTrack/isMetadata # is not yet exercised — both code paths assume they are mutually exclusive. with self.allocTmpSgpr(1, tag="initTDMDescriptor_tmpF6") as tmpF6: mod.add(SLShiftRightB32(sgpr(tmpF6.idx), hex(2), sgpr(sizeRefName(dim0Idx)), "F6: elements / 4")) mod.add(SMulI32(sgpr(tmpF6.idx), sgpr(tmpF6.idx), 3, "F6: * 3 = bytes")) mod.add(comp.setTensorDim0(descSgprName(1), tmpF6.idx, self, 0)) else: mod.add(comp.setTensorDim0(descSgprName(1), sizeRefName(dim0Idx), self, sizeShifter, False, isSparseTrack=isSparseTrack if dim0IsK else False, isMetadata=isMetadata if dim0IsK else False)) mod.add(comp.setTensorDim1(descSgprName(1), sizeRefName(dim1Idx), self, 0, False, isSparseTrack=isSparseTrack if not dim0IsK else False, isMetadata=isMetadata if not dim0IsK else False)) if is6bit: mod.add(comp.setTensorTile0(descSgprName(1), sizeTile0 * 3 // 4, self, 0)) else: mod.add(comp.setTensorTile0(descSgprName(1), sizeTile0, self, sizeShifter)) if isTdmIter: mod.add(comp.setTensorTile1(descSgprName(1), self._tdmIterTileDim1(kernel, tc, du, dtype), self)) else: mod.add(comp.setTensorTile1(descSgprName(1), sizeTile1 // numWaves // dim1Divisor, self)) # --- Tensor stride --- if isMetadata and not kernel["ProblemType"]["MetadataLayout"]: mod.add(comp.setTensorStride0Metadata(descSgprName(1), "SizeL")) elif isMetadataML1: ia = kernel["ProblemType"]["IndexAssignmentsMetadata"] mod.add(comp.setTensorStride0(descSgprName(1), f"Stride{tc}{self.states.indexChars[ia[1]]}", sizeShifter)) elif is6bit: with self.allocTmpSgpr(1, tag="initTDMDescriptor_tmpF62") as tmpF6: strRef = strideRefName() if isinstance(strRef, RegisterContainer): mod.add(SLShiftRightB32(sgpr(tmpF6.idx), hex(2), strRef, "F6: stride / 4")) else: mod.add(SLShiftRightB32(sgpr(tmpF6.idx), hex(2), sgpr(strRef), "F6: stride / 4")) mod.add(SMulI32(sgpr(tmpF6.idx), sgpr(tmpF6.idx), 3, "F6: * 3 = bytes")) mod.add(comp.setTensorStride0(descSgprName(1), tmpF6.idx, 0)) else: mod.add(comp.setTensorStride0(descSgprName(1), strideRefName(), sizeShifter)) if (kernel["TDMSplit"] and not ("MXS" in tc) and not kernel["ProblemType"]["Sparse"]): splitBoundary: int = self.tdmSplitLdsBoundary(kernel, tP) strideRefG, globalIncConst = self.tdmSplitGlobalInc(kernel, tP) mod.add(SMovB32(sgpr(f"tdm{tc}LdsSplitIncs"), splitBoundary, comment=f"tdm{tc} Lds Split Incs({round(mt * du * bpe // dim1Divisor)})")) mod.add(SMulI32(sgpr(f"tdm{tc}GlobalSplitIncs"), strideRefG, globalIncConst, comment=f"tdm{tc} Global Split Incs(stride * {mt * bpe // dim1Divisor})")) if isTdmIter: self._emitTdmIterateInit(mod, kernel, tc, dtype, du, mt, perIssueLoadRowDivisor=numWaves * dim1Divisor, descSgprName=descSgprName, strideRefName=strideRefName) return mod def initTDMDescriptorWaveSeparatedImpl(self, kernel, tP, waveIdxSgpr: int | str = "WaveIdx") -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tc: str = tP['tensorChar'] tlu: int = tP["tlu"] unrolledMajor = not tlu ti: int = tP["idx"] tileChar: str = tP["tileChar"] enableCluster = clusterEnabled(kernel["ClusterDim"]) mod = Module(f"Init TDM Descriptor {tc}") def descSgprName(idx: int) -> str: assert idx < 4 return f"tdm{tc}Group{idx}" def strideRefName() -> str | RegisterContainer: return self.strideRef(tc, ti) if unrolledMajor else self.strideRef(tc, 3) def sizeRefName(idx: int) -> str: idxChar= INDEX_CHARS[idx] return f"Size{idxChar}" def maskSgprName(tc: str) -> str: return f"MulticastMask" dtype: DataType = kernel["ProblemType"][f"DataType{tc}"] mt: int = kernel[f"MacroTile{ti}"] du: int = kernel["DepthU"] isMX: bool = tc.startswith("MX") duScale = kernel["ProblemType"][f"MXBlock{tc[-1]}"] if isMX else 1 du //= duScale if (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) or (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"] or tP["isM"]): du = du // 2 if tP["isM"]: du = du // 4 sizeTile0, sizeTile1 = (du, mt) if unrolledMajor else (mt, du) bpe: float = tP["bpeGR"] if not tP["isM"] else 1 #TODO: temp hack numWaves: int = kernel["NumWaves"] numComp: int = numWaves // 2 assert numComp & (numComp - 1) == 0 wavelen: int = kernel["WavefrontSize"] ldsConstOffset: int = kernel[f"LdsOffset{tc}"] ldsBlockSizePerPad: int = kernel[f"LdsBlockSizePerPad{tc}"] ldsPadSize: int = int(kernel[f"LdsPad{tc}"] * bpe) dim1Divisor = 2 if (kernel["TDMSplit"] and not ("MXS" in tc) and not kernel["ProblemType"]["Sparse"]) else 1 if ("MXS" in tc): subTc = tc[3] mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"] isTdmIter = (tc != "Metadata" and kernel.get("_TDMIterateMode%s" % tc, False)) # WS mode shares one load between A and B; when either tile is iterate, # the non-iterate side must also init Group2/3 (read by the shared load; # stale data corrupts the non-iterate load). MXS* never share SGPRs. needGroup23 = isTdmIter or (tc in ("A", "B") and self.isTdmWaveSeparated(kernel) and (kernel.get("_TDMIterateModeA", False) or kernel.get("_TDMIterateModeB", False))) group2Name = descSgprName(2) if needGroup23 else None # Group 3 aliases Group 2; skip the redundant alias-side zero-init. mod.add(comp.initOperands(descSgprName(0), descSgprName(1), group2Name, None)) mod.add(comp.setDataType(dtype, descSgprName(1), tc == "Metadata")) mod.add(comp.setGlobalAddr(descSgprName(0), f"Address{tc}")) if kernel["Multicast"] and enableCluster: mod.add(comp.setMulticastMask(descSgprName(1), maskSgprName(tc), self)) with self.allocTmpSgpr(2, tag="initTDMDescriptorWaveSeparatedImpl_tmpSgprRes") as tmpSgprRes: waveOffsetSgprIdx: int = tmpSgprRes.idx tmpPadSgprIdx: int = tmpSgprRes.idx + 1 mod.add(SLShiftRightB32(sgpr(waveOffsetSgprIdx), 1, sgpr(waveIdxSgpr), "wId=WaveIdx // 2 (each component covers 2 waves: numComp = numWaves // 2)")) dataBytes = mt // numComp * du * int(bpe * 4) // (4 * dim1Divisor) # Interleave rewrites only A/B; MX scales keep their own stride, relocated via ldsBaseMXS*. _segAB = bool(kernel.get("LDSSegmentInterleave") == 1) and tc in ("A", "B") _segFootprint = _segAB and kernel["LDSSegInterleaveOffsets"].get("footprintPacked", False) if _segAB: dataBytes = kernel["LDSSegInterleaveOffsets"]["writeStrideBytes"] mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), dataBytes, f"woffset = wId * (mt // numComp * du * bpe // dim1Divisor)")) # footprintPacked: writeStrideBytes is the post-pad footprint fA+fB; A/B tiles are packed # exactly, so the component jump must NOT be re-padded (tile-internal pad is set below). if ldsBlockSizePerPad != 0 and ldsPadSize != 0 and not _segFootprint: mod.add(SLShiftRightB32(sgpr(tmpPadSgprIdx), int(log2(ldsBlockSizePerPad)), sgpr(waveOffsetSgprIdx), \ f"numPadBlocks = woffset >> log2({ldsBlockSizePerPad=})")) mod.add(SMulI32(sgpr(tmpPadSgprIdx), sgpr(tmpPadSgprIdx), ldsPadSize, \ f"padBytes = numPadBlocks * ({ldsPadSize=})")) mod.add(SAddU32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), sgpr(tmpPadSgprIdx), \ "woffset += padBytes")) if kernel.get("LDSSegmentInterleave") == 1 and tc == "B": ldsConstOffset = kernel["LDSSegInterleaveOffsets"]["ldsBaseB"] elif kernel.get("LDSSegmentInterleave") == 1 and tc in ("MXSA", "MXSB"): _mxBase = kernel["LDSSegInterleaveOffsets"].get("ldsBase" + tc) if _mxBase is not None: ldsConstOffset = _mxBase # relocated MX scale base mod.add(SAddU32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), ldsConstOffset, "ldsOffset = woffset + ldsConstOffset")) mod.add(comp.setLdsAddr(descSgprName(0), sgpr(waveOffsetSgprIdx))) #TODO: refactor, currently special handling for FP4 along K-dim sizeShifter = 1 if dtype.isFloat4() else 0 sizeShifter += ceil(log2(duScale)) isSparseTrack: bool = (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) or (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) isMetadata: bool = tc == "Metadata" size, wgIdx = ("SizeI", "WorkGroup0") if tc[-1] == "A" else ("SizeJ", "WorkGroup1") with self.allocTmpSgpr(1, tag="initTDMDescriptorWaveSeparatedImpl_tmpSgprRes2") as tmpSgprRes: tmpSgprIdx: int = tmpSgprRes.idx mod.add(SMulI32(sgpr(tmpSgprIdx), mt, sgpr(wgIdx))) mod.add(SSubI32(sgpr(tmpSgprIdx), sgpr(size), sgpr(tmpSgprIdx))) mod.add(comp.setIterationEnabled(descSgprName(1), False)) if isTdmIter: # Iterate-mode supplies the pad via the LDS write stride; disable pad_interval. mod.add(comp.setPadding(descSgprName(1), 0, 0)) else: mod.add(comp.setPadding(descSgprName(1), ldsBlockSizePerPad, ldsPadSize)) if ("MXS" in tc): mxDU = kernel["DepthU"] // kernel["ProblemType"][f"MXBlock{subTc}"] numMxKGroups = mxDU // mxUnit dim0 = tmpSgprIdx dim1 = sizeRefName(3) with self.allocTmpSgpr(1, tag="initTDMDescriptorWaveSeparatedImpl_tmpSgpr2") as tmpSgpr: tmpSgprWaveOffset = tmpSgpr.idx if numMxKGroups < numComp: # M/N-splitting: offset within same k_group along tile dimension mod.add(VReadfirstlaneB32(sgpr(tmpSgprWaveOffset), vgpr("Serial"), "first tId")) mod.add(SLShiftRightB32(sgpr(tmpSgprWaveOffset), ceil(log2(wavelen)) + 1, sgpr(tmpSgprWaveOffset), "wId=fTid // wavelen // 2")) mod.add(SMulI32(sgpr(tmpSgprWaveOffset), sgpr(tmpSgprWaveOffset), round(mt // numComp), "woffset = wId * (mt // numComp)")) mod.add(SSubU32(sgpr(dim0), sgpr(dim0), sgpr(tmpSgprWaveOffset), "consider multiple waves")) mod.add(SCMovB32(sgpr(dim0), 0, "set to 0 for waves that no enough data to load")) mod.add(comp.setTensorDim0(descSgprName(1), dim0, self, ceil(log2(mxUnit)), True)) mod.add(comp.setTensorDim1(descSgprName(1), dim1, self, ceil(log2(duScale*mxUnit)), True)) # mod.add(comp.setTensorDim0(descSgprName(1), sizeRefName(ti), self, ceil(log2(mxUnit)), True)) # mod.add(comp.setTensorDim1(descSgprName(1), sizeRefName(3), self, ceil(log2(duScale*mxUnit)), True)) else: is6bit = dtype.is6bitFloat() dim0Idx, dim1Idx = (3, ti) if unrolledMajor else (ti, 3) dim0 = sizeRefName(dim0Idx) if unrolledMajor else tmpSgprIdx dim1 = tmpSgprIdx if unrolledMajor else sizeRefName(dim1Idx) if is6bit: with self.allocTmpSgpr(1, tag="initTDMDescriptorWaveSeparatedImpl_tmpF6") as tmpF6: mod.add(SLShiftRightB32(sgpr(tmpF6.idx), hex(2), sgpr(dim0), "F6: elements / 4")) mod.add(SMulI32(sgpr(tmpF6.idx), sgpr(tmpF6.idx), 3, "F6: * 3 = bytes")) mod.add(comp.setTensorDim0(descSgprName(1), tmpF6.idx, self, 0)) else: mod.add(comp.setTensorDim0(descSgprName(1), dim0, self, sizeShifter, False, isSparseTrack if unrolledMajor else False, isMetadata if unrolledMajor else False)) with self.allocTmpSgpr(1, tag="initTDMDescriptorWaveSeparatedImpl_tmpSgpr3") as tmpSgpr: tmpSgprWaveOffset = tmpSgpr.idx if unrolledMajor: mod.add(VReadfirstlaneB32(sgpr(tmpSgprWaveOffset), vgpr("Serial"), "first tId")) mod.add(SLShiftRightB32(sgpr(tmpSgprWaveOffset), ceil(log2(wavelen)) + 1, sgpr(tmpSgprWaveOffset), "wId=fTid // wavelen // 2")) mod.add(SMulI32(sgpr(tmpSgprWaveOffset), sgpr(tmpSgprWaveOffset), round(mt // numComp // dim1Divisor), "woffset = wId * (mt // numComp // dim1Divisor)")) mod.add(SSubU32(sgpr(dim1), sgpr(dim1), sgpr(tmpSgprWaveOffset), "consider multiple waves")) mod.add(SCMovB32(sgpr(dim1), 0, "set to 0 for waves that no enough data to load")) mod.add(comp.setTensorDim1(descSgprName(1), dim1, self, 0, False, isSparseTrack if not unrolledMajor else False, isMetadata if not unrolledMajor else False)) # The descriptor now holds the full per-wave dim1; TDMSplit recomputes # the half boundaries from it in globalReadDo. if tc.startswith("MX"): #reset to 0 since scale of sizeTile0 and stride for MX is not required sizeShifter = 1 if dtype.isFloat4() else 0 numMxKGroups = sizeTile0 // mxUnit if numMxKGroups >= numComp: # K-splitting: enough k_groups to divide among wave components mod.add(comp.setTensorTile0(descSgprName(1), sizeTile1 * mxUnit, self, sizeShifter)) mod.add(comp.setTensorTile1(descSgprName(1), numMxKGroups // numComp // dim1Divisor, self)) else: # M/N-splitting: not enough k_groups, split tile dimension instead mod.add(comp.setTensorTile0(descSgprName(1), sizeTile1 * mxUnit // numComp, self, sizeShifter)) mod.add(comp.setTensorTile1(descSgprName(1), numMxKGroups // dim1Divisor, self)) mod.add(comp.setTensorStride0(descSgprName(1), sizeRefName(ti), ceil(log2(mxUnit)), True)) else: is6bit = dtype.is6bitFloat() if is6bit: mod.add(comp.setTensorTile0(descSgprName(1), sizeTile0 * 3 // 4, self, 0)) else: mod.add(comp.setTensorTile0(descSgprName(1), sizeTile0, self, sizeShifter)) if isTdmIter: mod.add(comp.setTensorTile1(descSgprName(1), self._tdmIterTileDim1(kernel, tc, du, dtype), self)) else: mod.add(comp.setTensorTile1(descSgprName(1), sizeTile1 // numComp // dim1Divisor, self)) if is6bit: with self.allocTmpSgpr(1, tag="initTDMDescriptorWaveSeparatedImpl_tmpF62") as tmpF6: strRef = strideRefName() if isinstance(strRef, RegisterContainer): mod.add(SLShiftRightB32(sgpr(tmpF6.idx), hex(2), strRef, "F6: stride / 4")) else: mod.add(SLShiftRightB32(sgpr(tmpF6.idx), hex(2), sgpr(strRef), "F6: stride / 4")) mod.add(SMulI32(sgpr(tmpF6.idx), sgpr(tmpF6.idx), 3, "F6: * 3 = bytes")) mod.add(comp.setTensorStride0(descSgprName(1), tmpF6.idx, 0)) else: mod.add(comp.setTensorStride0(descSgprName(1), strideRefName(), sizeShifter)) # Multi-wave TDMSplit recomputes the LDS/global split increments transiently at # point of use (see _tdmSplitMultiWaveInc); nothing is initialized here. if isTdmIter: self._emitTdmIterateInit(mod, kernel, tc, dtype, du, mt, perIssueLoadRowDivisor=numComp * dim1Divisor, descSgprName=descSgprName, strideRefName=strideRefName) return mod def initTDMDescriptorWaveSeparated(self, kernel, tPA, tPB, waveIdxSgpr: int | str = "WaveIdx") -> Module: #TODO: TDM implement mod = Module("TDM Init Wave Separated") tcA: str = tPA["tensorChar"] tcB: str = tPB["tensorChar"] tdmInitLblA = Label(self.labels.getNameInc(f"TDMInit{tcA}"), "") tdmInitLblB = Label(self.labels.getNameInc(f"TDMInit{tcB}"), "") tdmInitLblEnd = Label(self.labels.getNameInc(f"TDMInit{tcA}{tcB}End"), "") mod.add(tdmInitLblA) mod.add(SBitcmp1B32(sgpr(waveIdxSgpr), 0, "Check parity of wId")) mod.add(SCBranchSCC1(tdmInitLblB.getLabelName(), "Jump to B if wId is odd")) mod.add(self.initTDMDescriptorWaveSeparatedImpl(kernel, tPA, waveIdxSgpr)) mod.add(SBranch(tdmInitLblEnd.getLabelName())) mod.add(tdmInitLblB) mod.add(self.initTDMDescriptorWaveSeparatedImpl(kernel, tPB, waveIdxSgpr)) mod.add(tdmInitLblEnd) return mod def tdmGlobalOffset(self, kernel: Mapping, tP: Mapping) -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tc: str = tP['tensorChar'] return comp.calculateStartAddr(self, kernel, tP, f"Address{tc}") def tdmGlobalOffsetWaveSeparated(self, kernel: Mapping, tPA: Mapping, tPB: Mapping, waveIdxSgpr: int | str = "WaveIdx") -> Module: mod = Module("TDM Global Offset Wave Separated") comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tcA: str = tPA["tensorChar"] tcB: str = tPB["tensorChar"] tdmGlobalOffsetLblA = Label(self.labels.getNameInc(f"TDMGlobalOffset{tcA}"), "") tdmGlobalOffsetLblB = Label(self.labels.getNameInc(f"TDMGlobalOffset{tcB}"), "") tdmGlobalOffsetLblEnd = Label(self.labels.getNameInc(f"TDMGlobalOffset{tcA}{tcB}End"), "") mod.add(tdmGlobalOffsetLblA) mod.add(SBitcmp1B32(sgpr(waveIdxSgpr), 0, "Check parity of wId")) mod.add(SCBranchSCC1(tdmGlobalOffsetLblB.getLabelName(), "Jump to B if wId is odd")) dstGroup0A = f"tdm{tcA}Group0" dstGroup0B = f"tdm{tcB}Group0" mod.add(comp.calculateStartAddrWaveSeparated(self, kernel, tPA, f"Address{tcA}", dstGroup0A, waveIdxSgpr)) mod.add(SBranch(tdmGlobalOffsetLblEnd.getLabelName())) mod.add(tdmGlobalOffsetLblB) mod.add(comp.calculateStartAddrWaveSeparated(self, kernel, tPB, f"Address{tcB}", dstGroup0B, waveIdxSgpr)) mod.add(tdmGlobalOffsetLblEnd) return mod def tdmApplyStreamKOffsetWaveSeparated(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: mod = Module("TDM StreamK K-offset Wave Separated") tcA: str = tPA["tensorChar"] tcB: str = tPB["tensorChar"] incSgprName = f"tdm{tcA}{tcB}Incs" group0Name = f"tdm{tcA}Group0" with self.allocTmpSgpr(1, tag="tdmApplyStreamKOffsetWaveSeparated_tmpSgprRes") as tmpSgprRes: tmpSgpr = tmpSgprRes.idx mod.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr("StreamKLocalStart"), src1=sgpr(incSgprName), comment="StreamK K-offset = localStart * increment")) mod.add(SAddU32(dst=sgpr(f"{group0Name}+2"), src0=sgpr(f"{group0Name}+2"), src1=sgpr(tmpSgpr), comment="Apply StreamK K-offset to TDM global addr")) return mod def tdmApplyStreamKTailOffsetWaveSeparated(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: mod = Module("TDM StreamK tail K-offset Wave Separated") tcA: str = tPA["tensorChar"] tcB: str = tPB["tensorChar"] incSgprName = f"tdm{tcA}{tcB}Incs" group0Name = f"tdm{tcA}Group0" with self.allocTmpSgpr(1) as tmpSgprRes: tmpSgpr = tmpSgprRes.idx mod.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr("StreamKLocalEnd"), src1=1, comment="tail iteration index within current StreamK tile")) mod.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=sgpr(incSgprName), comment="StreamK tail K-offset = (localEnd - 1) * increment")) mod.add(SAddU32(dst=sgpr(f"{group0Name}+2"), src0=sgpr(f"{group0Name}+2"), src1=sgpr(tmpSgpr), comment="Apply StreamK tail K-offset to TDM global addr")) return mod ############################################################################## # PAP TDM/DTL LDS-bank helpers. ############################################################################## def papTdmRecomputeWaveIdx(self, kernel: Mapping, waveIdxSgpr: int | str) -> Module: mod = Module("PAP TDM recompute WaveIdx") with self.allocTmpSgpr(1) as tmpSgprRes: mod.add(VReadfirstlaneB32(sgpr(tmpSgprRes.idx), vgpr("Serial"), "first tId")) mod.add(SLShiftRightB32(sgpr(waveIdxSgpr), ceil(log2(kernel["WavefrontSize"])), sgpr(tmpSgprRes.idx), "recompute WaveIdx for PAP TDM descriptor refresh")) return mod def papResetTDMDescriptorForTailWaveSeparated(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: mod = Module("PAP reset TDM descriptor for tail") with self.allocTmpSgpr(1) as waveIdxSgprRes: waveIdxSgpr = waveIdxSgprRes.idx mod.add(self.papTdmRecomputeWaveIdx(kernel, waveIdxSgpr)) mod.add(self.initTDMDescriptorWaveSeparated(kernel, tPA, tPB, waveIdxSgpr)) mod.add(self.tdmGlobalOffsetWaveSeparated(kernel, tPA, tPB, waveIdxSgpr)) mod.add(self.tdmApplyStreamKTailOffsetWaveSeparated(kernel, tPA, tPB)) mod.add(self.resetTDMDescriptorForTailWaveSeparated(kernel, tPA, tPB)) if kernel["LdsOffsetA_Blk"] == 0: return mod comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) with self.allocTmpSgpr(1) as tmpSgprRes: tailBankSgpr = tmpSgprRes.idx mod.add(self.papTdmSelectTailLdsBank(kernel, tailBankSgpr)) for tc in (tPA["tensorChar"], tPB["tensorChar"]): ldsAddrSgprName = comp.getLdsAddrSgprName(f"tdm{tc}Group0") mod.add(self.papTdmSetTailLdsBank(kernel, ldsAddrSgprName, tailBankSgpr)) return mod def papTdmUpdateDescriptor(self, kernel: Mapping, tPA: Mapping, tPB: Mapping, preservePapBank: bool=True) -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tcA: str = tPA["tensorChar"] tcB: str = tPB["tensorChar"] mod = Module(f"PAP TDM refresh descriptor ({tcA}/{tcB})") ldsAddrSgpr: str = comp.getLdsAddrSgprName(f"tdm{tcA}Group0") blkMask: int = kernel["LdsOffsetA_Blk"] with self.allocTmpSgpr(1) as waveIdxSgprRes: waveIdxSgpr = waveIdxSgprRes.idx mod.add(self.papTdmRecomputeWaveIdx(kernel, waveIdxSgpr)) if preservePapBank and blkMask != 0: with self.allocTmpSgpr(1) as tmpSgprRes: papBankSgpr = tmpSgprRes.idx mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr(ldsAddrSgpr), src1=hex(blkMask), comment=f"preserve PAP LDS bank before descriptor refresh ({blkMask:#x})")) mod.add(self.initTDMDescriptorWaveSeparated(kernel, tPA, tPB, waveIdxSgpr)) mod.add(SXorB32(dst=sgpr(ldsAddrSgpr), src0=sgpr(ldsAddrSgpr), src1=sgpr(papBankSgpr), comment="restore PAP LDS bank after descriptor refresh")) else: mod.add(self.initTDMDescriptorWaveSeparated(kernel, tPA, tPB, waveIdxSgpr)) mod.add(self.tdmGlobalOffsetWaveSeparated(kernel, tPA, tPB, waveIdxSgpr)) if kernel["StreamK"] > 0: mod.add(self.tdmApplyStreamKOffsetWaveSeparated(kernel, tPA, tPB)) return mod def papTdmSaveLdsBank(self, kernel: Mapping) -> Module: # Encode the PAP LDS bank in SkPrefetchPrimed so setupNewTile can align # local reads with the bank that the next-tile first-PGR group populated. mod = Module("TDM save PAP LDS bank in SkPrefetchPrimed") comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) ldsAddrSgpr: str = comp.getLdsAddrSgprName("tdmAGroup0") blkMask: int = kernel["LdsOffsetA_Blk"] if blkMask == 0: return mod with self.allocTmpSgpr(1) as tmpSgprRes: papBankSgpr = tmpSgprRes.idx mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr(ldsAddrSgpr), src1=hex(blkMask), comment=f"PAP bank = descriptor LDS addr & LdsOffsetA_Blk({blkMask:#x})")) mod.add(SOrB32(dst=sgpr("SkPrefetchPrimed"), src0=sgpr("SkPrefetchPrimed"), src1=sgpr(papBankSgpr), comment="encode PAP LDS bank in SkPrefetchPrimed")) return mod def papTdmRestoreLdsBank(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) mod = Module("TDM restore PAP LDS bank for primed path") skipLbl = Label(self.labels.getNameInc("SkipPapBankRestore"), "") blkMask: int = kernel["LdsOffsetA_Blk"] if blkMask == 0: return mod mod.add(SCmpEQU32(src0=sgpr("SkPrefetchPrimed"), src1=0, comment="primed?")) mod.add(SCBranchSCC1(labelName=skipLbl.getLabelName(), comment="not primed, skip bank restore")) with self.allocTmpSgpr(1) as tmpSgprRes: papBankSgpr = tmpSgprRes.idx mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr("SkPrefetchPrimed"), src1=hex(blkMask), comment="extract PAP LDS bank from SkPrefetchPrimed")) mod.add(SCmpEQU32(src0=sgpr(papBankSgpr), src1=0, comment="PAP wrote to bank 0?")) mod.add(SCBranchSCC1(labelName=skipLbl.getLabelName(), comment="bank 0, no adjustment needed")) ldsAddrSgpr: str = comp.getLdsAddrSgprName("tdmAGroup0") mod.add(SXorB32(dst=sgpr(ldsAddrSgpr), src0=sgpr(ldsAddrSgpr), src1=sgpr(papBankSgpr), comment="flip A/B descriptor to PAP bank")) if kernel["ProblemType"]["MXBlockA"] and kernel["ProblemType"]["MXBlockB"]: mxLdsAddrSgpr: str = comp.getLdsAddrSgprName("tdmMXSAGroup0") mod.add(SXorB32(dst=sgpr(mxLdsAddrSgpr), src0=sgpr(mxLdsAddrSgpr), src1=sgpr(papBankSgpr), comment="flip MX descriptor to PAP bank")) for tP in (tPA, tPB): tc = tP["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tc}"), src0=vgpr(f"LocalReadAddr{tc}"), src1=sgpr(papBankSgpr), comment=f"shift LocalReadAddr{tc} to PAP bank")) if kernel["ProblemType"]["MXBlockA"]: tcMX = tPA["MX"]["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tcMX}"), src0=vgpr(f"LocalReadAddr{tcMX}"), src1=sgpr(papBankSgpr), comment=f"shift LocalReadAddr{tcMX} to PAP bank")) if kernel["ProblemType"]["MXBlockB"]: tcMX = tPB["MX"]["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tcMX}"), src0=vgpr(f"LocalReadAddr{tcMX}"), src1=sgpr(papBankSgpr), comment=f"shift LocalReadAddr{tcMX} to PAP bank")) mod.add(skipLbl) return mod def papDtlSaveLdsBank(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: # DTL PAP uses the same SkPrefetchPrimed side channel as TDM. Bit 0 marks # primed; LDS bank bits are ORed in when the prefetch lands in the high bank. mod = Module("DTL save PAP LDS bank in SkPrefetchPrimed") blkMask: int = kernel["LdsOffsetA_Blk"] if blkMask == 0: return mod firstDtc = "A" if kernel["DirectToLdsA"] else "B" firstTP = tPA if firstDtc == "A" else tPB with self.allocTmpSgpr(1) as tmpSgprRes: papBankSgpr = tmpSgprRes.idx if self.states.IncLdsBufSwitch: mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr("LDSBufferWriteInc"), src1=hex(blkMask), comment="PAP DTL bank = LDSBufferWriteInc & LdsOffsetA_Blk")) elif self.states.useCommonSgprSwap: mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr("SwapCommon"), src1=hex(blkMask), comment="PAP DTL bank = SwapCommon & LdsOffsetA_Blk")) elif kernel["ExpandPointerSwap"]: mod.add(SMovB32(dst=sgpr(papBankSgpr), src=firstTP["localWriteSwapByteOffset"] & blkMask, comment="PAP DTL bank from local-write swap offset")) else: mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr(f"LocalWriteAddr{firstDtc}"), src1=hex(blkMask), comment=f"PAP DTL bank = LocalWriteAddr{firstDtc} & LdsOffsetA_Blk")) mod.add(SOrB32(dst=sgpr("SkPrefetchPrimed"), src0=sgpr("SkPrefetchPrimed"), src1=sgpr(papBankSgpr), comment="encode PAP DTL LDS bank in SkPrefetchPrimed")) return mod def papDtlRestoreLdsBank(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: mod = Module("DTL restore PAP LDS bank for primed path") skipLbl = Label(self.labels.getNameInc("SkipDtlPapBankRestore"), "") blkMask: int = kernel["LdsOffsetA_Blk"] if blkMask == 0: return mod mod.add(SCmpEQU32(src0=sgpr("SkPrefetchPrimed"), src1=0, comment="primed?")) mod.add(SCBranchSCC1(labelName=skipLbl.getLabelName(), comment="not primed, skip DTL bank restore")) with self.allocTmpSgpr(1) as tmpSgprRes: papBankSgpr = tmpSgprRes.idx mod.add(SAndB32(dst=sgpr(papBankSgpr), src0=sgpr("SkPrefetchPrimed"), src1=hex(blkMask), comment="extract PAP DTL LDS bank from SkPrefetchPrimed")) mod.add(SCmpEQU32(src0=sgpr(papBankSgpr), src1=0, comment="PAP DTL wrote to bank 0?")) mod.add(SCBranchSCC1(labelName=skipLbl.getLabelName(), comment="bank 0, no adjustment needed")) if self.states.IncLdsBufSwitch: mod.add(SMovB32(dst=sgpr("LDSBufferReadInc"), src=sgpr(papBankSgpr), comment="align local-read inc with PAP DTL write bank")) mod.add(SMovB32(dst=sgpr("LDSBufferWriteInc"), src=sgpr(papBankSgpr), comment="align local-write inc with PAP DTL write bank")) elif self.states.useCommonSgprSwap: mod.add(SMovB32(dst=sgpr("SwapCommon"), src=sgpr(papBankSgpr), comment="align common LDS swap with PAP DTL write bank")) elif not kernel["ExpandPointerSwap"]: for tP in (tPA, tPB): tc = tP["tensorChar"] if kernel["DirectToLds%s" % tc]: mod.add(SOrB32(dst=sgpr(f"LocalWriteAddr{tc}"), src0=sgpr(f"LocalWriteAddr{tc}"), src1=sgpr(papBankSgpr), comment=f"shift LocalWriteAddr{tc} to PAP DTL bank")) for tP in (tPA, tPB): tc = tP["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tc}"), src0=vgpr(f"LocalReadAddr{tc}"), src1=sgpr(papBankSgpr), comment=f"shift LocalReadAddr{tc} to PAP DTL bank")) if kernel["ProblemType"]["MXBlockA"]: tcMX = tPA["MX"]["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tcMX}"), src0=vgpr(f"LocalReadAddr{tcMX}"), src1=sgpr(papBankSgpr), comment=f"shift LocalReadAddr{tcMX} to PAP DTL bank")) if kernel["ProblemType"]["MXBlockB"]: tcMX = tPB["MX"]["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tcMX}"), src0=vgpr(f"LocalReadAddr{tcMX}"), src1=sgpr(papBankSgpr), comment=f"shift LocalReadAddr{tcMX} to PAP DTL bank")) mod.add(skipLbl) return mod def papTdmSelectTailLdsBank(self, kernel: Mapping, dstSgpr: int) -> Module: mod = Module("TDM select tail LDS bank") blkMask: int = kernel["LdsOffsetA_Blk"] mod.add(SAndB32(dst=sgpr(dstSgpr), src0=sgpr("SkPrefetchPrimed"), src1=hex(blkMask), comment="PAP LDS bank from primed state")) mod.add(SXorB32(dst=sgpr(dstSgpr), src0=sgpr(dstSgpr), src1=hex(blkMask), comment="tail uses opposite LDS bank from PAP")) mod.add(SCmpEQU32(src0=sgpr("SkPrefetchPrimed"), src1=0, comment="no PAP primed?")) mod.add(SCMovB32(dst=sgpr(dstSgpr), src=0, comment="no primed PAP, keep tail in bank 0")) return mod def papTdmSetTailLdsBank(self, kernel: Mapping, ldsAddrSgprName: str, tailBankSgpr: int) -> Module: mod = Module("TDM set tail LDS bank") blkMask: int = kernel["LdsOffsetA_Blk"] mod.add(SAndB32(dst=sgpr(ldsAddrSgprName), src0=sgpr(ldsAddrSgprName), src1=hex(~blkMask & 0xFFFFFFFF), comment="clear TDM LDS bank bit for tail loop")) mod.add(SOrB32(dst=sgpr(ldsAddrSgprName), src0=sgpr(ldsAddrSgprName), src1=sgpr(tailBankSgpr), comment="set TDM tail write bank opposite PAP")) return mod def papTdmShiftTailLdsBank(self, kernel: Mapping, tPA: Mapping, tPB: Mapping) -> Module: mod = Module("TDM shift tail local-read LDS bank") if kernel["LdsOffsetA_Blk"] == 0: return mod with self.allocTmpSgpr(1) as tmpSgprRes: tailBankSgpr = tmpSgprRes.idx mod.add(self.papTdmSelectTailLdsBank(kernel, tailBankSgpr)) for tP in (tPA, tPB): tc = tP["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tc}"), src0=vgpr(f"LocalReadAddr{tc}"), src1=sgpr(tailBankSgpr), comment=f"shift tail LocalReadAddr{tc} to TDM write bank")) if kernel["ProblemType"]["MXBlockA"]: tcMX = tPA["MX"]["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tcMX}"), src0=vgpr(f"LocalReadAddr{tcMX}"), src1=sgpr(tailBankSgpr), comment=f"shift tail LocalReadAddr{tcMX} to TDM write bank")) if kernel["ProblemType"]["MXBlockB"]: tcMX = tPB["MX"]["tensorChar"] mod.add(VAddU32(dst=vgpr(f"LocalReadAddr{tcMX}"), src0=vgpr(f"LocalReadAddr{tcMX}"), src1=sgpr(tailBankSgpr), comment=f"shift tail LocalReadAddr{tcMX} to TDM write bank")) return mod def tdmIncrementAB(self, kernel, tP, loopIdx=None, prefetchIndex=0) -> Module: # Subtile uses per-wave descriptors, so multi-wave is valid here if not kernel.get("UseSubtileImpl"): assert kernel["NumWaves"] == 1 comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tc: str = tP['tensorChar'] mod = Module("TDM increment") tdmGroup0 = f"tdm{tc}Group0" incSgprName = f"GlobalReadIncs{tc}" if loopIdx is not None and loopIdx == self.states.unrollIdx and self.states.staggerUCode: with self.allocTmpSgpr(2, tag="tdmIncrementAB_tmpSgprInfo") as tmpSgprInfo: incTmpLo = tmpSgprInfo.idx incTmpHi = tmpSgprInfo.idx + 1 if prefetchIndex: mod.add(SAddU32(dst=sgpr(incTmpLo), src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=prefetchIndex, comment="remove pf(%u)"%prefetchIndex)) mod.add(SCmpEQU32(src0=sgpr("StaggerUIter"), src1=sgpr(incTmpLo), comment="Is this wrapIter? (pf)")) else: mod.add(SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=sgpr("StaggerUIter"), comment="Is this the wrapIter?")) mod.add(SCSelectB32(dst=sgpr(incTmpLo), src0=sgpr(f"WrapU{tc}+0"), src1=sgpr(incSgprName), \ comment="select WrapU or normal inc (lo)")) mod.add(SCSelectB32(dst=sgpr(incTmpHi), src0=sgpr(f"WrapU{tc}+1"), src1=0, \ comment="select WrapU or normal inc (hi)")) mod.add(SAddU64(dst=sgpr(f"{tdmGroup0}+2", 2), src0=sgpr(f"{tdmGroup0}+2", 2), \ src1=sgpr(incTmpLo, 2), comment="TDM addr += inc (with wrap, 64-bit)")) else: mod.add(comp.incrementGlobalAddr(self, tdmGroup0, incSgprName)) if kernel["TDMSplit"] and not ("MXS" in tc) and not kernel["ProblemType"]["Sparse"]: mod.add(SSubU32(sgpr(f"{tdmGroup0}+2"), sgpr(f"{tdmGroup0}+2"), sgpr(f"tdm{tc}GlobalSplitIncs"), f"tdm{tc} Global Split Incs sub")) mod.add(SSubBU32(sgpr(f"{tdmGroup0}+3"), sgpr(f"{tdmGroup0}+3"), 0, f"tdm{tc} Global Split borrow")) mod.add(SSubU32(sgpr(f"{tdmGroup0}+1"), sgpr(f"{tdmGroup0}+1"), sgpr(f"tdm{tc}LdsSplitIncs"), f"tdm{tc} Lds Split Incs sub")) if (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) or \ (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]): if loopIdx is not None and loopIdx == self.states.unrollIdx and self.states.staggerUCode: with self.allocTmpSgpr(2, tag="tdmIncrementAB_tmpSgprInfo2") as tmpSgprInfo: incTmpLo = tmpSgprInfo.idx incTmpHi = tmpSgprInfo.idx + 1 if prefetchIndex: mod.add(SAddU32(dst=sgpr(incTmpLo), src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=prefetchIndex, comment="remove pf(%u)"%prefetchIndex)) mod.add(SCmpEQU32(src0=sgpr("StaggerUIter"), src1=sgpr(incTmpLo), comment="Is this wrapIter? (metadata pf)")) else: mod.add(SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=sgpr("StaggerUIter"), comment="Is this the wrapIter? (metadata)")) mod.add(SCSelectB32(dst=sgpr(incTmpLo), src0=sgpr("WrapUMetadata+0"), src1=sgpr("GlobalReadIncsMetadata"), \ comment="select WrapUMetadata or normal inc (lo)")) mod.add(SCSelectB32(dst=sgpr(incTmpHi), src0=sgpr("WrapUMetadata+1"), src1=0, \ comment="select WrapUMetadata or normal inc (hi)")) mod.add(SAddU32(dst=sgpr("tdmMetadataGroup0+2"), src0=sgpr("tdmMetadataGroup0+2"), \ src1=sgpr(incTmpLo), comment="TDM metadata addr += inc (with wrap, lo)")) mod.add(SAddCU32(dst=sgpr("tdmMetadataGroup0+3"), src0=sgpr("tdmMetadataGroup0+3"), \ src1=sgpr(incTmpHi), comment="TDM metadata addr += inc (with wrap, hi)")) else: mod.add(comp.incrementGlobalAddr(self, "tdmMetadataGroup0", "GlobalReadIncsMetadata")) return mod def tdmIncrementABWaveSperated(self, kernel, tPA, tPB, loopIdx=None, prefetchIndex=0) -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) #TODO: TDM replace by universal tdm group sgpr tcA: str = tPA['tensorChar'] tcB: str = tPB['tensorChar'] mod = Module("TDMGlobalIncrementsWaveSeparated") tdmGroup0 = f"tdm{tcA}Group0" incSgprName = f"tdm{tcA}{tcB}Incs" if loopIdx is not None and loopIdx == self.states.unrollIdx and self.states.staggerUCode: with self.allocTmpSgpr(4, tag="tdmIncrementABWaveSperated_tmpSgprInfo") as tmpSgprInfo: incTmpLo = tmpSgprInfo.idx incTmpHi = tmpSgprInfo.idx + 1 wrapTmpLo = tmpSgprInfo.idx + 2 wrapTmpHi = tmpSgprInfo.idx + 3 with self.allocTmpSgpr(1, tag="tdmIncrementABWaveSperated_tmpSgprInfo2") as waveIdTmp: self._emitTdmWaveParitySCC(mod, kernel, waveIdTmp.idx, "check wave parity") mod.add(SCSelectB32(dst=sgpr(wrapTmpLo), src0=sgpr(f"WrapU{tcB}+0"), src1=sgpr(f"WrapU{tcA}+0"), \ comment="select WrapU based on wave parity (lo)")) mod.add(SCSelectB32(dst=sgpr(wrapTmpHi), src0=sgpr(f"WrapU{tcB}+1"), src1=sgpr(f"WrapU{tcA}+1"), \ comment="select WrapU based on wave parity (hi)")) if prefetchIndex: mod.add(SAddU32(dst=sgpr(incTmpLo), src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=prefetchIndex, comment="remove pf(%u)"%prefetchIndex)) mod.add(SCmpEQU32(src0=sgpr("StaggerUIter"), src1=sgpr(incTmpLo), comment="Is this wrapIter? (pf)")) else: mod.add(SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=sgpr("StaggerUIter"), comment="Is this the wrapIter?")) mod.add(SCSelectB32(dst=sgpr(incTmpLo), src0=sgpr(wrapTmpLo), src1=sgpr(incSgprName), \ comment="select WrapU or normal inc (lo)")) mod.add(SCSelectB32(dst=sgpr(incTmpHi), src0=sgpr(wrapTmpHi), src1=0, \ comment="select WrapU or normal inc (hi)")) mod.add(SAddU64(dst=sgpr(f"{tdmGroup0}+2", 2), src0=sgpr(f"{tdmGroup0}+2", 2), \ src1=sgpr(incTmpLo, 2), comment="TDM addr += inc (with wrap, 64-bit)")) else: mod.add(comp.incrementGlobalAddr(self, tdmGroup0, incSgprName)) if kernel["TDMSplit"] and not (("MXS" in tcA) or ("MXS" in tcB)) and not kernel["ProblemType"]["Sparse"]: # Recompute the split increments transiently (see _tdmSplitMultiWaveInc). The # parity recompute clobbers SCC and runs before the sub chain, so the borrow # between the +2 subtract and the +3 borrow remains intact. with self.allocTmpSgpr(2, tag="tdmSplitIncSub") as incTmp: gIncIdx = incTmp.idx scratchIdx = incTmp.idx + 1 ldsIncOp = self._tdmSplitMultiWaveInc(mod, kernel, gIncIdx, scratchIdx, tPA, tPB) mod.add(SSubU32(sgpr(f"{tdmGroup0}+2"), sgpr(f"{tdmGroup0}+2"), sgpr(gIncIdx), "tdmAB Global Split Incs sub")) mod.add(SSubBU32(sgpr(f"{tdmGroup0}+3"), sgpr(f"{tdmGroup0}+3"), 0, "tdmAB Global Split borrow")) mod.add(SSubU32(sgpr(f"{tdmGroup0}+1"), sgpr(f"{tdmGroup0}+1"), ldsIncOp, "tdmAB Lds Split Incs sub")) if kernel["ProblemType"]["Sparse"] == 1 and tPA["is_sparse"] \ or kernel["ProblemType"]["Sparse"] == 2 and tPB["is_sparse"]: if loopIdx is not None and loopIdx == self.states.unrollIdx and self.states.staggerUCode: # Wave-uniform: all waves update tdmMetadataGroup0; only sparse-side waves actually # issue loads via this SRD. Mirrors the non-stagger metadata path above which is also # emitted wave-uniformly. with self.allocTmpSgpr(2, tag="tdmIncrementABWaveSperated_tmpSgprInfo3") as tmpSgprInfo: incTmpLo = tmpSgprInfo.idx incTmpHi = tmpSgprInfo.idx + 1 if prefetchIndex: mod.add(SAddU32(dst=sgpr(incTmpLo), src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=prefetchIndex, comment="remove pf(%u)"%prefetchIndex)) mod.add(SCmpEQU32(src0=sgpr("StaggerUIter"), src1=sgpr(incTmpLo), comment="Is this wrapIter? (metadata pf)")) else: mod.add(SCmpEQU32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=sgpr("StaggerUIter"), comment="Is this the wrapIter? (metadata)")) mod.add(SCSelectB32(dst=sgpr(incTmpLo), src0=sgpr("WrapUMetadata+0"), src1=sgpr("GlobalReadIncsMetadata"), \ comment="select WrapUMetadata or normal inc (lo)")) mod.add(SCSelectB32(dst=sgpr(incTmpHi), src0=sgpr("WrapUMetadata+1"), src1=0, \ comment="select WrapUMetadata or normal inc (hi)")) mod.add(SAddU32(dst=sgpr("tdmMetadataGroup0+2"), src0=sgpr("tdmMetadataGroup0+2"), \ src1=sgpr(incTmpLo), comment="TDM metadata addr += inc (with wrap, lo)")) mod.add(SAddCU32(dst=sgpr("tdmMetadataGroup0+3"), src0=sgpr("tdmMetadataGroup0+3"), \ src1=sgpr(incTmpHi), comment="TDM metadata addr += inc (with wrap, hi)")) else: mod.add(comp.incrementGlobalAddr(self, "tdmMetadataGroup0", "GlobalReadIncsMetadata")) return mod def tdmSetupIncrementWaveSeparated(self, kernel, tpA, tpB) -> Module: mod = Module() tcA: str = tpA["tensorChar"] tcB: str = tpB["tensorChar"] wavelen: int = kernel["WavefrontSize"] incSgprName = f"tdm{tcA}{tcB}Incs" mod.add(SBitcmp1B32(sgpr("WaveIdx"), 0, "Check parity of wId")) #TODO: should not directly use GRIA and GRIB mod.add(SCSelectB32(sgpr(incSgprName), sgpr(f"GlobalReadIncs{tcB}"), sgpr(f"GlobalReadIncs{tcA}"))) return mod def resetTDMDescriptorForTail(self, kernel: Mapping, tP: Mapping, tmpSgprWaveOffset = None) -> Module: comp: TensorDataMoverLoad = TensorDataMoverLoad.find(self) tc: str = tP['tensorChar'] tlu: int = tP["tlu"] unrolledMajor = not tlu numWaves: int = kernel["NumWaves"] numComp: int = numWaves // 2 mod = Module(f"Reset TDM Descriptor For Tail {tc}") def descSgprName(idx: int) -> str: assert idx < 2 return f"tdm{tc}Group{idx}" dtype: DataType = kernel["ProblemType"][f"DataType{tc}"] du: int = kernel["DepthU"] isMXS: bool = tc.startswith("MX") duScale = kernel["ProblemType"][f"MXBlock{tc[-1]}"] if isMXS else 1 if isMXS: subTc = tc[3] mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"] mxDU = kernel["DepthU"] // kernel["ProblemType"][f"MXBlock{subTc}"] numMxKGroups = mxDU // mxUnit mxKSplitting = numMxKGroups >= numComp if isMXS else False #TODO: refactor, currently special handling for FP4 along K-dim sizeShifter = 1 if dtype.isFloat4() else 0 sizeShifter += ceil(log2(duScale)) tdmDescIdx = 1 if (unrolledMajor and not isMXS) else 2 sizeShifter = ceil(log2(duScale*mxUnit)) if isMXS else sizeShifter # dim1 is in raw elements per Group 1 spec; only dim0 takes the FP4 shifter. if tdmDescIdx == 2 and not isMXS: sizeShifter = 0 # Reset TDM LDS write address (tdmGroup0+1) to the first buffer 0. # After an odd number of main loop iterations the address lands in the second # buffer (buffer 1). However, when numReadsIterCoalesced{A,B} > 1 ("wider local read"). # recalcLocalReadAddressesAB() performs this switch by recomputing the local-read # pointer to buffer 0; (e.g. ds_load_b128 covering 2 MI-K to ds_load_b64 per MI_K). # (needResetLROffsets or kernel["StreamK"]) in KernelWriter, keeping write/read consistent. needLdsReset = (kernel["StreamK"] or self.states.numReadsIterCoalescedA > 1 or self.states.numReadsIterCoalescedB > 1) if not kernel["1LDSBuffer"] and needLdsReset: ldsAddrSgprName: str = comp.getLdsAddrSgprName(descSgprName(0)) swapMask: int = kernel["LdsOffsetA_Blk"] mod.addComment("TDM tail: reset LDS write addr to buffer 0 (matches recalculated local-read ptr)") mod.add(SAndB32(sgpr(ldsAddrSgprName), sgpr(ldsAddrSgprName), hex(~swapMask & 0xFFFFFFFF), "clear swap bit so TDM writes to buffer 0, same half as tail local reads")) isSparseTrack: bool = (kernel["ProblemType"]["Sparse"] == 1 and tP["isA"]) or \ (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) with self.allocTmpSgpr(1, tag="resetTDMDescriptorForTail_tmpSgpr") as tmpSgpr: mod.add(SAndB32(sgpr(tmpSgpr.idx), sgpr("SizeL"), (du - 1))) if (not unrolledMajor or mxKSplitting) and tmpSgprWaveOffset != None: mod.add(SSubU32(sgpr(tmpSgpr.idx), sgpr(tmpSgpr.idx), sgpr(tmpSgprWaveOffset), "consider multiple waves")) mod.add(SCMovB32(sgpr(tmpSgpr.idx), 0, "set to 0 for waves that no enough data to load")) if dtype.is6bitFloat() and tdmDescIdx == 1 and not isMXS: # F6: 0.75 bytes/element — convert tail-K element count to bytes (elements * 3 / 4). # F6 currently dense-only; isSparseTrack is False here, so omitted. mod.add(SLShiftRightB32(sgpr(tmpSgpr.idx), hex(2), sgpr(tmpSgpr.idx), "F6 tail: / 4")) mod.add(SMulI32(sgpr(tmpSgpr.idx), sgpr(tmpSgpr.idx), 3, "F6 tail: * 3 = bytes")) mod.add(comp.resetTensorDimForTail(descSgprName(1), tmpSgpr.idx, tdmDescIdx, self, 0, isMXS)) else: mod.add(comp.resetTensorDimForTail(descSgprName(1), tmpSgpr.idx, tdmDescIdx, self, sizeShifter, isMXS, isSparseTrack)) return mod def resetTDMDescriptorForTailWaveSeparated(self, kernel, tPA, tPB) -> Module: #TODO: TDM implement mod = Module("TDM Init Wave Separated") tcA: str = tPA["tensorChar"] tcB: str = tPB["tensorChar"] numWaves: int = kernel["NumWaves"] wavelen: int = kernel["WavefrontSize"] numComp: int = numWaves // 2 tluA: int = tPA["tlu"] unrolledMajorA = not tluA tluB: int = tPB["tlu"] unrolledMajorB = not tluB du: int = kernel["DepthU"] isMXSA: bool = tcA.startswith("MX") isMXSB: bool = tcB.startswith("MX") if isMXSA: subTc = tcA[3] mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"] mxDU = kernel["DepthU"] // kernel["ProblemType"][f"MXBlock{subTc}"] numMxKGroups = mxDU // mxUnit mxKSplittingA = numMxKGroups >= numComp if isMXSA else False if isMXSB: subTc = tcB[3] mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"] mxDU = kernel["DepthU"] // kernel["ProblemType"][f"MXBlock{subTc}"] numMxKGroups = mxDU // mxUnit mxKSplittingB = numMxKGroups >= numComp if isMXSB else False assert numWaves > 1 tdmResetTailLblA = Label(f"TDMResetTail{tcA}", "") tdmResetTailLblB = Label(f"TDMResetTail{tcB}", "") tdmResetTailLblEnd = Label(f"TDMResetTail{tcA}{tcB}End", "") mod.add(tdmResetTailLblA) with self.allocTmpSgpr(2, tag="resetTDMDescriptorForTailWaveSeparated_tmpSgprRes") as tmpSgprRes: waveIdSgprIdx: int = tmpSgprRes.idx waveOfstSgprIdx = waveIdSgprIdx + 1 mod.add(VReadfirstlaneB32(sgpr(waveIdSgprIdx), vgpr("Serial"), "first tId")) mod.add(SLShiftRightB32(sgpr(waveIdSgprIdx), ceil(log2(wavelen)), sgpr(waveIdSgprIdx), "wId=fTid // wavelen")) mod.add(SBitcmp1B32(sgpr(waveIdSgprIdx), 0, "Check parity of wId")) mod.add(SCBranchSCC1(tdmResetTailLblB.getLabelName(), "Jump to B if wId is odd")) wOfstSgprId = None if unrolledMajorA and not isMXSA else waveOfstSgprIdx if not unrolledMajorA or mxKSplittingA: mod.add(SLShiftRightB32(sgpr(waveOfstSgprIdx), 1, sgpr(waveIdSgprIdx), "wOffset=wId // 2")) mod.add(SMulI32(sgpr(waveOfstSgprIdx), sgpr(waveOfstSgprIdx), int(du // numComp), "wOffset = wOffset * du // numpComp")) mod.add(self.resetTDMDescriptorForTail(kernel, tPA, wOfstSgprId)) mod.add(SBranch(tdmResetTailLblEnd.getLabelName())) mod.add(tdmResetTailLblB) wOfstSgprId = None if unrolledMajorB and not isMXSB else waveOfstSgprIdx if not unrolledMajorB or mxKSplittingB: mod.add(SLShiftRightB32(sgpr(waveOfstSgprIdx), 1, sgpr(waveIdSgprIdx), "wOffset=wId // 2")) mod.add(SMulI32(sgpr(waveOfstSgprIdx), sgpr(waveOfstSgprIdx), int(du // numComp), "wOffset = wOffset * du // numpComp")) mod.add(self.resetTDMDescriptorForTail(kernel, tPB, wOfstSgprId)) mod.add(tdmResetTailLblEnd) return mod def gl2PrefetchInit(self, kernel, tPA, tPB): comp = GL2PrefetchLoad.find(self) comp.init(self, kernel, tPA) comp.init(self, kernel, tPB) if kernel["ProblemType"]["MXBlockA"]: comp.init(self, kernel, tPA["MX"]) if kernel["ProblemType"]["MXBlockB"]: comp.init(self, kernel, tPB["MX"]) if kernel["enableTDMMetadata"]: comp.init(self, kernel, tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"]) def gl2PrefetchCalcAddr(self, kernel, tPA, tPB) -> Module: mod = Module("GL2 Prefetch Addresses Calculation") comp = GL2PrefetchLoad.find(self) tpList = [tPA, tPB] if kernel["ProblemType"]["MXBlockA"]: tpList.append(tPA["MX"]) if kernel["ProblemType"]["MXBlockB"]: tpList.append(tPB["MX"]) if kernel["enableTDMMetadata"]: tpList.append(tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"]) for tp in tpList: mod.add(comp.setIncrement(self, kernel, tp)) mod.add(comp.calculateStartAddr(self, kernel, tp)) return mod def gl2PrefetchIssueLoad(self, kernel, tPA, tPB) -> Module: mod = Module("GL2 Prefetch Issue Load") mod.addComment("GL2 Prefetch Issue Load") comp = GL2PrefetchLoad.find(self) mod.add(comp.issueLoad(self, kernel, tPA)) mod.add(comp.issueLoad(self, kernel, tPB)) if kernel["ProblemType"]["MXBlockA"]: mod.add(comp.issueLoad(self, kernel, tPA["MX"])) if kernel["ProblemType"]["MXBlockB"]: mod.add(comp.issueLoad(self, kernel, tPB["MX"])) if kernel["enableTDMMetadata"]: mod.add(comp.issueLoad(self, kernel, tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"])) return mod def gl2PrefetchIncrementAddr(self, kernel, tPA, tPB) -> Module: mod = Module("GL2 Prefetch Increment Address") mod.addComment("GL2 Prefetch Increment Address") comp = GL2PrefetchLoad.find(self) mod.add(comp.incrementAddr(self, kernel, tPA)) mod.add(comp.incrementAddr(self, kernel, tPB)) if kernel["ProblemType"]["MXBlockA"]: mod.add(comp.incrementAddr(self, kernel, tPA["MX"])) if kernel["ProblemType"]["MXBlockB"]: mod.add(comp.incrementAddr(self, kernel, tPB["MX"])) if kernel["enableTDMMetadata"]: mod.add(comp.incrementAddr(self, kernel, tPA["tpsMetadata"] if tPA["is_sparse"] else tPB["tpsMetadata"])) return mod def getHalfPLRGroups(self, kernel, lc, u): # 2: each subIter use 2 vgpr groups. 3: HalfPLR has 3 vgpr groups g0 = ((lc * kernel["LoopIters"] + u) * 2) % 3 g1 = (g0 + 1) % 3 return [g0, g1] def getHalfPLRValuStr(self, vgprGroups, offset, tc) -> str: assert(tc == 'A' or tc == 'B') offset = int(offset) numVgprPerGroup = (self.states.a.numVgprValuPerBlock if tc == 'A' else self.states.b.numVgprValuPerBlock) // 2 gIdx = offset // numVgprPerGroup newOffset = offset % numVgprPerGroup valuStr = "Valu%s_G%u+%u"%(tc, vgprGroups[gIdx], newOffset) return valuStr def graIncrementMask(self, kernel, tPA, tPB) -> Module: mod = Module("graIncrementMask") if not kernel["HalfPLR"]: return mod # Only support TDM mod.add(SCmpLeI32(src0=self.loopCounter(kernel, self.states.unrollIdx), \ src1=1, \ comment="%s"%"is this the last iters")) if kernel["NumWaves"] > 1: mod.add(SCMovB32(dst=sgpr("tdmABIncs"), src=0, comment="")) if kernel["ProblemType"]["MXBlockA"]: mod.add(SCMovB32(dst=sgpr("tdmMXSAMXSBIncs"), src=0, comment="")) else: mod.add(SCMovB32(dst=sgpr("GlobalReadIncsA"), src=0, comment="")) mod.add(SCMovB32(dst=sgpr("GlobalReadIncsB"), src=0, comment="")) if kernel["ProblemType"]["MXBlockA"]: mod.add(SCMovB32(dst=sgpr("GlobalReadIncsMXSA"), src=0, comment="")) if kernel["ProblemType"]["MXBlockB"]: mod.add(SCMovB32(dst=sgpr("GlobalReadIncsMXSB"), src=0, comment="")) return mod def _getEccOffset(totalWidth, bpr, bpe, glvw, idx, numVgprG2L): if totalWidth < 1: # Need extra offset if global read < 1 modVal = int(bpr / (bpe * glvw)) left = idx % modVal return numVgprG2L * left else: return 0 #TODO: remove this once upcoming compiler changes applied def _shiftLrElements(dstVgprStr: str, numVgprPerInput: int, tIdx: int) -> Module: mod = Module("BF6/FP6 element shift") paddedNumVgprs = 16 vgprPad = paddedNumVgprs - numVgprPerInput % paddedNumVgprs startVgpr = 0 if tIdx else 4 offset = 0 if tIdx else 1 for i, r in enumerate(range(startVgpr, paddedNumVgprs, 4)): for k in range(3): mod.add(VMovB32(vgpr(f"{dstVgprStr}+{r+k-(i+offset)}"), vgpr(f"{dstVgprStr}+{tIdx*vgprPad+r+k}"))) return mod