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
1853 lines
139 KiB
Python
1853 lines
139 KiB
Python
################################################################################
|
|
#
|
|
# 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.code import Module
|
|
from rocisa.container import DSModifiers, vgpr, sgpr, SDWAModifiers, VOP3PModifiers, ContinuousRegister
|
|
from rocisa.enum import HighBitSel, SelectBit, InstType
|
|
from rocisa.instruction import SMovB32, SWaitCnt, VOrB32, VPermB32, VLShiftLeftOrB32, \
|
|
VMovB32, VMovB64, VLShiftRightB32, VCvtFP8toF16, VCvtScalePkFP8toF16, VCvtFP8toF32, VCvtScaleFP8toF16, VCvtScalePkFP8toF16, \
|
|
VCvtPkF32toBF16, VCvtBF16toFP32, PVCvtBF16toFP32, VDot2CF32BF16, VSubF32, VSwapB32, MFMAInstruction, \
|
|
ECvtPkFP8toF32, ECvtF32toF16
|
|
|
|
from ..Component import LocalRead
|
|
|
|
from math import ceil
|
|
|
|
class LocalReadVALU(LocalRead):
|
|
kernel = {"EnableMatrixInstruction": False}
|
|
|
|
"""
|
|
Local Read: Do It A/B
|
|
iui = Inner Unroll Idx
|
|
epsi = expand pointer swap index. Only used for PAP
|
|
"""
|
|
def __call__(self, writer, kernel, bufferIdx, iui, epsi, tP):
|
|
tc = tP["tensorChar"]
|
|
|
|
if tc == "A":
|
|
writer.states.localReadDoCntA += 1
|
|
elif tc == "Metadata":
|
|
writer.states.localReadDoCntMetadata += 1
|
|
else:
|
|
writer.states.localReadDoCntB += 1
|
|
|
|
tile01 = tP["tile01Idx"]
|
|
imod = Module("LocalReadDo%s_I%s"%(tc,iui))
|
|
pack = Module("pack%s_I%s"%(tc,iui))
|
|
instruction = tP["localReadInstruction"]
|
|
numOffsets = instruction.numOffsets
|
|
blockWidth = instruction.blockWidth
|
|
offsetMultiplier = 1 # instruction.offsetMultiplier
|
|
valuIdx = 0
|
|
# dot2: currently only support unroll major LDS
|
|
if kernel["UseDotInstruction"]:
|
|
numVectorsPerTile = kernel["ThreadTile%u"%tile01]
|
|
numReadsPerVector = ceil((writer.states.lrvwUnrollA * tP["bpe"]) / (blockWidth*4)) # bytes/register
|
|
LdsPad = kernel["LdsPad%s"%tc] if kernel["LdsBlockSizePerPad%s"%tc] == 0 else 0
|
|
tileStride = kernel["_DepthU%s"%tc] + LdsPad if kernel["UnrollMajorLDS%s" % tP["tensorChar"]] else 1
|
|
else:
|
|
numVectorsPerTile = (kernel["ThreadTile%u"%tile01]//kernel["VectorWidthA"])
|
|
numReadsPerVector = ceil((kernel["VectorWidthA"] * tP["bpe"]) / (blockWidth*4)) # bytes/register
|
|
|
|
for vIdx in range(0, numVectorsPerTile):
|
|
for rIdx in range(0, int(numReadsPerVector)):
|
|
localReadCode = imod.add(Module("LocalRead%s Valu%u"%(tc,valuIdx)))
|
|
paramList = []
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuIdx), blockWidth)
|
|
|
|
# paramList.append(destVgpr)
|
|
# paramList.append(vgpr("LocalReadAddr%s"%tc))
|
|
|
|
for oIdx in range(0, numOffsets):
|
|
# dot2
|
|
if kernel["UseDotInstruction"]:
|
|
paramList.append(int(((rIdx*blockWidth + kernel["SubGroup%u"%tile01] * (vIdx*numOffsets+oIdx) * tileStride \
|
|
+ tP["localReadOffset"]) * tP["bpe"] + tP["localReadSwapByteOffset"]) // offsetMultiplier))
|
|
else:
|
|
paramList.append(int(((rIdx*blockWidth + kernel["SubGroup%u"%tile01] * (vIdx*numOffsets+oIdx)*kernel["VectorWidthA"] \
|
|
+ tP["localReadOffset"]) * tP["bpe"] + tP["localReadSwapByteOffset"]) // offsetMultiplier))
|
|
# print("Debug: Matrix{}, rIdx offset {}, vIdx offset {}, bpe {}, net offset {}".format( \
|
|
# tP["tensorChar"], \
|
|
# rIdx * blockWidth, \
|
|
# kernel["SubGroup%u" % tP["tensorIdx"]] * (vIdx * numOffsets + oIdx) * kernel["VectorWidth"] + tP["localReadOffset"], \
|
|
# tP["bpe"], \
|
|
# paramList[-1]))
|
|
# paramTuple = tuple(paramList)
|
|
num = paramList[0] // 65536
|
|
paramList[0] = paramList[0] - num * 65536
|
|
srcAddr = vgpr("LocalReadAddr%s+%d"%(tc, num))
|
|
|
|
if numOffsets == 1:
|
|
ds = DSModifiers(na=1, offset=paramList[0])
|
|
if numOffsets == 2:
|
|
ds = DSModifiers(na=2, offset0=paramList[0], offset1=paramList[1])
|
|
srcAddr=vgpr("LocalReadAddr%s"%tc)
|
|
LocalReadX = instruction.getInst()
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode)
|
|
valuIdx += blockWidth
|
|
|
|
# TODO - handle vector-load
|
|
with writer.allocTmpSgpr(1, tag="LocalReadVALU_tmpSgprInfo") as tmpSgprInfo:
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
if writer.db["CheckValue1%s" % tc]:
|
|
dbgVgpr = destVgpr
|
|
dbgVgprList = destVgpr.split("v[")
|
|
if len(dbgVgprList) == 1: # vIdx, no []
|
|
dbgVgpr = dbgVgprList[0]
|
|
else:
|
|
# We only check the first one now
|
|
# TODO: Handle vector, but need to take care the last one
|
|
dbgVgprList = (dbgVgprList[1].split("]")[0]).split(':')
|
|
dbgVgpr = "v[%s]"%dbgVgprList[0]
|
|
|
|
# localReadCode.addInst("s_waitcnt lgkmcnt(0)", "CheckValue1 wait for LDS read")
|
|
localReadCode.add(SWaitCnt(dscnt=0, comment="CheckValue1 wait for lds read"))
|
|
if writer.asmCaps["SeparateVscnt"]:
|
|
# localReadCode.addInst( "s_waitcnt_vscnt", -2, "0", "")
|
|
localReadCode.add(SWaitCnt(vlcnt=0, vscnt=0))
|
|
|
|
if kernel["ProblemType"]["DataType"].isHalf():
|
|
localReadCode.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(0x3c003c00), comment="CheckValue1: FP16")) # packed 1s
|
|
localReadCode.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
elif kernel["ProblemType"]["DataType"].isBFloat16():
|
|
localReadCode.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(0x3f803f80), comment="CheckValue1: BF16")) # packed 1s
|
|
localReadCode.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
# TODO - Check if this works
|
|
if kernel["ProblemType"]["DataType"].isInt8():
|
|
localReadCode.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(0x01010101), comment="CheckValue1: INT8")) # packed 1s
|
|
localReadCode.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
# TODO - Check if this works
|
|
elif kernel["ProblemType"]["DataType"].isInt8x4():
|
|
localReadCode.add(writer.assert_eq( dbgVgpr, 1))
|
|
|
|
elif kernel["ProblemType"]["DataType"].isSingle():
|
|
localReadCode.add(writer.assert_eq( dbgVgpr, 1.0) )
|
|
|
|
return imod, pack, Module()
|
|
|
|
class LocalReadMFMA(LocalRead):
|
|
kernel = {"EnableMatrixInstruction": True}
|
|
|
|
# LDS size is increased on gfx950. const offset is still 16-bit.
|
|
# this function handles both LDS size < 64K and LDS size >= 64K
|
|
def cal_offset_srcAddr(self, maxLDSConstOffset, tc, offset):
|
|
num = offset // maxLDSConstOffset
|
|
offset_val = offset - num * maxLDSConstOffset
|
|
srcAddr = vgpr("LocalReadAddr%s+%u" %(tc, num))
|
|
return offset_val, srcAddr
|
|
|
|
@staticmethod
|
|
def getMxsTileSpanInfo(kernel, tc, tile01, asmCaps):
|
|
"""
|
|
MX scale TileSpan scale-select: when MIWaveTile//VectorWidth is a positive even
|
|
multiple, LRA lays out the tile span so a single ds_load holds two scale blocks
|
|
(lower half-wave = block 2g, upper half-wave = partner block 2g+1). The consuming
|
|
WMMA then reads each block directly from that one register via the gfx1250
|
|
matrix_{a,b}_scale:N selector, so N scale ds_loads collapse to N/2.
|
|
Axis-neutral (tile01 selects the tile axis, so this serves either MX scale tensor).
|
|
Return tile-span layout info, or None.
|
|
|
|
This is the ds_load-halving gate and is independent of the wave count: LRA produces the
|
|
block 2g / partner 2g+1 half-wave layout for both MIWaveGroup==1 (non-split: nIdx = wtid)
|
|
and MIWaveGroup>1 (wave-split: nIdx = wtid % MI plus a hi offset). It MUST match
|
|
LraTileAssignment.tileSpan exactly (NOT tileSpanWaveSplit, which additionally requires
|
|
MIWaveGroup>1 only to pick between those two layouts).
|
|
|
|
Constraint: the half-wave scale-select selects lane i vs lane i+halfSpan within the
|
|
wave, so it only maps to this optimization when the two half-wave tile spans meet at
|
|
the wave midpoint, i.e. the tile-axis matrix instruction size equals WavefrontSize/2
|
|
(MatrixInstM for A, MatrixInstN for B).
|
|
"""
|
|
if "MXS" not in tc:
|
|
return None
|
|
# Single source of truth for the whole TileSpan feature. TileSpan is currently only
|
|
# supported on gfx1250: it needs the InMemorySwizzle half-wave scale layout AND the
|
|
# gfx1250 WMMA_V3 matrix_*_scale select. Gate explicitly on the arch (ISA (12,5,0))
|
|
# plus that consumer condition (see KernelWriterAssembly.mxsUsesScaleSel), so the load
|
|
# layout (LocalRead ds_load-halving + LraTileAssignment) never diverges from what the
|
|
# WMMA expects. Any other arch (e.g. gfx1151, where MatrixInst 16 == WavefrontSize/2
|
|
# would otherwise pass the geometry checks) never engages TileSpan.
|
|
if kernel.get("ISA") != (12, 5, 0) \
|
|
or not asmCaps.get("HasWMMA_V3", False) \
|
|
or kernel.get("MXScaleFormat") != "InMemorySwizzle":
|
|
return None
|
|
vectorWidth = kernel["VectorWidth%s" % tc]
|
|
miWaveTile = kernel["MIWaveTile"][tile01]
|
|
miWaveTileVectors = miWaveTile // vectorWidth
|
|
# The tile is handled as (miWaveTileVectors // 2) independent ratio-2 groups. Each
|
|
# group is one ds_load (lower half-wave = block 2g, upper half-wave = partner block
|
|
# 2g+1), so N ds_loads collapse to N/2. (ratio == 2 is the single-group special case.)
|
|
if miWaveTileVectors < 2 or (miWaveTileVectors % 2) != 0:
|
|
return None
|
|
matrixInstT = kernel["MatrixInstM"] if (tile01 == 0) else kernel["MatrixInstN"]
|
|
# The half-wave scale-select selects lane i vs i+halfSpan; only valid when the
|
|
# half-wave boundary is the wave midpoint (halfSpan == WavefrontSize/2).
|
|
if matrixInstT != kernel["WavefrontSize"] // 2:
|
|
return None
|
|
|
|
return {
|
|
"vectorWidth": vectorWidth,
|
|
"numGroups": miWaveTileVectors // 2,
|
|
}
|
|
|
|
# Vreg Value layout (assuming MIInputPerThread = 8)
|
|
# (1) local read dst
|
|
# T: index transpose case (lrvwTile>1 and not transposeCode)
|
|
# X: lrvwTile == 1 or not UseDirect32XEmulation
|
|
# interleave T reg case, (2)
|
|
# local read + lrvwTile == 2
|
|
# valOffset 0 - 1 (+16): X/T reg, 0 - 1 (+16)
|
|
# valOffset 2 - 3 (+16): X/T reg, 8 - 9 (+16)
|
|
# valOffset 4 - 5 (+16): X/T reg, 2 - 3 (+16)
|
|
# valOffset 6 - 7 (+16): X/T reg, 10 - 11 (+16)
|
|
# valOffset 8 - 9 (+16): X/T reg, 4 - 5 (+16)
|
|
# valOffset 10 - 11 (+16): X/T reg, 12 - 13 (+16)
|
|
# valOffset 12 - 13 (+16): X/T reg, 6 - 7 (+16)
|
|
# valOffset 14 - 15 (+16): X/T reg, 14 - 15 (+16)
|
|
# local read + lrvwTile == 4
|
|
# valOffset 0 - 3: X/T reg, 0 - 3
|
|
# valOffset 4 - 7: X/T reg, 8 - 11
|
|
# valOffset 8 - 11: X/T reg, 16 - 19
|
|
# valOffset 12 - 15: X/T reg, 24 - 27
|
|
# valOffset 16 - 19: X/T reg, 4 - 7
|
|
# valOffset 20 - 23: X/T reg, 12 - 15
|
|
# valOffset 24 - 27: X/T reg, 20 - 23
|
|
# valOffset 28 - 31: X/T reg, 28 - 31
|
|
# (2) interleave T reg (src of high value cvt, local read dest)
|
|
# For local read dst, apply this logic on top of above (1)
|
|
# valOffset 0 - 3: T reg, 0 - 3
|
|
# valOffset 4 - 7: X reg, 4 - 7
|
|
# valOffset 8 - 11: T reg, 4 - 7
|
|
# valOffset 12 - 15: X reg, 12 - 15
|
|
# valOffset 16 - 19: T reg, 8 - 11
|
|
# valOffset 20 - 23: X reg, 20 - 23
|
|
# valOffset 24 - 27: T reg, 12 - 15
|
|
# valOffset 28 - 31: X reg, 28 - 31
|
|
# (3) src of high value
|
|
# (3-1) not UseDirect32XEmulation case
|
|
# valOffset 0 - 31: X reg, 0 - 31
|
|
# (3-2) (not 3-1) and interleave T reg
|
|
# same as above (2)
|
|
# (3-3) (not 3-1) and not (interleave T reg) and transpose
|
|
# valOffset 0 - 31: T reg, 0 - 31
|
|
# (4) dst of cvt (high, low)
|
|
# (4-1) not (interleave T reg) and transpose
|
|
# valOffset 0 - 31: T reg, 0 - 31
|
|
# (4-2) (not 4-1)
|
|
# valOffset 0 - 31: X reg, 0 - 31
|
|
# Non transpose case
|
|
# src (dst = False)
|
|
# dst (dst = True)
|
|
# valOffset 0 - 31: X reg, 0 - 31
|
|
# Do index transpose if transpose is true
|
|
def TXInterleaveLayoutIdx(self, idx, miInputPerThread=8):
|
|
halfM = miInputPerThread // 2
|
|
retIdx = idx
|
|
XTchar = "X"
|
|
if idx % miInputPerThread < halfM:
|
|
XTchar = "T"
|
|
retIdxUpper = (idx // miInputPerThread) * halfM
|
|
retIdxLower = idx % halfM
|
|
retIdx = retIdxUpper + retIdxLower
|
|
return retIdx, XTchar
|
|
|
|
@staticmethod
|
|
def _genDsReadConvTable(miInputPerThread, lrvwTile):
|
|
halfM = miInputPerThread // 2
|
|
numRows = lrvwTile
|
|
colsPerRow = halfM // lrvwTile
|
|
table = []
|
|
for group in range(2):
|
|
for col in range(colsPerRow):
|
|
for row in range(numRows):
|
|
table.append(row * miInputPerThread + col * lrvwTile + group * halfM)
|
|
return table
|
|
|
|
def getTransposeXorTIndex(self, writer, kernel, idx, tc, lrvwTile, dst, isNext, localRead=False):
|
|
mipt = kernel["MIInputPerThread"]
|
|
assert lrvwTile <= 4, "lrvwTile is %d"%lrvwTile
|
|
abmatrixinfo = writer.states.a if tc == 'A' else writer.states.b
|
|
if isNext:
|
|
useTransposeCode = abmatrixinfo.useTransposeCodeNext
|
|
useDirect32XEmulation = abmatrixinfo.useDirect32XEmulationNext
|
|
else:
|
|
useTransposeCode = abmatrixinfo.useTransposeCodeThis
|
|
useDirect32XEmulation = abmatrixinfo.useDirect32XEmulationThis
|
|
TF32EmuInterleaveTreg = abmatrixinfo.TF32EmuInterleaveTreg
|
|
swapBlockSizeSub = abmatrixinfo.swapBlockSizeSub
|
|
blockH = mipt
|
|
blockW = lrvwTile
|
|
blockSize = blockH * blockW
|
|
# default: no change
|
|
XTchar = "X"
|
|
if localRead:
|
|
if lrvwTile > 1:
|
|
dsReadConvTable = self._genDsReadConvTable(mipt, lrvwTile)
|
|
blockIdx = idx // blockSize
|
|
idxInBlk = idx % blockSize
|
|
tableIdx = idxInBlk // blockW
|
|
idx = dsReadConvTable[tableIdx] + blockIdx * blockSize
|
|
if TF32EmuInterleaveTreg:
|
|
if writer.states.doFullPackCodePrefetch:
|
|
idx, XTchar = self.TXInterleaveLayoutIdx(idx, kernel["MIInputPerThread"])
|
|
else:
|
|
withinGroup = idx % 8
|
|
if withinGroup < 4:
|
|
XTchar = "T"
|
|
idx = (idx // 8) * 4 + withinGroup
|
|
elif writer.states.doFullPackCodePrefetch:
|
|
# full pack code prefetch case, dst of local read is always T reg
|
|
# use T as destination
|
|
XTchar = "T"
|
|
# always use same swapBlockSizeSub(=0)
|
|
idx = idx % swapBlockSizeSub
|
|
return idx, XTchar
|
|
if not dst:
|
|
# src case
|
|
if (not useDirect32XEmulation):
|
|
# temp reg case (both lrvwTile==1 and >1)
|
|
# no register conversion (use X 0-31)
|
|
pass
|
|
elif TF32EmuInterleaveTreg:
|
|
if writer.states.doFullPackCodePrefetch:
|
|
idx, XTchar = self.TXInterleaveLayoutIdx(idx, kernel["MIInputPerThread"])
|
|
else:
|
|
withinGroup = idx % 8
|
|
if withinGroup < 4:
|
|
XTchar = "T"
|
|
idx = (idx // 8) * 4 + withinGroup
|
|
elif writer.states.doFullPackCodePrefetch:
|
|
# full pack code prefetch case, dst of local read is always T reg
|
|
XTchar = "T" # use T reg for wider local read + transpose code
|
|
# always use same swapBlockSizeSub(=0)
|
|
idx = idx % swapBlockSizeSub
|
|
if lrvwTile > 1 and not useTransposeCode:
|
|
# do index transpose for src only
|
|
idx = self.getTransposeIndex(kernel, idx, lrvwTile)
|
|
else:
|
|
# dst case
|
|
# always "X"
|
|
pass
|
|
return idx, XTchar
|
|
|
|
# transpose vgprs for lrvwTile > 1
|
|
# From
|
|
# lrvwTile==2
|
|
# reg 0 reg 1 reg 16 reg 17
|
|
# reg 8 reg 9 reg 24 reg 25
|
|
# reg 2 reg 3 reg 18 reg 19
|
|
# reg 10 reg 11 reg 26 reg 27
|
|
# reg 4 reg 5 reg 20 reg 21
|
|
# reg 12 reg 13 reg 28 reg 29
|
|
# reg 6 reg 7 reg 22 reg 23
|
|
# reg 14 reg 15 reg 30 reg 31
|
|
# lrvwTile==4
|
|
# reg 0 reg 1 reg 2 reg 3
|
|
# reg 8 reg 9 reg 10 reg 11
|
|
# reg 16 reg 17 reg 18 reg 19
|
|
# reg 24 reg 25 reg 26 reg 27
|
|
# reg 4 reg 5 reg 6 reg 7
|
|
# reg 12 reg 13 reg 14 reg 15
|
|
# reg 20 reg 21 reg 22 reg 23
|
|
# reg 28 reg 29 reg 30 reg 31
|
|
# To (common for lrvwTile==2 and 4)
|
|
# reg 0 reg 8 reg 16 reg 24
|
|
# reg 1 reg 9 reg 17 reg 25
|
|
# reg 2 reg 10 reg 18 reg 26
|
|
# reg 3 reg 11 reg 19 reg 27
|
|
# reg 4 reg 12 reg 20 reg 28
|
|
# reg 5 reg 13 reg 21 reg 29
|
|
# reg 6 reg 14 reg 22 reg 30
|
|
# reg 7 reg 15 reg 23 reg 31
|
|
def getTransposeIndex(self, kernel, idx, lrvwTile):
|
|
if lrvwTile == 1:
|
|
return idx
|
|
mipt = kernel["MIInputPerThread"]
|
|
assert lrvwTile <= 4, "lrvwTile is %d"%lrvwTile
|
|
blockSize = mipt * lrvwTile
|
|
blockIdx = idx // blockSize
|
|
idxInBlk = idx % blockSize
|
|
dsTable = self._genDsReadConvTable(mipt, lrvwTile)
|
|
swappedIdx = dsTable[idxInBlk % mipt] + (idxInBlk // mipt) + blockIdx * blockSize
|
|
return swappedIdx
|
|
|
|
# if prefetch is for NextLoop or not
|
|
def getIsNext(self, writer, kernel, tc, bufferIdx):
|
|
numIterPerCoalescedRead = writer.states.numIterPerCoalescedReadA if tc == "A" else writer.states.numIterPerCoalescedReadB
|
|
isNext = bufferIdx < writer.states.numItersPLR * numIterPerCoalescedRead
|
|
# subIter case
|
|
if kernel["ForceUnrollSubIter"]:
|
|
isNext = writer.states.SubTileIdx == 0
|
|
return isNext
|
|
|
|
# Get vgpr string for Emu from mfmaIter
|
|
def getVgprStrForEmuMfma(self, writer, kernel, tc, bufferIdx, iui, valOffset, lrvwTile, u, dst=False, localRead=False):
|
|
# subIter case
|
|
if kernel["ForceUnrollSubIter"]:
|
|
# isNext means using prefetch local read for next iter
|
|
# MFMA case, isNext means smaller vreg index
|
|
if tc == "A":
|
|
isNext = (u & 1) == 0
|
|
else:
|
|
isNext = u < 2
|
|
else:
|
|
isNext = self.getIsNext(writer, kernel, tc, bufferIdx)
|
|
v0Idx, v0XTchar = self.getTransposeXorTIndex(writer, kernel, valOffset, tc, lrvwTile, dst, isNext, localRead=localRead)
|
|
return "Valu%s_%s%u_I%u+%u"%(tc, v0XTchar, bufferIdx, iui, v0Idx)
|
|
# Get vgpr string for Emu
|
|
def getVgprStrForEmu(self, writer, kernel, tc, bufferIdx, iui, valOffset, lrvwTile, dst=False, localRead=False):
|
|
isNext = self.getIsNext(writer, kernel, tc, bufferIdx)
|
|
v0Idx, v0XTchar = self.getTransposeXorTIndex(writer, kernel, valOffset, tc, lrvwTile, dst, isNext, localRead=localRead)
|
|
return "Valu%s_%s%u_I%u+%u"%(tc, v0XTchar, bufferIdx, iui, v0Idx)
|
|
# Get vgpr for Emu
|
|
def getVgprForEmu(self, writer, kernel, tc, bufferIdx, iui, valOffset, lrvwTile, vgprLen=1, dst=False, localRead=False):
|
|
v0 = vgpr(self.getVgprStrForEmu(writer, kernel, tc, bufferIdx, iui, valOffset, lrvwTile, dst=dst, localRead=localRead), vgprLen)
|
|
return v0
|
|
def get2VgprForEmu(self, writer, kernel, tc, bufferIdx, valOffset, iui, lrvwTile, vgprLen=1, dst=False, localRead=False):
|
|
# create 4 continuous vgpr set
|
|
v0 = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, valOffset + 0 * vgprLen, lrvwTile, vgprLen=vgprLen, dst=dst, localRead=localRead)
|
|
v1 = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, valOffset + 1 * vgprLen, lrvwTile, vgprLen=vgprLen, dst=dst, localRead=localRead)
|
|
return v0, v1
|
|
def get4VgprForEmu(self, writer, kernel, tc, bufferIdx, valOffset, iui, lrvwTile, vgprLen=1, dst=False, localRead=False):
|
|
# create 4 continuous vgpr set
|
|
v0 = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, valOffset + 0 * vgprLen, lrvwTile, vgprLen=vgprLen, dst=dst, localRead=localRead)
|
|
v1 = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, valOffset + 1 * vgprLen, lrvwTile, vgprLen=vgprLen, dst=dst, localRead=localRead)
|
|
v2 = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, valOffset + 2 * vgprLen, lrvwTile, vgprLen=vgprLen, dst=dst, localRead=localRead)
|
|
v3 = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, valOffset + 3 * vgprLen, lrvwTile, vgprLen=vgprLen, dst=dst, localRead=localRead)
|
|
return v0, v1, v2, v3
|
|
|
|
# do transpose with v_swap
|
|
def transposeLRVregs(self, kernel, module, tc, bufferIdx, iui, writer, lrvwTile, totalRegs, subTileIdx):
|
|
start = subTileIdx * totalRegs
|
|
last = start + totalRegs - 1
|
|
if lrvwTile == 1:
|
|
return module
|
|
done = [start, last]
|
|
for idx in range(start + 1,last):
|
|
vSrc = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, idx, 1, localRead=True)
|
|
newIdx = self.getTransposeIndex(kernel, idx, lrvwTile)
|
|
if idx in done or idx == newIdx:
|
|
# Already done or same vreg. No need to swap
|
|
done.append(idx)
|
|
continue
|
|
vDst = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, newIdx, 1, localRead=True)
|
|
commentStr = ""
|
|
if idx == last - 1:
|
|
commentStr = "__PACK_PRE_" + tc + "_I%d__"%(iui)
|
|
module.add(VSwapB32(dst=vDst, src=vSrc, comment="swap %d and %d"%(newIdx, idx) + commentStr))
|
|
done.append(idx)
|
|
done.append(newIdx)
|
|
|
|
def initTmpVregForPack(self, kernel, writer, tct, index, bufferIdx, baseValuiIdx, iui, module, lrvwTile, tmpvgpr, useDirect32XEmulation):
|
|
valOffset = baseValuiIdx + index
|
|
src0, src1 = self.get2VgprForEmu(writer, kernel, tct, bufferIdx, valOffset, iui, lrvwTile, vgprLen=2, dst=False)
|
|
if index % 8 == 0: # First half of registers use tmp registers
|
|
if (not useDirect32XEmulation):
|
|
# First half of the registers will be overwritten. Store FP32 values in tmp
|
|
overlap = True
|
|
while overlap == True:
|
|
val1 = writer.vgprPool.checkOutAligned(2, 2, "x32f temp val1")
|
|
val2 = writer.vgprPool.checkOutAligned(2, 2, "x32f temp val2")
|
|
overlap = False
|
|
if baseValuiIdx not in writer.states.tmpvgpr:
|
|
writer.states.tmpvgpr[baseValuiIdx] = []
|
|
if tct == "B":
|
|
overlap |= val1 in writer.states.tmpvgpr[baseValuiIdx]
|
|
overlap |= val2 in writer.states.tmpvgpr[baseValuiIdx]
|
|
tmpvgpr.append(val1)
|
|
tmpvgpr.append(val2)
|
|
module.add(VMovB64(dst=vgpr(val1, 2), src=src0))
|
|
module.add(VMovB64(dst=vgpr(val2, 2), src=src1))
|
|
|
|
def releaseTmpVregForPack(self, kernel, writer, tc, baseValuiIdx, tmpvgprFP32, useDirect32XEmulation):
|
|
if not (kernel["MatrixInstM"] == 16 and kernel["MatrixInstK"] == 16):
|
|
if not useDirect32XEmulation:
|
|
while len(tmpvgprFP32):
|
|
tmp = tmpvgprFP32.pop()
|
|
writer.vgprPool.checkIn(tmp)
|
|
# if A, write tmp vgpr to writer state
|
|
if tc == "A":
|
|
if baseValuiIdx not in writer.states.tmpvgpr:
|
|
writer.states.tmpvgpr[baseValuiIdx] = []
|
|
if tmp not in writer.states.tmpvgpr[baseValuiIdx]:
|
|
writer.states.tmpvgpr[baseValuiIdx].append(tmp)
|
|
# if B, free tmp vgpr from writer state
|
|
elif tc == "B":
|
|
if baseValuiIdx in writer.states.tmpvgpr:
|
|
writer.states.tmpvgpr[baseValuiIdx] = []
|
|
|
|
# 8 registers were read in fp32. Write to 2 of 4 registers with high bits
|
|
# Input: vgprValuA/B_X/T_ i ... i + 7 -> FP32
|
|
# Output: vgprValuA/B_X_ i .. i + 3 -> BF16 High bits
|
|
def pack4HiBits(self, kernel, writer, tct, index, bufferIdx, baseValuiIdx, iui, module, lrvwTile, commentForSchedule1, useDirect32XEmulation):
|
|
valOffset = baseValuiIdx + index
|
|
dstValOffset = baseValuiIdx + index // 2
|
|
v0, v1, v2, v3 = self.get4VgprForEmu(writer, kernel, tct, bufferIdx, valOffset, iui, lrvwTile, dst=False)
|
|
dst0, dst1 = self.get2VgprForEmu(writer, kernel, tct, bufferIdx, dstValOffset, iui, lrvwTile, dst=True)
|
|
module.add(VCvtPkF32toBF16(dst=dst0, src0=v0, src1=v1))
|
|
commentStr = commentForSchedule1
|
|
# do not put comment for scheduling in the following cases
|
|
# - UseMFMAF32XEmulation
|
|
# Delay the comment to the 2nd mfma
|
|
# - index % 8 == 0
|
|
# we need to put schedule comment at index % 8 == 4
|
|
# - not useDirect32XEmulation
|
|
# tmp vreg case, tmp vreg can be same between A and B
|
|
# put schedule comment only at final pack
|
|
#if kernel["UseMFMAF32XEmulation"] or (index % 8 == 0) or (not useDirect32XEmulation):
|
|
# commentStr = ""
|
|
module.add(VCvtPkF32toBF16(dst=dst1, src0=v2, src1=v3, comment=commentStr))
|
|
|
|
def pack4LowBitsStep1(self, kernel, writer, tc, valuiIdx, bufferIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, commentForSchedule1, useDirect32XEmulation):
|
|
baseValuiIdx = valuiIdx - valuiIdx % 8
|
|
if (valuiIdx % 8) == 0 and not useDirect32XEmulation:
|
|
# get 4 tmp vgprs
|
|
tmpIdx = len(tmpvgprFP32) - 2
|
|
vTBase = tmpvgprFP32[tmpIdx + 0] # str(v0t.regName) does not work for tmp vreg. Put vgpr index directly here
|
|
v0t = vgpr(tmpvgprFP32[tmpIdx + 0])
|
|
v1t = vgpr(tmpvgprFP32[tmpIdx] + 1)
|
|
v2t = vgpr(tmpvgprFP32[tmpIdx + 1])
|
|
v3t = vgpr(tmpvgprFP32[tmpIdx + 1] + 1)
|
|
else:
|
|
v0t, v1t, v2t, v3t = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, valuiIdx, iui, lrvwTile)
|
|
vTBase = str(v0t.regName)
|
|
srcVIdxHi = (valuiIdx-baseValuiIdx)//2 + baseValuiIdx
|
|
# vHi0, vHi1 is dst for Upper cvt
|
|
vHi0, vHi1 = self.get2VgprForEmu(writer, kernel, tc, bufferIdx, srcVIdxHi, iui, lrvwTile, dst=True)
|
|
# interleaveTreg case (means dst and src are same (dot2 or mfma))
|
|
if kernel["UseMFMAF32XEmulation"]:
|
|
vHiBase = str(vHi0.regName)
|
|
idMat = vgpr(writer.states.startVgprIdentityMatrix,2)
|
|
# We use a single MFMA 4x4x4_16b to perform 4 `vT - vHi` operations.
|
|
# - A is set to negative identity matrix
|
|
# - no need for DPP as B has the same layout as C & D
|
|
commentStr = ""
|
|
if commentForSchedule1 != "":
|
|
# add comment for scheduling (__TF32_1_)
|
|
# UseMFMAF32XEmulation case, put the mark at second mfma4x4. Skip here
|
|
commentStr = commentForSchedule1 + ": "
|
|
packCodeT.add(MFMAInstruction(instType=InstType.INST_BF16, accType=InstType.INST_F32, variant=[4,4,4,16], mfma1k=False,acc=vgpr(vTBase,4), a=idMat, b=vgpr(vHiBase,2), acc2=vgpr(vTBase,4),
|
|
comment="Calculate low bits for TF32 emulation%s"%commentStr))
|
|
else:
|
|
# Use per-tensor dedicated temp VGPR if available (avoids shared-temp corruption when SIA3 interleaves A and B pack code).
|
|
# Fall back to pool allocation for architectures that don't set it.
|
|
abmatrixinfo = writer.states.a if tc == 'A' else writer.states.b
|
|
useDedicatedTmp = abmatrixinfo.tmpVgprCvtSub >= 0
|
|
tmp = abmatrixinfo.tmpVgprCvtSub if useDedicatedTmp else writer.vgprPool.checkOut(1, "x32f tmp mod 4")
|
|
# Compute low bits = fp32(highBF16(A/B)) - fp32(A/B)
|
|
if kernel["UseDot2F32XEmulation"]:
|
|
packCodeT.add(VDot2CF32BF16(dst=v0t, src0=hex(0x8000bf80), src1=vHi0))
|
|
packCodeT.add(VDot2CF32BF16(dst=v1t, src0=hex(0xbf800000), src1=vHi0))
|
|
else:
|
|
packCodeT.add(PVCvtBF16toFP32(dst=vgpr(tmp), src=vHi0, comment="begin"+str(valuiIdx)))
|
|
packCodeT.add(VSubF32(dst=v0t, src0=v0t, src1=vgpr(tmp)))
|
|
packCodeT.add(VCvtBF16toFP32(dst=vgpr(tmp), src=vHi0, vgprMask=None, vi=1))
|
|
packCodeT.add(VSubF32(dst=v1t, src0=v1t, src1=vgpr(tmp)))
|
|
if kernel["UseDot2F32XEmulation"] and (valuiIdx % 8) == 0:
|
|
packCodeT.add(VDot2CF32BF16(dst=v2t, src0=hex(0x8000bf80), src1=vHi1))
|
|
packCodeT.add(VDot2CF32BF16(dst=v3t, src0=hex(0xbf800000), src1=vHi1))
|
|
else:
|
|
# We use cvt+sub pair since dot2 requires adding 4 wait states.
|
|
packCodeT.add(PVCvtBF16toFP32(dst=vgpr(tmp), src=vHi1))
|
|
packCodeT.add(VSubF32(dst=v2t, src0=v2t, src1=vgpr(tmp)))
|
|
packCodeT.add(VCvtBF16toFP32(dst=vgpr(tmp), src=vHi1, vgprMask=None, vi=1))
|
|
packCodeT.add(VSubF32(dst=v3t, src0=v3t, src1=vgpr(tmp), comment="end"))
|
|
if not useDedicatedTmp:
|
|
writer.vgprPool.checkIn(tmp)
|
|
|
|
def pack4LowBitsFinal(self, kernel, writer, tc, valuiIdx, bufferIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, useDirect32XEmulation, noComment=False):
|
|
baseValuiIdx = valuiIdx - valuiIdx % 8
|
|
# on last iteration, store lower bits in last 4 registers
|
|
if not (kernel["MatrixInstM"] == 16 and kernel["MatrixInstK"] == 16):
|
|
if useDirect32XEmulation:
|
|
v0, v1, v2, v3 = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, baseValuiIdx, iui, lrvwTile)
|
|
else:
|
|
#v0, v1, v2, v3 = self.get4TmpVgprForEmu(writer, kernel, tc, bufferIdx, baseValuiIdx, iui)
|
|
tmpIdx = len(tmpvgprFP32) - 2
|
|
v0 = vgpr(tmpvgprFP32[tmpIdx + 0])
|
|
v1 = vgpr(tmpvgprFP32[tmpIdx + 0] + 1)
|
|
v2 = vgpr(tmpvgprFP32[tmpIdx + 1])
|
|
v3 = vgpr(tmpvgprFP32[tmpIdx + 1] + 1)
|
|
v4, v5, v6, v7 = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, baseValuiIdx + 4, iui, lrvwTile)
|
|
# dst regs
|
|
# v4 and v4d should be different in useTransposeCode=False case
|
|
v4d, v5d, v6d, v7d = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, baseValuiIdx + 4, iui, lrvwTile, dst=True)
|
|
# remove s_nop by scheduling mfma4x4 right after first 4 cvt for high
|
|
# we will have mfma hi * hi + 4 final cvt A between mfma4x4 B and final conv B
|
|
# (more instructions between mfma4x4 B and final conv A)
|
|
# however, we still need nop in tail loop case (A, B scheduling is not interleaved in tail loop) and CMS
|
|
# we still need s_nop in main loop, NGLL, NLL if number MIWaveTileA/B are different.
|
|
# s_nop insertion is handled in _interleavePackAB
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v7d, src0=v6, src1=v7, comment="pack final begin"))
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v6d, src0=v4, src1=v5))
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v5d, src0=v2, src1=v3))
|
|
commentStr = "" if noComment else "__TF32_2_" + tc + "_%d pack final end"%(baseValuiIdx//8)
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v4d, src0=v0, src1=v1, comment=commentStr))
|
|
|
|
def localReadMX(self, writer, kernel, bufferIdx, iui, epsi, tP):
|
|
tc = tP["tensorChar"]
|
|
mxTc = tc[3]
|
|
imod = Module("LocalReadDo%s_I%s" % (tc,iui))
|
|
pack = Module("pack%s_I%s"%(tc,iui))
|
|
packPre = Module("pack%s_I%s Pre"%(tc,iui))
|
|
|
|
tile01 = tP["tile01Idx"]
|
|
instruction = tP["localReadInstruction"]
|
|
bpr = 4 # bytes/register
|
|
|
|
vectorWidth = kernel["VectorWidth%s"%tc]
|
|
mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{mxTc}"]
|
|
stridePerRead = instruction.blockWidth * bpr
|
|
tilePerRead = stridePerRead // mxUnit
|
|
MIWaveGroupShape = [ kernel["MatrixInstM"] * kernel["MatrixInstBM"] * kernel["MIWaveGroup"][0] * kernel["VectorWidthA"], \
|
|
kernel["MatrixInstN"] * kernel["MatrixInstBN"] * kernel["MIWaveGroup"][1] * kernel["VectorWidthB"]]
|
|
tileSpanInfo = self.getMxsTileSpanInfo(kernel, tc, tile01, writer.states.asmCaps)
|
|
mxsTileSpan = tileSpanInfo is not None
|
|
numVectorsPerTile = tileSpanInfo["numGroups"] if mxsTileSpan else kernel["MIWaveTile"][tile01] // vectorWidth
|
|
numReadsPerVector = int(vectorWidth // tilePerRead)
|
|
numVgpr = int(ceil(instruction.blockWidth))
|
|
|
|
valufIdx = 0
|
|
for vIdx in range(0, numVectorsPerTile):
|
|
tileSpanBaseValuiIdx = valufIdx
|
|
if mxsTileSpan:
|
|
localReadCode = imod.add(Module("LocalRead%s Valu%u"%(tc, tileSpanBaseValuiIdx)))
|
|
for eIdx in range(0, numReadsPerVector):
|
|
if mxsTileSpan:
|
|
valuiIdx = tileSpanBaseValuiIdx + eIdx * numVgpr
|
|
readModule = localReadCode
|
|
else:
|
|
valuiIdx = int(valufIdx)
|
|
readModule = imod.add(Module("LocalRead%s Valu%u"%(tc, valuiIdx)))
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx), numVgpr)
|
|
|
|
# load read instruction
|
|
paramList = []
|
|
|
|
offset_val = eIdx * stridePerRead
|
|
# Each tile-span group spans 2 scale blocks (the lower half-wave reads block
|
|
# 2*vIdx; the upper half-wave grabs partner block 2*vIdx+1 via the LRA
|
|
# wave-split), so the per-group stride is twice the per-vector stride.
|
|
blockStep = (2 * vIdx) if mxsTileSpan else vIdx
|
|
offset_val = offset_val + blockStep * MIWaveGroupShape[tile01] * mxUnit
|
|
offset_val = offset_val + tP["localReadOffset"]
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val = int(offset_val + (offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
offset_val = offset_val + tP["localReadSwapByteOffset"]
|
|
|
|
paramList.append(int(offset_val))
|
|
|
|
comment = "L -> Reg for MX"
|
|
|
|
addrIdx = paramList[0] // 65536
|
|
srcAddr=vgpr("LocalReadAddr%s+%u"%(tc, addrIdx))
|
|
paramList[0] -= addrIdx * 65536
|
|
|
|
ds = DSModifiers(na=1, offset=paramList[0])
|
|
LocalReadX = instruction.getInst()
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=readModule, comment=comment)
|
|
if not mxsTileSpan:
|
|
valufIdx += numVgpr
|
|
if mxsTileSpan:
|
|
# No permute: the partner block (upper half-wave of this single wave-split
|
|
# ds_load) is consumed directly by the WMMA via matrix_{a,b}_scale:1 (see
|
|
# mxsTileSpanScaleSel). Each group owns 2*vectorWidth scale *blocks*, but only
|
|
# the lower half (vectorWidth blocks) is actually loaded, spanning
|
|
# numReadsPerVector*numVgpr registers; the partner half is never loaded.
|
|
# We therefore pack the loaded groups contiguously (per-group stride ==
|
|
# lowerHalfSpan, not 2*lowerHalfSpan) and drop the upper-half vgprs entirely.
|
|
# mxsTileSpanScaleSel maps each logical block to this compacted register, and
|
|
# the MXS scale valu allocation is halved to match (see KernelWriter).
|
|
lowerHalfSpan = numReadsPerVector * numVgpr
|
|
valufIdx = tileSpanBaseValuiIdx + lowerHalfSpan
|
|
|
|
return imod, pack, packPre
|
|
|
|
|
|
"""
|
|
Local Read: Do It A/B
|
|
iui = Inner Unroll Idx
|
|
epsi = expand pointer swap index. Only used for PAP
|
|
"""
|
|
def __call__(self, writer, kernel, bufferIdx, iui, epsi, tP):
|
|
imod = Module("LocalReadDo%s_I%s" % (tP["tensorChar"],iui))
|
|
subTileIdx = writer.states.SubTileIdx if kernel["ForceUnrollSubIter"] else 0 # use SubTileIdx only for ForceUnrollSubIter
|
|
|
|
tc = tP["tensorChar"]
|
|
if tc == "A":
|
|
writer.states.localReadDoCntA += 1
|
|
elif tc == "MXSA":
|
|
writer.states.localReadDoCntMXSA += 1
|
|
elif tc == "Metadata":
|
|
writer.states.localReadDoCntMetadata += 1
|
|
elif tc == "B":
|
|
writer.states.localReadDoCntB += 1
|
|
elif tc == "MXSB":
|
|
writer.states.localReadDoCntMXSB += 1
|
|
else:
|
|
raise Exception(f"unsupport tc %s{tc}")
|
|
|
|
isgfx950 = kernel["ISA"][:2] == (9, 5)
|
|
isgfx950mx = isgfx950 and ("MXS" in tc)
|
|
|
|
if ("MXS" in tc and writer.states.asmCaps["HasWMMA_V3"]
|
|
and kernel["MXScaleFormat"] == "InMemorySwizzle"):
|
|
return self.localReadMX(writer, kernel, bufferIdx, iui, epsi, tP)
|
|
|
|
MacDataType = f"MacDataType{tc}" if(tc=="A" or tc=="B") else "DataType"
|
|
tile01 = tP["tile01Idx"]
|
|
instruction = tP["localReadInstruction"]
|
|
bpr = 4 # bytes/register
|
|
|
|
numOffsets = instruction.numOffsets
|
|
blockWidth = instruction.blockWidth
|
|
unrollBlockWidth = instruction.blockWidth if kernel["UnrollMajorLDS%s"%tc] else tP["bpeDS"]/4
|
|
tileBlockWidth = tP["bpeDS"]/4 if kernel["UnrollMajorLDS%s"%tc] else instruction.blockWidth
|
|
|
|
vectorWidth = kernel["VectorWidth%s"%tc]
|
|
|
|
numSubTiles = kernel["numSubTiles"]
|
|
MIWaveGroupShape = [ kernel["MatrixInstM"] * kernel["MatrixInstBM"] * kernel["MIWaveGroup"][0] * kernel["VectorWidthA"], \
|
|
kernel["MatrixInstN"] * kernel["MatrixInstBN"] * kernel["MIWaveGroup"][1] * kernel["VectorWidthB"]]
|
|
|
|
LdsPad = kernel["LdsPad%s"%tc] if kernel["LdsBlockSizePerPad%s"%tc] == 0 else 0
|
|
tileStride = 1
|
|
UnrollStride = kernel["MacroTile%s" % tP["tensorChar"]] + LdsPad
|
|
if kernel["UnrollMajorLDS%s" % tP["tensorChar"]]:
|
|
tileStride = kernel["_DepthU%s"%tc] + LdsPad
|
|
UnrollStride = 1
|
|
|
|
if "MXS" in tc:
|
|
subTc = tc[3]
|
|
mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"]
|
|
# MX scale LDS strides, gated by MXScaleFormat:
|
|
# - Swizzled (HostPreSwizzle/InMemorySwizzle):
|
|
# tileStride = mxUnit
|
|
# UnrollStride = MT * mxUnit (M-blocks interleaved on K)
|
|
# - NoSwizzle (canonical): LDS layout follows UnrollMajorLDS<tc>,
|
|
# which is mirrored from UnrollMajorLDS<A/B> = not TLU<A/B>:
|
|
# UMLDS=1 (K-major LDS): addr(m,k) = k + m*(_DepthU_MXS + LdsPad)
|
|
# tileStride = _DepthU_MXS + LdsPad
|
|
# UnrollStride = mxUnit
|
|
# UMLDS=0 (M-major LDS): addr(m,k) = m + k*(MT + LdsPad)
|
|
# tileStride = 1
|
|
# UnrollStride = MT + LdsPad
|
|
mxScaleFormat = kernel.get("MXScaleFormat", "NoSwizzle")
|
|
isMxSwizzled = mxScaleFormat in ("InMemorySwizzle", "HostPreSwizzle")
|
|
if isMxSwizzled:
|
|
tileStride = mxUnit
|
|
UnrollStride = kernel["MacroTile%s" % tP["tensorChar"]] * mxUnit
|
|
elif kernel["UnrollMajorLDS%s" % tP["tensorChar"]]:
|
|
tileStride = kernel["_DepthU%s" % tc] + LdsPad
|
|
UnrollStride = mxUnit
|
|
else:
|
|
tileStride = 1
|
|
UnrollStride = kernel["MacroTile%s" % tP["tensorChar"]] + LdsPad
|
|
|
|
enableLDSTr = tP["enableLDSTr"]
|
|
matrixInstT = kernel["MatrixInstM"] if (tile01 == 0) else kernel["MatrixInstN"]
|
|
matrixInstTO = min(kernel["MatrixInstM"], kernel["MatrixInstN"])
|
|
matrixInstTO = matrixInstT if ("MXS" in tc) else matrixInstTO
|
|
numTilePerInst = matrixInstT // matrixInstTO
|
|
MIInputPerThUnroll = kernel["MIInputPerThread%s"%tc] // numTilePerInst
|
|
numVectorsPerTile = kernel["MIWaveTile"][tile01] // vectorWidth
|
|
numReadsPerVector = int((vectorWidth * tP["bpeDS"]) / (tileBlockWidth * bpr))
|
|
# overloading numReadsPerUnroll for DirectToLds x2/x4 case when blockWidth of instruction < LocalReadVectorWidth
|
|
# fp64 TLU=1 reading 0.5element/lane/read..
|
|
# for TLU=0 case, blockWidth and LRVW should match
|
|
numReadsPerUnroll = ceil(tP["bpeDS"] * MIInputPerThUnroll / (unrollBlockWidth * bpr))
|
|
numVgpr = int(ceil(blockWidth))
|
|
tmpvgprFP32 = []
|
|
|
|
if tc == 'A':
|
|
lrvwTile = writer.states.lrvwTileA
|
|
abmatrixinfo = writer.states.a
|
|
elif tc == "MXSA":
|
|
lrvwTile = writer.states.lrvwTileMXSA
|
|
abmatrixinfo = writer.states.mxsa
|
|
elif tc == 'B':
|
|
lrvwTile = writer.states.lrvwTileB
|
|
abmatrixinfo = writer.states.b
|
|
elif tc == "MXSB":
|
|
lrvwTile = writer.states.lrvwTileMXSB
|
|
abmatrixinfo = writer.states.mxsb
|
|
elif tc == "Metadata":
|
|
lrvwTile = writer.states.lrvwTileMetadata
|
|
abmatrixinfo = writer.states.m
|
|
else:
|
|
raise Exception(f"unsupport tc %s{tc}")
|
|
|
|
numElementPerRead = 1 if kernel["ConvertAfterDS"] and not kernel["UseF32XEmulation"] else int(int(blockWidth * bpr) // tP['bpe'] // lrvwTile)
|
|
perpStride = abmatrixinfo.gNLCPerpStride
|
|
|
|
# pack register
|
|
if writer.states.archCaps["HasEccHalf"] or not writer.states.asmCaps["HasWMMA_V1"]:
|
|
needPack = tP["bpeDS"] < 4 and not kernel["UnrollMajorLDS%s"%tc] and not tP["isM"]
|
|
# specify I8 for the case that input number is equal to the localread blockwidth but need to split low and high bytes to different vgprs.
|
|
needPackMetadata = tP["isM"] and ((MIInputPerThUnroll * tP["bpeDS"] / (blockWidth * 4) > 1) or (kernel["ProblemType"][MacDataType].numBytes() == 1 and writer.states.lrvwTileMetadata > 1))
|
|
needPack |= needPackMetadata
|
|
else:
|
|
needPack = blockWidth == 0.25
|
|
needPack |= (kernel["ConvertAfterDS"] and (tP["bpe"] != tP["bpeDS"]))
|
|
needPack |= kernel["UseF32XEmulation"]
|
|
if isgfx950mx:
|
|
# Keep the gfx950 workaround, but allow normal MX packing on gfx1250.
|
|
needPack = False
|
|
|
|
pack = Module("pack%s_I%s"%(tc,iui))
|
|
packPre = Module("pack%s_I%s Pre"%(tc,iui))
|
|
|
|
isNext = self.getIsNext(writer, kernel, tc, bufferIdx)
|
|
if isNext:
|
|
useTransposeCode = writer.states.a.useTransposeCodeNext if tc == "A" else writer.states.b.useTransposeCodeNext
|
|
useDirect32XEmulation = writer.states.a.useDirect32XEmulationNext if tc == "A" else writer.states.b.useDirect32XEmulationNext
|
|
else:
|
|
useTransposeCode = writer.states.a.useTransposeCodeThis if tc == "A" else writer.states.b.useTransposeCodeThis
|
|
useDirect32XEmulation = writer.states.a.useDirect32XEmulationThis if tc == "A" else writer.states.b.useDirect32XEmulationThis
|
|
indexTranpose = lrvwTile > 1 and (not useTransposeCode)
|
|
|
|
numSplitMetadata = max(ceil((blockWidth * 4) // tP["bpeDS"]) - 1, 0) if tP["isM"] else 0
|
|
|
|
# caculate SMFMA layout
|
|
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 tP["isA"]) or (kernel["ProblemType"]["Sparse"] == 2 and tP["isB"]) or tP["isM"]
|
|
# 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 writer.states.asmCaps["HasSWMMAC_gfx1250"] and not tP["isM"]: blocksPerTGroupSMFMA = 2
|
|
if blocksPerTGroupSMFMA > 1:
|
|
threadGroups = kernel["WavefrontSize"] // matrixInstTO
|
|
elementsPerBlockSMFMA = MIInputPerThUnroll // blocksPerTGroupSMFMA # need adjust if blocks > 1 and is sparse track.
|
|
blockStride = elementsPerBlockSMFMA * threadGroups
|
|
blockOffsetSMFMA = blockStride - elementsPerBlockSMFMA
|
|
|
|
maxLDSConstOffset = writer.states.regCaps["maxLDSConstOffset"]
|
|
|
|
subIterLoadCount = 0
|
|
valufIdx = 0
|
|
halfPLR = (tP["isA"] or tP["isB"]) and kernel["HalfPLR%c"%tc]
|
|
if enableLDSTr:
|
|
numberMTilesPerWave = kernel["MIWaveTile"][tile01]
|
|
if writer.states.asmCaps["HasWMMA_V3"]:
|
|
if tP["bpeDS"] == 0.5:
|
|
LocalReadX = instruction.getInst(0)
|
|
wtRegStride = int(MIInputPerThUnroll * tP["bpeDS"] // bpr)
|
|
outerUnrolledIncrements = 64
|
|
innerUnrolledIncrements = 16
|
|
vwTrLoad = 16
|
|
|
|
for tIdx in range(numberMTilesPerWave):
|
|
for ti in range(0, numTilePerInst):
|
|
constOffset = int((tP["localReadOffset"] + matrixInstTO * ti + MIWaveGroupShape[tile01] * tIdx) * tP["bpeDS"])
|
|
for outerIdx in range(MIInputPerThUnroll//kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]):
|
|
for innerIdx in range(kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]//vwTrLoad):
|
|
paddedOffset = constOffset
|
|
paddedOffset += int((innerIdx * innerUnrolledIncrements + outerIdx * outerUnrolledIncrements) * UnrollStride * tP["bpeDS"])
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
paddedOffset += int((paddedOffset // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
paddedOffset, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, paddedOffset)
|
|
ds = DSModifiers(na=1, offset=paddedOffset)
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride * (tIdx * numTilePerInst+ti)+ 2 * (innerIdx + 2 * outerIdx), tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc, bufferIdx, iui, wtRegStride * (tIdx * numTilePerInst+ti), 2 * (innerIdx + 2 * outerIdx)), blockWidth)
|
|
localReadCode = imod.add(Module("LocalRead%s Valu%u"%(tc, int(valufIdx))))
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
elif tP["bpeDS"] == 0.75:
|
|
LocalReadX = instruction.getInst(0)
|
|
wtRegStride = (int(MIInputPerThUnroll * tP["bpeDS"] // bpr) + 15) // 16 * 16
|
|
outerUnrolledIncrements = 64
|
|
innerUnrolledIncrements = 16
|
|
vwTrLoad = 16
|
|
numVgprsPerLoad = 4 #use 3 for upcoming compiler change
|
|
|
|
for tIdx in range(numberMTilesPerWave):
|
|
constOffset = int((tP["localReadOffset"] + MIWaveGroupShape[tile01] * tIdx) * tP["bpeDS"])
|
|
for outerIdx in range(MIInputPerThUnroll//kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]):
|
|
for innerIdx in range(kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]//vwTrLoad):
|
|
paddedOffset = constOffset
|
|
paddedOffset += int((innerIdx * innerUnrolledIncrements + outerIdx * outerUnrolledIncrements) * UnrollStride * tP["bpeDS"])
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
paddedOffset += int((paddedOffset // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
paddedOffset, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, paddedOffset)
|
|
ds = DSModifiers(na=1, offset=paddedOffset)
|
|
vgprOffset = numVgprsPerLoad * (innerIdx + 2 * outerIdx)
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride*tIdx + vgprOffset, tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc, bufferIdx, iui, wtRegStride*tIdx, vgprOffset), blockWidth)
|
|
localReadCode = imod.add(Module("LocalRead%s Valu%u"%(tc, int(valufIdx))))
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
|
|
elif tP["bpeDS"] == 1:
|
|
LocalReadX = instruction.getInst(0)
|
|
wtRegStride = MIInputPerThUnroll * tP["bpeDS"] // bpr
|
|
numUnrolledIncrements = 32
|
|
vwTrLoad = 8
|
|
numberLRVWPerMIInput = MIInputPerThUnroll // kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]
|
|
for tIdx in range(numberMTilesPerWave):
|
|
offset = int((tP["localReadOffset"] + MIWaveGroupShape[tile01] * tIdx) * tP["bpeDS"])
|
|
if tP["isM"]:
|
|
numLoadTrPerMetadata = max(MIInputPerThUnroll // vwTrLoad, 1)
|
|
for v in range(numLoadTrPerMetadata):
|
|
incrementBytes = int((v * vwTrLoad ) * tP["bpeDS"] * UnrollStride)
|
|
paddedOffset = offset + incrementBytes
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
paddedOffset += int((paddedOffset // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
paddedOffset += tP["localReadSwapByteOffset"]
|
|
addrIdx = paddedOffset // 65536
|
|
srcAddr=vgpr("LocalReadAddr%s+%u"%(tc, addrIdx))
|
|
paddedOffset -= addrIdx * 65536
|
|
ds = DSModifiers(na=1, offset=paddedOffset)
|
|
wtRegStride = 2 if wtRegStride < 2 else wtRegStride # wtRegStride at least need to be 2.
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc, bufferIdx, iui, wtRegStride*tIdx, 2*v), blockWidth)
|
|
localReadCode: Module = imod.add(Module("LocalRead%s Valu%u"%(tc, int(valufIdx))))
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
else:
|
|
for i in range(MIInputPerThUnroll//kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]):
|
|
for v in range(kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]//vwTrLoad):
|
|
incrementBytes = int((v * vwTrLoad + i * numUnrolledIncrements) * tP["bpeDS"] * UnrollStride)
|
|
sparseDenseOffset = 0
|
|
if numberLRVWPerMIInput == 4 and kernel["ProblemType"]["Sparse"] != 0:
|
|
sparseDenseOffset = int((numUnrolledIncrements) * tP["bpeDS"] * UnrollStride) // 2
|
|
if i%2 == 1:
|
|
# For sparse, interleaved so that the 2nd part will have an offset.
|
|
paddedOffset = offset + incrementBytes - sparseDenseOffset
|
|
else:
|
|
paddedOffset = offset + incrementBytes
|
|
else:
|
|
paddedOffset = offset + incrementBytes
|
|
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
paddedOffset += int((paddedOffset // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
paddedOffset, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, paddedOffset)
|
|
ds = DSModifiers(na=1, offset=paddedOffset)
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride*tIdx + 2*v+4*i, tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc, bufferIdx, iui, wtRegStride*tIdx, 2*v+4*i), blockWidth)
|
|
localReadCode: Module = imod.add(Module("LocalRead%s Valu%u"%(tc, int(valufIdx))))
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
elif tP["bpeDS"] == 2:
|
|
numberLRVWPerMIInput = MIInputPerThUnroll // kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]
|
|
for tIdx in range(0, numberMTilesPerWave):
|
|
offset_val = int((tP["localReadOffset"]+MIWaveGroupShape[tile01]*tIdx) * tP["bpeDS"])
|
|
unpaddedOffset = offset_val
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val += int((offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
offset_split, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, offset_val)
|
|
ds = DSModifiers(na=1, offset=offset_split)
|
|
LocalReadX = instruction.getInst(0)
|
|
wtRegStride = int(MIInputPerThUnroll * tP["bpeDS"] // bpr)
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride*tIdx, tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+0"%(tc,bufferIdx,iui, wtRegStride*tIdx), blockWidth)
|
|
valuiIdx = int(valufIdx)
|
|
localReadCode = imod.add(Module("LocalRead%s Valu%u"%(tc,valuiIdx)))
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride*tIdx + blockWidth, tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc,bufferIdx,iui, wtRegStride*tIdx, blockWidth), blockWidth)
|
|
incrementBytes = int(numberLRVWPerMIInput*UnrollStride*kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]*tP["bpeDS"])
|
|
|
|
sparseDenseOffset = 0
|
|
if numberLRVWPerMIInput == 4:
|
|
# generally numberLRVWPerMIInput should be 2, for the dense matrix of sparse cases, it will be 4
|
|
incrementBytes = incrementBytes // 2
|
|
sparseDenseOffset = incrementBytes // 2 # for sparse dense matrix, we read the 2nd half of the data seperately.
|
|
offset_val = unpaddedOffset + incrementBytes - sparseDenseOffset
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val += int((offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
offset_split, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, offset_val)
|
|
ds = DSModifiers(na=1, offset=offset_split)
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
if numberLRVWPerMIInput == 4:
|
|
# for the dense case when sparse.
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride*tIdx + blockWidth * 2, tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc,bufferIdx,iui, wtRegStride*tIdx, blockWidth * 2), blockWidth)
|
|
offset_val = unpaddedOffset + incrementBytes * 2
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val += int((offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
offset_split, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, offset_val)
|
|
ds = DSModifiers(na=1, offset=offset_split)
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, wtRegStride*tIdx + blockWidth * 3, tc)
|
|
destVgpr = vgpr(valuStr, blockWidth)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc,bufferIdx,iui, wtRegStride*tIdx, blockWidth * 3), blockWidth)
|
|
offset_val = unpaddedOffset + incrementBytes * 2 + sparseDenseOffset
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val += int((offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
offset_split, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, offset_val)
|
|
ds = DSModifiers(na=1, offset=offset_split)
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
else:
|
|
assert False, f"Unhandled bpeDS: {tP['bpeDS']}"
|
|
else:
|
|
numOffsetsPerLoad = 2 * blocksPerTGroupSMFMA
|
|
numberMTilesPerWave = kernel["MIWaveTile"][tile01]
|
|
highBits = 0
|
|
totalLoads = numberMTilesPerWave * numOffsetsPerLoad
|
|
for tIdx in range(0, numberMTilesPerWave):
|
|
valuiIdx = int(valufIdx)
|
|
LocalReadX = instruction.getInst(highBits)
|
|
|
|
offset_val = (tP["localReadOffset"]+MIWaveGroupShape[tile01]*tIdx) * tP["bpeDS"] + tP["localReadSwapByteOffset"]
|
|
|
|
def applyPad(offset_val):
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val = offset_val + (offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"]
|
|
return offset_val
|
|
|
|
for oIdx in range(0,numOffsetsPerLoad):
|
|
if blocksPerTGroupSMFMA > 1 and oIdx % blocksPerTGroupSMFMA == 0:
|
|
offset_val += (kernel["MacroTile%s"%tc] * blockOffsetSMFMA) * tP["bpeDS"] * (oIdx // blocksPerTGroupSMFMA)
|
|
|
|
offset = int(applyPad(offset_val))
|
|
offset, srcAddr = self.cal_offset_srcAddr(maxLDSConstOffset, tc, offset)
|
|
ds = DSModifiers(na=1, offset=offset)
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u+%u"%(tc,bufferIdx,iui, 4*tIdx*blocksPerTGroupSMFMA, oIdx * 2), 2)
|
|
localReadCode = Module("LocalRead%s Valu%u"%(tc,valuiIdx))
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCode, comment="LDS Transpose")
|
|
if perpStride == 1:
|
|
inputPerThread = kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"] if not writer.states.inTailLoop else MIInputPerThUnroll
|
|
offset_val += (UnrollStride*inputPerThread) // (blocksPerTGroupSMFMA if writer.states.inTailLoop else 1)
|
|
else:
|
|
permBlock = kernel["MatrixInstK"]
|
|
perpStrideInv = permBlock // perpStride
|
|
inv4K = perpStrideInv * (4 % perpStride) + 4 // perpStride
|
|
offset_val += inv4K * kernel["MacroTile%s"%tc] * tP["bpeDS"]
|
|
splitPoint = max(totalLoads // numSubTiles, 1) if numSubTiles > 1 else totalLoads
|
|
if ((subTileIdx == 0 and subIterLoadCount < splitPoint) \
|
|
or (subTileIdx == 1 and subIterLoadCount >= splitPoint) \
|
|
or numSubTiles == 1) or writer.states.inTailLoop:
|
|
imod.add(localReadCode)
|
|
subIterLoadCount += 1
|
|
# Without enableLDSTr
|
|
else:
|
|
totalLoads = numVectorsPerTile * numReadsPerVector * numReadsPerUnroll
|
|
swapBlockSizeSub = (totalLoads * blockWidth)
|
|
if not writer.states.inTailLoop:
|
|
# divided by numSubTiles for non TailLoop case
|
|
# TailLoop case, code is not devided by numSubTiles
|
|
swapBlockSizeSub //= numSubTiles
|
|
# save the value for index transpose
|
|
abmatrixinfo.swapBlockSizeSub = swapBlockSizeSub
|
|
for vIdx in range(0, numVectorsPerTile):
|
|
for eIdx in range(0, numReadsPerVector):
|
|
valuiIdx = int(valufIdx)
|
|
baseValuiIdx = valuiIdx
|
|
localReadCode = imod.add(Module("LocalRead%s Valu%u"%(tc,valuiIdx)))
|
|
if needPack or numSplitMetadata:
|
|
packCode = pack.add(Module("packCode"))
|
|
packCodePre = packPre.add(Module("packCodePre"))
|
|
|
|
tmpvgpr = []
|
|
is_wmma_v3 = writer.states.asmCaps.get("HasWMMA_V3", False)
|
|
multiGroupXF32 = kernel["UseF32XEmulation"] and is_wmma_v3 and numVgpr * numReadsPerUnroll > 8
|
|
outerBaseValuiIdx = baseValuiIdx
|
|
for tiIdx in range(0, numTilePerInst):
|
|
for rIdx in range(0, numReadsPerUnroll):
|
|
valuiIdx = int(valufIdx)
|
|
baseValuiIdx = valuiIdx - (valuiIdx%8) # use multiple of 8
|
|
baseLRVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx), numVgpr)
|
|
if halfPLR:
|
|
valuStr = writer.getHalfPLRValuStr(writer.states.halfPLRGroups, valuiIdx, tc)
|
|
destVgpr = vgpr(valuStr, numVgpr)
|
|
else:
|
|
destVgpr = baseLRVgpr
|
|
highBitsForHalf = (blockWidth == 0.5) and ((rIdx % 2) == 1) # rIdx = 1
|
|
isHigh16Bits = (blockWidth == 0.25) and ( ((rIdx % 4) //2) == 1) # 2,3
|
|
|
|
packCodeT = Module() # Allocate temporary module for pack code
|
|
packCodePreT = Module() # Allocate temporary module for pack code Pre
|
|
localReadCodeT = Module()
|
|
|
|
if needPack or numSplitMetadata:
|
|
if kernel["UseF32XEmulation"]:
|
|
# Pack data 0-7 with layout:
|
|
# Val+0: bf16 high (0,1)
|
|
# Val+1: bf16 high (2,3)
|
|
# Val+2: bf16 high (4,5)
|
|
# Val+3: bf16 high (6,7)
|
|
# Val+4: bf16 low (0,1)
|
|
# Val+5: bf16 low (2,3)
|
|
# Val+6: bf16 low (4,5)
|
|
# Val+7: bf16 low (6,7)
|
|
|
|
if (valuiIdx % swapBlockSizeSub) == 0:
|
|
allPack4HiDone = False
|
|
allPack4LoDone = False
|
|
if useTransposeCode:
|
|
# generate Tranpose code (with v_swap) for wider local read + useTransposeCode
|
|
self.transposeLRVregs(kernel, packCodePreT, tc, bufferIdx, iui, writer, lrvwTile, swapBlockSizeSub, subTileIdx)
|
|
if writer.states.doFullPackCodePrefetch and useDirect32XEmulation:
|
|
# index transpose + doFullPackCodePrefetch case, finish all first conversion (high value) at once
|
|
allPack4HiDone = True
|
|
for idx in range(0, swapBlockSizeSub, 4):
|
|
newBaseValuiIdx = baseValuiIdx + idx - (idx % 8)
|
|
newIdx = (idx % 8)
|
|
# add comment at last one only
|
|
commentForSchedule1 = "__TF32_1_" + tc + "_%d"%(newBaseValuiIdx//8)
|
|
noComment = idx < swapBlockSizeSub - 4
|
|
commentForPack = "" if noComment else commentForSchedule1
|
|
self.pack4HiBits(kernel, writer, tc, newIdx, bufferIdx, newBaseValuiIdx, iui, packCodeT, lrvwTile, commentForPack, useDirect32XEmulation)
|
|
# do all low cvt
|
|
if useDirect32XEmulation:
|
|
allPack4LoDone = True
|
|
# allocate tmp vgpr first
|
|
self.initTmpVregForPack(kernel, writer, tc, 0, bufferIdx, baseValuiIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, useDirect32XEmulation)
|
|
# low pack 1st stage
|
|
for idx in range(0, swapBlockSizeSub, 4):
|
|
# add comment at last one only
|
|
commentForSchedule1 = "__TF32_1_" + tc + "_%d"%(newBaseValuiIdx//8)
|
|
noComment = idx < swapBlockSizeSub - 4
|
|
commentForPack = "" if noComment else commentForSchedule1
|
|
self.pack4LowBitsStep1(kernel, writer, tc, baseValuiIdx + idx, bufferIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, commentForPack, useDirect32XEmulation)
|
|
# low pack final
|
|
for idx in range(0, swapBlockSizeSub, 4):
|
|
# on last iteration, store lower bits in last 4 registers
|
|
if idx % 8 == 4:
|
|
self.pack4LowBitsFinal(kernel, writer, tc, baseValuiIdx + idx, bufferIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, useDirect32XEmulation, noComment=idx < swapBlockSizeSub - 4)
|
|
# release tmp regs
|
|
self.releaseTmpVregForPack(kernel, writer, tc, baseValuiIdx, tmpvgprFP32, useDirect32XEmulation)
|
|
|
|
# For every 8 read vgprs of fp32, pack high bits of bf16 into first 4 vgprs
|
|
if valuiIdx % 8 == 0 and not allPack4LoDone:
|
|
commentForSchedule1 = "__TF32_1_" + tc + "_%d"%(baseValuiIdx//8)
|
|
# allocate tmp vgpr first
|
|
self.initTmpVregForPack(kernel, writer, tc, 0, bufferIdx, baseValuiIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, useDirect32XEmulation)
|
|
if not allPack4HiDone:
|
|
commentForPack = ""
|
|
self.pack4HiBits(kernel, writer, tc, 0, bufferIdx, baseValuiIdx, iui, packCodeT, lrvwTile, commentForPack, useDirect32XEmulation)
|
|
commentForPack = commentForSchedule1
|
|
if kernel["UseMFMAF32XEmulation"] or (not useDirect32XEmulation):
|
|
commentForPack = ""
|
|
self.pack4HiBits(kernel, writer, tc, 4, bufferIdx, baseValuiIdx, iui, packCodeT, lrvwTile, commentForPack, useDirect32XEmulation)
|
|
|
|
do8PackAtOnce = indexTranpose and not allPack4HiDone
|
|
if (valuiIdx % 8) == 4 and do8PackAtOnce:
|
|
# index transpose case
|
|
# we need to keep both original values and transpose values
|
|
# do "Compute low bits" for 0-3 and 4-7 + final pack here
|
|
# do all at (valuiIdx % 8) == 4
|
|
abmatrixinfo = writer.states.a if tc == 'A' else writer.states.b
|
|
useDedicatedTmp = abmatrixinfo.tmpVgprCvtSub >= 0
|
|
tmp = abmatrixinfo.tmpVgprCvtSub if useDedicatedTmp else writer.vgprPool.checkOut(1, "x32f tmp")
|
|
valuiIdx0 = valuiIdx - 4 # for 0-3
|
|
valuiIdx1 = valuiIdx # for 4-7
|
|
# src (original value)
|
|
v0, v1, v2, v3 = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, valuiIdx0, iui, lrvwTile)
|
|
v4, v5, v6, v7 = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, valuiIdx1, iui, lrvwTile)
|
|
# dst (high)
|
|
v0t, v1t, v2t, v3t = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, valuiIdx0, iui, lrvwTile, dst=True)
|
|
v4t, v5t, v6t, v7t = self.get4VgprForEmu(writer, kernel, tc, bufferIdx, valuiIdx1, iui, lrvwTile, dst=True)
|
|
# low bits 0
|
|
packCodeT.add(PVCvtBF16toFP32(dst=vgpr(tmp), src=v0t, comment="begin"+str(valuiIdx0)))
|
|
packCodeT.add(VSubF32(dst=v4t, src0=v0, src1=vgpr(tmp)))
|
|
# low bits 1
|
|
packCodeT.add(VCvtBF16toFP32(dst=vgpr(tmp), src=v0t, vgprMask=None, vi=1))
|
|
packCodeT.add(VSubF32(dst=vgpr(tmp), src0=v1, src1=vgpr(tmp)))
|
|
# final 4
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v4t, src0=v4t, src1=vgpr(tmp), comment="pack final begin"))
|
|
# low bits 2
|
|
packCodeT.add(PVCvtBF16toFP32(dst=vgpr(tmp), src=v1t, comment="begin"+str(valuiIdx0+2)))
|
|
packCodeT.add(VSubF32(dst=v5t, src0=v2, src1=vgpr(tmp)))
|
|
# low bits 3
|
|
packCodeT.add(VCvtBF16toFP32(dst=vgpr(tmp), src=v1t, vgprMask=None, vi=1))
|
|
packCodeT.add(VSubF32(dst=vgpr(tmp), src0=v3, src1=vgpr(tmp)))
|
|
# final 5
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v5t, src0=v5t, src1=vgpr(tmp)))
|
|
# low bits 4
|
|
packCodeT.add(PVCvtBF16toFP32(dst=vgpr(tmp), src=v2t, comment="begin"+str(valuiIdx0+4)))
|
|
packCodeT.add(VSubF32(dst=v6t, src0=v4, src1=vgpr(tmp)))
|
|
# low bits 5
|
|
packCodeT.add(VCvtBF16toFP32(dst=vgpr(tmp), src=v2t, vgprMask=None, vi=1))
|
|
packCodeT.add(VSubF32(dst=vgpr(tmp), src0=v5, src1=vgpr(tmp)))
|
|
# final 6
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v6t, src0=v6t, src1=vgpr(tmp)))
|
|
# low bits 6
|
|
packCodeT.add(PVCvtBF16toFP32(dst=vgpr(tmp), src=v3t, comment="begin"+str(valuiIdx0+6)))
|
|
packCodeT.add(VSubF32(dst=v7t, src0=v6, src1=vgpr(tmp)))
|
|
# low bits 7
|
|
packCodeT.add(VCvtBF16toFP32(dst=vgpr(tmp), src=v3t, vgprMask=None, vi=1))
|
|
packCodeT.add(VSubF32(dst=vgpr(tmp), src0=v7, src1=vgpr(tmp)))
|
|
# final 7
|
|
commentStr = "__TF32_2_" + tc + "_%d pack final end"%(baseValuiIdx//8)
|
|
packCodeT.add(VCvtPkF32toBF16(dst=v7t, src0=v7t, src1=vgpr(tmp), comment=commentStr))
|
|
if not useDedicatedTmp:
|
|
writer.vgprPool.checkIn(tmp)
|
|
elif valuiIdx % 4 == 0 and (not do8PackAtOnce) and (not allPack4LoDone):
|
|
noComment = valuiIdx % 8 == 0
|
|
commentForPack = "" if noComment else commentForSchedule1
|
|
self.pack4LowBitsStep1(kernel, writer, tc, valuiIdx, bufferIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, commentForPack, useDirect32XEmulation)
|
|
# on last iteration, store lower bits in last 4 registers
|
|
if valuiIdx % 8 == 4:
|
|
self.pack4LowBitsFinal(kernel, writer, tc, valuiIdx, bufferIdx, iui, packCodeT, lrvwTile, tmpvgprFP32, useDirect32XEmulation)
|
|
if valuiIdx % 8 == 4 and (not allPack4LoDone):
|
|
self.releaseTmpVregForPack(kernel, writer, tc, baseValuiIdx, tmpvgprFP32, useDirect32XEmulation)
|
|
# For WMMA V3 multigroup XF32: rearrange packed BF16 data so all
|
|
# hi values are contiguous (X0+0..7) followed by lo values (X0+8..15).
|
|
# Must happen here in pack code (not MAC code) to avoid the scheduler
|
|
# interleaving the swap with residual/TF32_2 packing operations.
|
|
if multiGroupXF32 and rIdx == numReadsPerUnroll - 1:
|
|
halfGroup = 4
|
|
for i in range(halfGroup):
|
|
swapIdx1 = outerBaseValuiIdx + halfGroup + i
|
|
swapIdx2 = outerBaseValuiIdx + halfGroup * 2 + i
|
|
swapVgpr1 = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, swapIdx1))
|
|
swapVgpr2 = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, swapIdx2))
|
|
commentStr = "XF32 pack rearrange %s: swap [%d] <-> [%d]"%(tc, swapIdx1, swapIdx2)
|
|
# Tag the last swap with __TF32_2 so _packItemsConditional treats the swap block as a proper boundary,
|
|
# preventing _interleavePackAB from displacing swaps into the next group.
|
|
if i == halfGroup - 1:
|
|
commentStr += " __TF32_2_%s_%d_swap"%(tc, outerBaseValuiIdx // (halfGroup * 2))
|
|
packCodeT.add(VSwapB32(dst=swapVgpr1, src=swapVgpr2, comment=commentStr))
|
|
|
|
if kernel["ConvertAfterDS"] and (tP["bpe"] != tP["bpeDS"]):
|
|
if tP["bpe"] == 2 and tP["bpeDS"] == 4:
|
|
assert 0 # Doesn't support ConvertAfterDS
|
|
else:
|
|
highBitsForHalf = False
|
|
isHigh16Bits = False
|
|
if kernel["UnrollMajorLDS%s"%tc]:
|
|
cvtTimes = int(blockWidth * writer.states.bpr // tP["bpeDS"]) // MIInputPerThUnroll
|
|
for i in range(0, cvtTimes):
|
|
offset = cvtTimes - i - 1
|
|
if writer.states.asmCaps["Hascvtf16_fp8_sf32"]:
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+1+offset*2)),\
|
|
src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+offset)), scale=0x3f800000,\
|
|
vop3=VOP3PModifiers(op_sel=[1,0,0,0]), comment="convert fp8 to f16"))
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+0+offset*2)),\
|
|
src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+offset)), scale=0x3f800000,\
|
|
vop3=VOP3PModifiers(op_sel=[0,0,0,0]), comment="convert fp8 to f16"))
|
|
elif writer.states.asmCaps["HasCvtFP8toF16"]:
|
|
packCode.add(VCvtScalePkFP8toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+1+offset*2)),\
|
|
src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+offset)),\
|
|
vop3=VOP3PModifiers(op_sel=[1,0]),comment="convert F8 to F16"))
|
|
packCode.add(VCvtScalePkFP8toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+0+offset*2)),\
|
|
src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+offset)),\
|
|
vop3=VOP3PModifiers(op_sel=[0,0]), comment="convert fp8 to f16"))
|
|
else:
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+offset)),\
|
|
sel=HighBitSel.HIGH, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+1+offset*2)), src=vgpr("CvtTemp+0"),\
|
|
sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+1+offset*2)), src=vgpr("CvtTemp+1"),\
|
|
sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+offset)),\
|
|
sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+0+offset*2)), src=vgpr("CvtTemp+0"),\
|
|
sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx+0+offset*2)), src=vgpr("CvtTemp+1"),\
|
|
sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
elif (writer.states.lrvwTileA == 1 and tc == 'A') or (writer.states.lrvwTileB == 1 and tc == 'B'):
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, int(valufIdx*2)), numVgpr)
|
|
CvtDstVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, int(valufIdx*2)), numVgpr)
|
|
if needPack and (rIdx%4) != 0:
|
|
if (rIdx % 4) != 0:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%4, valuiIdx), numVgpr)
|
|
else:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valuiIdx), numVgpr)
|
|
if writer.states.asmCaps["Hascvtf16_fp8_sf32"]:
|
|
sel = [0,0,0,0] if (rIdx % 2 == 0) else [0,0,1,0]
|
|
packCodeT.add(VCvtScaleFP8toF16(dst=CvtDstVgpr, src=destVgpr, scale=0x3f800000, vop3=VOP3PModifiers(op_sel=sel), comment="convert fp8 to f16"))
|
|
elif writer.states.asmCaps["HasCvtFP8toF16"]:
|
|
sel = (rIdx % 2 == 0)
|
|
packCode.add(VCvtFP8toF16(dst=CvtDstVgpr, src=destVgpr, vop3=VOP3PModifiers(byte_sel=[sel]),comment="convert F8 to F16"))
|
|
else:
|
|
packCodeT.add(VCvtFP8toF32(dst=destVgpr, src=destVgpr, sdwa=SDWAModifiers(src0_sel=SelectBit.BYTE_0)))
|
|
packCodeT.add(ECvtF32toF16(dst=CvtDstVgpr, src=destVgpr, sel=HighBitSel.LOW if rIdx % 2 == 0 else HighBitSel.HIGH, comment="Convert to FP16"))
|
|
elif (writer.states.lrvwTileA == 2 and tc == 'A') or (writer.states.lrvwTileB == 2 and tc == 'B'):
|
|
if needPack or numSplitMetadata:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), vIdx*numVgpr), numVgpr)
|
|
for i in range(0, numVgpr):
|
|
cvtDstVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), vIdx*numVgpr), numVgpr)
|
|
if writer.states.asmCaps["Hascvtf16_fp8_sf32"]:
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=destVgpr, src=destVgpr,scale=0x3f800000,vop3=VOP3PModifiers(op_sel=[0,0,0,0]),comment="convert F8 to F16"))
|
|
elif writer.states.asmCaps["HasCvtFP8toF16"]:
|
|
packCode.add(VCvtScalePkFP8toF16(dst=destVgpr, src=destVgpr, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
else:
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=destVgpr, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=destVgpr, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=destVgpr, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
|
|
if rIdx == numReadsPerUnroll-1:
|
|
for i in range(0, numVgpr):
|
|
vgprIdx = int((vIdx * numVgpr + i) * tP["bpe"] * MIInputPerThUnroll // writer.states.bpr * min(writer.states.bpr // tP["bpe"], vectorWidth))
|
|
vgprOffset = 0
|
|
for vectorIdx in range(0, 2):
|
|
for elementIdx in range(0, int(tP["bpe"]*MIInputPerThUnroll//writer.states.bpr)):
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+vgprOffset)), \
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*2+1, i+vIdx*numVgpr)), \
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*2, i+vIdx*numVgpr)), \
|
|
src2=sgpr("PackKForV%u"%vectorIdx), \
|
|
comment="select K=%u%u for vector=%u"%(elementIdx*2, elementIdx*2+1, vectorIdx)))
|
|
vgprOffset += 1
|
|
elif (writer.states.lrvwTileA == 4 and tc == 'A') or (writer.states.lrvwTileB == 4 and tc == 'B'):
|
|
if needPack or numSplitMetadata:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2 * vIdx * numVgpr), numVgpr)
|
|
cvtDestVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2 * vIdx * numVgpr + 1), numVgpr)
|
|
packCodeT.add(VLShiftRightB32(dst=cvtDestVgpr, shiftHex=16, src=destVgpr, comment="shift 2 element to vgpr+1"))
|
|
if writer.states.asmCaps["Hascvtf16_fp8_sf32"]:
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=destVgpr, src=destVgpr,scale=0x3f800000, comment="convert F8 to F16"))
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr, src=cvtDestVgpr,scale=0x3f800000, comment="convert F8 to F16"))
|
|
elif writer.states.asmCaps["HasCvtFP8toF16"]:
|
|
packCode.add(VCvtScalePkFP8toF16(dst=destVgpr, src=destVgpr, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
packCode.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr, src=cvtDestVgpr, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
else:
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=destVgpr, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=destVgpr, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=destVgpr, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=cvtDestVgpr, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
|
|
if rIdx == numReadsPerUnroll-1:
|
|
for i in range(0, numVgpr*2):
|
|
vgprIdx = int((2 * vIdx * numVgpr + i) * tP["bpe"] * MIInputPerThUnroll // writer.states.bpr * min(writer.states.bpr // tP["bpe"], vectorWidth))
|
|
vgprOffset = 0
|
|
for vectorIdx in range(0, 2):
|
|
for elementIdx in range(0, int(tP["bpe"]*MIInputPerThUnroll//writer.states.bpr)):
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+vgprOffset)), \
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*2+1, i+2*vIdx*numVgpr)), \
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*2, i+2*vIdx*numVgpr)), \
|
|
src2=sgpr("PackKForV%u"%vectorIdx), \
|
|
comment="select K=%u%u for vector=%u"%(elementIdx*2, elementIdx*2+1, vectorIdx)))
|
|
vgprOffset += 1
|
|
elif (writer.states.lrvwTileA == 8 and tc == 'A') or (writer.states.lrvwTileB == 8 and tc == 'B'):
|
|
if needPack or numSplitMetadata:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2*vIdx*numVgpr), numVgpr)
|
|
cvtDestVgpr0 = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2*vIdx*numVgpr+0), 1)
|
|
cvtDestVgpr1 = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2*vIdx*numVgpr+1), 1)
|
|
cvtDestVgpr2 = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2*vIdx*numVgpr+2), 1)
|
|
cvtDestVgpr3 = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), 2*vIdx*numVgpr+3), 1)
|
|
packCodeT.add(VLShiftRightB32(dst=cvtDestVgpr3, shiftHex=16, src=cvtDestVgpr1, comment="shift 2 element to vgpr+3"))
|
|
|
|
packCodeT.add(VMovB32(dst=cvtDestVgpr2, src=cvtDestVgpr1))
|
|
packCodeT.add(VLShiftRightB32(dst=cvtDestVgpr1, shiftHex=16, src=cvtDestVgpr0, comment="shift 2 element to vgpr+1"))
|
|
if writer.states.asmCaps["Hascvtf16_fp8_sf32"]:
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr0, src=cvtDestVgpr0,scale=0x3f800000, comment="convert F8 to F16"))
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr1, src=cvtDestVgpr1,scale=0x3f800000, comment="convert F8 to F16"))
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr2, src=cvtDestVgpr2,scale=0x3f800000, comment="convert F8 to F16"))
|
|
packCodeT.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr3, src=cvtDestVgpr3,scale=0x3f800000, comment="convert F8 to F16"))
|
|
elif writer.states.asmCaps["HasCvtFP8toF16"]:
|
|
packCode.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr0, src=cvtDestVgpr0, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
packCode.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr1, src=cvtDestVgpr1, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
packCode.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr2, src=cvtDestVgpr2, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
packCode.add(VCvtScalePkFP8toF16(dst=cvtDestVgpr3, src=cvtDestVgpr3, vop3=VOP3PModifiers(op_sel=[0]),comment="convert F8 to F16"))
|
|
else:
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=cvtDestVgpr0, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr0, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr0, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=cvtDestVgpr1, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr1, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr1, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=cvtDestVgpr2, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr2, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr2, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtPkFP8toF32(dst=vgpr("CvtTemp", 2), src=cvtDestVgpr3, sel=HighBitSel.LOW, comment="convert to F32"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr3, src=vgpr("CvtTemp+0"), sel=HighBitSel.LOW, comment="Convert to FP16"))
|
|
packCodeT.add(ECvtF32toF16(dst=cvtDestVgpr3, src=vgpr("CvtTemp+1"), sel=HighBitSel.HIGH, comment="Convert to FP16"))
|
|
|
|
if rIdx == numReadsPerUnroll-1:
|
|
for i in range(0, numVgpr*2):
|
|
vgprIdx = int((2 * vIdx * numVgpr + i) * tP["bpe"] * MIInputPerThUnroll // writer.states.bpr * min(writer.states.bpr // tP["bpe"], vectorWidth))
|
|
vgprOffset = 0
|
|
for vectorIdx in range(0, 2):
|
|
for elementIdx in range(0, int(tP["bpe"]*MIInputPerThUnroll//writer.states.bpr)):
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+vgprOffset)), \
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*2+1, i+2*vIdx*numVgpr)), \
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*2, i+2*vIdx*numVgpr)), \
|
|
src2=sgpr("PackKForV%u"%vectorIdx), \
|
|
comment="select K=%u%u for vector=%u"%(elementIdx*2, elementIdx*2+1, vectorIdx)))
|
|
vgprOffset += 1
|
|
else:
|
|
pass
|
|
elif lrvwTile > 1 and not kernel["UseF32XEmulation"]:
|
|
highBitsForHalf = 0
|
|
isHigh8Bits = 0
|
|
isHigh16Bits = 0
|
|
numElementPerReg = int(writer.states.bpr//tP["bpe"])
|
|
|
|
needPackK16 = False
|
|
needPackK8Lw = False
|
|
if kernel["ProblemType"][MacDataType].isHalf() or kernel["ProblemType"][MacDataType].isBFloat16():
|
|
if writer.states.lrvwTileA > 1 or writer.states.lrvwTileB > 1:
|
|
needPackK16 = True
|
|
if writer.states.lrvwTileMetadata > 1:
|
|
needPackK8Lw = True
|
|
|
|
tPackM = "M" if needPackK16 and needPackK8Lw else ""
|
|
|
|
if needPack or numSplitMetadata:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), vIdx*numVgpr), numVgpr)
|
|
if rIdx == numReadsPerUnroll-1:
|
|
for i in range(0, numVgpr):
|
|
# convert from [tile][MiInputPerThread][vector] to [tile][vector][MiInputPerThread]
|
|
vgprIdx = int((vIdx*numVgpr+i)*tP["bpeDS"]*MIInputPerThUnroll//writer.states.bpr*min(writer.states.bpr//tP["bpeDS"],vectorWidth))
|
|
if numSplitMetadata:
|
|
vgprIdx = (vIdx*numVgpr+i)*ceil(tP["bpeDS"]*MIInputPerThUnroll / writer.states.bpr)*min(writer.states.bpr//tP["bpeDS"],vectorWidth)
|
|
if MIInputPerThUnroll == 8:
|
|
vgprOffset = 0
|
|
for elementIdx in range(0, numSplitMetadata+1):
|
|
if elementIdx >= writer.states.bpr:
|
|
break
|
|
packCode.add(VPermB32(
|
|
dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + elementIdx * 2)),
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 1, i+vIdx*numVgpr)),
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 0, i+vIdx*numVgpr)),
|
|
src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)),
|
|
comment="1 select K=%u,%u for vector=%u"%(0, 1, vgprOffset)
|
|
))
|
|
packCode.add(VPermB32(
|
|
dst=vgpr("PackTemp"),
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 3, i+vIdx*numVgpr)),
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 2, i+vIdx*numVgpr)),
|
|
src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)),
|
|
comment="1 select K=%u, %u for vector=%u"%(2, 3, vgprOffset)
|
|
))
|
|
packCode.add(VLShiftLeftOrB32(
|
|
dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + elementIdx * 2)),
|
|
src0=vgpr("PackTemp"),
|
|
shiftHex=16,
|
|
src1=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + elementIdx * 2)),
|
|
comment="pack two half Vgpr to one Vgpr"
|
|
))
|
|
packCode.add(VPermB32(
|
|
dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + elementIdx * 2 + 1)),
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 5, i+vIdx*numVgpr)),
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 4, i+vIdx*numVgpr)),
|
|
src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)),
|
|
comment="2 select K=%u,%u for vector=%u"%(4, 5, vgprOffset)
|
|
))
|
|
packCode.add(VPermB32(
|
|
dst=vgpr("PackTemp"),
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 7, i+vIdx*numVgpr)),
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 6, i+vIdx*numVgpr)),
|
|
src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)),
|
|
comment="2 select K=%u,%u for vector=%u"%(6, 7, vgprOffset)
|
|
))
|
|
packCode.add(VLShiftLeftOrB32(
|
|
dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + elementIdx * 2 + 1)),
|
|
src0=vgpr("PackTemp"),
|
|
shiftHex=16,
|
|
src1=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + elementIdx * 2 + 1)),
|
|
comment="pack two half Vgpr to one Vgpr"
|
|
))
|
|
vgprOffset += 1
|
|
elif MIInputPerThUnroll == 4:
|
|
vgprOffset = 0
|
|
for elementIdx in range(0, numSplitMetadata+1):
|
|
if elementIdx >= writer.states.bpr:
|
|
break
|
|
# since the number of input thread is 4, so will alwasy be D0, D1, D2, D3
|
|
packCode.add(VPermB32(
|
|
dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)),
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 1, i+vIdx*numVgpr)),
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 0, i+vIdx*numVgpr)),
|
|
src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)),
|
|
comment="1 select K=%u%u for vector=%u"%(0, 1, vgprOffset)
|
|
))
|
|
packCode.add(VPermB32(
|
|
dst=vgpr("PackTemp"),
|
|
src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 3, i+vIdx*numVgpr)),
|
|
src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 2, i+vIdx*numVgpr)),
|
|
src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)),
|
|
comment="1 select K=%u%u for vector=%u"%(2, 3, vgprOffset)
|
|
))
|
|
packCode.add(VLShiftLeftOrB32(
|
|
dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)),
|
|
src0=vgpr("PackTemp"),
|
|
shiftHex=16,
|
|
src1=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)),
|
|
comment="pack two half Vgpr to one Vgpr"
|
|
))
|
|
vgprOffset += 1
|
|
elif MIInputPerThUnroll == 2:
|
|
vgprOffset = 0
|
|
for elementIdx in range(0, numSplitMetadata+1):
|
|
if elementIdx >= writer.states.bpr:
|
|
break
|
|
# since the number of input thread is 2, so will alwasy be D0 and D1
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)), src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 1, i+vIdx*numVgpr)), src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, 0, i+vIdx*numVgpr)), src2=sgpr("PackKFor%sV%u"%(tPackM, vgprOffset)), \
|
|
comment="select K=%u%u for vector=%u"%(0, 1, vgprOffset)))
|
|
vgprOffset += 1
|
|
elif MIInputPerThUnroll == 1:
|
|
destVgpr_ = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%(MIInputPerThUnroll), vIdx*numVgpr + i))
|
|
bitShift = 0
|
|
for elementIdx in range(0, numSplitMetadata+1):
|
|
# go to next vgpr
|
|
if elementIdx >= writer.states.bpr:
|
|
break
|
|
comment_ = "another VGPR storing lshr %d-bit value %d %d" %(bitShift, vgprIdx, elementIdx) if bitShift != 0 else ""
|
|
packCodeT.add(VMovB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)), src=destVgpr_, comment=comment_))
|
|
if bitShift != 0:
|
|
packCodeT.add(VLShiftRightB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)), shiftHex=hex(bitShift), src=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx)), comment="ValuMetadata Vpgr >> %d" % bitShift))
|
|
bitShift += 8
|
|
else:
|
|
assert False
|
|
elif tP["isM"]:
|
|
vgprOffset = 0
|
|
for elementIdx in range(0, MIInputPerThUnroll):
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+elementIdx+vIdx*2)), src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, vgprOffset*2 + 1 , i+vIdx*numVgpr)), src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, vgprOffset*2, i+vIdx*numVgpr)), src2=sgpr("PackKForV%u"%elementIdx), \
|
|
comment="select K=%u%u for vector=%u"%(vgprOffset*2+1, vgprOffset*2, elementIdx)))
|
|
vgprOffset += (1 if elementIdx % 2 == 1 else 0)
|
|
elif kernel["ProblemType"][MacDataType].isHalf() or kernel["MFMA_BF16_1K"] or kernel["ProblemType"][MacDataType].isBFloat16():
|
|
vgprOffset = 0
|
|
for vectorIdx in range(0, numElementPerReg):
|
|
for elementIdx in range(0, int(tP["bpe"]*MIInputPerThUnroll//writer.states.bpr)):
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+vgprOffset)), src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*numElementPerReg+1, i+vIdx*numVgpr)), src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*numElementPerReg, i+vIdx*numVgpr)), src2=sgpr("PackKForV%u"%vectorIdx), \
|
|
comment="select K=%u%u for vector=%u"%(elementIdx*numElementPerReg, elementIdx*numElementPerReg+1, vectorIdx)))
|
|
vgprOffset += 1
|
|
elif kernel["ProblemType"][MacDataType].isInt8() or kernel["ProblemType"][MacDataType].is8bitFloat():
|
|
vgprOffset = 0
|
|
# vertorIdx 2,3 is for the case vectorWidth > 2
|
|
for vectorIdx in range(0, numElementPerReg):
|
|
if vectorWidth <= 2 and vectorIdx > 1:
|
|
break
|
|
for elementIdx in range(0, int(tP["bpe"]*MIInputPerThUnroll//writer.states.bpr)):
|
|
packCodeT.add(VPermB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx+vgprOffset)), src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*numElementPerReg+1, i+vIdx*numVgpr)), src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*numElementPerReg, i+vIdx*numVgpr)), src2=sgpr("PackKForV%u"%vectorIdx), \
|
|
comment="select K=%u%u for vector=%u"%(elementIdx*4, elementIdx*4+1, vectorIdx)))
|
|
packCodeT.add(VPermB32(dst=vgpr("PackTemp"), src0=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*numElementPerReg+3, i+vIdx*numVgpr)), src1=vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, elementIdx*numElementPerReg+2, i+vIdx*numVgpr)), src2=sgpr("PackKForV%u"%vectorIdx), \
|
|
comment="select K=%u%u for vector=%u"%(elementIdx*4+2, elementIdx*4+3, vectorIdx)))
|
|
packCodeT.add(VLShiftLeftOrB32(dst=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + vgprOffset)), src0=vgpr("PackTemp"), shiftHex=16, src1=vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, vgprIdx + vgprOffset)), comment="pack two half Vgpr to one Vgpr"))
|
|
vgprOffset += 1
|
|
|
|
else:
|
|
isHigh8Bits = (blockWidth == 0.25) and ( ((rIdx % 4) % 2) == 1) # 1,3
|
|
# pack for blockWidth 0.5 type
|
|
if tP["isM"]:
|
|
elementsPerPack = kernel["MIInputPerThreadMetadata"]
|
|
vgprsPerPack = ceil(kernel["MIInputPerThreadMetadata"] / writer.states.bpr)
|
|
dstVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valufIdx//elementsPerPack*vgprsPerPack), numVgpr)
|
|
if (blockWidth == 0.25):
|
|
if ((rIdx % elementsPerPack) == (elementsPerPack-1)): #rIdx = 1, 3,...
|
|
for rIdx_ in range(0, ceil(rIdx/2)):
|
|
if rIdx_:
|
|
lowVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx_*2, valufIdx//elementsPerPack), numVgpr)
|
|
else:
|
|
lowVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valufIdx//elementsPerPack*vgprsPerPack), numVgpr)
|
|
highVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx_*2 + 1, valufIdx//elementsPerPack), numVgpr)
|
|
packCode.add(VLShiftLeftOrB32(dst=lowVgpr, src0=highVgpr, shiftHex=8, src1=lowVgpr, comment="pack two int8 Vgpr to one half Vgpr"))
|
|
if elementsPerPack >= 4:
|
|
for rIdx_ in range(0, ceil(rIdx/4)):
|
|
if rIdx_:
|
|
lowVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx_*4, valufIdx//elementsPerPack), numVgpr)
|
|
else:
|
|
lowVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valufIdx//elementsPerPack*vgprsPerPack), numVgpr)
|
|
highVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx_*4 + 2, valufIdx//elementsPerPack), numVgpr)
|
|
dstVgpr = vgpr("Valu%s_X%u_I%u+%u"%(tc, bufferIdx, iui, valufIdx//elementsPerPack*vgprsPerPack+rIdx_), numVgpr)
|
|
packCode.add(VLShiftLeftOrB32(dst=dstVgpr, src0=highVgpr, shiftHex=hex(0x10), src1=lowVgpr, comment="pack two int8x2 Vgpr to one Vgpr"))
|
|
if needPackMetadata:
|
|
destVgpr = dstVgpr
|
|
if rIdx % elementsPerPack:
|
|
destVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx % elementsPerPack, valufIdx//elementsPerPack), numVgpr)
|
|
#don't use the ds read high
|
|
isHigh8Bits = False
|
|
isHigh16Bits = False
|
|
elif writer.states.archCaps["HasEccHalf"] or not writer.states.asmCaps["HasWMMA_V1"]: # ECC pack
|
|
if highBitsForHalf:
|
|
highVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%2, valuiIdx), numVgpr)
|
|
if writer.states.archCaps["DSLow16NotPreserve"]:
|
|
packCodeT.add(VLShiftLeftOrB32(dst=baseLRVgpr, src0=highVgpr, shiftHex=hex(0x10), src1=baseLRVgpr, comment="pack two half Vgpr to one Vgpr"))
|
|
else:
|
|
packCodeT.add(VOrB32(dst=baseLRVgpr, src0=baseLRVgpr, src1=highVgpr, comment="pack two half Vgpr to one Vgpr"))
|
|
destVgpr = highVgpr
|
|
# pack for blockWidth 0.25 type
|
|
if rIdx != 0:
|
|
if isHigh8Bits or isHigh16Bits:
|
|
highVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%4, valuiIdx), numVgpr)
|
|
destVgpr = highVgpr
|
|
if isHigh8Bits:
|
|
lowVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, (rIdx%4)-1, valuiIdx), numVgpr) if isHigh16Bits else baseLRVgpr
|
|
packCodeT.add(VLShiftLeftOrB32(dst=lowVgpr, src0=highVgpr, shiftHex=8, src1=lowVgpr, comment="pack two int8 Vgpr to one half Vgpr"))
|
|
if isHigh16Bits:
|
|
if writer.states.archCaps["DSLow16NotPreserve"]:
|
|
packCodeT.add(VLShiftLeftOrB32(dst=baseLRVgpr, src0=lowVgpr, shiftHex=hex(0x10), src1=baseLRVgpr, comment="pack two half Vgpr to one Vgpr"))
|
|
else:
|
|
packCodeT.add(VOrB32(dst=baseLRVgpr, src0=baseLRVgpr, src1=lowVgpr, comment="pack two half Vgpr to one Vgpr"))
|
|
else: # no ECC pack
|
|
# pack for No ECC blockwidth 0.25 type
|
|
if rIdx != 0:
|
|
if isHigh8Bits:
|
|
highVgpr = vgpr("Valu%s_X%u_I%u_D%u+%u"%(tc, bufferIdx, iui, rIdx%2, valuiIdx), numVgpr)
|
|
destVgpr = highVgpr
|
|
if isHigh8Bits and isHigh16Bits:
|
|
packCodeT.add(VLShiftLeftOrB32(dst=baseLRVgpr, src0=highVgpr, shiftHex=hex(0x8), src1=baseLRVgpr, comment="pack two int8x2 Vgpr to one Vgpr"))
|
|
|
|
if kernel["ConvertAfterDS"] and kernel["UnrollMajorLDS%s"%tc]:
|
|
valufIdx += blockWidth * (tP["bpe"] // tP["bpeDS"]) if (not tP["isM"]) else 1
|
|
# workaround for gfx950 MX
|
|
# need to increment valufIdx by 1
|
|
elif isgfx950mx:
|
|
valufIdx += 1
|
|
else:
|
|
valufIdx += blockWidth if (not tP["isM"]) else (numVgpr if writer.states.asmCaps["HasSWMMAC_gfx1250"] else 1)
|
|
|
|
# load read instruction
|
|
paramList = []
|
|
|
|
# gfx1250 LDS offset formula shared by XF32 and BF16/Half/FP8/etc paths.
|
|
# The WMMA V3 LDS layout uses a *2 factor on the unroll stride.
|
|
def calcGfx1250LdsOffset():
|
|
if "MXS" in tc:
|
|
# rIdx walks the K-scales packed into one MFMA-K sub-iter.
|
|
# Step is UnrollStride, which equals mxUnit for K-major LDS
|
|
# (UMLDS=1) and MT+LdsPad for M-major LDS (UMLDS=0). For
|
|
# UMLDS=1, numReadsPerUnroll is typically 1 so this loop
|
|
# degenerates to the pre-existing incOffset=0 behavior.
|
|
incOffset = rIdx * numElementPerRead * UnrollStride
|
|
elif kernel["UnrollMajorLDS%s" % tP["tensorChar"]]:
|
|
incOffset = rIdx * numElementPerRead * UnrollStride * 2
|
|
incOffset += tiIdx * matrixInstTO * vectorWidth * tileStride
|
|
else:
|
|
vw = kernel[f"LocalReadVectorWidth{tc if('MXS' not in tc) else 'MXS'}"]
|
|
incOffset = (rIdx // vw) * UnrollStride * vw
|
|
incOffset += rIdx * numElementPerRead * UnrollStride
|
|
return int((incOffset + offset_val + tP["localReadOffset"]) * tP["bpeDS"])
|
|
|
|
for oIdx in range(0, numOffsets):
|
|
# segment-interleave: a vIdx group can cross the LDS component
|
|
# boundary. Split it into within-component (vCols) + component jump
|
|
# (segCompByteOff, added post-pad below).
|
|
segCompByteOff = 0
|
|
vCols = (vIdx * numOffsets + oIdx) * MIWaveGroupShape[tile01]
|
|
if (kernel.get("LDSSegmentInterleave") == 1
|
|
and kernel["LDSSegInterleaveOffsets"].get("footprintPacked")
|
|
and tc in ("A", "B")):
|
|
numComp = kernel["NumWaves"] // 2
|
|
compCols = kernel["MacroTile%u" % tile01] // numComp
|
|
if compCols > 0 and MIWaveGroupShape[tile01] > 0:
|
|
segCompByteOff = (vCols // compCols) * kernel["LDSSegInterleaveOffsets"]["writeStrideBytes"]
|
|
vCols = vCols % compCols
|
|
if perpStride > 1 and kernel["ProblemType"]["TLU%s"%tc] == 0:
|
|
permBlock = kernel["MatrixInstK"] if kernel["ProblemType"]["TLU%s"%tc] == 1 else kernel["VectorWidth%s"%tc] * kernel["MatrixInstM"]
|
|
perpStrideInv = permBlock // perpStride
|
|
offset_val = (eIdx * (perpStrideInv) + vCols) * tileStride
|
|
else:
|
|
offset_val = (eIdx + vCols) * tileStride
|
|
|
|
if kernel["ProblemType"]["Sparse"] != 0:
|
|
if blocksPerTGroupSMFMA > 1:
|
|
blockId = (rIdx * numElementPerRead) // elementsPerBlockSMFMA # block 0 or block 1
|
|
if kernel["UnrollMajorLDS%s"%(tc)]:
|
|
offset_val = offset_val + (blockOffsetSMFMA * blockId)
|
|
else:
|
|
offset_val = offset_val + (blockOffsetSMFMA * blockId) * UnrollStride
|
|
offset_val = int((rIdx * numElementPerRead * UnrollStride + offset_val + tP["localReadOffset"]) * tP["bpeDS"])
|
|
elif writer.states.asmCaps["HasMFMA_f8f6f4"] and kernel["ProblemType"][MacDataType].is8bitFloat() and kernel["MatrixInstK"] > 32 and not isgfx950mx:
|
|
incOffset = 0
|
|
midIdx = numReadsPerUnroll // 2
|
|
if rIdx >= midIdx:
|
|
if kernel["UnrollMajorLDS%s" % tP["tensorChar"]] == False:
|
|
# TODO: why are these the offsets???
|
|
if kernel["MatrixInstM"] == 32:
|
|
incOffset = midIdx * numElementPerRead * UnrollStride
|
|
elif kernel["MatrixInstM"] == 16:
|
|
incOffset = 3 * midIdx * numElementPerRead * UnrollStride
|
|
else:
|
|
if kernel["MatrixInstM"] == 32:
|
|
incOffset = 16
|
|
elif kernel["MatrixInstM"] == 16:
|
|
incOffset = 48
|
|
incOffset = rIdx * numElementPerRead * UnrollStride + incOffset
|
|
offset_val = (incOffset + offset_val + tP["localReadOffset"]) * tP["bpeDS"]
|
|
elif kernel["UseF32XEmulation"]:
|
|
# Previously a single ds_read could be used to load all inputs for mfma
|
|
# For emulated TF32, 2x ds_read is required along with a different mfma layout
|
|
# so we need to adjust the offsets accordingly for the second ds_read.
|
|
if tuple(kernel["ISA"][:2]) == (12, 5):
|
|
# gfx1250 WMMA V3: shared offset formula with BF16/Half (see calcGfx1250LdsOffset)
|
|
offset_val = calcGfx1250LdsOffset()
|
|
else:
|
|
# Numbers here are specific to the mfma layout (gfx950)
|
|
incOffset = 0
|
|
midIdx = numReadsPerUnroll // 2
|
|
if rIdx >= midIdx:
|
|
if kernel["UnrollMajorLDS%s" % tP["tensorChar"]] == False:
|
|
if kernel["MatrixInstM"] == 32:
|
|
incOffset = midIdx * numElementPerRead * UnrollStride
|
|
elif kernel["MatrixInstM"] == 16:
|
|
incOffset = 3 * midIdx * numElementPerRead * UnrollStride
|
|
else:
|
|
if kernel["MatrixInstM"] == 32 and kernel["MatrixInstK"] == 16:
|
|
incOffset = 4
|
|
elif kernel["MatrixInstM"] == 16 and kernel["MatrixInstK"] == 32:
|
|
incOffset = 12
|
|
incOffset = rIdx * numElementPerRead * UnrollStride + incOffset
|
|
offset_val = (incOffset + offset_val + tP["localReadOffset"]) * tP["bpeDS"]
|
|
# For wmma_v3, the maximum number of bytes per read is 16 bytes in 4 vgprs, which happens in the case of fp16/bf16/fp8/bf8/fp6/bf6/f4.
|
|
elif tuple(kernel["ISA"][:2]) == (12, 5) \
|
|
and (kernel["ProblemType"][MacDataType].is8bitFloat() or kernel["ProblemType"][MacDataType].isBFloat16() \
|
|
or kernel["ProblemType"][MacDataType].isHalf() or kernel["ProblemType"][MacDataType].isFloat4() \
|
|
or kernel["ProblemType"][MacDataType].is6bitFloat() or kernel["ProblemType"][MacDataType].isInt8()):
|
|
offset_val = calcGfx1250LdsOffset()
|
|
else:
|
|
offset_val = int((rIdx * numElementPerRead * UnrollStride + offset_val + tP["localReadOffset"]) * tP["bpeDS"])
|
|
|
|
if (kernel["LdsBlockSizePerPad%s"%tc] != 0) and (kernel["LdsPad%s"%tc] != 0):
|
|
offset_val = int(offset_val + (offset_val // kernel["LdsBlockSizePerPad%s"%tc]) * kernel["LdsPad%s"%tc] * tP["bpeDS"])
|
|
# component jump is post-pad: writeStrideBytes already includes pad.
|
|
offset_val = int(offset_val + segCompByteOff)
|
|
offset_val = offset_val + tP["localReadSwapByteOffset"]
|
|
# TODO: Add NLC>1 offset calcs here?
|
|
if (kernel["DirectToLds%s" % tc] and \
|
|
kernel["GlobalReadVectorWidth%s"%tc] * tP["bpeDS"] > 4) and not kernel["UseGeneralizedNLCOne%s"%tc]:
|
|
# another address conversion for DirectToLds + NumLoadsCoalesced > 1
|
|
dummy, offset_val = writer.lraOffsetConversionForDTLandNLC(kernel, tP, offset_val)
|
|
|
|
paramList.append(int(offset_val))
|
|
|
|
comment = "L -> Reg lro=%d swapByteOffset=%u ti=%u vIdx=%u eIdx=%u rIdx=%u oIdx=%u buffer=%u iui=%u" \
|
|
% (tP["localReadOffset"], tP["localReadSwapByteOffset"], MIWaveGroupShape[tile01], vIdx, eIdx, rIdx, oIdx, bufferIdx, iui)
|
|
|
|
highBits = 0 if writer.states.archCaps["DSLow16NotPreserve"] else highBitsForHalf or isHigh16Bits
|
|
|
|
addrIdx = paramList[0] // 65536
|
|
srcAddr=vgpr("LocalReadAddr%s+%u"%(tc, addrIdx))
|
|
tdmFullLdsOffset = paramList[0]
|
|
paramList[0] -= addrIdx * 65536
|
|
|
|
if numOffsets == 1:
|
|
ds = DSModifiers(na=1, offset=paramList[0])
|
|
else:
|
|
ds = DSModifiers(na=2, offset0=paramList[0], offset1=paramList[1])
|
|
LocalReadX = instruction.getInst(highBits)
|
|
if kernel["UseF32XEmulation"]:
|
|
index = valuiIdx
|
|
# UseF32XEmulation case, convert index in transpose case.
|
|
# getVgprForEmu handles both transose and non-transpose cases.
|
|
# dest of local read is source of TF32 conv. Need to specify dest=False for getVgprForEmu
|
|
# indexTranpose case, disable index conversion for local read
|
|
destVgpr = self.getVgprForEmu(writer, kernel, tc, bufferIdx, iui, index, lrvwTile, vgprLen=numVgpr, dst=False, localRead=True)
|
|
|
|
# When numVectorsPerTile==1 the per-wave reads never cross the TDMSplit
|
|
# half boundary (only vIdx=0 exists), so the byte-offset half classifier
|
|
# would tag every read half0 and leave the half1 tensor_load un-waited.
|
|
# Such reads' combined region depends on BOTH half loads -> carry both
|
|
# half tokens.
|
|
tdmBothHalves = (kernel["TDMSplit"] and not kernel["ProblemType"]["Sparse"]
|
|
and not tP.get("isM", False) and numVectorsPerTile == 1)
|
|
self._emitLdsRead(writer, kernel, tP, LocalReadX, dst=destVgpr, src=srcAddr, ds=ds, module=localReadCodeT, ldsByteOffset=tdmFullLdsOffset, bothHalves=tdmBothHalves, comment=comment)
|
|
# TODO - handle vector-load
|
|
with writer.allocTmpSgpr(1, tag="LocalReadVALU_tmpSgprInfo2") as tmpSgprInfo:
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
if writer.db["CheckValue1%s"%tc] and not writer.inTailLoop:
|
|
|
|
dbgVgpr = destVgpr
|
|
dbgVgprList = destVgpr.split("v[")
|
|
if len(dbgVgprList) == 1: # vIdx, no []
|
|
dbgVgpr = dbgVgprList[0]
|
|
else:
|
|
# We only check the first one now
|
|
# TODO: Handle vector, but need to take care the last one
|
|
dbgVgprList = (dbgVgprList[1].split("]")[0]).split(':')
|
|
dbgVgpr = "v[%s]"%dbgVgprList[0]
|
|
localReadCodeT.add(SWaitCnt(dscnt=0, vscnt=0, comment="CheckValue1 wait for LDS read"))
|
|
|
|
if kernel["ProblemType"][MacDataType].isHalf():
|
|
hexValue = hex(0x3c003c00) # packed 1s
|
|
if needPack:
|
|
hexValue = hex(0x3c000000) if highBitsForHalf else hex(0x00003c00)
|
|
localReadCodeT.add(SMovB32(dst=sgpr(tmpSgpr), src=hexValue, comment="CheckValue1: FP16"))
|
|
localReadCodeT.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
elif kernel["ProblemType"][MacDataType].isBFloat16():
|
|
hexValue = hex(0x3f803f80) # packed 1s
|
|
if needPack:
|
|
hexValue = hex(0x3f800000) if highBitsForHalf else hex(0x00003f80)
|
|
localReadCodeT.add(SMovB32(dst=sgpr(tmpSgpr), src=hexValue, comment="CheckValue1: BF16"))
|
|
localReadCodeT.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
if kernel["ProblemType"][MacDataType].isInt8():
|
|
if needPack:
|
|
hexValue = hex(0x00010000) if isHigh16Bits else hex(0x00000001)
|
|
localReadCodeT.add(SMovB32(dst=sgpr(tmpSgpr), src=hexValue, comment="CheckValue1: INT8"))
|
|
localReadCodeT.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
# TODO - Check if this works. But need this? MFMA would use INT8
|
|
elif kernel["ProblemType"][MacDataType].isInt8x4():
|
|
localReadCodeT.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(0x01010101), comment="CheckValue1: INT8x4"))
|
|
localReadCodeT.add(writer.assert_eq( dbgVgpr, sgpr(tmpSgpr)))
|
|
|
|
elif kernel["ProblemType"][MacDataType].isSingle():
|
|
localReadCodeT.add(writer.assert_eq( dbgVgpr, 1.0) )
|
|
|
|
addPackLR = False
|
|
splitPoint = max(totalLoads // numSubTiles, 1) if numSubTiles > 1 else totalLoads
|
|
if ((subTileIdx == 0 and subIterLoadCount < splitPoint) \
|
|
or (subTileIdx == 1 and subIterLoadCount >= splitPoint) \
|
|
or numSubTiles == 1) or writer.states.inTailLoop:
|
|
addPackLR = True
|
|
|
|
if addPackLR:
|
|
if needPack or numSplitMetadata:
|
|
packCode.add(packCodeT)
|
|
packCodePre.add(packCodePreT)
|
|
localReadCode.add(localReadCodeT)
|
|
|
|
subIterLoadCount += 1
|
|
# End of loop3
|
|
if needPack:
|
|
if tP["isA"]:
|
|
writer.states.a.numPackCvt = len(packCode.flatitems())
|
|
elif tP["isB"]:
|
|
writer.states.b.numPackCvt = len(packCode.flatitems())
|
|
# End of loop2
|
|
# End of loop1
|
|
# DTV case, do not return local read code.
|
|
if (tP["isA"] or tP["isB"]) and kernel["DirectToVgpr%s"%tc]:
|
|
imod = Module("LocalReadDo%s_I%s (Empty)" % (tP["tensorChar"],iui))
|
|
|
|
# DTV and Tr Load case, do not return pack code
|
|
if (tP["isA"] or tP["isB"]) and kernel["enableGLTr%s"%tc]:
|
|
pack = Module("Pack%s_I%s (Empty)" % (tP["tensorChar"],iui))
|
|
|
|
# free any remaining tmp vgprs from emulation
|
|
if useDirect32XEmulation:
|
|
while len(tmpvgprFP32):
|
|
tmp = tmpvgprFP32.pop()
|
|
writer.vgprPool.checkIn(tmp)
|
|
|
|
return imod, pack, packPre
|