derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
20058 lines
1.1 MiB
Plaintext
20058 lines
1.1 MiB
Plaintext
################################################################################
|
|
#
|
|
# 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<LoopIters-1> (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<LoopIters-1> (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 <h,h,h,h,"h,h"> +"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 = <prev-incs> - 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 + <prev-incs>"%(tc, loopChar, tc, loopChar) ))
|
|
module.add(SSubI32(dst=sgpr(graInc), \
|
|
src0=0, \
|
|
src1=sgpr(graInc), \
|
|
comment="incr%s%s = - (stride%s%s + <prev-incs>)"%(tc, loopChar, tc, loopChar) ))
|
|
elif isMirrorIdxPrev:
|
|
module.add(SAddI32(dst=sgpr(graInc), \
|
|
src0=stride, \
|
|
src1=sgpr(graInc), \
|
|
comment="incr%s%s = stride%s%s + <prev-incs>"%(tc, loopChar, tc, loopChar) ))
|
|
else:
|
|
module.add(SSubI32(dst=sgpr(graInc), \
|
|
src0=stride, \
|
|
src1=sgpr(graInc), \
|
|
comment="incr%s%s = stride%s%s - <prev-incs>"%(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]<InType>
|
|
miOutInstType = dataTypeNameAbbrevToInstType(dataTypeA.MIOutputTypeNameAbbrev()) # v_mfma_<OutType>..
|
|
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<tc>:
|
|
# 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<<log2(bpe) for elt-1's real advance.
|
|
# `forceinitrow0=1` opens the same gate inside incrementToNextRow.
|
|
# Legacy path (CompactLoopStore=False) keeps the original `addrCalc.rowInc`
|
|
# gate -- bit-for-bit unchanged because `False and ...` short-circuits and
|
|
# forceinitrow0 only matters when we actually enter the function.
|
|
if (ss.optSrdIncForRow and (addrCalc.rowInc or (kernel["CompactLoopStore"] and elementIdx == 0 and batchIdx == 0))):
|
|
module.add(addrCalc.incrementToNextRow(kernel, "D", ss, tmpS01, forceinitrow0=1,
|
|
overrideAfterPrimerRows=overrideAfterPrimerRows))
|
|
|
|
dataType = kernel["ProblemType"]["DestDataType"]
|
|
globalOffset = addrCalc.globalOffset
|
|
elif tc == 'TD':
|
|
isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal(
|
|
self.states.asmCaps, kernel["NonTemporalD"], _temporalHint(kernel, "TD"), _nonVolatile(kernel, "TD"))
|
|
isGlc, isSlc, isNT, scope, th, nv = forceCoherentNonTemporal(
|
|
self.states.asmCaps, isNT, _temporalHint(kernel, "TD"), _nonVolatile(kernel, "TD"), scope)
|
|
|
|
bps = int(kernel["ProblemType"]["DestDataType"].numBytes()) * ss.cfg.gwvw
|
|
rpv = int(kernel["ProblemType"]["DestDataType"].numBytes()) * ss.cfg.gwvw / self.states.bpr
|
|
if kernel["BufferStore"]:
|
|
addr0 = vgpr(addrCalc.addrGSUSyncVgprs)
|
|
addr1 = sgpr("SrdTD", 4)
|
|
else:
|
|
addr0 = vgpr(addrCalc.addrGSUSyncVgprs,2)
|
|
addr1 = ""
|
|
if ss.optSrdIncForRow and addrCalc.rowInc:
|
|
module.add(addrCalc.incrementToNextRow(kernel, "TD", ss, tmpS01))
|
|
dataType = kernel["ProblemType"]["DestDataType"]
|
|
globalOffset = addrCalc.globalOffset
|
|
globalOffset = int((globalOffset/self.states.bpeCexternal) * self.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters())
|
|
elif tc == 'WS':
|
|
isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal(
|
|
self.states.asmCaps, kernel["NonTemporalWS"], _temporalHint(kernel, "WS"), _nonVolatile(kernel, "WS"))
|
|
isGlc, isSlc, isNT, scope, th, nv = forceCoherentNonTemporal(
|
|
self.states.asmCaps, isNT, _temporalHint(kernel, "WS"), _nonVolatile(kernel, "WS"), scope)
|
|
isDlc = True
|
|
|
|
bps = self.states.bpeCinternal * ss.cfg.gwvw
|
|
rpv = self.states.bpeCinternal * ss.cfg.gwvw / self.states.bpr
|
|
if kernel["BufferStore"]:
|
|
addr0 = vgpr(addrCalc.addrDVgpr)
|
|
addr1 = sgpr("SrdWS", 4)
|
|
else:
|
|
addr0 = vgpr(addrCalc.addrDVgpr,2)
|
|
addr1 = ""
|
|
dataType = kernel["ProblemType"]["ComputeDataType"]
|
|
globalOffset = 0
|
|
elif tc == 'Bias':
|
|
bps = self.states.bpeCinternal * ss.cfg.gwvw
|
|
rpv = self.states.bpeCinternal * ss.cfg.gwvw / self.states.bpr
|
|
|
|
if kernel["BufferStore"]:
|
|
addr0 = vgpr(addrCalc.addrBiasVgpr)
|
|
addr1 = sgpr("Srd%s"%tc, 4)
|
|
else:
|
|
addr0 = vgpr(addrCalc.addrBiasVgpr,2)
|
|
addr1 = ""
|
|
if ss.optSrdIncForRow and addrCalc.rowInc:
|
|
module.add(addrCalc.incrementToNextRow(kernel, tc, ss, tmpS01, bpeType=self.states.bpeCinternal))
|
|
dataType = kernel["ProblemType"]["ComputeDataType"]
|
|
globalOffset = addrCalc.globalOffsetInternal
|
|
elif tc == 'E':
|
|
bps = self.states.bpeE * ss.cfg.gwvw
|
|
rpv = self.states.bpeE * ss.cfg.gwvw / self.states.bpr
|
|
|
|
if kernel["BufferStore"]:
|
|
addr0 = vgpr(addrCalc.addrEVgpr)
|
|
addr1 = sgpr("Srd%s"%tc, 4)
|
|
else:
|
|
addr0 = vgpr(addrCalc.addrEVgpr,2)
|
|
addr1 = ""
|
|
if ss.optSrdIncForRow and addrCalc.rowInc:
|
|
module.add(addrCalc.incrementToNextRow(kernel, tc, ss, tmpS01, bpeType=self.states.bpeE))
|
|
dataType = kernel["ProblemType"]["DataTypeE"]
|
|
globalOffset = addrCalc.globalOffsetE
|
|
else:
|
|
printExit("Unsupported store tc %s"%tc)
|
|
|
|
useBuffer = kernel["BufferStore"]
|
|
if dataType.isHalf() or dataType.isBFloat16():
|
|
if not kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
# (H,H,H,H,H,H), internal H
|
|
if self.states.asmCaps["HasWMMA_V1"] and kernel["EnableMatrixInstruction"]:
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, hi16=0, comment=comment, scope=scope, th=th, nv=nv))
|
|
else:
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx//2, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, hi16=sumIdx%2, comment=comment, scope=scope, th=th, nv=nv))
|
|
else:
|
|
# (B,B,B,B,S,S), internal S
|
|
# (H,H,H,H,H,H), internal S
|
|
# (H,H,H,H,S,S), internal S
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, hi16=0, comment=comment, scope=scope, th=th, nv=nv))
|
|
elif dataType.isInt32() or dataType.isSingle():
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, comment=comment, scope=scope, th=th, nv=nv))
|
|
elif dataType.isDouble() or dataType.isSingleComplex():
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx*2, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, comment=comment, scope=scope, th=th, nv=nv))
|
|
elif dataType.isDoubleComplex():
|
|
rps = dataType.numRegisters()
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx*rps, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, comment=comment, scope=scope, th=th, nv=nv))
|
|
elif dataType.isInt8() or dataType.isAnyFloat8() or dataType.isAnyBFloat8() or dataType.isAnyFloat8BFloat8() or dataType.isAnyBFloat8Float8():
|
|
if kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(self.chooseGlobalWrite(useBuffer, bps, sumIdx, rpv, \
|
|
addr0, addr1, globalOffset, soffset=wsOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, dlc=isDlc, comment=comment, scope=scope, th=th, nv=nv))
|
|
return module
|
|
|
|
##############################################################################
|
|
# Global Read Input
|
|
##############################################################################
|
|
# readInput: global read into VGPRs for tc in {'C','E','WS', ...}.
|
|
# CompactLoopStore: `elementIdx` is used so elt-0 (rowInc=0) also emits
|
|
# incrementToNextRow as a chain seed (matches addStore). For tc == 'C' we
|
|
# additionally route the chain through s[tmpS01+1] so the C-load primer
|
|
# does not trash the D-store chain (which uses tmpS01).
|
|
def readInput(self, kernel, ss, tc: str, dataType, addrCalc, vc0, data, gwvw, addr, tmpS01,
|
|
elementIdx=None, batchIdx=None, overrideAfterPrimerRows=0):
|
|
module = Module("read%sInput"%tc)
|
|
bps = dataType.numBytes() * gwvw
|
|
useBuffer = kernel["BufferStore"]
|
|
|
|
if kernel["BufferStore"]:
|
|
addr0 = vgpr(addr)
|
|
addr1 = sgpr("Srd%s"%tc, 4)
|
|
else:
|
|
addr0 = vgpr(addr,2)
|
|
addr1 = ""
|
|
|
|
isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal(
|
|
self.states.asmCaps, kernel["NonTemporal%s"%tc], _temporalHint(kernel, tc), _nonVolatile(kernel, tc))
|
|
|
|
soffset = 0
|
|
if tc == 'E':
|
|
globalOffset = addrCalc.globalOffsetE
|
|
bpeType = self.states.bpeE
|
|
elif tc == 'Gate':
|
|
globalOffset = addrCalc.globalOffsetGate
|
|
bpeType = self.states.bpeGate
|
|
elif tc == 'WS':
|
|
soffset = tmpS01
|
|
globalOffset = 0
|
|
bpeType = self.states.bpeCinternal
|
|
else:
|
|
if dataType == kernel["ProblemType"]["ComputeDataType"]:
|
|
globalOffset = addrCalc.globalOffsetInternal
|
|
bpeType = self.states.bpeCinternal
|
|
else:
|
|
globalOffset = addrCalc.globalOffset
|
|
bpeType = self.states.bpeCexternal
|
|
if tc == 'C':
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
globalOffset = int((globalOffset/self.states.bpeCexternal) * self.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters())
|
|
|
|
isWorkspace = tc == 'WS'
|
|
# CompactLoopStore: drop rowInc gate so elt-0 also emits chain seed
|
|
# (preserves the s[stmp] chain for the delayed-primer pattern). tc=='C' +
|
|
# CLS uses tmpS01+1 so the C-load chain primer does not trash D-store's
|
|
# chain (which uses tmpS01). Non-CLS uses tmpS01 for all tcs (= baseline).
|
|
# `forceinitrow0=1` opens the same gate inside incrementToNextRow.
|
|
if (ss.optSrdIncForRow and (addrCalc.rowInc or (kernel["CompactLoopStore"] and elementIdx == 0 and batchIdx == 0))) and not isWorkspace:
|
|
_stmp = (tmpS01 + 1) if (tc == 'C' and kernel["CompactLoopStore"]) else tmpS01
|
|
module.add(addrCalc.incrementToNextRow(kernel, tc, ss, _stmp, forceinitrow0=1, bpeType=bpeType,
|
|
overrideAfterPrimerRows=overrideAfterPrimerRows))
|
|
|
|
if dataType.isHalf():
|
|
hi16 = 0 if (self.states.HHH_WMMA or tc == 'Gate') else (vc0 % 2)
|
|
module.add(self.chooseGlobalRead(useBuffer, bps, data, \
|
|
addr0, addr1, soffset=soffset, offset=globalOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, lds=False, hi16=hi16, \
|
|
comment="load %s"%tc, scope=scope, th=th, nv=nv))
|
|
elif dataType.isInt8() or dataType.is8bitFloat():
|
|
module.add(self.chooseGlobalRead(useBuffer, bps, data, \
|
|
addr0, addr1, soffset=soffset, offset=globalOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, lds=False, \
|
|
#hi16=vc0 % 4,
|
|
comment="load %s"%tc, scope=scope, th=th, nv=nv))
|
|
elif dataType.isBFloat16() or \
|
|
dataType.isInt32() or \
|
|
dataType.isSingle() or \
|
|
dataType.isDouble() or \
|
|
dataType.isSingleComplex() or \
|
|
dataType.isDoubleComplex():
|
|
module.add(self.chooseGlobalRead(useBuffer, bps, data, \
|
|
addr0, addr1, soffset=soffset, offset=globalOffset, \
|
|
glc=isGlc, slc=isSlc, nt=isNT, lds=False, \
|
|
comment="load %s"%tc, scope=scope, th=th, nv=nv))
|
|
|
|
return module
|
|
|
|
##############################################################################
|
|
# Global Write Batch
|
|
##############################################################################
|
|
def globalWriteBatch(self, kernel, tPA, tPB, activation, ss: StoreState, batchIdx, \
|
|
applyAlpha, beta, edge, atomic, gwvw, atomicW, \
|
|
batchElements, addrE, addrD, addrC, addrBias, \
|
|
addrScaleAVec, addrScaleBVec, addrScaleAlphaVec, biasLocalBarrierInit: bool, \
|
|
tmpVgpr, tmpVgprDynamic, cvtVgprStruct, activationSetPCStruct, activationTypeStr, \
|
|
batchElementSgprs, tmpSgpr, codeAccVgprRead, codeMulAlpha, factorDim, numBatches, \
|
|
inter_iter_rowInc=0, direct_next_rowInc=0) -> 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
|