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
4318 lines
260 KiB
Python
4318 lines
260 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.enum import CacheScope
|
|
from rocisa.code import Module, Label
|
|
from rocisa.container import vgpr, sgpr, mgpr, SMEMModifiers, MUBUFModifiers, GLOBALModifiers, replaceHolder, EXEC,\
|
|
VOP3PModifiers, ContinuousRegister, DSModifiers
|
|
from rocisa.instruction import GlobalInv, GlobalWb, SAddCU32, SAddU32, SAndB32, SBarrier, \
|
|
SBranch, SCBranchSCC0, SCBranchSCC1, SCMovB32, SCSelectB32, SCmpEQU32, SCmpEQU64, \
|
|
SCmpGeU32, SCmpGtU32, SCmpLeU32, SCmpLtU32, SLShiftLeftB32, SLShiftLeftB64, SLShiftRightB32, VLShiftLeftB32, SLoadB32, \
|
|
SMaxI32, SMinU32, SMovB32, SMovB64, SMulI32, SNop, SOrB32, SSleep, SStoreB32, SSubU32, \
|
|
SWaitCnt, SWaitXCnt, VAddF32, VAddF64, VAddPKF16, VAddU32, VSubU32, VLShiftRightB32, VMovB32, \
|
|
VReadfirstlaneB32, VCmpXEqU32, VCvtBF16toFP32, GlobalAtomicIncU32Saddr, BufferLoadB32, BufferStoreB32, \
|
|
SAtomicInc, DSLoadB32, DSStoreB32, SLongBranch, SLongBranchPositive
|
|
from rocisa.functions import scalarStaticDivideAndRemainder, sMagicDiv2, \
|
|
vectorStaticMultiply, BranchIfNotZero, scalarUInt24DivideAndRemainder, scalarUInt32DivideAndRemainder
|
|
|
|
from .Subtile.SubtileLREmit import localReadResetOffsetsSubtile
|
|
|
|
from ..Common import print2, ceilDivide, log2, clusterEnabled
|
|
from ..Component import Component
|
|
from ..AsmStoreState import StoreState, VectorDataTypes
|
|
from ..AsmAddressCalculation import AddrCalculation
|
|
import abc
|
|
|
|
from copy import deepcopy
|
|
|
|
class XCCMapping(Component):
|
|
"""
|
|
XCC mapping code.
|
|
"""
|
|
|
|
class XCCMappingOff(XCCMapping):
|
|
kernel = {"StreamKXCCMapping": 0}
|
|
|
|
def __call__(self, writer, kernel):
|
|
module = Module("XCCMapping Off")
|
|
return module
|
|
|
|
class XCCMappingOn(XCCMapping):
|
|
|
|
@classmethod
|
|
def matches(cls, writer, debug=False):
|
|
return writer.states.kernel["StreamKXCCMapping"] > 0
|
|
|
|
def __call__(self, writer, kernel):
|
|
module = Module("XCCMapping On")
|
|
|
|
with writer.allocTmpSgpr(4, tag="StreamKXCCMappingOn_tmpSgprRes") as tmpSgprRes:
|
|
sXCC = tmpSgprRes.idx
|
|
sGridC = tmpSgprRes.idx + 1
|
|
sGridF = tmpSgprRes.idx + 2
|
|
sGridM = tmpSgprRes.idx + 3
|
|
sTmp = None
|
|
sTmpRes = None
|
|
sqTmp = writer.sgprPool.checkOut(1, "sqTmp")
|
|
divisor = kernel["StreamKXCCMapping"]
|
|
if ((divisor & (divisor - 1)) != 0): # Need temp registers if not power of 2
|
|
sTmp = writer.sgprPool.checkOutAligned(2, 2, "sTmp", preventOverflow=not kernel.get("UseSubtileImpl", False))
|
|
sTmpRes = ContinuousRegister(idx=sTmp, size=2)
|
|
|
|
# sGridC = ceil(grid / xccm)
|
|
sGrid = writer.acquireStreamKConstSgpr(kernel, "skGrid")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sGrid), src=vgpr(writer.states.skConstVgprs["skGrid"])))
|
|
module.add(SAddU32(dst=sgpr(sGridC), src0=sgpr(sGrid), src1=hex(kernel["StreamKXCCMapping"] - 1), comment="ceil(grid/xccm)"))
|
|
module.add(scalarStaticDivideAndRemainder(qReg=sGridC, rReg=-1, dReg=sGridC, divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes, doRemainder=0))
|
|
# sGridF = floor(grid / xccm)
|
|
# sGridM = grid % xccm
|
|
module.add(scalarStaticDivideAndRemainder(qReg=sGridF, rReg=sGridM, dReg=sGrid, divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes))
|
|
writer.releaseStreamKConstSgpr(sGrid)
|
|
# sXCC = wg0 % xccm
|
|
# sqtmp is temp register for quotient for non-power-of-2 case
|
|
# sqtmp overlaps temp registers, works in this case and output is discarded
|
|
module.add(scalarStaticDivideAndRemainder(qReg=sqTmp, rReg=sXCC, dReg="WorkGroup0", divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes, doRemainder=2))
|
|
# Check if current XCC requires a remainder WG or not
|
|
module.add(SCmpLtU32(src0=sgpr(sXCC), src1=sgpr(sGridM), comment="XCCM < Remainder"))
|
|
module.add(SCSelectB32(dst=sgpr(sGridC), src0=sgpr(sGridC), src1=sgpr(sGridF), comment="Select multiplier"))
|
|
module.add(SCSelectB32(dst=sgpr(sGridM), src0=0, src1=sgpr(sGridM), comment="Select remainder"))
|
|
# WG = floor(wg0 / xccm) * xccm + XCCoffset + optional remainder
|
|
module.add(scalarStaticDivideAndRemainder(qReg="WorkGroup0", rReg=-1, dReg="WorkGroup0", divisor=kernel["StreamKXCCMapping"], tmpSgprRes=sTmpRes, doRemainder=0))
|
|
module.add(SMulI32(dst=sgpr(sXCC), src0=sgpr(sXCC), src1=sgpr(sGridC), comment="XCC group id"))
|
|
module.add(SAddU32(dst=sgpr("WorkGroup0"), src0=sgpr("WorkGroup0"), src1=sgpr(sXCC), comment="Add XCC group offset"))
|
|
module.add(SAddU32(dst=sgpr("WorkGroup0"), src0=sgpr("WorkGroup0"), src1=sgpr(sGridM), comment="Add remainder offset"))
|
|
|
|
writer.sgprPool.checkIn(sqTmp)
|
|
if sTmp is not None:
|
|
writer.sgprPool.checkIn(sTmp)
|
|
|
|
return module
|
|
|
|
|
|
class StreamKMemoryOrdering(Component):
|
|
"""
|
|
Memory-ordering fences and flag accessors for the StreamK partial-tile
|
|
handshake.
|
|
|
|
StreamK uses a producer/consumer protocol: one workgroup writes a partial
|
|
tile to a workspace and signals completion via a flag, and other
|
|
workgroups poll the flag and read the partials. The required cross-CU
|
|
memory ordering depends on the target ISA:
|
|
|
|
- Most arches: `s_waitcnt vscnt(0)` before the flag store and an SMEM
|
|
flag load with `glc/dlc/scope:SCOPE_DEV` are sufficient because
|
|
ordering between L1/L2 and the device-scope coherence point is
|
|
implicit.
|
|
|
|
- gfx1250: the L2 has independent partitions and SMEM is not coherent
|
|
with the VMEM flag store, so an explicit `global_wb scope:SCOPE_DEV`
|
|
is required on the release side and a `global_inv scope:SCOPE_DEV`
|
|
on the acquire side. Additionally, XNACK-replay can reorder a
|
|
volatile/atomic VMEM op past in-flight VMEM, so `s_wait_xcnt 0` must
|
|
precede such ops. The flag itself must be read via VMEM (not SMEM)
|
|
to observe the producer's release-side fence.
|
|
|
|
Selection is driven by the `HasInvWbDevFences` arch capability. The
|
|
XNACK-replay drain in `preVolatileVmem` is gated separately on
|
|
`RequiresXCntForVolatileVMEM` and lives on the abstract base so a
|
|
future arch needing only one of the two can be supported by adding a
|
|
single capability flag.
|
|
"""
|
|
def __call__(self):
|
|
assert(0)
|
|
|
|
def preVolatileVmem(self, writer, comment="") -> Module:
|
|
"""Drain in-flight VMEM (XNACK-replay) before a volatile/atomic VMEM op.
|
|
|
|
Required on arches with `RequiresXCntForVolatileVMEM`. No-op
|
|
elsewhere.
|
|
"""
|
|
module = Module("StreamK pre-volatile VMEM drain")
|
|
if writer.states.archCaps["RequiresXCntForVolatileVMEM"]:
|
|
module.add(SWaitXCnt(xcnt=0, comment=comment))
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def releaseFence(self, writer) -> Module:
|
|
"""Memory fence ordering prior partial-tile stores before the flag store."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def acquireFence(self, writer) -> Module:
|
|
"""Memory fence after observing the flag and before reading partials."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def readFlag(self, writer, dst, soffset) -> Module:
|
|
"""Read the StreamK completion flag into SGPR `dst` for compare."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def flagBufferMubuf(self) -> MUBUFModifiers:
|
|
"""MUBUF modifiers for buffer load/store of the flag word."""
|
|
pass
|
|
|
|
|
|
class StreamKMemoryOrderingDefault(StreamKMemoryOrdering):
|
|
"""No-op cross-CU fences; SMEM flag with glc/dlc/SCOPE_DEV.
|
|
|
|
Used on every arch that does not require explicit cross-L2 fences.
|
|
"""
|
|
archCaps = {"HasInvWbDevFences": False}
|
|
|
|
def releaseFence(self, writer) -> Module:
|
|
module = Module("StreamK release fence (default)")
|
|
module.add(SWaitCnt(vscnt=0, comment="wait for data store"))
|
|
return module
|
|
|
|
def acquireFence(self, writer) -> Module:
|
|
return Module("StreamK acquire fence (default, no-op)")
|
|
|
|
def readFlag(self, writer, dst, soffset) -> Module:
|
|
module = Module("StreamK read flag (SMEM)")
|
|
module.add(SLoadB32(dst=sgpr(dst), base=sgpr("AddressFlags", 2),
|
|
soffset=soffset,
|
|
smem=SMEMModifiers(glc=True, dlc=True,
|
|
scope=CacheScope.SCOPE_DEV),
|
|
comment="get flag"))
|
|
module.add(SWaitCnt(kmcnt=0, comment="wait for flag load"))
|
|
return module
|
|
|
|
def flagBufferMubuf(self) -> MUBUFModifiers:
|
|
return MUBUFModifiers(offen=True, glc=True, dlc=True,
|
|
scope=CacheScope.SCOPE_DEV)
|
|
|
|
|
|
class StreamKMemoryOrderingDevScopeFences(StreamKMemoryOrdering):
|
|
"""Explicit cross-L2 release/acquire fences via global_wb/global_inv
|
|
scope:SCOPE_DEV plus a VMEM-coherent flag read.
|
|
|
|
Selected on arches whose L2 is partitioned across CUs/XCDs and whose
|
|
SMEM is not coherent with the VMEM flag write (e.g. gfx1250).
|
|
"""
|
|
archCaps = {"HasInvWbDevFences": True}
|
|
|
|
def releaseFence(self, writer) -> Module:
|
|
module = Module("StreamK release fence (dev-scope)")
|
|
module.add(SWaitCnt(vlcnt=0,
|
|
comment="release: drain in-flight loads before global_wb"))
|
|
module.add(SWaitCnt(vscnt=0, comment="wait for data store"))
|
|
module.add(GlobalWb(scope=CacheScope.SCOPE_DEV,
|
|
comment="release: writeback partials to L2-coherent point"))
|
|
module.add(SWaitCnt(vlcnt=0, vscnt=0,
|
|
comment="release: wait for global_wb"))
|
|
return module
|
|
|
|
def acquireFence(self, writer) -> Module:
|
|
# Drop stale dev-scope cache lines so the next dependent read (the flag
|
|
# word in getFlagValue, or the partials after the flag is observed) is
|
|
# re-fetched from the L2-coherent point.
|
|
module = Module("StreamK acquire fence (dev-scope)")
|
|
module.add(GlobalInv(scope=CacheScope.SCOPE_DEV,
|
|
comment="acquire: invalidate before dependent dev-scope read"))
|
|
module.add(SWaitCnt(vlcnt=0, comment="acquire: wait for global_inv"))
|
|
return module
|
|
|
|
def readFlag(self, writer, dst, soffset) -> Module:
|
|
streamk = Component.StreamK.find(writer)
|
|
module = Module("StreamK read flag (VMEM)")
|
|
flagVgpr = writer.vgprPool.checkOut(1, "flagAcq")
|
|
module.add(streamk.getFlagValue(writer, dst=vgpr(flagVgpr),
|
|
soffset=soffset, comment="acquire: get flag (VMEM)"))
|
|
module.add(SWaitCnt(vlcnt=0, comment="acquire: wait VMEM flag load"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(dst), src=vgpr(flagVgpr),
|
|
comment="move VMEM flag to SGPR for compare"))
|
|
writer.vgprPool.checkIn(flagVgpr)
|
|
return module
|
|
|
|
def flagBufferMubuf(self) -> MUBUFModifiers:
|
|
return MUBUFModifiers(offen=True, scope=CacheScope.SCOPE_DEV)
|
|
|
|
|
|
class StreamK(Component):
|
|
"""
|
|
StreamK code.
|
|
"""
|
|
# --- Variant feature flags. Each flag is snapshotted onto
|
|
# writer.states.streamK by KernelWriter._initKernel and queried by
|
|
# call sites that previously branched on the integer value of
|
|
# kernel["StreamK"]. Defaults are False; each concrete StreamK*
|
|
# subclass overrides the flags it sets.
|
|
#
|
|
# emitsParallelReductionSgprAliases: emit the SkSplit/skTiles +
|
|
# SkPartialIdx/Beta SGPR aliases pre-epilogue
|
|
# borrowsSrdWsInEpilogue: epilogue may borrow the SrdWS SGPR pool
|
|
# emitsWorkspaceReductionBpe: epilogue allocates dtype-aware Log2Bpe
|
|
# SGPRs for the workspace reduction
|
|
# requiresWorkspaceReductionStorePath: the global-write elements
|
|
# emit must emit the workspace-reduction store branch
|
|
# (disables noGSUBranch fast path)
|
|
# keepsConstantsInSgpr: the dynamic per-XCD path references SK
|
|
# kernarg constants directly, so they cannot be cached in VGPRs
|
|
# on gfx1250
|
|
# supportsSubtileImpl: variant is accepted by UseSubtileImpl=1
|
|
emitsParallelReductionSgprAliases: bool = False
|
|
borrowsSrdWsInEpilogue: bool = False
|
|
emitsWorkspaceReductionBpe: bool = False
|
|
requiresWorkspaceReductionStorePath: bool = False
|
|
keepsConstantsInSgpr: bool = False
|
|
supportsSubtileImpl: bool = False
|
|
|
|
def __call__(self):
|
|
assert(0)
|
|
|
|
@staticmethod
|
|
def _depthUForTc(kernel, tc):
|
|
"""Return the per-StreamK-iteration K-stride (element count) for a tensor.
|
|
|
|
StreamK counts iterations in full DepthU units, so non-sparse data
|
|
tensors always use DepthU even in multi-DU mode (where _DepthU{A,B} is
|
|
the smaller per-uid swizzle sub-stride, not a compression).
|
|
|
|
For MXSA/MXSB (MX swizzled/pre-shuffle case), the swizzled block size
|
|
is 32 * 256 so an additional *32 multiplier is needed.
|
|
|
|
For Sparse problems the compressed data operand and the Metadata
|
|
tensor genuinely hold fewer elements per DepthU of computation, so
|
|
they advance by their per-tensor _DepthU{A,B,Metadata} stride (the
|
|
develop behavior); using full DepthU there would over-advance the SRD.
|
|
"""
|
|
if tc in ("MXSA", "MXSB"):
|
|
key = "_DepthU%s" % tc
|
|
if key in kernel:
|
|
_DepthU = kernel[key]
|
|
if kernel.get("UseSubtileImpl"):
|
|
_DepthU = (_DepthU * 32)
|
|
return _DepthU
|
|
return kernel["DepthU"]
|
|
if kernel["ProblemType"]["Sparse"]:
|
|
key = "_DepthU%s" % tc
|
|
if key in kernel:
|
|
return kernel[key]
|
|
return kernel["DepthU"]
|
|
|
|
def shiftSrd(self, writer, srdIdx) -> Module:
|
|
module = Module("shiftSrd")
|
|
if writer.states.version[:2] == (12, 5):
|
|
with writer.allocTmpSgpr(1, tag="shiftSrd_tmpSgprRes") as stmpRes:
|
|
module.addComment("Shift num records for gfx125x")
|
|
module.add(SAndB32(sgpr(stmpRes.idx), sgpr(srdIdx+2), 0x7F))
|
|
module.add(SLShiftLeftB32(sgpr(stmpRes.idx), 25, sgpr(stmpRes.idx)))
|
|
module.add(SAndB32(sgpr(srdIdx+1), sgpr(srdIdx+1), 0x1FFFFFF))
|
|
module.add(SOrB32(sgpr(srdIdx+1), sgpr(srdIdx+1), sgpr(stmpRes.idx)))
|
|
module.add(SLShiftRightB32(sgpr(srdIdx+2), 7, sgpr(srdIdx+2)))
|
|
|
|
return module
|
|
|
|
def _fetchNextWorkItem(self, writer, kernel, sWorkItemIdx, sAddress) -> Module:
|
|
"""Atomic fetch-and-increment for the dynamic work-queue counter.
|
|
|
|
Targets with scalar memory atomics use ``s_atomic_inc`` directly; targets
|
|
without them (e.g. gfx12/gfx1250, where ``HasSAtomic`` is false) issue a
|
|
returning vector atomic from lane 0 instead.
|
|
"""
|
|
module = Module("fetchNextWorkItem")
|
|
|
|
if writer.states.asmCaps["HasSAtomic"]:
|
|
module.add(SAtomicInc(dst=sgpr(sWorkItemIdx), base=sgpr(sAddress, 2), soffset=0,
|
|
smem=SMEMModifiers(glc=True), comment="Fetch next work item index"))
|
|
module.add(SWaitCnt(kmcnt=0, comment="Wait for scalar memory op"))
|
|
return module
|
|
|
|
# No scalar memory atomic: issue the wrapping fetch-and-increment as a returning
|
|
# vector atomic from lane 0 only. Use ``global_atomic_inc_u32`` in SADDR form
|
|
# (scalar 64-bit base + per-lane 32-bit offset).
|
|
memOrder = Component.StreamKMemoryOrdering.find(writer)
|
|
vZeroOffset = writer.vgprPool.checkOut(1, "AtomicZeroOffset")
|
|
vWrapValue = writer.vgprPool.checkOut(1, "AtomicWrapValue")
|
|
vFetchedIdx = writer.vgprPool.checkOut(1, "AtomicFetchedIdx")
|
|
sSavedExec = writer.sgprPool.checkOutAligned(writer.states.laneSGPRCount,
|
|
writer.states.laneSGPRCount,
|
|
"SavedExec")
|
|
execMovInst = SMovB32 if kernel["WavefrontSize"] == 32 else SMovB64
|
|
|
|
module.add(VMovB32(dst=vgpr(vWrapValue), src=sgpr(sWorkItemIdx),
|
|
comment="Queue wrap threshold (atomic_inc src)"))
|
|
module.add(VMovB32(dst=vgpr(vZeroOffset), src=0,
|
|
comment="Zero per-lane offset; queue base stays in saddr"))
|
|
module.add(memOrder.preVolatileVmem(writer, comment="drain xnacks before dynamic queue atomic"))
|
|
module.add(execMovInst(dst=sgpr(sSavedExec, writer.states.laneSGPRCount),
|
|
src=EXEC(), comment="save exec mask"))
|
|
module.add(VCmpXEqU32(dst=EXEC(), src0=vgpr("Serial"), src1=0,
|
|
comment="lane 0 fetches next work item"))
|
|
module.add(GlobalAtomicIncU32Saddr(dst=vgpr(vFetchedIdx),
|
|
vaddr=vgpr(vZeroOffset),
|
|
data=vgpr(vWrapValue),
|
|
saddr=sgpr(sAddress, 2),
|
|
modifier=GLOBALModifiers(scope=CacheScope.SCOPE_DEV),
|
|
comment="Fetch next work item index"))
|
|
module.add(SWaitCnt(vlcnt=0, comment="Wait for VMEM atomic return (loadcnt; global needs no dscnt)"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sWorkItemIdx), src=vgpr(vFetchedIdx),
|
|
comment="Read fetched work item index"))
|
|
module.add(execMovInst(dst=EXEC(),
|
|
src=sgpr(sSavedExec, writer.states.laneSGPRCount),
|
|
comment="restore exec mask"))
|
|
writer.sgprPool.checkIn(sSavedExec)
|
|
writer.vgprPool.checkIn(vFetchedIdx)
|
|
writer.vgprPool.checkIn(vWrapValue)
|
|
writer.vgprPool.checkIn(vZeroOffset)
|
|
return module
|
|
|
|
def _skv(self, writer, name):
|
|
"""Return the VGPR index holding a StreamK constant."""
|
|
return writer.states.skConstVgprs[name]
|
|
|
|
# ------------------------------------------------------------------
|
|
# Single-hop next-neighbor work stealing (codegen-time, off by default)
|
|
#
|
|
# Queues are one per XCD: numQueues = archCaps["NumXCD"], a power of two so
|
|
# queue mapping uses shift/AND fast masking (queueIdx = StreamKIdx & mask).
|
|
# A queue whose home fetch is empty steals once from its next neighbor
|
|
# s = (q+1) & mask; each queue has exactly one predecessor p = (q-1) & mask.
|
|
#
|
|
# AddressFlags buffer layout (per problem):
|
|
# [0, numQueues*stride) per-queue counters, one per cache line
|
|
# [numQueues*stride, ...) partials/fixup ready flags (one word per tile)
|
|
# Counter stride == archCaps["CacheLineBytes"] (128B on gfx942/gfx950), so
|
|
# each per-XCD counter sits on its own line. Each counter's atomic_inc uses a
|
|
# static predecessor-inclusive auto-reset bound, so it self-zeroes every
|
|
# launch -- there is no explicit end-of-kernel reset.
|
|
|
|
def _wsQueueConstants(self, writer, kernel):
|
|
"""Return (numQueues, mask, log2Queues, cacheLineLog2) for this arch.
|
|
|
|
``numQueues`` = archCaps["NumXCD"] and the counter stride =
|
|
archCaps["CacheLineBytes"]; both must be powers of two for the shift/AND
|
|
queue masking and queue-address shift to be valid (asserted below).
|
|
"""
|
|
numQueues = writer.states.archCaps["NumXCD"]
|
|
assert numQueues > 0 and (numQueues & (numQueues - 1)) == 0, (
|
|
"StreamK dynamic-queue fast masking requires a power-of-two queue count "
|
|
"(got %d for ISA %s)" % (numQueues, tuple(kernel["ISA"][:2])))
|
|
strideBytes = writer.states.archCaps["CacheLineBytes"]
|
|
assert strideBytes > 0 and (strideBytes & (strideBytes - 1)) == 0, (
|
|
"StreamK per-queue counter stride must be a power-of-two cache-line "
|
|
"size (got %d for ISA %s)" % (strideBytes, tuple(kernel["ISA"][:2])))
|
|
return numQueues, numQueues - 1, log2(numQueues), log2(strideBytes)
|
|
|
|
def _wsFlagsBaseOffset(self, writer, kernel):
|
|
"""Byte offset where the partials/fixup ready flags begin.
|
|
|
|
The flags region starts right after the per-queue counters, i.e. after
|
|
``numQueues * strideBytes`` bytes (8 * 128 = 1024 on gfx942/gfx950).
|
|
"""
|
|
numQueues, _, _, cacheLineLog2 = self._wsQueueConstants(writer, kernel)
|
|
return numQueues << cacheLineLog2
|
|
|
|
def _wsStructuralCount(self, mod, mask, log2Queues, sDst, sTotal, sQueue, sTmp, comment):
|
|
"""Emit sDst = (sTotal >> log2Queues) + [sQueue < (sTotal & mask)].
|
|
|
|
Reuses the shift/and(mask)/cmp/cselect idiom already used for
|
|
tilesInQueue / workgroupsInQueue in graWorkGroup so the per-queue
|
|
structural share (tiles or workgroups) can be recomputed for an
|
|
arbitrary queue index. ``sTmp`` is a caller-owned scratch SGPR.
|
|
"""
|
|
mod.add(SLShiftRightB32(dst=sgpr(sDst), src=sgpr(sTotal), shiftHex=log2Queues, comment=comment))
|
|
mod.add(SAndB32(dst=sgpr(sTmp), src0=sgpr(sTotal), src1=mask, comment="Remainder"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sQueue), src1=sgpr(sTmp), comment="Queue gets a structural extra?"))
|
|
mod.add(SCSelectB32(dst=sgpr(sTmp), src0=1, src1=0))
|
|
mod.add(SAddU32(dst=sgpr(sDst), src0=sgpr(sDst), src1=sgpr(sTmp)))
|
|
|
|
def streamKWorkStealingHomeBound(self, writer, mod, kernel, sBound, sQueueIdx, sGrid):
|
|
"""Fold the predecessor's workgroup count into the home auto-reset bound.
|
|
|
|
Adds the predecessor term W_p to the ``tiles_q + W_q - 1`` already in
|
|
``sBound``, giving the stealing bound ``tiles_q + W_q + W_p - 1`` (queue q
|
|
also absorbs W_p increments from its one predecessor p = (q-1) & mask).
|
|
Caller gates on kernel["StreamKWorkStealing"] and passes the grid SGPR
|
|
name ("skGrid" for SK4, "SKGrid" for SK5-dynamic); ``sQueueIdx`` is
|
|
preserved. Exact only when W_q >= 1 whenever tiles_q >= 1 (skGrid >=
|
|
numQueues); the Solution layer rejects debug overrides that break this.
|
|
"""
|
|
_, mask, log2Queues, _ = self._wsQueueConstants(writer, kernel)
|
|
sPred = writer.sgprPool.checkOut(1, "wsPredQueue")
|
|
sWp = writer.sgprPool.checkOut(1, "wsPredWorkgroups")
|
|
sTmp = writer.sgprPool.checkOut(1, "wsPredTmp")
|
|
# p = (q - 1) & mask (wraps 0 -> numQueues-1 for unsigned subtract)
|
|
mod.add(SSubU32(dst=sgpr(sPred), src0=sgpr(sQueueIdx), src1=1, comment="Predecessor queue (q-1)"))
|
|
mod.add(SAndB32(dst=sgpr(sPred), src0=sgpr(sPred), src1=mask, comment="Wrap predecessor index"))
|
|
# W_p = (skGrid >> log2) + [p < (skGrid & mask)]
|
|
self._wsStructuralCount(mod, mask, log2Queues, sWp, sGrid, sPred, sTmp,
|
|
comment="Predecessor workgroups W_(q-1)")
|
|
mod.add(SAddU32(dst=sgpr(sBound), src0=sgpr(sBound), src1=sgpr(sWp),
|
|
comment="Home auto-reset bound += predecessor workgroups (next-neighbor steal)"))
|
|
writer.sgprPool.checkIn(sTmp)
|
|
writer.sgprPool.checkIn(sWp)
|
|
writer.sgprPool.checkIn(sPred)
|
|
|
|
def streamKWorkStealingSteal(self, writer, mod, kernel, sQueueIdx, sWorkItemIdx, sGrid, mkLabel):
|
|
"""Single-hop next-neighbor steal on the per-XCD queue topology.
|
|
|
|
On entry sQueueIdx holds home queue q and sWorkItemIdx holds the home
|
|
fetch result (both live). If the home fetch was valid (index <
|
|
TotalItems) this is a no-op; otherwise one s_atomic_inc steals from the
|
|
next neighbor s = (q+1) & mask and the global tile index is recomputed
|
|
from s. A lost race leaves sWorkItemIdx >= TotalItems, so the downstream
|
|
valid-index check turns this WG into a no-op. sQueueIdx is clobbered
|
|
(advanced to s). Caller gates on kernel["StreamKWorkStealing"] and passes
|
|
the grid SGPR name ("skGrid" for SK4, "SKGrid" for SK5-dynamic). The
|
|
steal atomic uses the stolen queue's bound ``tiles_s + W_s + W_q - 1``.
|
|
"""
|
|
_, mask, log2Queues, cacheLineLog2 = self._wsQueueConstants(writer, kernel)
|
|
skFetchDone = mkLabel("SK_FetchDone")
|
|
mod.add(SCmpLtU32(src0=sgpr(sWorkItemIdx), src1=sgpr("TotalItems"), comment="Home fetch valid?"))
|
|
mod.add(SCBranchSCC1(labelName=skFetchDone.getLabelName(), comment="Valid work fetched; no steal"))
|
|
|
|
# Build the steal auto-reset bound tiles_s + W_s + W_q - 1 into
|
|
# sWorkItemIdx (dead here). W_q is the stealer's own workgroup count and
|
|
# must be computed while sQueueIdx still holds q, before advancing to s.
|
|
sTmp = writer.sgprPool.checkOut(1, "wsStealTmp")
|
|
sWq = writer.sgprPool.checkOut(1, "wsStealerWorkgroups")
|
|
self._wsStructuralCount(mod, mask, log2Queues, sWq, sGrid, sQueueIdx, sTmp,
|
|
comment="Stealer workgroups W_q")
|
|
|
|
# Walk to the immediate next queue (wrap within the per-XCD queues, single-hop next-neighbor).
|
|
mod.add(SAddU32(dst=sgpr(sQueueIdx), src0=sgpr(sQueueIdx), src1=1, comment="Next queue"))
|
|
mod.add(SAndB32(dst=sgpr(sQueueIdx), src0=sgpr(sQueueIdx), src1=mask, comment="Wrap queue index"))
|
|
|
|
# tiles_s into sWorkItemIdx, then += W_s and += W_q, then -1.
|
|
self._wsStructuralCount(mod, mask, log2Queues, sWorkItemIdx, "TotalItems", sQueueIdx, sTmp,
|
|
comment="Stolen-queue tiles tiles_s")
|
|
sWs = writer.sgprPool.checkOut(1, "wsStolenWorkgroups")
|
|
self._wsStructuralCount(mod, mask, log2Queues, sWs, sGrid, sQueueIdx, sTmp,
|
|
comment="Stolen-queue workgroups W_s")
|
|
mod.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=sgpr(sWs), comment="tiles_s + W_s"))
|
|
writer.sgprPool.checkIn(sWs)
|
|
mod.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=sgpr(sWq), comment="+ W_q (stealer)"))
|
|
mod.add(SSubU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=1, comment="Steal auto-reset bound"))
|
|
writer.sgprPool.checkIn(sWq)
|
|
writer.sgprPool.checkIn(sTmp)
|
|
|
|
# One atomic on the neighbor's counter with the static self-reset bound.
|
|
sAddress = writer.sgprPool.checkOutAligned(2, 2, "wsStealAddress")
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sAddress), src=sgpr(sQueueIdx), shiftHex=cacheLineLog2, comment="Stride queues to cache lines (stolen queue)"))
|
|
mod.add(SAddU32(dst=sgpr(sAddress+0), src0=sgpr(sAddress+0), src1=sgpr("AddressFlags+0")))
|
|
mod.add(SAddCU32(dst=sgpr(sAddress+1), src0=0, src1=sgpr("AddressFlags+1")))
|
|
mod.add(SAtomicInc(dst=sgpr(sWorkItemIdx), base=sgpr(sAddress, 2), soffset=0, smem=SMEMModifiers(glc=True), comment="Fetch stolen work item index"))
|
|
mod.add(SWaitCnt(kmcnt=0, comment="Wait for scalar memory op"))
|
|
writer.sgprPool.checkIn(sAddress)
|
|
# Recompute global tile index from the neighbor's queue.
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sWorkItemIdx), src=sgpr(sWorkItemIdx), shiftHex=log2Queues))
|
|
mod.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=sgpr(sQueueIdx)))
|
|
mod.add(skFetchDone)
|
|
|
|
@abc.abstractmethod
|
|
def preLoop(self, writer, kernel):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def graWorkGroup(self, writer, kernel, tPA, tPB):
|
|
pass
|
|
|
|
def prefetchAcrossPersistentSetupNextTile(self, writer, kernel, tPA, tPB, skipLroReset=False):
|
|
"""Recompute StreamK tile locals and map tile index to WorkGroup* for the *next* tile.
|
|
|
|
After each persistent iteration's main body, ``StreamKIter`` already holds the starting
|
|
global iteration index for the next chunk (set at the beginning of ``graWorkGroup``).
|
|
Running ``skTileIndex`` + ``skIndexToWG`` + WGM remapping here matches the start of the
|
|
next ``setupNewTile`` / ``graWorkGroup`` (without advancing ``StreamKIter`` again), so
|
|
SGPRs are warm before the persistent back-edge.
|
|
|
|
When ``skipLroReset`` is True the local-read-offset reset inside
|
|
``skTileIndex`` is suppressed. This is needed when PAP runs *before*
|
|
the NLL body: the NLL still needs the current tile's read pointers."""
|
|
from Tensile.Components.WorkGroupMappingAlgos import DefaultWGM, SpaceFillingCurveWalk
|
|
|
|
module = Module("StreamK prefetchAcrossPersistentSetupNextTile")
|
|
with writer.allocTmpSgpr(4, 2, "SKPrefetchTemp") as sTmpRes:
|
|
sTmp = sTmpRes.idx
|
|
module.add(self.skTileIndex(writer, kernel, sTmp, tPA, tPB, skipLroReset=skipLroReset))
|
|
module.add(self.skIndexToWG(writer, kernel, sTmp))
|
|
if len(kernel["SpaceFillingAlgo"]):
|
|
writer.states.WGMTransformLevels = len(kernel["SpaceFillingAlgo"])
|
|
module.add(SpaceFillingCurveWalk(writer, kernel, "WGM"))
|
|
else:
|
|
module.add(DefaultWGM(writer, kernel, "WGM"))
|
|
return module
|
|
|
|
def computeTotalTiles(self, writer, kernel, dstSgpr):
|
|
"""Compute totalTiles = NumWorkGroups0 * NumWorkGroups1 * batchCount into dstSgpr."""
|
|
module = Module("StreamK computeTotalTiles")
|
|
module.add(SMulI32(dst=sgpr(dstSgpr), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"), comment="totalTiles = nwg0 * nwg1"))
|
|
for i in range(kernel["ProblemType"]["NumIndicesC"] - kernel["ProblemType"]["NumIndicesFree"]):
|
|
batchIdx = kernel["ProblemType"]["NumIndicesFree"] + i
|
|
module.add(SMulI32(dst=sgpr(dstSgpr), src0=sgpr(dstSgpr), src1=sgpr("SizesFree+%u" % batchIdx), comment="totalTiles *= batch dim %u" % i))
|
|
return module
|
|
|
|
def computeTotalIters(self, writer, kernel, dstSgpr):
|
|
"""Compute totalIters = NumWorkGroups0 * NumWorkGroups1 * batchCount * ItersPerTile into dstSgpr."""
|
|
module = Module("StreamK computeTotalIters")
|
|
module.add(self.computeTotalTiles(writer, kernel, dstSgpr))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMulI32(dst=sgpr(dstSgpr), src0=sgpr(dstSgpr), src1=sgpr(sIpt), comment="totalIters = totalTiles * itersPerTile"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
return module
|
|
|
|
def skTileIndex(self, writer, kernel, sTmp, tPA, tPB, skipLroReset=False):
|
|
module = Module("StreamK skTileIndex")
|
|
skConstsInVgprs = writer.isStreamKConstantsToVgprEnabled(kernel)
|
|
|
|
# Always reset pointers to handle odd-exit case which moves LRO to the upper bank.
|
|
# Skipped when PAP calls this before the NLL body: the current
|
|
# tile's local read pointers must stay intact for the remaining MACs.
|
|
if kernel["PrefetchGlobalRead"] and not skipLroReset:
|
|
if not kernel["UseSubtileImpl"]:
|
|
module.add(writer.localReadResetOffsets(kernel, tPA))
|
|
if kernel["ProblemType"]["MXBlockA"] and "MX" in tPA:
|
|
module.add(writer.localReadResetOffsets(kernel, tPA["MX"]))
|
|
if kernel["ProblemType"]["MXBlockB"] and "MX" in tPB:
|
|
module.add(writer.localReadResetOffsets(kernel, tPB["MX"]))
|
|
module.add(writer.localReadResetOffsets(kernel, tPB))
|
|
else:
|
|
module.add(localReadResetOffsetsSubtile(writer, kernel))
|
|
|
|
module.addComment0("StreamK calculate tile idx and map to WG")
|
|
|
|
# sTmp = tile index
|
|
sMagicNum = writer.acquireStreamKConstSgpr(kernel, "MagicNumberItersPerTile")
|
|
sMagicShift = writer.acquireStreamKConstSgpr(kernel, "MagicShiftItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sMagicNum), src=vgpr(writer.states.skConstVgprs["MagicNumberItersPerTile"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sMagicShift), src=vgpr(writer.states.skConstVgprs["MagicShiftItersPerTile"])))
|
|
# SK5: mode bit (30) already cleared at preLoop; keep magic add + shift.
|
|
if kernel["StreamK"] == 5:
|
|
module.add(SAndB32(dst=sgpr(sMagicShift), src0=sgpr(sMagicShift), src1=hex(0x8000001F),
|
|
comment="SK5: keep magic add bit (31) + 5-bit shift, drop mode bit (30)"))
|
|
module.add(sMagicDiv2(sgpr(sTmp), sgpr(sTmp+1), sgpr("StreamKIter"), sgpr(sMagicNum), sgpr(sMagicShift), sgpr(sTmp+2)))
|
|
writer.releaseStreamKConstSgpr(sMagicNum)
|
|
writer.releaseStreamKConstSgpr(sMagicShift)
|
|
# sTmp+1 = tile start, sTmp+2 = tile end
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr(sTmp), src1=sgpr(sIpt), comment="Tile start iteration"))
|
|
module.add(SAddU32(dst=sgpr(sTmp+2), src0=sgpr(sTmp+1), src1=sgpr(sIpt), comment="Tile end iteration"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
# local start
|
|
module.add(SSubU32(dst=sgpr("StreamKLocalStart"), src0=sgpr("StreamKIter"), src1=sgpr(sTmp+1), comment="Local iteration start"))
|
|
# local end (SK tile)
|
|
module.add(SMinU32(dst=sgpr("StreamKLocalEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(sTmp+2), comment="1. (Local) iteration end (SK tile)"))
|
|
module.add(SSubU32(dst=sgpr("StreamKLocalEnd"), src0=sgpr("StreamKLocalEnd"), src1=sgpr(sTmp+1), comment="2. Local iteration end (SK tile)"))
|
|
|
|
return module
|
|
|
|
def skIndexToWG(self, writer, kernel, sTmp):
|
|
# Note: There's one unused sgpr passed with sTmp.
|
|
module = Module("StreamK skIndexToWG")
|
|
|
|
# Map StreamK tile index to wg0/1
|
|
module.addComment0("Map StreamK tile index to wg0/1/2")
|
|
module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"), comment="Total tiles"))
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
module.add(scalarUInt32DivideAndRemainder(qReg="WorkGroup2", dReg=sTmp, divReg=sTmp+1, rReg=sTmp+2, tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=True, comment="TileID // nWG0*nWG1"))
|
|
|
|
# Store tileID for use later in general WGM algo
|
|
if kernel["SpaceFillingAlgo"]:
|
|
module.add(SNop(waitState=1, comment=""))
|
|
module.add(SMovB32(dst=sgpr("StreamKTileID"), src=sgpr(sTmp+2), comment=""))
|
|
|
|
module.add(scalarUInt32DivideAndRemainder(qReg="WorkGroup1", dReg=sTmp+2, divReg="NumWorkGroups0", rReg="WorkGroup0", tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=True, comment="TileID // nWG0"))
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
module.addSpaceLine()
|
|
|
|
return module
|
|
|
|
def skExtraIters(self, writer, kernel, sSkExtraIters, sTmp):
|
|
# skExtraIters = skTiles * ItersPerTile - SKItersPerWG * skGrid
|
|
# Use sSkExtraIters/sTmp as readfirstlane destinations to reduce SGPR pressure
|
|
skConstsInVgprs = writer.isStreamKConstantsToVgprEnabled(kernel)
|
|
module = Module("StreamK skExtraIters")
|
|
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sSkExtraIters), src=vgpr(writer.states.skConstVgprs["skTiles"])))
|
|
sT = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sT), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
skTilesSrc = sSkExtraIters
|
|
else:
|
|
skTilesSrc = "skTiles"
|
|
|
|
module.add(SMulI32(dst=sgpr(sSkExtraIters), src0=sgpr(skTilesSrc), src1=sgpr(sT)))
|
|
writer.releaseStreamKConstSgpr(sT)
|
|
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sTmp), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
sT = writer.acquireStreamKConstSgpr(kernel, "skGrid")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sT), src=vgpr(writer.states.skConstVgprs["skGrid"])))
|
|
skItersPerWgSrc = sTmp
|
|
else:
|
|
skItersPerWgSrc = "SKItersPerWG"
|
|
|
|
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(skItersPerWgSrc), src1=sgpr(sT)))
|
|
writer.releaseStreamKConstSgpr(sT)
|
|
|
|
module.add(SSubU32(dst=sgpr(sSkExtraIters), src0=sgpr(sSkExtraIters), src1=sgpr(sTmp), comment="skTiles * ItersPerTile - SKItersPerWG * skGrid"))
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def computeLoadSrd(self, writer, kernel, tP, sTmp):
|
|
pass
|
|
|
|
def computeLoadSrdCommon(self, writer, kernel, tP, sTmp):
|
|
module = Module("StreamK Common computeLoadSrd")
|
|
|
|
tileStart = sTmp + 2
|
|
tc = tP["tensorChar"]
|
|
depthU = self._depthUForTc(kernel, tc)
|
|
# StreamK partial tile - offset to tile start index
|
|
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("StreamKLocalStart"), src1=depthU, comment="StreamK tile start offset"))
|
|
strideL = writer.strideRef(tc, kernel["ProblemType"]["IndicesSummation"][0])
|
|
module.add(writer.s_mul_u64_u32(sgpr(sTmp), sgpr(sTmp+1), sgpr(sTmp), strideL, comment="StreamK tile start offset"))
|
|
# Overflow check removed
|
|
# if kernel["CheckDimOverflow"] >=2:
|
|
# kStr += self.assert_eq(sgpr(sTmp+1),0)
|
|
module.add(SAddU32(dst=sgpr(tileStart+0), src0=sgpr(tileStart+0), src1=sgpr(sTmp+0), comment="accum GsuOffset term to tilestart"))
|
|
module.add(SAddCU32(dst=sgpr(tileStart+1), src0=sgpr(tileStart+1), src1=sgpr(sTmp+1), comment="accum GsuOffset term to tilestart"))
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def computeStoreSrdStart(self, writer, kernel):
|
|
pass
|
|
|
|
def computeStoreSrdStartCommon(self, writer, kernel):
|
|
module = Module("StreamK Common computeStoreSrdStart")
|
|
if kernel["StreamKForceDPOnly"]:
|
|
return module
|
|
skConstsInVgprs = writer.isStreamKConstantsToVgprEnabled(kernel)
|
|
|
|
# Check for parallel reduction
|
|
# Paralell reduction stores to SrdD in split format, fixup happens in post kernel
|
|
skSplitSrd = Label("SK_SplitSrd", "")
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=skSplitSrd.getLabelName(), comment="Skip this block if using single-kernel stream-k fixup"))
|
|
# Alpha/Beta will be applied in post kernel if necessary
|
|
# module.add(SMovB32(dst=sgpr("Alpha"), src=1.0, comment="For parallel reduction, alpha applied in post kernel"))
|
|
# module.add(SMovB32(dst=sgpr("Beta"), src=0.0, comment="For parallel reduction, beta applied in post kernel"))
|
|
|
|
indices = list(range(0, kernel["ProblemType"]["NumIndicesC"]))
|
|
numDim = len(indices)
|
|
|
|
sSkt = writer.acquireStreamKConstSgpr(kernel, "skTiles")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(writer.states.skConstVgprs["skTiles"])))
|
|
module.add(SCmpEQU32(src0=sgpr(sSkt), src1=1, comment="split == 1 ?"))
|
|
writer.releaseStreamKConstSgpr(sSkt)
|
|
module.add(SCBranchSCC1(labelName=skSplitSrd.getLabelName(), comment="branch if split == 1"))
|
|
# Parallel reduction: adjust output buffer address to per split buffer
|
|
with writer.allocTmpSgpr(4, alignment=1, tag="computeStoreSrdStartCommon_tmpSgprInfo") as tmpSgprInfo:
|
|
if tmpSgprInfo.idx % 2 == 0:
|
|
tmpSgprX2 = tmpSgprInfo.idx+0
|
|
tmpSgpr0 = tmpSgprInfo.idx+0
|
|
tmpSgpr1 = tmpSgprInfo.idx+1
|
|
tmpSgpr2 = tmpSgprInfo.idx+2
|
|
tmpSgpr3 = tmpSgprInfo.idx+3
|
|
else:
|
|
tmpSgprX2 = tmpSgprInfo.idx+1
|
|
tmpSgpr0 = tmpSgprInfo.idx+1
|
|
tmpSgpr1 = tmpSgprInfo.idx+2
|
|
tmpSgpr2 = tmpSgprInfo.idx+0
|
|
tmpSgpr3 = tmpSgprInfo.idx+3
|
|
module.addComment("Split Output Buffer offset: Free0 + (Free1-1)*StrideC1J + (Free2-1)*StrideCK * SplitIdx * bpe%s")
|
|
# PartialIdx was saved in sgprBeta for re-use
|
|
module.addModuleAsFlatItems(writer.s_mul_u64_u32(sgpr(tmpSgpr0), sgpr(tmpSgpr1), sgpr("SizesFree+0"), sgpr("SkPartialIdx"), comment="Free0"))
|
|
for i in range(1, numDim):
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr2), src0=sgpr("SizesFree+%u"%i), src1=1, comment="Free%u" % i))
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr2), src0=sgpr(tmpSgpr2), src1=sgpr("SkPartialIdx"), comment="Free%u" % i))
|
|
module.addModuleAsFlatItems(writer.s_mul_u64_u32(sgpr(tmpSgpr2), sgpr(tmpSgpr3), sgpr(tmpSgpr2), sgpr("StrideC%s"%writer.states.indexChars[i]), comment="Free%u" % i))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr0), src0=sgpr(tmpSgpr0), src1=sgpr(tmpSgpr2), comment="Free%u" % i))
|
|
module.add(SAddCU32(dst=sgpr(tmpSgpr1), src0=sgpr(tmpSgpr1), src1=sgpr(tmpSgpr3), comment="Free%u" % i))
|
|
module.add(SLShiftLeftB64(dst=sgpr(tmpSgprX2,2), src=sgpr(tmpSgprX2,2), shiftHex=log2(writer.states.bpeCinternal), comment="scale by bpe"))
|
|
module.add(SAddU32(dst=sgpr("SrdD+0"), src0=sgpr("SrdD+0"), src1=sgpr(tmpSgprX2), comment="add lo GSU offset to SRD"))
|
|
module.add(SAddCU32(dst=sgpr("SrdD+1"), src0=sgpr("SrdD+1"), src1=sgpr(tmpSgpr1), comment="add hi GSU offset to SRD"))
|
|
|
|
module.add(skSplitSrd)
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def graAddresses(self, writer, kernel, tP, vTmp):
|
|
pass
|
|
|
|
def graAddressesCommon(self, writer, kernel, tP, vTmp):
|
|
module = Module("StreamK Common graAddresses")
|
|
|
|
tc = tP["tensorChar"]
|
|
depthU = self._depthUForTc(kernel, tc)
|
|
# StreamK partial tile - offset to tile start index
|
|
tmpOffset = writer.sgprPool.checkOut(2, "skStartOffset")
|
|
module.add(SMulI32(dst=sgpr(tmpOffset), src0=sgpr("StreamKLocalStart"), src1=int(depthU * tP["bpe"]), comment="StreamK tile start offset"))
|
|
strideL = writer.strideRef(tc, kernel["ProblemType"]["IndicesSummation"][0])
|
|
module.add(writer.s_mul_u64_u32(sgpr(tmpOffset), sgpr(tmpOffset+1), sgpr(tmpOffset), strideL, comment="StreamK tile start offset"))
|
|
# Overflow check removed
|
|
# if kernel["CheckDimOverflow"] >=2:
|
|
# kStr += self.assert_eq(sgpr(tmpOffset+1),0)
|
|
module.add(SAddU32(dst=sgpr(tmpOffset+0), src0=sgpr(tmpOffset+0), src1=sgpr("Address%s+0" % tc), comment="accum skOffset term to tilestart"))
|
|
module.add(SAddCU32(dst=sgpr(tmpOffset+1), src0=sgpr(tmpOffset+1), src1=sgpr("Address%s+1" % tc), comment="accum skOffset term to tilestart"))
|
|
module.add(VMovB32(dst=vgpr(vTmp+0), src=sgpr(tmpOffset+0)))
|
|
module.add(VMovB32(dst=vgpr(vTmp+1), src=sgpr(tmpOffset+1)))
|
|
writer.sgprPool.checkIn(tmpOffset)
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def declareStaggerParms(self, writer, kernel):
|
|
pass
|
|
|
|
def declareStaggerParmsCommon(self, writer, kernel):
|
|
module = Module("StreamK Common declareStaggerParms")
|
|
|
|
# Set stagger=0 for partial tiles to avoid using stagger larger than workload
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpGtU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
module.add(SCMovB32(dst=sgpr("StaggerUIter"), src=0, comment="set stagger=0 for partial tiles"))
|
|
module.add(SCmpLtU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr(sIpt), comment="does wg finish tile?"))
|
|
module.add(SCMovB32(dst=sgpr("StaggerUIter"), src=0, comment="set stagger=0 for partial tiles"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def tailLoopNumIter(self, writer, kernel, loopCounter):
|
|
pass
|
|
|
|
def tailLoopNumIterCommon(self, writer, kernel, loopCounter):
|
|
module = Module("StreamK Common tailLoopNumIter")
|
|
|
|
# skip tail loop if StreamK WG not processing final iteration
|
|
# Check if tile finished
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpLtU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr(sIpt), comment="Check if WG processes final iteration of tile"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SCMovB32(dst=loopCounter, src=0, comment="This WG not completing tile"))
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def calculateLoopNumIter(self, writer, kernel, loopCounterName, loopIdx, tmpSgprInfo):
|
|
pass
|
|
|
|
def calculateLoopNumIterCommon(self, writer, kernel, loopCounterName, loopIdx, tmpSgprInfo):
|
|
module = Module("StreamK Common calculateLoopNumIter")
|
|
|
|
# Use StreamK params for loop count
|
|
module.add(SSubU32(dst=sgpr(loopCounterName), src0=sgpr("StreamKLocalEnd"), src1=sgpr("StreamKLocalStart"), comment="StreamK loop counter = localEnd - localStart"))
|
|
# Short circuit if alpha==0 (set loopCounter to 0 to skip main loop)
|
|
alphaLabel2 = Label(writer.labels.getNameInc("SKAlphaCheck"), "")
|
|
module.add(BranchIfNotZero("Alpha", kernel["ProblemType"]["ComputeDataType"].toEnum(), alphaLabel2))
|
|
module.add(SMovB32(dst=sgpr(loopCounterName), src=0, comment="Skip iterations"))
|
|
module.add(alphaLabel2)
|
|
|
|
# Adjust loop count for tail loop
|
|
if not kernel["NoTailLoop"]:
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
unrollIdx = writer.states.unrollIdx
|
|
loopChar = writer.states.indexChars[kernel["ProblemType"]["IndicesSummation"][unrollIdx]]
|
|
|
|
assert kernel["DepthU"] % 2 == 0 # Assuming DepthU is power of 2, if odd DepthU were supported this divide would need 2 more temp registers for divide
|
|
maxUnit = writer.states.tailloopInNllmaxUnit
|
|
# tailloopInNll + maxUnit == 1 case, tailloopInNll is always used and no need to adjust loopCounter
|
|
if not (writer.states.tailloopInNll and maxUnit == 1):
|
|
if ((kernel["DepthU"] & (kernel["DepthU"] - 1)) == 0):
|
|
module.add(scalarStaticDivideAndRemainder(qReg=tmpSgpr, rReg=tmpSgpr+1, dReg=("SizesSum+%u" % unrollIdx), divisor=kernel["DepthU"], tmpSgprRes=None, doRemainder=2))
|
|
else:
|
|
with writer.allocTmpSgpr(4, tag="calculateLoopNumIterCommon_tmpSgpr1") as tmpSgpr1:
|
|
module.add(scalarStaticDivideAndRemainder(qReg=tmpSgpr, rReg=tmpSgpr+1, dReg=("SizesSum+%u" % unrollIdx), divisor=kernel["DepthU"], tmpSgprRes=tmpSgpr1, doRemainder=2))
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+1), src1=0, comment="numIter%s == 0"%loopChar ))
|
|
module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=0, src1=1, comment="check if size uses tail loop"))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr(sIpt), comment="Check if WG processes final iteration of tile"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=0, comment="this WG runs tail loop"))
|
|
|
|
if writer.states.tailloopInNll and maxUnit > 1:
|
|
# tailloopInNll + maxUnit > 1 case, we need to check if SizesSum is multiple of maxUnit at runtime.
|
|
# if SizesSum is not multiple of maxUnit, we do not use tailloopInNll and need to decrement loopCounter for StreamK
|
|
# if SizesSum is multiple of maxUnit, we use tailloopInNll and need to increment loopCounter by 1.
|
|
# With considering both increment and decrement, we do not need to adjust loopCounter.
|
|
module.add(SAndB32(dst=sgpr(tmpSgpr+2), src0=sgpr("SizesSum+%u" % unrollIdx), src1=maxUnit-1, \
|
|
comment="if summation is not multiple of %u, skip tailloopInNll"%maxUnit))
|
|
module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=0, comment="do not decrement in tailloopInNll case"))
|
|
|
|
module.add(SSubU32(dst=sgpr(loopCounterName), src0=sgpr(loopCounterName), src1=sgpr(tmpSgpr), comment="Adjust loop counter for tail loop"))
|
|
module.add(SMaxI32(dst=sgpr(loopCounterName), src0=sgpr(loopCounterName), src1=0, comment="Avoid setting negative value to loopCounter"))
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def storeBranches(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct):
|
|
pass
|
|
|
|
def storeBranchesCommon(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct):
|
|
module = Module("StreamK Common storeBranches")
|
|
|
|
# No branches when no StreamK partial/fixup path can be reached.
|
|
if kernel["StreamKAtomic"] or kernel["StreamKForceDPOnly"]:
|
|
return module
|
|
|
|
memOrder = Component.StreamKMemoryOrdering.find(writer)
|
|
skConstsInVgprs = writer.isStreamKConstantsToVgprEnabled(kernel)
|
|
skStoreLabel = Label(label=writer.labels.getNameInc("SK_Store"), comment="")
|
|
|
|
if kernel["StreamKFixupTreeReduction"] == 1:
|
|
skFixupTreeLabel = Label(label=writer.labels.getNameInc("SK_Fixup_Tree"), comment="")
|
|
skFixupTreeLoopStart = Label(label=writer.labels.getNameInc("SK_Fixup_TreeLoop_Start"), comment="")
|
|
skFixupWaitForFlag = Label(label=writer.labels.getNameInc("SK_Fixup_Wait_Flag"), comment="")
|
|
endFixupLoop = Label(label=writer.labels.getNameInc("endFixupLoop"), comment="")
|
|
skFixupCalcPartialIdx = Label(label=writer.labels.getNameInc("SK_Fixup_CalcPartialIdx"), comment="")
|
|
|
|
# sIter = writer.sgprPool.checkOut(2, "SKIter")
|
|
sPartialIdx = writer.sgprPool.checkOut(1, "SK_Fixup_Partial_idx")
|
|
|
|
sSkExtraIters = writer.sgprPool.checkOut(1, "extraIters")
|
|
tmpSgpr = writer.sgprPool.checkOut(1, tag="StreamKCommon_storeBranches_tmpSgpr")
|
|
module.add(self.skExtraIters(writer, kernel, sSkExtraIters, tmpSgpr))
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
# Skip to global write if WG started and finished tile
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
module.add(SCBranchSCC0(labelName=skFixupTreeLabel.getLabelName(), comment="If we didn't start the tile, always to SK Tree fixup"))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr(sIpt), comment="does wg finish tile?"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SCBranchSCC1(labelName=skStoreLabel.getLabelName(), comment="Branch if started and finished tile, go to regular store code"))
|
|
|
|
# Start Tree Fixup
|
|
module.add(skFixupTreeLabel)
|
|
|
|
# Calculate partialIdx
|
|
tmpVgpr = writer.vgprPool.checkOutAligned(4, 2, tag="StreamKCommon_storeBranches_tmpVgpr")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=4)
|
|
|
|
# Get Iter of the start of the tile
|
|
tmpSgpr = writer.sgprPool.checkOut(3, tag="StreamKCommon_storeBranches_tmpSgpr2")
|
|
|
|
# Compute dpSectionSize = (totalTiles - skTiles) * ItersPerTile
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
sSkt = writer.acquireStreamKConstSgpr(kernel, "skTiles")
|
|
sIpw = writer.acquireStreamKConstSgpr(kernel, "SKItersPerWG")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(writer.states.skConstVgprs["skTiles"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpw), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
module.add(self.computeTotalTiles(writer, kernel, tmpSgpr))
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=sgpr(sSkt), comment="dpTiles = totalTiles - skTiles"))
|
|
writer.releaseStreamKConstSgpr(sSkt)
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=sgpr(sIpt), comment="Offset to first SK tile"))
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr("StreamKIter"), src1=sgpr(tmpSgpr), comment="Iter relative to starting SK iter"))
|
|
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr+1), src0=sgpr(tmpSgpr), src1=1, comment="minus 1 to get Iter in current tile"))
|
|
module.add(scalarUInt24DivideAndRemainder(qReg=tmpSgpr+0, dReg=tmpSgpr+1, divReg=sIpt, rReg=-1, tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=False, comment="wgCount = tileStart / (itersPerTile)"))
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr+1), src0=sgpr(tmpSgpr+0), src1=sgpr(sIpt), comment="tileStart=tileIdx * ItersPerTile"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr+0), src0=sgpr(sIpw), src1=1, comment="ItersPerWG w/ extraIter"))
|
|
module.add(scalarUInt24DivideAndRemainder(qReg=tmpSgpr+2, dReg=tmpSgpr+1, divReg=tmpSgpr+0, rReg=-1, tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=False, comment="wgCount = tileStart / (itersPerWG+1)"))
|
|
module.add(SCmpLtU32(src0=sgpr(tmpSgpr+2), src1=sgpr(sSkExtraIters), comment="find co-op group start"))
|
|
module.add(SCBranchSCC1(labelName=skFixupCalcPartialIdx.getLabelName(), comment="All WG have extra iter so far, skip following calcs"))
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr+0), src0=sgpr(tmpSgpr+1), src1=sgpr(sSkExtraIters), comment="tileStart - extraIters"))
|
|
module.add(scalarUInt24DivideAndRemainder(qReg=tmpSgpr+2, dReg=tmpSgpr+0, divReg=sIpw, rReg=-1, tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=False, comment="wgExtraIters = (tileStart - extraIters) / itersPerWG"))
|
|
writer.releaseStreamKConstSgpr(sIpw)
|
|
module.add(skFixupCalcPartialIdx)
|
|
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(SSubU32(dst=sgpr(sPartialIdx), src0=sgpr(sIdx), src1=sgpr(tmpSgpr+2), comment="partialIdx = streamkidx - coopGroupStart"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
|
|
sFlagIdx = writer.sgprPool.checkOut(1, "FlagIdx")
|
|
sIdxOffset = writer.sgprPool.checkOut(1, "IdxOffset")
|
|
module.add(SMovB32(dst=sgpr(sIdxOffset), src=1, comment="Init IdxOffset=1"))
|
|
|
|
module.add(skFixupTreeLoopStart) # start tree fixup loop
|
|
|
|
# First, jump to partial write if (partialIdx//2)*2 != partialIdx, i.e. branch if last bit is 1
|
|
module.add(SAndB32(dst=sgpr(tmpSgpr+0), src0=sgpr(sPartialIdx), src1=1))
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+0), src1=1, comment="partialIdx&1==1?"))
|
|
module.add(writer.longBranchScc1(skPartialsLabel, posNeg=1))
|
|
module.add(SLShiftRightB32(dst=sgpr(sPartialIdx), src=sgpr(sPartialIdx), shiftHex=log2(2), comment="sPartialIdx // 2"))
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(SAddU32(dst=sgpr(sFlagIdx), src0=sgpr(sIdx), src1=sgpr(sIdxOffset), comment="flagIdx=StreamKIdx+IdxOffset"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
|
|
# If the flag we're waiting for is past this tile we can finish the fixup step
|
|
# LocalEnd + 1 + (sIdxOffset-1) * SKItersPerWG + (Extras)
|
|
sIpw = writer.acquireStreamKConstSgpr(kernel, "SKItersPerWG")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpw), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr+1), src0=sgpr(sIdxOffset), src1=1, comment="Starting on next WG so offset-1"))
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr+2), src0=sgpr(sIpw), src1=sgpr(tmpSgpr+1), comment="Before extra iters"))
|
|
writer.releaseStreamKConstSgpr(sIpw)
|
|
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr+0), src0=sgpr(sFlagIdx), src1=sgpr(sSkExtraIters), comment="TargetWG-ExtraIters"))
|
|
module.add(SMinU32(dst=sgpr(tmpSgpr+0), src0=sgpr(tmpSgpr+0), src1=sgpr(tmpSgpr+1), comment="min of above and (offset-1)"))
|
|
module.add(SCmpLtU32(src0=sgpr(sFlagIdx), src1=sgpr(sSkExtraIters), comment="TargetWG < extraIters?"))
|
|
module.add(SCSelectB32(dst=sgpr(tmpSgpr+0), src0=0, src1=sgpr(tmpSgpr+0), comment="If True, don't sub any iters"))
|
|
module.add(SSubU32(dst=sgpr(tmpSgpr+1), src0=sgpr(tmpSgpr+1), src1=sgpr(tmpSgpr+0), comment="extras = (offset-1) - (possible extras)"))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr+2), src0=sgpr(tmpSgpr+2), src1=sgpr(tmpSgpr+1), comment="Add possible extra iters"))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr+0), src0=sgpr("StreamKLocalEnd"), src1=1, comment="Start of next wg"))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr+2), src0=sgpr(tmpSgpr+0), src1=sgpr(tmpSgpr+2)))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpGtU32(src0=sgpr(tmpSgpr+2), src1=sgpr(sIpt)))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SCBranchSCC1(labelName=endFixupLoop.getLabelName()))
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
# check flag
|
|
tmpSgpr = writer.sgprPool.checkOut(2, "globalWriteElements")
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(sFlagIdx), shiftHex=log2(4), comment="flag offset based on wg index"))
|
|
|
|
module.add(skFixupWaitForFlag) # loop to wait for flag
|
|
module.add(memOrder.readFlag(writer, dst=tmpSgpr+1, soffset=sgpr(tmpSgpr)))
|
|
if kernel["DebugStreamK"] & 2 == 0: # Don't wait for partials if not being written
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+1), src1=1, comment="check if ready"))
|
|
module.add(SCBranchSCC0(labelName=skFixupWaitForFlag.getLabelName(), comment="if flag not set, wait and check again"))
|
|
module.add(memOrder.acquireFence(writer))
|
|
|
|
module.add(SBarrier(comment="wait for all workgroups before resetting flag"))
|
|
skipFlagReset = Label(label=writer.labels.getNameInc("SK_SkipFlagReset"), comment="")
|
|
module.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr+2), src=vgpr("Serial"), comment="Wave 0 updates flags"))
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=0, comment="Check for wave 0"))
|
|
module.add(SCBranchSCC0(labelName=skipFlagReset.getLabelName(), comment="Skip flag reset"))
|
|
if writer.states.asmCaps["HasScalarStore"]:
|
|
# (tmpSgpr+2) contains a vlue of 0, use it to reset the flag
|
|
module.add(SStoreB32(src=sgpr(tmpSgpr+2), base=sgpr("AddressFlags", 2), soffset=sgpr(tmpSgpr), smem=SMEMModifiers(glc=True), comment="reset flag"))
|
|
else:
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr), src=0, comment="move 0 to tmpVgpr"))
|
|
module.add(self.setFlagValue(writer, src=vgpr(tmpVgpr), soffset=sgpr(tmpSgpr), comment="reset flag"))
|
|
module.add(skipFlagReset)
|
|
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
# fixup step
|
|
if kernel["DebugStreamK"] & 1 == 0: # Skip fixup reads if set, need to do the loop if partial writes are enabled
|
|
fixupEdge = [False] # Test no edge variant
|
|
module.add(self.fixupStep(writer, kernel, vectorWidths, elements, fixupEdge, tmpVgpr, cvtVgprStruct, sFlagIdx))
|
|
|
|
# Could branch if our new offset puts us off the tile, but we essentially do that when calculating if our target wg is off the tile earlier
|
|
module.add(SLShiftLeftB32(dst=sgpr(sIdxOffset), src=sgpr(sIdxOffset), shiftHex=log2(2), comment="IdxOffset *= 2 for Tree reduction"))
|
|
module.add(SBranch(labelName=skFixupTreeLoopStart.getLabelName(), comment="Branch to continue fixup loop"))
|
|
|
|
# If we started the tile, we reduced the partial results to that WG, so global write
|
|
# Otherwise, partial write
|
|
module.add(endFixupLoop)
|
|
# Done fixup loop
|
|
writer.sgprPool.checkIn(sIdxOffset)
|
|
writer.sgprPool.checkIn(sFlagIdx)
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
module.add(writer.longBranchScc0(skPartialsLabel, posNeg=1))
|
|
else: # linear reduction
|
|
skFixupLabel = Label(label=writer.labels.getNameInc("SK_Fixup"), comment="")
|
|
|
|
# StreamK store branches
|
|
# if we're doing parallel reduction, jump to global write
|
|
# module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
# module.add(SCBranchSCC1(labelName=skStoreLabel.getLabelName(), comment="Branch if using parallel reduction, go to regular store code"))
|
|
|
|
tmpSgpr = writer.sgprPool.checkOut(4, "globalWriteElements")
|
|
# if we did not start the tile, store partials
|
|
# branch to beta == 0 store path
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
module.add(writer.longBranchScc0(skPartialsLabel, posNeg=1))
|
|
# module.add(SCBranchSCC0(labelName=skPartialsLabel.getLabelName(), comment="Branch if not start tile, store partials"))
|
|
|
|
if kernel["DebugStreamK"] & 1 == 0:
|
|
# if we started and finished the tile, regular store code
|
|
# branch to regular store code, skip fixup step
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr(sIpt), comment="does wg finish tile?"))
|
|
module.add(SCBranchSCC1(labelName=skStoreLabel.getLabelName(), comment="Branch if started and finished tile, go to regular store code"))
|
|
|
|
# if we started the tile but did not finish it, fix up step
|
|
# run fixup code before regular store code
|
|
sCtaIdx = writer.sgprPool.checkOut(1, "CtaIdx") # self.defineSgpr("CtaIdx", 1)
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(SAddU32(dst=sgpr(sCtaIdx), src0=sgpr(sIdx), src1=1, comment="input partial tile index"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
|
|
sFixupEnd = writer.sgprPool.checkOut(1, "FixupEnd") # self.defineSgpr("CtaEnd", 1)
|
|
sMagicNum = writer.acquireStreamKConstSgpr(kernel, "MagicNumberItersPerTile")
|
|
sMagicShift = writer.acquireStreamKConstSgpr(kernel, "MagicShiftItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sMagicNum), src=vgpr(writer.states.skConstVgprs["MagicNumberItersPerTile"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sMagicShift), src=vgpr(writer.states.skConstVgprs["MagicShiftItersPerTile"])))
|
|
module.add(sMagicDiv2(sgpr(tmpSgpr), sgpr(tmpSgpr+1), sgpr("StreamKIterEnd"), sgpr(sMagicNum), sgpr(sMagicShift), sgpr(tmpSgpr+2)))
|
|
writer.releaseStreamKConstSgpr(sMagicNum)
|
|
writer.releaseStreamKConstSgpr(sMagicShift)
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=sgpr(sIpt), comment="start iteration of partial tile"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SSubU32(dst=sgpr(sFixupEnd), src0=sgpr("StreamKIterEnd"), src1=sgpr(tmpSgpr), comment="calc iterations completed by this WG"))
|
|
|
|
module.add(skFixupLabel)
|
|
|
|
# Check flag
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(sCtaIdx), shiftHex=log2(4), comment="flag offset based on CTA index"))
|
|
module.add(memOrder.readFlag(writer, dst=tmpSgpr+2, soffset=sgpr(tmpSgpr)))
|
|
if kernel["DebugStreamK"] & 2 == 0:
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=1, comment="check if ready"))
|
|
module.add(SCBranchSCC0(labelName=skFixupLabel.getLabelName(), comment="if flag not set, wait and check again"))
|
|
module.add(memOrder.acquireFence(writer))
|
|
|
|
# TODO Barrier here to sync all threads in workgroup, but maybe better to have separate flag for each wavefront (to be tested)
|
|
module.add(SBarrier(comment="wait for all workgroups before resetting flag"))
|
|
skipFlagReset = Label(label=writer.labels.getNameInc("SK_SkipFlagReset"), comment="")
|
|
module.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr+2), src=vgpr("Serial"), comment="Wave 0 updates flags"))
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=0, comment="Check for wave 0"))
|
|
module.add(SCBranchSCC0(labelName=skipFlagReset.getLabelName(), comment="Skip flag reset"))
|
|
if writer.states.asmCaps["HasScalarStore"]:
|
|
# (tmpSgpr+2) contains a vlue of 0, use it to reset the flag
|
|
module.add(SStoreB32(src=sgpr(tmpSgpr+2), base=sgpr("AddressFlags", 2), soffset=sgpr(tmpSgpr), smem=SMEMModifiers(glc=True), comment="reset flag"))
|
|
else:
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr), src=0, comment="move 0 to tmpVgpr"))
|
|
module.add(self.setFlagValue(writer, src=vgpr(tmpVgpr), soffset=sgpr(tmpSgpr), comment="reset flag"))
|
|
module.add(skipFlagReset)
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
fixupEdge = [False] # Test no edge variant
|
|
# Fixup writes to workspace (no bias LDS barriers), safe to defer.
|
|
deferFixup = (
|
|
kernel.get("UseSubtileImpl")
|
|
)
|
|
if deferFixup:
|
|
fixupDeferredLabel = Label(label=writer.labels.getNameInc("Fixup_E0_Deferred"), comment="")
|
|
fixupReturnLabel = Label(label=writer.labels.getNameInc("Fixup_E0_Deferred_Return"), comment="")
|
|
# Keep original Fixup_E0 label inline as a stub
|
|
fixupInlineLabel = Label(label=writer.labels.getNameInc("Fixup_E%u" % 0), comment="")
|
|
module.add(fixupInlineLabel)
|
|
with writer.allocTmpSgpr(3, tag="StreamKOn_fixupInline_tmpSgprInfo") as tmpSgprInfo:
|
|
module.add(SLongBranchPositive(fixupDeferredLabel, tmpSgprInfo, comment="jump to deferred fixup block"))
|
|
module.addComment0("=" * 60)
|
|
module.addComment0(" Fixup block deferred to after persistent loop")
|
|
module.addComment0(" (would have been inline here in non-deferred version)")
|
|
module.addComment0("=" * 60)
|
|
module.add(fixupReturnLabel)
|
|
# Collect fixup code in deferred module
|
|
fixupModule = Module("Fixup_DeferredBlock")
|
|
fixupModule.add(fixupDeferredLabel)
|
|
fixupModule.add(self.fixupStep(writer, kernel, vectorWidths, elements, fixupEdge, tmpVgpr, cvtVgprStruct, sCtaIdx))
|
|
with writer.allocTmpSgpr(3, tag="StreamKOn_fixupDeferred_tmpSgprInfo") as tmpSgprInfo:
|
|
posLabel = writer.labels.getNameInc("FixupDeferredReturnDir")
|
|
fixupModule.add(SLongBranch(fixupReturnLabel, tmpSgprInfo, posLabel, comment="return from deferred fixup block"))
|
|
writer.states.deferredFixupModule = fixupModule
|
|
else:
|
|
fixupModule = None
|
|
module.add(self.fixupStep(writer, kernel, vectorWidths, elements, fixupEdge, tmpVgpr, cvtVgprStruct, sCtaIdx))
|
|
|
|
if kernel["StreamK"] in (3, 4, 5):
|
|
sSkExtraIters = writer.sgprPool.checkOut(1, "extraIters")
|
|
sIterCount = writer.sgprPool.checkOut(1, "iterCount")
|
|
module.add(self.skExtraIters(writer, kernel, sSkExtraIters, sIterCount)) # sIterCount is a temp register
|
|
sIpw = writer.acquireStreamKConstSgpr(kernel, "SKItersPerWG")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpw), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
module.add(SAddU32(dst=sgpr(sIterCount), src0=sgpr(sIpw), src1=1, comment="Add extra iter"))
|
|
module.add(SCmpLtU32(src0=sgpr(sCtaIdx), src1=sgpr(sSkExtraIters), comment="Check if next WG had an extra iteration"))
|
|
module.add(SCSelectB32(dst=sgpr(sIterCount), src0=sgpr(sIterCount), src1=sgpr(sIpw), comment="Select correct number of iterations for next WG"))
|
|
writer.releaseStreamKConstSgpr(sIpw)
|
|
module.add(SAddU32(dst=sgpr(sFixupEnd), src0=sgpr(sFixupEnd), src1=sgpr(sIterCount), comment="next partial tile iteration"))
|
|
writer.sgprPool.checkIn(sSkExtraIters)
|
|
writer.sgprPool.checkIn(sIterCount)
|
|
module.add(SAddU32(dst=sgpr(sCtaIdx), src0=sgpr(sCtaIdx), src1=1, comment="next partial tile index"))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SCmpLtU32(src0=sgpr(sFixupEnd), src1=sgpr(sIpt), comment="done loading partial tiles?"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SCBranchSCC1(labelName=skFixupLabel.getLabelName(), comment="Branch to continue fixup loop"))
|
|
|
|
writer.sgprPool.checkIn(sFixupEnd)
|
|
writer.sgprPool.checkIn(sCtaIdx)
|
|
|
|
module.add(skStoreLabel)
|
|
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def writePartials(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel):
|
|
pass
|
|
|
|
def writePartialsCommon(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel):
|
|
module = Module("StreamK Common writePartials")
|
|
|
|
# No partial writes for atomic or DP-only StreamK.
|
|
if kernel["StreamKAtomic"] or kernel["StreamKForceDPOnly"]:
|
|
return module
|
|
|
|
module.add(skPartialsLabel)
|
|
if kernel["DebugStreamK"] & 2 != 0:
|
|
return module
|
|
|
|
# fixupEdge = [False] # Temporary hack to test no edge variant
|
|
edges = [False]
|
|
|
|
partialsLabels = {}
|
|
for edge in edges:
|
|
partialsLabels[edge] = Label(writer.labels.getNameInc("GW_Partials_E%u" % ( 1 if edge else 0)), comment="")
|
|
|
|
if False in edges and True in edges:
|
|
with self.allocTmpSgpr(4, tag="StreamKCommon_writePartials_tmpSgprInfo") as tmpSgprInfo:
|
|
module.add(writer.checkIsEdge(kernel, tmpSgprInfo, partialsLabels[True], partialsLabels[True]))
|
|
|
|
# WritePartials writes to workspace (no bias LDS barriers), safe to defer.
|
|
deferPartials = (
|
|
kernel.get("UseSubtileImpl")
|
|
)
|
|
if deferPartials:
|
|
partialsDeferredLabel = Label(label=writer.labels.getNameInc("GW_Partials_E0_Deferred"), comment="")
|
|
partialsReturnLabel = Label(label=writer.labels.getNameInc("GW_Partials_E0_Deferred_Return"), comment="")
|
|
# Inline stub
|
|
for edge in edges:
|
|
module.add(partialsLabels[edge])
|
|
with writer.allocTmpSgpr(3, tag="StreamKCommon_writePartials_tmpSgprInfo2") as tmpSgprInfo:
|
|
module.add(SLongBranchPositive(partialsDeferredLabel, tmpSgprInfo, comment="writePartials (deferred)"))
|
|
module.addComment0("=" * 60)
|
|
module.addComment0(" WritePartials block deferred to after persistent loop")
|
|
module.addComment0(" (would have been inline here in non-deferred version)")
|
|
module.addComment0("=" * 60)
|
|
module.add(partialsReturnLabel)
|
|
module.add(SBranch(labelName=endLabel.getLabelName(), comment="jump to end"))
|
|
# Deferred block
|
|
partialsModule = Module("Partials_DeferredBlock")
|
|
partialsModule.add(partialsDeferredLabel)
|
|
for edge in edges:
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
partialsModule.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
partialsModule.add(self.computeWorkspaceSrd(writer, kernel, sgpr(sIdx)))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
partialsModule.add(self.partialsWriteProcedure(writer, kernel, vectorWidths, elements, False, False, edge, tmpVgpr, cvtVgprStruct, partialsReturnLabel))
|
|
writer.states.deferredPartialsModule = partialsModule
|
|
else:
|
|
for edge in edges:
|
|
module.add(partialsLabels[edge])
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(self.computeWorkspaceSrd(writer, kernel, sgpr(sIdx)))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
module.add(self.partialsWriteProcedure(writer, kernel, vectorWidths, elements, False, False, edge, tmpVgpr, cvtVgprStruct, endLabel))
|
|
|
|
return module
|
|
|
|
def computeWorkspaceSrd(self, writer, kernel, sPartialIdx, tmpSgpr = None):
|
|
module = Module("StreamK Common computeWorkspaceSrd")
|
|
|
|
# Base Address
|
|
module.add(SMovB64(dst=sgpr("SrdWS", 2), src=sgpr("AddressWS", 2)))
|
|
module.add(SMovB32(dst=sgpr("SrdWS+2"), src="BufferOOB"))
|
|
module.add(SMovB32(dst=sgpr("SrdWS+3"), src="Srd127_96"))
|
|
module.add(writer.shiftSrd("WS"))
|
|
|
|
tmpLocal = None
|
|
if tmpSgpr == None:
|
|
tmpLocal = writer.sgprPool.checkOut(1, "SKMappingTemp")
|
|
tmpSgpr = tmpLocal
|
|
|
|
assert kernel["BufferStore"]
|
|
module.addSpaceLine()
|
|
module.add(SMulI32(dst=sgpr(tmpSgpr), src0=hex(kernel["MacroTile0"]*kernel["MacroTile1"]*writer.states.bpeCinternal), src1=sPartialIdx, comment="Offset to correct partials tile"))
|
|
module.add(SAddU32(dst=sgpr("SrdWS+0"), src0=sgpr("SrdWS+0"), src1=sgpr(tmpSgpr), comment="add lo to SRD"))
|
|
module.add(SAddCU32(dst=sgpr("SrdWS+1"), src0=sgpr("SrdWS+1"), src1=0, comment="add hi to SRD"))
|
|
|
|
if tmpLocal is not None:
|
|
writer.sgprPool.checkIn(tmpLocal)
|
|
|
|
return module
|
|
|
|
def partialsWriteProcedure(self, writer, kernel, vectorWidths, elements, alpha, beta, edge, tmpVgpr, cvtVgprStruct, endLabel):
|
|
module = Module("StreamK Common partialsWriteProcedure")
|
|
memOrder = Component.StreamKMemoryOrdering.find(writer)
|
|
|
|
# PreLoopVmcntCaseStr = ""
|
|
# # not generate Case 2 if StoreCInUnroll with StoreVectorWidth==1 (Case 2 will be same as Case 3)
|
|
# if self.canOptimizePreLoopLWVmcnt:
|
|
# if beta:
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.OrdNLL_B1_Store
|
|
# elif edge or (kernel["StoreCInUnroll"] and kernel["StoreVectorWidth"]==1):
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.OrdNLL_E1_Store
|
|
# else:
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.OptNLL_Store
|
|
# PreLoopVmcntCaseStr = inst("s_mov_b32", sgpr("PreLoopLWVmcntCase"), hex(self.currPreLoopVmcntCase.value), \
|
|
# "for optimizing next PreLoop LW vmcnt, set to Case%u"%self.currPreLoopVmcntCase.value)
|
|
# # reset vmcnt if the dict has this key (OptNLL_Store, OrdNLL_E1_Store),
|
|
# # OrdNLL_B1_Store is excluded
|
|
# if self.currPreLoopVmcntCase in self.preLoopVmcntDict:
|
|
# self.preLoopVmcntDict[self.currPreLoopVmcntCase] = 0
|
|
|
|
edgeI = edge
|
|
#edgeI = True # set to True to disable vector stores
|
|
gwvw = vectorWidths[edgeI]
|
|
#print "globalWriteElements: edge=", edge, "beta=", beta, "atomic=", atomic
|
|
|
|
########################################
|
|
# Calculate Vgprs for Write Batching
|
|
########################################
|
|
|
|
vectorDataTypes = VectorDataTypes()
|
|
ss = StoreState(writer, kernel, gwvw, edge, beta, False, elements[edgeI], vectorDataTypes, dim=0, isWorkspace=True)
|
|
|
|
#print self.vgprPool.state()
|
|
# Use VGPR up to next occupancy threshold:
|
|
maxVgprs, _ = writer.getMaxRegsForOccupancy(kernel["NumThreads"], writer.vgprPool.size(), writer.sgprPool.size(), \
|
|
writer.getLdsSize(kernel), writer.agprPool.size(), writer.states.doubleVgpr)
|
|
if writer.states.serializedStore: # get aggressive when serializedStore is on; not necessarily exclusive to this parameter
|
|
# len(elements[edgeI])
|
|
# tl = []
|
|
# for i in range(self.vgprPool.size()-self.vgprPool.available(), maxVgprs):
|
|
# tl.append(self.vgprPool.checkOut(1, "grow-pool up to next occupancy for GlobalWrite"))
|
|
# for t in tl:
|
|
# self.vgprPool.checkIn(t)
|
|
writer.vgprPool.growPool(writer.vgprPool.size()-writer.vgprPool.available(), maxVgprs, 1, \
|
|
"grow-pool up to next occupancy for GlobalWrite")
|
|
# align = 1
|
|
# # align adjustment
|
|
# if self.ss.cfg.numVgprsPerAddr > 1:
|
|
# align = max(align, self.ss.cfg.numVgprsPerAddr)
|
|
# if self.ss.cfg.numVgprPerValuC*gwvw > 1:
|
|
# align = max(align, self.ss.cfg.numVgprPerValuC*gwvw)
|
|
# if int(ceil(self.ss.cfg.numVgprsPerDataPerVI * gwvw)) > 1:
|
|
# align = max(align, int(ceil(self.ss.cfg.numVgprsPerDataPerVI * gwvw)))
|
|
numVgprAvailable = writer.vgprPool.availableBlock(ss.numVgprsPerElement, ss.align)
|
|
|
|
# Grow the register pool if needed - we need enough regs for at least one element
|
|
# Unfortunate since this means the write logic is setting the VGPR requirement
|
|
# for the entire kernel but at least we have a functional kernel.
|
|
# Before growing the pool, see if we can shrink the write vector width instead?
|
|
# TODO : the vgprSerial is needed for-ever and if we grow here will split the
|
|
# range of the tmps. Maybe want to move vgprSerial to first vgpr?
|
|
|
|
# TODO: Minimum elems for StoreRemap
|
|
# TODO: Which of DataType or DestDataType is in a better sense? 0114: Check Using DestDataType + HSS
|
|
minElements = 1
|
|
if kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16():
|
|
minElements = 2
|
|
elif kernel["ProblemType"]["DataType"].is8bitFloat():
|
|
# TODO STREAM-K check if needed
|
|
minElements = 4
|
|
minNeeded = minElements * ss.numVgprsPerElement
|
|
|
|
shrinkDb = 0
|
|
if shrinkDb:
|
|
print("numVgprAvailable=", numVgprAvailable, "minElements=", minElements, "minNeeded=", minNeeded)
|
|
|
|
if numVgprAvailable < minNeeded:
|
|
gwvwOrig = gwvw
|
|
currentOccupancy = writer.getOccupancy(kernel["NumThreads"], writer.vgprPool.size(), \
|
|
writer.sgprPool.size(), writer.getLdsSize(kernel), writer.agprPool.size(), writer.states.doubleVgpr)
|
|
futureOccupancy = writer.getOccupancy(kernel["NumThreads"], writer.vgprPool.size() - numVgprAvailable + minNeeded, \
|
|
writer.sgprPool.size(), writer.getLdsSize(kernel), writer.agprPool.size(), writer.states.doubleVgpr)
|
|
|
|
if shrinkDb:
|
|
print("currentOccupancy=%u futureOccupancy=%u VGPRs=%u numVgprAvail=%u vgprPerElem=%u" \
|
|
% (currentOccupancy, futureOccupancy, writer.vgprPool.size(), \
|
|
numVgprAvailable, minElements*ss.numVgprsPerElement))
|
|
if futureOccupancy > currentOccupancy:
|
|
if shrinkDb:
|
|
print("warning: %s growing VGPR for GlobalWrite batching - this may bloat VGPR usage" % \
|
|
(writer.states.kernelName))
|
|
print(" numVgprAvailable=", numVgprAvailable, \
|
|
"numVgprsPerElement=", ss.numVgprsPerElement, \
|
|
"beta=", beta, "gwvw=", gwvw)
|
|
elif gwvw != gwvwOrig:
|
|
ss.gwvw = gwvw # make both representations consistent
|
|
if shrinkDb:
|
|
print2("info: %s shrank gwvw from %u to %u but kept occupancy same=%u." \
|
|
% (writer.states.kernelName, gwvwOrig, gwvw, currentOccupancy))
|
|
|
|
if numVgprAvailable < minElements*ss.numVgprsPerElement:
|
|
print2("info: growing pool += %d * %d for GlobalWrite\n" \
|
|
% (minElements,ss.numVgprsPerElement))
|
|
print2(writer.vgprPool.state())
|
|
# tl = []
|
|
# for i in range(0,minElements):
|
|
# tl.append(self.vgprPool.checkOut(numVgprsPerElement, "grow-pool for GlobalWrite"))
|
|
# for t in tl:
|
|
# self.vgprPool.checkIn(t)
|
|
writer.vgprPool.growPool(0, minElements, ss.numVgprsPerElement, \
|
|
"grow-pool for GlobalWrite")
|
|
numVgprAvailable = writer.vgprPool.available()
|
|
print2(writer.vgprPool.state())
|
|
|
|
# set atomicW after we potentially resize GWVW
|
|
# atomicW = min(gwvw, kernel["VectorAtomicWidth"])
|
|
atomicW = min(gwvw, writer.getVectorAtomicWidth(kernel))
|
|
|
|
# print("NumVgprAvailable", numVgprAvailable)
|
|
if ss.numVgprsPerElement:
|
|
numElementsPerBatch = numVgprAvailable // ss.numVgprsPerElement
|
|
else:
|
|
numElementsPerBatch = len(elements[edgeI]) # max, do 'em all
|
|
|
|
# Cap batch size to align on MIWaveTile[0] boundaries (see refineOccupancy).
|
|
if kernel.get("UseSubtileImpl") and kernel.get("EnableMatrixInstruction"):
|
|
miwt0 = kernel["MIWaveTile"][0]
|
|
totalElems = kernel["MIWaveTile"][0] * kernel["MIWaveTile"][1]
|
|
if numElementsPerBatch >= totalElems:
|
|
numElementsPerBatch = totalElems
|
|
elif miwt0 > 1 and numElementsPerBatch >= miwt0:
|
|
numElementsPerBatch = (numElementsPerBatch // miwt0) * miwt0
|
|
|
|
# assert(writer.states.numVgprValuC % gwvw == 0) # sanity check
|
|
|
|
numElementsPerBatch = numElementsPerBatch if not kernel["NumElementsPerBatchStore"] else min(kernel["NumElementsPerBatchStore"],numElementsPerBatch)
|
|
|
|
if shrinkDb:
|
|
print("NumElementsPerBatch=", numElementsPerBatch, "LimitedBySgprs=", ss.cfg.numElementsPerBatchLimitedBySgprs, \
|
|
"WARNING" if ss.cfg.numElementsPerBatchLimitedBySgprs < numElementsPerBatch else "okay")
|
|
if ss.cfg.numElementsPerBatchLimitedBySgprs < numElementsPerBatch:
|
|
numElementsPerBatch = ss.cfg.numElementsPerBatchLimitedBySgprs
|
|
|
|
# TODO: Which of DataType or DestDataType is in a better sense? 0114: Check Using DestDataType + HSS
|
|
if (kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16()):
|
|
# only do an even number of halves - since these share hi/lo pieces of some registers?
|
|
if numElementsPerBatch > 1:
|
|
numElementsPerBatch = int(numElementsPerBatch/2)*2
|
|
elif not kernel["EnableMatrixInstruction"]:
|
|
# (excluding MFMA+LSU case. It can work without an issue)
|
|
# The globalWriteBatch routine below can't handle odd elements per batch
|
|
# and 0 elements per batch is illegal.
|
|
# so if we don't have *GPR resources to handle a larger batch then need
|
|
# to mark overflowedResources rather than generate a kernel that won't work.
|
|
# It might be possible to fix globalWriteBatch to handle this case but these
|
|
# are likely to be low-performing so likely not worth optimizing.
|
|
if shrinkDb:
|
|
print("WARNING: half requires at least two elements per batch")
|
|
self.overflowedResources = 3
|
|
#elif kernel["ProblemType"]["DataType"].is8bitFloat():
|
|
# if numElementsPerBatch > 1:
|
|
# numElementsPerBatch = int(numElementsPerBatch/4)*4
|
|
|
|
assert numElementsPerBatch > 0, "numElementsPerBatch=0 for %s"%writer.states.kernelName
|
|
|
|
#numElementsPerBatch=min(2,numElementsPerBatch) # hack to control number of batches
|
|
# if atomic and (ss.optSingleColVgpr or ss.optSharedColVgpr):
|
|
# # hack to avoid re-using address vgpr across rows
|
|
# # atomics need to perform several memory operations
|
|
# # if the batch spans multiple rows, need multiple address vgpr
|
|
# # which is not currently supported in the two opt*ColVgpr modes
|
|
# firstRow = [e for e in elements[edgeI] if e[0]==0 and e[2]==0]
|
|
# numElementsPerBatch=min(len(firstRow),numElementsPerBatch)
|
|
|
|
numBatches = max(1, ceilDivide(len(elements[edgeI]),numElementsPerBatch))
|
|
|
|
numSgprs = ss.cfg.numTempSgprPerBatch + ss.cfg.numMaskSgprPerBatch + ss.cfg.numMaskSgprPerElement * numElementsPerBatch
|
|
|
|
# TODO STREAM-K activation code
|
|
|
|
if writer.db["PrintStoreRegisterDb"]:
|
|
print("edgeI", edgeI, "NumBatches", numBatches, "NumElementsPerBatch", numElementsPerBatch, "numVgprsPerElement", ss.numVgprsPerElement, "len(elements[edgeI])", len(elements[edgeI]))
|
|
print("numSgprs=", numSgprs, "sgprPool.size()=", writer.sgprPool.size(), "numTempSgprPerBatch=", ss.cfg.numTempSgprPerBatch,
|
|
"numMaskSgprPerBatch=", ss.cfg.numMaskSgprPerBatch, "numMaskSgprPerElement=", ss.cfg.numMaskSgprPerElement)
|
|
print(writer.sgprPool.state())
|
|
module.addComment1("edge=%d, allocate %u sgpr. perBatchTmpS=%u perBatchMaskS=%u perElementMaskS=%u elementsPerBatch=%u" %
|
|
(edgeI, numSgprs, ss.cfg.numTempSgprPerBatch, ss.cfg.numMaskSgprPerBatch, ss.cfg.numMaskSgprPerElement, numElementsPerBatch))
|
|
#kStr += "// storeStats, %d, %d, %d\n"% (edgeI, numSgprs, numElementsPerBatch)
|
|
# so if we don't have *GPR resources to handle a larger batch then need
|
|
# to mark overflowedResources rather than generate a kernel that won't work.
|
|
with writer.allocTmpSgpr(numSgprs, 2, tag="StreamKCommon_partialsWriteBatch_tmpSgprRes") as tmpSgprRes:
|
|
tmpSgpr = tmpSgprRes.idx
|
|
elementSgprs = tmpSgpr + ss.cfg.numTempSgprPerBatch
|
|
|
|
codeAccVgprRead = deepcopy(writer.codes.accVgprRead) if writer.states.serializedStore else None
|
|
# TODO STREAM-K remove this?
|
|
useCodeMulAlpha = kernel["MIArchVgpr"] and alpha and not (kernel["GlobalSplitU"] > 1 or kernel["GlobalSplitU"] == -1)
|
|
if useCodeMulAlpha: # do not set codeAccVgprRead=None if GSU>1
|
|
codeAccVgprRead = None
|
|
|
|
for batchIdx in range(0, numBatches):
|
|
elementStartIdx = batchIdx * numElementsPerBatch
|
|
elementStopIdx = min(elementStartIdx + numElementsPerBatch, len(elements[edgeI]))
|
|
elementsThisBatch = elements[edgeI][elementStartIdx:elementStopIdx]
|
|
#print("BATCH[%u/%u]: elements[edgeI][%u:%u] VGPRs=%u" % (batchIdx, numBatches, elementStartIdx, elementStopIdx,numVgprsPerElement ))
|
|
# elementVgprs can be large and should be perfectly tuned to the number of available
|
|
# VGPRS. We do not want to accidentally overflow and grow the pool here:
|
|
|
|
module.add(self.partialsWriteBatch(writer, kernel, ss, batchIdx, alpha, beta, edge, gwvw, atomicW, \
|
|
elementsThisBatch, writer.vgprs.addrD, writer.vgprs.addrC, \
|
|
tmpVgpr, cvtVgprStruct, \
|
|
elementSgprs, tmpSgpr, codeAccVgprRead))
|
|
# delay PreLoopVmcntCase code after globalWrite
|
|
# if self.canOptimizePreLoopLWVmcnt:
|
|
# kStr += PreLoopVmcntCaseStr
|
|
|
|
# Set flag
|
|
module.add(memOrder.releaseFence(writer))
|
|
module.add(SBarrier(comment="store all data before setting flag"))
|
|
|
|
if kernel["StreamK"] == 4:
|
|
# TODO modularize this section into abstract function
|
|
module.add(self.calculatePartialIdx(tmpSgpr))
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(tmpSgpr), shiftHex=log2(4), comment="flag offset based on partial index"))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=self._wsFlagsBaseOffset(writer, kernel), comment="Offset flags to come after the work queues"))
|
|
elif kernel["StreamK"] == 5:
|
|
# SK5 hybrid: dispatch on StreamKHybridMode bit
|
|
# (0 = static SK3 -> use StreamKIdx, 1 = dynamic SK4 -> use calculatePartialIdx).
|
|
sk5FlagStatic = Label(writer.labels.getNameInc("SK5_PartialsFlagStatic"), "")
|
|
sk5FlagDone = Label(writer.labels.getNameInc("SK5_PartialsFlagDone"), "")
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKHybridMode"), src1=0,
|
|
comment="SK5: mode bit == 0 -> SK3 (static) flag offset"))
|
|
module.add(SCBranchSCC1(labelName=sk5FlagStatic.getLabelName(),
|
|
comment="SK5: branch to static flag offset"))
|
|
# SK4 (dynamic) flag offset
|
|
module.add(self.calculatePartialIdx(tmpSgpr))
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(tmpSgpr), shiftHex=log2(4),
|
|
comment="SK5/SK4: flag offset based on partial index"))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=self._wsFlagsBaseOffset(writer, kernel),
|
|
comment="SK5/SK4: offset flags to come after the work queues"))
|
|
module.add(SBranch(labelName=sk5FlagDone.getLabelName(),
|
|
comment="SK5: skip static flag offset"))
|
|
# SK3 (static) flag offset
|
|
module.add(sk5FlagStatic)
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr("StreamKIdx"), shiftHex=log2(4),
|
|
comment="SK5/SK3: flag offset based on CTA index"))
|
|
module.add(sk5FlagDone)
|
|
else:
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(sIdx), shiftHex=log2(4), comment="flag offset based on CTA index"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
|
|
with writer.allocTmpSgpr(1, tag="StreamKCommon_setFlag_tmpSgprRes") as flagSgprRes:
|
|
flagSgpr = flagSgprRes.idx
|
|
skipFlagSet = Label(label=writer.labels.getNameInc("SK_SkipFlagSet"), comment="")
|
|
module.add(VReadfirstlaneB32(dst=sgpr(flagSgpr), src=vgpr("Serial"), comment="Wave 0 updates flags"))
|
|
module.add(SCmpEQU32(src0=sgpr(flagSgpr), src1=0, comment="Check for wave 0"))
|
|
module.add(SCBranchSCC0(labelName=skipFlagSet.getLabelName(), comment="Skip flag set"))
|
|
if writer.states.asmCaps["HasScalarStore"]:
|
|
module.add(SMovB32(dst=sgpr(flagSgpr), src=1, comment="flag data"))
|
|
module.add(SStoreB32(src=sgpr(flagSgpr), base=sgpr("AddressFlags", 2), soffset=sgpr(tmpSgpr), smem=SMEMModifiers(glc=True), comment="set flag"))
|
|
else:
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr), src=1, comment="move 1 to tmpVgpr"))
|
|
module.add(self.setFlagValue(writer, src=vgpr(tmpVgpr), soffset=sgpr(tmpSgpr), comment="set flag"))
|
|
module.add(skipFlagSet)
|
|
module.add(SWaitCnt(kmcnt=0, comment="wait for flag")) # TODO just for testing
|
|
|
|
if "Deferred" in endLabel.getLabelName():
|
|
posLabel = writer.labels.getNameInc("PartialsDeferredReturnDir")
|
|
with writer.allocTmpSgpr(3, tag="StreamKCommon_partialsDeferredReturn_tmpSgprInfo") as tmpSgprInfo:
|
|
module.add(SLongBranch(endLabel, tmpSgprInfo, posLabel, comment="jump to end"))
|
|
else:
|
|
module.add(SBranch(labelName=endLabel.getLabelName(), comment="jump to end"))
|
|
|
|
# Finish one write path, reset currPreLoopVmcntCase to Undefined
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.Undefined
|
|
|
|
return module
|
|
|
|
def setFlagValue(self, writer, src, soffset, comment=""):
|
|
module = Module("Buffer Store Flag Value")
|
|
memOrder = Component.StreamKMemoryOrdering.find(writer)
|
|
tmpSgprBuffer = writer.sgprPool.checkOutAligned(4, 4, tag="StreamKCommon_setFlagValue_tmpSgprBuffer", preventOverflow=False)
|
|
tmpVgprOff = writer.vgprPool.checkOut(1, "vaddr_off")
|
|
module.add(VMovB32(dst=vgpr(tmpVgprOff), src=0, comment="zero vaddr offset"))
|
|
module.add(SMovB64(dst=sgpr(tmpSgprBuffer, 2), src=sgpr("AddressFlags", 2)))
|
|
module.add(SMovB32(dst=sgpr(tmpSgprBuffer+2), src="BufferOOB"))
|
|
module.add(SMovB32(dst=sgpr(tmpSgprBuffer+3), src="Srd127_96"))
|
|
module.add(self.shiftSrd(writer, tmpSgprBuffer))
|
|
module.add(memOrder.preVolatileVmem(writer, comment="drain xnacks before volatile VMEM store"))
|
|
module.add(BufferStoreB32(src=src, vaddr=vgpr(tmpVgprOff), saddr=sgpr(tmpSgprBuffer, 4), soffset=soffset,
|
|
mubuf=memOrder.flagBufferMubuf(), comment=comment))
|
|
# Release the flag store: drain the store and (on dev-scope arches) global_wb
|
|
# the flag word to the L2-coherent point so a peer's acquire can observe it.
|
|
# On other targets, releaseFence is just the s_wait vscnt 0 we'd emit anyway.
|
|
module.add(memOrder.releaseFence(writer))
|
|
writer.vgprPool.checkIn(tmpVgprOff)
|
|
writer.sgprPool.checkIn(tmpSgprBuffer)
|
|
|
|
return module
|
|
|
|
def getFlagValue(self, writer, dst, soffset, comment=""):
|
|
"""Buffer-load primitive for the StreamK flag.
|
|
|
|
Used by `StreamKMemoryOrderingDevScopeFences.readFlag` to perform a
|
|
VMEM-coherent flag load. Default arches read the flag via SMEM
|
|
directly in `StreamKMemoryOrderingDefault.readFlag` and never call
|
|
this helper.
|
|
"""
|
|
module = Module("Buffer Load Flag Value")
|
|
memOrder = Component.StreamKMemoryOrdering.find(writer)
|
|
# Acquire before the read so this (and every spin re-read) sees device memory.
|
|
module.add(memOrder.acquireFence(writer))
|
|
tmpSgprBuffer = writer.sgprPool.checkOutAligned(4, 4, tag="StreamKCommon_getFlagValue_tmpSgprBuffer", preventOverflow=False)
|
|
tmpVgprOff = writer.vgprPool.checkOut(1, "vaddr_off")
|
|
module.add(VMovB32(dst=vgpr(tmpVgprOff), src=0, comment="zero vaddr offset"))
|
|
module.add(SMovB64(dst=sgpr(tmpSgprBuffer, 2), src=sgpr("AddressFlags", 2)))
|
|
module.add(SMovB32(dst=sgpr(tmpSgprBuffer+2), src="BufferOOB"))
|
|
module.add(SMovB32(dst=sgpr(tmpSgprBuffer+3), src="Srd127_96"))
|
|
module.add(self.shiftSrd(writer, tmpSgprBuffer))
|
|
module.add(memOrder.preVolatileVmem(writer, comment="drain xnacks before volatile VMEM load"))
|
|
module.add(BufferLoadB32(dst=dst, vaddr=vgpr(tmpVgprOff), saddr=sgpr(tmpSgprBuffer, 4), soffset=soffset,
|
|
mubuf=memOrder.flagBufferMubuf(),
|
|
comment=comment))
|
|
writer.vgprPool.checkIn(tmpVgprOff)
|
|
writer.sgprPool.checkIn(tmpSgprBuffer)
|
|
|
|
return module
|
|
|
|
def partialsWriteBatch(self, writer, kernel, ss, batchIdx, applyAlpha, beta, edge, gwvw, atomicW, \
|
|
batchElements, addrD, addrC, \
|
|
tmpVgpr, cvtVgprStruct, batchElementSgprs, tmpSgpr, codeAccVgprRead):
|
|
module = Module("StreamK Common partialsWriteBatch")
|
|
|
|
module.addComment0("optSingleColVgpr=%u optSharedColVgpr=%u optSGPRUsage=%s optSrdIncForRow=%u" % \
|
|
(ss.optSingleColVgpr, ss.optSharedColVgpr, ss.optSGPRUsage, ss.optSrdIncForRow))
|
|
|
|
if kernel["StoreSyncOpt"]:
|
|
module.add(SSleep(kernel["StoreSyncOpt"] - 1, "optimization: sync and wait"))
|
|
module.add(SBarrier())
|
|
|
|
# comment tt1, tt0, vc1, vc0
|
|
# tt = thread tile, vc=vector component
|
|
commentStr = "Partials Write%s%s%s Batch #%u (d1,d0,vc1,vc0) =\n " \
|
|
% (" Alpha" if applyAlpha else "", " Beta" if beta else "", " Edge" if edge else "", batchIdx)
|
|
for elementIdx in range(0, len(batchElements)):
|
|
element = batchElements[elementIdx]
|
|
commentStr += "(%u,%u,%u,%u:vw%u)" % (element[0], element[1], element[2], element[3], gwvw)
|
|
if elementIdx < len(batchElements)-1:
|
|
commentStr += "; "
|
|
module.addComment2(commentStr)
|
|
|
|
# allow expanding vgpr pool for OptNLL
|
|
# preventOverflow = (not isOptNLL)
|
|
# ss.setupStoreElementsForBatch(kernel, gwvw, batchElements, batchElementSgprs, isOptNLL=isOptNLL, isWorkspace=True)
|
|
ss.setupStoreElementsForBatch(kernel, gwvw, batchElements, batchElementSgprs, isOptNLL=False, factorDim=0, isWorkspace=True)
|
|
|
|
storesIssued = 0
|
|
tmpS01 = tmpSgpr # scratch sgprs
|
|
|
|
########################################
|
|
# calculate addr and masks
|
|
module.addComment1("calc coords, apply mask, and issue loads (if necessary)")
|
|
# On input, coord0 and coord1 are VGPRs computed in the pre-batch code, based
|
|
# on the thread and tid number. These are ELEMENT offsets from start of tensor C
|
|
# for the top-left corner this thread will write. These are not changed
|
|
# across all the store loop iters.
|
|
if writer.db["ConservativeWaitCnt"] & 0x10:
|
|
module.add(SBarrier(comment="debug"))
|
|
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
|
|
module.add(SBarrier(comment="debug"))
|
|
|
|
if not edge and writer.db["ForceEdgeStores"]>=2:
|
|
module.add(writer.getBomb()) # should not get here
|
|
if edge and writer.db["AssertNoEdge"]:
|
|
module.add(writer.getBomb()) # should not get here
|
|
|
|
## create code Module to push mov vgpr,acc instructions
|
|
# if kernel["StoreCInUnroll"] and not edge:
|
|
# accVgprRead = Code.Module("movaccVgpr")
|
|
# self.StoreCUnrollLoadCWaitComment = "waitcnt for LoadC" # this will be used later to identify waitcnt for loadC
|
|
|
|
########################################
|
|
# AccVgpr read
|
|
# if kernel.enabledSetPrioSplitLDS:
|
|
# kStr += inst("s_setprio", "0", "")
|
|
if codeAccVgprRead is not None and kernel["LocalSplitU"] == 1:
|
|
# CompactLoopStore makes accVgprRead use v_movrelsd_2_b32 (src VGPR
|
|
# index offset by M0) so the same body can cover the full thread
|
|
# tile inside a CLS loop. The StreamK partials-write path runs
|
|
# OUTSIDE that CLS loop, but reuses the same precomputed
|
|
# writer.codes.accVgprRead module -- M0 still holds whatever the
|
|
# last CLS loop left in it (sgprWorkGroup2 + step), so the source
|
|
# VGPR index gets a random offset and the partials wrote into D
|
|
# come out scrambled. Force M0=0 here so v_movrelsd_2_b32 behaves
|
|
# like the plain v_mov_b32 the non-CLS path used to emit.
|
|
if kernel.get("CompactLoopStore", False):
|
|
module.add(SMovB32(dst=mgpr(0), src=0,
|
|
comment="reset M0 for v_movrelsd_2_b32 outside CLS loop"))
|
|
regsPerScalar = writer.states.bpeCinternal // writer.states.bpr # register per scalar
|
|
# loop over store instructions within one batch
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# loop over scalars within one store instruction
|
|
for vi in range(0, gwvw):
|
|
# loop over registers within one scalar
|
|
for rIdx in range(0, regsPerScalar):
|
|
startVgprValuOffset = 0 if kernel.get("UseSubtileImpl") else writer.states.c.startVgprValu
|
|
module.add(replaceHolder(codeAccVgprRead.popFirstItem(), ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx - startVgprValuOffset))
|
|
# if kernel["StoreCInUnroll"] and not edge:
|
|
# tempStr = tempStr.replace("__placeholder__",str(elementIdx*gwvw*regsPerScalar + regsPerScalar*vi + rIdx))
|
|
# accVgprRead.addCode(tempStr.replace("ValuC","L2GC"))
|
|
|
|
if not kernel["MIArchVgpr"]:
|
|
module.add(SNop(1, "2 wait states required before reading vgpr"))
|
|
|
|
########################################
|
|
# Not Atomic
|
|
########################################
|
|
# else:
|
|
# edge has v_cndmask so loads or stores may not issue, hard to track vmcnt:
|
|
for elementIdx in range(len(batchElements)):
|
|
for vi in range(gwvw):
|
|
sumIdxV = ss.elementSumIdx[elementIdx] + vi
|
|
# TODO STREAM-K is start value needed now?
|
|
# TODO KUPO!!!!!!!!!!!!!!!!
|
|
# newSumIdxV = sumIdxV - writer.states.c.startVgprValu
|
|
# covers sgemm, gemm_ex(HHS/HSS/BBS/BSS (HPA=T)), int8 (int8x4?)
|
|
if kernel["ProblemType"]["ComputeDataType"].isInt32() or kernel["ProblemType"]["ComputeDataType"].isSingle():
|
|
if writer.db["ForceExpectedValue"]:
|
|
module.add(VMovB32(dst=vgpr("ValuC+%u"%sumIdxV), src=writer.db["ValueCExpectedValue"], comment="force expected value"))
|
|
# module.add(VMovB32(dst=vgpr("ValuC+%u"%newSumIdxV), src=self.debugConfig["ValueCExpectedValue"], comment="force expected value" ))
|
|
if writer.db["ForceVSerial"]:
|
|
module.add(VMovB32(dst=vgpr("ValuC+%u"%sumIdxV), src=vgpr("Serial"), comment="force expected value to serial"))
|
|
# module.add(VMovB32(dst=vgpr("ValuC+%u"%newSumIdxV), src=vgpr("Serial"), comment="force expected value to serial" ))
|
|
if writer.db["CheckValueC"]:
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=writer.db["ValueCExpectedValue"], comment="Move expected value"))
|
|
module.add(writer.getCmpAssert(writer.asmAssert.eq, vgpr("ValuC+%u"%sumIdxV), sgpr(tmpS01)))
|
|
|
|
module.addComment1("apply mask, calc new C and issue writes")
|
|
|
|
# if kernel["ProblemType"]["DestDataType"].isBFloat16() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
# vgprBf16Temp = tmpCVTVgpr
|
|
# vgprBf16Mask = vgprBf16Temp + 1
|
|
# vgprFp32Nan = vgprBf16Temp + 2
|
|
# vgprBf16Inc = vgprBf16Temp + 3
|
|
# kStr += inst("v_mov_b32", vgpr(vgprBf16Mask), "0xffff0000", comment="mask for pack two bfloat16 element to 32bit" )
|
|
# kStr += inst("v_mov_b32", vgpr(vgprFp32Nan), "0x7fff0000", comment="fp32 Nan" )
|
|
# kStr += inst("v_mov_b32", vgpr(vgprBf16Inc), "0x7fff", comment="rounding bias for bfloat16" )
|
|
if kernel["ProblemType"]["DestDataType"].isBFloat16() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBf16Mask), "0xffff0000", comment="mask for pack two bfloat16 element to 32bit" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp32Nan), "0x7fff0000", comment="fp32 Nan" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBf16Inc), "0x7fff", comment="rounding bias for bfloat16" ))
|
|
elif kernel["ProblemType"]["DestDataType"].isFloat8_fnuz() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8NanInf), "0x207", comment="Nan and +/- inf" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Max), "0x43700000", comment="Fp8 Max value 240 as float32" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Min), "0xc3700000", comment="Fp8 Min value -240 as float32" ))
|
|
elif kernel["ProblemType"]["DestDataType"].isFloat8() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8NanInf), "0x207", comment="Nan and +/- inf" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Max), "0x43E00000", comment="Fp8 Max value 448 as float32" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Min), "0xc3E00000", comment="Fp8 Min value -448 as float32" ))
|
|
elif kernel["ProblemType"]["DestDataType"].isAnyBFloat8() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8NanInf), "0x207", comment="Nan and +/- inf" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8Max), "0x47600000", comment="BF8 Max value 57344 as float32" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8Min), "0xc7600000", comment="BF8 Min value -57344 as float32" ))
|
|
|
|
if kernel["EnableMatrixInstruction"]:
|
|
WaveNum = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] * kernel["WorkGroup"][2]
|
|
else:
|
|
WaveNum = kernel["NumThreads"] // kernel["WavefrontSize"]
|
|
|
|
storeCode = Module("Partials GroupLoadStore")
|
|
for elementIdx in range(len(batchElements)):
|
|
element = batchElements[elementIdx]
|
|
addrCalc: AddrCalculation = ss.elementAddr[elementIdx]
|
|
addr = addrCalc.addrDVgpr
|
|
# For UseSubtileImpl, vgprValuC is remapped; add the base offset so the
|
|
# WS store reads from the correct accumulator VGPRs. For the regular path
|
|
# (non-subtile), startVgprValu is already accounted for by the vgprValuC
|
|
# assembler macro, so no offset is needed (matches rebase behaviour).
|
|
if kernel.get("UseSubtileImpl"):
|
|
sumIdx = ss.elementSumIdx[elementIdx] + writer.states.c.startVgprValu
|
|
else:
|
|
sumIdx = ss.elementSumIdx[elementIdx]
|
|
storeWidth = gwvw # pitch must match store/load width gwvw, not StoreVectorWidth (differ on source kernels)
|
|
# storeWidth = 2
|
|
if batchIdx == 0 and elementIdx == 0:
|
|
tmpSgprRes = ContinuousRegister(idx=tmpS01, size=1)
|
|
module.add(vectorStaticMultiply(vgpr(addr), vgpr("Serial"), storeWidth * writer.states.bpeCinternal, tmpSgprRes))
|
|
# kStr += inst("v_mul_lo_u32", , "Partials buffer address")
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=0, comment="Init sgpr offset"))
|
|
else:
|
|
increment = (kernel["WavefrontSize"] * WaveNum) * storeWidth * writer.states.bpeCinternal
|
|
# module.addComment1("WavefrontSize={}, WaveNum={}, storeWidth={}, bpeC={}".format(kernel["WavefrontSize"], WaveNum, storeWidth, writer.states.bpeCinternal))
|
|
module.add(SAddU32(dst=sgpr(tmpS01), src0=sgpr(tmpS01), src1=increment, comment="Inc sgpr offset"))
|
|
|
|
# TODO StreamK need this packing code???
|
|
# if self.asmCaps["HasWMMA"] and kernel["EnableMatrixInstructionStore"] and kernel["ProblemType"]["DestDataType"].isHalf() and (not kernel["ProblemType"]["HighPrecisionAccumulate"]):
|
|
# for vi in range(0, gwvw):
|
|
# sumIdxV = ss.elementSumIdx[elementIdx] + vi
|
|
# if vi%2 == 1:
|
|
# d = ss.elementSumIdx[elementIdx] + vi//2
|
|
# kStr += inst("v_pack_b32_f16", vgpr(d), vgpr("ValuC+%u"%(sumIdxV-1)), vgpr("ValuC+%u"%sumIdxV), "Pack with neighbor" )
|
|
|
|
# if not kernel["StoreRemapVectorWidth"]:
|
|
tmpStoreCode = writer.addStore(kernel, ss, 'WS', addrCalc, sumIdx, tmpS01, edge, wsOffset=sgpr(tmpS01))
|
|
if kernel["GroupLoadStore"]:
|
|
storeCode.add(tmpStoreCode)
|
|
else:
|
|
module.add(tmpStoreCode)
|
|
storesIssued += 1
|
|
|
|
module.add(storeCode)
|
|
|
|
# return registers to pool:
|
|
lastData = -1
|
|
for elementIdx in range(0, len(batchElements)):
|
|
if not ss.sharedColDVgprs:
|
|
addrCalc: AddrCalculation = ss.elementAddr[elementIdx]
|
|
addrDVgpr = addrCalc.addrDVgpr
|
|
addrCVgpr = addrCalc.addrCVgpr
|
|
writer.vgprPool.checkIn(addrDVgpr)
|
|
if addrCVgpr != addrDVgpr:
|
|
writer.vgprPool.checkIn(addrCVgpr)
|
|
|
|
data = ss.elementData[elementIdx]
|
|
if data != 0:
|
|
if data != lastData:
|
|
writer.vgprPool.checkIn(data)
|
|
lastData = data
|
|
|
|
ss.firstBatch = False
|
|
ss.checkInTempVgprC()
|
|
|
|
if writer.states.serializedStore:
|
|
module.add(SNop(0, "1 wait state required when next inst writes vgprs held by previous dwordx4 store inst"))
|
|
|
|
# Update the store cnt to preLoopVmcntDict for Case2/3
|
|
# (No need to update for Case0:'Undefined' or Case4:'OrdNLL_B1_Store')
|
|
# TODO STREAM-K Need this?
|
|
# if self.currPreLoopVmcntCase in self.preLoopVmcntDict:
|
|
# if not self.archCaps["SeparateVscnt"]:
|
|
# self.preLoopVmcntDict[self.currPreLoopVmcntCase] += storesIssued
|
|
|
|
return module
|
|
|
|
def fixupStep(self, writer, kernel, vectorWidths, elements, edges, tmpVgpr, cvtVgprStruct, sPartialIdx):
|
|
module = Module("StreamK Common fixupStep")
|
|
|
|
fixupLabels = {}
|
|
for edge in edges:
|
|
fixupLabels[edge] = Label(writer.labels.getNameInc("Fixup_E%u" % ( 1 if edge else 0)), comment="")
|
|
|
|
# branch if Edge0 or Edge1
|
|
if False in edges and True in edges:
|
|
module.add(writer.checkIsEdge(kernel, tmpSgprInfo, fixupLabels[True], fixupLabels[True]))
|
|
|
|
# by now we either jumped to E1 or stayed at E0
|
|
for edge in edges:
|
|
# write label for batch case
|
|
module.add(fixupLabels[edge])
|
|
|
|
# PreLoopVmcntCaseStr = ""
|
|
# # not generate Case 2 if StoreCInUnroll with StoreVectorWidth==1 (Case 2 will be same as Case 3)
|
|
# if self.canOptimizePreLoopLWVmcnt:
|
|
# if edge or (kernel["StoreCInUnroll"] and kernel["StoreVectorWidth"]==1):
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.OrdNLL_E1_Store
|
|
# else:
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.OptNLL_Store
|
|
# PreLoopVmcntCaseStr = inst("s_mov_b32", sgpr("PreLoopLWVmcntCase"), hex(self.currPreLoopVmcntCase.value), \
|
|
# "for optimizing next PreLoop LW vmcnt, set to Case%u"%self.currPreLoopVmcntCase.value)
|
|
# # reset vmcnt if the dict has this key (OptNLL_Store, OrdNLL_E1_Store),
|
|
# # OrdNLL_B1_Store is excluded
|
|
# if self.currPreLoopVmcntCase in self.preLoopVmcntDict:
|
|
# self.preLoopVmcntDict[self.currPreLoopVmcntCase] = 0
|
|
|
|
edgeI = edge
|
|
#edgeI = True # set to True to disable vector stores
|
|
gwvw = vectorWidths[edgeI]
|
|
|
|
########################################
|
|
# Calculate Vgprs for Write Batching
|
|
########################################
|
|
|
|
vectorDataTypes = VectorDataTypes()
|
|
ss = StoreState(writer, kernel, gwvw, edge, True, False, elements[edgeI], vectorDataTypes, dim=0, isWorkspace=True)
|
|
|
|
# how many vgprs are needed for zero elements
|
|
# 2 for addressC in vgpr for addition - already checked out
|
|
# 2 for coord0,1 of thread - already checked out
|
|
# 2 for tmp - already checked out
|
|
|
|
# 5 = how many vgprs are needed per element (flat)
|
|
# - 2 for addr
|
|
# - 3 for GLOBAL_OFFSET_C calculation (can overlap below, therefore max)
|
|
# - if beta gwvw*rpe for new value
|
|
# - if atomic 2*rpe for old and cmp values
|
|
|
|
# print("numVgprsPerAddr=%u, numVgprsPerDataPerVI=%u, numVgprPerValuC=%u"%(self.ss.cfg.numVgprsPerAddr, self.ss.cfg.numVgprsPerDataPerVI, self.ss.cfg.numVgprPerValuC))
|
|
# numVgprsPerElement = self.ss.cfg.numVgprPerValuC*gwvw + self.ss.cfg.numVgprsPerAddr + int(ceil(self.ss.cfg.numVgprsPerDataPerVI * gwvw))
|
|
|
|
# if kernel["GroupLoadStore"] and kernel["ProblemType"]["UseBeta"]:
|
|
# numVgprsPerElement += self.ss.cfg.numVgprsPerAddr
|
|
|
|
#print self.vgprPool.state()
|
|
# Use VGPR up to next occupancy threshold:
|
|
maxVgprs, _ = writer.getMaxRegsForOccupancy(kernel["NumThreads"], writer.vgprPool.size(), writer.sgprPool.size(), \
|
|
writer.getLdsSize(kernel), writer.agprPool.size(), writer.states.doubleVgpr)
|
|
if writer.states.serializedStore: # get aggressive when serializedStore is on; not necessarily exclusive to this parameter
|
|
# len(elements[edgeI])
|
|
# tl = []
|
|
# for i in range(self.vgprPool.size()-self.vgprPool.available(), maxVgprs):
|
|
# tl.append(self.vgprPool.checkOut(1, "grow-pool up to next occupancy for GlobalWrite"))
|
|
# for t in tl:
|
|
# self.vgprPool.checkIn(t)
|
|
writer.vgprPool.growPool(writer.vgprPool.size()-writer.vgprPool.available(), maxVgprs, 1, \
|
|
"grow-pool up to next occupancy for GlobalWrite")
|
|
# align = 1
|
|
# # align adjustment
|
|
# if self.ss.cfg.numVgprsPerAddr > 1:
|
|
# align = max(align, self.ss.cfg.numVgprsPerAddr)
|
|
# if self.ss.cfg.numVgprPerValuC*gwvw > 1:
|
|
# align = max(align, self.ss.cfg.numVgprPerValuC*gwvw)
|
|
# if int(ceil(self.ss.cfg.numVgprsPerDataPerVI * gwvw)) > 1:
|
|
# align = max(align, int(ceil(self.ss.cfg.numVgprsPerDataPerVI * gwvw)))
|
|
numVgprAvailable = writer.vgprPool.availableBlock(ss.numVgprsPerElement, ss.align)
|
|
|
|
# Grow the register pool if needed - we need enough regs for at least one element
|
|
# Unfortunate since this means the write logic is setting the VGPR requirement
|
|
# for the entire kernel but at least we have a functional kernel.
|
|
# Before growing the pool, see if we can shrink the write vector width instead?
|
|
# TODO : the vgprSerial is needed for-ever and if we grow here will split the
|
|
# range of the tmps. Maybe want to move vgprSerial to first vgpr?
|
|
|
|
# TODO: Minimum elems for StoreRemap
|
|
# TODO: Which of DataType or DestDataType is in a better sense? 0114: Check Using DestDataType + HSS
|
|
minElements = 1
|
|
if kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16():
|
|
minElements = 2
|
|
elif kernel["ProblemType"]["DataType"].is8bitFloat():
|
|
minElements = 4
|
|
minNeeded = minElements * ss.numVgprsPerElement
|
|
|
|
shrinkDb = 0
|
|
if shrinkDb:
|
|
print("numVgprAvailable=", numVgprAvailable, "minElements=", minElements, "minNeeded=", minNeeded)
|
|
|
|
if numVgprAvailable < minNeeded:
|
|
gwvwOrig = gwvw
|
|
currentOccupancy = writer.getOccupancy(kernel["NumThreads"], writer.vgprPool.size(), \
|
|
writer.sgprPool.size(), writer.getLdsSize(kernel), writer.agprPool.size(), writer.states.doubleVgpr)
|
|
futureOccupancy = writer.getOccupancy(kernel["NumThreads"], writer.vgprPool.size() - numVgprAvailable + minNeeded, \
|
|
writer.sgprPool.size(), writer.getLdsSize(kernel), writer.agprPool.size(), writer.states.doubleVgpr)
|
|
|
|
if shrinkDb:
|
|
print("currentOccupancy=%u futureOccupancy=%u VGPRs=%u numVgprAvail=%u vgprPerElem=%u" \
|
|
% (currentOccupancy, futureOccupancy, writer.vgprPool.size(), \
|
|
numVgprAvailable, minElements*ss.numVgprsPerElement))
|
|
if futureOccupancy > currentOccupancy:
|
|
if shrinkDb:
|
|
print("warning: %s growing VGPR for GlobalWrite batching - this may bloat VGPR usage" % \
|
|
(writer.states.kernelName))
|
|
print(" numVgprAvailable=", numVgprAvailable, \
|
|
"numVgprsPerElement=", ss.numVgprsPerElement, \
|
|
"gwvw=", gwvw)
|
|
elif gwvw != gwvwOrig:
|
|
ss.gwvw = gwvw # make both representations consistent
|
|
if shrinkDb:
|
|
print2("info: %s shrank gwvw from %u to %u but kept occupancy same=%u." \
|
|
% (writer.states.kernelName, gwvwOrig, gwvw, currentOccupancy))
|
|
|
|
if numVgprAvailable < minElements*ss.numVgprsPerElement:
|
|
print2("info: growing pool += %d * %d for GlobalWrite\n" \
|
|
% (minElements,ss.numVgprsPerElement))
|
|
print2(writer.vgprPool.state())
|
|
# tl = []
|
|
# for i in range(0,minElements):
|
|
# tl.append(self.vgprPool.checkOut(numVgprsPerElement, "grow-pool for GlobalWrite"))
|
|
# for t in tl:
|
|
# self.vgprPool.checkIn(t)
|
|
writer.vgprPool.growPool(0, minElements, ss.numVgprsPerElement, \
|
|
"grow-pool for GlobalWrite")
|
|
numVgprAvailable = writer.vgprPool.available()
|
|
print2(writer.vgprPool.state())
|
|
|
|
# print("NumVgprAvailable", numVgprAvailable)
|
|
if ss.numVgprsPerElement:
|
|
numElementsPerBatch = numVgprAvailable // ss.numVgprsPerElement
|
|
else:
|
|
numElementsPerBatch = len(elements[edgeI]) # max, do 'em all
|
|
|
|
# assert(self.numVgprValuC % gwvw == 0) # sanity check
|
|
|
|
numElementsPerBatch = numElementsPerBatch if not kernel["NumElementsPerBatchStore"] else min(kernel["NumElementsPerBatchStore"],numElementsPerBatch)
|
|
|
|
if shrinkDb:
|
|
print("NumElementsPerBatch=", numElementsPerBatch, "LimitedBySgprs=", ss.cfg.numElementsPerBatchLimitedBySgprs, \
|
|
"WARNING" if ss.cfg.numElementsPerBatchLimitedBySgprs < numElementsPerBatch else "okay")
|
|
if ss.cfg.numElementsPerBatchLimitedBySgprs < numElementsPerBatch:
|
|
numElementsPerBatch = ss.cfg.numElementsPerBatchLimitedBySgprs
|
|
|
|
# TODO: Which of DataType or DestDataType is in a better sense? 0114: Check Using DestDataType + HSS
|
|
if (kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16()):
|
|
# only do an even number of halves - since these share hi/lo pieces of some registers?
|
|
if numElementsPerBatch > 1:
|
|
numElementsPerBatch = int(numElementsPerBatch/2)*2
|
|
elif not kernel["EnableMatrixInstruction"]:
|
|
# (excluding MFMA+LSU case. It can work without an issue)
|
|
# The globalWriteBatch routine below can't handle odd elements per batch
|
|
# and 0 elements per batch is illegal.
|
|
# so if we don't have *GPR resources to handle a larger batch then need
|
|
# to mark overflowedResources rather than generate a kernel that won't work.
|
|
# It might be possible to fix globalWriteBatch to handle this case but these
|
|
# are likely to be low-performing so likely not worth optimizing.
|
|
if shrinkDb:
|
|
print("WARNING: half requires at least two elements per batch")
|
|
self.overflowedResources = 3
|
|
#elif kernel["ProblemType"]["DataType"].is8bitFloat():
|
|
# if numElementsPerBatch > 1:
|
|
# numElementsPerBatch = int(numElementsPerBatch/4)*4
|
|
|
|
assert numElementsPerBatch > 0, "numElementsPerBatch=0 for %s"%writer.states.kernelName
|
|
|
|
# if no atomics and no edge, then write whole vectors
|
|
# ERROR commented out in globalWriteELements, causes numVectorsPerBatch to not be int
|
|
# if not edge: # not atomic and
|
|
# numVectorsPerBatch = numElementsPerBatch / kernel["GlobalWriteVectorWidth"]
|
|
# #print " NumVectorsPerBatch", numVectorsPerBatch
|
|
# numElementsPerBatch = numVectorsPerBatch * kernel["GlobalWriteVectorWidth"]
|
|
numBatches = max(1, ceilDivide(len(elements[edgeI]),numElementsPerBatch))
|
|
|
|
numSgprs = ss.cfg.numTempSgprPerBatch + ss.cfg.numMaskSgprPerBatch + ss.cfg.numMaskSgprPerElement * numElementsPerBatch
|
|
|
|
if writer.db["PrintStoreRegisterDb"]:
|
|
print("edgeI", edgeI, "NumBatches", numBatches, "NumElementsPerBatch", numElementsPerBatch, "numVgprsPerElement", ss.numVgprsPerElement, "len(elements[edgeI])", len(elements[edgeI]))
|
|
print ("numSgprs=", numSgprs, "sgprPool.size()=", writer.sgprPool.size(), "numTempSgprPerBatch=", ss.cfg.numTempSgprPerBatch,
|
|
"numMaskSgprPerBatch=", ss.cfg.numMaskSgprPerBatch, "numMaskSgprPerElement=", ss.cfg.numMaskSgprPerElement)
|
|
print(writer.sgprPool.state())
|
|
module.addComment1("edge=%d, allocate %u sgpr. perBatchTmpS=%u perBatchMaskS=%u perElementMaskS=%u elementsPerBatch=%u" %
|
|
(edgeI, numSgprs, ss.cfg.numTempSgprPerBatch, ss.cfg.numMaskSgprPerBatch, ss.cfg.numMaskSgprPerElement, numElementsPerBatch))
|
|
#kStr += "// storeStats, %d, %d, %d\n"% (edgeI, numSgprs, numElementsPerBatch)
|
|
# so if we don't have *GPR resources to handle a larger batch then need
|
|
# to mark overflowedResources rather than generate a kernel that won't work.
|
|
|
|
with writer.allocTmpSgpr(numSgprs, 2, tag="StreamKCommon_fixupStep_tmpSgprRes") as tmpSgprRes:
|
|
tmpSgpr = tmpSgprRes.idx
|
|
elementSgprs = tmpSgpr + ss.cfg.numTempSgprPerBatch
|
|
|
|
codeAccVgprRead = deepcopy(writer.codes.accVgprRead) if writer.states.serializedStore else None
|
|
# codeAccVgprRead = deepcopy(writer.codes.codeAccVgprRead) if writer.states.serializedStore else None
|
|
codeAccVgprWrite = deepcopy(writer.codes.accVgprWrite) if writer.states.serializedStore else None
|
|
|
|
module.add(self.computeWorkspaceSrd(writer, kernel, sgpr(sPartialIdx), tmpSgpr))
|
|
|
|
for batchIdx in range(0, numBatches):
|
|
elementStartIdx = batchIdx * numElementsPerBatch
|
|
elementStopIdx = min(elementStartIdx + numElementsPerBatch, len(elements[edgeI]))
|
|
elementsThisBatch = elements[edgeI][elementStartIdx:elementStopIdx]
|
|
#print("BATCH[%u/%u]: elements[edgeI][%u:%u] VGPRs=%u" % (batchIdx, numBatches, elementStartIdx, elementStopIdx,numVgprsPerElement ))
|
|
# elementVgprs can be large and should be perfectly tuned to the number of available
|
|
# VGPRS. We do not want to accidentally overflow and grow the pool here:
|
|
|
|
module.add(self.fixupBatch(writer, kernel, ss, batchIdx, edge, gwvw, \
|
|
elementsThisBatch, writer.vgprs.addrD, writer.vgprs.addrC, \
|
|
tmpVgpr, cvtVgprStruct, \
|
|
elementSgprs, tmpSgpr, codeAccVgprRead, codeAccVgprWrite,
|
|
elementStartIdx))
|
|
# delay PreLoopVmcntCase code after globalWrite
|
|
# if self.canOptimizePreLoopLWVmcnt:
|
|
# kStr += PreLoopVmcntCaseStr
|
|
|
|
# Finish one write path, reset currPreLoopVmcntCase to Undefined
|
|
# self.currPreLoopVmcntCase = PreLoopVmcntCase.Undefined
|
|
|
|
# kStr += inst("s_branch", skStoreLabel, "jump to store")
|
|
|
|
return module
|
|
|
|
def fixupBatch(self, writer, kernel, ss, batchIdx, edge, gwvw, \
|
|
batchElements, addrD, addrC, \
|
|
tmpVgpr, cvtVgprStruct, batchElementSgprs, tmpSgpr, codeAccVgprRead, codeAccVgprWrite,
|
|
elementStartIdx=0):
|
|
module = Module("StreamK Common fixupBatch")
|
|
|
|
module.addComment0("optSingleColVgpr=%u optSharedColVgpr=%u optSGPRUsage=%s optSrdIncForRow=%u" % \
|
|
(ss.optSingleColVgpr, ss.optSharedColVgpr, ss.optSGPRUsage, ss.optSrdIncForRow))
|
|
|
|
if kernel["StoreSyncOpt"]:
|
|
module.add(SSleep(kernel["StoreSyncOpt"] - 1, "optimization: sync and wait"))
|
|
module.add(SBarrier())
|
|
|
|
# comment tt1, tt0, vc1, vc0
|
|
# tt = thread tile, vc=vector component
|
|
commentStr = "Fixup%s Batch #%u (d1,d0,vc1,vc0) =\n " \
|
|
% (" Edge" if edge else "", batchIdx)
|
|
for elementIdx in range(0, len(batchElements)):
|
|
element = batchElements[elementIdx]
|
|
commentStr += "(%u,%u,%u,%u:vw%u)" % (element[0], element[1], element[2], element[3], gwvw)
|
|
if elementIdx < len(batchElements)-1:
|
|
commentStr += "; "
|
|
module.addComment2(commentStr)
|
|
# print(self.kernelName)
|
|
# print(commentStr)
|
|
|
|
# allow expanding vgpr pool for OptNLL
|
|
# preventOverflow = True #(not isOptNLL)
|
|
# ss.setupStoreElementsForBatch(kernel, gwvw, batchElements, batchElementSgprs, preventOverflow=preventOverflow, isWorkspace=True)
|
|
ss.setupStoreElementsForBatch(kernel, gwvw, batchElements, batchElementSgprs, False, 0, True, elementStartIdx)
|
|
|
|
loadsIssued = 0
|
|
storesIssued = 0
|
|
tmpS01 = tmpSgpr # scratch sgprs
|
|
|
|
# laneSGPRC = writer.states.laneSGPRCount
|
|
# always use gwvw for buffer load C for atomic_cmpswap
|
|
# bpm = self.bpeCexternal * atomicW
|
|
# bpm = self.bpeCexternal * gwvw
|
|
# vgprLoadDW = 1*(bpm//4)
|
|
# atomic oparation width. 1 for b32, 2 for b64
|
|
# atomicOpW = (atomicW * self.bpeCexternal) // 4
|
|
# if atomicOpW > 2:
|
|
# # should not exceeding 2.
|
|
# atomicOpW = 2
|
|
|
|
########################################
|
|
# calculate addr and masks
|
|
module.addComment1("calc coords, apply mask, and issue loads (if necessary)")
|
|
# On input, coord0 and coord1 are VGPRs computed in the pre-batch code, based
|
|
# on the thread and tid number. These are ELEMENT offsets from start of tensor C
|
|
# for the top-left corner this thread will write. These are not changed
|
|
# across all the store loop iters.
|
|
if writer.db["ConservativeWaitCnt"] & 0x10:
|
|
module.add(SBarrier(comment="debug"))
|
|
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
|
|
module.add(SBarrier(comment="debug"))
|
|
|
|
if not edge and writer.db["ForceEdgeStores"]>=2:
|
|
module.add(writer.getBomb()) # should not get here
|
|
if edge and writer.db["AssertNoEdge"]:
|
|
module.add(writer.getBomb()) # should not get here
|
|
|
|
# atomicAddC = kernel["AtomicAddC"] and not edge
|
|
|
|
## create code Module to push mov vgpr,acc instructions
|
|
# if kernel["StoreCInUnroll"] and not edge:
|
|
# accVgprRead = Code.Module("movaccVgpr")
|
|
# self.StoreCUnrollLoadCWaitComment = "waitcnt for LoadC" # this will be used later to identify waitcnt for loadC
|
|
|
|
if kernel["EnableMatrixInstruction"]:
|
|
WaveNum = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] * kernel["WorkGroup"][2]
|
|
else:
|
|
WaveNum = kernel["NumThreads"] // kernel["WavefrontSize"]
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
element = batchElements[elementIdx]
|
|
addrCVgpr = ss.elementAddr[elementIdx].addrCVgpr
|
|
# addrDVgpr = ss.elementAddr[elementIdx].addrDVgpr
|
|
addrCalc = ss.elementAddr[elementIdx]
|
|
data = ss.elementData[elementIdx]
|
|
# mask = ss.elementMask[elementIdx]
|
|
# sumIdx = ss.elementSumIdx[elementIdx]
|
|
# d1 = element[0]
|
|
# d0 = element[1]
|
|
# vc1 = element[2]
|
|
vc0 = element[3]
|
|
|
|
storeWidth = gwvw # pitch must match store/load width gwvw, not StoreVectorWidth (differ on source kernels)
|
|
# storeWidth = 2
|
|
if batchIdx == 0 and elementIdx == 0:
|
|
tmpS01Res = ContinuousRegister(idx=tmpS01, size=1)
|
|
module.add(vectorStaticMultiply(vgpr(addrCVgpr), vgpr("Serial"), storeWidth * writer.states.bpeCinternal, tmpS01Res))
|
|
# kStr += inst("v_mul_lo_u32", , "Partials buffer address")
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=0, comment="Init sgpr offset"))
|
|
else:
|
|
increment = (kernel["WavefrontSize"] * WaveNum) * storeWidth * writer.states.bpeCinternal
|
|
# module.addComment1("WavefrontSize={}, WaveNum={}, storeWidth={}, bpeC={}".format(kernel["WavefrontSize"], WaveNum, storeWidth, writer.states.bpeCinternal))
|
|
module.add(SAddU32(dst=sgpr(tmpS01), src0=sgpr(tmpS01), src1=increment, comment="Inc sgpr offset"))
|
|
|
|
module.add(writer.readInput(kernel, ss, 'WS', kernel["ProblemType"]["ComputeDataType"], addrCalc, vc0, data, gwvw, addrCVgpr, sgpr(tmpS01)))
|
|
loadsIssued += 1
|
|
|
|
########################################
|
|
# AccVgpr read
|
|
# if kernel.enabledSetPrioSplitLDS:
|
|
# kStr += inst("s_setprio", "0", "")
|
|
if codeAccVgprRead is not None and kernel["LocalSplitU"] == 1:
|
|
# Same M0 reset as partialsWriteBatch: the SK fixup path runs after
|
|
# the CLS loop has left M0 = sgprWorkGroup2+step. accVgprRead is the
|
|
# precomputed v_movrelsd_2_b32 module (when CompactLoopStore=True),
|
|
# which would pick up that stale M0 and reorder accumulator vregs.
|
|
if kernel.get("CompactLoopStore", False):
|
|
module.add(SMovB32(dst=mgpr(0), src=0,
|
|
comment="reset M0 for v_movrelsd_2_b32 outside CLS loop"))
|
|
regsPerScalar = writer.states.bpeCinternal // writer.states.bpr # register per scalar
|
|
# loop over store instructions within one batch
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# loop over scalars within one store instruction
|
|
for vi in range(0, gwvw):
|
|
# loop over registers within one scalar
|
|
for rIdx in range(0, regsPerScalar):
|
|
module.add(replaceHolder(codeAccVgprRead.popFirstItem(), ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx - writer.states.c.startVgprValu))
|
|
# tempStr = str(codeAccVgprRead.popFirstItem())
|
|
# kStr += tempStr.replace("__placeholder__", str(ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx))
|
|
# if kernel["StoreCInUnroll"] and not edge:
|
|
# tempStr = tempStr.replace("__placeholder__",str(elementIdx*gwvw*regsPerScalar + regsPerScalar*vi + rIdx))
|
|
# accVgprRead.addCode(tempStr.replace("ValuC","L2GC"))
|
|
|
|
if not kernel["MIArchVgpr"]:
|
|
module.add(SNop(1, "2 wait states required before reading vgpr"))
|
|
|
|
########################################
|
|
# Not Atomic
|
|
########################################
|
|
# edge has v_cndmask so loads or stores may not issue, hard to track vmcnt:
|
|
interleaveStoreVmcnt = writer.states.interleaveStoreVmcnt and not edge
|
|
for elementIdx in range(0, len(batchElements)):
|
|
for vi in range(0, gwvw):
|
|
sumIdxV = ss.elementSumIdx[elementIdx] + vi
|
|
# covers sgemm, gemm_ex(HHS/HSS/BBS/BSS (HPA=T)), int8 (int8x4?)
|
|
if kernel["ProblemType"]["ComputeDataType"].isInt32() or kernel["ProblemType"]["ComputeDataType"].isSingle():
|
|
if writer.db["ForceExpectedValue"]:
|
|
module.add(VMovB32(dst=vgpr("ValuC+%u"%sumIdxV), src=writer.db["ValueCExpectedValue"], comment="force expected value"))
|
|
if writer.db["ForceVSerial"]:
|
|
module.add(VMovB32(dst=vgpr("ValuC+%u"%sumIdxV), src=vgpr("Serial"), comment="force expected value to serial"))
|
|
if writer.db["CheckValueC"]:
|
|
module.add(SMovB32(dst=sgpr(tmpS01), src=writer.db["ValueCExpectedValue"], comment="Move expected value"))
|
|
module.add(writer.getCmpAssert(writer.asmAssert.eq, vgpr("ValuC+%u"%sumIdxV), sgpr(tmpS01)))
|
|
|
|
########################################
|
|
# wait for batched load
|
|
if not interleaveStoreVmcnt: # beta and
|
|
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="wait C"))
|
|
|
|
# PreLoop LWVmcnt: When a vmcnt(cnt) is inserted here, means the GlobalLoad for PAP is finished
|
|
# So the preLoopVmcntDict value is meaningless since we no longer need to wait in next PreLoop
|
|
# And this only occurs when beta=true, so case must not be 2 or 3
|
|
# assert self.currPreLoopVmcntCase not in self.preLoopVmcntDict, \
|
|
# "PreLoopVmcntCase 2 or 3 shouldn't enter the beta true case"
|
|
|
|
module.addComment1("apply mask, calc new C and issue writes")
|
|
|
|
# if kernel["ProblemType"]["DestDataType"].isBFloat16() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
# vgprBf16Temp = tmpCVTVgpr
|
|
# vgprBf16Mask = vgprBf16Temp + 1
|
|
# vgprFp32Nan = vgprBf16Temp + 2
|
|
# vgprBf16Inc = vgprBf16Temp + 3
|
|
# kStr += inst("v_mov_b32", vgpr(vgprBf16Mask), "0xffff0000", comment="mask for pack two bfloat16 element to 32bit" )
|
|
# kStr += inst("v_mov_b32", vgpr(vgprFp32Nan), "0x7fff0000", comment="fp32 Nan" )
|
|
# kStr += inst("v_mov_b32", vgpr(vgprBf16Inc), "0x7fff", comment="rounding bias for bfloat16" )
|
|
if kernel["ProblemType"]["DestDataType"].isBFloat16() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBf16Mask), "0xffff0000", comment="mask for pack two bfloat16 element to 32bit" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp32Nan), "0x7fff0000", comment="fp32 Nan" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBf16Inc), "0x7fff", comment="rounding bias for bfloat16" ))
|
|
elif kernel["ProblemType"]["DestDataType"].isFloat8() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8NanInf), "0x207", comment="Nan and +/- inf" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Max), "0x43E00000", comment="OCP Fp8 Max value 448 as float32" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Min), "0xc3E00000", comment="OCP Fp8 Min value -448 as float32" ))
|
|
elif kernel["ProblemType"]["DestDataType"].isFloat8_fnuz() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8NanInf), "0x207", comment="Nan and +/- inf" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Max), "0x43700000", comment="Fp8 Max value 240 as float32" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprFp8Min), "0xc3700000", comment="Fp8 Min value -240 as float32" ))
|
|
elif kernel["ProblemType"]["DestDataType"].isAnyBFloat8() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8NanInf), "0x207", comment="Nan and +/- inf" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8Max), "0x47600000", comment="BF8 Max value 57344 as float32" ))
|
|
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8Min), "0xc7600000", comment="BF8 Min value -57344 as float32" ))
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
element = batchElements[elementIdx]
|
|
addr = ss.elementAddr[elementIdx].addrDVgpr
|
|
mask = ss.elementMask[elementIdx]
|
|
addrCalc = ss.elementAddr[elementIdx]
|
|
# d1 = element[0]
|
|
# d0 = element[1]
|
|
# vc1 = element[2]
|
|
vc0 = element[3]
|
|
sumIdx = ss.elementSumIdx[elementIdx]
|
|
|
|
# apply in-bounds exec mask
|
|
if edge and not kernel["BufferStore"]:
|
|
module.add(writer.getEdgeMovInstType()(EXEC(), sgpr(mask, writer.states.laneSGPRC), "sgprs -> exec"))
|
|
# kStr += inst("s_mov_b{}".format(wavelen), self.exec, sgpr(mask,laneSGPRC), "sgprs -> exec" )
|
|
|
|
# if beta:
|
|
# if GWVW=1 the half path still assumes we have
|
|
# at least two stores so does some combining across VI -
|
|
# for example assuming we can have two elements and can use pk_mul
|
|
# here:
|
|
if interleaveStoreVmcnt: # beta and
|
|
vlcnt = loadsIssued - elementIdx - 1
|
|
# we are waiting for loads to finish, so no need to wait for stores if counted separately
|
|
if writer.states.asmCaps["SeparateVscnt"] or writer.states.asmCaps["SeparateVMcnt"]:
|
|
vscnt = -1
|
|
vmComment = "{} = {} - {} - 1".format(vlcnt, loadsIssued, elementIdx)
|
|
else:
|
|
waitStoreCnt = storesIssued if not kernel["GroupLoadStore"] else 0
|
|
vscnt = waitStoreCnt
|
|
vmComment = "{} = {} - {} + {} - 1".format(vlcnt, loadsIssued, elementIdx, waitStoreCnt)
|
|
|
|
#print "wmvcnt=", vmcnt
|
|
module.addSpaceLine()
|
|
# if not atomicAddC:
|
|
module.add(SWaitCnt(vlcnt=vlcnt, vscnt=vscnt, comment="wait C (interleaved) {}".format(vmComment)))
|
|
|
|
# PreLoop LWVmcnt: When a vmcnt(cnt) is inserted here, means the GlobalLoad for PAP is finished
|
|
# So the preLoopVmcntDict value is meaningless since we no longer need to wait in next PreLoop
|
|
# And this only occurs when beta=true, so case must not be 2 or 3
|
|
# assert self.currPreLoopVmcntCase not in self.preLoopVmcntDict, "PreLoopVmcntCase 2 or 3 shouldn't enter the beta true case"
|
|
|
|
for vi in range(0, gwvw):
|
|
dataV = ss.elementData[elementIdx] + int(vi*ss.cfg.numVgprsPerDataPerVI)
|
|
sumIdxV = ss.elementSumIdx[elementIdx] + vi
|
|
if kernel["ProblemType"]["ComputeDataType"].isHalf():
|
|
if not kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
if writer.states.asmCaps["HasWMMA"] and kernel["EnableMatrixInstructionStore"]:
|
|
dataV = ss.elementData[elementIdx] + int(vi / 2 * ss.cfg.numVgprsPerDataPerVI)
|
|
# if (vi % 2) == 0:
|
|
# kStr += inst("v_pk_mul_f16", vgpr(dataV), sgpr("Beta"), vgpr(dataV+0), \
|
|
# "%s = C*beta ei=%u vi=%u"%(vgpr(dataV),elementIdx, vi))
|
|
# else:
|
|
if (vi % 2) != 0:
|
|
module.add(VLShiftRightB32(dst=vgpr(dataV), shiftHex=16, src=vgpr(dataV), \
|
|
comment="shift 16bit to get next half of packed ValueC"))
|
|
# dataV+0 = new c = old c*beta + rC
|
|
module.add(VAddPKF16(dst=vgpr("ValuC+%u"%(sumIdxV)), src0=vgpr(dataV), src1=vgpr("ValuC+%u"%(sumIdxV)), \
|
|
comment="sum*alpha + C*beta"))
|
|
elif sumIdxV%2==0 or (not ss.cfg.halfDataRegPerVI and gwvw==1):
|
|
newSumIdxV = sumIdxV // 2 - writer.states.c.startVgprValu
|
|
# dataV+0 = new c = old c*beta
|
|
# kStr += inst("v_pk_mul_f16", vgpr(dataV), sgpr("Beta"), vgpr(dataV+0), \
|
|
# "%s = C*beta ei=%u vi=%u"%(vgpr(dataV),elementIdx, vi))
|
|
# dataV+0 = new c = old c*beta + rC
|
|
module.add(VAddPKF16(dst=vgpr("ValuC+%u"%(newSumIdxV)), src0=vgpr(dataV), src1=vgpr("ValuC+%u"%(newSumIdxV)), \
|
|
comment="sum*alpha + C*beta"))
|
|
else:
|
|
pass # add will have been done previously
|
|
else: # HPA
|
|
newSumIdxV = sumIdxV - writer.states.c.startVgprValu
|
|
# dataV+0 = new c = old c*beta + rC
|
|
# src0 = beta = f32 = opsel 00
|
|
# src1 = dataV = f16.lo = opsel 10 or 11 depending on even/odd
|
|
# src2 = sumIdxV = f32 = opsel 00
|
|
dataCExternal = ss.elementData[elementIdx] + vi//2
|
|
hi16 = (vi + gwvw*vc0) % 2
|
|
# TODO try to replace with add? need opsel for f16 src
|
|
# kStr += inst(self.mixinst, vgpr("ValuC+%u"%sumIdxV), sgpr("Beta"), \
|
|
# module.add(writer.states.mixinst(dst=vgpr("ValuC+%u"%newSumIdxV), src0=sgpr("Beta"), \
|
|
# src1=vgpr(dataCExternal), src2=vgpr("ValuC+%u"%newSumIdxV), \
|
|
# vop3=VOP3PModifiers(op_sel=[0,hi16,0], op_sel_hi=[0,1,0]),
|
|
# comment="//C*=beta"))
|
|
module.add(writer.states.mixinst(dst=vgpr("ValuC+%u"%newSumIdxV), src0=1, \
|
|
src1=vgpr(dataCExternal), src2=vgpr("ValuC+%u"%newSumIdxV), \
|
|
vop3=VOP3PModifiers(op_sel=[0,hi16,0], op_sel_hi=[0,1,0]),
|
|
comment="//C*=beta"))
|
|
# kStr += inst(self.mixinst, vgpr("ValuC+%u"%sumIdxV), 1, \
|
|
# vgpr(dataCExternal), vgpr("ValuC+%u"%sumIdxV), \
|
|
# "op_sel:[0,%u,0] op_sel_hi:[0,1,0]" % (hi16), \
|
|
# "//C*=beta")
|
|
|
|
elif kernel["ProblemType"]["ComputeDataType"].isBFloat16():
|
|
if kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
# dataV+0 = new c = old c*beta + rC
|
|
# src0 = beta = f32 = opsel 00
|
|
# src1 = dataV = f16.lo = opsel 10 or 11 depending on even/odd
|
|
# src2 = sumIdxV = f32 = opsel 00
|
|
dataCExternal = ss.elementData[elementIdx] + vi//2
|
|
# if (vi%2) == 1:
|
|
# kStr += inst("v_and_b32", vgpr(tmpVgpr), vgpr(dataCExternal), vgpr(vgprBf16Mask), "convert bf16 to fp32")
|
|
# else:
|
|
# kStr += inst("v_lshlrev_b32", vgpr(tmpVgpr), "16", vgpr(dataCExternal), "convert bf16 to fp32" )
|
|
module.add(VCvtBF16toFP32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), vgprMask=vgpr(cvtVgprStruct.vgprBf16Mask), vi=(vi)))
|
|
newSumIdxV = sumIdxV - writer.states.c.startVgprValu
|
|
module.add(VAddF32(dst=vgpr("ValuC+%u"%sumIdxV), src0=vgpr("ValuC+%u"%sumIdxV), src1=vgpr(tmpVgpr), comment="accum partials"))
|
|
|
|
elif kernel["ProblemType"]["ComputeDataType"].isSingle():
|
|
if kernel["ProblemType"]["DataType"].isInt8():
|
|
newSumIdxV = sumIdxV - writer.states.c.startVgprValu
|
|
module.add(VAddU32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(dataV+0), src1=vgpr("ValuC+%u"%newSumIdxV), comment="accum partials"))
|
|
else:
|
|
newSumIdxV = sumIdxV - writer.states.c.startVgprValu
|
|
module.add(VAddF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr("ValuC+%u"%newSumIdxV), src1=vgpr(dataV+0), comment="accum partials"))
|
|
|
|
elif kernel["ProblemType"]["ComputeDataType"].isInt32():
|
|
newSumIdxV = sumIdxV - writer.states.c.startVgprValu
|
|
# assume we will need to replace v_mac_f32 with v_add_u32 and s_mul_lo_i32
|
|
# v_mad_i32_i24
|
|
module.add(VAddU32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(dataV+0), src1=vgpr("ValuC+%u"%newSumIdxV), comment="accum partials"))
|
|
|
|
elif kernel["ProblemType"]["ComputeDataType"].isDouble():
|
|
newSumIdxV = sumIdxV * 2 - writer.states.c.startVgprValu
|
|
# dataV+0 = new c = old c*beta
|
|
module.add(VAddF64(dst=vgpr("ValuC+%u"%(newSumIdxV),2), src0=vgpr("ValuC+%u"%(newSumIdxV),2), src1=vgpr(dataV+0,2), comment="accum partials"))
|
|
|
|
# single precision complex
|
|
elif kernel["ProblemType"]["ComputeDataType"].isSingleComplex():
|
|
newSumIdxV = sumIdxV * 2 - writer.states.c.startVgprValu
|
|
module.add(VAddF32(dst=vgpr("ValuC+%u"%(newSumIdxV)), src0=vgpr("ValuC+%u"%(newSumIdxV)), src1=vgpr(dataV+0), comment="accum partials real"))
|
|
module.add(VAddF32(dst=vgpr("ValuC+%u"%(newSumIdxV+1)), src0=vgpr("ValuC+%u"%(newSumIdxV+1)), src1=vgpr(dataV+1), comment="accum partials imag"))
|
|
|
|
# double precision complex
|
|
elif kernel["ProblemType"]["ComputeDataType"].isDoubleComplex():
|
|
newSumIdxV = sumIdxV * 4 - writer.states.c.startVgprValu
|
|
module.add(VAddF64(dst=vgpr("ValuC+%u"%(newSumIdxV+0),2), src0=vgpr("ValuC+%u"%(newSumIdxV+0),2), src1=vgpr(dataV+0,2), comment="accum partials real"))
|
|
module.add(VAddF64(dst=vgpr("ValuC+%u"%(newSumIdxV+2),2), src0=vgpr("ValuC+%u"%(newSumIdxV+2),2), src1=vgpr(dataV+2,2), comment="accum partials imag"))
|
|
|
|
########################################
|
|
# AccVgpr write
|
|
# if kernel.enabledSetPrioSplitLDS:
|
|
# kStr += inst("s_setprio", "0", "")
|
|
if codeAccVgprWrite is not None and kernel["LocalSplitU"] == 1:
|
|
regsPerScalar = writer.states.bpeCinternal // writer.states.bpr # register per scalar
|
|
# loop over store instructions within one batch
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# loop over scalars within one store instruction
|
|
for vi in range(0, gwvw):
|
|
# loop over registers within one scalar
|
|
for rIdx in range(0, regsPerScalar):
|
|
module.add(replaceHolder(codeAccVgprWrite.popFirstItem(), ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx - writer.states.c.startVgprValu))
|
|
# tempStr = str(codeAccVgprWrite.popFirstItem())
|
|
# kStr += tempStr.replace("__placeholder__", str(ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx))
|
|
# if kernel["StoreCInUnroll"] and not edge:
|
|
# tempStr = tempStr.replace("__placeholder__",str(elementIdx*gwvw*regsPerScalar + regsPerScalar*vi + rIdx))
|
|
# accVgprRead.addCode(tempStr.replace("ValuC","L2GC"))
|
|
|
|
if not kernel["MIArchVgpr"]:
|
|
module.add(SNop(1, "2 wait states required before reading vgpr"))
|
|
|
|
# if self.db["CheckStoreC"]>=0:
|
|
# useBuffer = kernel["BufferStore"]
|
|
# # Note - CheckStoreC won't work for EDGE store cases since they load 0 for OOB, would need more sophisticated check
|
|
# # Note - TODO- CheckStoreC also won't work for StoreRemap
|
|
# kStr += inst("s_waitcnt", "vmcnt(0)", "CheckStoreC, wait for stores to complete" )
|
|
# if self.archCaps["SeparateVscnt"]:
|
|
# kStr += inst("s_waitcnt_vscnt", -2, "0", "writes")
|
|
# for elementIdx in range(0, len(batchElements)):
|
|
# addr = ss.elementAddr[elementIdx].addrDVgpr
|
|
# sumIdx = ss.elementSumIdx[elementIdx]
|
|
|
|
# bps = kernel["ProblemType"]["DestDataType"].numBytes() * gwvw
|
|
# if kernel["BufferStore"]:
|
|
# addr0 = vgpr(addr)
|
|
# addr1 = sgpr("SrdC", 4)
|
|
# else:
|
|
# addr0 = vgpr(addr,2)
|
|
# addr1 = ""
|
|
|
|
# if kernel["ProblemType"]["DestDataType"].isHalf() or kernel["ProblemType"]["DestDataType"].isBFloat16():
|
|
# if not kernel["ProblemType"]["HighPrecisionAccumulate"]:
|
|
# kStr += self.chooseGlobalRead(useBuffer, bps, sumIdx//2, \
|
|
# addr0, addr1, soffset=0, offset=0, extraFields="", dtlNoDestVgpr=False, hi16=sumIdx%2).toStr()
|
|
# else:
|
|
# kStr += self.chooseGlobalRead(useBuffer, bps, sumIdx, \
|
|
# addr0, addr1, soffset=0, offset=0, extraFields="", dtlNoDestVgpr=False, hi16=0).toStr()
|
|
# elif kernel["ProblemType"]["DestDataType"].isInt32() or kernel["ProblemType"]["DestDataType"].isSingle():
|
|
# kStr += self.chooseGlobalRead(useBuffer, bps, sumIdx, \
|
|
# addr0, addr1, soffset=0, offset=0, extraFields="", dtlNoDestVgpr=False).toStr()
|
|
# elif kernel["ProblemType"]["DestDataType"].isDouble() or kernel["ProblemType"]["DestDataType"].isSingleComplex() :
|
|
# kStr += self.chooseGlobalRead(useBuffer, bps, sumIdx*2, \
|
|
# addr0, addr1, soffset=0, offset=0, extraFields="", dtlNoDestVgpr=False).toStr()
|
|
# elif kernel["ProblemType"]["DestDataType"].isDoubleComplex():
|
|
# kStr += self.chooseGlobalRead(useBuffer, bps, sumIdx*4, \
|
|
# addr0, addr1, soffset=0, offset=0, extraFields="", dtlNoDestVgpr=False).toStr()
|
|
# kStr += inst("s_waitcnt", "vmcnt(0)", "CheckStoreC, wait for stores to complete" )
|
|
# if self.archCaps["SeparateVscnt"]:
|
|
# kStr += inst("s_waitcnt_vscnt", -2, "0", "writes")
|
|
|
|
# # Add checks for expected values:
|
|
# kStr += inst("s_mov_b32", sgpr(tmpS01), self.db["CheckStoreC"], "expected value")
|
|
# for elementIdx in range(0, len(batchElements)):
|
|
# sumIdx = ss.elementSumIdx[elementIdx]
|
|
# # Need to fix for other types:
|
|
# assert (kernel["ProblemType"]["DestDataType"].isSingle() or kernel["ProblemType"]["DestDataType"].isInt32())
|
|
# kStr += self.assert_eq(vgpr(sumIdx), sgpr(tmpS01))
|
|
|
|
|
|
if edge and (not kernel["BufferStore"]): # atomic or
|
|
# subsequent batch must start with full exec mask
|
|
# BufferStore doesn't need exec since it used buffer range checking when
|
|
# possible
|
|
module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))
|
|
|
|
if writer.db["ConservativeWaitCnt"] & 0x40:
|
|
module.add(SBarrier(comment="debug"))
|
|
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
|
|
module.add(SBarrier(comment="debug"))
|
|
|
|
########################################
|
|
# End Not Atomic
|
|
########################################
|
|
|
|
# return registers to pool:
|
|
lastData = -1
|
|
for elementIdx in range(0, len(batchElements)):
|
|
if not ss.sharedColDVgprs:
|
|
addrCalc: AddrCalculation = ss.elementAddr[elementIdx]
|
|
addrDVgpr = addrCalc.addrDVgpr
|
|
addrCVgpr = addrCalc.addrCVgpr
|
|
writer.vgprPool.checkIn(addrDVgpr)
|
|
if addrCVgpr != addrDVgpr:
|
|
writer.vgprPool.checkIn(addrCVgpr)
|
|
|
|
data = ss.elementData[elementIdx]
|
|
if data != 0:
|
|
if data != lastData:
|
|
writer.vgprPool.checkIn(data)
|
|
lastData = data
|
|
|
|
ss.firstBatch = False
|
|
ss.checkInTempVgprC()
|
|
|
|
if writer.states.serializedStore:
|
|
module.add(SNop(0, "1 wait state required when next inst writes vgprs held by previous dwordx4 store inst"))
|
|
|
|
# Update the store cnt to preLoopVmcntDict for Case2/3
|
|
# (No need to update for Case0:'Undefined' or Case4:'OrdNLL_B1_Store')
|
|
# if self.currPreLoopVmcntCase in self.preLoopVmcntDict:
|
|
# if not self.archCaps["SeparateVscnt"]:
|
|
# self.preLoopVmcntDict[self.currPreLoopVmcntCase] += storesIssued
|
|
|
|
return module
|
|
|
|
def stridedBatchOrGeneralBatch(self, stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel):
|
|
module = Module("StreamK stridedBatchOrGeneralBatch")
|
|
if kernel["ProblemType"]["SupportUserArgs"]:
|
|
module.add(SCmpEQU32(src0=sgpr("ArgType"), src1=3, comment="ArgType == 3 for General Batched GEMM"))
|
|
module.add(SCBranchSCC0(labelName=stridedBatchedGemmLoad.getLabelName()))
|
|
# Check for StreamK Kernel when ArgType == 3 (General Batched GEMM)
|
|
# AddressFlags == 0, then parallel reduction in StreamK and SrdC/D is not dereferenced as pointer array
|
|
# AddressFlags != 0, then not parallel reduction in StreamK and SrdC/D is dereferenced as pointer array
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=generalBatchedGemmLoad.getLabelName()))
|
|
return module
|
|
|
|
@abc.abstractmethod
|
|
def initializeSrdAddressFlagsCheck(self, GeneralBatchedGemmSrdInitiation):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def routeToGeneralBatchedOrStridedBatched(self, stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def kernelEnd(self, writer, kernel):
|
|
pass
|
|
|
|
class StreamKOff(StreamK):
|
|
kernel = {"StreamK": 0}
|
|
|
|
def preLoop(self, writer, kernel):
|
|
module = Module("StreamK Off openLoop")
|
|
return module
|
|
|
|
def graWorkGroup(self, writer, kernel, tPA, tPB):
|
|
module = Module("StreamK Off graWorkGroup")
|
|
return module
|
|
|
|
def prefetchAcrossPersistentSetupNextTile(self, writer, kernel, tPA, tPB, skipLroReset=False):
|
|
module = Module("StreamK Off prefetchAcrossPersistentSetupNextTile")
|
|
return module
|
|
|
|
def computeLoadSrd(self, writer, kernel, tP, sTmp):
|
|
module = Module("StreamK Off computeLoadSrd")
|
|
return module
|
|
|
|
def computeStoreSrdStart(self, writer, kernel):
|
|
module = Module("StreamK Off computeStoreSrdStart")
|
|
return module
|
|
|
|
def graAddresses(self, writer, kernel, tP, vTmp):
|
|
module = Module("StreamK Off graAddresses")
|
|
|
|
tc = tP["tensorChar"]
|
|
module.add(VMovB32(dst=vgpr(vTmp+0), src=sgpr("Address%s+0" % tc)))
|
|
module.add(VMovB32(dst=vgpr(vTmp+1), src=sgpr("Address%s+1" % tc)))
|
|
|
|
return module
|
|
|
|
def declareStaggerParms(self, writer, kernel):
|
|
module = Module("StreamK Off declareStaggerParms")
|
|
return module
|
|
|
|
def tailLoopNumIter(self, writer, kernel, loopCounter):
|
|
module = Module("StreamK Off tailLoopNumIter")
|
|
return module
|
|
|
|
def calculateLoopNumIter(self, writer, kernel, loopCounterName, loopIdx, tmpSgprInfo):
|
|
module = Module("StreamK Off calculateLoopNumIter")
|
|
|
|
quotient = loopCounterName
|
|
dividend = "SizesSum+%u" % loopIdx #sumSize = self.sumSize(kernel, loopIdx)
|
|
divisor = kernel["DepthU"]
|
|
|
|
module.add(scalarStaticDivideAndRemainder(qReg=quotient, rReg=-1, dReg=dividend, divisor=divisor, tmpSgprRes=tmpSgprInfo, doRemainder=False))
|
|
if writer.states.tailloopInNll:
|
|
maxUnit = writer.states.tailloopInNllmaxUnit
|
|
sgprSizesSum = sgpr("SizesSum+%u" % writer.states.unrollIdx)
|
|
depthU = kernel["DepthU"]
|
|
tmpSgpr = tmpSgprInfo.idx
|
|
assert (depthU > 1 and (depthU & (depthU - 1) == 0)), "DepthU should be power of 2 with tailloopInNll"
|
|
# We need to increment loop counter by 1 if we use tailloopInNll code
|
|
# We do not use tailloopInNll code if
|
|
# summation is multiple of depthU, or
|
|
# maxUnit > 1 and summation is not multiple of maxUnit
|
|
module.add(SAndB32(dst=sgpr(tmpSgpr), src0=sgprSizesSum, src1=depthU-1, \
|
|
comment="tailloopInNll: check if summation is multiple of DepthU(%u)"%depthU))
|
|
module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=1, src1=0, \
|
|
comment="tailloopInNll: select loopcounter increment value (0 if summation is multiple of DepthU else 1"))
|
|
if maxUnit > 1:
|
|
# maxUnit > 1 case, check if summation is multiple of maxUnit or not
|
|
# if summation is multiple of maxUnit or not, tailloopInNll is used and increment loopCounter
|
|
module.add(SAndB32(dst=sgpr(tmpSgpr+1), src0=sgprSizesSum, src1=maxUnit-1, \
|
|
comment="if summation is multiple of %u, use tailloopInNll"%maxUnit))
|
|
module.add(SCSelectB32(dst=sgpr(tmpSgpr), src0=0, src1=sgpr(tmpSgpr), \
|
|
comment="select loopcounter (0 if summation is multiple of %u)"%maxUnit))
|
|
if kernel["GlobalSplitU"] != 0:
|
|
# skip tailloopInNll code if GSU>1
|
|
module.add(SAndB32(dst=sgpr(tmpSgpr+1), src0=sgpr("GSU"), src1=writer.gsuMaskHex(kernel), comment="Restore GSU"))
|
|
module.add(SCmpGtU32(src0=sgpr(tmpSgpr+1), src1=1, comment="GSU > 1 ?"))
|
|
module.add(SCMovB32(dst=sgpr(tmpSgpr), src=0, comment="do not increment loopcounter if GSU > 1"))
|
|
module.add(SAddU32(dst=sgpr(loopCounterName), src0=sgpr(loopCounterName), \
|
|
src1=sgpr(tmpSgpr), comment="increment loopcounter for tailloopInNll" ))
|
|
|
|
|
|
return module
|
|
|
|
def storeBranches(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct):
|
|
module = Module("StreamK Off storeBranches")
|
|
return module
|
|
|
|
def writePartials(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel):
|
|
module = Module("StreamK Off writePartials")
|
|
return module
|
|
|
|
def initializeSrdAddressFlagsCheck(self, GeneralBatchedGemmSrdInitiation):
|
|
module = Module("StreamK Off initializeSrdAddressFlagsCheck")
|
|
module.add(SBranch(labelName=GeneralBatchedGemmSrdInitiation.getLabelName(), comment="General Batched GEMM, Srd initialized to 0"))
|
|
return module
|
|
|
|
def routeToGeneralBatchedOrStridedBatched(self, stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel):
|
|
module = Module("StreamK Off routeToGeneralBatchedOrStridedBatched")
|
|
return module
|
|
|
|
def kernelEnd(self, writer, kernel):
|
|
module = Module("StreamK Off kernelEnd")
|
|
return module
|
|
|
|
class StreamKTwoTileDPFirst(StreamK):
|
|
kernel = {"StreamK": 3}
|
|
emitsParallelReductionSgprAliases = True
|
|
borrowsSrdWsInEpilogue = True
|
|
emitsWorkspaceReductionBpe = True
|
|
requiresWorkspaceReductionStorePath = True
|
|
supportsSubtileImpl = True
|
|
|
|
def preLoop(self, writer, kernel):
|
|
module = Module("StreamK TwoTileDPFirst openLoop")
|
|
skConstsInVgprs = writer.isStreamKConstantsToVgprEnabled(kernel)
|
|
|
|
xccMapping = Component.XCCMapping.find(writer)
|
|
module.add(xccMapping(writer, kernel))
|
|
|
|
# Skip the gfx12 ttmp reread under clustering: defineAndResources already left
|
|
# the cluster-decoded rank in WorkGroup0/1/2, and rereading ttmp9 (cluster_x here)
|
|
# would collide StreamKIdx across the cluster.
|
|
if writer.states.archCaps["WorkGroupIdFromTTM"] and not clusterEnabled(kernel["ClusterDim"]):
|
|
module.add(SMovB32(dst=sgpr("WorkGroup0"), src="ttmp9", comment="workaround"))
|
|
module.add(SAndB32(dst=sgpr("WorkGroup1"), src0=hex(0xFFFF), src1="ttmp7", comment="workaround"))
|
|
module.add(SLShiftRightB32(dst=sgpr("WorkGroup2"), shiftHex=hex(0x10), src="ttmp7", comment="workaround"))
|
|
|
|
if skConstsInVgprs:
|
|
module.add(VMovB32(dst=vgpr(self._skv(writer, "StreamKIdx")), src=sgpr("WorkGroup0"),
|
|
comment="Save original StreamK index to VGPR"))
|
|
else:
|
|
module.add(SMovB32(dst=sgpr("StreamKIdx"), src=sgpr("WorkGroup0"),
|
|
comment="Save original StreamK index"))
|
|
|
|
if kernel["StreamKForceDPOnly"]:
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(sIdx), src1=sgpr(sIpt), comment="DP starting iteration"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as sTmpRes:
|
|
sTmp = sTmpRes.idx
|
|
module.add(self.computeTotalTiles(writer, kernel, sTmp))
|
|
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(sTmp), src1=sgpr(sIpt), comment="totalIters = totalTiles * itersPerTile"))
|
|
module.add(SMovB32(dst=sgpr("StreamKIterEnd"), src=sgpr(sTmp), comment="DP ending iteration"))
|
|
module.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr(sTmp), comment="Make sure there's work to do"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
return module
|
|
|
|
# Two-tile SK (DP first)
|
|
# Do DP tiles before SK
|
|
skInitDone = Label("SK_InitDone", "")
|
|
|
|
# Choose reduction strategy
|
|
# If synchronizer buffer exists, then do single-kernel stream-k fixup step with tree reduction
|
|
# If there's no synchronizer, parallel reduction is done in a post-kernel
|
|
skSplitInit = Label("SK_SplitInit", "")
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=skSplitInit.getLabelName(), comment="Jump to single kernel init"))
|
|
|
|
################
|
|
# Parallel reduction init
|
|
################
|
|
# WGsPerTile = skTiles (would be WGsPerTile = grid / tiles)
|
|
# tile = Idx / WGsPerTile
|
|
# partialIndex = Idx % WGsPerTile
|
|
stmpTileIdx = writer.sgprPool.checkOut(1, "TileIdx")
|
|
stmpPartialIdx = writer.sgprPool.checkOut(1, "PartialIdx")
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(scalarUInt32DivideAndRemainder(qReg=stmpTileIdx, dReg=sIdx, divReg="SkSplit", rReg=stmpPartialIdx, tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=True, comment="TileIdx = SKIdx // WGsPerTile, PartialIdx = SKIdx % WGsPerTile"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
|
|
# if (partialIdx < extraIters) then (skIter = partialIdx * (itersPerWG + 1)) else (skIter = partialIdx * itersPerWG + extraIters)
|
|
skHasExtraLabel = Label("SK_HasExtra", "")
|
|
skDoneExtraLabel = Label("SK_DoneExtra", "")
|
|
|
|
# PartialIdx = itersPerTile % skSplit (skSplit is passed as SkSplit)
|
|
# extraIters = ItersPerTile - SkSplit * skItersPerWG
|
|
sSkExtraIters = writer.sgprPool.checkOut(1, "extraIters")
|
|
sIpw = writer.acquireStreamKConstSgpr(kernel, "SKItersPerWG")
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpw), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMulI32(dst=sgpr(sSkExtraIters), src0=sgpr("SkSplit"), src1=sgpr(sIpw)))
|
|
module.add(SSubU32(dst=sgpr(sSkExtraIters), src0=sgpr(sIpt), src1=sgpr(sSkExtraIters), comment="extraIters = itersPerTile - SkSplit * skItersPerWG"))
|
|
|
|
module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(stmpPartialIdx), src1=sgpr(sIpw), comment="StreamK starting iteration (case: after extra iters)"))
|
|
module.add(SCmpLtU32(src0=sgpr(stmpPartialIdx), src1=sgpr(sSkExtraIters), comment="Check if WG gets an extra iteration"))
|
|
module.add(SCBranchSCC1(labelName=skHasExtraLabel.getLabelName(), comment="Has extra iter"))
|
|
# No extra
|
|
module.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"), src1=sgpr(sSkExtraIters), comment="This WG does not have an extra iteration"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"), src1=sgpr(sIpw), comment="StreamK ending iteration (case: after extra iters)"))
|
|
module.add(SBranch(labelName=skDoneExtraLabel.getLabelName(), comment="Done init for parallel reduction"))
|
|
# Has extra
|
|
module.add(skHasExtraLabel)
|
|
module.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"), src1=sgpr(stmpPartialIdx), comment="This WG has an extra iteration"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"), src1=sgpr(sIpw), comment="StreamK ending iteration (case: after extra iters)"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=1, comment="StreamK ending iteration (case: after extra iters)"))
|
|
module.add(skDoneExtraLabel)
|
|
writer.releaseStreamKConstSgpr(sIpw)
|
|
# Offset to tile
|
|
module.add(SMulI32(dst=sgpr(stmpTileIdx), src0=sgpr(stmpTileIdx), src1=sgpr(sIpt), comment="Tile offset = tilesIdx * itersPerTile"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"), src1=sgpr(stmpTileIdx), comment="Offset to correct tile"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(stmpTileIdx), comment="Offset to correct tile"))
|
|
# Save partial idx for later
|
|
module.add(SMovB32(dst=sgpr("SkPartialIdx"), src=sgpr(stmpPartialIdx), comment="Save partial idx for SrdD calculation"))
|
|
# Done init
|
|
module.add(SBranch(labelName=skInitDone.getLabelName(), comment="Done init for parallel reduction"))
|
|
|
|
# # Save PratialIdx for later, skExtraIters is unused for partial reduction
|
|
# module.add(SMovB32(dst=sgpr("skExtraIters"), src=sgpr(stmpPartialIdx), comment="Save partial idx for SrdD calculation"))
|
|
# # StreamKIter = tile * itersPerTile + itersPerWG * partialIndex
|
|
# module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(stmpTileIdx), src1=sgpr("ItersPerTile"), comment="Tile offset = tilesIdx * itersPerTile"))
|
|
# module.add(SMulI32(dst=sgpr(stmpPartialIdx), src0=sgpr("SKItersPerWG"), src1=sgpr(stmpPartialIdx), comment="Offset within tile = itersPerWG * partialIdx"))
|
|
# module.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"), src1=sgpr(stmpPartialIdx), comment="StreamKIter = tileIdx * itersPerTile + partialIdx * itersPerWG"))
|
|
# # if itersPerWG * partialIndex > itersPerTile jump to end
|
|
# module.add(SCmpLtU32(src0=sgpr(stmpPartialIdx), src1=sgpr("ItersPerTile"), comment="Make sure there's work to do"))
|
|
# module.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
# # StreamKIterEnd = StreamKIter + itersPerWG
|
|
# module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"), src1=sgpr("SKItersPerWG"), comment="StreamKIterEnd = StreamKIter + itersPerWG"))
|
|
# # tileEnd = (tile + 1) * itersPerTile
|
|
# module.add(SAddU32(dst=sgpr(stmpTileIdx), src0=sgpr(stmpTileIdx), src1=1, comment="Find end of tile"))
|
|
# module.add(SMulI32(dst=sgpr(stmpTileIdx), src0=sgpr(stmpTileIdx), src1=sgpr("ItersPerTile"), comment="Find end of tile"))
|
|
# # StreamKIterEnd = min(StreamKIterEnd, tileEnd)
|
|
# # TODO SMin instruciton
|
|
# module.add(SCmpLtU32(src0=sgpr("StreamKIterEnd"), src1=sgpr(stmpTileIdx), comment="StreamKIterEnd = min(StreamKIterEnd, tileEnd)"))
|
|
# module.add(SCSelectB32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(stmpTileIdx), comment="Set start iter"))
|
|
# # Done init
|
|
# module.add(SBranch(labelName=skInitDone.getLabelName(), comment="Done init for parallel reduction"))
|
|
module.add(skSplitInit)
|
|
writer.sgprPool.checkIn(sSkExtraIters)
|
|
writer.sgprPool.checkIn(stmpPartialIdx)
|
|
writer.sgprPool.checkIn(stmpTileIdx)
|
|
|
|
################
|
|
# Tree reduction init
|
|
################
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(sIdx), src1=sgpr(sIpt), comment="DP starting iteration (case: DP work to do)"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as sTmpRes:
|
|
sTmp = sTmpRes.idx
|
|
module.add(self.computeTotalIters(writer, kernel, sTmp))
|
|
module.add(SMovB32(dst=sgpr("StreamKIterEnd"), src=sgpr(sTmp), comment="DP ending iteration (case: only DP work to do)"))
|
|
sSkt = writer.acquireStreamKConstSgpr(kernel, "skTiles")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(writer.states.skConstVgprs["skTiles"])))
|
|
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(sSkt), src1=sgpr(sIpt), comment="Total SK iters"))
|
|
writer.releaseStreamKConstSgpr(sSkt)
|
|
module.add(SCmpLtU32(src0=sgpr(sTmp), src1=sgpr("StreamKIterEnd"), comment="Check if there are DP tiles to do"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SCBranchSCC1(labelName=skInitDone.getLabelName(), comment="Done init"))
|
|
|
|
# If there are no DP tiles to do, regular SK init
|
|
# iter count after all extra iters have been distributed
|
|
with writer.allocTmpSgpr(1, tag="extraIters") as extraItersRes, \
|
|
writer.allocTmpSgpr(2, alignment=1, tag="SKIter") as skIterRes:
|
|
sSkExtraIters = extraItersRes.idx
|
|
sIter = skIterRes.idx
|
|
module.add(self.skExtraIters(writer, kernel, sSkExtraIters, sIter)) # sIter used as tmp
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
sIpw = writer.acquireStreamKConstSgpr(kernel, "SKItersPerWG")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpw), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(sIdx), src1=sgpr(sIpw), comment="StreamK starting iteration (case: after extra iters)"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"), src1=sgpr(sSkExtraIters), comment="Add extra iters"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"), src1=sgpr(sIpw), comment="StreamK ending iteration (case: after extra iters)"))
|
|
# iter count before all extra iters have been distributed
|
|
module.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr(sIpw), src1=1, comment="Spread out extra iterations"))
|
|
writer.releaseStreamKConstSgpr(sIpw)
|
|
module.add(SMulI32(dst=sgpr(sIter), src0=sgpr(sIdx), src1=sgpr(sIter+1), comment="StreamK starting iteration (case: before extra iters)"))
|
|
module.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr(sIter), src1=sgpr(sIter+1), comment="StreamK ending iteration (case: before extra iters)"))
|
|
# select correct start/end iteration index
|
|
module.add(SCmpLtU32(src0=sgpr(sIdx), src1=sgpr(sSkExtraIters), comment="Check if lane gets an extra iteration"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
module.add(SCSelectB32(dst=sgpr("StreamKIter"), src0=sgpr(sIter), src1=sgpr("StreamKIter"), comment="Set start iter"))
|
|
module.add(SCSelectB32(dst=sgpr("StreamKIterEnd"), src0=sgpr(sIter+1), src1=sgpr("StreamKIterEnd"), comment="Set end iter"))
|
|
sTmp = writer.sgprPool.checkOut(1, "TotalSKIters")
|
|
sSkt = writer.acquireStreamKConstSgpr(kernel, "skTiles")
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(writer.states.skConstVgprs["skTiles"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr(sSkt), src1=sgpr(sIpt), comment="Total SK iters"))
|
|
writer.releaseStreamKConstSgpr(sSkt)
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SMinU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(sTmp), comment="Cap ending iter at total SK iters"))
|
|
writer.sgprPool.checkIn(sTmp)
|
|
|
|
module.add(skInitDone)
|
|
# check if this WG has no work to do
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as sTmpRes:
|
|
sTmp = sTmpRes.idx
|
|
module.add(self.computeTotalIters(writer, kernel, sTmp))
|
|
module.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr(sTmp), comment="Make sure there's work to do"))
|
|
module.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
|
|
return module
|
|
|
|
def graWorkGroup(self, writer, kernel, tPA, tPB):
|
|
module = Module("StreamK TwoTileDPFirst graWorkGroup")
|
|
skConstsInVgprs = writer.isStreamKConstantsToVgprEnabled(kernel)
|
|
|
|
# StreamK workgroup mapping
|
|
sTmp = writer.sgprPool.checkOutAligned(4, 2, "SKMappingTemp", preventOverflow=not kernel.get("UseSubtileImpl", False))
|
|
|
|
if kernel["StreamKForceDPOnly"]:
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(self.computeTotalTiles(writer, kernel, sTmp+3))
|
|
module.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr(sIpt), comment="dpSectionSize = totalTiles * ItersPerTile"))
|
|
module.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr(sTmp+3), comment="Make sure there's DP work to do"))
|
|
module.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
|
|
module.add(self.skTileIndex(writer, kernel, sTmp, tPA, tPB))
|
|
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
sGrid = writer.acquireStreamKConstSgpr(kernel, "skGrid")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sGrid), src=vgpr(writer.states.skConstVgprs["skGrid"])))
|
|
module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr(sGrid), src1=sgpr(sIpt), comment="DP iterations shift"))
|
|
writer.releaseStreamKConstSgpr(sGrid)
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("StreamKIter"), comment="Add DP shift"))
|
|
module.add(SMovB32(dst=sgpr("StreamKIter"), src=sgpr(sTmp+1), comment="Store next DP iteration"))
|
|
|
|
module.add(self.skIndexToWG(writer, kernel, sTmp))
|
|
|
|
alphaLabel = Label(writer.labels.getNameInc("SKAlphaCheck"), "")
|
|
module.add(BranchIfNotZero("Alpha", kernel["ProblemType"]["ComputeDataType"].toEnum(), alphaLabel))
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
skCloseLoopLabel = Label("SK_CloseLoop", "")
|
|
module.add(writer.longBranchScc0(skCloseLoopLabel, posNeg=1))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr(sIpt), comment="Skip iterations"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(alphaLabel)
|
|
|
|
writer.sgprPool.checkIn(sTmp)
|
|
return module
|
|
|
|
module.add(self.skTileIndex(writer, kernel, sTmp, tPA, tPB))
|
|
|
|
skUpdateDone = Label("SK_UpdateDone", "")
|
|
|
|
# Choose reduction strategy
|
|
skSplitUpdate = Label("SK_SplitUpdate", "")
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=skSplitUpdate.getLabelName(), comment="Jump to single kernel update"))
|
|
# Parallel reduction doesn't cross tile boundaries, move to end
|
|
module.add(SMovB32(dst=sgpr(sTmp+1), src=sgpr("StreamKIterEnd"), comment="Parallel reduction, work contained to single partial tile"))
|
|
# Done update
|
|
module.add(SBranch(labelName=skUpdateDone.getLabelName(), comment="Done update for parallel reduction"))
|
|
module.add(skSplitUpdate)
|
|
|
|
module.add(self.computeTotalTiles(writer, kernel, sTmp+3))
|
|
sSkt = writer.acquireStreamKConstSgpr(kernel, "skTiles")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sSkt), src=vgpr(writer.states.skConstVgprs["skTiles"])))
|
|
module.add(SSubU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr(sSkt), comment="dpTiles = totalTiles - skTiles"))
|
|
writer.releaseStreamKConstSgpr(sSkt)
|
|
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
sGrid = writer.acquireStreamKConstSgpr(kernel, "skGrid")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sGrid), src=vgpr(writer.states.skConstVgprs["skGrid"])))
|
|
module.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr(sIpt), comment="dpSectionSize = dpTiles * ItersPerTile"))
|
|
|
|
# If in DP, add dpShift
|
|
module.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr(sGrid), src1=sgpr(sIpt), comment="DP iterations shift"))
|
|
writer.releaseStreamKConstSgpr(sGrid)
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("StreamKIter"), comment="Add DP shift"))
|
|
# if sTmp+1 < sTmp+3, continue DP (add dpShift)
|
|
module.add(SCmpLtU32(src0=sgpr(sTmp+1), src1=sgpr(sTmp+3), comment="Check if still in DP section"))
|
|
module.add(SCBranchSCC1(labelName=skUpdateDone.getLabelName(), comment="Done update"))
|
|
# if StreamKIter >= sTmp+3, continue SK (add skShift?)
|
|
module.add(SMovB32(dst=sgpr(sTmp+1), src=sgpr(sTmp+2), comment="SK iterations shift"))
|
|
module.add(SCmpLeU32(src0=sgpr(sTmp+3), src1=sgpr("StreamKIter"), comment="Check if continuing in SK section"))
|
|
module.add(SCBranchSCC1(labelName=skUpdateDone.getLabelName(), comment="Done update"))
|
|
# if sTmp+1 > sTmp+3 and StreamKIter < sTmp+3, switch from DP to SK (add dpShift)
|
|
# iter count after all extra iters have been distributed
|
|
with writer.allocTmpSgpr(1, tag="extraIters") as extraItersRes, \
|
|
writer.allocTmpSgpr(2, alignment=1, tag="SKIter") as skIterRes:
|
|
sSkExtraIters = extraItersRes.idx
|
|
sIter = skIterRes.idx
|
|
module.add(self.skExtraIters(writer, kernel, sSkExtraIters, sIter)) # sIter used as tmp
|
|
sIdx = writer.acquireStreamKConstSgpr(kernel, "StreamKIdx")
|
|
if skConstsInVgprs:
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIdx), src=vgpr(writer.states.skConstVgprs["StreamKIdx"])))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIter), src=vgpr(writer.states.skConstVgprs["SKItersPerWG"])))
|
|
sItersPerWg = sIter
|
|
else:
|
|
sItersPerWg = "SKItersPerWG"
|
|
module.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(sIdx), src1=sgpr(sItersPerWg), comment="StreamK starting iteration (case: after extra iters)"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"), src1=sgpr(sSkExtraIters), comment="Add extra iters"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"), src1=sgpr(sItersPerWg), comment="StreamK ending iteration (case: after extra iters)"))
|
|
# iter count before all extra iters have been distributed
|
|
module.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr(sItersPerWg), src1=1, comment="Spread out extra iterations"))
|
|
module.add(SMulI32(dst=sgpr(sIter), src0=sgpr(sIdx), src1=sgpr(sIter+1), comment="StreamK starting iteration (case: before extra iters)"))
|
|
module.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr(sIter), src1=sgpr(sIter+1), comment="StreamK ending iteration (case: before extra iters)"))
|
|
# select correct start/end iteration index
|
|
module.add(SCmpLtU32(src0=sgpr(sIdx), src1=sgpr(sSkExtraIters), comment="Check if lane gets an extra iteration"))
|
|
writer.releaseStreamKConstSgpr(sIdx)
|
|
module.add(SCSelectB32(dst=sgpr("StreamKIter"), src0=sgpr(sIter), src1=sgpr("StreamKIter"), comment="Set start iter"))
|
|
module.add(SCSelectB32(dst=sgpr("StreamKIterEnd"), src0=sgpr(sIter+1), src1=sgpr("StreamKIterEnd"), comment="Set end iter"))
|
|
module.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr("StreamKIter"), src1=sgpr(sTmp+3), comment="Offset to start of SK section"))
|
|
module.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(sTmp+3), comment="Offset to start of SK section"))
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as tmpTotalIters:
|
|
sTotalIters = tmpTotalIters.idx
|
|
module.add(self.computeTotalIters(writer, kernel, sTotalIters))
|
|
# TODO maybe remove clamp, since extra iters code should guarantee total iterations match
|
|
module.add(SMinU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"), src1=sgpr(sTotalIters), comment="Cap ending iter at total SK iters"))
|
|
# check if this WG has no work to do
|
|
module.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr(sTotalIters), comment="Make sure there's work to do"))
|
|
module.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
|
|
# If in SK, next iteration is sTmp+2
|
|
# Increment StreamK iteration
|
|
module.add(skUpdateDone)
|
|
module.add(SMovB32(dst=sgpr("StreamKIter"), src=sgpr(sTmp+1), comment="Store current iteration"))
|
|
|
|
# Map SK index to WG
|
|
module.add(self.skIndexToWG(writer, kernel, sTmp))
|
|
|
|
# Short circuit if alpha==0 (skip main loop and reading A/B, only do beta * C)
|
|
# To skip main loop in stream-k, we check if this WG is responsible for writing results (ie: WG starts tile)
|
|
# If WG starts tile then set LocalEnd=ItersPerTile to skip fixup step, and set loopCounter to 0 to skip main loop
|
|
# If WG does not start tile, skip to end of persistent loop to check for other SK tile
|
|
alphaLabel = Label(writer.labels.getNameInc("SKAlphaCheck"), "")
|
|
module.add(BranchIfNotZero("Alpha", kernel["ProblemType"]["ComputeDataType"].toEnum(), alphaLabel))
|
|
# Skip to end if not doing the global write
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
skCloseLoopLabel = Label("SK_CloseLoop", "")
|
|
module.add(writer.longBranchScc0(skCloseLoopLabel, posNeg=1))
|
|
sIpt = writer.acquireStreamKConstSgpr(kernel, "ItersPerTile")
|
|
if writer.isStreamKConstantsToVgprEnabled(kernel):
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sIpt), src=vgpr(writer.states.skConstVgprs["ItersPerTile"])))
|
|
module.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr(sIpt), comment="Skip iterations"))
|
|
writer.releaseStreamKConstSgpr(sIpt)
|
|
module.add(alphaLabel)
|
|
|
|
writer.sgprPool.checkIn(sTmp)
|
|
|
|
return module
|
|
|
|
def computeLoadSrd(self, writer, kernel, tP, sTmp):
|
|
module = Module("StreamK TwoTileDPFirst computeLoadSrd")
|
|
module.add(self.computeLoadSrdCommon(writer, kernel, tP, sTmp))
|
|
return module
|
|
|
|
def computeStoreSrdStart(self, writer, kernel):
|
|
module = Module("StreamK TwoTileDPFirst computeStoreSrdStart")
|
|
module.add(self.computeStoreSrdStartCommon(writer, kernel))
|
|
return module
|
|
|
|
def graAddresses(self, writer, kernel, tP, vTmp):
|
|
module = Module("StreamK TwoTileDPFirst graAddresses")
|
|
module.add(self.graAddressesCommon(writer, kernel, tP, vTmp))
|
|
return module
|
|
|
|
def declareStaggerParms(self, writer, kernel):
|
|
module = Module("StreamK TwoTileDPFirst declareStaggerParms")
|
|
module.add(self.declareStaggerParmsCommon(writer, kernel))
|
|
return module
|
|
|
|
def tailLoopNumIter(self, writer, kernel, loopCounter):
|
|
module = Module("StreamK TwoTileDPFirst tailLoopNumIter")
|
|
module.add(self.tailLoopNumIterCommon(writer, kernel, loopCounter))
|
|
return module
|
|
|
|
def calculateLoopNumIter(self, writer, kernel, loopCounterName, loopIdx, tmpSgprInfo):
|
|
module = Module("StreamK TwoTileDPFirst calculateLoopNumIter")
|
|
module.add(self.calculateLoopNumIterCommon(writer, kernel, loopCounterName, loopIdx, tmpSgprInfo))
|
|
return module
|
|
|
|
def storeBranches(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct):
|
|
module = Module("StreamK TwoTileDPFirst storeBranches")
|
|
module.add(self.storeBranchesCommon(writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct))
|
|
return module
|
|
|
|
def writePartials(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel):
|
|
module = Module("StreamK TwoTileDPFirst writePartials")
|
|
module.add(self.writePartialsCommon(writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel))
|
|
return module
|
|
|
|
def initializeSrdAddressFlagsCheck(self, GeneralBatchedGemmSrdInitiation):
|
|
module = Module("StreamK TwoTileDPFirst initializeSrdAddressFlagsCheck")
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=GeneralBatchedGemmSrdInitiation.getLabelName(), comment="Parallel Reduction for General Batched GEMM, Srd initialized to workspace"))
|
|
return module
|
|
|
|
def routeToGeneralBatchedOrStridedBatched(self, stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel):
|
|
module = Module("StreamK TwoTileDPFirst routeToGeneralBatchedOrStridedBatched")
|
|
module.add(self.stridedBatchOrGeneralBatch(stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel))
|
|
return module
|
|
|
|
def kernelEnd(self, writer, kernel):
|
|
module = Module("StreamK TwoTileDPFirst kernelEnd")
|
|
return module
|
|
|
|
class StreamKDynamic(StreamK):
|
|
kernel = {"StreamK": 4}
|
|
requiresWorkspaceReductionStorePath = True
|
|
keepsConstantsInSgpr = True
|
|
supportsSubtileImpl = True
|
|
|
|
def preLoop(self, writer, kernel):
|
|
module = Module("StreamK Dynamic openLoop")
|
|
|
|
xccMapping = Component.XCCMapping.find(writer)
|
|
module.add(xccMapping(writer, kernel))
|
|
|
|
# Skip the gfx12 ttmp reread under clustering: defineAndResources already left
|
|
# the cluster-decoded rank in WorkGroup0/1/2, and rereading ttmp9 (cluster_x here)
|
|
# would collide StreamKIdx across the cluster.
|
|
if writer.states.archCaps["WorkGroupIdFromTTM"] and not clusterEnabled(kernel["ClusterDim"]):
|
|
module.add(SMovB32(dst=sgpr("WorkGroup0"), src="ttmp9", comment="workaround"))
|
|
module.add(SAndB32(dst=sgpr("WorkGroup1"), src0=hex(0xFFFF), src1="ttmp7", comment="workaround"))
|
|
module.add(SLShiftRightB32(dst=sgpr("WorkGroup2"), shiftHex=hex(0x10), src="ttmp7", comment="workaround"))
|
|
|
|
module.add(SMovB32(dst=sgpr("StreamKIdx"), src=sgpr("WorkGroup0"), comment="Save original StreamK index"))
|
|
# Work stealing: this WG has not yet seen its home queue empty.
|
|
if kernel["StreamKWorkStealing"]:
|
|
module.add(SMovB32(dst=sgpr("StreamKStickyEmpty"), src=0, comment="WS: home not yet empty"))
|
|
# Two-tile SK (DP first)
|
|
# Do DP tiles before SK
|
|
skInitDone = Label("SK_InitDone", "")
|
|
module.add(skInitDone)
|
|
|
|
return module
|
|
|
|
def graWorkGroup(self, writer, kernel, tPA, tPB):
|
|
module = Module("StreamK Dynamic graWorkGroup")
|
|
|
|
skFullTile = Label("SK_FullTile", "")
|
|
skPartialTile = Label("SK_PartialTile", "")
|
|
skDone = Label("SK_Done", "")
|
|
|
|
# Local address for sharing work id
|
|
vLocalAddress = writer.vgprPool.checkOut(1, "LocalAddress")
|
|
# TODO reorganize to reduce waits
|
|
module.add(VMovB32(dst=vgpr(vLocalAddress), src=vgpr("Serial"), comment="Move local address to vgpr"))
|
|
module.add(VLShiftLeftB32(dst=vgpr(vLocalAddress), src=vgpr(vLocalAddress), shiftHex=log2(4), comment="Scale by BPE"))
|
|
sFirstLane = writer.sgprPool.checkOut(1, "FirstLane")
|
|
module.add(SNop(waitState=4, comment="4 wait required between VALU op and readfirstlane using the value"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sFirstLane), src=vgpr(vLocalAddress), comment="Read first lane of local address"))
|
|
module.add(SNop(waitState=2, comment="2 wait required between readfirstlane and VALU op using the value"))
|
|
module.add(VSubU32(dst=vgpr(vLocalAddress), src0=vgpr(vLocalAddress), src1=sgpr(sFirstLane)))
|
|
writer.sgprPool.checkIn(sFirstLane)
|
|
|
|
# Only first wave reads next work item index
|
|
skSkipWorkItem = Label("SK_SkipWorkItem", "")
|
|
sWave = writer.sgprPool.checkOut(1, "Wave")
|
|
# module.add(SNop(waitState=4, comment="4 wait required between VALU op and readfirstlane using the value"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sWave), src=vgpr("Serial"), comment="Wave 0 updates flags"))
|
|
# module.add(SNop(waitState=2, comment="2 wait required between VALU op and readfirstlane using the value"))
|
|
module.add(SCmpEQU32(src0=sgpr(sWave), src1=0, comment="Check for wave 0"))
|
|
module.add(SCBranchSCC0(labelName=skSkipWorkItem.getLabelName(), comment="Skip work item"))
|
|
writer.sgprPool.checkIn(sWave)
|
|
|
|
# Per-arch dynamic-queue fast-mask constants: log2(numQueues) for the
|
|
# StreamKIdx/queue divisions, log2(cache-line size) for the counter stride.
|
|
_, _, wsLog2Queues, wsCacheLineLog2 = self._wsQueueConstants(writer, kernel)
|
|
|
|
# Default queue index
|
|
sQueueIdx = writer.sgprPool.checkOut(1, "QueueIdx")
|
|
module.add(SLShiftRightB32(dst=sgpr(sQueueIdx), src=sgpr("StreamKIdx"), shiftHex=wsLog2Queues))
|
|
module.add(SLShiftLeftB32(dst=sgpr(sQueueIdx), src=sgpr(sQueueIdx), shiftHex=wsLog2Queues))
|
|
module.add(SSubU32(dst=sgpr(sQueueIdx), src0=sgpr("StreamKIdx"), src1=sgpr(sQueueIdx), comment="Default queue index"))
|
|
|
|
# Queue address
|
|
sAddress = writer.sgprPool.checkOutAligned(2, 2, "Address")
|
|
module.add(SLShiftLeftB32(dst=sgpr(sAddress), src=sgpr(sQueueIdx), shiftHex=wsCacheLineLog2, comment="Stride queues to different cache lines"))
|
|
module.add(SAddU32(dst=sgpr(sAddress+0), src0=sgpr(sAddress+0), src1=sgpr("AddressFlags+0")))
|
|
module.add(SAddCU32(dst=sgpr(sAddress+1), src0=0, src1=sgpr("AddressFlags+1")))
|
|
|
|
# Tiles in queue
|
|
sTilesInQueue = writer.sgprPool.checkOut(1, "tilesInQueue")
|
|
module.add(SLShiftRightB32(dst=sgpr(sTilesInQueue), src=sgpr("TotalItems"), shiftHex=wsLog2Queues))
|
|
sRemainder = writer.sgprPool.checkOut(1, "remainder tiles")
|
|
module.add(SLShiftLeftB32(dst=sgpr(sRemainder), src=sgpr(sTilesInQueue), shiftHex=wsLog2Queues))
|
|
module.add(SSubU32(dst=sgpr(sRemainder), src0=sgpr("TotalItems"), src1=sgpr(sRemainder), comment="Remainder tiles"))
|
|
module.add(SCmpLtU32(src0=sgpr(sQueueIdx), src1=sgpr(sRemainder), comment="Check if queue gets an extra tile"))
|
|
module.add(SCSelectB32(dst=sgpr(sRemainder), src0=1, src1=0))
|
|
module.add(SAddU32(dst=sgpr(sTilesInQueue), src0=sgpr(sTilesInQueue), src1=sgpr(sRemainder)))
|
|
writer.sgprPool.checkIn(sRemainder)
|
|
|
|
# Workgroups in queue
|
|
sWorkgroupsInQueue = writer.sgprPool.checkOut(1, "workgroupsInQueue")
|
|
module.add(SLShiftRightB32(dst=sgpr(sWorkgroupsInQueue), src=sgpr("skGrid"), shiftHex=wsLog2Queues))
|
|
sRemainder = writer.sgprPool.checkOut(1, "remainder workgroups")
|
|
module.add(SLShiftLeftB32(dst=sgpr(sRemainder), src=sgpr(sWorkgroupsInQueue), shiftHex=wsLog2Queues))
|
|
module.add(SSubU32(dst=sgpr(sRemainder), src0=sgpr("skGrid"), src1=sgpr(sRemainder), comment="Remainder workgroups"))
|
|
module.add(SCmpLtU32(src0=sgpr(sQueueIdx), src1=sgpr(sRemainder), comment="Check if queue gets an extra tile"))
|
|
module.add(SCSelectB32(dst=sgpr(sRemainder), src0=1, src1=0))
|
|
module.add(SAddU32(dst=sgpr(sWorkgroupsInQueue), src0=sgpr(sWorkgroupsInQueue), src1=sgpr(sRemainder)))
|
|
writer.sgprPool.checkIn(sRemainder)
|
|
|
|
# Fetch next work item index
|
|
sWorkItemIdx = writer.sgprPool.checkOut(1, "nextWorkItemIdx")
|
|
module.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sTilesInQueue), src1=sgpr(sWorkgroupsInQueue), comment="Queue reset"))
|
|
module.add(SSubU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=1))
|
|
writer.sgprPool.checkIn(sTilesInQueue)
|
|
writer.sgprPool.checkIn(sWorkgroupsInQueue)
|
|
|
|
# Work stealing: fold the predecessor's workgroup count into the home
|
|
# auto-reset bound so the counter still self-resets under next-neighbor stealing.
|
|
if kernel["StreamKWorkStealing"]:
|
|
self.streamKWorkStealingHomeBound(writer, module, kernel, sWorkItemIdx, sQueueIdx, "skGrid")
|
|
|
|
# Work stealing: once this WG has seen its home queue empty (sticky), it
|
|
# never touches the home counter again -- skip the home fetch and force
|
|
# the steal path with an invalid sentinel index (>= TotalItems).
|
|
if kernel["StreamKWorkStealing"]:
|
|
skStealOnly = Label(writer.labels.getNameInc("SK_StealOnly"), "")
|
|
skHomeFetched = Label(writer.labels.getNameInc("SK_HomeFetched"), "")
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKStickyEmpty"), src1=0, comment="Home not yet empty?"))
|
|
module.add(SCBranchSCC0(labelName=skStealOnly.getLabelName(), comment="Sticky: skip home fetch, steal only"))
|
|
|
|
# Fetch next work item
|
|
module.add(self._fetchNextWorkItem(writer, kernel, sWorkItemIdx, sAddress))
|
|
writer.sgprPool.checkIn(sAddress)
|
|
|
|
# Convert to global work item index
|
|
module.add(SLShiftLeftB32(dst=sgpr(sWorkItemIdx), src=sgpr(sWorkItemIdx), shiftHex=wsLog2Queues))
|
|
module.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=sgpr(sQueueIdx)))
|
|
|
|
# Work stealing: latch the sticky-empty flag on the first empty home
|
|
# fetch, then fall through to the steal (or, when already sticky, jump
|
|
# straight to the steal with the sentinel index).
|
|
if kernel["StreamKWorkStealing"]:
|
|
module.add(SCmpGeU32(src0=sgpr(sWorkItemIdx), src1=sgpr("TotalItems"), comment="Home fetch empty?"))
|
|
module.add(SCSelectB32(dst=sgpr("StreamKStickyEmpty"), src0=1, src1=0, comment="Latch sticky-empty on empty home"))
|
|
module.add(SBranch(labelName=skHomeFetched.getLabelName(), comment="Home fetched; try one steal"))
|
|
module.add(skStealOnly)
|
|
module.add(SMovB32(dst=sgpr(sWorkItemIdx), src=sgpr("TotalItems"), comment="Sentinel index (>= TotalItems) forces steal"))
|
|
module.add(skHomeFetched)
|
|
self.streamKWorkStealingSteal(writer, module, kernel, sQueueIdx, sWorkItemIdx, "skGrid", lambda base: Label(writer.labels.getNameInc(base), ""))
|
|
writer.sgprPool.checkIn(sQueueIdx)
|
|
|
|
# Share work item index with all waves
|
|
vWaveWorkItemIdx = writer.vgprPool.checkOut(1, "WaveWorkItemIdx")
|
|
module.add(VMovB32(dst=vgpr(vWaveWorkItemIdx), src=sgpr(sWorkItemIdx), comment="Move work item index to vgpr"))
|
|
module.add(DSStoreB32(dstAddr=vgpr(vLocalAddress), src=vgpr(vWaveWorkItemIdx), ds=DSModifiers(offset=0)))
|
|
module.add(SWaitCnt(dscnt=0))
|
|
|
|
module.add(skSkipWorkItem)
|
|
module.add(SBarrier())
|
|
|
|
module.add(DSLoadB32(dst=vgpr(vWaveWorkItemIdx), src=vgpr(vLocalAddress), ds=DSModifiers(offset=0)))
|
|
module.add(SWaitCnt(dscnt=0))
|
|
# module.add(SNop(waitState=4, comment="4 wait states required before reading vgpr by lane"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr(sWorkItemIdx), src=vgpr(vWaveWorkItemIdx), comment="Read work item index from vgpr"))
|
|
# module.add(SNop(waitState=2, comment="2 wait states required before reading vgpr by lane"))
|
|
module.add(SBarrier())
|
|
|
|
writer.vgprPool.checkIn(vLocalAddress)
|
|
writer.vgprPool.checkIn(vWaveWorkItemIdx)
|
|
|
|
# Check if work item index is valid
|
|
module.add(SCmpLtU32(src0=sgpr(sWorkItemIdx), src1=sgpr("TotalItems"), comment="Check if work item index is valid"))
|
|
# If work item index is not valid, skip to end of kernel
|
|
module.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
|
|
# Check if work item is a full tile. The full-tile work-item count
|
|
# spans all batches (as TotalItems does), so it must use the
|
|
# batch-inclusive total tile count (nWG0 * nWG1 * batchCount).
|
|
sFullTile = writer.sgprPool.checkOut(1, "fullTile")
|
|
module.add(self.computeTotalTiles(writer, kernel, sFullTile))
|
|
module.add(SSubU32(dst=sgpr(sFullTile), src0=sgpr(sFullTile), src1=sgpr("skTiles"), comment="Get number of full-tile work items (across all batches)"))
|
|
module.add(SCmpLtU32(src0=sgpr(sWorkItemIdx), src1=sgpr(sFullTile), comment="Check if work item is a full tile"))
|
|
module.add(SCBranchSCC0(labelName=skPartialTile.getLabelName(), comment="Work item is a partial tile"))
|
|
|
|
# Calculate iteration range for full tile
|
|
module.add(skFullTile)
|
|
module.add(SMovB32(dst=sgpr("StreamKTileIdx"), src=sgpr(sWorkItemIdx), comment="StreamKTileIdx = nextWorkItemIdx"))
|
|
module.add(SMovB32(dst=sgpr("StreamKLocalStart"), src=0, comment="StreamKLocalStart = 0"))
|
|
module.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr("ItersPerTile"), comment="StreamKLocalEnd = ItersPerTile"))
|
|
module.add(SBranch(labelName=skDone.getLabelName(), comment="Done"))
|
|
|
|
# Calculate iteration range for partial tile
|
|
module.add(skPartialTile)
|
|
# Calculate tile index of partial work item = floor((WorkItem - FullTiles) / skSplit) + FullTiles
|
|
module.add(SSubU32(dst=sgpr("StreamKTileIdx"), src0=sgpr(sWorkItemIdx), src1=sgpr(sFullTile), comment="Tile index of partial work item"))
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
module.add(scalarUInt32DivideAndRemainder(qReg="StreamKTileIdx", dReg="StreamKTileIdx", divReg="SKSplit", rReg="StreamKPartialIdx", tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=True))
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
module.add(SAddU32(dst=sgpr("StreamKTileIdx"), src0=sgpr("StreamKTileIdx"), src1=sgpr(sFullTile), comment="Offset to first partial tile"))
|
|
module.add(SMulI32(dst=sgpr("StreamKLocalStart"), src0=sgpr("StreamKPartialIdx"), src1=sgpr("SKItersPerWI"), comment="StreamKLocalStart = PartialIdx * SKItersPerWI"))
|
|
module.add(SAddU32(dst=sgpr("StreamKLocalEnd"), src0=sgpr("StreamKLocalStart"), src1=sgpr("SKItersPerWI"), comment="StreamKLocalEnd = StreamKLocalStart + SKItersPerWI"))
|
|
module.add(SMinU32(dst=sgpr("StreamKLocalEnd"), src0=sgpr("StreamKLocalEnd"), src1=sgpr("ItersPerTile"), comment="Cap ending iter at ItersPerTile"))
|
|
|
|
module.add(skDone)
|
|
writer.sgprPool.checkIn(sFullTile)
|
|
writer.sgprPool.checkIn(sWorkItemIdx)
|
|
|
|
# Map StreamK tile index to wg0/1
|
|
module.addComment0("Map StreamK tile index to wg0/1/2")
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
sRemainder = writer.sgprPool.checkOut(1, "StreamKTileIdxRemainder")
|
|
# Per-batch tile count (NOT batch-inclusive): splits the global tile
|
|
# index into batch (WorkGroup2) and the in-batch tile.
|
|
sTilesPerBatch = writer.sgprPool.checkOut(1, "TilesPerBatch")
|
|
module.add(SMulI32(dst=sgpr(sTilesPerBatch), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"), comment="tiles per batch = nWG0 * nWG1"))
|
|
module.add(scalarUInt32DivideAndRemainder(qReg="WorkGroup2", dReg="StreamKTileIdx", divReg=sTilesPerBatch, rReg=sRemainder, tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=True, comment="TileID // nWG0*nWG1"))
|
|
# Store tileID for use later in general WGM algo
|
|
# if kernel["SpaceFillingAlgo"]:
|
|
# module.add(SNop(waitState=4, comment=""))
|
|
# module.add(SMovB32(dst=sgpr("StreamKTileID"), src=sgpr(sTmp+2), comment=""))
|
|
module.add(scalarUInt32DivideAndRemainder(qReg="WorkGroup1", dReg=sRemainder, divReg="NumWorkGroups0", rReg="WorkGroup0", tmpVgprRes=tmpVgprRes, wavewidth=kernel["WavefrontSize"], doRemainder=True, comment="TileID // nWG0"))
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
writer.sgprPool.checkIn(sRemainder)
|
|
module.addSpaceLine()
|
|
|
|
writer.sgprPool.checkIn(sTilesPerBatch)
|
|
|
|
# Map SK index to WG
|
|
# module.add(self.skIndexToWG(writer, kernel, sTmp))
|
|
|
|
# Short circuit if alpha==0 (skip main loop and reading A/B, only do beta * C)
|
|
# To skip main loop in stream-k, we check if this WG is responsible for writing results (ie: WG starts tile)
|
|
# If WG starts tile then set LocalEnd=ItersPerTile to skip fixup step, and set loopCounter to 0 to skip main loop
|
|
# If WG does not start tile, skip to end of persistent loop to check for other SK tile
|
|
# TODO verify alpha check is correct for dynamic + streamk
|
|
# Use getNameInc (like the other SKAlphaCheck sites) so this label is
|
|
# unique: calculateLoopNumIterCommon also emits an "SKAlphaCheck" label
|
|
# in the same kernel, and a hardcoded name here collides with it
|
|
# ("symbol already defined") on the dynamic StreamK path.
|
|
alphaLabel = Label(writer.labels.getNameInc("SKAlphaCheck"), "")
|
|
module.add(BranchIfNotZero("Alpha", kernel["ProblemType"]["ComputeDataType"].toEnum(), alphaLabel))
|
|
# Skip to end if not doing the global write
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
skCloseLoopLabel = Label("SK_CloseLoop", "")
|
|
module.add(writer.longBranchScc0(skCloseLoopLabel, posNeg=1))
|
|
module.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr("ItersPerTile"), comment="Skip iterations"))
|
|
module.add(alphaLabel)
|
|
|
|
# writer.sgprPool.checkIn(sTmp)
|
|
|
|
return module
|
|
|
|
def computeLoadSrd(self, writer, kernel, tc, sTmp):
|
|
module = Module("StreamK Dynamic computeLoadSrd")
|
|
module.add(self.computeLoadSrdCommon(writer, kernel, tc, sTmp))
|
|
return module
|
|
|
|
def computeStoreSrdStart(self, writer, kernel):
|
|
module = Module("StreamK Dynamic computeStoreSrdStart")
|
|
module.add(self.computeStoreSrdStartCommon(writer, kernel))
|
|
return module
|
|
|
|
def graAddresses(self, writer, kernel, tP, vTmp):
|
|
module = Module("StreamK Dynamic graAddresses")
|
|
module.add(self.graAddressesCommon(writer, kernel, tP, vTmp))
|
|
return module
|
|
|
|
def declareStaggerParms(self, writer, kernel):
|
|
module = Module("StreamK Dynamic declareStaggerParms")
|
|
module.add(self.declareStaggerParmsCommon(writer, kernel))
|
|
return module
|
|
|
|
def tailLoopNumIter(self, writer, kernel, loopCounter):
|
|
module = Module("StreamK Dynamic tailLoopNumIter")
|
|
module.add(self.tailLoopNumIterCommon(writer, kernel, loopCounter))
|
|
return module
|
|
|
|
def calculateLoopNumIter(self, writer, kernel, loopCounterName, loopIdx, tmpSgprInfo):
|
|
module = Module("StreamK Dynamic calculateLoopNumIter")
|
|
module.add(self.calculateLoopNumIterCommon(writer, kernel, loopCounterName, loopIdx, tmpSgprInfo))
|
|
return module
|
|
|
|
def calculateFirstPartialIdx(self, sPartialIdx):
|
|
module = Module("StreamK Dynamic calculateFirstPartialIdx")
|
|
|
|
module.add(SMulI32(dst=sgpr(sPartialIdx), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"), comment="Total tiles"))
|
|
module.add(SSubU32(dst=sgpr(sPartialIdx), src0=sgpr(sPartialIdx), src1=sgpr("skTiles"), comment="Number of full tiles"))
|
|
module.add(SSubU32(dst=sgpr(sPartialIdx), src0=sgpr("StreamKTileIdx"), src1=sgpr(sPartialIdx), comment="PartialTile = (TileIdx - #FullTiles)"))
|
|
module.add(SMulI32(dst=sgpr(sPartialIdx), src0=sgpr(sPartialIdx), src1=sgpr("SKSplit"), comment="PartialIdxBase = PartialTile * SKSplit"))
|
|
|
|
return module
|
|
|
|
def calculatePartialIdx(self, sPartialIdx):
|
|
module = Module("StreamK Dynamic calculatePartialIdx")
|
|
|
|
module.add(self.calculateFirstPartialIdx(sPartialIdx))
|
|
module.add(SAddU32(dst=sgpr(sPartialIdx), src0=sgpr(sPartialIdx), src1=sgpr("StreamKPartialIdx"), comment="Offset to correct partials tile"))
|
|
|
|
return module
|
|
|
|
def storeBranches(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct):
|
|
module = Module("StreamK Dynamic storeBranches")
|
|
|
|
# No branches for atomic mode
|
|
if kernel["StreamKAtomic"]:
|
|
return module
|
|
|
|
skStoreLabel = Label(label=writer.labels.getNameInc("SK_Store"), comment="")
|
|
skFixupLabel = Label(label=writer.labels.getNameInc("SK_Fixup"), comment="")
|
|
|
|
# StreamK store branches
|
|
# if we're doing parallel reduction, jump to global write
|
|
# module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
# module.add(SCBranchSCC1(labelName=skStoreLabel.getLabelName(), comment="Branch if using parallel reduction, go to regular store code"))
|
|
|
|
tmpSgpr = writer.sgprPool.checkOut(4, "globalWriteElements")
|
|
# if we did not finish the tile, store partials
|
|
# branch to beta == 0 store path
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr("ItersPerTile"), comment="does wg finish tile?"))
|
|
module.add(writer.longBranchScc0(skPartialsLabel, posNeg=1))
|
|
|
|
if kernel["DebugStreamK"] & 1 == 0:
|
|
# if we started and finished the tile, regular store code
|
|
# branch to regular store code, skip fixup step
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0, comment="does wg start tile?"))
|
|
module.add(SCBranchSCC1(labelName=skStoreLabel.getLabelName(), comment="Branch if started and finished tile, go to regular store code"))
|
|
|
|
# if we finished the tile but did not start it, fix up step
|
|
# run fixup code before regular store code
|
|
sPartialIdx = writer.sgprPool.checkOut(1, "PartialIdx")
|
|
module.add(self.calculateFirstPartialIdx(sPartialIdx))
|
|
|
|
sFixupEnd = writer.sgprPool.checkOut(1, "FixupEnd")
|
|
module.add(SAddU32(dst=sgpr(sFixupEnd), src0=sgpr(sPartialIdx), src1=sgpr("StreamKPartialIdx"), comment="Final partial tile index"))
|
|
|
|
module.add(skFixupLabel)
|
|
|
|
# Check flag
|
|
module.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(sPartialIdx), shiftHex=log2(4), comment="flag offset based on partial index"))
|
|
module.add(SAddU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=self._wsFlagsBaseOffset(writer, kernel), comment="Offset flags to come after the work queues"))
|
|
module.add(SLoadB32(dst=sgpr(tmpSgpr+2), base=sgpr("AddressFlags", 2), soffset=sgpr(tmpSgpr), smem=SMEMModifiers(glc=True, dlc=True, scope=CacheScope.SCOPE_DEV), comment="get flag"))
|
|
|
|
module.add(SWaitCnt(kmcnt=0, comment="wait for flag load"))
|
|
if kernel["DebugStreamK"] & 2 == 0:
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=1, comment="check if ready"))
|
|
module.add(SCBranchSCC0(labelName=skFixupLabel.getLabelName(), comment="if flag not set, wait and check again"))
|
|
|
|
# TODO Barrier here to sync all threads in workgroup, but maybe better to have separate flag for each wavefront (to be tested)
|
|
module.add(SBarrier(comment="wait for all workgroups before resetting flag"))
|
|
skipFlagReset = Label(label=writer.labels.getNameInc("SK_SkipFlagReset"), comment="")
|
|
module.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr+2), src=vgpr("Serial"), comment="Wave 0 updates flags"))
|
|
module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=0, comment="Check for wave 0"))
|
|
module.add(SCBranchSCC0(labelName=skipFlagReset.getLabelName(), comment="Skip flag reset"))
|
|
if writer.states.asmCaps["HasScalarStore"]:
|
|
# (tmpSgpr+2) contains a vlue of 0, use it to reset the flag
|
|
module.add(SStoreB32(src=sgpr(tmpSgpr+2), base=sgpr("AddressFlags", 2), soffset=sgpr(tmpSgpr), smem=SMEMModifiers(glc=True), comment="reset flag"))
|
|
else:
|
|
module.add(VMovB32(dst=vgpr(tmpVgpr), src=0, comment="move 0 to tmpVgpr"))
|
|
module.add(self.setFlagValue(writer, src=vgpr(tmpVgpr), soffset=sgpr(tmpSgpr), comment="reset flag"))
|
|
module.add(skipFlagReset)
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
fixupEdge = [False] # Test no edge variant
|
|
module.add(self.fixupStep(writer, kernel, vectorWidths, elements, fixupEdge, tmpVgpr, cvtVgprStruct, sPartialIdx))
|
|
|
|
module.add(SAddU32(dst=sgpr(sPartialIdx), src0=sgpr(sPartialIdx), src1=1, comment="next partial tile index"))
|
|
module.add(SCmpLtU32(src0=sgpr(sPartialIdx), src1=sgpr(sFixupEnd), comment="done loading partial tiles?"))
|
|
module.add(SCBranchSCC1(labelName=skFixupLabel.getLabelName(), comment="Branch to continue fixup loop"))
|
|
|
|
writer.sgprPool.checkIn(sFixupEnd)
|
|
writer.sgprPool.checkIn(sPartialIdx)
|
|
|
|
module.add(skStoreLabel)
|
|
|
|
return module
|
|
|
|
def writePartials(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel):
|
|
# TODO this can be combined with common case, only part that's different is workspace SRD
|
|
module = Module("StreamK Dynamic writePartials")
|
|
|
|
# No partials for atomic mode
|
|
if kernel["StreamKAtomic"]:
|
|
return module
|
|
|
|
module.add(skPartialsLabel)
|
|
if kernel["DebugStreamK"] & 2 != 0:
|
|
return module
|
|
|
|
# fixupEdge = [False] # Temporary hack to test no edge variant
|
|
edges = [False]
|
|
|
|
partialsLabels = {}
|
|
for edge in edges:
|
|
partialsLabels[edge] = Label(writer.labels.getNameInc("GW_Partials_E%u" % ( 1 if edge else 0)), comment="")
|
|
|
|
if False in edges and True in edges:
|
|
with self.allocTmpSgpr(4, tag="StreamKDynamic_writePartials_tmpSgprInfo") as tmpSgprInfo:
|
|
module.add(writer.checkIsEdge(kernel, tmpSgprInfo, partialsLabels[True], partialsLabels[True]))
|
|
|
|
for edge in edges:
|
|
module.add(partialsLabels[edge])
|
|
sPartialIdx = writer.sgprPool.checkOut(1, "PartialIdx")
|
|
module.add(self.calculatePartialIdx(sPartialIdx))
|
|
module.add(self.computeWorkspaceSrd(writer, kernel, sgpr(sPartialIdx)))
|
|
writer.sgprPool.checkIn(sPartialIdx)
|
|
module.add(self.partialsWriteProcedure(writer, kernel, vectorWidths, elements, False, False, edge, tmpVgpr, cvtVgprStruct, endLabel))
|
|
|
|
return module
|
|
|
|
def initializeSrdAddressFlagsCheck(self, GeneralBatchedGemmSrdInitiation):
|
|
module = Module("StreamK Dynamic initializeSrdAddressFlagsCheck")
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=GeneralBatchedGemmSrdInitiation.getLabelName(), comment="Parallel Reduction for General Batched GEMM, Srd initialized to workspace"))
|
|
return module
|
|
|
|
def routeToGeneralBatchedOrStridedBatched(self, stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel):
|
|
module = Module("StreamK Dynamic routeToGeneralBatchedOrStridedBatched")
|
|
module.add(self.stridedBatchOrGeneralBatch(stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel))
|
|
return module
|
|
|
|
def kernelEnd(self, writer, kernel):
|
|
module = Module("StreamK Dynamic kernelEnd")
|
|
|
|
# Per-queue atomic_inc auto-resets; no kernelEnd reset needed.
|
|
|
|
return module
|
|
|
|
def _extract_hybrid_mode():
|
|
"""Extract SK5 mode bit 30 into StreamKHybridMode; clear it in MagicShiftItersPerTile."""
|
|
module = Module("SK5 mode extraction")
|
|
module.add(SLShiftRightB32(dst=sgpr("StreamKHybridMode"),
|
|
src=sgpr("MagicShiftItersPerTile"),
|
|
shiftHex=hex(30),
|
|
comment="SK5: shift mode bit (bit 30) down"))
|
|
module.add(SAndB32(dst=sgpr("StreamKHybridMode"),
|
|
src0=sgpr("StreamKHybridMode"),
|
|
src1=hex(0x1),
|
|
comment="SK5: isolate mode bit -> StreamKHybridMode"))
|
|
module.add(SAndB32(dst=sgpr("MagicShiftItersPerTile"),
|
|
src0=sgpr("MagicShiftItersPerTile"),
|
|
src1=hex(0xBFFFFFFF),
|
|
comment="SK5: clear bit 30; keep shift bits and magic add bit"))
|
|
return module
|
|
|
|
class StreamKHybrid(StreamK):
|
|
"""
|
|
Hybrid SK3 + SK4: emits both the static (TwoTileDPFirst) and dynamic
|
|
(Dynamic work-queue) code paths in a single kernel. A runtime mode bit
|
|
packed into bit 30 of the MagicShiftItersPerTile kernel arg selects
|
|
which path executes. The bit is extracted once at preLoop entry into
|
|
the StreamKHybridMode SGPR; every divergent SK3-vs-SK4 callsite emits
|
|
both fragments back-to-back gated by an s_cmp_eq_u32 + s_cbranch on
|
|
that single SGPR.
|
|
|
|
Kernel-argument layout (see Tensile/Components/Signature.py SK5 branch
|
|
and tensilelite/src/ContractionSolution.cpp SK5 branch):
|
|
|
|
Slot SK3 (primary, defineSgpr) SK4 (RegSet alias)
|
|
---- -------------------------- ---------------------
|
|
0 ItersPerTile ItersPerTile (shared)
|
|
1 MagicNumberItersPerTile TotalItems
|
|
2 MagicShiftItersPerTile SKTiles
|
|
3 SKItersPerWG SKSplit
|
|
4 skGrid SKItersPerWI
|
|
5 skTiles SKGrid
|
|
|
|
The host pushes only the 6 args matching the active mode; the inactive
|
|
path's code is dead (never executed at runtime) but still references
|
|
the SK4 names, which are resolved to the SK3 slots via RegSet aliases
|
|
emitted in KernelWriterAssembly.py (SK5 block, line ~1502).
|
|
"""
|
|
kernel = {"StreamK": 5}
|
|
emitsParallelReductionSgprAliases = True
|
|
borrowsSrdWsInEpilogue = True
|
|
emitsWorkspaceReductionBpe = True
|
|
requiresWorkspaceReductionStorePath = True
|
|
keepsConstantsInSgpr = True
|
|
supportsSubtileImpl = True
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers
|
|
# ------------------------------------------------------------------
|
|
def _emitModeExtraction(self, writer, kernel):
|
|
return _extract_hybrid_mode()
|
|
|
|
def _emitSk3Sk4Branch(self, writer, module, tag, emitDynamic, emitStatic):
|
|
"""Emit mode-gated dual path: dynamic (SK4) first, static (SK3) second."""
|
|
sk5Static = Label(writer.labels.getNameInc(f"SK5_Static{tag}"), "")
|
|
sk5Done = Label(writer.labels.getNameInc(f"SK5_{tag}Done"), "")
|
|
|
|
module.add(SCmpEQU32(src0=sgpr("StreamKHybridMode"), src1=0,
|
|
comment=f"SK5: mode bit == 0 -> SK3 (static) {tag}"))
|
|
module.add(SCBranchSCC1(labelName=sk5Static.getLabelName(),
|
|
comment=f"SK5: branch to static {tag}"))
|
|
|
|
module.addComment2(f"SK5 dynamic (SK4) {tag}")
|
|
emitDynamic(module)
|
|
|
|
module.add(SBranch(labelName=sk5Done.getLabelName(),
|
|
comment=f"SK5: skip static {tag}"))
|
|
|
|
module.add(sk5Static)
|
|
module.addComment2(f"SK5 static (SK3) {tag}")
|
|
emitStatic(module)
|
|
|
|
module.add(sk5Done)
|
|
|
|
# ------------------------------------------------------------------
|
|
# preLoop
|
|
# ------------------------------------------------------------------
|
|
def preLoop(self, writer, kernel):
|
|
module = Module("StreamK Hybrid openLoop")
|
|
|
|
# ----- Common prologue: XCC mapping, gfx12 workaround, save WG0 -----
|
|
xccMapping = Component.XCCMapping.find(writer)
|
|
module.add(xccMapping(writer, kernel))
|
|
|
|
# Skip the gfx12 ttmp reread under clustering: defineAndResources already left
|
|
# the cluster-decoded rank in WorkGroup0/1/2, and rereading ttmp9 (cluster_x here)
|
|
# would collide StreamKIdx across the cluster.
|
|
if writer.states.archCaps["WorkGroupIdFromTTM"] and not clusterEnabled(kernel["ClusterDim"]):
|
|
module.add(SMovB32(dst=sgpr("WorkGroup0"), src="ttmp9", comment="workaround"))
|
|
module.add(SAndB32(dst=sgpr("WorkGroup1"), src0=hex(0xFFFF), src1="ttmp7", comment="workaround"))
|
|
module.add(SLShiftRightB32(dst=sgpr("WorkGroup2"), shiftHex=hex(0x10), src="ttmp7", comment="workaround"))
|
|
|
|
# SK5 always has isStreamKConstantsToVgprEnabled(kernel) == False,
|
|
# so save directly to the StreamKIdx SGPR (no VGPR-cache path).
|
|
module.add(SMovB32(dst=sgpr("StreamKIdx"), src=sgpr("WorkGroup0"),
|
|
comment="SK5: save original StreamK index"))
|
|
# Work stealing: this WG has not yet seen its home queue empty.
|
|
if kernel["StreamKWorkStealing"]:
|
|
module.add(SMovB32(dst=sgpr("StreamKStickyEmpty"), src=0,
|
|
comment="WS: home not yet empty"))
|
|
|
|
# ----- Extract the mode bit once for the whole kernel -----
|
|
module.add(self._emitModeExtraction(writer, kernel))
|
|
|
|
def emitDynamicPreLoop(mod):
|
|
sk4InitDone = Label(writer.labels.getNameInc("SK_InitDone"), "")
|
|
mod.add(sk4InitDone)
|
|
|
|
def emitStaticPreLoop(mod):
|
|
sk3InitDone = Label(writer.labels.getNameInc("SK_InitDone"), "")
|
|
sk3SplitInit = Label(writer.labels.getNameInc("SK_SplitInit"), "")
|
|
|
|
# Choose reduction strategy: parallel (no synchronizer) vs tree (synchronizer)
|
|
mod.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0),
|
|
comment="Check for synchronizer"))
|
|
mod.add(SCBranchSCC0(labelName=sk3SplitInit.getLabelName(),
|
|
comment="Jump to single kernel init"))
|
|
|
|
# ---- Parallel reduction init ----
|
|
stmpTileIdx = writer.sgprPool.checkOut(1, "TileIdx")
|
|
stmpPartialIdx = writer.sgprPool.checkOut(1, "PartialIdx")
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
mod.add(scalarUInt32DivideAndRemainder(
|
|
qReg=stmpTileIdx, dReg="StreamKIdx", divReg="SkSplit",
|
|
rReg=stmpPartialIdx, tmpVgprRes=tmpVgprRes,
|
|
wavewidth=kernel["WavefrontSize"], doRemainder=True,
|
|
comment="TileIdx = SKIdx // WGsPerTile, PartialIdx = SKIdx % WGsPerTile"))
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
|
|
skHasExtraLabel = Label(writer.labels.getNameInc("SK_HasExtra"), "")
|
|
skDoneExtraLabel = Label(writer.labels.getNameInc("SK_DoneExtra"), "")
|
|
|
|
sSkExtraIters = writer.sgprPool.checkOut(1, "extraIters")
|
|
mod.add(SMulI32(dst=sgpr(sSkExtraIters),
|
|
src0=sgpr("SkSplit"), src1=sgpr("SKItersPerWG")))
|
|
mod.add(SSubU32(dst=sgpr(sSkExtraIters),
|
|
src0=sgpr("ItersPerTile"), src1=sgpr(sSkExtraIters),
|
|
comment="extraIters = itersPerTile - SkSplit * skItersPerWG"))
|
|
|
|
mod.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr(stmpPartialIdx),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK starting iteration (case: after extra iters)"))
|
|
mod.add(SCmpLtU32(src0=sgpr(stmpPartialIdx), src1=sgpr(sSkExtraIters),
|
|
comment="Check if WG gets an extra iteration"))
|
|
mod.add(SCBranchSCC1(labelName=skHasExtraLabel.getLabelName(),
|
|
comment="Has extra iter"))
|
|
# No extra
|
|
mod.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr(sSkExtraIters),
|
|
comment="This WG does not have an extra iteration"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK ending iteration (case: after extra iters)"))
|
|
mod.add(SBranch(labelName=skDoneExtraLabel.getLabelName(),
|
|
comment="Done init for parallel reduction"))
|
|
# Has extra
|
|
mod.add(skHasExtraLabel)
|
|
mod.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr(stmpPartialIdx),
|
|
comment="This WG has an extra iteration"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK ending iteration (case: after extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"),
|
|
src1=1,
|
|
comment="StreamK ending iteration (case: after extra iters)"))
|
|
mod.add(skDoneExtraLabel)
|
|
|
|
# Offset to tile
|
|
mod.add(SMulI32(dst=sgpr(stmpTileIdx), src0=sgpr(stmpTileIdx),
|
|
src1=sgpr("ItersPerTile"),
|
|
comment="Tile offset = tilesIdx * itersPerTile"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr(stmpTileIdx), comment="Offset to correct tile"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"),
|
|
src1=sgpr(stmpTileIdx), comment="Offset to correct tile"))
|
|
# Save partial idx for SrdD calculation
|
|
mod.add(SMovB32(dst=sgpr("SkPartialIdx"), src=sgpr(stmpPartialIdx),
|
|
comment="Save partial idx for SrdD calculation"))
|
|
mod.add(SBranch(labelName=sk3InitDone.getLabelName(),
|
|
comment="Done init for parallel reduction"))
|
|
|
|
mod.add(sk3SplitInit)
|
|
writer.sgprPool.checkIn(sSkExtraIters)
|
|
writer.sgprPool.checkIn(stmpPartialIdx)
|
|
writer.sgprPool.checkIn(stmpTileIdx)
|
|
|
|
# ---- Tree reduction init ----
|
|
mod.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIdx"),
|
|
src1=sgpr("ItersPerTile"),
|
|
comment="DP starting iteration (case: DP work to do)"))
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as sTmpRes:
|
|
sTmp = sTmpRes.idx
|
|
mod.add(self.computeTotalIters(writer, kernel, sTmp))
|
|
mod.add(SMovB32(dst=sgpr("StreamKIterEnd"), src=sgpr(sTmp),
|
|
comment="DP ending iteration (case: only DP work to do)"))
|
|
mod.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("skTiles"),
|
|
src1=sgpr("ItersPerTile"), comment="Total SK iters"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sTmp), src1=sgpr("StreamKIterEnd"),
|
|
comment="Check if there are DP tiles to do"))
|
|
mod.add(SCBranchSCC1(labelName=sk3InitDone.getLabelName(),
|
|
comment="Done init"))
|
|
|
|
# No DP tiles to do, regular SK init
|
|
sSkExtraIters = writer.sgprPool.checkOut(1, "extraIters")
|
|
sIter = writer.sgprPool.checkOut(2, "SKIter")
|
|
mod.add(self.skExtraIters(writer, kernel, sSkExtraIters, sIter))
|
|
mod.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIdx"),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK starting iteration (case: after extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr(sSkExtraIters), comment="Add extra iters"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK ending iteration (case: after extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr("SKItersPerWG"), src1=1,
|
|
comment="Spread out extra iterations"))
|
|
mod.add(SMulI32(dst=sgpr(sIter), src0=sgpr("StreamKIdx"),
|
|
src1=sgpr(sIter+1),
|
|
comment="StreamK starting iteration (case: before extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr(sIter), src1=sgpr(sIter+1),
|
|
comment="StreamK ending iteration (case: before extra iters)"))
|
|
mod.add(SCmpLtU32(src0=sgpr("StreamKIdx"), src1=sgpr(sSkExtraIters),
|
|
comment="Check if lane gets an extra iteration"))
|
|
mod.add(SCSelectB32(dst=sgpr("StreamKIter"), src0=sgpr(sIter),
|
|
src1=sgpr("StreamKIter"), comment="Set start iter"))
|
|
mod.add(SCSelectB32(dst=sgpr("StreamKIterEnd"), src0=sgpr(sIter+1),
|
|
src1=sgpr("StreamKIterEnd"), comment="Set end iter"))
|
|
writer.sgprPool.checkIn(sSkExtraIters)
|
|
writer.sgprPool.checkIn(sIter)
|
|
sTmp = writer.sgprPool.checkOut(1, "TotalSKIters")
|
|
mod.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("skTiles"),
|
|
src1=sgpr("ItersPerTile"), comment="Total SK iters"))
|
|
mod.add(SMinU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"),
|
|
src1=sgpr(sTmp), comment="Cap ending iter at total SK iters"))
|
|
writer.sgprPool.checkIn(sTmp)
|
|
|
|
mod.add(sk3InitDone)
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as sTmpRes:
|
|
sTmp = sTmpRes.idx
|
|
mod.add(self.computeTotalIters(writer, kernel, sTmp))
|
|
mod.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr(sTmp),
|
|
comment="Make sure there's work to do"))
|
|
mod.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
|
|
self._emitSk3Sk4Branch(writer, module, "PreLoop", emitDynamicPreLoop, emitStaticPreLoop)
|
|
return module
|
|
|
|
# ------------------------------------------------------------------
|
|
# graWorkGroup
|
|
# ------------------------------------------------------------------
|
|
def graWorkGroup(self, writer, kernel, tPA, tPB):
|
|
module = Module("StreamK Hybrid graWorkGroup")
|
|
|
|
def emitDynamicGRA(mod):
|
|
|
|
skFullTile = Label(writer.labels.getNameInc("SK_FullTile"), "")
|
|
skPartialTile = Label(writer.labels.getNameInc("SK_PartialTile"), "")
|
|
skDone = Label(writer.labels.getNameInc("SK_Done"), "")
|
|
|
|
# Local address for sharing work id
|
|
vLocalAddress = writer.vgprPool.checkOut(1, "LocalAddress")
|
|
mod.add(VMovB32(dst=vgpr(vLocalAddress), src=vgpr("Serial"),
|
|
comment="Move local address to vgpr"))
|
|
mod.add(VLShiftLeftB32(dst=vgpr(vLocalAddress), src=vgpr(vLocalAddress),
|
|
shiftHex=log2(4), comment="Scale by BPE"))
|
|
sFirstLane = writer.sgprPool.checkOut(1, "FirstLane")
|
|
mod.add(SNop(waitState=4,
|
|
comment="4 wait required between VALU op and readfirstlane using the value"))
|
|
mod.add(VReadfirstlaneB32(dst=sgpr(sFirstLane), src=vgpr(vLocalAddress),
|
|
comment="Read first lane of local address"))
|
|
mod.add(SNop(waitState=2,
|
|
comment="2 wait required between readfirstlane and VALU op using the value"))
|
|
mod.add(VSubU32(dst=vgpr(vLocalAddress), src0=vgpr(vLocalAddress),
|
|
src1=sgpr(sFirstLane)))
|
|
writer.sgprPool.checkIn(sFirstLane)
|
|
|
|
# Only first wave reads next work item index
|
|
skSkipWorkItem = Label(writer.labels.getNameInc("SK_SkipWorkItem"), "")
|
|
sWave = writer.sgprPool.checkOut(1, "Wave")
|
|
mod.add(VReadfirstlaneB32(dst=sgpr(sWave), src=vgpr("Serial"),
|
|
comment="Wave 0 updates flags"))
|
|
mod.add(SCmpEQU32(src0=sgpr(sWave), src1=0, comment="Check for wave 0"))
|
|
mod.add(SCBranchSCC0(labelName=skSkipWorkItem.getLabelName(),
|
|
comment="Skip work item"))
|
|
writer.sgprPool.checkIn(sWave)
|
|
|
|
# Per-arch dynamic-queue fast-mask constants: log2(numQueues) for the
|
|
# StreamKIdx/queue divisions, log2(cache-line size) for the counter stride.
|
|
_, _, wsLog2Queues, wsCacheLineLog2 = self._wsQueueConstants(writer, kernel)
|
|
|
|
# Default queue index
|
|
sQueueIdx = writer.sgprPool.checkOut(1, "QueueIdx")
|
|
mod.add(SLShiftRightB32(dst=sgpr(sQueueIdx), src=sgpr("StreamKIdx"),
|
|
shiftHex=wsLog2Queues))
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sQueueIdx), src=sgpr(sQueueIdx),
|
|
shiftHex=wsLog2Queues))
|
|
mod.add(SSubU32(dst=sgpr(sQueueIdx), src0=sgpr("StreamKIdx"),
|
|
src1=sgpr(sQueueIdx), comment="Default queue index"))
|
|
|
|
# Queue address
|
|
sAddress = writer.sgprPool.checkOutAligned(2, 2, "Address")
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sAddress), src=sgpr(sQueueIdx),
|
|
shiftHex=wsCacheLineLog2,
|
|
comment="Stride queues to different cache lines"))
|
|
mod.add(SAddU32(dst=sgpr(sAddress+0), src0=sgpr(sAddress+0),
|
|
src1=sgpr("AddressFlags+0")))
|
|
mod.add(SAddCU32(dst=sgpr(sAddress+1), src0=0, src1=sgpr("AddressFlags+1")))
|
|
|
|
# Tiles in queue
|
|
sTilesInQueue = writer.sgprPool.checkOut(1, "tilesInQueue")
|
|
mod.add(SLShiftRightB32(dst=sgpr(sTilesInQueue), src=sgpr("TotalItems"),
|
|
shiftHex=wsLog2Queues))
|
|
sRemainder = writer.sgprPool.checkOut(1, "remainder tiles")
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sRemainder), src=sgpr(sTilesInQueue),
|
|
shiftHex=wsLog2Queues))
|
|
mod.add(SSubU32(dst=sgpr(sRemainder), src0=sgpr("TotalItems"),
|
|
src1=sgpr(sRemainder), comment="Remainder tiles"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sQueueIdx), src1=sgpr(sRemainder),
|
|
comment="Check if queue gets an extra tile"))
|
|
mod.add(SCSelectB32(dst=sgpr(sRemainder), src0=1, src1=0))
|
|
mod.add(SAddU32(dst=sgpr(sTilesInQueue), src0=sgpr(sTilesInQueue),
|
|
src1=sgpr(sRemainder)))
|
|
writer.sgprPool.checkIn(sRemainder)
|
|
|
|
# Workgroups in queue
|
|
sWorkgroupsInQueue = writer.sgprPool.checkOut(1, "workgroupsInQueue")
|
|
# SK5: SKGrid is the SK4-dedicated grid SGPR (uppercase).
|
|
mod.add(SLShiftRightB32(dst=sgpr(sWorkgroupsInQueue), src=sgpr("SKGrid"),
|
|
shiftHex=wsLog2Queues))
|
|
sRemainder = writer.sgprPool.checkOut(1, "remainder workgroups")
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sRemainder), src=sgpr(sWorkgroupsInQueue),
|
|
shiftHex=wsLog2Queues))
|
|
mod.add(SSubU32(dst=sgpr(sRemainder), src0=sgpr("SKGrid"),
|
|
src1=sgpr(sRemainder), comment="Remainder workgroups"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sQueueIdx), src1=sgpr(sRemainder),
|
|
comment="Check if queue gets an extra tile"))
|
|
mod.add(SCSelectB32(dst=sgpr(sRemainder), src0=1, src1=0))
|
|
mod.add(SAddU32(dst=sgpr(sWorkgroupsInQueue),
|
|
src0=sgpr(sWorkgroupsInQueue), src1=sgpr(sRemainder)))
|
|
writer.sgprPool.checkIn(sRemainder)
|
|
|
|
# Fetch next work item index
|
|
sWorkItemIdx = writer.sgprPool.checkOut(1, "nextWorkItemIdx")
|
|
mod.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sTilesInQueue),
|
|
src1=sgpr(sWorkgroupsInQueue), comment="Queue reset"))
|
|
mod.add(SSubU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx), src1=1))
|
|
writer.sgprPool.checkIn(sTilesInQueue)
|
|
writer.sgprPool.checkIn(sWorkgroupsInQueue)
|
|
|
|
# Work stealing: fold the predecessor's workgroup count into the
|
|
# home auto-reset bound so the counter still self-resets under
|
|
# next-neighbor stealing.
|
|
if kernel["StreamKWorkStealing"]:
|
|
self.streamKWorkStealingHomeBound(writer, mod, kernel, sWorkItemIdx,
|
|
sQueueIdx, "SKGrid")
|
|
|
|
# Work stealing: once this WG has seen its home queue empty (sticky),
|
|
# never touch the home counter again -- skip the home fetch and force
|
|
# the steal path with an invalid sentinel index (>= TotalItems).
|
|
if kernel["StreamKWorkStealing"]:
|
|
skStealOnly = Label(writer.labels.getNameInc("SK_StealOnly"), "")
|
|
skHomeFetched = Label(writer.labels.getNameInc("SK_HomeFetched"), "")
|
|
mod.add(SCmpEQU32(src0=sgpr("StreamKStickyEmpty"), src1=0,
|
|
comment="Home not yet empty?"))
|
|
mod.add(SCBranchSCC0(labelName=skStealOnly.getLabelName(),
|
|
comment="Sticky: skip home fetch, steal only"))
|
|
|
|
mod.add(self._fetchNextWorkItem(writer, kernel, sWorkItemIdx, sAddress))
|
|
writer.sgprPool.checkIn(sAddress)
|
|
|
|
# Convert to global work item index
|
|
mod.add(SLShiftLeftB32(dst=sgpr(sWorkItemIdx), src=sgpr(sWorkItemIdx),
|
|
shiftHex=wsLog2Queues))
|
|
mod.add(SAddU32(dst=sgpr(sWorkItemIdx), src0=sgpr(sWorkItemIdx),
|
|
src1=sgpr(sQueueIdx)))
|
|
|
|
# Work stealing: latch the sticky-empty flag on the first empty home
|
|
# fetch, then fall through to one steal (or, when already sticky,
|
|
# jump straight to the steal with the sentinel index).
|
|
if kernel["StreamKWorkStealing"]:
|
|
mod.add(SCmpGeU32(src0=sgpr(sWorkItemIdx), src1=sgpr("TotalItems"),
|
|
comment="Home fetch empty?"))
|
|
mod.add(SCSelectB32(dst=sgpr("StreamKStickyEmpty"), src0=1, src1=0,
|
|
comment="Latch sticky-empty on empty home"))
|
|
mod.add(SBranch(labelName=skHomeFetched.getLabelName(),
|
|
comment="Home fetched; try one steal"))
|
|
mod.add(skStealOnly)
|
|
mod.add(SMovB32(dst=sgpr(sWorkItemIdx), src=sgpr("TotalItems"),
|
|
comment="Sentinel index (>= TotalItems) forces steal"))
|
|
mod.add(skHomeFetched)
|
|
self.streamKWorkStealingSteal(writer, mod, kernel, sQueueIdx, sWorkItemIdx,
|
|
"SKGrid",
|
|
lambda base: Label(writer.labels.getNameInc(base), ""))
|
|
writer.sgprPool.checkIn(sQueueIdx)
|
|
|
|
# Share work item index with all waves
|
|
vWaveWorkItemIdx = writer.vgprPool.checkOut(1, "WaveWorkItemIdx")
|
|
mod.add(VMovB32(dst=vgpr(vWaveWorkItemIdx), src=sgpr(sWorkItemIdx),
|
|
comment="Move work item index to vgpr"))
|
|
mod.add(DSStoreB32(dstAddr=vgpr(vLocalAddress), src=vgpr(vWaveWorkItemIdx),
|
|
ds=DSModifiers(offset=0)))
|
|
mod.add(SWaitCnt(dscnt=0))
|
|
|
|
mod.add(skSkipWorkItem)
|
|
mod.add(SBarrier())
|
|
|
|
mod.add(DSLoadB32(dst=vgpr(vWaveWorkItemIdx), src=vgpr(vLocalAddress),
|
|
ds=DSModifiers(offset=0)))
|
|
mod.add(SWaitCnt(dscnt=0))
|
|
mod.add(VReadfirstlaneB32(dst=sgpr(sWorkItemIdx), src=vgpr(vWaveWorkItemIdx),
|
|
comment="Read work item index from vgpr"))
|
|
mod.add(SBarrier())
|
|
|
|
writer.vgprPool.checkIn(vLocalAddress)
|
|
writer.vgprPool.checkIn(vWaveWorkItemIdx)
|
|
|
|
# Check if work item index is valid
|
|
mod.add(SCmpLtU32(src0=sgpr(sWorkItemIdx), src1=sgpr("TotalItems"),
|
|
comment="Check if work item index is valid"))
|
|
mod.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
|
|
# Full tile vs partial tile. The full-tile work-item count spans all
|
|
# batches (as TotalItems does), so it must use the batch-inclusive
|
|
# total tile count (nWG0 * nWG1 * batchCount). SK5: SKTiles is the
|
|
# SK4-dedicated tiles SGPR (uppercase).
|
|
sFullTile = writer.sgprPool.checkOut(1, "fullTile")
|
|
mod.add(self.computeTotalTiles(writer, kernel, sFullTile))
|
|
mod.add(SSubU32(dst=sgpr(sFullTile), src0=sgpr(sFullTile), src1=sgpr("SKTiles"),
|
|
comment="Get number of full-tile work items (across all batches)"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sWorkItemIdx), src1=sgpr(sFullTile),
|
|
comment="Check if work item is a full tile"))
|
|
mod.add(SCBranchSCC0(labelName=skPartialTile.getLabelName(),
|
|
comment="Work item is a partial tile"))
|
|
|
|
# Full tile
|
|
mod.add(skFullTile)
|
|
mod.add(SMovB32(dst=sgpr("StreamKTileIdx"), src=sgpr(sWorkItemIdx),
|
|
comment="StreamKTileIdx = nextWorkItemIdx"))
|
|
mod.add(SMovB32(dst=sgpr("StreamKLocalStart"), src=0,
|
|
comment="StreamKLocalStart = 0"))
|
|
mod.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr("ItersPerTile"),
|
|
comment="StreamKLocalEnd = ItersPerTile"))
|
|
mod.add(SBranch(labelName=skDone.getLabelName(), comment="Done"))
|
|
|
|
# Partial tile
|
|
mod.add(skPartialTile)
|
|
mod.add(SSubU32(dst=sgpr("StreamKTileIdx"), src0=sgpr(sWorkItemIdx),
|
|
src1=sgpr(sFullTile),
|
|
comment="Tile index of partial work item"))
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
mod.add(scalarUInt32DivideAndRemainder(
|
|
qReg="StreamKTileIdx", dReg="StreamKTileIdx", divReg="SKSplit",
|
|
rReg="StreamKPartialIdx", tmpVgprRes=tmpVgprRes,
|
|
wavewidth=kernel["WavefrontSize"], doRemainder=True))
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
mod.add(SAddU32(dst=sgpr("StreamKTileIdx"), src0=sgpr("StreamKTileIdx"),
|
|
src1=sgpr(sFullTile), comment="Offset to first partial tile"))
|
|
mod.add(SMulI32(dst=sgpr("StreamKLocalStart"), src0=sgpr("StreamKPartialIdx"),
|
|
src1=sgpr("SKItersPerWI"),
|
|
comment="StreamKLocalStart = PartialIdx * SKItersPerWI"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKLocalEnd"), src0=sgpr("StreamKLocalStart"),
|
|
src1=sgpr("SKItersPerWI"),
|
|
comment="StreamKLocalEnd = StreamKLocalStart + SKItersPerWI"))
|
|
mod.add(SMinU32(dst=sgpr("StreamKLocalEnd"), src0=sgpr("StreamKLocalEnd"),
|
|
src1=sgpr("ItersPerTile"),
|
|
comment="Cap ending iter at ItersPerTile"))
|
|
|
|
mod.add(skDone)
|
|
writer.sgprPool.checkIn(sFullTile)
|
|
writer.sgprPool.checkIn(sWorkItemIdx)
|
|
|
|
# Map StreamK tile index to wg0/1/2
|
|
mod.addComment0("Map StreamK tile index to wg0/1/2")
|
|
tmpVgpr = writer.vgprPool.checkOut(2, "div")
|
|
tmpVgprRes = ContinuousRegister(idx=tmpVgpr, size=2)
|
|
sRemainder = writer.sgprPool.checkOut(1, "StreamKTileIdxRemainder")
|
|
# Per-batch tile count (NOT batch-inclusive): splits the global tile
|
|
# index into batch (WorkGroup2) and the in-batch tile.
|
|
sTilesPerBatch = writer.sgprPool.checkOut(1, "TilesPerBatch")
|
|
mod.add(SMulI32(dst=sgpr(sTilesPerBatch), src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"), comment="tiles per batch = nWG0 * nWG1"))
|
|
mod.add(scalarUInt32DivideAndRemainder(
|
|
qReg="WorkGroup2", dReg="StreamKTileIdx", divReg=sTilesPerBatch,
|
|
rReg=sRemainder, tmpVgprRes=tmpVgprRes,
|
|
wavewidth=kernel["WavefrontSize"], doRemainder=True,
|
|
comment="TileID // nWG0*nWG1"))
|
|
mod.add(scalarUInt32DivideAndRemainder(
|
|
qReg="WorkGroup1", dReg=sRemainder, divReg="NumWorkGroups0",
|
|
rReg="WorkGroup0", tmpVgprRes=tmpVgprRes,
|
|
wavewidth=kernel["WavefrontSize"], doRemainder=True,
|
|
comment="TileID // nWG0"))
|
|
tmpVgprRes = None
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
writer.sgprPool.checkIn(sRemainder)
|
|
mod.addSpaceLine()
|
|
|
|
writer.sgprPool.checkIn(sTilesPerBatch)
|
|
|
|
# alpha == 0 short-circuit
|
|
alphaLabelD = Label(writer.labels.getNameInc("SKAlphaCheck"), "")
|
|
mod.add(BranchIfNotZero("Alpha",
|
|
kernel["ProblemType"]["ComputeDataType"].toEnum(),
|
|
alphaLabelD))
|
|
mod.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0,
|
|
comment="does wg start tile?"))
|
|
skCloseLoopLabelD = Label("SK_CloseLoop", "")
|
|
mod.add(writer.longBranchScc0(skCloseLoopLabelD, posNeg=1))
|
|
mod.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr("ItersPerTile"),
|
|
comment="Skip iterations"))
|
|
mod.add(alphaLabelD)
|
|
|
|
|
|
def emitStaticGRA(mod):
|
|
|
|
sTmp = writer.sgprPool.checkOutAligned(4, 2, "SKMappingTemp")
|
|
|
|
mod.add(self.skTileIndex(writer, kernel, sTmp, tPA, tPB))
|
|
|
|
skUpdateDone = Label(writer.labels.getNameInc("SK_UpdateDone"), "")
|
|
skSplitUpdate = Label(writer.labels.getNameInc("SK_SplitUpdate"), "")
|
|
|
|
mod.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0),
|
|
comment="Check for synchronizer"))
|
|
mod.add(SCBranchSCC0(labelName=skSplitUpdate.getLabelName(),
|
|
comment="Jump to single kernel update"))
|
|
# Parallel reduction
|
|
mod.add(SMovB32(dst=sgpr(sTmp+1), src=sgpr("StreamKIterEnd"),
|
|
comment="Parallel reduction, work contained to single partial tile"))
|
|
mod.add(SBranch(labelName=skUpdateDone.getLabelName(),
|
|
comment="Done update for parallel reduction"))
|
|
mod.add(skSplitUpdate)
|
|
|
|
mod.add(self.computeTotalTiles(writer, kernel, sTmp+3))
|
|
mod.add(SSubU32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr("skTiles"),
|
|
comment="dpTiles = totalTiles - skTiles"))
|
|
|
|
mod.add(SMulI32(dst=sgpr(sTmp+3), src0=sgpr(sTmp+3), src1=sgpr("ItersPerTile"),
|
|
comment="dpSectionSize = dpTiles * ItersPerTile"))
|
|
|
|
mod.add(SMulI32(dst=sgpr(sTmp+1), src0=sgpr("skGrid"), src1=sgpr("ItersPerTile"),
|
|
comment="DP iterations shift"))
|
|
mod.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr(sTmp+1), src1=sgpr("StreamKIter"),
|
|
comment="Add DP shift"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sTmp+1), src1=sgpr(sTmp+3),
|
|
comment="Check if still in DP section"))
|
|
mod.add(SCBranchSCC1(labelName=skUpdateDone.getLabelName(),
|
|
comment="Done update"))
|
|
mod.add(SMovB32(dst=sgpr(sTmp+1), src=sgpr(sTmp+2),
|
|
comment="SK iterations shift"))
|
|
mod.add(SCmpLeU32(src0=sgpr(sTmp+3), src1=sgpr("StreamKIter"),
|
|
comment="Check if continuing in SK section"))
|
|
mod.add(SCBranchSCC1(labelName=skUpdateDone.getLabelName(),
|
|
comment="Done update"))
|
|
|
|
# Switch from DP to SK
|
|
sSkExtraIters = writer.sgprPool.checkOut(1, "extraIters")
|
|
sIter = writer.sgprPool.checkOut(2, "SKIter")
|
|
mod.add(self.skExtraIters(writer, kernel, sSkExtraIters, sIter))
|
|
mod.add(SMulI32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIdx"),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK starting iteration (case: after extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIter"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr(sSkExtraIters), comment="Add extra iters"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIter"),
|
|
src1=sgpr("SKItersPerWG"),
|
|
comment="StreamK ending iteration (case: after extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr("SKItersPerWG"), src1=1,
|
|
comment="Spread out extra iterations"))
|
|
mod.add(SMulI32(dst=sgpr(sIter), src0=sgpr("StreamKIdx"), src1=sgpr(sIter+1),
|
|
comment="StreamK starting iteration (case: before extra iters)"))
|
|
mod.add(SAddU32(dst=sgpr(sIter+1), src0=sgpr(sIter), src1=sgpr(sIter+1),
|
|
comment="StreamK ending iteration (case: before extra iters)"))
|
|
mod.add(SCmpLtU32(src0=sgpr("StreamKIdx"), src1=sgpr(sSkExtraIters),
|
|
comment="Check if lane gets an extra iteration"))
|
|
writer.sgprPool.checkIn(sSkExtraIters)
|
|
mod.add(SCSelectB32(dst=sgpr("StreamKIter"), src0=sgpr(sIter),
|
|
src1=sgpr("StreamKIter"), comment="Set start iter"))
|
|
mod.add(SCSelectB32(dst=sgpr("StreamKIterEnd"), src0=sgpr(sIter+1),
|
|
src1=sgpr("StreamKIterEnd"), comment="Set end iter"))
|
|
writer.sgprPool.checkIn(sIter)
|
|
mod.add(SAddU32(dst=sgpr(sTmp+1), src0=sgpr("StreamKIter"), src1=sgpr(sTmp+3),
|
|
comment="Offset to start of SK section"))
|
|
mod.add(SAddU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"),
|
|
src1=sgpr(sTmp+3), comment="Offset to start of SK section"))
|
|
with writer.allocTmpSgpr(1, tag="TotalIters") as tmpTotalIters:
|
|
sTotalIters = tmpTotalIters.idx
|
|
mod.add(self.computeTotalIters(writer, kernel, sTotalIters))
|
|
mod.add(SMinU32(dst=sgpr("StreamKIterEnd"), src0=sgpr("StreamKIterEnd"),
|
|
src1=sgpr(sTotalIters),
|
|
comment="Cap ending iter at total SK iters"))
|
|
mod.add(SCmpLtU32(src0=sgpr("StreamKIter"), src1=sgpr(sTotalIters),
|
|
comment="Make sure there's work to do"))
|
|
mod.add(writer.longBranchScc0(Label("KernelEnd", ""), posNeg=1))
|
|
|
|
mod.add(skUpdateDone)
|
|
mod.add(SMovB32(dst=sgpr("StreamKIter"), src=sgpr(sTmp+1),
|
|
comment="Store current iteration"))
|
|
|
|
# Map SK index to WG
|
|
mod.add(self.skIndexToWG(writer, kernel, sTmp))
|
|
|
|
# alpha == 0 short-circuit (static path)
|
|
alphaLabelS = Label(writer.labels.getNameInc("SKAlphaCheck"), "")
|
|
mod.add(BranchIfNotZero("Alpha",
|
|
kernel["ProblemType"]["ComputeDataType"].toEnum(),
|
|
alphaLabelS))
|
|
mod.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0,
|
|
comment="does wg start tile?"))
|
|
skCloseLoopLabelS = Label("SK_CloseLoop", "")
|
|
mod.add(writer.longBranchScc0(skCloseLoopLabelS, posNeg=1))
|
|
mod.add(SMovB32(dst=sgpr("StreamKLocalEnd"), src=sgpr("ItersPerTile"),
|
|
comment="Skip iterations"))
|
|
mod.add(alphaLabelS)
|
|
|
|
writer.sgprPool.checkIn(sTmp)
|
|
|
|
|
|
self._emitSk3Sk4Branch(writer, module, "GRA", emitDynamicGRA, emitStaticGRA)
|
|
return module
|
|
|
|
# ------------------------------------------------------------------
|
|
# Common delegations
|
|
# ------------------------------------------------------------------
|
|
def computeLoadSrd(self, writer, kernel, tP, sTmp):
|
|
module = Module("StreamK Hybrid computeLoadSrd")
|
|
module.add(self.computeLoadSrdCommon(writer, kernel, tP, sTmp))
|
|
return module
|
|
|
|
def computeStoreSrdStart(self, writer, kernel):
|
|
module = Module("StreamK Hybrid computeStoreSrdStart")
|
|
module.add(self.computeStoreSrdStartCommon(writer, kernel))
|
|
return module
|
|
|
|
def graAddresses(self, writer, kernel, tP, vTmp):
|
|
module = Module("StreamK Hybrid graAddresses")
|
|
module.add(self.graAddressesCommon(writer, kernel, tP, vTmp))
|
|
return module
|
|
|
|
def declareStaggerParms(self, writer, kernel):
|
|
module = Module("StreamK Hybrid declareStaggerParms")
|
|
module.add(self.declareStaggerParmsCommon(writer, kernel))
|
|
return module
|
|
|
|
def tailLoopNumIter(self, writer, kernel, loopCounter):
|
|
module = Module("StreamK Hybrid tailLoopNumIter")
|
|
module.add(self.tailLoopNumIterCommon(writer, kernel, loopCounter))
|
|
return module
|
|
|
|
def calculateLoopNumIter(self, writer, kernel, loopCounterName, loopIdx, tmpSgprInfo):
|
|
module = Module("StreamK Hybrid calculateLoopNumIter")
|
|
module.add(self.calculateLoopNumIterCommon(writer, kernel, loopCounterName, loopIdx, tmpSgprInfo))
|
|
return module
|
|
|
|
# ------------------------------------------------------------------
|
|
# SK4-style partial-index helpers (used by the dynamic side of
|
|
# partialsWriteProcedure and the dynamic SRD setup in writePartials).
|
|
# Note: SK5 uses SKTiles (uppercase) as the SK4-dedicated tile count.
|
|
# ------------------------------------------------------------------
|
|
def calculateFirstPartialIdx(self, sPartialIdx):
|
|
module = Module("StreamK Hybrid calculateFirstPartialIdx")
|
|
module.add(SMulI32(dst=sgpr(sPartialIdx),
|
|
src0=sgpr("NumWorkGroups0"), src1=sgpr("NumWorkGroups1"),
|
|
comment="Total tiles"))
|
|
module.add(SSubU32(dst=sgpr(sPartialIdx),
|
|
src0=sgpr(sPartialIdx), src1=sgpr("SKTiles"),
|
|
comment="Number of full tiles"))
|
|
module.add(SSubU32(dst=sgpr(sPartialIdx),
|
|
src0=sgpr("StreamKTileIdx"), src1=sgpr(sPartialIdx),
|
|
comment="PartialTile = (TileIdx - #FullTiles)"))
|
|
module.add(SMulI32(dst=sgpr(sPartialIdx),
|
|
src0=sgpr(sPartialIdx), src1=sgpr("SKSplit"),
|
|
comment="PartialIdxBase = PartialTile * SKSplit"))
|
|
return module
|
|
|
|
def calculatePartialIdx(self, sPartialIdx):
|
|
module = Module("StreamK Hybrid calculatePartialIdx")
|
|
module.add(self.calculateFirstPartialIdx(sPartialIdx))
|
|
module.add(SAddU32(dst=sgpr(sPartialIdx),
|
|
src0=sgpr(sPartialIdx), src1=sgpr("StreamKPartialIdx"),
|
|
comment="Offset to correct partials tile"))
|
|
return module
|
|
|
|
# ------------------------------------------------------------------
|
|
# storeBranches: runtime dispatch between SK4 inlined body and SK3
|
|
# storeBranchesCommon. Both paths terminate with their own internal
|
|
# SK_Store label and fall through to the actual store sequence, so we
|
|
# need an explicit SBranch over the static body after the dynamic
|
|
# body completes.
|
|
# ------------------------------------------------------------------
|
|
def storeBranches(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct):
|
|
module = Module("StreamK Hybrid storeBranches")
|
|
|
|
if kernel["StreamKAtomic"]:
|
|
return module
|
|
|
|
def emitDynamicStore(mod):
|
|
skStoreLabel = Label(writer.labels.getNameInc("SK_Store"), "")
|
|
skFixupLabel = Label(writer.labels.getNameInc("SK_Fixup"), "")
|
|
|
|
tmpSgpr = writer.sgprPool.checkOut(4, "globalWriteElements")
|
|
mod.add(SCmpEQU32(src0=sgpr("StreamKLocalEnd"), src1=sgpr("ItersPerTile"),
|
|
comment="does wg finish tile?"))
|
|
mod.add(writer.longBranchScc0(skPartialsLabel, posNeg=1))
|
|
|
|
if kernel["DebugStreamK"] & 1 == 0:
|
|
mod.add(SCmpEQU32(src0=sgpr("StreamKLocalStart"), src1=0,
|
|
comment="does wg start tile?"))
|
|
mod.add(SCBranchSCC1(labelName=skStoreLabel.getLabelName(),
|
|
comment="Branch if started and finished tile, go to regular store code"))
|
|
|
|
sPartialIdx = writer.sgprPool.checkOut(1, "PartialIdx")
|
|
mod.add(self.calculateFirstPartialIdx(sPartialIdx))
|
|
|
|
sFixupEnd = writer.sgprPool.checkOut(1, "FixupEnd")
|
|
mod.add(SAddU32(dst=sgpr(sFixupEnd), src0=sgpr(sPartialIdx),
|
|
src1=sgpr("StreamKPartialIdx"),
|
|
comment="Final partial tile index"))
|
|
|
|
mod.add(skFixupLabel)
|
|
|
|
mod.add(SLShiftLeftB32(dst=sgpr(tmpSgpr), src=sgpr(sPartialIdx),
|
|
shiftHex=log2(4),
|
|
comment="flag offset based on partial index"))
|
|
mod.add(SAddU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=self._wsFlagsBaseOffset(writer, kernel),
|
|
comment="Offset flags to come after the work queues"))
|
|
mod.add(SLoadB32(dst=sgpr(tmpSgpr+2), base=sgpr("AddressFlags", 2),
|
|
soffset=sgpr(tmpSgpr),
|
|
smem=SMEMModifiers(glc=True, dlc=True, scope=CacheScope.SCOPE_DEV),
|
|
comment="get flag"))
|
|
|
|
mod.add(SWaitCnt(kmcnt=0, comment="wait for flag load"))
|
|
if kernel["DebugStreamK"] & 2 == 0:
|
|
mod.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=1, comment="check if ready"))
|
|
mod.add(SCBranchSCC0(labelName=skFixupLabel.getLabelName(),
|
|
comment="if flag not set, wait and check again"))
|
|
|
|
mod.add(SBarrier(comment="wait for all workgroups before resetting flag"))
|
|
skipFlagReset = Label(writer.labels.getNameInc("SK_SkipFlagReset"), "")
|
|
mod.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr+2), src=vgpr("Serial"),
|
|
comment="Wave 0 updates flags"))
|
|
mod.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=0, comment="Check for wave 0"))
|
|
mod.add(SCBranchSCC0(labelName=skipFlagReset.getLabelName(),
|
|
comment="Skip flag reset"))
|
|
if writer.states.asmCaps["HasScalarStore"]:
|
|
mod.add(SStoreB32(src=sgpr(tmpSgpr+2), base=sgpr("AddressFlags", 2),
|
|
soffset=sgpr(tmpSgpr),
|
|
smem=SMEMModifiers(glc=True), comment="reset flag"))
|
|
else:
|
|
mod.add(VMovB32(dst=vgpr(tmpVgpr), src=0, comment="move 0 to tmpVgpr"))
|
|
mod.add(self.setFlagValue(writer, src=vgpr(tmpVgpr), soffset=sgpr(tmpSgpr),
|
|
comment="reset flag"))
|
|
mod.add(skipFlagReset)
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
fixupEdge = [False]
|
|
mod.add(self.fixupStep(writer, kernel, vectorWidths, elements,
|
|
fixupEdge, tmpVgpr, cvtVgprStruct, sPartialIdx))
|
|
|
|
mod.add(SAddU32(dst=sgpr(sPartialIdx), src0=sgpr(sPartialIdx), src1=1,
|
|
comment="next partial tile index"))
|
|
mod.add(SCmpLtU32(src0=sgpr(sPartialIdx), src1=sgpr(sFixupEnd),
|
|
comment="done loading partial tiles?"))
|
|
mod.add(SCBranchSCC1(labelName=skFixupLabel.getLabelName(),
|
|
comment="Branch to continue fixup loop"))
|
|
|
|
writer.sgprPool.checkIn(sFixupEnd)
|
|
writer.sgprPool.checkIn(sPartialIdx)
|
|
else:
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
mod.add(skStoreLabel)
|
|
|
|
def emitStaticStore(mod):
|
|
mod.add(self.storeBranchesCommon(writer, kernel, skPartialsLabel,
|
|
vectorWidths, elements, tmpVgpr, cvtVgprStruct))
|
|
|
|
self._emitSk3Sk4Branch(writer, module, "Store", emitDynamicStore, emitStaticStore)
|
|
return module
|
|
|
|
# ------------------------------------------------------------------
|
|
# writePartials: runtime dispatch only for the workspace SRD setup;
|
|
# partialsWriteProcedure itself handles the SK5 flag-offset branch
|
|
# internally (modified above), so we call it once per edge.
|
|
# ------------------------------------------------------------------
|
|
def writePartials(self, writer, kernel, skPartialsLabel, vectorWidths, elements, tmpVgpr, cvtVgprStruct, endLabel):
|
|
module = Module("StreamK Hybrid writePartials")
|
|
|
|
if kernel["StreamKAtomic"]:
|
|
return module
|
|
|
|
module.add(skPartialsLabel)
|
|
if kernel["DebugStreamK"] & 2 != 0:
|
|
return module
|
|
|
|
edges = [False]
|
|
partialsLabels = {}
|
|
for edge in edges:
|
|
partialsLabels[edge] = Label(writer.labels.getNameInc("GW_Partials_E%u" % (1 if edge else 0)), comment="")
|
|
|
|
for edge in edges:
|
|
module.add(partialsLabels[edge])
|
|
|
|
def emitDynamicSrd(mod):
|
|
sPartialIdx = writer.sgprPool.checkOut(1, "PartialIdx")
|
|
mod.add(self.calculatePartialIdx(sPartialIdx))
|
|
mod.add(self.computeWorkspaceSrd(writer, kernel, sgpr(sPartialIdx)))
|
|
writer.sgprPool.checkIn(sPartialIdx)
|
|
|
|
def emitStaticSrd(mod):
|
|
mod.add(self.computeWorkspaceSrd(writer, kernel, sgpr("StreamKIdx")))
|
|
|
|
self._emitSk3Sk4Branch(writer, module, "PartialsSrd", emitDynamicSrd, emitStaticSrd)
|
|
|
|
module.add(self.partialsWriteProcedure(writer, kernel, vectorWidths, elements,
|
|
False, False, edge, tmpVgpr, cvtVgprStruct,
|
|
endLabel))
|
|
|
|
return module
|
|
|
|
def initializeSrdAddressFlagsCheck(self, GeneralBatchedGemmSrdInitiation):
|
|
module = Module("StreamK Hybrid initializeSrdAddressFlagsCheck")
|
|
module.add(SCmpEQU64(src0=sgpr("AddressFlags", 2), src1=hex(0), comment="Check for synchronizer"))
|
|
module.add(SCBranchSCC0(labelName=GeneralBatchedGemmSrdInitiation.getLabelName(), comment="Parallel Reduction for General Batched GEMM, Srd initialized to workspace"))
|
|
return module
|
|
|
|
def routeToGeneralBatchedOrStridedBatched(self, stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel):
|
|
module = Module("StreamK Hybrid routeToGeneralBatchedOrStridedBatched")
|
|
module.add(self.stridedBatchOrGeneralBatch(stridedBatchedGemmLoad, generalBatchedGemmLoad, kernel))
|
|
return module
|
|
|
|
def kernelEnd(self, writer, kernel):
|
|
module = Module("StreamK Hybrid kernelEnd")
|
|
|
|
# Per-queue atomic_inc auto-resets; no kernelEnd reset needed.
|
|
|
|
return module
|
|
|
|
|
|
# Mapping from kernel["StreamK"] int -> variant class. Lets non-
|
|
# KernelWriter consumers (e.g. Solution validation) read variant
|
|
# feature flags without an instantiated KernelWriter.
|
|
_STREAMK_VARIANT_BY_INT = {
|
|
0: StreamKOff,
|
|
3: StreamKTwoTileDPFirst,
|
|
4: StreamKDynamic,
|
|
5: StreamKHybrid,
|
|
}
|
|
|
|
|
|
def streamKVariantClass(streamK):
|
|
return _STREAMK_VARIANT_BY_INT[streamK]
|