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
263 lines
15 KiB
Python
263 lines
15 KiB
Python
from ..Component import GL2Prefetch
|
|
from ..Common import INDEX_CHARS
|
|
from typing import Mapping
|
|
from rocisa.code import Module
|
|
from rocisa.instruction import SMulI32, SAddU64, VMovB32, VAddU32, VAddCOU32, \
|
|
VAddCCOU32, VAddNCU64, VLShiftRightB32, VMulLOU32, VMulHIU32, GlobalPrefetchB8, \
|
|
VCmpGtU32, VCndMaskB32, SSubI32, SMovB32, SAddU32, SAddCU32
|
|
from rocisa.container import sgpr, vgpr, RegisterContainer, VCC, GLOBALModifiers, ContinuousRegister
|
|
from rocisa.functions import vectorMultiply64Bpe, scalarMultiplyBpe, vectorStaticDivideAndRemainder, \
|
|
scalarStaticRemainder
|
|
from rocisa.enum import TemporalHint, CacheScope
|
|
from math import log2, ceil
|
|
|
|
class GL2PrefetchLoad(GL2Prefetch):
|
|
asmCaps = {"HasGlobalPrefetch": True}
|
|
globalModifiers = GLOBALModifiers(th=TemporalHint.TH_NT, scope=CacheScope.SCOPE_SE)
|
|
|
|
def __call__(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping):
|
|
pass
|
|
|
|
def init(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping):
|
|
globalPrefetchSize: int = writer.states.regCaps["GlobalPrefetchSize"]
|
|
tc: str = tp["tensorChar"]
|
|
isMX: bool = tc.startswith("MX")
|
|
isM: bool = tp.get("isM", False)
|
|
# Cooperative prefetch spans the *whole* cluster: every workgroup in the
|
|
# cluster contributes threads, and together they cover all the distinct
|
|
# macro-tiles the cluster consumes rather than only the single tile one
|
|
# workgroup uses for its own computation. Along the MT-selector axis
|
|
# (WorkGroup0 for A, WorkGroup1 for B) the cluster spans numTileWGs
|
|
# contiguous macro-tiles, so the tile dimension of the prefetched block
|
|
# is scaled accordingly.
|
|
numCooperativeWGs: int = kernel["ClusterDim"][0] * kernel["ClusterDim"][1]
|
|
numCooperativeThreads: int = numCooperativeWGs * kernel["NumThreads"]
|
|
|
|
subTc: str = tc if isM else tc[-1]
|
|
mt: int = kernel["MacroTile%s" % subTc]
|
|
numTileWGs: int = kernel["ClusterDim"][tp["idx"]] if isM else (kernel["ClusterDim"][0] if subTc == "A" else kernel["ClusterDim"][1])
|
|
bpe: float = tp["bpeGR"]
|
|
|
|
if isMX:
|
|
coalescedDim = mt * numTileWGs * kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"]
|
|
perpendicularDim = kernel["DepthU"] // kernel["MatrixInstK"]
|
|
else:
|
|
du: int = kernel["_DepthU%s" % subTc]
|
|
coalescedDim, perpendicularDim = (mt * numTileWGs, du) if tp["tlu"] else (du, mt * numTileWGs)
|
|
|
|
tp["gl2ncp"] = perpendicularDim
|
|
tp["gl2ncc"] = max(1, round(coalescedDim * bpe) // globalPrefetchSize)
|
|
tp["gl2nc"] = tp["gl2ncp"] * tp["gl2ncc"]
|
|
tp["gl2nl"] = max(1, ceil(tp["gl2nc"] / numCooperativeThreads))
|
|
|
|
def setIncrement(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping) -> Module:
|
|
mod = Module()
|
|
tc: str = tp["tensorChar"]
|
|
tIdx: int = tp['idx']
|
|
isM: bool = tp.get("isM", False)
|
|
subTc: str = tc if isM else tc[-1]
|
|
bpe: float = tp["bpeGR"]
|
|
du: int = kernel["_DepthU%s" % subTc]
|
|
if tc.startswith("MX"):
|
|
mod.add(SMulI32(sgpr(f"GL2PrefetchInc{tc}"), sgpr("Size%s"%INDEX_CHARS[tIdx]), \
|
|
round(kernel["DepthU"] // kernel["ProblemType"][f"MXBlock{subTc}"] * bpe), comment="addr increment"))
|
|
elif tp["tlu"]:
|
|
perpStride: str | RegisterContainer = writer.strideRef(subTc, 3)
|
|
mod.add(SMulI32(sgpr(f"GL2PrefetchInc{tc}"), perpStride, round(du * bpe), comment="addr increment"))
|
|
else:
|
|
mod.add(SMovB32(dst=sgpr(f"GL2PrefetchInc{tc}"), src=round(du * bpe), comment="addr increment"))
|
|
return mod
|
|
|
|
def calculateStartAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping) -> Module:
|
|
mod = Module()
|
|
globalPrefetchSize: int = writer.states.regCaps["GlobalPrefetchSize"]
|
|
tc: str = tp["tensorChar"]
|
|
tIdx: int = tp['idx']
|
|
tlu: bool = tp["tlu"]
|
|
isMX: bool = tc.startswith("MX")
|
|
isM: bool = tp.get("isM", False)
|
|
subTc: str = tc if isM else tc[-1]
|
|
mt: int = kernel["MacroTile%s" % subTc]
|
|
bpe: float = tp["bpeGR"]
|
|
tileStride: str | RegisterContainer = writer.strideRef(subTc, tIdx)
|
|
unrollStride: str | RegisterContainer = writer.strideRef(subTc, 3)
|
|
perpStride: str | RegisterContainer = unrollStride if tlu else tileStride
|
|
# WorkGroup{tIdx} selects the macro-tile; the other cluster axis is the
|
|
# cooperative sharing axis. The whole cluster cooperates on the prefetch.
|
|
sgprTileWgName: str = f"WorkGroup{tIdx}"
|
|
sgprShareWgName: str = f"WorkGroup{1 - tIdx}"
|
|
sgprSizeFreeName: str = f"Size{INDEX_CHARS[tIdx]}"
|
|
numThreads: int = kernel["NumThreads"]
|
|
vgprAddrBaseName: str = f"GL2PrefetchAddr{tc}"
|
|
vgprAddrName0: str = f"{vgprAddrBaseName}_0"
|
|
numTileWGs: int = kernel["ClusterDim"][tIdx]
|
|
numShareWGs: int = kernel["ClusterDim"][1 - tIdx]
|
|
numCooperativeWGs: int = numTileWGs * numShareWGs
|
|
numCooperativeThreads: int = numCooperativeWGs * numThreads
|
|
ncc: int = tp["gl2ncc"]
|
|
nc: int = tp["gl2nc"]
|
|
nl: int = tp["gl2nl"]
|
|
ncPerInst: int = ceil(nc / tp["gl2nl"])
|
|
inactiveShiftBits: int = int(log2(numCooperativeThreads // ncPerInst))
|
|
numTmpSgpr = 4
|
|
tmpVgprIdx = writer.vgprPool.checkOutAligned(2, 2)
|
|
tmpVgprCoalIdx = writer.vgprPool.checkOutAligned(1, 1)
|
|
if isMX:
|
|
mxUnit: int = kernel["MatrixInstK"] // kernel["ProblemType"][f"MXBlock{subTc}"]
|
|
|
|
mod.addComment(f"gl2 prefetch calc start addr of {tc}")
|
|
with writer.allocTmpSgpr(numTmpSgpr, 2) as tmpSgprRes:
|
|
tmpSgprIdx0 = tmpSgprRes.idx
|
|
tmpSgprIdx1 = tmpSgprRes.idx + 1
|
|
tmpSgprIdx2 = tmpSgprRes.idx + 2
|
|
tmpSgprIdx3 = tmpSgprRes.idx + 3
|
|
# Cooperative thread index over the whole cluster. Flatten this
|
|
# workgroup's cluster-local (tile, share) position into a single index
|
|
# and offset the wave's Serial by it, so the cluster's threads jointly
|
|
# enumerate all cooperative chunks. tmpSgprIdx3 keeps the cluster-local
|
|
# tile index; the cluster's base macro-tile (WorkGroup{tIdx} minus it)
|
|
# is recovered from it for the MT offset below.
|
|
mod.add(scalarStaticRemainder(tmpSgprIdx0, tmpSgprIdx3, sgprTileWgName, numTileWGs, \
|
|
tmpSgprRes, comment="cluster-local tile idx"))
|
|
mod.add(scalarStaticRemainder(tmpSgprIdx0, tmpSgprIdx0, sgprShareWgName, numShareWGs, \
|
|
tmpSgprRes, comment="cluster-local share idx"))
|
|
mod.add(SMulI32(sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx3), numShareWGs, \
|
|
comment="tile idx * shareWGs"))
|
|
mod.add(SAddU32(sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx0), \
|
|
comment="flattened cluster WG idx"))
|
|
mod.add(SMulI32(sgpr(tmpSgprIdx0), sgpr(tmpSgprIdx1), numThreads, \
|
|
comment="cluster WG idx * numThreads"))
|
|
mod.add(VAddU32(vgpr(vgprAddrName0), vgpr("Serial"), sgpr(tmpSgprIdx0), \
|
|
comment="cooperative thread idx"))
|
|
if inactiveShiftBits > 0:
|
|
assert nl == 1, "Should only have one inst if inactiveShiftBits > 0"
|
|
mod.add(VLShiftRightB32(vgpr(vgprAddrName0), inactiveShiftBits, vgpr(vgprAddrName0), \
|
|
comment="shift inactive index"))
|
|
else:
|
|
for i in range(1, nl):
|
|
src = f"{vgprAddrBaseName}_{i-1}"
|
|
dst = f"{vgprAddrBaseName}_{i}"
|
|
mod.add(VAddU32(vgpr(dst), vgpr(src), ncPerInst, comment="inst index"))
|
|
# the last inst may contain overflow address, we need to mask it
|
|
vgprAddrNameLast = f"{vgprAddrBaseName}_{(nl-1)}"
|
|
mod.add(VCmpGtU32(VCC(), vgpr(vgprAddrNameLast), nc-1, comment="overflow number of needed cachelines?"))
|
|
mod.add(VCndMaskB32(vgpr(vgprAddrNameLast), vgpr(vgprAddrNameLast), nc-1, VCC()))
|
|
|
|
# MT offset & edge limit (in units of elements). The offset is the
|
|
# cluster's base macro-tile (WorkGroup{tIdx} floored to the cluster,
|
|
# i.e. minus the cluster-local tile idx kept in tmpSgprIdx3), since the
|
|
# cooperative block now spans all numTileWGs tiles the cluster covers.
|
|
mod.add(SSubI32(sgpr(tmpSgprIdx0), sgpr(sgprTileWgName), sgpr(tmpSgprIdx3), \
|
|
comment="cluster base tile"))
|
|
if isMX:
|
|
mod.add(SMulI32(sgpr(tmpSgprIdx0), sgpr(tmpSgprIdx0), mxUnit * mt, \
|
|
comment=f"clusterBaseTile * mxUnit({mxUnit}) * MT({mt})"))
|
|
mod.add(SSubI32(sgpr(tmpSgprIdx1), sgpr(sgprSizeFreeName), 1))
|
|
mod.add(SMulI32(sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx1), mxUnit))
|
|
mod.add(SSubI32(sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx0), comment="max offset inside cluster tiles"))
|
|
else:
|
|
mod.add(SMulI32(sgpr(tmpSgprIdx0), sgpr(tmpSgprIdx0), mt, comment=f"clusterBaseTile * MT({mt})"))
|
|
mod.add(SSubI32(sgpr(tmpSgprIdx1), sgpr(sgprSizeFreeName), 1))
|
|
mod.add(SSubI32(sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx0), comment="max offset inside cluster tiles"))
|
|
|
|
# will we have MX stride later?
|
|
if isMX:
|
|
perpStride = sgpr(tmpSgprIdx2)
|
|
mod.add(SMulI32(perpStride, sgpr(sgprSizeFreeName), mxUnit, f"MX perp stride"))
|
|
for i in range(nl):
|
|
vgprAddrName = f"{vgprAddrBaseName}_{i}"
|
|
vgprAddrNameHi = vgprAddrName + "+1"
|
|
if ncc > 1:
|
|
mod.add(VMovB32(vgpr(tmpVgprCoalIdx), vgpr(vgprAddrName)))
|
|
mod.add(vectorStaticDivideAndRemainder(vgprAddrName, tmpVgprCoalIdx, tmpVgprCoalIdx, \
|
|
ncc, ContinuousRegister(tmpVgprIdx, 2), comment="coal/perp index calc"))
|
|
mod.add(VMulLOU32(vgpr(tmpVgprCoalIdx), vgpr(tmpVgprCoalIdx), round(globalPrefetchSize / bpe), \
|
|
comment="coal * globalPrefetchSize / bpe"))
|
|
else:
|
|
mod.add(VMovB32(vgpr(tmpVgprCoalIdx), 0, comment="coalesced index"))
|
|
|
|
# edge protection
|
|
if isMX or tlu:
|
|
mod.add(VCmpGtU32(VCC(), vgpr(tmpVgprCoalIdx), sgpr(tmpSgprIdx1), comment="> edge limit?"))
|
|
mod.add(VCndMaskB32(vgpr(tmpVgprCoalIdx), vgpr(tmpVgprCoalIdx), sgpr(tmpSgprIdx1), VCC()))
|
|
else:
|
|
mod.add(VCmpGtU32(VCC(), vgpr(vgprAddrName), sgpr(tmpSgprIdx1), comment="> edge limit?"))
|
|
mod.add(VCndMaskB32(vgpr(vgprAddrName), vgpr(vgprAddrName), sgpr(tmpSgprIdx1), VCC()))
|
|
# perp stride
|
|
mod.add(VMulHIU32(vgpr(vgprAddrNameHi), vgpr(vgprAddrName), perpStride, comment="perp *= stride"))
|
|
mod.add(VMulLOU32(vgpr(vgprAddrName), vgpr(vgprAddrName), perpStride))
|
|
# coal + perp
|
|
mod.add(VAddCOU32(vgpr(vgprAddrName), VCC(), vgpr(vgprAddrName), vgpr(tmpVgprCoalIdx), comment="coal + perp"))
|
|
mod.add(VAddCCOU32(vgpr(vgprAddrNameHi), VCC(), vgpr(vgprAddrNameHi), 0, VCC()))
|
|
mod.add(vectorMultiply64Bpe(vgprAddrName, vgprAddrName, bpe, tmpVgprIdx, comment="scale by bpe"))
|
|
|
|
# base address + MT offset (in units of bytes)
|
|
mod.add(scalarMultiplyBpe(tmpSgprIdx0, tmpSgprIdx0, bpe))
|
|
if isMX or tlu:
|
|
mod.add(SAddU32(sgpr(tmpSgprIdx0), sgpr("Address%s"%tc), sgpr(tmpSgprIdx0), comment="base address + MT offset"))
|
|
mod.add(SAddCU32(sgpr(tmpSgprIdx1), sgpr("Address%s+1"%tc), 0))
|
|
else:
|
|
mod.addModuleAsFlatItems(writer.s_mul_u64_u32(
|
|
sgpr(tmpSgprIdx0), sgpr(tmpSgprIdx1),
|
|
sgpr(tmpSgprIdx0), perpStride,
|
|
tmpVgprIdx, comment="*= stride"))
|
|
mod.add(SAddU64(sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx0, 2), sgpr("Address%s"%tc, 2), comment="base address + MT offset"))
|
|
|
|
# strided batch offset
|
|
if kernel["ProblemType"]["Batched"]:
|
|
assert kernel["ProblemType"]["StridedBatched"], "Currently GL2Prefetch does not support general batch"
|
|
for batchIdx in kernel["ProblemType"]["IndicesBatch"]:
|
|
# packed index check
|
|
if batchIdx in kernel["ProblemType"]["IndicesFree"] or batchIdx not in tp['ia']:
|
|
continue
|
|
assert(batchIdx==2) # can only have one wg2 with a batch. Other dimensions should be packed into wg0/wg1
|
|
batchStrideName = "Stride%s%s"%(tc, writer.states.indexChars[batchIdx])
|
|
mod.add(scalarMultiplyBpe(tmpSgprIdx2, batchStrideName, bpe, comment="batchStride * bpe"))
|
|
mod.addModuleAsFlatItems(writer.s_mul_u64_u32(
|
|
sgpr(tmpSgprIdx2), sgpr(tmpSgprIdx3),
|
|
sgpr("WorkGroup2"), sgpr(tmpSgprIdx2),
|
|
tmpVgprIdx, comment="batch offset * wg2"))
|
|
mod.add(SAddU64(sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx2, 2)))
|
|
# skip PGR loads (uses GSU-adjusted increment)
|
|
if kernel["PrefetchGlobalRead"] > 0:
|
|
if kernel["PrefetchGlobalRead"] > 1:
|
|
mod.addModuleAsFlatItems(writer.s_mul_u64_u32(
|
|
sgpr(tmpSgprIdx2), sgpr(tmpSgprIdx3),
|
|
sgpr(f"GL2PrefetchInc{tc}"), kernel["PrefetchGlobalRead"],
|
|
tmpVgprIdx, comment="*= PGR"))
|
|
mod.add(SAddU64(sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx0, 2), sgpr(tmpSgprIdx2, 2), \
|
|
comment="skip PGR loads"))
|
|
else:
|
|
mod.add(SAddU32(sgpr(tmpSgprIdx0), sgpr(tmpSgprIdx0), sgpr(f"GL2PrefetchInc{tc}"), \
|
|
comment="skip PGR loads"))
|
|
mod.add(SAddCU32(sgpr(tmpSgprIdx1), sgpr(tmpSgprIdx1), 0, \
|
|
comment="skip PGR loads"))
|
|
|
|
# add all together
|
|
for i in range(tp["gl2nl"]):
|
|
dst = f"{vgprAddrBaseName}_{i}"
|
|
mod.add(VAddNCU64(vgpr(dst, 2), vgpr(dst, 2), sgpr(tmpSgprIdx0, 2)))
|
|
|
|
writer.vgprPool.checkIn(tmpVgprIdx)
|
|
writer.vgprPool.checkIn(tmpVgprCoalIdx)
|
|
return mod
|
|
|
|
def issueLoad(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping) -> Module:
|
|
mod = Module()
|
|
tc: str = tp["tensorChar"]
|
|
for i in range(tp["gl2nl"]):
|
|
addrName = f"GL2PrefetchAddr{tc}_{i}"
|
|
mod.add(GlobalPrefetchB8(vgpr(addrName, 2), sgpr("off", isOff=True), self.globalModifiers))
|
|
return mod
|
|
|
|
def incrementAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp: Mapping) -> Module:
|
|
mod = Module()
|
|
tc: str = tp["tensorChar"]
|
|
inc = sgpr(f"GL2PrefetchInc{tc}")
|
|
for i in range(tp["gl2nl"]):
|
|
addrName = f"GL2PrefetchAddr{tc}_{i}"
|
|
addrNameHi = addrName + "+1"
|
|
mod.add(VAddCOU32(vgpr(addrName), VCC(), vgpr(addrName), inc))
|
|
mod.add(VAddCCOU32(vgpr(addrNameHi), VCC(), vgpr(addrNameHi), 0, VCC()))
|
|
|
|
return mod |