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
610 lines
35 KiB
Python
610 lines
35 KiB
Python
################################################################################
|
|
#
|
|
# Copyright (C) 2024-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 cop-
|
|
# ies 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 IM-
|
|
# PLIED, 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 CONNE-
|
|
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
################################################################################
|
|
|
|
from rocisa.code import Module, Label, RegSet
|
|
from rocisa.container import DSModifiers, ContinuousRegister
|
|
from rocisa.instruction import DSLoadB128, DSLoadB32, DSLoadB64, DSStoreB128, \
|
|
DSStoreB32, DSStoreB64, SAndB32, SCBranchSCC0, SCmpEQU32, SMovB32, SNop, SWaitCnt, \
|
|
VAddF32, VAddI32, VAddU32, VAndB32, VLShiftLeftAddU32, VMovB32, VMulLOU32, VAddF64
|
|
from rocisa.functions import vectorStaticDivide
|
|
from copy import deepcopy
|
|
from ..Common import log2, ceilDivide, DataDirection
|
|
from ..Component import Component
|
|
from ..KernelWriterModules import *
|
|
from ..AsmStoreState import StoreState, VectorDataTypes
|
|
#import abc
|
|
|
|
class LSU(Component):
|
|
"""
|
|
LSU block.
|
|
"""
|
|
|
|
class LSUOff(LSU):
|
|
kernel = {"LocalSplitU": 1}
|
|
|
|
class LSUOn(LSU):
|
|
|
|
@classmethod
|
|
def matches(cls, writer, debug=False):
|
|
return writer.states.kernel["LocalSplitU"] > 1
|
|
|
|
def __call__(self):
|
|
assert(0)
|
|
|
|
def splitOutputData(self, writer, kernel):
|
|
self.LSUelemCoord0 = []
|
|
self.LSUelemCoord1 = []
|
|
self.LSUelements = []
|
|
self.LSUfullVws = []
|
|
# NOTE: notLocalFullTileElements returns vector widths in descending order, for example: [8, 4, 2, 1]
|
|
# first-longest vector width and corresponding element are used through LSU component
|
|
(vwsdummy, elesdummy, self.LSUfullVws, self.LSUelements) = writer.notLocalFullTileElements(kernel)
|
|
atomic = False # atomic is for GSU > 1
|
|
beta = True
|
|
vectorDataTypes = VectorDataTypes()
|
|
ss = StoreState(writer, kernel, self.LSUfullVws[0], False, beta, atomic, self.LSUelements[0], vectorDataTypes, dim=0)
|
|
self.LSUelemCoord0, self.LSUelemCoord1 = ss.getStoreElementsInfoForBatch(kernel, self.LSUelements[0])
|
|
|
|
# search for valid lsu wave offset
|
|
maxtt1 = 0
|
|
maxtt0 = 0
|
|
maxvc1 = 0
|
|
maxvc0 = 0
|
|
validOffset = -1
|
|
validOffset0 = -1
|
|
validOffset1 = -1
|
|
self.LSUelementsArchIdx = [[] for i in range(4)]
|
|
self.LSUelementsPerLSUWave = []
|
|
self.LSUelemCoord0PerLSUWave = []
|
|
self.LSUelemCoord1PerLSUWave = []
|
|
|
|
# Check valid LSU/VW combination
|
|
if len(self.LSUelements[0]) >= kernel["LocalSplitU"]:
|
|
if kernel["LocalSplitU"] == 4:
|
|
idxGrp = 1
|
|
for idxGrp in range(1, len(self.LSUelements[0])//4 + 1):
|
|
for i in range(idxGrp):
|
|
i0 = i
|
|
i1 = i + 1 * idxGrp
|
|
i2 = i + 2 * idxGrp
|
|
i3 = i + 3 * idxGrp
|
|
offset0 = self.LSUelemCoord0[i0] + self.LSUelemCoord1[i0] * kernel["MacroTile0"]
|
|
offset1 = self.LSUelemCoord0[i1] + self.LSUelemCoord1[i1] * kernel["MacroTile0"]
|
|
offset2 = self.LSUelemCoord0[i2] + self.LSUelemCoord1[i2] * kernel["MacroTile0"]
|
|
offset3 = self.LSUelemCoord0[i3] + self.LSUelemCoord1[i3] * kernel["MacroTile0"]
|
|
if (offset3 - offset2 == offset2 - offset1) and (offset2 - offset1 == offset1 - offset0):
|
|
validOffset0 = self.LSUelemCoord0[i1] - self.LSUelemCoord0[i0]
|
|
validOffset1 = self.LSUelemCoord1[i1] - self.LSUelemCoord1[i0]
|
|
if self.LSUelemCoord0[i2] - self.LSUelemCoord0[i1] == validOffset0 \
|
|
and self.LSUelemCoord0[i3] - self.LSUelemCoord0[i2] == validOffset0 \
|
|
and self.LSUelemCoord1[i2] - self.LSUelemCoord1[i1] == validOffset1 \
|
|
and self.LSUelemCoord1[i3] - self.LSUelemCoord1[i2] == validOffset1:
|
|
validOffset = offset1 - offset0
|
|
break
|
|
if validOffset != -1:
|
|
break
|
|
for idx in range(0, len(self.LSUelements[0]), 4*idxGrp):
|
|
for idx2 in range(idxGrp):
|
|
self.LSUelementsArchIdx[0].append(self.LSUfullVws[0]*(idx + idx2))
|
|
self.LSUelementsArchIdx[1].append(self.LSUfullVws[0]*(idx + 1*idxGrp + idx2))
|
|
self.LSUelementsArchIdx[2].append(self.LSUfullVws[0]*(idx + 2*idxGrp + idx2))
|
|
self.LSUelementsArchIdx[3].append(self.LSUfullVws[0]*(idx + 3*idxGrp + idx2))
|
|
self.LSUelementsPerLSUWave.append(self.LSUelements[0][idx + idx2])
|
|
self.LSUelemCoord0PerLSUWave.append(self.LSUelemCoord0[idx + idx2])
|
|
self.LSUelemCoord1PerLSUWave.append(self.LSUelemCoord1[idx + idx2])
|
|
elif kernel["LocalSplitU"] == 2:
|
|
i = 0
|
|
offset0 = self.LSUelemCoord0[i] + self.LSUelemCoord1[i] * kernel["MacroTile0"]
|
|
offset1 = self.LSUelemCoord0[i + 1] + self.LSUelemCoord1[i + 1] * kernel["MacroTile0"]
|
|
validOffset = offset1 - offset0
|
|
validOffset0 = self.LSUelemCoord0[i + 1] - self.LSUelemCoord0[i]
|
|
validOffset1 = self.LSUelemCoord1[i + 1] - self.LSUelemCoord1[i]
|
|
for idx in range(0, len(self.LSUelements[0]), 2):
|
|
self.LSUelementsArchIdx[0].append(self.LSUfullVws[0]*idx)
|
|
self.LSUelementsArchIdx[1].append(self.LSUfullVws[0]*(idx+1))
|
|
self.LSUelementsPerLSUWave.append(self.LSUelements[0][idx])
|
|
self.LSUelemCoord0PerLSUWave.append(self.LSUelemCoord0[idx])
|
|
self.LSUelemCoord1PerLSUWave.append(self.LSUelemCoord1[idx])
|
|
else:
|
|
assert 0, "No valid LSU offset found."
|
|
|
|
if validOffset == -1:
|
|
assert 0, "No valid LSU offset found."
|
|
self.LSUValidOffset0 = validOffset0
|
|
self.LSUValidOffset1 = validOffset1
|
|
return validOffset
|
|
|
|
##############################################################################
|
|
# LocalSplitU: Local Write, Read, and Reduction
|
|
##############################################################################
|
|
def writeReadReduction(self, writer, kernel):
|
|
module = Module("localSplitULocalWriteAndRead")
|
|
module.addComment2("LocalSplitU Reduction")
|
|
module.add(writer._syncThreads(kernel))
|
|
module.add(Label("localSplitULocalWriteAndRead", ""))
|
|
|
|
acc2arch, arch2acc = accToArchMapper(kernel)
|
|
accImOffset = accVgprImagNumOffset(kernel)
|
|
complexMultiplier = 2 if kernel["ProblemType"]["DataType"].isComplex() else 1
|
|
# prepare the data that is to be Reduction in this wave
|
|
# the output LSUelementsArchIdx has all arch-indices.
|
|
validOffset = self.splitOutputData(writer, kernel)
|
|
|
|
numAccIdx = len(self.LSUelementsArchIdx[0])
|
|
numSetAccIdx = ceilDivide(numAccIdx, kernel["LocalSplitUReuseLDS"])
|
|
maxLDSConstOffset = writer.states.regCaps["maxLDSConstOffset"]
|
|
# computeStoreVgprs
|
|
if kernel["EnableMatrixInstruction"]:
|
|
module.add(writer.computeStoreVgprs(kernel))
|
|
else:
|
|
# new method. output self.vgprs.coord0InMT/coord1InMT
|
|
module.add(writer.computeStoreVgprs(kernel, \
|
|
divisor = kernel["MacroTile0"] // kernel["GlobalWriteVectorWidth"], \
|
|
tid0Scale = kernel["GlobalWriteVectorWidth"], \
|
|
tid1Scale = 1))
|
|
|
|
# Checkout local read resource
|
|
bpr = 4 #bytes per register
|
|
regsPerElem = kernel["ProblemType"]["ComputeDataType"].numRegisters()
|
|
bytesPerElem = int(kernel["ProblemType"]["ComputeDataType"].numBytes())
|
|
bytesPerVector = self.LSUfullVws[0] * bytesPerElem
|
|
numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1]
|
|
regsPerStep = int((bytesPerVector+3)//4)
|
|
elementStep = bytesPerVector // bytesPerElem
|
|
numTotalAccVgprLdsReduction = len(self.LSUelements[0])*regsPerStep*(self.LSUfullVws[0]//elementStep)
|
|
assert (numTotalAccVgprLdsReduction%kernel["LocalSplitU"]) == 0, "AccVgprLdsRedcution % LSU != 0"
|
|
numTotalAccVgprLdsReduction = numTotalAccVgprLdsReduction // kernel["LocalSplitU"]
|
|
self.accVgprLdsReduction = writer.vgprPool.checkOutAligned(numTotalAccVgprLdsReduction, 4, "LsuReduction")
|
|
module.add(RegSet("v", "vgprLsuReduction", self.accVgprLdsReduction))
|
|
module.addComment0("Size of vgprLsuReduction is %u"%(numTotalAccVgprLdsReduction))
|
|
writer.states.c.startVgprValu = self.accVgprLdsReduction
|
|
|
|
# Local Read VGPR idx
|
|
localReadVgprIdx = 0
|
|
|
|
lsu_id = writer.vgprPool.checkOut(1,"lsu_id")
|
|
wave_id = writer.vgprPool.checkOut(1,"wave_id")
|
|
tmpVgpr = writer.vgprPool.checkOutAligned(2, 2, tag="LSUOn_writeReadReduction_tmpVgpr")
|
|
tmpVgprRes = ContinuousRegister(tmpVgpr, 2)
|
|
|
|
module.add(vectorStaticDivide(wave_id, "Serial", \
|
|
kernel["WavefrontSize"], tmpVgprRes))
|
|
module.add(vectorStaticDivide(lsu_id, wave_id, numWaves, tmpVgprRes, \
|
|
comment="Get LSU wave ID"))
|
|
module.add(VAndB32(vgpr(wave_id), hex(numWaves - 1), vgpr(wave_id), \
|
|
comment="Get wave ID"))
|
|
|
|
# accVgprs read in GSUAMBSK
|
|
codeAccVgprRead = []
|
|
if kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel":
|
|
codeAccVgprRead = deepcopy(writer.codes.accVgprRead) if writer.states.serializedStore else None
|
|
|
|
for reUseIdx in range(kernel["LocalSplitUReuseLDS"]):
|
|
module.addComment1("LocalSplitU: local write %d/%d"%(reUseIdx+1,kernel["LocalSplitUReuseLDS"]))
|
|
module.add(Label("localSplitULocalWriteAndRead_%d"%(reUseIdx+1), ""))
|
|
|
|
startLSUaccIdxSet = reUseIdx * numSetAccIdx
|
|
endLSUaccIdxSet = min(numAccIdx, startLSUaccIdxSet + numSetAccIdx)
|
|
|
|
#scan the needed accVGPRIdx
|
|
neededAccVGPRIdx = [[] for i in range(kernel["LocalSplitU"])]
|
|
numAccVgpr = 0
|
|
miRegPerOut = kernel["MIRegPerOut"]
|
|
for lsu in range(kernel["LocalSplitU"]):
|
|
for i in range(startLSUaccIdxSet, endLSUaccIdxSet):
|
|
for j in range(self.LSUfullVws[0]):
|
|
for cm in range(complexMultiplier):
|
|
for r in range(miRegPerOut):
|
|
accIdx = arch2acc[self.LSUelementsArchIdx[lsu][i] + j] * miRegPerOut + r + (cm*accImOffset)
|
|
neededAccVGPRIdx[lsu].append(accIdx)
|
|
numAccVgpr += 1
|
|
|
|
# lsuProcessOffset is used when local read
|
|
numVgprPerLSU = len(neededAccVGPRIdx[0])
|
|
lsuProcessOffset = numVgprPerLSU * kernel["WavefrontSize"] * 4
|
|
|
|
# break here instead of assert to release resource and do a post-rejection
|
|
# caused by range(startLSUaccIdxSet, endLSUaccIdxSet) = range(1,1)
|
|
if numAccVgpr <= 0:
|
|
writer.states.invalidLSUCode = True
|
|
break
|
|
|
|
assert numAccVgpr > 0,"startLSUaccIdxSet=%u,endLSUaccIdxSet=%u,numAccIdx=%u"%(startLSUaccIdxSet,endLSUaccIdxSet,numAccIdx)
|
|
|
|
dataPerWave = numAccVgpr * kernel["WavefrontSize"] * 4
|
|
ldsStride = dataPerWave * numWaves
|
|
# Prepare Write/Read instruction info
|
|
if bytesPerVector % 16 == 0:
|
|
DSStoreBX = DSStoreB128
|
|
DSLoadBX = DSLoadB128
|
|
numInstPerVW = bytesPerVector // 16
|
|
regsPerStore = 4
|
|
elif bytesPerVector % 8 == 0:
|
|
DSStoreBX = DSStoreB64
|
|
DSLoadBX = DSLoadB64
|
|
numInstPerVW = bytesPerVector // 8
|
|
regsPerStore = 2
|
|
else:
|
|
DSStoreBX = DSStoreB32
|
|
DSLoadBX = DSLoadB32
|
|
numInstPerVW = bytesPerVector // 4
|
|
regsPerStore = 1
|
|
|
|
accVgprRes = writer.vgprPool.checkOutAligned(regsPerStore, 4, "accLSUVgprRes")
|
|
|
|
regsPerVector = self.LSUfullVws[0] * regsPerElem
|
|
maxOffset = (kernel["LocalSplitU"] -1) * ldsStride + ((numVgprPerLSU // regsPerVector -1) * numInstPerVW + (numInstPerVW -1)) * regsPerStore * (bpr * kernel["WavefrontSize"])
|
|
numAddr = maxOffset // maxLDSConstOffset + 1
|
|
addr = writer.vgprPool.checkOut(numAddr,"addr")
|
|
with writer.allocTmpSgpr(1, tag="LSUOn_writeReadReduction_tmpSgprInfo") as tmpSgprInfo:
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(dataPerWave), \
|
|
comment="dataPerWave (%d)"%dataPerWave))
|
|
module.add(VAndB32(vgpr(addr), hex(kernel["WavefrontSize"]-1), vgpr("Serial"), \
|
|
comment="initial addr"))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(tmpSgpr), src1=vgpr(wave_id), \
|
|
comment="tmp = waveId * dataPerWave"))
|
|
module.add(VLShiftLeftAddU32(dst=vgpr(addr), shiftHex=log2(regsPerStore * bpr), src0=vgpr(addr), src1=vgpr(tmpVgpr), \
|
|
comment="addr = initial addr + tmp"))
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(ldsStride), \
|
|
comment="ldsStride = waveNum * dataPerWave (%d)"%ldsStride))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(tmpSgpr), src1=vgpr(lsu_id), \
|
|
comment="tmp = (waveNum * dataPerWave) * lsu_id"))
|
|
module.add(VAddU32(vgpr(addr), vgpr(tmpVgpr), vgpr(addr), \
|
|
comment="addr += tmp"))
|
|
|
|
if reUseIdx > 0:
|
|
module.add(SWaitCnt(dscnt=0, comment="wait for all reads"))
|
|
module.add(writer._syncThreads(kernel, "pre-lsu local write"))
|
|
|
|
# # Do Local Write
|
|
for lsu in range(kernel["LocalSplitU"]):
|
|
for i in range(numVgprPerLSU // regsPerStore):
|
|
for r in range(regsPerStore):
|
|
srcIdx = neededAccVGPRIdx[lsu][i*regsPerStore + r]
|
|
readInst = writer.accVgprReadWriteFunction(kernel, srcIdx, True)
|
|
srcVgpr = writer.accVgprReadWriteIndex(kernel, srcIdx)
|
|
module.add(readInst(dst=vgpr(accVgprRes+r),
|
|
src=srcVgpr,
|
|
comment="copy acc[%u] to vreg[%u], LSU%u will process" % (srcIdx, accVgprRes+r, lsu)))
|
|
|
|
module.add(SNop(waitState=1, comment="2 wait states required before reading vgpr"))
|
|
dstIdx = lsu * (numVgprPerLSU // regsPerStore) + i
|
|
module.add(DSStoreBX(dstAddr=vgpr(addr), src=vgpr(accVgprRes, regsPerStore), \
|
|
ds=DSModifiers(offset=(dstIdx * regsPerStore * (bpr * kernel["WavefrontSize"]))), \
|
|
comment="arch[%d]"%dstIdx))
|
|
|
|
# Release local write resource
|
|
writer.vgprPool.checkIn(accVgprRes)
|
|
|
|
module.addComment1("LocalSplitU: local read %d/%d"%(reUseIdx+1,kernel["LocalSplitUReuseLDS"]))
|
|
|
|
# Calculate offset for wave id and lsu id
|
|
with writer.allocTmpSgpr(1, tag="LSUOn_writeReadReduction_tmpSgprInfo2") as tmpSgprInfo:
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
module.add(VAndB32(vgpr(addr), hex(kernel["WavefrontSize"]-1), vgpr("Serial"), \
|
|
comment="initial addr"))
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=hex(dataPerWave), \
|
|
comment="wave offset (%d)"%dataPerWave))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(tmpSgpr), src1=vgpr(wave_id), \
|
|
comment="wave offset = wave_id * wave offset"))
|
|
module.add(VLShiftLeftAddU32(dst=vgpr(addr), shiftHex=log2(regsPerStore * bpr), src0=vgpr(addr), src1=vgpr(tmpVgpr), \
|
|
comment="addr = initial addr + wave offset"))
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), \
|
|
src=hex(lsuProcessOffset), comment="LSU Process Offset %d"%(lsuProcessOffset)))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(tmpSgpr), src1=vgpr(lsu_id), \
|
|
comment="lsu offset = lsu_id * LSU Process Offset"))
|
|
module.add(VAddU32(dst=vgpr(addr), src0=vgpr(addr), src1=vgpr(tmpVgpr), \
|
|
comment="addr += lsu offset"))
|
|
for i in range(1,numAddr):
|
|
module.add(VAddU32(vgpr(addr+i), maxLDSConstOffset*i, vgpr(addr), \
|
|
comment="addr += maxLDSConstOffset*%u"%(i)))
|
|
|
|
module.add(SWaitCnt(dscnt=0, comment="wait for all writes"))
|
|
module.add(writer._syncThreads(kernel, "post-lsu local write"))
|
|
module.add(Label("localSplitULocalRead_%d"%(reUseIdx+1), ""))
|
|
|
|
moduleReduction = Module("LocalSplitU_Reduction")
|
|
inLoopTmpVgpr = writer.vgprPool.checkOutAligned(numVgprPerLSU*(kernel["LocalSplitU"]-1), 4, "TempLsuReduction")
|
|
|
|
# Do Local Read
|
|
for i in range(0, numVgprPerLSU // regsPerVector):
|
|
for v in range(numInstPerVW):
|
|
for r in range(0, kernel["LocalSplitU"]):
|
|
regIdx = (i * numInstPerVW + v) * regsPerStore
|
|
offset = r * ldsStride + regIdx * (bpr * kernel["WavefrontSize"])
|
|
num = offset // maxLDSConstOffset
|
|
offset -= num * maxLDSConstOffset
|
|
srcvgpr = vgpr(addr+num)
|
|
if r == 0:
|
|
vgprStr = "LsuReduction+%u"%(localReadVgprIdx)
|
|
else:
|
|
vgprStr = inLoopTmpVgpr + (numVgprPerLSU * (r - 1) + regIdx)
|
|
module.add(DSLoadBX(dst=vgpr(vgprStr, regsPerStore), src=srcvgpr, \
|
|
ds=DSModifiers(offset=(offset)), \
|
|
comment="r=%u i=%u, from acc[%d]"%(r, (i * numInstPerVW + v), neededAccVGPRIdx[0][(i * numInstPerVW + v)])))
|
|
# Generate Reduction code at the same time.
|
|
if r == 0:
|
|
# Insert waitcnt code here
|
|
numTotalInst = numVgprPerLSU // regsPerVector * numInstPerVW * kernel["LocalSplitU"]
|
|
numPassedInst = (i * numInstPerVW + (v + 1)) * kernel["LocalSplitU"]
|
|
numLRWaitCnt = numTotalInst - numPassedInst
|
|
moduleReduction.add(SWaitCnt(dscnt=numLRWaitCnt, comment="wait count is (%u-%u)"%(numTotalInst, numPassedInst)))
|
|
if r > 0:
|
|
for regToAdd in range(0, regsPerStore, regsPerElem):
|
|
if kernel["ProblemType"]["ComputeDataType"].isSingle():
|
|
moduleReduction.add(VAddF32(dst=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd)), src0=vgpr(vgprStr+regToAdd), \
|
|
src1=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd)), comment=""))
|
|
elif kernel["ProblemType"]["ComputeDataType"].isInt32():
|
|
moduleReduction.add(VAddI32(dst=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd)), src0=vgpr(vgprStr+regToAdd), \
|
|
src1=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd)), comment=""))
|
|
elif kernel["ProblemType"]["ComputeDataType"].isSingleComplex():
|
|
moduleReduction.add(VAddF32(dst=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+0)), src0=vgpr(vgprStr+regToAdd+0), \
|
|
src1=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+0)), comment=""))
|
|
moduleReduction.add(VAddF32(dst=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+1)), src0=vgpr(vgprStr+regToAdd+1), \
|
|
src1=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+1)), comment=""))
|
|
elif kernel["ProblemType"]["ComputeDataType"].isDoubleComplex():
|
|
moduleReduction.add(VAddF64(dst=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+0), 2), src0=vgpr(vgprStr+regToAdd+0, 2), \
|
|
src1=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+0), 2), comment=""))
|
|
moduleReduction.add(VAddF64(dst=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+2), 2), src0=vgpr(vgprStr+regToAdd+2, 2), \
|
|
src1=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd+2), 2), comment=""))
|
|
else:
|
|
# TODO: hpa_half, int8
|
|
assert(0) # unsupported data type, need to modify here and LSU write/read code
|
|
|
|
if kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel":
|
|
if not kernel["MIArchVgpr"]:
|
|
# write to accvgpr
|
|
for regToAdd in range(regsPerStore):
|
|
codeAccVgprReadInst = codeAccVgprRead.popFirstItem() # v_accvgpr_read_b32
|
|
accIdx = codeAccVgprReadInst.srcs[0].regIdx
|
|
writeInst = writer.accVgprReadWriteFunction(kernel, accIdx, False)
|
|
dstVgpr = writer.accVgprReadWriteIndex(kernel, accIdx)
|
|
moduleReduction.add(writeInst(dst=dstVgpr, src=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd)),
|
|
comment="copy vreg[%u] to acc[%u]" % (localReadVgprIdx+regToAdd, accIdx)))
|
|
else:
|
|
# write to vgpr
|
|
for regToAdd in range(regsPerStore):
|
|
codeAccVgprReadInst = codeAccVgprRead.popFirstItem() # v_mov_b32
|
|
dstVgpr = codeAccVgprReadInst.srcs[0]
|
|
moduleReduction.add(VMovB32(dst=dstVgpr, src=vgpr("LsuReduction+%u"%(localReadVgprIdx+regToAdd)),
|
|
comment="copy lsu results to vreg"))
|
|
localReadVgprIdx += regsPerStore
|
|
|
|
# Release write/read resource
|
|
writer.vgprPool.checkIn(addr)
|
|
|
|
# Do Reduction
|
|
module.add(moduleReduction)
|
|
|
|
# Release reduction resource
|
|
writer.vgprPool.checkIn(inLoopTmpVgpr)
|
|
|
|
# Release all resource
|
|
writer.vgprPool.checkIn(lsu_id)
|
|
writer.vgprPool.checkIn(wave_id)
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
# reset vgprValuC register
|
|
module.add(RegSet("v", "vgprValuC", self.accVgprLdsReduction))
|
|
|
|
return module
|
|
|
|
##############################################################################
|
|
# LocalSplitU: Global Write Indices
|
|
##############################################################################
|
|
def globalWriteIndices(self, writer, kernel):
|
|
module = Module("localSplitUGlobalWriteIndices")
|
|
|
|
# Add LSU Offset back
|
|
packedC1 = kernel["PackedC1IndicesX"]
|
|
strideC1 = "StrideC%s" % (writer.states.indexChars[packedC1[0]])
|
|
strideD1 = "StrideD%s" % (writer.states.indexChars[packedC1[0]])
|
|
wave_id = writer.vgprPool.checkOut(1, "tmpWaveID")
|
|
tmpVgpr1 = writer.vgprPool.checkOutAligned(2,2,"tmpVgpr1")
|
|
tmpVgpr1Res = ContinuousRegister(tmpVgpr1, 2)
|
|
module.add(vectorStaticDivide(wave_id, "Serial", kernel["WavefrontSize"], tmpVgpr1Res))
|
|
numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1]
|
|
module.add(vectorStaticDivide(wave_id, wave_id, numWaves, tmpVgpr1Res))
|
|
|
|
with writer.allocTmpSgpr(1, tag="LSUOn_globalWriteIndices_tmpSgprInfo") as tmpSgprInfo:
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
if self.LSUValidOffset0 > 0:
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), \
|
|
src=hex(self.LSUValidOffset0), comment="a valid offset"))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr1), src0=vgpr(wave_id), src1=sgpr(tmpSgpr), comment="wave LSU offset"))
|
|
module.add(VAddU32(dst=vgpr(writer.vgprs.coord0), src0=vgpr(tmpVgpr1), src1=vgpr(writer.vgprs.coord0), comment="coord0 += LSU offset0"))
|
|
else:
|
|
module.addComment0("valid offset coord0 is zero.")
|
|
|
|
if self.LSUValidOffset1 > 0:
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), \
|
|
src=hex(self.LSUValidOffset1), comment="a valid offset"))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr1), src0=vgpr(wave_id), src1=sgpr(tmpSgpr), comment="wave LSU offset"))
|
|
module.add(VAddU32(dst=vgpr(writer.vgprs.coord1), src0=vgpr(tmpVgpr1), src1=vgpr(writer.vgprs.coord1), comment="coord1 += LSU offset1"))
|
|
module.add(VAddU32(dst=vgpr(writer.vgprs.coord1InMT), src0=vgpr(tmpVgpr1), src1=vgpr(writer.vgprs.coord1InMT), comment="coord1InMT += LSU offset1"))
|
|
|
|
# this code is from CouputeStoreVgprs. coord 1 : offset part
|
|
# cinRowPtr/coutRowPtrD are only allocated for BufferStore (SRD-based row offsets)
|
|
if kernel["BufferStore"]:
|
|
packedC1 = kernel["PackedC1IndicesX"]
|
|
strideC1 = "StrideC%s" % (writer.states.indexChars[packedC1[0]])
|
|
strideD1 = "StrideD%s" % (writer.states.indexChars[packedC1[0]])
|
|
module.add(VMulLOU32(dst=vgpr(writer.vgprs.cinRowPtr), src0=vgpr(writer.vgprs.coord1InMT), src1=sgpr(strideC1), comment=" offset 1"))
|
|
module.add(VMulLOU32(dst=vgpr(writer.vgprs.coutRowPtrD), src0=vgpr(writer.vgprs.coord1InMT), src1=sgpr(strideD1), comment=" offset 1"))
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
module.add(VMovB32(dst=vgpr(writer.vgprs.coutRowPtrE), src=vgpr(writer.vgprs.coord1InMT), comment=" save offset 1 for E"))
|
|
if writer.vgprs.coutRowPtrGate != -1:
|
|
module.add(VMovB32(dst=vgpr(writer.vgprs.coutRowPtrGate), src=vgpr(writer.vgprs.coord1InMT), comment=" save offset 1 for Gate"))
|
|
if writer.vgprs.coutRowPtrBias != -1:
|
|
index = packedC1[0] - 1
|
|
strideW1 = "Size%s" % "I" if index == 0 else ("J" if index == 1 else (writer.states.indexChars[index]))
|
|
module.add(VMulLOU32(dst=vgpr(writer.vgprs.coutRowPtrBias), src0=vgpr(writer.vgprs.coord1InMT), src1=sgpr(strideW1), comment=" offset 1"))
|
|
else:
|
|
module.addComment0("valid offset coord1 is zero.")
|
|
|
|
writer.vgprPool.checkIn(tmpVgpr1)
|
|
writer.vgprPool.checkIn(wave_id)
|
|
writer.vgprPool.checkIn(writer.vgprs.coord0InMT)
|
|
writer.vgprPool.checkIn(writer.vgprs.coord1InMT)
|
|
|
|
if kernel["BufferStore"]:
|
|
#print "----AddressC-LocalSplitU"
|
|
#print self.vgprPool.state()
|
|
writer.vgprs.addrE = -1
|
|
writer.vgprs.addrD = -1
|
|
writer.vgprs.addrC = -1
|
|
writer.vgprs.addrBias = -1
|
|
writer.vgprs.addrScaleAVec = -1
|
|
writer.vgprs.addrScaleBVec = -1
|
|
writer.vgprs.addrScaleAlphaVec = -1
|
|
else:
|
|
writer.vgprs.addrE = -1
|
|
writer.vgprs.addrBias = -1
|
|
writer.vgprs.addrScaleAVec = -1
|
|
writer.vgprs.addrScaleBVec = -1
|
|
writer.vgprs.addrScaleAlphaVec = -1
|
|
writer.vgprs.addrD = writer.vgprPool.checkOut(2, tag="LSUOn_globalWriteIndices_addrD")
|
|
module.add(VMovB32(
|
|
dst=vgpr(writer.vgprs.addrD+0), \
|
|
src=sgpr("AddressD+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32(
|
|
dst=vgpr(writer.vgprs.addrD+1), \
|
|
src=sgpr("AddressD+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
writer.vgprs.addrC = writer.vgprPool.checkOut(2, tag="LSUOn_globalWriteIndices_addrC")
|
|
module.add(VMovB32(
|
|
dst=vgpr(writer.vgprs.addrC+0), \
|
|
src=sgpr("AddressC+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32(
|
|
dst=vgpr(writer.vgprs.addrC+1), \
|
|
src=sgpr("AddressC+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
|
|
if kernel["GlobalSplitU"] != 0:
|
|
gsuLabel = Label(label=writer.labels.getNameInc("GSU"), comment="")
|
|
with writer.allocTmpSgpr(1, tag="LSUOn_globalWriteIndices_tmpSgprGSU") as tmpSgprGSU:
|
|
module.add(SAndB32(dst=sgpr(tmpSgprGSU.idx), src0=sgpr("GSU"), src1=writer.gsuMaskHex(kernel), comment="Restore GSU"))
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgprGSU.idx), src1=1, comment="GSU == 1 ?"))
|
|
module.add(SCBranchSCC0(labelName=gsuLabel.getLabelName(), comment="branch if GSU != 1"))
|
|
if kernel["ProblemType"]["UseE"]:
|
|
writer.vgprs.addrE = writer.vgprPool.checkOut(2, 'addrE')
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrE+0), \
|
|
src=sgpr("AddressE+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrE+1), \
|
|
src=sgpr("AddressE+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
if writer.states.useBias == DataDirection.READ:
|
|
writer.vgprs.addrBias = writer.vgprPool.checkOut(2, 'addrBias')
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrBias+0), \
|
|
src=sgpr("AddressBias+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrBias+1), \
|
|
src=sgpr("AddressBias+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector"):
|
|
writer.vgprs.addrScaleAVec = writer.vgprPool.checkOut(2, 'addrScaleAVec')
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrScaleAVec+0), \
|
|
src=sgpr("AddressScaleA+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrScaleAVec+1), \
|
|
src=sgpr("AddressScaleA+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
writer.vgprs.addrScaleBVec = writer.vgprPool.checkOut(2, 'addrScaleVVec')
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrScaleBVec+0), \
|
|
src=sgpr("AddressScaleB+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrScaleBVec+1), \
|
|
src=sgpr("AddressScaleB+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"]:
|
|
writer.vgprs.addrScaleAlphaVec = writer.vgprPool.checkOut(2, 'addrScaleAlphaVec')
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrScaleAlphaVec+0), \
|
|
src=sgpr("AddressScaleAlphaVec+0"), \
|
|
comment="sgpr -> vgpr"))
|
|
module.add(VMovB32( \
|
|
dst=vgpr(writer.vgprs.addrScaleAlphaVec+1), \
|
|
src=sgpr("AddressScaleAlphaVec+1"), \
|
|
comment="sgpr -> vgpr"))
|
|
if kernel["GlobalSplitU"] != 0:
|
|
module.add(gsuLabel)
|
|
|
|
return module
|
|
|
|
|
|
##############################################################################
|
|
# LocalSplitU: Global Write
|
|
##############################################################################
|
|
def globalWrite(self, writer, kernel, tPA, tPB):
|
|
if not writer.do["PostLoop"]: return ""
|
|
|
|
elements_0 = [[] for y in range(2)]
|
|
elements_1 = [[] for y in range(2)]
|
|
elements_f0 = [[] for y in range(2)]
|
|
elements_f1 = [[] for y in range(2)]
|
|
# TODO: LocalSplitU has special global write elements, we will not use all StoreVectorWidths
|
|
(fullVws, elementss_0, fullVws_1, elementss_1) = writer.notLocalFullTileElements(kernel)
|
|
(fullVw, elements_0[False], fullVw_1, elements_1[False]) = (fullVws[0], elementss_0[0], fullVws_1[0], elementss_1[0])
|
|
(edgeVw, elements_0[True], edgeVw_1, elements_1[True] ) = (fullVws[-1], elementss_0[-1], fullVws_1[-1], elementss_1[-1])
|
|
edgeScaled_0 = len(elements_0[True]) // len(elements_1[False])
|
|
edgeScaled_1 = len(elements_1[True]) // len(elements_1[False])
|
|
noEgScaled_0 = len(elements_0[False]) // len(elements_1[False])
|
|
|
|
for i in range(0, len(elements_1[False])):
|
|
element = elements_1[False][i]
|
|
if element in self.LSUelementsPerLSUWave:
|
|
elements_f1[False].append(element)
|
|
for j in range(0, edgeScaled_0):
|
|
# in general, edge will affect vc0 dimension.
|
|
element = elements_0[True][i*edgeScaled_0+j]
|
|
elements_f0[True].append(element)
|
|
for j in range(0, edgeScaled_1):
|
|
# in general, edge will affect vc0 dimension.
|
|
element = elements_1[True][i*edgeScaled_1+j]
|
|
elements_f1[True].append(element)
|
|
for j in range(0, noEgScaled_0):
|
|
# in general, edge will affect vc0 dimension.
|
|
element = elements_0[False][i*noEgScaled_0+j]
|
|
elements_f0[False].append(element)
|
|
|
|
vectorWidths = [fullVw, edgeVw]
|
|
vectorWidths_1 = [fullVw_1, edgeVw_1]
|
|
|
|
noGSUBranch = (kernel["GlobalSplitU"] == 0 and not writer.states.streamK.requiresWorkspaceReductionStorePath)
|
|
module = Module("localSplitUGlobalWrite")
|
|
storeModule, _ = writer.globalWriteElements(kernel, tPA, tPB, vectorWidths, vectorWidths_1, elements_f0, elements_f1, noGSUBranch=noGSUBranch)
|
|
module.add(storeModule)
|
|
writer.cleanupGlobalWrite(kernel)
|
|
writer.vgprPool.checkIn(self.accVgprLdsReduction)
|
|
return module
|