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
1186 lines
68 KiB
Python
1186 lines
68 KiB
Python
################################################################################
|
|
#
|
|
# Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved.
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
|
|
# ies of the Software, and to permit persons to whom the Software is furnished
|
|
# to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
|
|
# PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
|
|
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
################################################################################
|
|
|
|
from .AsmAddressCalculation import AddrCalculation
|
|
from .Common import DataDirection, isSubtileMultiDU
|
|
from .Common.DataType import DataType
|
|
|
|
from math import ceil, trunc, modf
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
@dataclass
|
|
class VectorUnit:
|
|
dataType: Optional[DataType] = None
|
|
dstVgpr: int = -1
|
|
offsetVgpr: int = -1
|
|
turn: int = 0
|
|
ldsOffset: int = 0
|
|
|
|
@dataclass
|
|
class VectorDataTypes:
|
|
scaleA: VectorUnit = field(default_factory=VectorUnit)
|
|
scaleB: VectorUnit = field(default_factory=VectorUnit)
|
|
biasM: VectorUnit = field(default_factory=VectorUnit)
|
|
biasN: VectorUnit = field(default_factory=VectorUnit)
|
|
scaleAlphaM: VectorUnit = field(default_factory=VectorUnit)
|
|
scaleAlphaN: VectorUnit = field(default_factory=VectorUnit)
|
|
|
|
def bias(self, dim):
|
|
return self.biasM if dim == 0 else self.biasN
|
|
|
|
def scaleAlpha(self, dim):
|
|
return self.scaleAlphaM if dim == 0 else self.scaleAlphaN
|
|
|
|
def isValid(self):
|
|
return self.scaleA.dataType != None or self.scaleB.dataType != None or self.biasM.dataType != None \
|
|
or self.scaleAlphaM.dataType != None or self.biasN.dataType != None or self.scaleAlphaN.dataType != None
|
|
|
|
##############################################################################
|
|
# StoreState
|
|
# tracks state that is preserved across globalWriteBatch calls:
|
|
# init is called before globalWriteBatch
|
|
# the KernelWriter object
|
|
##############################################################################
|
|
class StoreState:
|
|
|
|
##############################################################################
|
|
# Setup store config for number of sgpr and vgpr needed
|
|
# These are set based on edge, atomic, etc - do not change during
|
|
# the generation of the store code.
|
|
##############################################################################
|
|
class StoreConstConfig:
|
|
def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic, isWorkspace=False):
|
|
self.gwvw = gwvw
|
|
self.lsu = kernel["LocalSplitU"]
|
|
|
|
if ss.optSingleColVgpr:
|
|
# use one vgpr (allocated in ss.sharedColDVgprs) for all addressing
|
|
# - need 0 additional vgpr per element.
|
|
self.numVgprsPerAddr = 0
|
|
else:
|
|
self.numVgprsPerAddr = kernelWriter.states.rpgo if kernel["BufferStore"] else kernelWriter.states.rpga
|
|
|
|
if ss.optSGPRUsage == 'BufferLoad_Mask':
|
|
self.numMaskSgprPerElement = 0
|
|
self.numMaskSgprPerBatch = 0
|
|
self.numTempSgprPerBatch = (2 if getattr(kernelWriter.states, 'storeAlign8', False) else 1) * kernelWriter.states.laneSGPRCount
|
|
elif ss.optSGPRUsage == 'BufferLoad_Edge_Mask':
|
|
self.numMaskSgprPerElement = 0
|
|
self.numMaskSgprPerBatch = kernelWriter.states.laneSGPRCount
|
|
self.numTempSgprPerBatch = 2 * kernelWriter.states.laneSGPRCount
|
|
else:
|
|
self.numMaskSgprPerElement = kernelWriter.states.laneSGPRCount
|
|
self.numMaskSgprPerBatch = 0
|
|
self.numTempSgprPerBatch = 2 * kernelWriter.states.laneSGPRCount
|
|
|
|
if self.numMaskSgprPerElement:
|
|
numSgprAvailable = kernelWriter.states.regCaps["MaxSgpr"] - kernelWriter.sgprPool.size() + kernelWriter.sgprPool.availableBlockAtEnd()
|
|
numSgprAvailable = numSgprAvailable & ~0x1 # make sure it's aligned
|
|
#print("numSgprAvailable=", numSgprAvailable)
|
|
self.numElementsPerBatchLimitedBySgprs = (numSgprAvailable - self.numTempSgprPerBatch - self.numMaskSgprPerBatch) // self.numMaskSgprPerElement
|
|
else:
|
|
self.numElementsPerBatchLimitedBySgprs = 9999 # no limit
|
|
|
|
if self.numElementsPerBatchLimitedBySgprs<=0:
|
|
kernelWriter.overflowedResources = 2
|
|
self.numElementsPerBatchLimitedBySgprs = 1 # dummy value
|
|
#assert self.numElementsPerBatchLimitedBySgprs > 0, "numElementsPerBatchLimitedBySgprs=0 for %s"%self.kernelName
|
|
|
|
bpeC = kernelWriter.states.bpeCexternal
|
|
if isWorkspace:
|
|
bpeC = kernelWriter.states.bpeCinternal
|
|
|
|
if atomic:
|
|
# flat atomics have another VGPR to allow different data for return#
|
|
regsPerElement = 2 if kernel["BufferStore"] else (3 + 1) # + 1 for alignment
|
|
# The atomic loop processes multiple elements in single instruction
|
|
# so will use VGPR from consec elements? TODO
|
|
self.numVgprsPerDataPerVI = (1.0 * regsPerElement * bpeC) / kernelWriter.states.bpr
|
|
elif beta:
|
|
self.numVgprsPerDataPerVI = (1.0 * bpeC) / kernelWriter.states.bpr
|
|
else:
|
|
self.numVgprsPerDataPerVI = 0.0
|
|
|
|
if kernelWriter.states.serializedStore:
|
|
#self.numVgprPerValuC = kernel["MIRegPerOut"]
|
|
self.numVgprPerValuC = kernelWriter.states.bpeCinternal//kernelWriter.states.bpr # vgpr needed from register pool
|
|
else:
|
|
self.numVgprPerValuC = 0 # null since they are already declared in macro part of assembly kernel
|
|
|
|
# indicates each vector element is actually half -
|
|
# changes vgpr allocation so two elements share a data vgpr
|
|
# Really only used if gwvw=1 - edge cases
|
|
# exception: data vgpr cannot be shared if UseInitialStridesCD is enabled and card enable EccHalf,
|
|
# since each buffer_load_short would overwrite undefined 16bit as zero.
|
|
self.halfDataRegPerVI = gwvw*self.numVgprsPerDataPerVI == 0.5 and not (kernel["ProblemType"]["UseInitialStridesCD"] and (kernelWriter.states.archCaps["HasEccHalf"] or not kernelWriter.states.asmCaps["HasWMMA_V1"])) and not (kernel["ProblemType"]["DestDataType"].numRegisters() == 0.25)
|
|
# indicates the VGPRs index offset from LSU Reduction.
|
|
# Used for multi-batch/Edge cases.
|
|
self.lsuStartVgprOffset = 0
|
|
|
|
# StoreState constructor:
|
|
def __init__(self, kernelWriter, kernel, gwvw, edge, beta, atomic, elements, vectorDataTypes, dim, isWorkspace=False):
|
|
self.kernelWriter = kernelWriter
|
|
self.kernel = kernel
|
|
self.lsu = kernel["LocalSplitU"]
|
|
self.lsuStartVgprOffset = 0
|
|
self.factorDim = dim
|
|
self.vectorDataTypes = vectorDataTypes
|
|
|
|
self.isReset = False
|
|
#--
|
|
# Optimizations for coord0/column address calculations:
|
|
#
|
|
# optSingleColVgpr:
|
|
# - works in cases where the data is written row by row to memory.
|
|
# In this case we can use a single vgpr for addressing:
|
|
# - Use the load/store instruction offset (fixed at compile-time)
|
|
# - the horizontal addresses are fixed offsets from the base
|
|
# - as we move to a new row, increment the appropriate SRDs
|
|
|
|
# optSharedColVgpr:
|
|
# - Each col gets it's own address, but elements in later rows with the same col will share VGPR.
|
|
# - allows cols to be non-adjacent
|
|
# - this is mutually exclusive with optSingleColVgpr - not as optimal but provides
|
|
# more flexibility.
|
|
|
|
# optSrdIncForRow: optimize coord1/row address calculations:
|
|
# - Move the SRD between memory operations to get to new row
|
|
# atomic needs to reset the SRD to handle retry loop. Then might work.
|
|
|
|
self.optSingleColVgpr = 0
|
|
self.optSharedColVgpr = 0
|
|
self.optSrdIncForRow = 0
|
|
|
|
# opt*ColVgpr doesn't work for edge since each element requires own addr VGPR so
|
|
# we can perform bounds check and set to -1 for OOB accesses.
|
|
# if optSingleColVgpr = optSharedColVgpr = 0, then each element gets
|
|
# 1-2 VGPRs to track address. Address calcs are performed independently
|
|
# for each element.
|
|
|
|
# atomic contains multiple memory operations which need to preserve
|
|
# the address for each load. Memops in same row can use offsets
|
|
# and share a base register but Memops in different rows need
|
|
# different registers or need to inteligently reset the SRD.
|
|
if kernel["BufferStore"] and not edge and not atomic:
|
|
if len(kernel["PackedC0IndicesX"]) > 1:
|
|
# packed mode needs a unique VGPR address calc for each column.
|
|
self.optSharedColVgpr = 1
|
|
elif len(kernel["PackedC1IndicesX"]) > 1:
|
|
self.optSharedColVgpr = 0
|
|
self.optSingleColVgpr = 0
|
|
else:
|
|
self.optSingleColVgpr = 1
|
|
|
|
if not atomic and len(kernel["PackedC1IndicesX"]) == 1:
|
|
self.optSrdIncForRow = 1
|
|
|
|
if kernel["StoreRemapVectorWidth"] and not isWorkspace:
|
|
self.optSrdIncForRow = 1
|
|
|
|
if kernel["ProblemType"]["UseInitialStridesCD"]:
|
|
self.optSingleColVgpr = 0 # BOZO, hack to disable this
|
|
self.optSharedColVgpr = 0# BOZO, hack to disable this
|
|
|
|
self.optSGPRUsage = None
|
|
if kernel["BufferStore"] and (not atomic):
|
|
self.optSGPRUsage = 'BufferLoad_Edge_Mask' if edge else 'BufferLoad_Mask'
|
|
|
|
# can't have both of these enabled:
|
|
assert (not (self.optSingleColVgpr and self.optSharedColVgpr))
|
|
|
|
|
|
self.cfg = self.StoreConstConfig(kernelWriter, kernel, self, gwvw, edge, beta, atomic, isWorkspace)
|
|
|
|
# Use to detect new rows:
|
|
self.lastCoordOffset1 = 0
|
|
|
|
# vgpr holding current coord, setup initial state
|
|
self.coord1Vgpr = kernelWriter.vgprs.coord1
|
|
|
|
# epilogue related
|
|
self.useBias = kernelWriter.states.useBias
|
|
if kernel["StreamK"] > 0 and isWorkspace:
|
|
self.useBias = DataDirection.NONE
|
|
|
|
self.useGateResidual = kernelWriter.states.useGateResidual
|
|
if kernel["StreamK"] > 0 and isWorkspace:
|
|
self.useGateResidual = False
|
|
|
|
isSingleKernel = ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or (kernel["StreamK"] > 0 and not isWorkspace)
|
|
self.referenceVgprDim = [[], []]
|
|
if self.useBias == DataDirection.READ:
|
|
self.referenceVgprDim[self.factorDim].append("Bias")
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
self.referenceVgprDim[self.factorDim].append("ScaleAlpha")
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
self.referenceVgprDim[0].append("ScaleA")
|
|
self.referenceVgprDim[1].append("ScaleB")
|
|
|
|
if self.optSharedColVgpr:
|
|
numCols = len([e for e in elements if e[0] == 0 and e[2] == 0]) # count #elements with row d1=v1==0
|
|
self.numAddrVgpr = numCols
|
|
self.sharedColDVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColDVgprs for packed elements")
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
self.sharedGSUSyncVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedGSUSyncVgprs for packed elements")
|
|
else:
|
|
self.sharedGSUSyncVgprs = None
|
|
if kernel["GroupLoadStore"] and kernel["ProblemType"]["UseBeta"]:
|
|
self.sharedColCVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColCVgprs for packed elements")
|
|
else:
|
|
self.sharedColCVgprs = self.sharedColDVgprs
|
|
if self.useBias != DataDirection.NONE:
|
|
self.sharedColBiasVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColBiasVgprs for packed elements")
|
|
else:
|
|
self.sharedColBiasVgprs = None
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
self.sharedColEVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColEVgprs for packed elements")
|
|
else:
|
|
self.sharedColEVgprs = None
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
self.sharedColGateVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColGateVgprs for packed elements")
|
|
else:
|
|
self.sharedColGateVgprs = None
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
# Only allocate a separate ScaleAlphaVec column-address vgpr for
|
|
# multi-DU kernels. Non-multi-DU reuses the Bias column address
|
|
# (localReferenceVgpr) for the SAV LDS read, so a dedicated SAV
|
|
# address would only force a redundant per-read recompute in
|
|
# emitLdChange.
|
|
if (kernel.get("UseSubtileImpl") and self.useBias == DataDirection.READ and isSubtileMultiDU(kernel)) or \
|
|
(self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha"):
|
|
self.sharedColScaleAlphaVecVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColScaleAlphaVecVgprs for packed elements")
|
|
else:
|
|
self.sharedColScaleAlphaVecVgprs = None
|
|
else:
|
|
self.sharedColScaleAlphaVecVgprs = None
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
self.sharedColScaleAVecVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColScaleAVecVgprs for packed elements")
|
|
else:
|
|
self.sharedColScaleAVecVgprs = None
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
self.sharedColScaleBVecVgprs = kernelWriter.vgprPool.checkOut(self.numAddrVgpr, "sharedColScaleBVecVgprs for packed elements")
|
|
else:
|
|
self.sharedColScaleBVecVgprs = None
|
|
else:
|
|
self.sharedColScaleAVecVgprs = None
|
|
self.sharedColScaleBVecVgprs = None
|
|
elif self.optSingleColVgpr:
|
|
self.numAddrVgpr = 1
|
|
self.sharedColDVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColDVgprs")
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
self.sharedGSUSyncVgprs = kernelWriter.vgprPool.checkOut(1, "sharedGSUSyncVgprs")
|
|
else:
|
|
self.sharedGSUSyncVgprs = None
|
|
self.singleColBiasAddrUpdated = False
|
|
self.singleColEAddrUpdated = False
|
|
self.singleColDAddrUpdated = False
|
|
self.singleColTDAddrUpdated = False
|
|
self.singleColCAddrUpdated = False
|
|
self.singleColGateAddrUpdated = False
|
|
if kernel["ProblemType"]["UseBeta"]:
|
|
self.sharedColCVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColCVgprs")
|
|
else:
|
|
self.sharedColCVgprs = self.sharedColDVgprs
|
|
if self.useBias != DataDirection.NONE:
|
|
self.sharedColBiasVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColBiasVgprs for packed elements")
|
|
else:
|
|
self.sharedColBiasVgprs = None
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
self.sharedColEVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColEVgprs for packed elements")
|
|
else:
|
|
self.sharedColEVgprs = None
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
self.sharedColGateVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColGateVgprs")
|
|
else:
|
|
self.sharedColGateVgprs = None
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
# Only allocate a separate ScaleAlphaVec column-address vgpr for
|
|
# multi-DU kernels (see note above).
|
|
if (kernel.get("UseSubtileImpl") and self.useBias == DataDirection.READ and isSubtileMultiDU(kernel)) or \
|
|
(self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha"):
|
|
self.sharedColScaleAlphaVecVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColScaleAlphaVecVgprs for packed elements")
|
|
else:
|
|
self.sharedColScaleAlphaVecVgprs = None
|
|
else:
|
|
self.sharedColScaleAlphaVecVgprs = None
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
self.sharedColScaleAVecVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColScaleAVecVgprs for packed elements")
|
|
else:
|
|
self.sharedColScaleAVecVgprs = None
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
self.sharedColScaleBVecVgprs = kernelWriter.vgprPool.checkOut(1, "sharedColScaleBVecVgprs for packed elements")
|
|
else:
|
|
self.sharedColScaleBVecVgprs = None
|
|
else:
|
|
self.sharedColScaleAVecVgprs = None
|
|
self.sharedColScaleBVecVgprs = None
|
|
else:
|
|
self.numAddrVgpr = 0
|
|
self.sharedColEVgprs = None
|
|
self.sharedColDVgprs = None
|
|
self.sharedGSUSyncVgprs = None
|
|
self.sharedColCVgprs = None
|
|
self.sharedColBiasVgprs = None
|
|
self.sharedColGateVgprs = None
|
|
self.sharedColScaleAVecVgprs = None
|
|
self.sharedColScaleBVecVgprs = None
|
|
self.sharedColScaleAlphaVecVgprs = None
|
|
|
|
# For detecting when we are running first batch
|
|
self.firstBatch = True
|
|
|
|
# TODO code migrated from globalWrite, comments outdated with respect to features like activation
|
|
# 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
|
|
|
|
# Calculate numVgprsPerElement
|
|
# print("numVgprsPerAddr=%u, numVgprsPerDataPerVI=%u, numVgprPerValuC=%u"%(self.cfg.numVgprsPerAddr, self.cfg.numVgprsPerDataPerVI, self.cfg.numVgprPerValuC))
|
|
self.numVgprsPerElement = self.cfg.numVgprPerValuC*gwvw + self.cfg.numVgprsPerAddr + int(ceil(self.cfg.numVgprsPerDataPerVI * gwvw))
|
|
# TODO STREAM-K Update numVgprsPerElement for workspace
|
|
if kernel["GroupLoadStore"] and kernel["ProblemType"]["UseBeta"]:
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr # E address
|
|
# Only needed in gradient activation
|
|
if (kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["ActivationType"] != 'none'):
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
self.numVgprsPerElement += numVgprs * gwvw # Loaded data
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
self.numVgprsPerElement += numVgprs * gwvw # Loaded data
|
|
# We will use the same vgpr + ds_offset to load the vec addr
|
|
|
|
if self.useBias != DataDirection.NONE:
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr # Bias address
|
|
if self.useBias == DataDirection.READ:
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
self.numVgprsPerElement += numVgprs * gwvw if self.factorDim == 0 else min(gwvw, 2) # Loaded data
|
|
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
if self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha":
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr # ScaleAlphaVec address
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
self.numVgprsPerElement += numVgprs * gwvw if self.factorDim == 0 else min(gwvw, 2) # Loaded data
|
|
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr # ScaleAVec address
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr # ScaleBVec address
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
self.numVgprsPerElement += numVgprs * gwvw + (numVgprs * min(gwvw, 2)) # Loaded data
|
|
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
self.numVgprsPerElement += self.cfg.numVgprsPerAddr # addrGSUSyncVgprs
|
|
|
|
# Calculate align
|
|
self.align = 1
|
|
# align adjustment
|
|
if self.cfg.numVgprsPerAddr > 1:
|
|
self.align = max(self.align, self.cfg.numVgprsPerAddr)
|
|
if self.cfg.numVgprPerValuC*gwvw > 1:
|
|
self.align = max(self.align, self.cfg.numVgprPerValuC*gwvw)
|
|
if int(ceil(self.cfg.numVgprsPerDataPerVI * gwvw)) > 1:
|
|
self.align = max(self.align, int(ceil(self.cfg.numVgprsPerDataPerVI * gwvw)))
|
|
|
|
##############################################################################
|
|
# Setup data structures to feed store loops:
|
|
# self.elementAddr, self.elementData, self.elementMask, self.elementSumIdx
|
|
# batchElements is a list of (d0,d1,v0,v1) for which stores to perform
|
|
# batchElementSgprs is SGPRs to use for mask. If None, elementMask is
|
|
# not initialized.
|
|
#
|
|
# Also create an AddrCalc for each memory operation.
|
|
##############################################################################
|
|
def getStoreElementsInfoForBatch(self, kernel, batchElements):
|
|
self.elementCoord0 = []
|
|
self.elementCoord1 = []
|
|
self.elementSumIdx = []
|
|
|
|
kw = self.kernelWriter
|
|
|
|
if kernel["EnableMatrixInstruction"]:
|
|
matrixInstT = min(kernel["MatrixInstM"], kernel["MatrixInstN"])
|
|
matrixInstBM = kernel["MatrixInstM"] // matrixInstT
|
|
matrixInstBN = kernel["MatrixInstN"] // matrixInstT
|
|
|
|
matrixInstM = (kernel["MatrixInstM"] * kernel["MatrixInstBM"]) if (kernel["MatrixInstM"] == 4) else matrixInstT
|
|
matrixInstN = (kernel["MatrixInstN"] * kernel["MatrixInstBN"]) if (kernel["MatrixInstN"] == 4) else matrixInstT
|
|
matrixInstBM = 1 if (kernel["MatrixInstM"] == 4) else kernel["MatrixInstBM"] * matrixInstBM
|
|
matrixInstBN = 1 if (kernel["MatrixInstN"] == 4) else kernel["MatrixInstBN"] * matrixInstBN
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
|
|
element = batchElements[elementIdx]
|
|
(d1,d0,vc1,vc0) = element
|
|
|
|
coordOffset1 = 0
|
|
if kernel["EnableMatrixInstruction"]:
|
|
vectorWidth = kernel["VectorWidthB"] # TODO: nonSwap VectorWidth
|
|
MIOutputVectorWidth = kernel["MIOutputVectorWidth"]
|
|
MFMAContinuousOutputs = MIOutputVectorWidth if kernel["SourceSwap"] else 1
|
|
OutputsPerMIMN = (matrixInstM * matrixInstN // self.kernel["WavefrontSize"]) if kernel["SourceSwap"] else 1
|
|
|
|
eIdx1 = d1 % (OutputsPerMIMN // MFMAContinuousOutputs)
|
|
remain_d1 = d1 // (OutputsPerMIMN // MFMAContinuousOutputs)
|
|
|
|
bIdx1 = remain_d1 % matrixInstBN
|
|
remain_d1 = remain_d1 // matrixInstBN
|
|
wtIdex = remain_d1 % kernel["MIWaveTile"][1]
|
|
|
|
coordOffset1 = eIdx1 * (self.kernel["WavefrontSize"] // matrixInstN) * MFMAContinuousOutputs
|
|
coordOffset1 += bIdx1 * matrixInstN
|
|
# Subtile kernels: successive wave tiles step by MIBShape1 (not MIBShape1 * MIWaveGroup[1]).
|
|
wtStep1 = matrixInstN * matrixInstBN if kernel.get("UseSubtileImpl") else matrixInstN * matrixInstBN * kernel["MIWaveGroup"][1]
|
|
coordOffset1 += wtIdex * wtStep1
|
|
coordOffset1 = coordOffset1 * vectorWidth + vc1
|
|
else: # mac instruction
|
|
if kernel["LocalSplitU"] > 1:
|
|
strideD1 = (kernel["NumThreads"]*kernel["VectorWidthB"]//kernel["MacroTile0"])
|
|
else:
|
|
strideD1 = (kernel["SubGroup1"] * kernel["VectorWidthB"])
|
|
coordOffset1 = d1 * strideD1 + vc1
|
|
|
|
newCoord1 = (self.firstBatch and elementIdx==0) or (coordOffset1 != self.lastCoordOffset1)
|
|
self.elementCoord1.append(coordOffset1)
|
|
|
|
# gpr and offset assignments for element
|
|
coordOffset0 = 0
|
|
if kernel["EnableMatrixInstruction"]:
|
|
vectorWidth = kernel["VectorWidthA"]
|
|
MFMAContinuousOutputs = 1 if kernel["SourceSwap"] else kernel["MIOutputVectorWidth"]
|
|
OutputsPerMIMN = 1 if kernel["SourceSwap"] else matrixInstM * matrixInstN // self.kernel["WavefrontSize"]
|
|
|
|
eIdx0 = d0 % (OutputsPerMIMN // MFMAContinuousOutputs)
|
|
remain_d0 = d0 // (OutputsPerMIMN // MFMAContinuousOutputs)
|
|
|
|
bIdx0 = remain_d0 % matrixInstBM
|
|
remain_d0 = remain_d0 // matrixInstBM
|
|
wtIdex = remain_d0 % kernel["MIWaveTile"][0]
|
|
|
|
coordOffset0 = eIdx0 * (self.kernel["WavefrontSize"] // matrixInstM) * MFMAContinuousOutputs
|
|
coordOffset0 += bIdx0 * matrixInstM
|
|
# Subtile kernels: each wave owns a contiguous block of rows, so successive
|
|
# wave tiles step by MIBShape0 (not MIBShape0 * MIWaveGroup[0]).
|
|
wtStep = matrixInstM * matrixInstBM if kernel.get("UseSubtileImpl") else matrixInstM * matrixInstBM * kernel["MIWaveGroup"][0]
|
|
coordOffset0 += wtIdex * wtStep
|
|
coordOffset0 = coordOffset0 * vectorWidth + vc0
|
|
else: # mac instruction
|
|
coordOffset0 = d0 * kernel["SubGroup0"]*kernel["VectorWidthA"] + vc0
|
|
self.elementCoord0.append(coordOffset0)
|
|
|
|
return self.elementCoord0, self.elementCoord1
|
|
|
|
def setupStoreElementsForBatchWihoutVgprCheckOut(self, kernel, gwvw, batchElements, batchElementSgprs, isOptNLL, factorDim, isWorkspace=False, elementStartIdx=0):
|
|
|
|
self.elementAddr = []
|
|
self.elementDataE = []
|
|
self.elementDataGate = []
|
|
self.elementData = [] # VGPR to use for element data, needed for atomic or beta
|
|
self.elementDataBias = []
|
|
self.elementDataScaleAVec = []
|
|
self.elementDataScaleBVec = []
|
|
self.elementDataScaleAlphaVec = []
|
|
self.elementMask = [] # SGPR to use for element mask
|
|
self.elementSumIdx = []
|
|
self.elementCoord0 = []
|
|
self.elementCoord1 = []
|
|
|
|
#generate elementCoord0/1
|
|
self.getStoreElementsInfoForBatch(kernel, batchElements)
|
|
kw = self.kernelWriter
|
|
dataType = kernel["ProblemType"]["DestDataType"]
|
|
if isWorkspace:
|
|
dataType = kernel["ProblemType"]["ComputeDataType"]
|
|
|
|
biasVgprMap = {}
|
|
scaleAVecVgprMap = {}
|
|
scaleBVecVgprMap = {}
|
|
scaleAlphaVecVgprMap = {}
|
|
lastData = 0
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# Create the AddrCalc for each memory load/store
|
|
# This is the control code that sets up the dest, source, offsets, etc and
|
|
# identifies cases where the AddrCalc is a new row and therefore needs some
|
|
# additional math. Each AddrCalc contains isolated state sufficient to
|
|
# perform any needed range checks and address calculations for the element.
|
|
#
|
|
# The AddrCalc creation code here maintains state across elements (including
|
|
# across write batches) to remove replicated calculations.
|
|
#
|
|
# Later the AddrCalc::emitAddressSetupCode will emit the necessary code
|
|
# Also allocate VGPR resources here, if needed.
|
|
|
|
element = batchElements[elementIdx]
|
|
coordOffset0 = self.elementCoord0[elementIdx]
|
|
coordOffset1 = self.elementCoord1[elementIdx]
|
|
newCoord1 = (self.firstBatch and elementIdx==0) or (coordOffset1 != self.lastCoordOffset1)
|
|
|
|
(d1,d0,vc1,vc0) = element
|
|
|
|
#print "Edge=", edge, element
|
|
sumIdx = 0
|
|
if kernel["LocalSplitU"] > 1:
|
|
if len(self.elementSumIdx) == 0:
|
|
sumIdx = kw.states.c.startVgprValu // self.cfg.numVgprPerValuC + elementStartIdx * gwvw
|
|
else:
|
|
sumIdx = self.elementSumIdx[-1] + self.cfg.gwvw
|
|
else:
|
|
bestVw = kernel["VectorWidthA"]
|
|
elementsLoadedPerVw = kernel["NumThreads"] * bestVw
|
|
elementsLoadedPerbestVw = kernel["NumThreads"] * kernel["StoreVectorWidth"]
|
|
|
|
if elementsLoadedPerVw < elementsLoadedPerbestVw:
|
|
bestVw = kernel["StoreVectorWidth"]
|
|
|
|
if kernel["EnableMatrixInstruction"]:
|
|
alignment = self.cfg.numVgprPerValuC * self.cfg.gwvw
|
|
sumIdx = kw.vgprPool.checkOutAligned(self.cfg.numVgprPerValuC*self.cfg.gwvw, alignment, "vgprValuC") // self.cfg.numVgprPerValuC
|
|
else:
|
|
sumIdx = kw.states.c.startVgprValu + vc0 + d0*kernel["VectorWidthA"] + vc1*kernel["ThreadTile0"] + d1*kernel["VectorWidthA"]*kernel["ThreadTile0"]
|
|
self.elementSumIdx.append(sumIdx) # sumIdx is an element idx, need to div/2 for half
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# Create the AddrCalc for each memory load/store
|
|
# This is the control code that sets up the dest, source, offsets, etc and
|
|
# identifies cases where the AddrCalc is a new row and therefore needs some
|
|
# additional math. Each AddrCalc contains isolated state sufficient to
|
|
# perform any needed range checks and address calculations for the element.
|
|
#
|
|
# The AddrCalc creation code here maintains state across elements (including
|
|
# across write batches) to remove replicated calculations.
|
|
#
|
|
# Later the AddrCalc::emitAddressSetupCode will emit the necessary code
|
|
# Also allocate VGPR resources here, if needed.
|
|
|
|
element = batchElements[elementIdx]
|
|
coordOffset0 = self.elementCoord0[elementIdx]
|
|
coordOffset1 = self.elementCoord1[elementIdx]
|
|
newCoord1 = (self.firstBatch and elementIdx==0) or (coordOffset1 != self.lastCoordOffset1)
|
|
|
|
(d1,d0,vc1,vc0) = element
|
|
|
|
# if numVgprsPerDataPerVI == 0.5, then two consecutive elements
|
|
# should have same data pointer, next should move.
|
|
if self.cfg.numVgprsPerDataPerVI > 0:
|
|
if self.cfg.halfDataRegPerVI:
|
|
# TODO- check (H,H,H,H,S,S)
|
|
if kernel["ProblemType"]["HighPrecisionAccumulate"] and \
|
|
(dataType.isBFloat16() or dataType.isHalf()):
|
|
data = kw.vgprPool.checkOutAligned(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
int(ceil(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw))), "writeBatch-data for ei=%u and ei=%u"%(elementIdx,elementIdx+1), preventOverflow=not isOptNLL)
|
|
else:
|
|
if elementIdx%2 == 0:
|
|
# allocate for two elements:
|
|
data = kw.vgprPool.checkOutAligned(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
int(ceil(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw))), "writeBatch-data for ei=%u and ei=%u"%(elementIdx,elementIdx+1), preventOverflow=not isOptNLL)
|
|
lastData = data
|
|
else:
|
|
data = lastData
|
|
del lastData
|
|
else:
|
|
if self.cfg.numVgprsPerDataPerVI == 0.5 or self.cfg.numVgprsPerDataPerVI == 0.25:
|
|
data = kw.vgprPool.checkOutAligned(int(ceil(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw)), \
|
|
int(ceil(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw)), "writeBatch-data for ei=%u"%elementIdx, preventOverflow=False)
|
|
else:
|
|
data = kw.vgprPool.checkOutAligned(int(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
int(ceil(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw)), "writeBatch-data for ei=%u"%elementIdx, preventOverflow=False)
|
|
#data = kw.vgprPool.checkOut(int(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
# "writeBatch-data for ei=%u"%elementIdx, preventOverflow=False)
|
|
else:
|
|
data = 0
|
|
self.elementData.append(data)
|
|
|
|
# if self.useBias == DataDirection.READ:
|
|
# coordOffset = coordOffset0 if factorDim == 0 else coordOffset1
|
|
# if coordOffset in biasVgprMap:
|
|
# dataBias = biasVgprMap[coordOffset]
|
|
# else:
|
|
# gwvw = self.cfg.gwvw if factorDim == 0 else min(self.cfg.gwvw, 2)
|
|
# numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
# dataBias = kw.vgprPool.checkOutAligned(int(numVgprs*gwvw), \
|
|
# int(ceil(numVgprs*gwvw)), "bias data for ei=%u"%elementIdx, preventOverflow=False)
|
|
# biasVgprMap[coordOffset] = dataBias
|
|
# else:
|
|
# dataBias = 0
|
|
# self.elementDataBias.append(dataBias)
|
|
# Only needed in gradient activation
|
|
# if (kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["ActivationType"] != 'none' and kernel["ProblemType"]["UseE"]) and (kernel["GlobalSplitU"] == 1):
|
|
# numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
# dataE = kw.vgprPool.checkOutAligned(int(numVgprs*self.cfg.gwvw), \
|
|
# int(ceil(numVgprs*self.cfg.gwvw)), "e data for ei=%u"%elementIdx, preventOverflow=False)
|
|
# else:
|
|
# dataE = 0
|
|
# self.elementDataE.append(dataE)
|
|
|
|
# if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and ((kernel["GlobalSplitU"] == 1) or (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel")):
|
|
# if coordOffset0 in scaleAVecVgprMap:
|
|
# dataScaleAVec = scaleAVecVgprMap[coordOffset0]
|
|
# else:
|
|
# numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
# dataScaleAVec = kw.vgprPool.checkOutAligned(int(numVgprs*self.cfg.gwvw), \
|
|
# int(ceil(numVgprs*self.cfg.gwvw)), "scaleAVec data for ei=%u"%elementIdx, preventOverflow=False)
|
|
# scaleAVecVgprMap[coordOffset0] = dataScaleAVec
|
|
# if coordOffset1 in scaleBVecVgprMap:
|
|
# dataScaleBVec = scaleBVecVgprMap[coordOffset1]
|
|
# else:
|
|
# gwvw = min(self.cfg.gwvw, 2)
|
|
# numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
# dataScaleBVec = kw.vgprPool.checkOutAligned(int(numVgprs*gwvw), \
|
|
# int(ceil(numVgprs*gwvw)), "scaleBVec data for ei=%u"%elementIdx, preventOverflow=False)
|
|
# scaleBVecVgprMap[coordOffset1] = dataScaleBVec
|
|
# else:
|
|
# dataScaleAVec = 0
|
|
# dataScaleBVec = 0
|
|
# self.elementDataScaleAVec.append(dataScaleAVec)
|
|
# self.elementDataScaleBVec.append(dataScaleBVec)
|
|
# if kernel["ProblemType"]["UseScaleAlphaVec"] and ((kernel["GlobalSplitU"] == 1) or (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel")):
|
|
# coordOffset = coordOffset0 if factorDim == 0 else coordOffset1
|
|
# gwvw = self.cfg.gwvw if factorDim == 0 else min(self.cfg.gwvw, 2)
|
|
# if coordOffset in scaleAlphaVecVgprMap:
|
|
# dataScaleAlphaVec = scaleAlphaVecVgprMap[coordOffset]
|
|
# else:
|
|
# numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
# dataScaleAlphaVec = kw.vgprPool.checkOutAligned(int(numVgprs*gwvw), \
|
|
# int(ceil(numVgprs*gwvw)), "scaleAlphaVec data for ei=%u"%elementIdx, preventOverflow=False)
|
|
# scaleAlphaVecVgprMap[coordOffset] = dataScaleAlphaVec
|
|
# else:
|
|
# dataScaleAlphaVec = 0
|
|
# self.elementDataScaleAlphaVec.append(dataScaleAlphaVec)
|
|
if batchElementSgprs != None:
|
|
if self.optSGPRUsage:
|
|
mask = batchElementSgprs
|
|
else:
|
|
mask = batchElementSgprs + self.cfg.numMaskSgprPerBatch + elementIdx * self.cfg.numMaskSgprPerElement
|
|
self.elementMask.append(mask)
|
|
|
|
if self.optSingleColVgpr:
|
|
# use same address vgpr for all
|
|
addrEVgpr = self.sharedColEVgprs
|
|
addrDVgpr = self.sharedColDVgprs
|
|
addrGSUSyncVgprs = self.sharedGSUSyncVgprs
|
|
addrCVgpr = self.sharedColCVgprs
|
|
addrBiasVgpr = self.sharedColBiasVgprs
|
|
addrGateVgpr = self.sharedColGateVgprs
|
|
addrScaleAVecVgpr = self.sharedColScaleAVecVgprs
|
|
addrScaleBVecVgpr = self.sharedColScaleBVecVgprs
|
|
addrScaleAlphaVecVgpr = self.sharedColScaleAlphaVecVgprs
|
|
elif self.optSharedColVgpr:
|
|
if kernel["EnableMatrixInstruction"]:
|
|
elementCol = (d0 * kernel["MIOutputVectorWidth"] + vc0) / gwvw
|
|
else:
|
|
elementCol = (d0 * kernel["VectorWidthA"] + vc0) / gwvw
|
|
assert (modf(elementCol)[0] < 0.001)
|
|
elementCol = trunc(elementCol)
|
|
addrDVgpr = self.sharedColDVgprs+elementCol
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
addrGSUSyncVgprs = self.sharedGSUSyncVgprs+elementCol
|
|
else:
|
|
addrGSUSyncVgprs = None
|
|
addrCVgpr = self.sharedColCVgprs+elementCol
|
|
if self.useBias != DataDirection.NONE:
|
|
addrBiasVgpr = self.sharedColBiasVgprs+elementCol
|
|
else:
|
|
addrBiasVgpr = None
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1):
|
|
addrEVgpr = self.sharedColEVgprs+elementCol
|
|
else:
|
|
addrEVgpr = None
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
addrGateVgpr = self.sharedColGateVgprs+elementCol
|
|
else:
|
|
addrGateVgpr = None
|
|
#print ("d0=", d0, "vc0=", vc0, "elementCol=", elementCol)
|
|
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and (kernel["GlobalSplitU"] == 1):
|
|
if self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha":
|
|
addrScaleAlphaVecVgpr = self.sharedColScaleAlphaVecVgprs+elementCol
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and (kernel["GlobalSplitU"] == 1):
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
addrScaleAVecVgpr = self.sharedColScaleAVecVgprs+elementCol
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
addrScaleBVecVgpr = self.sharedColScaleBVecVgprs+elementCol
|
|
else:
|
|
addrScaleBVecVgpr = None
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
addrScaleBVecVgpr = None
|
|
else:
|
|
# allocate new VGPR for each element:
|
|
addrDVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "writeDBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
addrGSUSyncVgprs = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "writeDBatch-addr for ei=%u"%(elementIdx), preventOverflow=False)
|
|
else:
|
|
addrGSUSyncVgprs = None
|
|
if kernel["GroupLoadStore"] and kernel["ProblemType"]["UseBeta"]:
|
|
addrCVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadCBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrCVgpr = addrDVgpr
|
|
if self.useBias != DataDirection.NONE:
|
|
addrBiasVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadBiasBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrBiasVgpr = None
|
|
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1):
|
|
addrEVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadEBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrEVgpr = None
|
|
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
# No-opt path: reuse D's addr
|
|
addrGateVgpr = addrDVgpr
|
|
else:
|
|
addrGateVgpr = None
|
|
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and ((kernel["GlobalSplitU"] == 1) or (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel")):
|
|
if self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha":
|
|
addrScaleAlphaVecVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadScaleAlphaVecBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and ((kernel["GlobalSplitU"] == 1) or (kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel")):
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
addrScaleAVecVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadScaleAVecBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
addrScaleBVecVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadScaleAVecBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrScaleBVecVgpr = None
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
addrScaleBVecVgpr = None
|
|
self.elementAddr.append(AddrCalculation(kw, self, addrCVgpr, addrDVgpr, addrGSUSyncVgprs, addrEVgpr, addrBiasVgpr, addrScaleAVecVgpr, addrScaleBVecVgpr, addrScaleAlphaVecVgpr, element, coordOffset0, \
|
|
self.kernelWriter.vgprs.coord1, coordOffset1, coordOffset1 - self.lastCoordOffset1, newCoord1, self.vectorDataTypes, addrGateVgpr=addrGateVgpr))
|
|
self.lastCoordOffset1 = coordOffset1
|
|
|
|
# reset flag
|
|
self.isReset = False
|
|
|
|
def setupStoreElementsForBatch(self, kernel, gwvw, batchElements, batchElementSgprs, isOptNLL, factorDim, isWorkspace=False, elementStartIdx=0):
|
|
|
|
self.elementAddr = []
|
|
self.elementDataE = []
|
|
self.elementDataGate = []
|
|
self.elementData = [] # VGPR to use for element data, needed for atomic or beta
|
|
self.elementDataBias = []
|
|
self.elementDataScaleAVec = []
|
|
self.elementDataScaleBVec = []
|
|
self.elementDataScaleAlphaVec = []
|
|
self.elementMask = [] # SGPR to use for element mask
|
|
self.elementSumIdx = []
|
|
self.elementCoord0 = []
|
|
self.elementCoord1 = []
|
|
|
|
#generate elementCoord0/1
|
|
self.getStoreElementsInfoForBatch(kernel, batchElements)
|
|
kw = self.kernelWriter
|
|
dataType = kernel["ProblemType"]["DestDataType"]
|
|
if isWorkspace:
|
|
dataType = kernel["ProblemType"]["ComputeDataType"]
|
|
|
|
biasVgprMap = {}
|
|
scaleAVecVgprMap = {}
|
|
scaleBVecVgprMap = {}
|
|
scaleAlphaVecVgprMap = {}
|
|
lastData = 0
|
|
|
|
isSingleKernel = ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or (kernel["StreamK"] > 0 and not isWorkspace)
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# Create the AddrCalc for each memory load/store
|
|
# This is the control code that sets up the dest, source, offsets, etc and
|
|
# identifies cases where the AddrCalc is a new row and therefore needs some
|
|
# additional math. Each AddrCalc contains isolated state sufficient to
|
|
# perform any needed range checks and address calculations for the element.
|
|
#
|
|
# The AddrCalc creation code here maintains state across elements (including
|
|
# across write batches) to remove replicated calculations.
|
|
#
|
|
# Later the AddrCalc::emitAddressSetupCode will emit the necessary code
|
|
# Also allocate VGPR resources here, if needed.
|
|
|
|
element = batchElements[elementIdx]
|
|
coordOffset0 = self.elementCoord0[elementIdx]
|
|
coordOffset1 = self.elementCoord1[elementIdx]
|
|
newCoord1 = (self.firstBatch and elementIdx==0) or (coordOffset1 != self.lastCoordOffset1)
|
|
|
|
(d1,d0,vc1,vc0) = element
|
|
|
|
#print "Edge=", edge, element
|
|
sumIdx = 0
|
|
if kernel["LocalSplitU"] > 1:
|
|
if len(self.elementSumIdx) == 0:
|
|
sumIdx = kw.states.c.startVgprValu // self.cfg.numVgprPerValuC + elementStartIdx * gwvw
|
|
else:
|
|
sumIdx = self.elementSumIdx[-1] + self.cfg.gwvw
|
|
else:
|
|
bestVw = kernel["VectorWidthA"]
|
|
elementsLoadedPerVw = kernel["NumThreads"] * bestVw
|
|
elementsLoadedPerbestVw = kernel["NumThreads"] * kernel["StoreVectorWidth"]
|
|
|
|
if elementsLoadedPerVw < elementsLoadedPerbestVw:
|
|
bestVw = kernel["StoreVectorWidth"]
|
|
|
|
if kernel["EnableMatrixInstruction"]:
|
|
alignment = self.cfg.numVgprPerValuC * self.cfg.gwvw
|
|
#print(self.cfg.numVgprPerValuC, self.cfg.gwvw)
|
|
#exit(1)
|
|
sumIdx = kw.vgprPool.checkOutAligned(self.cfg.numVgprPerValuC*self.cfg.gwvw, alignment, "vgprValuC") // self.cfg.numVgprPerValuC
|
|
else:
|
|
sumIdx = kw.states.c.startVgprValu + vc0 + d0*kernel["VectorWidthA"] + vc1*kernel["ThreadTile0"] + d1*kernel["VectorWidthA"]*kernel["ThreadTile0"]
|
|
self.elementSumIdx.append(sumIdx) # sumIdx is an element idx, need to div/2 for half
|
|
|
|
for elementIdx in range(0, len(batchElements)):
|
|
# Create the AddrCalc for each memory load/store
|
|
# This is the control code that sets up the dest, source, offsets, etc and
|
|
# identifies cases where the AddrCalc is a new row and therefore needs some
|
|
# additional math. Each AddrCalc contains isolated state sufficient to
|
|
# perform any needed range checks and address calculations for the element.
|
|
#
|
|
# The AddrCalc creation code here maintains state across elements (including
|
|
# across write batches) to remove replicated calculations.
|
|
#
|
|
# Later the AddrCalc::emitAddressSetupCode will emit the necessary code
|
|
# Also allocate VGPR resources here, if needed.
|
|
|
|
element = batchElements[elementIdx]
|
|
coordOffset0 = self.elementCoord0[elementIdx]
|
|
coordOffset1 = self.elementCoord1[elementIdx]
|
|
newCoord1 = (self.firstBatch and elementIdx==0) or (coordOffset1 != self.lastCoordOffset1)
|
|
|
|
(d1,d0,vc1,vc0) = element
|
|
|
|
# if numVgprsPerDataPerVI == 0.5, then two consecutive elements
|
|
# should have same data pointer, next should move.
|
|
if self.cfg.numVgprsPerDataPerVI > 0:
|
|
if self.cfg.halfDataRegPerVI:
|
|
# TODO- check (H,H,H,H,S,S)
|
|
if kernel["ProblemType"]["HighPrecisionAccumulate"] and \
|
|
(dataType.isBFloat16() or dataType.isHalf()):
|
|
data = kw.vgprPool.checkOutAligned(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
int(ceil(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw))), "writeBatch-data for ei=%u and ei=%u"%(elementIdx,elementIdx+1), preventOverflow=not isOptNLL)
|
|
else:
|
|
if elementIdx%2 == 0:
|
|
# allocate for two elements:
|
|
data = kw.vgprPool.checkOutAligned(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
int(ceil(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw))), "writeBatch-data for ei=%u and ei=%u"%(elementIdx,elementIdx+1), preventOverflow=not isOptNLL)
|
|
lastData = data
|
|
else:
|
|
data = lastData
|
|
del lastData
|
|
else:
|
|
if self.cfg.numVgprsPerDataPerVI == 0.5 or self.cfg.numVgprsPerDataPerVI == 0.25:
|
|
data = kw.vgprPool.checkOutAligned(int(ceil(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw)), \
|
|
int(ceil(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw)), "writeBatch-data for ei=%u"%elementIdx, preventOverflow=False)
|
|
else:
|
|
data = kw.vgprPool.checkOutAligned(int(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
int(ceil(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw)), "writeBatch-data for ei=%u"%elementIdx, preventOverflow=False)
|
|
#data = kw.vgprPool.checkOut(int(self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
|
|
# "writeBatch-data for ei=%u"%elementIdx, preventOverflow=False)
|
|
else:
|
|
data = 0
|
|
self.elementData.append(data)
|
|
|
|
if self.useBias == DataDirection.READ:
|
|
coordOffset = coordOffset0 if factorDim == 0 else coordOffset1
|
|
if coordOffset in biasVgprMap:
|
|
dataBias = biasVgprMap[coordOffset]
|
|
else:
|
|
gwvw = self.cfg.gwvw if factorDim == 0 else min(self.cfg.gwvw, 2)
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
dataBias = kw.vgprPool.checkOutAligned(int(numVgprs*gwvw), \
|
|
int(ceil(numVgprs*gwvw)), "bias data for ei=%u"%elementIdx, preventOverflow=False)
|
|
biasVgprMap[coordOffset] = dataBias
|
|
else:
|
|
dataBias = 0
|
|
self.elementDataBias.append(dataBias)
|
|
# Only needed in gradient activation
|
|
if (kernel["ProblemType"]["Gradient"] and kernel["ProblemType"]["ActivationType"] != 'none' and kernel["ProblemType"]["UseE"]) and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
dataE = kw.vgprPool.checkOutAligned(int(numVgprs*self.cfg.gwvw), \
|
|
int(ceil(numVgprs*self.cfg.gwvw)), "e data for ei=%u"%elementIdx, preventOverflow=False)
|
|
else:
|
|
dataE = 0
|
|
self.elementDataE.append(dataE)
|
|
|
|
# Per-element gate data VGPR
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
dataGate = kw.vgprPool.checkOutAligned(int(numVgprs*self.cfg.gwvw), \
|
|
int(ceil(numVgprs*self.cfg.gwvw)), "gate data for ei=%u"%elementIdx, preventOverflow=False)
|
|
else:
|
|
dataGate = 0
|
|
self.elementDataGate.append(dataGate)
|
|
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if coordOffset0 in scaleAVecVgprMap:
|
|
dataScaleAVec = scaleAVecVgprMap[coordOffset0]
|
|
else:
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
dataScaleAVec = kw.vgprPool.checkOutAligned(int(numVgprs*self.cfg.gwvw), \
|
|
int(ceil(numVgprs*self.cfg.gwvw)), "scaleAVec data for ei=%u"%elementIdx, preventOverflow=False)
|
|
scaleAVecVgprMap[coordOffset0] = dataScaleAVec
|
|
if coordOffset1 in scaleBVecVgprMap:
|
|
dataScaleBVec = scaleBVecVgprMap[coordOffset1]
|
|
else:
|
|
gwvw = min(self.cfg.gwvw, 2)
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
dataScaleBVec = kw.vgprPool.checkOutAligned(int(numVgprs*gwvw), \
|
|
int(ceil(numVgprs*gwvw)), "scaleBVec data for ei=%u"%elementIdx, preventOverflow=False)
|
|
scaleBVecVgprMap[coordOffset1] = dataScaleBVec
|
|
else:
|
|
dataScaleAVec = 0
|
|
dataScaleBVec = 0
|
|
self.elementDataScaleAVec.append(dataScaleAVec)
|
|
self.elementDataScaleBVec.append(dataScaleBVec)
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
coordOffset = coordOffset0 if factorDim == 0 else coordOffset1
|
|
gwvw = self.cfg.gwvw if factorDim == 0 else min(self.cfg.gwvw, 2)
|
|
if coordOffset in scaleAlphaVecVgprMap:
|
|
dataScaleAlphaVec = scaleAlphaVecVgprMap[coordOffset]
|
|
else:
|
|
numVgprs = int(ceil(kernel["ProblemType"]["ComputeDataType"].numRegisters()))
|
|
dataScaleAlphaVec = kw.vgprPool.checkOutAligned(int(numVgprs*gwvw), \
|
|
int(ceil(numVgprs*gwvw)), "scaleAlphaVec data for ei=%u"%elementIdx, preventOverflow=False)
|
|
scaleAlphaVecVgprMap[coordOffset] = dataScaleAlphaVec
|
|
else:
|
|
dataScaleAlphaVec = 0
|
|
self.elementDataScaleAlphaVec.append(dataScaleAlphaVec)
|
|
if batchElementSgprs != None:
|
|
if self.optSGPRUsage:
|
|
mask = batchElementSgprs
|
|
else:
|
|
mask = batchElementSgprs + self.cfg.numMaskSgprPerBatch + elementIdx * self.cfg.numMaskSgprPerElement
|
|
self.elementMask.append(mask)
|
|
|
|
if self.optSingleColVgpr:
|
|
# use same address vgpr for all
|
|
addrEVgpr = self.sharedColEVgprs
|
|
addrDVgpr = self.sharedColDVgprs
|
|
addrGSUSyncVgprs = self.sharedGSUSyncVgprs
|
|
addrCVgpr = self.sharedColCVgprs
|
|
addrBiasVgpr = self.sharedColBiasVgprs
|
|
addrGateVgpr = self.sharedColGateVgprs
|
|
addrScaleAVecVgpr = self.sharedColScaleAVecVgprs
|
|
addrScaleBVecVgpr = self.sharedColScaleBVecVgprs
|
|
addrScaleAlphaVecVgpr = self.sharedColScaleAlphaVecVgprs
|
|
elif self.optSharedColVgpr:
|
|
if kernel["EnableMatrixInstruction"]:
|
|
elementCol = (d0 * kernel["MIOutputVectorWidth"] + vc0) / gwvw
|
|
else:
|
|
elementCol = (d0 * kernel["VectorWidthA"] + vc0) / gwvw
|
|
assert (modf(elementCol)[0] < 0.001)
|
|
elementCol = trunc(elementCol)
|
|
addrDVgpr = self.sharedColDVgprs+elementCol
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
addrGSUSyncVgprs = self.sharedGSUSyncVgprs+elementCol
|
|
else:
|
|
addrGSUSyncVgprs = None
|
|
addrCVgpr = self.sharedColCVgprs+elementCol
|
|
if self.useBias != DataDirection.NONE:
|
|
addrBiasVgpr = self.sharedColBiasVgprs+elementCol
|
|
else:
|
|
addrBiasVgpr = None
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
addrEVgpr = self.sharedColEVgprs+elementCol
|
|
else:
|
|
addrEVgpr = None
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
addrGateVgpr = self.sharedColGateVgprs+elementCol
|
|
else:
|
|
addrGateVgpr = None
|
|
#print ("d0=", d0, "vc0=", vc0, "elementCol=", elementCol)
|
|
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or (kernel["StreamK"] > 0 and not isWorkspace)):
|
|
if self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha":
|
|
addrScaleAlphaVecVgpr = self.sharedColScaleAlphaVecVgprs+elementCol
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or (kernel["StreamK"] > 0 and not isWorkspace)):
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
addrScaleAVecVgpr = self.sharedColScaleAVecVgprs+elementCol
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
addrScaleBVecVgpr = self.sharedColScaleBVecVgprs+elementCol
|
|
else:
|
|
addrScaleBVecVgpr = None
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
addrScaleBVecVgpr = None
|
|
else:
|
|
# allocate new VGPR for each element:
|
|
# Allow pool growth for per-element D addrs; refineOccupancy still caps total VGPR.
|
|
addrDVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "writeDBatch-addr for ei=%u"%(elementIdx), preventOverflow=False)
|
|
if kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
|
|
addrGSUSyncVgprs = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "writeDBatch-addr for ei=%u"%(elementIdx), preventOverflow=False)
|
|
else:
|
|
addrGSUSyncVgprs = None
|
|
if kernel["GroupLoadStore"] and kernel["ProblemType"]["UseBeta"]:
|
|
addrCVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadCBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrCVgpr = addrDVgpr
|
|
if self.useBias != DataDirection.NONE:
|
|
addrBiasVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadBiasBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrBiasVgpr = None
|
|
|
|
if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1):
|
|
addrEVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadEBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrEVgpr = None
|
|
|
|
if self.useGateResidual and (kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1 or kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
|
|
# No-opt path: Gate reuses D's per-element addr VGPR.
|
|
addrGateVgpr = addrDVgpr
|
|
else:
|
|
addrGateVgpr = None
|
|
|
|
if kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
|
|
if self.referenceVgprDim[self.factorDim] and self.referenceVgprDim[self.factorDim][0] == "ScaleAlpha":
|
|
addrScaleAlphaVecVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadScaleAlphaVecBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
else:
|
|
addrScaleAlphaVecVgpr = None
|
|
|
|
if (kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
|
|
if self.referenceVgprDim[0] and self.referenceVgprDim[0][0] == "ScaleA":
|
|
addrScaleAVecVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadScaleAVecBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
if self.referenceVgprDim[1] and self.referenceVgprDim[1][0] == "ScaleB":
|
|
addrScaleBVecVgpr = kw.vgprPool.checkOutAligned(self.cfg.numVgprsPerAddr, \
|
|
int(ceil(self.cfg.numVgprsPerAddr)), "loadScaleAVecBatch-addr for ei=%u"%(elementIdx), preventOverflow=not isOptNLL)
|
|
else:
|
|
addrScaleBVecVgpr = None
|
|
else:
|
|
addrScaleAVecVgpr = None
|
|
addrScaleBVecVgpr = None
|
|
self.elementAddr.append(AddrCalculation(kw, self, addrCVgpr, addrDVgpr, addrGSUSyncVgprs, addrEVgpr, addrBiasVgpr, addrScaleAVecVgpr, addrScaleBVecVgpr, addrScaleAlphaVecVgpr, element, coordOffset0, \
|
|
self.kernelWriter.vgprs.coord1, coordOffset1, coordOffset1 - self.lastCoordOffset1, newCoord1, self.vectorDataTypes, addrGateVgpr=addrGateVgpr))
|
|
self.lastCoordOffset1 = coordOffset1
|
|
|
|
# reset flag
|
|
self.isReset = False
|
|
|
|
def checkInTempVgprC(self):
|
|
if not hasattr(self, 'elementSumIdx'):
|
|
return # object was not fully initialized
|
|
if self.kernelWriter.states.serializedStore is False:
|
|
return # early exit; currently only serializedStore==True checks out C-tile from register pool
|
|
|
|
if len(self.elementSumIdx) > 0 and self.lsu == 1:
|
|
for i in self.elementSumIdx:
|
|
self.kernelWriter.vgprPool.checkIn(i * self.cfg.numVgprPerValuC)
|
|
# print("checked in vgpr %u"%i)
|
|
self.elementSumIdx = []
|
|
|
|
def resetState(self):
|
|
if not self.isReset:
|
|
# Init part
|
|
if self.optSharedColVgpr:
|
|
pass # Nothing to reset
|
|
elif self.optSingleColVgpr:
|
|
self.singleColBiasAddrUpdated = False
|
|
self.singleColEAddrUpdated = False
|
|
self.singleColDAddrUpdated = False
|
|
self.singleColTDAddrUpdated = False
|
|
self.singleColCAddrUpdated = False
|
|
self.singleColGateAddrUpdated = False
|
|
else:
|
|
pass # Nothing to reset
|
|
# setup store element
|
|
self.lastCoordOffset1 = 0
|
|
self.isReset = True
|
|
|
|
def __del__(self):
|
|
|
|
if (self.sharedColEVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColEVgprs)
|
|
self.sharedColEVgprs = None
|
|
if (self.sharedColDVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColDVgprs)
|
|
if (self.sharedColCVgprs != self.sharedColDVgprs and self.sharedColCVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColCVgprs)
|
|
self.sharedColCVgprs = None
|
|
self.sharedColDVgprs = None
|
|
if (self.sharedGSUSyncVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedGSUSyncVgprs)
|
|
if (self.sharedColBiasVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColBiasVgprs)
|
|
self.sharedColBiasVgprs = None
|
|
if (self.sharedColGateVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColGateVgprs)
|
|
self.sharedColGateVgprs = None
|
|
if (self.sharedColScaleAVecVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColScaleAVecVgprs)
|
|
self.sharedColScaleAVecVgprs = None
|
|
if (self.sharedColScaleBVecVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColScaleBVecVgprs)
|
|
self.sharedColScaleBVecVgprs = None
|
|
if (self.sharedColScaleAlphaVecVgprs != None):
|
|
self.kernelWriter.vgprPool.checkIn(self.sharedColScaleAlphaVecVgprs)
|
|
self.sharedColScaleAlphaVecVgprs = None
|
|
self.checkInTempVgprC()
|
|
self.resetState()
|