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
1035 lines
66 KiB
Python
1035 lines
66 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 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
|
|
from rocisa.container import EXEC, VCC, vgpr, sgpr
|
|
from rocisa.instruction import MacroInstruction, SAddCU32, SAddU32, SAndB32, \
|
|
SAndB64, SLShiftLeftB32, SMovB32, SMovB64, SMulI32, SSubBU32, SSubU32, \
|
|
VAddCCOU32, VAddCOU32, VAddI32, VAddU32, VAndB32, VBfiB32, \
|
|
VCmpEQU32, VCmpGtU32, VCmpLtU32, VCmpXNeU32, VCndMaskB32, VLShiftLeftB32, \
|
|
VMadI32I24, VMovB32, VMulLOU32, VSubI32, VSubU32
|
|
from rocisa.functions import vectorAddMultiplyBpe, vectorMultiply64Bpe
|
|
from .Common import INDEX_CHARS, DataDirection, log2
|
|
|
|
##############################################################################
|
|
# Fields associated with computing address
|
|
##############################################################################
|
|
class AddrCalculation:
|
|
# rowInc is number of rows to add to the base address
|
|
# coord0Vgpr : This is VGPR that holds coord0. Coord0 is element-space
|
|
# packed index for the 0 coordinate of the C/D matrix.
|
|
# coord1Vgpr : VGPR which tracks the last coord1 calculation.
|
|
# If this is new coord1, just overwrite it with latest calc.
|
|
def __init__(self, kernelWriter, ss, addrCVgpr, addrDVgpr, addrGSUSyncVgprs, addrEVgpr, addrBiasVgpr, addrScaleAVecVgpr, addrScaleBVecVgpr, addrScaleAlphaVecVgpr, element, \
|
|
coordOffset0, coord1Vgpr, coordOffset1, rowInc, newCoord1, vectorDataTypes, addrGateVgpr=None):
|
|
self.kernelWriter = kernelWriter
|
|
|
|
# vgprs for address, could be more than one (for flat)
|
|
self.addrEVgpr = addrEVgpr
|
|
self.addrDVgpr = addrDVgpr
|
|
self.addrGSUSyncVgprs = addrGSUSyncVgprs
|
|
self.addrCVgpr = addrCVgpr
|
|
self.addrBiasVgpr = addrBiasVgpr
|
|
self.addrGateVgpr = addrGateVgpr
|
|
self.addrScaleAVecVgpr = addrScaleAVecVgpr
|
|
self.addrScaleBVecVgpr = addrScaleBVecVgpr
|
|
self.addrScaleAlphaVecVgpr = addrScaleAlphaVecVgpr
|
|
self.coord1Vgpr = coord1Vgpr # vgpr that stores coord1Vgpr
|
|
|
|
self.element = element
|
|
self.coordOffset0 = coordOffset0
|
|
self.coordOffset1 = coordOffset1
|
|
self.rowInc = rowInc
|
|
self.rowIncDirtyRowPtr = 0 # rowInc was used to modify rowPtr, need to recompute addr
|
|
self.newCoord1 = newCoord1 # vgpr that stores newCoord1
|
|
self.vectorDataTypes = vectorDataTypes
|
|
|
|
self.biasOffset = [0, 0]
|
|
self.scaleAVecOffset = 0
|
|
self.scaleBVecOffset = 0
|
|
self.scaleAlphaVecOffset = [0, 0]
|
|
if ss.optSingleColVgpr:
|
|
# optimized stores use the load offset for coordOffset0 calculations.
|
|
self.biasOffset[0] = coordOffset0 * kernelWriter.states.bpeCinternal + self.vectorDataTypes.bias(0).ldsOffset
|
|
self.biasOffset[1] = coordOffset1 * kernelWriter.states.bpeCinternal + self.vectorDataTypes.bias(1).ldsOffset
|
|
self.scaleAVecOffset = coordOffset0 * kernelWriter.states.bpeCinternal + self.vectorDataTypes.scaleA.ldsOffset
|
|
self.scaleBVecOffset = coordOffset1 * kernelWriter.states.bpeCinternal + self.vectorDataTypes.scaleB.ldsOffset
|
|
self.scaleAlphaVecOffset[0] = coordOffset0 * kernelWriter.states.bpeCinternal + self.vectorDataTypes.scaleAlpha(0).ldsOffset
|
|
self.scaleAlphaVecOffset[1] = coordOffset1 * kernelWriter.states.bpeCinternal + self.vectorDataTypes.scaleAlpha(1).ldsOffset
|
|
self.globalOffset = coordOffset0 * kernelWriter.states.bpeCexternal
|
|
self.globalOffsetE = coordOffset0 * kernelWriter.states.bpeE
|
|
self.globalOffsetGate = coordOffset0 * kernelWriter.states.bpeGate
|
|
self.globalOffsetInternal = coordOffset0 * kernelWriter.states.bpeCinternal
|
|
else:
|
|
# else non-opt stores include the coord0 offset into VGPR address calcs
|
|
self.biasOffset[0] = self.vectorDataTypes.bias(0).ldsOffset
|
|
self.biasOffset[1] = self.vectorDataTypes.bias(1).ldsOffset
|
|
self.scaleAVecOffset = self.vectorDataTypes.scaleA.ldsOffset
|
|
self.scaleBVecOffset = self.vectorDataTypes.scaleB.ldsOffset
|
|
self.scaleAlphaVecOffset[0] = self.vectorDataTypes.scaleAlpha(0).ldsOffset
|
|
self.scaleAlphaVecOffset[1] = self.vectorDataTypes.scaleAlpha(1).ldsOffset
|
|
self.globalOffset = 0
|
|
self.globalOffsetE = 0
|
|
self.globalOffsetGate = 0
|
|
self.globalOffsetInternal = 0
|
|
|
|
def addScaled(self, destV, src0, src1, scale1, tmpS01, comment="", comment1="", comment2="scale stride"):
|
|
"""
|
|
Use minimally efficient instructions to add stride*scale
|
|
"""
|
|
|
|
module = Module("addScaled")
|
|
if scale1 == 1:
|
|
module.add(VAddU32(dst=destV, src0=src0, src1=src1, comment=comment1 if comment1!="" else comment))
|
|
else:
|
|
module.add(SMulI32(dst=sgpr(tmpS01), src0=src1, src1=scale1, comment=comment2))
|
|
if comment1!="":
|
|
module.add(VAddU32(dst=destV, src0=src0, src1=sgpr(tmpS01), comment=comment))
|
|
else:
|
|
module.add(VAddI32(dst=destV, src0=src0, src1=sgpr(tmpS01), comment=comment))
|
|
return module
|
|
|
|
def emitCoord1Advance(self, rowInc, tmpS01, d0=0, vc0=0, comment=None, scomment=None):
|
|
"""
|
|
Advance the shared coord1 VGPR (kernelWriter.vgprs.coord1) in place by
|
|
`rowInc` rows. This is the single canonical "coord1 += rowInc" emitter,
|
|
shared by the per-element cross-row advance (emitAddressCoordIncrement)
|
|
and the CLS look-ahead at batch end (GlobalWriteBatch.emit()).
|
|
rowInc == 0 -> nothing
|
|
0 < rowInc <= 64 -> v_add_co_u32 with immediate
|
|
rowInc > 64 -> s_mov + v_add_co_u32
|
|
rowInc < 0 -> s_mov + v_add_nc_i32
|
|
`comment`/`scomment` override the v_add / s_mov comments (used by the CLS
|
|
look-ahead to keep its "coord1.la ... (look-ahead)" marker); when None the
|
|
canonical per-element comments (coord1.1/.2/.3) are emitted so the
|
|
per-element / non-CLS .s stays byte-identical.
|
|
"""
|
|
module = Module("coord1Advance")
|
|
coord1 = self.kernelWriter.vgprs.coord1
|
|
if scomment is None:
|
|
scomment = "rowInc d1=%u vc1=%u"%(d0, vc0)
|
|
if rowInc == 0:
|
|
pass
|
|
elif 0 < rowInc <= 64:
|
|
module.add(VAddCOU32(dst=vgpr(coord1), dst1=VCC(), src0=vgpr(coord1), \
|
|
src1=rowInc, comment=comment if comment else "coord1.1: coord1Vgpr += d1*sg1*VW + vc1"))
|
|
elif rowInc > 64:
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=rowInc, comment=scomment))
|
|
module.add(VAddCOU32(dst=vgpr(coord1), dst1=VCC(), src0=vgpr(coord1), \
|
|
src1=sgpr(tmpS01), comment=comment if comment else "coord1.2: coord1 += d1*sg1*VW + vc1"))
|
|
else: # rowInc < 0
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=rowInc, comment=scomment))
|
|
module.add(VAddI32(dst=vgpr(coord1), src0=vgpr(coord1), \
|
|
src1=sgpr(tmpS01), comment=comment if comment else "coord1.3: coord1 += d1*sg1*VW + vc1"))
|
|
return module
|
|
|
|
def emitRowPtrAdvance(self, kernel, ss, tmpS01, rowInc, lookahead=False):
|
|
"""
|
|
Advance the row-pointer VGPRs (cinRowPtr / coutRowPtrD / coutRowPtrE /
|
|
coutRowPtrBias, or packed-C1 extract) by `rowInc` rows. This is the single
|
|
canonical row-pointer advance, shared by:
|
|
- the per-element cross-row advance in emitAddressSetupCode (lookahead=False)
|
|
- the CLS look-ahead at batch end in GlobalWriteBatch.emit() (lookahead=True)
|
|
Same spirit as the delayed-primer reuse of incrementToNextRow: the look-ahead
|
|
does NOT re-implement the advance, it produces it from the same source -- so it
|
|
automatically covers E / Bias / packed-C1 with the exact same conditions.
|
|
|
|
`lookahead` only switches comment text (.la markers) and skips the per-element
|
|
bookkeeping (rowIncDirtyRowPtr + the "Fix for UseInitialStridesCD" banner). The
|
|
emitted instructions are otherwise identical, so non-CLS / per-element .s stays
|
|
byte-identical.
|
|
"""
|
|
module = Module("emitRowPtrAdvance")
|
|
kw = self.kernelWriter
|
|
if not lookahead:
|
|
self.rowIncDirtyRowPtr = 1
|
|
#assert (not kernel["ProblemType"]["UseInitialStridesCD"])
|
|
module.addComment1("Fix for UseInitialStridesCD, emitAddressSetupCode")
|
|
|
|
if len(kernel["PackedC1IndicesX"]) == 1:
|
|
strideChar = self.kernelWriter.states.indexChars[kernel["PackedC1IndicesX"][0]]
|
|
if lookahead:
|
|
module.add(self.addScaled(vgpr(kw.vgprs.cinRowPtr), vgpr(kw.vgprs.cinRowPtr), \
|
|
sgpr("StrideC%s"%strideChar), rowInc, tmpS01, \
|
|
comment="cinRowPtr.la: += StrideC%s*rowInc"%strideChar, \
|
|
comment1="cinRowPtr.la: += StrideC%s"%strideChar, \
|
|
comment2="scale StrideC%s for look-ahead"%strideChar))
|
|
module.add(self.addScaled(vgpr(kw.vgprs.coutRowPtrD), vgpr(kw.vgprs.coutRowPtrD), \
|
|
sgpr("StrideD%s"%strideChar), rowInc, tmpS01, \
|
|
comment="coutRowPtrD.la: += StrideD%s*rowInc"%strideChar, \
|
|
comment1="coutRowPtrD.la: += StrideD%s"%strideChar, \
|
|
comment2="scale StrideD%s for look-ahead"%strideChar))
|
|
else:
|
|
module.add(self.addScaled(vgpr(kw.vgprs.cinRowPtr), vgpr(kw.vgprs.cinRowPtr), \
|
|
sgpr("StrideC%s"%strideChar), rowInc, tmpS01, "ROWINC- Move cinRowPtr to next row"))
|
|
module.add(self.addScaled(vgpr(kw.vgprs.coutRowPtrD), vgpr(kw.vgprs.coutRowPtrD), \
|
|
sgpr("StrideD%s"%strideChar), rowInc, tmpS01, "Move coutRowPtrD to next row"))
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
ecomment = "coutRowPtrE.la: += StrideE%s*rowInc"%strideChar if lookahead else "Move coutRowPtrE to next row"
|
|
module.add(self.addScaled(vgpr(kw.vgprs.coutRowPtrE), vgpr(kw.vgprs.coutRowPtrE), \
|
|
sgpr("StrideE%s"%strideChar), rowInc, tmpS01, ecomment))
|
|
if kw.vgprs.coutRowPtrBias != -1:
|
|
index = kernel["PackedC1IndicesX"][0] - 1
|
|
strideW1 = "Size%s" % "I" if index == 0 else ("J" if index == 1 else (kw.states.indexChars[index]))
|
|
bcomment = "coutRowPtrBias.la: += %s*rowInc"%strideW1 if lookahead else "Move coutRowPtrBias to next row"
|
|
module.add(self.addScaled(vgpr(kw.vgprs.coutRowPtrBias), vgpr(kw.vgprs.coutRowPtrBias), \
|
|
sgpr(strideW1), rowInc, tmpS01, bcomment))
|
|
if kw.vgprs.coutRowPtrGate != -1:
|
|
module.add(self.addScaled(vgpr(kw.vgprs.coutRowPtrGate), vgpr(kw.vgprs.coutRowPtrGate), \
|
|
sgpr("GateStride+0"), self.rowInc, tmpS01, "Move coutRowPtrGate to next row"))
|
|
elif len(kernel["PackedC1IndicesX"]) > 1:
|
|
module.add(kw.extractPackedCoord1ToRowStart(kernel, kernel["PackedC1IndicesX"] , self.coord1Vgpr, 'D'))
|
|
return module
|
|
|
|
|
|
def emitAddressCoordIncrement(self, kernel, ss, tmpVgpr, tmpS01, updateCoord1,
|
|
skipCrossBatchCoord1Inc=False):
|
|
"""
|
|
Emit code that computes the coord0 and coord1 for this element
|
|
sets self.coord0Vgpr with the address that holds the coord0 value for this element.
|
|
Input:
|
|
- tmpVgpr is a 1 temporary VGPR used for coord0 calculation on edges
|
|
"""
|
|
|
|
module = Module("emitAddressCoordIncrement")
|
|
kw = self.kernelWriter
|
|
(d1,d0,vc1,vc0) = self.element
|
|
self.coord0Vgpr = None # will set below
|
|
|
|
# module.addComment0("store addr=v%u coordOffset0=%u"% \
|
|
# (self.addr, self.coordOffset0))
|
|
module.addComment0("(d1,vc1,d0,vc0)=(%u,%u,%u,%u)"\
|
|
% (d1,vc1,d0,vc0))
|
|
if ss.optSingleColVgpr:
|
|
self.coord0Vgpr = kw.vgprs.coord0
|
|
elif not ss.optSharedColVgpr or (d1 == vc1 == 0):
|
|
# not share mode or first row always does the address calc math:
|
|
|
|
if self.coordOffset0 == 0:
|
|
self.coord0Vgpr = kw.vgprs.coord0
|
|
elif self.coordOffset0 <= 64:
|
|
self.coord0Vgpr = tmpVgpr
|
|
module.add(VAddCOU32(dst=vgpr(self.coord0Vgpr), dst1=VCC(), src0=vgpr(kw.vgprs.coord0), src1=self.coordOffset0, \
|
|
comment="coord0.1: coord0 += d0*sg0*VW + vc0"))
|
|
else:
|
|
self.coord0Vgpr = tmpVgpr
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=self.coordOffset0, comment="coordOffset0 d0=%u vc0=%u"%(d0, vc0)))
|
|
module.add(VAddCOU32(dst=vgpr(self.coord0Vgpr), dst1=VCC(), src0=vgpr(kw.vgprs.coord0), src1=sgpr(tmpS01), \
|
|
comment="coord0.2: coord0 += d0*sg0*VW + vc0"))
|
|
|
|
if self.newCoord1:
|
|
if not kernel["BufferStore"] or updateCoord1:
|
|
if self.rowInc== 0:
|
|
None
|
|
elif skipCrossBatchCoord1Inc:
|
|
module.addComment0("coord1Vgpr cross-batch inc skipped (CLS look-ahead)")
|
|
else:
|
|
# coord1 += rowInc (canonical, shared with CLS look-ahead).
|
|
# self.coord1Vgpr == kernelWriter.vgprs.coord1 (set in
|
|
# AsmStoreState), so this is an in-place advance.
|
|
module.add(self.emitCoord1Advance(self.rowInc, tmpS01, d0, vc0))
|
|
return module
|
|
|
|
def getRowPtr(self, kw, tc):
|
|
if tc == 'C':
|
|
return kw.vgprs.cinRowPtr
|
|
elif tc == 'E':
|
|
return kw.vgprs.coutRowPtrE
|
|
elif tc == 'Bias':
|
|
return kw.vgprs.coutRowPtrBias
|
|
elif tc == 'Gate':
|
|
return kw.vgprs.coutRowPtrGate
|
|
else:
|
|
return kw.vgprs.coutRowPtrD
|
|
|
|
def getAddrVgpr(self, kw, tc):
|
|
if tc == 'C':
|
|
return self.addrCVgpr
|
|
elif tc == 'E':
|
|
return self.addrEVgpr
|
|
elif tc == 'Bias':
|
|
return self.addrBiasVgpr
|
|
elif tc == 'Gate':
|
|
return self.addrGateVgpr
|
|
elif tc == 'TD':
|
|
return self.addrGSUSyncVgprs
|
|
else:
|
|
return self.addrDVgpr
|
|
|
|
# storeChar is 'C' or 'D'
|
|
# elementVgpr is coord0Vgpr*strideCD0, or optimized to just coord0Vgpr if strideCD0 is unit const
|
|
def emitExtractAndScalePackedDims(self, kernel, ss, tmpVgpr, storeChar):
|
|
module = Module("emitExtractAndScalePackedDims")
|
|
kw = self.kernelWriter
|
|
packedIndices = kernel["PackedC0IndicesX"]
|
|
packedBits = self.coord0Vgpr # start with coord0, will move to temp below
|
|
rowPtr = self.getRowPtr(kw, storeChar)
|
|
addrVgpr = self.getAddrVgpr(kw, storeChar)
|
|
bpe = kw.states.bpeCinternal if (tc == 'Bias') else (kw.states.bpeE if (tc == 'E') else kw.states.bpeCexternal)
|
|
|
|
for i,idx in enumerate(packedIndices[:-1]):
|
|
# vgprTmp assignments:
|
|
# - tmp+0 may be the incoming packed coordinate 0, used on replay too
|
|
# - tmp+1 is DIV output
|
|
# - tmp+2 is scratch
|
|
idxChar= INDEX_CHARS[idx]
|
|
module.addComment0("extract %s"%kw.sizeRef(idx))
|
|
assert(tmpVgpr+1 != packedBits) # bad since we still need packedBits below for remainder (can't overwrite here)
|
|
if kw.states.asmCaps["HasVgprMSB"]:
|
|
module.add(VMagicDiv(kernel["MagicDivAlg"], tmpVgpr+1, vgpr(packedBits), sgpr("MagicNumberSize%s"%idxChar), \
|
|
sgpr("MagicShiftSize%s"%idxChar), sgpr("MagicAbitSize%s"%idxChar)))
|
|
else:
|
|
module.add(MacroInstruction("V_MAGIC_DIV", \
|
|
args=[tmpVgpr+1, vgpr(packedBits), sgpr("MagicNumberSize%s"%idxChar), \
|
|
sgpr("MagicShiftSize%s"%idxChar), sgpr("MagicAbitSize%s"%idxChar) if kernel["MagicDivAlg"]==2 else "0"]))
|
|
# tmpVgpr+1 returns the quotient, tmpVgpr+2 is overwritten
|
|
|
|
# compute remainder, packedBits % sizeIdx - this is the 'extracted' index that must be scaled
|
|
# remainder is mul and sub
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr+2), src0=vgpr(tmpVgpr+1), src1=kw.sizeRef(idx), \
|
|
comment="remainder part 1"))
|
|
module.add(VSubU32(dst=vgpr(tmpVgpr+2), src0=vgpr(packedBits), src1=vgpr(tmpVgpr+2),
|
|
comment="remainder part 2"))
|
|
|
|
if i==0:
|
|
module.add(VMulLOU32(dst=vgpr(addrVgpr), src0=vgpr(tmpVgpr+2), \
|
|
src1=kw.strideRef(storeChar, idx), comment="addrCalc <- scaled extracted dim"))
|
|
else:
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr+2), src0=vgpr(tmpVgpr+2), \
|
|
src1=kw.strideRef(storeChar, idx), comment="scale extracted dim"))
|
|
module.add(VAddU32(dst=vgpr(addrVgpr), src0=vgpr(addrVgpr), \
|
|
src1=vgpr(tmpVgpr+2), comment="addrCalc += scaled extracted dim "))
|
|
|
|
if i < len(packedIndices)-2:
|
|
# TODO - might be able to eliminate this
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr+0), src=vgpr(tmpVgpr+1), \
|
|
comment="Copy remaining bits for next divide"))
|
|
packedBits = tmpVgpr+0
|
|
|
|
if len(packedIndices)>1:
|
|
# if we unpacked something, then scale it to BPE
|
|
module.addComment0("extract final %s"%kw.sizeRef(packedIndices[-1]))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr+2), src0=vgpr(tmpVgpr+1), \
|
|
src1=kw.strideRef(storeChar, packedIndices[-1]), comment="scale final extracted dim"))
|
|
module.add(VAddU32(dst=vgpr(addrVgpr), src0=vgpr(addrVgpr), \
|
|
src1=vgpr(tmpVgpr+2), comment="addrCalc += scaled extracted dim "))
|
|
|
|
module.add(vectorAddMultiplyBpe(addrVgpr, rowPtr, addrVgpr, bpe, \
|
|
comment="packed: add rowPtr and scaleToBpe"))
|
|
|
|
return module
|
|
|
|
def emitScaleToBpe(self, kernel, ss, tmpVgpr, tmpSgpr, singleUpdate, tc, dim):
|
|
"""
|
|
Needs 3 temporary VGPRs
|
|
"""
|
|
|
|
module = Module("emitScaleToBpe")
|
|
kw = self.kernelWriter
|
|
(d1,d0,vc1,vc0) = self.element
|
|
rowPtr = self.getRowPtr(kw, tc)
|
|
addrVgpr = self.getAddrVgpr(kw, tc)
|
|
if tc == 'Gate' and addrVgpr is None:
|
|
addrVgpr = self.addrDVgpr
|
|
if tc == 'Bias':
|
|
bpe = kw.states.bpeCinternal
|
|
elif tc == 'E':
|
|
bpe = kw.states.bpeE
|
|
elif tc == 'Gate':
|
|
bpe = kw.states.bpeGate
|
|
else:
|
|
bpe = kw.states.bpeCexternal
|
|
if (tc == 'C' or tc == 'TD'):
|
|
bpe = bpe if (kernel["_GlobalAccumulation"] != "MultipleBufferSingleKernel") else kw.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters()
|
|
# set when we generate code that updates the address
|
|
# optSingleColVgpr and optSharedColVgpr attempt to minimize these updates
|
|
updatedAddr = False
|
|
|
|
# scale and set final address:
|
|
# Gate uses D's stride for col-direction (gate is per-element along col like D).
|
|
if tc == 'Gate':
|
|
stride0 = kw.strideRef('D', 0)
|
|
elif (tc == 'Bias') or (tc == 'ScaleAlphaVec') or (tc == 'ScaleAVec') or (tc == 'ScaleBVec'):
|
|
stride0 = kw.strideRef('D', 0)
|
|
else:
|
|
stride0 = kw.strideRef(tc, 0)
|
|
if kw.isConstUnitStride(stride0):
|
|
elementVgpr = self.coord0Vgpr
|
|
else:
|
|
module.add(VMulLOU32(dst=vgpr(addrVgpr), \
|
|
src0=vgpr(self.coord0Vgpr), \
|
|
src1=stride0, \
|
|
comment="scale element by non-unit stride"))
|
|
elementVgpr = addrVgpr
|
|
|
|
isSingleKernel = ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or kernel["StreamK"] > 0
|
|
if ss.optSingleColVgpr:
|
|
# This is first element in the first batch, create a byte address that will
|
|
# be re-used by subsequent elements:
|
|
# if this element is firstInBatch - may need to set up a bpe-scaled row pointer for the batch:
|
|
# - need row-ptr start of each batch
|
|
assert (kw.vgprs.coord0 == self.coord0Vgpr) # elementAddr assignment above assumes these are the same
|
|
if singleUpdate:
|
|
updatedAddr = True
|
|
if tc == 'C':
|
|
singleColAddrUpdated = ss.singleColCAddrUpdated
|
|
elif tc == 'E':
|
|
singleColAddrUpdated = ss.singleColEAddrUpdated
|
|
elif tc == 'Bias':
|
|
singleColAddrUpdated = ss.singleColBiasAddrUpdated
|
|
elif tc == 'TD':
|
|
singleColAddrUpdated = ss.singleColTDAddrUpdated
|
|
elif tc == 'Gate':
|
|
singleColAddrUpdated = ss.singleColGateAddrUpdated
|
|
else:
|
|
singleColAddrUpdated = ss.singleColDAddrUpdated
|
|
if not singleColAddrUpdated or not ss.optSrdIncForRow:
|
|
if tc == 'Bias' and kw.states.useBias == DataDirection.READ:
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
module.add(VSubU32(dst=vgpr(self.addrBiasVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrBiasVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrBiasVgpr), \
|
|
comment="Bias address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrBiasVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrBiasVgpr), \
|
|
comment="add lds offset"))
|
|
ss.singleColBiasAddrUpdated = True
|
|
return module
|
|
if tc == 'ScaleAlphaVec' and kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
if self.addrScaleAlphaVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAlphaVecVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="ScaleAlpha address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleAVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleAVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile0"], src1=sgpr("WorkGroup0"), comment="wgp0 * MT0"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAVecVgpr), src0=vgpr(self.coord0Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="ScaleAVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleBVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleBVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile1"], src1=sgpr("WorkGroup1"), comment="wgp1 * MT1"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleBVecVgpr), src0=vgpr(self.coord1Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="ScaleBVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'C':
|
|
ss.singleColCAddrUpdated = True
|
|
elif tc == 'E':
|
|
ss.singleColEAddrUpdated = True
|
|
elif tc == 'Bias':
|
|
ss.singleColBiasAddrUpdated = True
|
|
elif tc == 'TD':
|
|
ss.singleColTDAddrUpdated = True
|
|
elif tc == 'Gate':
|
|
ss.singleColGateAddrUpdated = True
|
|
else:
|
|
ss.singleColDAddrUpdated = True
|
|
module.add(vectorAddMultiplyBpe(addrVgpr, rowPtr, elementVgpr, bpe, \
|
|
comment="optSingleColVgpr scaleToBpe: sharedAddrVgpr <- cinRowPtr + coord0, scaled by BPE. BSHERE:coord0=%d, coord0Vgpr=%d"%(kw.vgprs.coord0, self.coord0Vgpr)))
|
|
elif ss.optSharedColVgpr:
|
|
# Need an address calculation for the first address in each row:
|
|
if d1==0 and vc1==0:
|
|
if tc == 'Bias' and kw.states.useBias == DataDirection.READ:
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
module.add(VSubU32(dst=vgpr(self.addrBiasVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrBiasVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrBiasVgpr), \
|
|
comment="Bias address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrBiasVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrBiasVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleAlphaVec' and kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
if self.addrScaleAlphaVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAlphaVecVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="ScaleAlpha address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleAVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleAVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile0"], src1=sgpr("WorkGroup0"), comment="wgp0 * MT0"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAVecVgpr), src0=vgpr(self.coord0Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="ScaleAVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleBVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleBVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile1"], src1=sgpr("WorkGroup1"), comment="wgp1 * MT1"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleBVecVgpr), src0=vgpr(self.coord1Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="ScaleBVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
packedIndices = kernel["PackedC0IndicesX"]
|
|
if len(packedIndices) > 1:
|
|
updatedAddr = True
|
|
module.add(self.emitExtractAndScalePackedDims(kernel, ss, tmpVgpr, tc))
|
|
else:
|
|
updatedAddr = True
|
|
module.add(vectorAddMultiplyBpe(addrVgpr, rowPtr, elementVgpr, bpe,
|
|
comment="optSharedColVgpr scaleToBpe for first row: col addr <- cinRowPtr + coord0, scaled by BPE"))
|
|
else:
|
|
# Generate final address calculation (to bytes) for each element
|
|
# The unpacking takes 8-10 instructions so could be worth optimizing someday :
|
|
# each col has same offset so could create a class to hold column-specific state including
|
|
# the byte address offset for that col and the mask in/out.
|
|
if tc == 'Bias' and kw.states.useBias == DataDirection.READ:
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
module.add(VSubU32(dst=vgpr(self.addrBiasVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrBiasVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrBiasVgpr), \
|
|
comment="Bias address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrBiasVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrBiasVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleAlphaVec' and kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
if self.addrScaleAlphaVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAlphaVecVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="ScaleAlpha address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleAVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleAVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile0"], src1=sgpr("WorkGroup0"), comment="wgp0 * MT0"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAVecVgpr), src0=vgpr(self.coord0Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="ScaleAVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
if tc == 'ScaleBVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleBVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile1"], src1=sgpr("WorkGroup1"), comment="wgp1 * MT1"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleBVecVgpr), src0=vgpr(self.coord1Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
shiftHex=hex(log2(kw.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="ScaleBVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="add lds offset"))
|
|
return module
|
|
packedIndices = kernel["PackedC0IndicesX"]
|
|
if len(packedIndices) > 1:
|
|
updatedAddr = True
|
|
module.add(self.emitExtractAndScalePackedDims(kernel, ss, tmpVgpr, tc))
|
|
else:
|
|
updatedAddr = True
|
|
module.add(vectorAddMultiplyBpe(addrVgpr, rowPtr, elementVgpr, bpe, \
|
|
comment="scaleToBpe: accumulate d0 lower and *= bpe into Cin addr"))
|
|
|
|
# if not optSrdIncForRow then we may have moved the row pointer
|
|
# and depending on paths above may not have refreshed addrVgpr already.
|
|
# if so - do it here:
|
|
if self.rowIncDirtyRowPtr and not updatedAddr:
|
|
module.add(vectorAddMultiplyBpe(addrVgpr, rowPtr, kw.vgprs.coord0, bpe, \
|
|
comment="scaleToBpe: Update address with new rowPtr"))
|
|
|
|
return module
|
|
|
|
def edgeProtectCode(self, kernel, edge, beta, atomic, mask, tmpSgpr):
|
|
"""
|
|
Generate code to protect address offset in edge case
|
|
"""
|
|
|
|
module = Module("edgeProtectCode")
|
|
kw = self.kernelWriter
|
|
tmpS01 = tmpSgpr
|
|
tmpS23 = tmpSgpr+self.kernelWriter.states.laneSGPRCount
|
|
|
|
laneSGPRCount = self.kernelWriter.states.laneSGPRCount
|
|
wavefrontSize = kernel["WavefrontSize"]
|
|
|
|
# Now do the edge check and compute the address in bytes:
|
|
if kernel["BufferStore"]:
|
|
if edge and (not kernel["StoreRemapVectorWidth"] or (kernel["StoreRemapVectorWidth"] and (beta or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"))):
|
|
# Set address to -1 if OOB on either dimension
|
|
# and only check the x/coord0 index here, save a couple inst
|
|
sizeBoundary = [0,0]
|
|
sizeBoundary[0] = \
|
|
sgpr("PackedSize0") if len(kernel["PackedC0IndicesX"]) > 1 \
|
|
else kw.sizeRef(kernel["ProblemType"]["Index0"])
|
|
sizeBoundary[1] = \
|
|
sgpr("PackedSize1") if len(kernel["PackedC1IndicesX"]) > 1 \
|
|
else kw.sizeRef(kernel["ProblemType"]["Index1"])
|
|
|
|
module.add(VCmpLtU32(dst=sgpr(tmpS01,laneSGPRCount), src0=vgpr(self.coord0Vgpr), src1=sizeBoundary[0], comment="coord0 < size0" ))
|
|
module.add(VCmpLtU32(dst=sgpr(mask,laneSGPRCount), src0=vgpr(self.coord1Vgpr), src1=sizeBoundary[1], comment="coord1 < size1" ))
|
|
SAndX = SAndB64 if wavefrontSize == 64 else SAndB32
|
|
module.add(SAndX(dst=sgpr(mask,laneSGPRCount), src0=sgpr(tmpS01,laneSGPRCount), src1=sgpr(mask,laneSGPRCount), comment="in0 && in1" ))
|
|
else:
|
|
module.add(VCmpLtU32(dst=sgpr(tmpS01,laneSGPRCount), src0=vgpr(self.coord0Vgpr), src1=sgpr("SizesFree+0"), comment="coord0 < size0" ))
|
|
module.add(VCmpLtU32(dst=sgpr(tmpS23,laneSGPRCount), src0=vgpr(self.coord1Vgpr), src1=sgpr("SizesFree+1"), comment="coord1 < size1" ))
|
|
SAndX = SAndB64 if wavefrontSize == 64 else SAndB32
|
|
module.add(SAndX(dst=sgpr(mask,laneSGPRCount), src0=sgpr(tmpS01,laneSGPRCount), src1=sgpr(tmpS23,laneSGPRCount), comment="in0 && in1" ))
|
|
|
|
if (beta or atomic):
|
|
SMovX = SMovB64 if wavefrontSize == 64 else SMovB32
|
|
module.add(SMovX(dst=EXEC(), src=sgpr(mask,laneSGPRCount), comment="sgprs -> exec" ))
|
|
|
|
return module
|
|
|
|
# TODO - mask should be part of AddrCalc state not passed as parm
|
|
def emitAddressSetupCode(self, kernel, tPB, ss, tmpVgpr, tmpS01, edge, beta, atomic, elementIdx, addrVgpr,
|
|
skipCrossBatchAdvance=False):
|
|
"""
|
|
Generate code to set up the address vgpr
|
|
Input:
|
|
tmpVgpr : two temp vgprs
|
|
Output:
|
|
Returns kStr with appropriate setup code
|
|
Sets self.coord0Vgpr with vgpr that contains the coord0 for this element. This enables
|
|
optimization - if no setup code is required the coord0 can be the input.
|
|
"""
|
|
|
|
module = Module("emitAddressSetupCode")
|
|
kw = self.kernelWriter
|
|
|
|
updateCoord1 = (edge or len(kernel["PackedC1IndicesX"]) > 1)
|
|
module.add(self.emitAddressCoordIncrement(kernel, ss, tmpVgpr, tmpS01, updateCoord1,
|
|
skipCrossBatchCoord1Inc=skipCrossBatchAdvance))
|
|
|
|
# calculate flat load offset
|
|
if not kernel["BufferStore"]:
|
|
# flat: in-bounds exec mask
|
|
# global offset macro (requires 3 tmpVgpr)
|
|
# final address = C + index*bytes
|
|
params = ["%u" % addrVgpr]
|
|
for i in range(0, kernel["ProblemType"]["NumIndicesC"]):
|
|
if i == kernel["ProblemType"]["Index0"]:
|
|
params.append("%s" % (self.coord0Vgpr))
|
|
elif i == kernel["ProblemType"]["Index1"]:
|
|
params.append("%s" % (self.coord1Vgpr))
|
|
else: # just a group index
|
|
params.append("sgprWorkGroup%u"%i)
|
|
params.append("%s" % (tmpVgpr+2))
|
|
if kw.states.asmCaps["HasVgprMSB"]:
|
|
module.add(kw.globalOffset(kernel, None, "C", params))
|
|
else:
|
|
module.add(MacroInstruction(name="GLOBAL_OFFSET_C", args=params))
|
|
module.add(vectorMultiply64Bpe(addrVgpr, addrVgpr, kw.states.bpeCexternal, tmpVgpr))
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr+2), src=vgpr(addrVgpr+0), comment="temp store offset 0"))
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr+3), src=vgpr(addrVgpr+1), comment="temp store offset 1"))
|
|
|
|
# Move the row ptr VGPR
|
|
# optSrdIncForRow moves the SRD so don't move here
|
|
if not ss.optSrdIncForRow and kernel["BufferStore"]:
|
|
if self.rowInc != 0 and skipCrossBatchAdvance:
|
|
self.rowIncDirtyRowPtr = 1
|
|
module.addComment1("rowPtr cross-batch advance skipped (CLS look-ahead)")
|
|
elif self.rowInc != 0:
|
|
# Per-element cross-row advance via the canonical single-source helper
|
|
# (also used by the CLS look-ahead in GlobalWriteBatch.emit()).
|
|
module.add(self.emitRowPtrAdvance(kernel, ss, tmpS01, self.rowInc, lookahead=False))
|
|
|
|
# Shift Pointer for MFMA:
|
|
# For MFMA shift pointer, correct data is stored in another thread.
|
|
# Therefore, MFMA cannot use v_mov to amend store data
|
|
# It needs to modify the coord1 of thread directly.
|
|
# (Use ShiftVectorComponentsMFMA now)
|
|
if 0 and (not kernel["SourceSwap"]) and (not kernel["GuaranteeNoPartialB"]) and tPB["rtv"] and kernel["EnableMatrixInstruction"] and edge:
|
|
(d1,d0,vc1,vc0) = self.element
|
|
if (d1 == vc1 == d0 == vc0 == 0) or self.newCoord1:
|
|
sgprCnt = self.kernelWriter.states.laneSGPRCount
|
|
waveSize = kernel["WavefrontSize"]
|
|
packedC1 = kernel["PackedC1IndicesX"]
|
|
strideC1 = "StrideC%s" % (kw.states.indexChars[packedC1[0]])
|
|
strideD1 = "StrideD%s" % (kw.states.indexChars[packedC1[0]])
|
|
strideE1 = "StrideE%s" % (kw.states.indexChars[packedC1[0]])
|
|
index = packedC1[0] - 1
|
|
strideW1 = "Size%s" % "I" if index == 0 else ("J" if index == 1 else (kw.states.indexChars[index]))
|
|
|
|
module.addComment1("shift vector components d1")
|
|
vw = kernel["GlobalReadVectorWidthB"]
|
|
vTmp1 = tmpVgpr
|
|
vTmp2 = tmpVgpr+1
|
|
sTmp1 = tmpS01
|
|
sTmp2 = tmpS01+sgprCnt
|
|
# check conditions
|
|
module.add(VBfiB32(dst=vgpr(vTmp1), src0=vw-1, src1=0, src2=vgpr(self.coord1Vgpr), comment="coord1 & ~(vw-1)"))
|
|
module.add(VBfiB32(dst=vgpr(vTmp2), src0=vw-1, src1=0, src2=sgpr("SizesFree+%u"%tPB["idx"]), comment="sizeFree & ~(vw-1)"))
|
|
module.add(VCmpEQU32(dst=sgpr(sTmp1,sgprCnt), src0=vgpr(vTmp1), src1=vgpr(vTmp2), comment="if coord1 is in edge glvw"))
|
|
module.add(VAndB32(dst=vgpr(vTmp2), src0=sgpr("SizesFree+%u"%tPB["idx"]), src1=vw-1, comment="sizeFree mod VW"))
|
|
module.add(VCmpGtU32(dst=sgpr(sTmp2,sgprCnt), src0=vgpr(vTmp2), src1=0, comment="this problem is not multiple size of glvw"))
|
|
SAndBX = SAndB64 if waveSize == 64 else SAndB32
|
|
module.add(SAndBX(dst=sgpr(sTmp1,sgprCnt), src0=sgpr(sTmp1,sgprCnt), src1=sgpr(sTmp2,sgprCnt), comment="AND both conditions"))
|
|
# calculate new coord
|
|
module.add(VAddU32(dst=vgpr(vTmp1), src0=vgpr(self.coord1Vgpr), src1=vgpr(vTmp2), comment="shift coord1"))
|
|
module.add(VBfiB32(dst=vgpr(vTmp1), src0=vw-1, src1=vgpr(vTmp1), src2=sgpr("SizesFree+%u"%tPB["idx"]), comment="new coord1 = (shift coord1 & (vw-1)) | (sizeFree & ~(vw-1))"))
|
|
module.add(VSubI32(dst=vgpr(vTmp2), src0=vgpr(vTmp1), src1=vgpr(self.coord1Vgpr), comment="shift how many column"))
|
|
module.add(VCndMaskB32(dst=vgpr(self.coord1Vgpr), src0=vgpr(self.coord1Vgpr), src1=vgpr(vTmp1), \
|
|
src2=sgpr(sTmp1,sgprCnt), comment="set new coord1 if meet conditions" ))
|
|
|
|
module.add(VMadI32I24(dst=vgpr(vTmp1), src0=sgpr(strideC1), src1=vgpr(vTmp2), src2=vgpr(kw.vgprs.cinRowPtr), \
|
|
comment="new rowStart address += shift column * StridesC"))
|
|
module.add(VCndMaskB32(dst=vgpr(kw.vgprs.cinRowPtr), src0=vgpr(kw.vgprs.cinRowPtr), src1=vgpr(vTmp1), src2=sgpr(sTmp1,sgprCnt), \
|
|
comment="set new rowStart if meet conditions" ))
|
|
module.add(VMadI32I24(dst=vgpr(vTmp1), src0=sgpr(strideD1), src1=vgpr(vTmp2), src2=vgpr(kw.vgprs.coutRowPtrD), \
|
|
comment="new rowStart address += shift column * StridesD"))
|
|
module.add(VCndMaskB32(dst=vgpr(kw.vgprs.coutRowPtrD), src0=vgpr(kw.vgprs.coutRowPtrD), src1=vgpr(vTmp1), src2=sgpr(sTmp1,sgprCnt), \
|
|
comment="set new rowStart if meet conditions" ))
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
module.add(VMadI32I24(dst=vgpr(vTmp1), src0=sgpr(strideE1), src1=vgpr(vTmp2), src2=vgpr(kw.vgprs.coutRowPtrE), \
|
|
comment="new rowStart address += shift column * StridesE"))
|
|
module.add(VCndMaskB32(dst=vgpr(kw.vgprs.coutRowPtrE), src0=vgpr(kw.vgprs.coutRowPtrE), src1=vgpr(vTmp1), src2=sgpr(sTmp1,sgprCnt), \
|
|
comment="set new rowStart if meet conditions" ))
|
|
if kw.vgprs.coutRowPtrBias != -1:
|
|
module.add(VMadI32I24(dst=vgpr(vTmp1), src0=sgpr(strideW1), src1=vgpr(vTmp2), src2=vgpr(kw.vgprs.coutRowPtrBias), \
|
|
comment="new rowStart address += shift column * StridesW"))
|
|
module.add(VCndMaskB32(dst=vgpr(kw.vgprs.coutRowPtrBias), src0=vgpr(kw.vgprs.coutRowPtrBias), src1=vgpr(vTmp1), src2=sgpr(sTmp1,sgprCnt), \
|
|
comment="set new rowStart if meet conditions" ))
|
|
|
|
if kernel["StoreRemapVectorWidth"]:
|
|
ldsPad = max(kernel["StoreRemapVectorWidth"],kernel["MIOutputVectorWidth"])
|
|
module.add(VMovB32(dst=vgpr(vTmp1), src=hex((kernel["MacroTile0"]+ldsPad)*kw.states.bpeCexternal), \
|
|
comment="lds byte stride = (MT0 + PAD) * bpe"))
|
|
module.add(VMadI32I24(dst=vgpr(vTmp1), src0=vgpr(vTmp1), src1=vgpr(vTmp2), src2=vgpr(kw.vgprs.storeRemapLW), \
|
|
comment="new lds write address += shift column * Lds byte Stride"))
|
|
module.add(VCndMaskB32(dst=vgpr(kw.vgprs.storeRemapLW), src0=vgpr(kw.vgprs.storeRemapLW), src1=vgpr(vTmp1), \
|
|
src2=sgpr(sTmp1,sgprCnt), comment="set new rowStart if meet conditions" ))
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
printExit("Output E does not support StoreRemapVectorWidth")
|
|
if kw.vgprs.coutRowPtrBias != -1:
|
|
printExit("Bias reduction does not support StoreRemapVectorWidth")
|
|
module.addSpaceLine()
|
|
|
|
return module
|
|
|
|
def emitLdChange(self, kernel, ss, tc, edge, beta, mask, bufferOOB, singleUpdate, tmpVgpr, tmpSgpr, addrVgpr, BufAddr, dim):
|
|
"""
|
|
Generate code for final C read/D write address
|
|
"""
|
|
|
|
laneSGPRCount = self.kernelWriter.states.laneSGPRCount
|
|
module = Module("emitLdChange")
|
|
isSingleKernel = ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or kernel["StreamK"] > 0
|
|
if kernel["BufferStore"]:
|
|
module.add(self.emitScaleToBpe(kernel, ss, tmpVgpr, tmpSgpr, singleUpdate, tc, dim))
|
|
if edge and (not kernel["StoreRemapVectorWidth"] or (kernel["StoreRemapVectorWidth"] and (beta or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"))) and \
|
|
(tc != 'ScaleAlphaVec' and (not (tc == 'Bias' and self.kernelWriter.states.useBias == DataDirection.READ)) and tc != 'ScaleAVec' and tc != 'ScaleBVec'):
|
|
module.add(VCndMaskB32(dst=vgpr(addrVgpr), src0=vgpr(bufferOOB), src1=vgpr(addrVgpr), \
|
|
src2=sgpr(mask,laneSGPRCount), comment="LD%s clip if OOB. offset" % tc ))
|
|
# dot2: for every NumWaveSplitK contiguous threads, only the rightmost (with the largest threadIdx) one need to write value to D matrix.
|
|
if kernel["NumWaveSplitK"] > 1 and (tc == 'D' or tc == 'TD'):
|
|
kidx = bufferOOB-1
|
|
module.add(VAndB32(dst=vgpr(kidx), src0=(kernel["NumWaveSplitK"]-1), src1=vgpr("Serial"), comment="kidx = Serial % NumWaveSplitK"))
|
|
module.add(VCmpXNeU32(dst=EXEC(), src0=(kernel["NumWaveSplitK"]-1), src1=vgpr(kidx), comment="kidx == NumWaveSplitK-1" ))
|
|
module.add(VMovB32(dst=vgpr(addrVgpr), src=vgpr(bufferOOB), comment="LD%s clip if kidx != NumWaveSplitK-1." % tc))
|
|
module.add(SMovB64(dst=EXEC(), src=-1, comment="Reset exec mask"))
|
|
|
|
else:
|
|
if tc == 'Bias' and kernel["ProblemType"]["UseBias"] and isSingleKernel:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(VSubU32(dst=vgpr(self.addrBiasVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrBiasVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrBiasVgpr), \
|
|
comment="Bias address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrBiasVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrBiasVgpr), \
|
|
comment="add lds offset"))
|
|
elif tc == 'ScaleAlphaVec' and kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
if self.addrScaleAlphaVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile%u"%dim], src1=sgpr("WorkGroup%u"%dim), comment="wgp%u * MT%u"%(dim, dim)))
|
|
coordVgpr = self.coord0Vgpr if dim == 0 else self.coord1Vgpr
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAlphaVecVgpr), src0=vgpr(coordVgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="ScaleAlpha address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAlphaVecVgpr), \
|
|
comment="add lds offset"))
|
|
elif tc == 'ScaleAVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleAVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile0"], src1=sgpr("WorkGroup0"), comment="wgp0 * MT0"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleAVecVgpr), src0=vgpr(self.coord0Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="ScaleAVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleAVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleAVecVgpr), \
|
|
comment="add lds offset"))
|
|
elif tc == 'ScaleBVec' and (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.addrScaleBVecVgpr:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=kernel["MacroTile1"], src1=sgpr("WorkGroup1"), comment="wgp1 * MT1"))
|
|
module.add(VSubU32(dst=vgpr(self.addrScaleBVecVgpr), src0=vgpr(self.coord1Vgpr), src1=sgpr(tmpSgpr)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
shiftHex=hex(log2(self.kernelWriter.states.bpeCinternal)), \
|
|
src=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="ScaleBVec address scaled by BPE"))
|
|
if kernel["LdsOffsetBias"] != 0:
|
|
module.add(VAddU32(dst=vgpr(self.addrScaleBVecVgpr), \
|
|
src0=(kernel["LdsOffsetBias"]), \
|
|
src1=vgpr(self.addrScaleBVecVgpr), \
|
|
comment="add lds offset"))
|
|
else:
|
|
# store a copy of the offset in 2 of the tmpVgpr for D
|
|
module.add(VAddCOU32(dst=vgpr(addrVgpr+0), dst1=VCC(), src0=vgpr(BufAddr+0), src1=vgpr(tmpVgpr+2), \
|
|
comment="addrVgpr = C(D) + index*bytes (lo)" ))
|
|
module.add(VAddCCOU32(dst=vgpr(addrVgpr+1), dst1=VCC(), src0=vgpr(BufAddr+1), src1=vgpr(tmpVgpr+3), \
|
|
src2=VCC(), comment="addrVgpr = C(D) + index*bytes (hi)"))
|
|
return module
|
|
|
|
@staticmethod
|
|
def incrementSrdMultipleRows(srcDstBaseSgpr: str, strideSgpr: str, tmpSgpr: str, numRows: int, bpe: int) -> Module:
|
|
module = Module("incrementToNextRows")
|
|
if numRows > 1:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), \
|
|
src0=sgpr(strideSgpr), \
|
|
src1=numRows*bpe, \
|
|
comment="scale %s *= numRows(%u) * bpe"%(strideSgpr, numRows)))
|
|
elif numRows < 0:
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), \
|
|
src0=sgpr(strideSgpr), \
|
|
src1=(-numRows)*bpe, \
|
|
comment="scale %s *= numRows(%u) * bpe"%(strideSgpr, numRows)))
|
|
else:
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), \
|
|
src=sgpr(strideSgpr), \
|
|
shiftHex=log2(bpe), \
|
|
comment="incToNextRow: Scale by BPE"))
|
|
dstLow = f"{srcDstBaseSgpr}+0"
|
|
dstHigh = f"{srcDstBaseSgpr}+1"
|
|
|
|
if numRows >= 0:
|
|
module.add(SAddU32(dst=sgpr(dstLow), \
|
|
src0=sgpr(dstLow), \
|
|
src1=sgpr(tmpSgpr), \
|
|
comment="incToNextRow: gra SRD += inc(lower)" ))
|
|
module.add(SAddCU32(dst=sgpr(dstHigh), \
|
|
src0=sgpr(dstHigh), \
|
|
src1=0, \
|
|
comment="incToNextRow: gra SRD += inc(upper)" ))
|
|
else:
|
|
module.add(SSubU32(dst=sgpr(dstLow), \
|
|
src0=sgpr(dstLow), \
|
|
src1=sgpr(tmpSgpr), \
|
|
comment="incToNextRow: gra SRD -= inc(lower)" ))
|
|
module.add(SSubBU32(dst=sgpr(dstHigh), \
|
|
src0=sgpr(dstHigh), \
|
|
src1=0, \
|
|
comment="incToNextRow: gra SRD -= inc(upper)" ))
|
|
return module
|
|
|
|
def incrementToNextRow(self, kernel, tc, ss, stmp, forceinitrow0=0,
|
|
overrideAfterPrimerRows=0, bpeType=None, dst=-1):
|
|
"""
|
|
Generate code to move to the next row(s)
|
|
If optSrdIncForRow, this will move the SRD forward
|
|
If not, this could generate some other instructions
|
|
|
|
CompactLoopStore (CLS-loop mode) uses a DELAYED primer pattern: s_add
|
|
uses s[stmp] prepared by the PREVIOUS call, and the stride compute
|
|
(s_lshl/s_mul) at the END primes s[stmp] for the NEXT call. Every elt
|
|
must emit (including rowInc==0) to preserve the s[stmp] chain --
|
|
elt-0's s_add is a no-op (s[stmp]=0 from preamble init) + s_lshl that
|
|
primes s[stmp]=StrideCD1<<log2(bpe) for elt-1's real advance. Caller
|
|
passes `forceinitrow0=1` to allow this chain-seed emit when rowInc==0.
|
|
|
|
`overrideAfterPrimerRows` > 0 replaces the AFTER-primer's numRows with
|
|
the given value. The caller (GlobalWriteBatch) computes this as the
|
|
rowInc of the NEXT EMITTING elt (look-ahead), so the primer that
|
|
elt-N writes to s[stmp] is exactly the increment elt-(N+1) needs.
|
|
Without override, the primer would carry elt-N's own rowInc, which
|
|
elt-(N+1) would then erroneously consume -- the off-by-one root cause.
|
|
0 means no override (default delayed chain, primer = own numRows).
|
|
|
|
Legacy path (CompactLoopStore=False): emit stride compute BEFORE s_add
|
|
(each call self-contained, gated by rowInc != 0). `forceinitrow0=0`
|
|
and `overrideAfterPrimerRows=0` defaults keep legacy behaviour.
|
|
|
|
"""
|
|
|
|
module = Module("incrementToNextRow")
|
|
numRows = self.rowInc
|
|
tmpBpe = bpeType if bpeType else self.kernelWriter.states.bpeCexternal
|
|
if (tc == 'C' or tc == 'TD') and (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
tmpBpe = int(self.kernelWriter.states.bpr * kernel["ProblemType"]["DestDataType"].numRegisters())
|
|
if ss.optSrdIncForRow:
|
|
# CLS chain seed: when caller asks for `forceinitrow0=1` we must
|
|
# also emit on rowInc==0 to keep s[stmp] primed; otherwise legacy
|
|
# behaviour: skip when no advance is needed.
|
|
if numRows or forceinitrow0:
|
|
packedC1 = kernel["PackedC1IndicesX"]
|
|
assert(len(packedC1) == 1) # would need to extract each dim and scale
|
|
if tc == 'Bias' and (not kernel["WorkGroupReduction"]):
|
|
index = packedC1[0] - 1
|
|
strideCD1 = "Size%s" % "I" if index == 0 else ("J" if index == 1 else (self.kernelWriter.states.indexChars[index]))
|
|
elif tc == 'Gate':
|
|
# Gate uses its own GateStride+0 (gate's col stride) for row advance.
|
|
strideCD1 = "GateStride+0"
|
|
else:
|
|
td = "D" if tc == 'TD' else tc
|
|
strideCD1 = "Stride%s%s"%(td ,self.kernelWriter.states.indexChars[packedC1[0]])
|
|
|
|
# Build a strideCompute block for a given numRows. Used for
|
|
# BOTH the legacy "before s_add" emit (with the call's own
|
|
# numRows) and the CLS delayed "after s_add" primer.
|
|
def _buildStrideCompute(nr):
|
|
sc = Module("strideCompute")
|
|
if nr > 1:
|
|
sc.add(SMulI32(dst=sgpr(stmp), \
|
|
src0=sgpr(strideCD1), \
|
|
src1=nr*tmpBpe, \
|
|
comment="scale Stride%s *= numRows(%u) * bpe"%(tc,nr)))
|
|
elif nr < 0:
|
|
sc.add(SMulI32(dst=sgpr(stmp), \
|
|
src0=sgpr(strideCD1), \
|
|
src1=(-nr)*tmpBpe, \
|
|
comment="scale Stride%s *= numRows(%u) * bpe"%(tc,nr)))
|
|
else:
|
|
sc.add(SLShiftLeftB32(dst=sgpr(stmp), \
|
|
src=sgpr(strideCD1), \
|
|
shiftHex=log2(tmpBpe), \
|
|
comment="incToNextRow: Scale by BPE"))
|
|
return sc
|
|
|
|
# Legacy (non-CompactLoopStore): stride compute BEFORE s_add,
|
|
# using the call's own numRows.
|
|
if not kernel["CompactLoopStore"]:
|
|
module.add(_buildStrideCompute(numRows))
|
|
|
|
if dst == -1:
|
|
dstLow = "Srd%s+0"%(tc)
|
|
dstHigh = "Srd%s+1"%(tc)
|
|
else:
|
|
dstLow = dst+0
|
|
dstHigh = dst+1
|
|
|
|
if numRows >= 0:
|
|
module.add(SAddU32(dst=sgpr(dstLow), \
|
|
src0=sgpr(dstLow), \
|
|
src1=sgpr(stmp), \
|
|
comment="incToNextRow: gra SRD += inc(lower)" ))
|
|
module.add(SAddCU32(dst=sgpr(dstHigh), \
|
|
src0=sgpr(dstHigh), \
|
|
src1=0, \
|
|
comment="incToNextRow: gra SRD += inc(upper)" ))
|
|
else: # numRows < 0
|
|
module.add(SSubU32(dst=sgpr(dstLow), \
|
|
src0=sgpr(dstLow), \
|
|
src1=sgpr(stmp), \
|
|
comment="incToNextRow: gra SRD -= inc(lower)" ))
|
|
module.add(SSubBU32(dst=sgpr(dstHigh), \
|
|
src0=sgpr(dstHigh), \
|
|
src1=0, \
|
|
comment="incToNextRow: gra SRD -= inc(upper)" ))
|
|
|
|
# CompactLoopStore: stride compute AFTER s_add (primes s[stmp]
|
|
# for the NEXT call's s_add). When caller passes
|
|
# `overrideAfterPrimerRows > 0`, use it as the primer rows --
|
|
# this is the look-ahead value = NEXT EMITTING elt's rowInc,
|
|
# so elt-(N+1)'s s_add reads exactly its own advance. Without
|
|
# override, primer uses elt-N's own numRows (off-by-one for
|
|
# CLS, but kept as the documented fallback).
|
|
if kernel["CompactLoopStore"]:
|
|
primerRows = overrideAfterPrimerRows if overrideAfterPrimerRows else numRows
|
|
module.add(_buildStrideCompute(primerRows))
|
|
|
|
return module
|