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
1311 lines
75 KiB
Python
1311 lines
75 KiB
Python
################################################################################
|
|
#
|
|
# Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved.
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
|
|
# ies of the Software, and to permit persons to whom the Software is furnished
|
|
# to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
|
|
# PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
|
|
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
################################################################################
|
|
|
|
from rocisa import countInstruction, countGlobalRead, countLocalWrite, \
|
|
countDSStoreB192, countDSStoreB128, countVMovB32
|
|
from rocisa.base import Item
|
|
from rocisa.code import Module, TextBlock
|
|
from rocisa.container import DSModifiers, HolderContainer, replaceHolder
|
|
|
|
from rocisa.instruction import SWaitCnt, SWaitAlu, DSStoreB128, DSStoreB64, DSStoreB32, TensorLoadToLds
|
|
|
|
from ..Common import roundUp, print2
|
|
from ..Component import SIA
|
|
|
|
from copy import deepcopy
|
|
from typing import Tuple
|
|
PRECISION = 100
|
|
class SIA3(SIA):
|
|
kernel = {"_ScheduleIterAlg": 3}
|
|
def __call__(self):
|
|
assert(0)
|
|
|
|
def schedIntoIteration(self, writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter, firstIter, lastLoop, lastLc, globalReadIncACode, globalReadIncBCode, isNGLL):
|
|
# Get schedule information
|
|
numMfmaBetweenLWandBarrier, latencyLeft = getLocalWriteMFMAEnd(writer, kernel, tensorParametersA, tensorParametersB)
|
|
#########
|
|
# Internally assign an optimized LWPM value for PGR2
|
|
#########
|
|
# strategy is to distribute LW/GR as wide as possible to avoid hitting vmem FIFO
|
|
# LWPM = (LW_End - LW_Start) / numLW
|
|
if kernel["LocalWritePerMfma"] == -1:
|
|
lwStartMfmaIndex = getLocalWriteMFMAStart(writer, kernel, tensorParametersA, tensorParametersB, latencyLeft)
|
|
numLocalWriteModPerMfma = getNumLocalWritePerMfma(writer, kernel, lwStartMfmaIndex)
|
|
else:
|
|
numLocalWriteModPerMfma = roundUp(kernel["LocalWritePerMfma"]*PRECISION)
|
|
# safe guard
|
|
# should not be smaller than the value calculated with lwStartMfmaIndex=0
|
|
numLocalWriteModPerMfmaMin = getNumLocalWritePerMfma(writer, kernel, 0)
|
|
numLocalWriteModPerMfma = max(numLocalWriteModPerMfma, numLocalWriteModPerMfmaMin)
|
|
|
|
writer.states.numGlobalReadInsPerMfma, writer.states.numLocalWriteModPerMfma = calculateGRPMandLWPM(writer, kernel, numLocalWriteModPerMfma)
|
|
localWriteEndIter = fixLocalWriteEndMfmaIndex(writer, kernel, tensorParametersA, tensorParametersB, \
|
|
globalReadIncACode, globalReadIncBCode, numMfmaBetweenLWandBarrier, lastLoop)
|
|
numGlobalReadInsPerIter, numLocalWriteModPerIter, numEmptyGlobalReadIncCode = getScheduleParamMfma(writer)
|
|
numLocalWritesPerSched = writer.states.numLocalWriteModPerMfma
|
|
# Schedule global read
|
|
if not writer.states.scheduleGlobalRead:
|
|
itemsGRToSchedLater, lastLoadIter = noSchedGlobalRead(writer, kernel, globalReadIncACode, globalReadIncBCode)
|
|
writer.states.grEndMfmaIndex = 0
|
|
else:
|
|
itemsGRToSched, itemsGRToSchedLater = prepareGRInstToSched(writer, kernel, isNGLL)
|
|
itemsGRIncToSched = appendInstToSchedSIA3(writer, kernel, numEmptyGlobalReadIncCode, globalReadIncACode, globalReadIncBCode)
|
|
schedNumForIter0, endIter = getSchedNumForIter0SIA3(writer, kernel, itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter)
|
|
lastLoadIter = schedGlobalRead(writer, itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, schedNumForIter0, endIter)
|
|
# Schedule local write
|
|
if not writer.states.scheduleLocalWrite:
|
|
noSchedLocalWrite(writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter)
|
|
writer.states.lwStartMfmaIndex = writer.states.lwEndMfmaIndex
|
|
if kernel["1LDSBuffer"] or kernel["DirectToLds"]:
|
|
writer.states.sync1LdsMfmaIndex = max(writer.states.lwStartMfmaIndex - 1, 0)
|
|
else:
|
|
itemsLWToSched, numWritesToSched = prepareLWInstToSched(writer, kernel, numLocalWritesPerSched, isNGLL=isNGLL)
|
|
startIter = assignLWSchedIndexSIA3(writer, kernel, numLocalWritesPerSched, localWriteEndIter, numWritesToSched)
|
|
readsToWait, readsToWaitNGLL = getReadsToWait(writer, kernel)
|
|
# add waitcnt for DirectToVgpr + PGR1. Delaying wait for DirectToVgpr global read
|
|
if (kernel["DirectToVgprA"] or kernel["DirectToVgprB"]) and kernel["PrefetchGlobalRead"] == 1:
|
|
# DirectToVgpr + swapGlobalRead case, actual DTV load is in self.globalReadBCode (due to swap).
|
|
# Need to check self.globalReadBCode
|
|
readsToWaitDTV = len(list(writer.codes.globalReadB.middle.items()))
|
|
readsToWait += readsToWaitDTV
|
|
readsToWaitNGLL += readsToWaitDTV
|
|
# make sure numLocalWriteModPerIter is enough to schedule localwrite
|
|
startIterItem = numLocalWriteModPerIter - (writer.states.lwStartMfmaIndex % writer.states.numMfmaPerIter) * numLocalWritesPerSched
|
|
schedLocalWrite(writer, kernel, numLocalWriteModPerIter, numLocalWritesPerSched, localWriteEndIter, \
|
|
itemsGRToSchedLater, itemsLWToSched, startIter, readsToWait, readsToWaitNGLL, \
|
|
firstIter, lastLc, isNGLL, startIterItem)
|
|
|
|
class SIA2(SIA):
|
|
kernel = {"_ScheduleIterAlg": 2}
|
|
def __call__(self):
|
|
assert(0)
|
|
|
|
def schedIntoIteration(self, writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter, firstIter, lastLoop, lastLc, globalReadIncACode, globalReadIncBCode, isNGLL):
|
|
# Get schedule information
|
|
numGlobalReadInsPerIter, numLocalWriteModPerIter, numEmptyGlobalReadIncCode = getScheduleParams(kernel)
|
|
numLocalWritesPerSched = numLocalWriteModPerIter
|
|
# Schedule global read
|
|
if not writer.states.scheduleGlobalRead:
|
|
itemsGRToSchedLater, lastLoadIter = noSchedGlobalRead(writer, kernel, globalReadIncACode, globalReadIncBCode)
|
|
else:
|
|
itemsGRToSched, itemsGRToSchedLater = prepareGRInstToSched(writer, kernel, isNGLL)
|
|
itemsGRIncToSched = appendInstToSchedDefault(numEmptyGlobalReadIncCode, globalReadIncACode, globalReadIncBCode)
|
|
schedNumForIter0, endIter = getSchedNumForIter0Default(itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, localWriteEndIter)
|
|
lastLoadIter = schedGlobalRead(writer, itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, schedNumForIter0, endIter)
|
|
# Schedule local write
|
|
if not writer.states.scheduleLocalWrite:
|
|
noSchedLocalWrite(writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter)
|
|
else:
|
|
itemsLWToSched, numWritesToSched = prepareLWInstToSched(writer, kernel, numLocalWritesPerSched)
|
|
startIter = assignLWSchedIndexDefault(writer, kernel, numLocalWritesPerSched, localWriteEndIter, lastLoadIter, numWritesToSched)
|
|
readsToWait, readsToWaitNGLL = getReadsToWait(writer, kernel)
|
|
schedLocalWrite(writer, kernel, numLocalWriteModPerIter, numLocalWritesPerSched, localWriteEndIter, \
|
|
itemsGRToSchedLater, itemsLWToSched, startIter, readsToWait, readsToWaitNGLL, \
|
|
firstIter, lastLc, isNGLL)
|
|
|
|
class SIA1(SIA):
|
|
kernel = {"_ScheduleIterAlg": 1}
|
|
def __call__(self):
|
|
assert(0)
|
|
|
|
def schedIntoIteration(self, writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter, firstIter, lastLoop, lastLc, globalReadIncACode, globalReadIncBCode, isNGLL):
|
|
# Get schedule information
|
|
numGlobalReadInsPerIter, numLocalWriteModPerIter, numEmptyGlobalReadIncCode = getScheduleParams(kernel)
|
|
numLocalWritesPerSched = numLocalWriteModPerIter
|
|
# Schedule global read
|
|
if not writer.states.scheduleGlobalRead:
|
|
itemsGRToSchedLater, lastLoadIter = noSchedGlobalRead(writer, kernel, globalReadIncACode, globalReadIncBCode)
|
|
else:
|
|
itemsGRToSched, itemsGRToSchedLater = prepareGRInstToSched(writer, kernel, isNGLL)
|
|
itemsGRIncToSched = appendInstToSchedDefault(numEmptyGlobalReadIncCode, globalReadIncACode, globalReadIncBCode)
|
|
schedNumForIter0, endIter = getSchedNumForIter0Default(itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, localWriteEndIter)
|
|
lastLoadIter = schedGlobalRead(writer, itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, schedNumForIter0, endIter)
|
|
# Schedule local write
|
|
if not writer.states.scheduleLocalWrite:
|
|
noSchedLocalWrite(writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter)
|
|
else:
|
|
itemsLWToSched, numWritesToSched = prepareLWInstToSched(writer, kernel, numLocalWritesPerSched)
|
|
startIter = assignLWSchedIndexDefault(writer, kernel, numLocalWritesPerSched, localWriteEndIter, lastLoadIter, numWritesToSched)
|
|
readsToWait, readsToWaitNGLL = getReadsToWait(writer, kernel)
|
|
schedLocalWrite(writer, kernel, numLocalWriteModPerIter, numLocalWritesPerSched, localWriteEndIter, \
|
|
itemsGRToSchedLater, itemsLWToSched, startIter, readsToWait, readsToWaitNGLL, \
|
|
firstIter, lastLc, isNGLL)
|
|
|
|
class SIA0(SIA):
|
|
kernel = {"_ScheduleIterAlg": 0}
|
|
def __call__(self):
|
|
assert False
|
|
|
|
def schedIntoIteration(self, writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter, firstIter, lastLoop, lastLc, globalReadIncACode, globalReadIncBCode, isNGLL):
|
|
# Get schedule information
|
|
numGlobalReadInsPerIter, numLocalWriteModPerIter, numEmptyGlobalReadIncCode = getScheduleParams(kernel)
|
|
numLocalWritesPerSched = numLocalWriteModPerIter
|
|
itemsGRToSchedLater, lastLoadIter = noSchedGlobalRead(writer, kernel, globalReadIncACode, globalReadIncBCode)
|
|
# Schedule local write
|
|
noSchedLocalWrite(writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter)
|
|
|
|
################################################################################
|
|
################################################################################
|
|
###
|
|
### Schedule Parameters
|
|
###
|
|
################################################################################
|
|
################################################################################
|
|
|
|
def checkLocalReadFIFO(localReadFIFO, miLatency, numWaves, numMFMA, blockWidth):
|
|
# Add space to avoid LR FIFO stall
|
|
# lrStallLatencyBuffer:
|
|
# 40 quad-cycle - 4 x miLatency for b128
|
|
# 20 quad-cycle - 4 x miLatency for b64 (equal to one miLatency)
|
|
# 10 quad-cycle - 4 x miLatency for b32 (no stall)
|
|
# so no stall happen for b64/b32/b16
|
|
if blockWidth < 4:
|
|
return False
|
|
# The FIFO length is 16 so that each wave has 16/numWaves buffer.
|
|
lrStallLatencyBuffer = roundUp(blockWidth) * 10 - ((16 / numWaves) * miLatency)
|
|
if len(localReadFIFO) < (16 / numWaves):
|
|
localReadFIFO.append(numMFMA)
|
|
else:
|
|
oldNumMFMA = localReadFIFO[0]
|
|
if (numMFMA - oldNumMFMA) * miLatency >= lrStallLatencyBuffer:
|
|
localReadFIFO.pop(0)
|
|
localReadFIFO.append(numMFMA)
|
|
else:
|
|
# FIFO is full
|
|
return True
|
|
return False
|
|
|
|
def getScheduleParams(kernel):
|
|
numGlobalReadInsPerIter = roundUp(kernel["GlobalReadPerMfma"] * PRECISION) if kernel["GlobalReadPerMfma"] > 0 else PRECISION
|
|
numLocalWriteModPerIter = roundUp(kernel["LocalWritePerMfma"] * PRECISION) if kernel["LocalWritePerMfma"] > 0 else PRECISION
|
|
numEmptyGlobalReadIncCode = numGlobalReadInsPerIter - 1
|
|
return numGlobalReadInsPerIter, numLocalWriteModPerIter, numEmptyGlobalReadIncCode
|
|
|
|
def getLocalWriteMFMAEnd(writer, kernel, tensorParametersA, tensorParametersB):
|
|
numMfmaPerIter = writer.states.numMfmaPerIter
|
|
#########
|
|
# Get localWriteEnd
|
|
#########
|
|
# assign parameter
|
|
# 1. we calculate number of mfma to prefetch localReads for next loop
|
|
# 2. we put a barrier before the last mfma
|
|
# 3. we put last localWrite before 2~3 mfma, then the barrier
|
|
# localReads followed following sequence to be scheduled
|
|
# ds_read[A][0], ds_read[B][0], ds_read[A][1:], ds_read[B][1:]
|
|
# NOTE: we need this sequence for new feature "breaking waitcnt"
|
|
# TODO: breaking waitcnt
|
|
writer.states.numMfmaForLR = 1
|
|
latencyLeft = writer.states.miLatencyLeft
|
|
miLatencyLeft = writer.states.miLatencyLeft
|
|
tPM = tensorParametersA["tpsMetadata"] if tensorParametersA["is_sparse"] else tensorParametersB["tpsMetadata"]
|
|
|
|
# we can skip some LR waitcnt
|
|
# Since the first mfma only use B[:1], so we only wait for B[0]
|
|
isLocalReadsOpt = False
|
|
tmpLatencyLeft = 0
|
|
tmpNumMfmaForLR = 0
|
|
localReadsNotWaited = writer.states.numReadsPerIterB // kernel["InnerUnroll"] - writer.states.numReadsPerUnrollB
|
|
if kernel["UnrollMajorLDSB"] and localReadsNotWaited > 0:
|
|
isLocalReadsOpt = True
|
|
|
|
# check instruction FIFO
|
|
localReadFIFO = []
|
|
numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] * kernel["LocalSplitU"]
|
|
|
|
def calculateLatencyLeft(numReads, localReadBlockWidth, localReadLatency):
|
|
nonlocal localReadFIFO, numWaves, latencyLeft, miLatencyLeft
|
|
for _ in range(numReads):
|
|
while(checkLocalReadFIFO(localReadFIFO, writer.states.miLatency, numWaves, writer.states.numMfmaForLR, localReadBlockWidth)):
|
|
writer.states.numMfmaForLR += 1
|
|
latencyLeft -= localReadLatency*2
|
|
if latencyLeft < 0:
|
|
writer.states.numMfmaForLR += 1
|
|
latencyLeft = max(miLatencyLeft - localReadLatency*2,0)
|
|
|
|
for iui in range(kernel["InnerUnroll"]):
|
|
# ds_read[A][0]
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollA, tensorParametersA["localReadInstruction"].blockWidth, tensorParametersA["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSA][0]
|
|
if kernel["ProblemType"]["MXBlockA"]:
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollMXSA, tensorParametersA["MX"]["localReadInstruction"].blockWidth, tensorParametersA["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSB][0]
|
|
if kernel["ProblemType"]["MXBlockB"]:
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollMXSB, tensorParametersB["MX"]["localReadInstruction"].blockWidth, tensorParametersB["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[M][0]
|
|
if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]:
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollMetadata, tPM["localReadInstruction"].blockWidth, tPM["localReadInstruction"].issueLatency)
|
|
# ds_read[B][0]
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollB, tensorParametersB["localReadInstruction"].blockWidth, tensorParametersB["localReadInstruction"].issueLatency)
|
|
|
|
# ds_read[A][1:]
|
|
calculateLatencyLeft((writer.states.numReadsPerIterA//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollA), tensorParametersA["localReadInstruction"].blockWidth, tensorParametersA["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSA][1:]
|
|
if kernel["ProblemType"]["MXBlockA"]:
|
|
calculateLatencyLeft(writer.states.numReadsPerIterMXSA//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollMXSA, tensorParametersA["MX"]["localReadInstruction"].blockWidth, tensorParametersA["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[M][1:]
|
|
if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]:
|
|
calculateLatencyLeft((writer.states.numReadsPerIterMetadata//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollMetadata), tPM["localReadInstruction"].blockWidth, tPM["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSB][1:]
|
|
if kernel["ProblemType"]["MXBlockB"]:
|
|
calculateLatencyLeft((writer.states.numReadsPerIterMXSB//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollMXSB), tensorParametersB["MX"]["localReadInstruction"].blockWidth, tensorParametersB["MX"]["localReadInstruction"].issueLatency)
|
|
# get the latency before B[:1]
|
|
if isLocalReadsOpt:
|
|
tmpLatencyLeft = latencyLeft
|
|
tmpNumMfmaForLR = writer.states.numMfmaForLR
|
|
# ds_read[B][1:]
|
|
calculateLatencyLeft((writer.states.numReadsPerIterB//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollB), tensorParametersB["localReadInstruction"].blockWidth, tensorParametersB["localReadInstruction"].issueLatency)
|
|
|
|
# to calculate number of mfma we need to wait before data arrive from lds to vgpr.
|
|
# latency: 40 quad-cycle for 4 word, 20 quad-cycle for 2 word, 10 quad-cycle for 1 word / half word
|
|
writer.states.numMfmaForNextLoopLR = writer.states.numMfmaForLR
|
|
latencyForLR = roundUp(tensorParametersB["localReadInstruction"].blockWidth)*10
|
|
if isLocalReadsOpt:
|
|
latencyForLR = min(latencyForLR, (writer.states.numMfmaForLR - tmpNumMfmaForLR) * writer.states.miLatency)
|
|
latencyForLR -= max(tmpLatencyLeft,0) # remaining latency in mfma
|
|
else:
|
|
latencyForLR -= max(latencyLeft,0) # remaining latency in mfma
|
|
latencyForLR -= writer.states.miLatency # last LR will have 1 mfma latency
|
|
# add extra latency
|
|
latencyForLR += kernel["ExtraLatencyForLR"]
|
|
while latencyForLR > 0:
|
|
latencyForLR -= writer.states.miLatency
|
|
writer.states.numMfmaForNextLoopLR += 1
|
|
# final index definition
|
|
writer.states.numMfmaForNextLoopLR = min(writer.states.numMfmaForNextLoopLR,numMfmaPerIter-1)
|
|
writer.states.syncPlrMfmaIndex = numMfmaPerIter*(kernel["LoopIters"]-writer.states.numItersPLR+1) - writer.states.numMfmaForNextLoopLR - 1 if writer.states.numItersPLR else 0
|
|
if writer.states.doFullPackCodePrefetch and kernel["ForceUnrollSubIter"]:
|
|
# move syncPlrMfmaIndex one iteration ahead for doFullPackCodePrefetch and subIter
|
|
writer.states.syncPlrMfmaIndex = max(0, writer.states.syncPlrMfmaIndex - numMfmaPerIter)
|
|
elif writer.states.doPackPreSchedulingNextLoop:
|
|
# doPackPreSchedulingNextLoop only case (not doFullPackCodePrefetch), move syncPlrMfmaIndex to the top of
|
|
writer.states.syncPlrMfmaIndex = numMfmaPerIter*(kernel["LoopIters"]-writer.states.numItersPLR)
|
|
|
|
numMfmaBetweenLWandBarrier = 2 if kernel["MatrixInstM"] == 32 else 3
|
|
if kernel["NoLdsWriteCode"]:
|
|
# no interval needed in no ds_write case
|
|
numMfmaBetweenLWandBarrier = 0
|
|
if writer.states.scheduleGROverBarrier or writer.states.numItersPLR == 0:
|
|
writer.states.lwEndMfmaIndex = writer.states.numMfmaPerIter*kernel["LoopIters"] - 1
|
|
else:
|
|
writer.states.lwEndMfmaIndex = max(writer.states.syncPlrMfmaIndex - numMfmaBetweenLWandBarrier,0)
|
|
return numMfmaBetweenLWandBarrier, latencyLeft
|
|
|
|
def getLocalWriteMFMAStart(writer, kernel, tensorParametersA, tensorParametersB, latencyLeft):
|
|
numMfmaPerIter = writer.states.numMfmaPerIter
|
|
|
|
tPM = tensorParametersA["tpsMetadata"] if tensorParametersA["is_sparse"] else tensorParametersB["tpsMetadata"]
|
|
#########
|
|
# Get localWriteStart
|
|
#########
|
|
if not writer.states.oneBufferScheduling:
|
|
# TODO: replace here for real number of globalReadIncInst
|
|
# numGRIncInst = 18 # Always on. Original logic: 12 if not kernel["StaggerU"] else 18
|
|
numGRIncInst = 12 if not writer.states.staggerUCode else 18
|
|
if kernel["ProblemType"]["MXBlockA"] and kernel["ProblemType"]["MXBlockB"]:
|
|
# MXSA+MXSB case, double the number
|
|
numGRIncInst *= 2
|
|
numInstPerMfma = max(roundUp(writer.states.miLatencyLeft/2),1)
|
|
numMfmaToSched = roundUp(numGRIncInst/numInstPerMfma)
|
|
lwStartMfmaIndex = 1 + numMfmaToSched
|
|
else:
|
|
# check instruction FIFO
|
|
localReadFIFO = []
|
|
numWaves = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1] * kernel["LocalSplitU"]
|
|
# for 1LDSB, we have to issue localwrites after localreads
|
|
if kernel["ClusterLocalRead"]:
|
|
# we have enough vgprBuffer to schedule localReads in the front of loop
|
|
numMfmaForCurrentLoopLR = 1
|
|
latencyLeft = writer.states.miLatencyLeft
|
|
subIter = kernel["ForceUnrollSubIter"]
|
|
# subIter case, LoopIters is increased to 4, but we should use 1 for latency calculation
|
|
lookRange = 1 if subIter else kernel["LoopIters"] - writer.states.numItersPLR
|
|
for u in range(lookRange):
|
|
doReadA = ((u < kernel["LoopIters"] // writer.states.numIterPerCoalescedReadA - writer.states.numItersPLR) and not kernel["DirectToVgprA"]) or subIter
|
|
doReadMXSA = False
|
|
if kernel["ProblemType"]["MXBlockA"] > 0:
|
|
doReadMXSA = ((u < kernel["LoopIters"] // writer.states.numIterPerCoalescedReadMXSA - writer.states.numItersPLR) and not kernel["DirectToVgprMXSA"]) or subIter
|
|
doReadB = ((u < kernel["LoopIters"] // writer.states.numIterPerCoalescedReadB - writer.states.numItersPLR) and not kernel["DirectToVgprB"]) or subIter
|
|
doReadMXSB = False
|
|
if kernel["ProblemType"]["MXBlockB"] > 0:
|
|
doReadMXSB = ((u < kernel["LoopIters"] // writer.states.numIterPerCoalescedReadMXSB - writer.states.numItersPLR) and not kernel["DirectToVgprMXSB"]) or subIter
|
|
doReadM = ((u < kernel["LoopIters"] // writer.states.numIterPerCoalescedReadMetadata - writer.states.numItersPLR))
|
|
doReadMXSA = doReadMXSA and kernel["ProblemType"]["MXBlockA"] and (not kernel["DirectToVgprMXSA"])
|
|
doReadMXSB = doReadMXSB and kernel["ProblemType"]["MXBlockB"] and (not kernel["DirectToVgprMXSB"])
|
|
doReadM = doReadM and (kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"])
|
|
|
|
def calculateLatencyLeft(numReads, localReadBlockWidth, localReadLatency):
|
|
nonlocal localReadFIFO, numWaves, latencyLeft, numMfmaForCurrentLoopLR
|
|
for _ in range(numReads):
|
|
while(checkLocalReadFIFO(localReadFIFO, writer.states.miLatency, numWaves, numMfmaForCurrentLoopLR, localReadBlockWidth)):
|
|
numMfmaForCurrentLoopLR += 1
|
|
latencyLeft -= localReadLatency*2
|
|
if latencyLeft < 0:
|
|
numMfmaForCurrentLoopLR += 1
|
|
latencyLeft = max(writer.states.miLatencyLeft - localReadLatency*2,0)
|
|
|
|
for iui in range(kernel["InnerUnroll"]):
|
|
# ds_read[A][0]
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollA * doReadA, tensorParametersA["localReadInstruction"].blockWidth, tensorParametersA["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSA][0]
|
|
if doReadMXSA:
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollMXSA * doReadMXSA, tensorParametersA["MX"]["localReadInstruction"].blockWidth, tensorParametersA["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[M][0]
|
|
if doReadM:
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollMetadata * doReadM, tPM["localReadInstruction"].blockWidth, tPM["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSB][0]
|
|
if doReadMXSB:
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollMXSB * doReadMXSB, tensorParametersB["MX"]["localReadInstruction"].blockWidth, tensorParametersB["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[B][0]
|
|
calculateLatencyLeft(writer.states.numReadsPerUnrollB * doReadB, tensorParametersB["localReadInstruction"].blockWidth, tensorParametersB["localReadInstruction"].issueLatency)
|
|
|
|
# ds_read[A][1:]
|
|
calculateLatencyLeft((writer.states.numReadsPerIterA//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollA) * doReadA, tensorParametersA["localReadInstruction"].blockWidth, tensorParametersA["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSA][1:]
|
|
if doReadMXSA:
|
|
calculateLatencyLeft((writer.states.numReadsPerIterMXSA//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollMXSA) * doReadMXSA, tensorParametersA["MX"]["localReadInstruction"].blockWidth, tensorParametersA["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[M][1:]
|
|
if doReadM:
|
|
calculateLatencyLeft((writer.states.numReadsPerIterMetadata//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollMetadata) * doReadM, tPM["localReadInstruction"].blockWidth, tPM["localReadInstruction"].issueLatency)
|
|
# ds_read[MXSB][1:]
|
|
if doReadMXSB:
|
|
calculateLatencyLeft((writer.states.numReadsPerIterMXSB//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollMXSB) * doReadMXSB, tensorParametersB["MX"]["localReadInstruction"].blockWidth, tensorParametersB["MX"]["localReadInstruction"].issueLatency)
|
|
# ds_read[B][1:]
|
|
calculateLatencyLeft((writer.states.numReadsPerIterB//kernel["InnerUnroll"] - writer.states.numReadsPerUnrollB) * doReadB, tensorParametersB["localReadInstruction"].blockWidth, tensorParametersB["localReadInstruction"].issueLatency)
|
|
lwStartMfmaIndex = numMfmaForCurrentLoopLR
|
|
else:
|
|
lwStartMfmaIndex = numMfmaPerIter * (kernel["LoopIters"] - 1 - writer.states.numItersPLR) + writer.states.numMfmaForLR
|
|
# to calculate number of mfma we need to wait before data arrive from lds to vgpr.
|
|
# latency: 40 quad-cycle for 4 word, 20 quad-cycle for 2 word, 10 quad-cycle for 1 word / half word
|
|
if writer.states.numIterPerCoalescedReadB > writer.states.numIterPerCoalescedReadA:
|
|
latencyForLR = roundUp(tensorParametersA["localReadInstruction"].blockWidth) * 10
|
|
else:
|
|
latencyForLR = roundUp(tensorParametersB["localReadInstruction"].blockWidth) * 10
|
|
latencyForLR -= max(latencyLeft,0) # remaining latency in mfma
|
|
while latencyForLR > 0:
|
|
latencyForLR -= writer.states.miLatency
|
|
lwStartMfmaIndex += 1
|
|
|
|
if lwStartMfmaIndex > writer.states.lwEndMfmaIndex:
|
|
lwStartMfmaIndex = writer.states.lwEndMfmaIndex
|
|
return lwStartMfmaIndex
|
|
|
|
def getNumLocalWritePerMfma(writer, kernel, lwStartMfmaIndex):
|
|
#########
|
|
# Get LocalWritePerMfma
|
|
#########
|
|
numMfmaCanSched = writer.states.lwEndMfmaIndex - lwStartMfmaIndex + 1
|
|
numLoadsA = kernel["DepthU"]*kernel["MacroTileA"]//kernel["GlobalReadVectorWidthA"]//kernel["NumThreads"]
|
|
if kernel["ProblemType"]["MXBlockA"]:
|
|
numLoadsA += kernel["_DepthUMXSA"]*kernel["MacroTileA"]//kernel["GlobalReadVectorWidthMXSA"]//kernel["NumThreads"]
|
|
numLoadsB = kernel["DepthU"]*kernel["MacroTileB"]//kernel["GlobalReadVectorWidthB"]//kernel["NumThreads"]
|
|
if kernel["ProblemType"]["MXBlockB"]:
|
|
numLoadsB += kernel["_DepthUMXSB"]*kernel["MacroTileB"]//kernel["GlobalReadVectorWidthMXSB"]//kernel["NumThreads"]
|
|
if kernel["ProblemType"]["Sparse"] and not kernel["DirectToVgprSparseMetadata"]:
|
|
macroTile = kernel["MacroTileB"] if kernel["ProblemType"]["Sparse"] == 2 else kernel["MacroTileA"]
|
|
numLoadsM = kernel["DepthU"]*macroTile//kernel["GlobalReadVectorWidthMetadata"]//kernel["NumThreads"]
|
|
else:
|
|
numLoadsM = 0
|
|
writesToSched = (numLoadsA + numLoadsB + numLoadsM- 1) * PRECISION
|
|
oldValue = 0
|
|
newValue = PRECISION
|
|
loop = 0
|
|
# 1. number of padded writesToSched is (numWrites - 1) * 100 + 1
|
|
# LW ---99--- LW ---99--- LW
|
|
# 2. we need to pad it to multiple of LWPM
|
|
# LW ---99--- LW ---99--- LW --?--
|
|
# | ------- multiple of LWPM ---- |
|
|
# 3. if LWPM is not multiple of 100, we need extra empty instructions to schedule GR for PGR2
|
|
# LW ---99--- LW ---99--- LW --?-- --?--
|
|
# | ------- multiple of LWPM ---- |-LWPM-|
|
|
# 4. then we put GR into padded writesToSched
|
|
# put GR after LW + LWPM of empty inst, so that we can offset GR 1 mfma with LW if possible
|
|
# Ex. LWPM = 0.25
|
|
# LW --24- GR ------74------ LW --24- GR ------74------ LW --24- GR --24-
|
|
# mfma--24-mfma--24-mfma--24-mfma--24-mfma--24-mfma--24-mfma--24-mfma--24-mfma
|
|
# we need LWPM to get precise LWPM
|
|
# so we iterate formula 10 times to get LWPM
|
|
while oldValue != newValue and newValue > 0 and loop < 10:
|
|
loop += 1
|
|
oldValue = newValue
|
|
newValue = roundUp((writesToSched+1 + (oldValue - (writesToSched+1) % oldValue) + oldValue%PRECISION) / numMfmaCanSched)
|
|
return newValue
|
|
|
|
def calculateGRPMandLWPM(writer, kernel, numLocalWriteModPerMfma) -> Tuple[int, int]:
|
|
#####
|
|
# Assign GRPM and LWPM
|
|
#####
|
|
# HOW THIS WORK
|
|
# padding each globalReadInstruction to 100 with empty instruction,
|
|
# each mfma will schedule intructions GRPM*100 times from padded globalReadInstruction.
|
|
# Ex. GRPM = 0.5
|
|
# GR ---------99--------- GR --------99---------- GR
|
|
# mfma --49-- mfma --49-- mfma --49-- mfma --49-- mfma --49--
|
|
numGlobalReadInsPerMfma = roundUp(kernel["GlobalReadPerMfma"]*PRECISION)
|
|
# HOW THIS WORK
|
|
# padding each globalReadInstruction to 100 with empty instruction,
|
|
# each mfma will schedule intructions GRPM*100 times from padded globalReadInstruction.
|
|
# Ex. LWPM = 0.5
|
|
# LW ---------99--------- LW --------99---------- LW
|
|
# mfma --49-- mfma --49-- mfma --49-- mfma --49-- mfma --49--
|
|
if kernel["PrefetchGlobalRead"] == 1:
|
|
# In PGR1:
|
|
# Larger LWPM can provide more latency to hide global read
|
|
# However, larger LWPM may cause mfma bubbles
|
|
# we set LWPM to 1 unless it requires larger LWPM to enable 1LDSB
|
|
if kernel["1LDSBuffer"]:
|
|
numLocalWriteModPerMfma = max(numLocalWriteModPerMfma, PRECISION)
|
|
else:
|
|
numLocalWriteModPerMfma = PRECISION
|
|
|
|
return numGlobalReadInsPerMfma, numLocalWriteModPerMfma
|
|
|
|
def getScheduleParamMfma(writer):
|
|
numMfmaPerIter = writer.states.numMfmaPerIter
|
|
numGlobalReadInsPerIter = numMfmaPerIter * writer.states.numGlobalReadInsPerMfma
|
|
numLocalWriteModPerIter = numMfmaPerIter * writer.states.numLocalWriteModPerMfma
|
|
# if numGlobalReadInsPerMfma>1, we still want to schedule only 1 GlobalReadIncCode per mfma
|
|
# inserting empty CodeModule so that generator will schedule 1 GlobalReadIncCode 1 empty CodeModule if numGlobalReadInsPerMfma=2
|
|
numEmptyGlobalReadIncCode = writer.states.numGlobalReadInsPerMfma - 1
|
|
return numGlobalReadInsPerIter, numLocalWriteModPerIter, numEmptyGlobalReadIncCode
|
|
|
|
def fixLocalWriteEndMfmaIndex(writer, kernel, tPA, tPB, globalReadIncACode, globalReadIncBCode, numMfmaBetweenLWandBarrier, lastLoop):
|
|
numMfmaPerIter = writer.states.numMfmaPerIter
|
|
# If numLocalWriteModPerMfma is not multiple of 100,
|
|
# last globalread will be scheduled at lwEndMfmaIndex,
|
|
# and last localwrite will be scheduled at lwEndMfmaIndex - 1
|
|
# so we offset lwEndMfmaIndex by 1 mfma
|
|
if kernel["PrefetchGlobalRead"] >= 2 and writer.states.numLocalWriteModPerMfma % PRECISION != 0 and not kernel["NoLdsWriteCode"]:
|
|
numMfmaBetweenLWandBarrier -= 1
|
|
|
|
if writer.states.scheduleGROverBarrier or writer.states.numItersPLR == 0:
|
|
writer.states.lwEndMfmaIndex = numMfmaPerIter*kernel["LoopIters"] - 1
|
|
else:
|
|
writer.states.lwEndMfmaIndex = max(writer.states.syncPlrMfmaIndex - numMfmaBetweenLWandBarrier,0)
|
|
localWriteEndIter = writer.states.lwEndMfmaIndex//numMfmaPerIter
|
|
localWriteEndIter = min(kernel["LoopIters"] - 1, localWriteEndIter)
|
|
assert localWriteEndIter < kernel["LoopIters"]
|
|
assert writer.states.lwEndMfmaIndex < numMfmaPerIter*kernel["LoopIters"]
|
|
return localWriteEndIter
|
|
|
|
################################################################################
|
|
################################################################################
|
|
###
|
|
### Schedule Global Read
|
|
###
|
|
################################################################################
|
|
################################################################################
|
|
|
|
def _splitTdmLoad(grCode):
|
|
"""Split a globalRead module into (non-TDM items module, TDM-load-only module)."""
|
|
tdmLoadMod = Module("deferredTdmLoad")
|
|
nonTdmMod = Module()
|
|
grAItems = grCode.items() if hasattr(grCode, 'items') else []
|
|
for item in grAItems:
|
|
hasTdm = isinstance(item, TensorLoadToLds) or \
|
|
(hasattr(item, 'flatitems') and any(isinstance(f, TensorLoadToLds) for f in item.flatitems()))
|
|
if hasTdm:
|
|
tdmLoadMod.add(item)
|
|
else:
|
|
nonTdmMod.add(item)
|
|
return nonTdmMod, tdmLoadMod
|
|
|
|
|
|
def noSchedGlobalRead(writer, kernel, globalReadIncACode, globalReadIncBCode):
|
|
# For TDM+SIA0: defer tensor_load_to_lds to after the LDS swap to avoid writing
|
|
# into the LDS buffer still being read by ds_loads
|
|
tdmDeferLoad = kernel["enableTDMA"] and kernel["enableTDMB"]
|
|
localWriteEndIter = kernel["LoopIters"] - writer.states.numItersPLR - 1
|
|
tdmLoadIter = min(localWriteEndIter + 1, kernel["LoopIters"] - 1)
|
|
|
|
if kernel["PrefetchGlobalRead"] == 2:
|
|
# SIA0 does not schedule GR/LW instruction-by-instruction. If global reads
|
|
# are emitted at iter 0, they clobber vgprG2L* before the later local-write
|
|
# iteration stores the previous prefetch to LDS. Place non-TDM reads in the
|
|
# local-write iteration; SIA0 sub-iteration order emits localWriteCode before
|
|
# globalReadCode, preserving the G2L value until it is consumed.
|
|
grIter = localWriteEndIter if kernel["_ScheduleIterAlg"] == 0 and not kernel["NoLdsWriteCode"] and not tdmDeferLoad else 0
|
|
imod = writer.codes.perIterGlobalRead[grIter].add(Module())
|
|
imod.addComment1("Global Read IncA")
|
|
imod.add(globalReadIncACode)
|
|
imod.addComment1("Global Read IncB")
|
|
imod.add(globalReadIncBCode)
|
|
if tdmDeferLoad:
|
|
imod.addComment1("Global Read A")
|
|
imod.add(writer.codes.dtlsM0UpdateA)
|
|
nonTdmMod, tdmLoadModA = _splitTdmLoad(writer.codes.globalReadA)
|
|
imod.add(nonTdmMod)
|
|
imod.addComment1("Global Read MXSA")
|
|
imod.add(writer.codes.dtlsM0UpdateMXSA)
|
|
# MXSA/MXSB share A/B's XOR double-buffered MX-scale LDS, so their
|
|
# tensor_load_to_lds must defer in lockstep with A/B (after the LDS swap).
|
|
# Emitting them eagerly lets the async TDM overwrite the MX-scale LDS while
|
|
# this wave's own MX ds_loads are still reading it (self-wave WAR; s_wait_dscnt
|
|
# does not gate TDM, which retires on tensorcnt).
|
|
nonTdmModMXSA, tdmLoadModMXSA = _splitTdmLoad(writer.codes.globalReadMXSA)
|
|
imod.add(nonTdmModMXSA)
|
|
imod.addComment1("Global Read MXSB")
|
|
imod.add(writer.codes.dtlsM0UpdateMXSB)
|
|
nonTdmModMXSB, tdmLoadModMXSB = _splitTdmLoad(writer.codes.globalReadMXSB)
|
|
imod.add(nonTdmModMXSB)
|
|
imod.addComment1("Global Read B")
|
|
imod.add(writer.codes.dtlsM0UpdateB)
|
|
nonTdmModB, tdmLoadModB = _splitTdmLoad(writer.codes.globalReadB)
|
|
imod.add(nonTdmModB)
|
|
# Defer both A and B tensor_load_to_lds to after the LDS swap (at tdmLoadIter).
|
|
# This ensures both loads go into the post-swap buffer, matching the ds_load reads.
|
|
deferMod = writer.codes.perIterGlobalRead[tdmLoadIter].add(Module())
|
|
if tdmLoadModA.itemsSize() > 0:
|
|
deferMod.addComment1("Global Read A (TDM deferred after LDS swap)")
|
|
deferMod.add(tdmLoadModA)
|
|
if tdmLoadModMXSA.itemsSize() > 0:
|
|
deferMod.addComment1("Global Read MXSA (TDM deferred after LDS swap)")
|
|
deferMod.add(tdmLoadModMXSA)
|
|
if tdmLoadModMXSB.itemsSize() > 0:
|
|
deferMod.addComment1("Global Read MXSB (TDM deferred after LDS swap)")
|
|
deferMod.add(tdmLoadModMXSB)
|
|
if tdmLoadModB.itemsSize() > 0:
|
|
# TODO: For the 1-wave case, schedule B's TDM load independently for better tensor load balance.
|
|
deferMod.addComment1("Global Read B (TDM deferred after LDS swap)")
|
|
deferMod.add(tdmLoadModB)
|
|
# Metadata TDM also writes to a double-buffered LDS region that is XOR-swapped
|
|
# alongside A's (see s_xor on sgprtdmMetadataGroup0+1). Defer its tensor_load_to_lds
|
|
# to tdmLoadIter so the load lands in the post-swap buffer, matching the
|
|
# ds_load_tr8_b64 reads from the new side; otherwise the async TDM write races
|
|
# against the current iter's metadata local reads from the same LDS region.
|
|
if kernel["ProblemType"]["Sparse"]:
|
|
nonTdmModM, tdmLoadModM = _splitTdmLoad(writer.codes.globalReadMetadata)
|
|
imod.addComment1("Global Read Metadata")
|
|
imod.add(nonTdmModM)
|
|
if tdmLoadModM.itemsSize() > 0:
|
|
deferMod.addComment1("Global Read Metadata (TDM deferred after LDS swap)")
|
|
deferMod.add(tdmLoadModM)
|
|
imod.add(writer.codes.gl2PrefetchIncrement)
|
|
imod.add(writer.codes.gl2Prefetch)
|
|
else:
|
|
imod.addComment1("Global Read A")
|
|
imod.add(writer.codes.dtlsM0UpdateA)
|
|
imod.add(writer.codes.globalReadA)
|
|
imod.addComment1("Global Read MXSA")
|
|
imod.add(writer.codes.dtlsM0UpdateMXSA)
|
|
imod.add(writer.codes.globalReadMXSA)
|
|
imod.addComment1("Global Read MXSB")
|
|
imod.add(writer.codes.dtlsM0UpdateMXSB)
|
|
imod.add(writer.codes.globalReadMXSB)
|
|
imod.addComment1("Global Read B")
|
|
imod.add(writer.codes.dtlsM0UpdateB)
|
|
imod.add(writer.codes.globalReadB)
|
|
if kernel["ProblemType"]["Sparse"]:
|
|
imod.addComment1("Global Read Metadata")
|
|
imod.add(writer.codes.globalReadMetadata)
|
|
imod.add(writer.codes.gl2PrefetchIncrement)
|
|
imod.add(writer.codes.gl2Prefetch)
|
|
else:
|
|
# put everything in the header (original behavior for PGR=0/1):
|
|
writer.codes.unrollLoopHeader.add(writer.codes.dtlsM0UpdateA)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadA)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.dtlsM0UpdateMXSA)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadMXSA)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.dtlsM0UpdateMXSB)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadMXSB)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.dtlsM0UpdateB)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadB)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadMetadata) if kernel["ProblemType"]["Sparse"] else None
|
|
writer.codes.unrollLoopHeader.add(globalReadIncACode)
|
|
writer.codes.unrollLoopHeader.add(globalReadIncBCode)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.gl2PrefetchIncrement)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.gl2Prefetch)
|
|
# Dummy
|
|
itemsGRToSchedLater = []
|
|
lastLoadIter = 0
|
|
return itemsGRToSchedLater, lastLoadIter
|
|
|
|
def prepareGRInstToSched(writer, kernel, isNGLL):
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadA.header)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadMXSA.header)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadMXSB.header)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadB.header)
|
|
writer.codes.unrollLoopHeader.add(writer.codes.globalReadMetadata.header) if kernel["ProblemType"]["Sparse"] else None
|
|
|
|
# Add all loads from middle as individual schedulable items
|
|
# when using PGR2, put global read instruction right after corresponding localWrite instruction
|
|
if isNGLL and (kernel["UnrollLoopSwapGlobalReadOrder"] == 1 and not (kernel["DirectToLdsA"] and kernel["DirectToLdsB"])):
|
|
itemsGRToSched = []
|
|
itemsGRToSchedLater = []
|
|
elif kernel["PrefetchGlobalRead"] >= 2:
|
|
itemsGRToSched = []
|
|
itemsGRToSchedLater = list(writer.codes.globalReadA.middle.items()) + \
|
|
list(writer.codes.globalReadMXSA.middle.items()) + \
|
|
list(writer.codes.globalReadMXSB.middle.items()) + \
|
|
list(writer.codes.globalReadB.middle.items()) + \
|
|
list(writer.codes.globalReadMetadata.middle.items())
|
|
else:
|
|
itemsGRToSched = list(writer.codes.globalReadA.middle.items()) + \
|
|
list(writer.codes.globalReadMXSA.middle.items()) + \
|
|
list(writer.codes.globalReadMXSB.middle.items()) + \
|
|
list(writer.codes.globalReadB.middle.items()) + \
|
|
list(writer.codes.globalReadMetadata.middle.items())
|
|
itemsGRToSchedLater = []
|
|
|
|
itemsGRToSchedTemp = []
|
|
for i in range(len(itemsGRToSched)):
|
|
itemsGRToSchedTemp.append(itemsGRToSched.pop(0))
|
|
for j in range(PRECISION-1):
|
|
itemsGRToSchedTemp.append(Module())
|
|
itemsGRToSched = itemsGRToSchedTemp
|
|
return itemsGRToSched, itemsGRToSchedLater
|
|
|
|
|
|
def appendInstToSchedSIA3(writer, kernel, numEmptyGlobalReadIncCode, globalReadIncACode, globalReadIncBCode):
|
|
itemsGRIncToSched = []
|
|
# for SIA3, we can break GlobalReadIncCode to avoid mfma bubbles
|
|
if kernel["PrefetchGlobalRead"] >= 2:
|
|
# skip to schedule global read for PGR2 first mfma
|
|
for i in range(numEmptyGlobalReadIncCode+1):
|
|
imod = Module()
|
|
itemsGRIncToSched.append(imod)
|
|
numInst = countInstruction(globalReadIncACode) + countInstruction(globalReadIncBCode)
|
|
numInstPerMfma = writer.states.numInstPerMfma
|
|
|
|
globalReadIncItems = globalReadIncACode.flatitems() + globalReadIncBCode.flatitems()
|
|
numMfmaToSched = roundUp(numInst/numInstPerMfma)
|
|
for j in range(numMfmaToSched):
|
|
imod = Module()
|
|
count = 0
|
|
while globalReadIncItems and count < numInstPerMfma:
|
|
tempInst = globalReadIncItems.pop(0)
|
|
imod.add(tempInst)
|
|
if countInstruction(tempInst):
|
|
count += 1
|
|
itemsGRIncToSched.append(imod)
|
|
for i in range(numEmptyGlobalReadIncCode):
|
|
imod = Module()
|
|
itemsGRIncToSched.append(imod)
|
|
return itemsGRIncToSched
|
|
|
|
def appendInstToSchedDefault(numEmptyGlobalReadIncCode, globalReadIncACode, globalReadIncBCode):
|
|
itemsGRIncToSched = []
|
|
itemsGRIncToSched.append(globalReadIncACode)
|
|
for i in range(numEmptyGlobalReadIncCode):
|
|
imod = Module()
|
|
itemsGRIncToSched.append(imod)
|
|
itemsGRIncToSched.append(globalReadIncBCode)
|
|
for i in range(numEmptyGlobalReadIncCode):
|
|
imod = Module()
|
|
itemsGRIncToSched.append(imod)
|
|
return itemsGRIncToSched
|
|
|
|
def getSchedNumForIter0SIA3(writer, kernel, itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter):
|
|
numMfmaPerIter = writer.states.numMfmaPerIter
|
|
# Loop in PGR1: GlobalRead -> GlobalReadInc -> LocalWrite
|
|
# but GlobalReadInc shouldn't block LocalWrite so we count them out
|
|
# Loop in PGR2: GlobalReadInc -> LocalWrite/GlobalRead pair
|
|
# since LocalWrite/GlobalRead pair depends on GlobalReadInc, we count in only GlobalReadInc
|
|
if kernel["PrefetchGlobalRead"] >= 2:
|
|
loadsToSched = len(itemsGRIncToSched)
|
|
else:
|
|
loadsToSched = len(itemsGRToSched)
|
|
|
|
# Here is to adjust scheduling silently in order to have validation pass.
|
|
# Better way is to use larger globalReadPerMfma.
|
|
## schedule more instructions at first iteration if no enough mfma to schedule globalRead
|
|
writer.states.grEndMfmaIndex = max(0, roundUp(loadsToSched/writer.states.numGlobalReadInsPerMfma) - 1)
|
|
endIndex = writer.states.lwEndMfmaIndex
|
|
if writer.states.scheduleGROverBarrier:
|
|
# scheduleGROverBarrier case, lwEnd can be after barrier sync
|
|
# grEnd must be before barrier sync in that case
|
|
endIndex = min(writer.states.lwEndMfmaIndex, writer.states.syncPlrMfmaIndex)
|
|
if writer.states.IncLdsBufSwitch:
|
|
# IncLdsBufSwitch case, we use s operation for local read swap and need to finish GR Inc before that
|
|
if endIndex > writer.states.numMfmaPerIter - 1:
|
|
# adjust endIndex
|
|
endIndex = writer.states.numMfmaPerIter - 1
|
|
if writer.states.grEndMfmaIndex > endIndex:
|
|
schedNumForIter0 = numGlobalReadInsPerIter + (writer.states.grEndMfmaIndex - endIndex) * writer.states.numGlobalReadInsPerMfma
|
|
writer.states.grEndMfmaIndex = endIndex
|
|
else:
|
|
schedNumForIter0 = numGlobalReadInsPerIter
|
|
if kernel["PrefetchGlobalRead"] == 1:
|
|
globalReadIncEndMfmaIndex = writer.states.grEndMfmaIndex + roundUp(len(itemsGRIncToSched)/writer.states.numGlobalReadInsPerMfma)
|
|
endIter = roundUp((globalReadIncEndMfmaIndex+1)/numMfmaPerIter)
|
|
else:
|
|
endIter = roundUp((writer.states.grEndMfmaIndex+1)/numMfmaPerIter)
|
|
## schedule more instructions at first iteration if no enough mfma to schedule globalRead + globalReadInc
|
|
if endIter > kernel["LoopIters"]:
|
|
endIter = kernel["LoopIters"]
|
|
if kernel["PrefetchGlobalRead"] == 1:
|
|
schedNumForIter0 += (globalReadIncEndMfmaIndex+1 - kernel["LoopIters"]*numMfmaPerIter) * writer.states.numGlobalReadInsPerMfma
|
|
return schedNumForIter0, endIter
|
|
|
|
# schedNumForIter0 SIA 1 or 2
|
|
# distribute the instructions in itemsGRToSched evenly as possible to iterations: perIterGlobalReadCode[0,endIter)
|
|
# last one is perIterGlobalReadCode[endIter-1],
|
|
# Ideally: endIter <= localWriteEndIter,
|
|
# then put M0 updateCode (if any) and first 'schedNumForIter0' GR-inst in perIterGlobalReadCode[0]
|
|
# put every numGlobalReadInsPerIter GR-insts in perIterGlobalReadCode[1]~[endIter-1]
|
|
# corner case: endIter > localWriteEndIter, set endIter = localWriteEndIter,in this case, schedNumForIter0 will > 1
|
|
# and perIterGlobalReadCode[0] would need to schedule more instructions
|
|
def getSchedNumForIter0Default(itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, localWriteEndIter):
|
|
# reads and incs are scheduled in iters range(0..endIter)
|
|
endIter = roundUp((len(itemsGRToSched) + len(itemsGRIncToSched)) / numGlobalReadInsPerIter)
|
|
# FIXME:
|
|
# above formula precisely count number of GR + GRInc
|
|
# however it has regression issue with tuned yaml with default GRPM.
|
|
# below formula follows old logic to add 2 to the instruction count, so it may has larger schedNumForIter0
|
|
# we should use above formula with GRPM tuning for better performance
|
|
# NOTE: both formula pass validation test
|
|
endIter = roundUp((len(itemsGRToSched) + len(itemsGRIncToSched) + 2*PRECISION) / numGlobalReadInsPerIter)
|
|
if endIter > localWriteEndIter:
|
|
# Front-load some of the buffer loads if we don't have enough loop iters:
|
|
# could use a different/smarter algorithm to space out the loads?
|
|
schedNumForIter0 = (endIter-(localWriteEndIter) + 1) * numGlobalReadInsPerIter
|
|
endIter = localWriteEndIter
|
|
else:
|
|
# schedule b2b for readCnt > 2 (True for bigger TT)
|
|
schedNumForIter0 = numGlobalReadInsPerIter
|
|
return schedNumForIter0, endIter
|
|
|
|
def schedGlobalRead(writer, itemsGRToSched, itemsGRIncToSched, numGlobalReadInsPerIter, schedNumForIter0, endIter):
|
|
# insert dtlsM0UpdateACode dtlsM0UpdateBCode code
|
|
if writer.codes.globalReadA.middle.items():
|
|
writer.codes.globalReadA.middle.getItem(0).add(writer.codes.dtlsM0UpdateA, 0)
|
|
if writer.codes.globalReadMXSA.middle.items():
|
|
writer.codes.globalReadMXSA.middle.getItem(0).add(writer.codes.dtlsM0UpdateMXSA, 0)
|
|
if writer.codes.globalReadMXSB.middle.items():
|
|
writer.codes.globalReadMXSB.middle.getItem(0).add(writer.codes.dtlsM0UpdateMXSB, 0)
|
|
if writer.codes.globalReadB.middle.items():
|
|
writer.codes.globalReadB.middle.getItem(0).add(writer.codes.dtlsM0UpdateB, 0)
|
|
# TODO: if metadata supports directToLds
|
|
# if writer.codes.globalReadMetadata.middle.items():
|
|
# writer.codes.globalReadMetadata.middle.getItem(0).add(writer.codes.dtlsM0UpdateMetadata, 0)
|
|
|
|
itemsGRToSched.extend(itemsGRIncToSched)
|
|
# append 'n' global load at a time
|
|
# append global load(S) first 'number of global load(s)' determined by schedNumForIter0
|
|
for item in itemsGRToSched[:schedNumForIter0]:
|
|
writer.codes.perIterGlobalRead[0].add(item)
|
|
itemsGRToSched = itemsGRToSched[schedNumForIter0:] # trim the scheduled GRs, do the rest in the following loop
|
|
|
|
lastLoadIter = 0
|
|
for u in range(1, endIter):
|
|
# append itemPerIter GR for each iteration,
|
|
# and trim the scheduled ones at the end of loop
|
|
itemPerIter = 1 * numGlobalReadInsPerIter
|
|
try:
|
|
for item in itemsGRToSched[:itemPerIter]:
|
|
writer.codes.perIterGlobalRead[u].add(item)
|
|
lastLoadIter = u
|
|
itemsGRToSched = itemsGRToSched[itemPerIter:]
|
|
except IndexError:
|
|
break # itemsGRToSched is 0-length, no code left to schedule
|
|
|
|
assert not itemsGRToSched # should have scheduled everything already, itemsGRToSched should be empty
|
|
|
|
writer.codes.perIterGlobalRead[endIter-1].add(writer.codes.globalReadA.footer)
|
|
writer.codes.perIterGlobalRead[endIter-1].add(writer.codes.globalReadMXSA.footer)
|
|
writer.codes.perIterGlobalRead[endIter-1].add(writer.codes.globalReadMXSB.footer)
|
|
writer.codes.perIterGlobalRead[endIter-1].add(writer.codes.globalReadB.footer)
|
|
writer.codes.perIterGlobalRead[endIter-1].add(writer.codes.globalReadMetadata.footer)
|
|
return lastLoadIter
|
|
|
|
################################################################################
|
|
################################################################################
|
|
###
|
|
### Schedule Local Write
|
|
###
|
|
################################################################################
|
|
################################################################################
|
|
|
|
def noSchedLocalWrite(writer, kernel, tensorParametersA, tensorParametersB, localWriteEndIter):
|
|
# if no scheduleLocalWrite - just add writes to localWritelocalWriteEndIter
|
|
# If PGR=0, writes have to be done immediately following the loads - no opportunity to schedule
|
|
# so don't add to schedule, these will be added separately and before the first iter
|
|
if kernel["PrefetchGlobalRead"] and not kernel["NoLdsWriteCode"]:
|
|
# do we need a module here? That would prevent these from being scheduled
|
|
imod = writer.codes.perIterLocalWrite[localWriteEndIter][1].add(Module())
|
|
imod.add(
|
|
writer._wait(kernel, tensorParametersA, tensorParametersB, 0, -1, -1, \
|
|
"1wait for global read"))
|
|
imod.addComment1("local write A")
|
|
imod.add(writer.codes.localWriteA)
|
|
imod.addComment1("local write MXSA")
|
|
imod.add(writer.codes.localWriteMXSA)
|
|
imod.addComment1("local write MXSB")
|
|
imod.add(writer.codes.localWriteMXSB)
|
|
imod.addComment1("local write B")
|
|
imod.add(writer.codes.localWriteB)
|
|
|
|
if kernel["PrefetchGlobalRead"] == 2 and writer.codes.perIterLocalWriteCodeNGLL is not None: # TODO: check condition
|
|
# do we need a module here? That would prevent these from being scheduled
|
|
imod = writer.codes.perIterLocalWriteCodeNGLL[localWriteEndIter][1].add(Module())
|
|
imod.add(
|
|
writer._wait(kernel, tensorParametersA, tensorParametersB, 0, -1, -1, \
|
|
"1wait for global read"))
|
|
imod.addComment1("local write A")
|
|
imod.add(writer.codes.localWriteA)
|
|
imod.addComment1("local write MXSA")
|
|
imod.add(writer.codes.localWriteMXSA)
|
|
imod.addComment1("local write MXSB")
|
|
imod.add(writer.codes.localWriteMXSB)
|
|
imod.addComment1("local write B")
|
|
imod.add(writer.codes.localWriteB)
|
|
|
|
def prepareLWInstToSched(writer, kernel, numLocalWritesPerSched, isNGLL=False):
|
|
#################
|
|
# create a plan #
|
|
#################
|
|
itemsLWToSched = list(writer.codes.localWriteA.items()) + \
|
|
list(writer.codes.localWriteMXSA.items()) + \
|
|
list(writer.codes.localWriteMXSB.items()) + \
|
|
list(writer.codes.localWriteB.items())
|
|
numDummy = 0
|
|
insertDummyTop = False
|
|
if kernel["PrefetchGlobalRead"] >= 2:
|
|
# PrefetchGlobalRead + DirectToLds/DirectToVgpr case, need to add dummy list to insert global read
|
|
lenA = len(list(writer.codes.globalReadA.middle.items()))
|
|
lenMXSA = len(list(writer.codes.globalReadMXSA.middle.items()))
|
|
lenMXSB = len(list(writer.codes.globalReadMXSB.middle.items()))
|
|
lenB = len(list(writer.codes.globalReadB.middle.items()))
|
|
|
|
lenAFooter = len(list(writer.codes.globalReadA.footer.items()))
|
|
lenMXSAFooter = len(list(writer.codes.globalReadMXSA.footer.items()))
|
|
lenMXSBFooter = len(list(writer.codes.globalReadMXSB.footer.items()))
|
|
lenBFooter = len(list(writer.codes.globalReadB.footer.items()))
|
|
|
|
# A/B swap check for DTV. NGLL case, no swap
|
|
swapped = writer.isSwapGlobalReadOrderForDtvOrDtl(kernel) and (not isNGLL)
|
|
insertDummyTop = True
|
|
if swapped:
|
|
# swap A and B (SwapGlobalReadOrder case, the actual content is swapped (B is in globalReadACode). Need adjustment)
|
|
lenA, lenB = lenB, lenA
|
|
lenMXSA, lenMXSB = lenMXSB, lenMXSA
|
|
|
|
if kernel["DirectToLdsA"] or kernel["DirectToVgprA"]:
|
|
if kernel["DirectToLdsA"]:
|
|
# PGR2 + DTLcase, footer code is added in middle. Need to subtract 1 (for footer inst)
|
|
lenA -= lenAFooter
|
|
numDummy += lenA
|
|
insertDummyTop = (not swapped)
|
|
if kernel["ProblemType"]["MXBlockA"] and (kernel["DirectToLdsMXSA"] or kernel["DirectToVgprMXSA"]):
|
|
if kernel["DirectToLdsMXSA"]:
|
|
# PGR2 + DTLcase, footer code is added in middle. Need to subtract 1 (for footer inst)
|
|
lenMXSA -= lenMXSAFooter
|
|
numDummy += lenMXSA
|
|
insertDummyTop = (not swapped)
|
|
if kernel["ProblemType"]["MXBlockB"] and (kernel["DirectToLdsMXSB"] or kernel["DirectToVgprMXSB"]):
|
|
if kernel["DirectToLdsMXSB"]:
|
|
# PGR2 + DTLcase, footer code is added in middle. Need to subtract 1 (for footer inst)
|
|
lenMXSB -= lenMXSBFooter
|
|
numDummy += lenMXSB
|
|
insertDummyTop = swapped
|
|
if kernel["DirectToLdsB"] or kernel["DirectToVgprB"]:
|
|
if kernel["DirectToLdsB"]:
|
|
# PGR2 + DTLcase, footer code is added in middle. Need to subtract 1 (for footer inst)
|
|
lenB -= lenBFooter
|
|
numDummy += lenB
|
|
insertDummyTop = swapped
|
|
if kernel["ProblemType"]["Sparse"]:
|
|
if kernel["DirectToLdsMetadata"]:
|
|
# DirectToLdsMetadata M0 update
|
|
numDummy -= 1
|
|
else:
|
|
# Workaround, Sparse B + (DTLB only) don't need this since there's no meta GR in globalReadB
|
|
sparseTc = 'B' if kernel["ProblemType"]["Sparse"] == 2 else 'A'
|
|
if not (sparseTc == 'B' and kernel["DirectToLdsB"] and not kernel["DirectToLdsA"]) and kernel["DirectToLds%s"%sparseTc]:
|
|
numMetaGR = kernel["NumLoadsPerpendicularMetadata"] * kernel["NumLoadsCoalescedMetadata"]
|
|
numDummy -= numMetaGR
|
|
# extend localWrite by inserting empty Module
|
|
# See getNumLocalWritePerMfma for how this work
|
|
itemsLWToSchedTemp = []
|
|
counter = 0
|
|
itemsLWToSchedLength_1 = len(itemsLWToSched) - 1
|
|
for i in range(itemsLWToSchedLength_1 + numDummy):
|
|
if insertDummyTop:
|
|
if i < numDummy:
|
|
item = None
|
|
else:
|
|
item = itemsLWToSched.pop(0)
|
|
itemsLWToSchedTemp.append([counter, item])
|
|
else:
|
|
if i < itemsLWToSchedLength_1:
|
|
item = itemsLWToSched.pop(0)
|
|
itemsLWToSchedTemp.append([counter, item])
|
|
else:
|
|
item = None
|
|
counter += 1
|
|
skip = kernel["PrefetchGlobalRead"] >= 2 and kernel["ProblemType"]["Sparse"] and kernel["DirectToVgprSparseMetadata"] \
|
|
and item.name.startswith("MetadataWrite") and countVMovB32(item)
|
|
if not skip:
|
|
for _ in range(PRECISION-1):
|
|
counter += 1
|
|
if itemsLWToSched or (numDummy and itemsLWToSchedLength_1 == -1):
|
|
if itemsLWToSched:
|
|
itemsLWToSchedTemp.append([counter, itemsLWToSched.pop(0)])
|
|
counter += 1
|
|
for i in range(numLocalWritesPerSched + numLocalWritesPerSched % PRECISION - counter % numLocalWritesPerSched):
|
|
counter += 1
|
|
itemsLWToSched = itemsLWToSchedTemp
|
|
if not itemsLWToSched:
|
|
itemsLWToSched.append([max(0, counter - 1), None])
|
|
elif itemsLWToSched[-1][0] != (counter - 1):
|
|
itemsLWToSched.append([counter - 1, None]) # end of the list if not equal to counter
|
|
# This counts the number of modules which contain a ds_write
|
|
# Scheduler below keeps all writes in the same module in same iteration
|
|
# so this is better match to what it is trying to do
|
|
# numWritesToSched = sum(1 for item in itemsLWToSched if countLocalWrite(item)
|
|
numWritesToSched = itemsLWToSched[-1][0]
|
|
return itemsLWToSched, numWritesToSched
|
|
|
|
def assignLWSchedIndexSIA3(writer, kernel, numLocalWritesPerSched, localWriteEndIter, numWritesToSched):
|
|
numMfmaPerIter = writer.states.numMfmaPerIter
|
|
writer.states.lwStartMfmaIndex = writer.states.lwEndMfmaIndex - max(1,roundUp(numWritesToSched/numLocalWritesPerSched)) + 1
|
|
if writer.states.lwStartMfmaIndex < writer.states.grEndMfmaIndex:
|
|
# adjust lwStartMfmaIndex for PGR1
|
|
writer.states.lwStartMfmaIndex = writer.states.grEndMfmaIndex
|
|
if kernel["1LDSBuffer"] or kernel["DirectToLds"]:
|
|
writer.states.sync1LdsMfmaIndex = max(writer.states.lwStartMfmaIndex - 1, 0)
|
|
startIter = writer.states.lwStartMfmaIndex//numMfmaPerIter
|
|
assert startIter < localWriteEndIter+1 # startIter should be at or before the endIter
|
|
return startIter
|
|
|
|
def assignLWSchedIndexDefault(writer, kernel, numLocalWritesPerSched, localWriteEndIter, lastLoadIter, numWritesToSched):
|
|
startIter = localWriteEndIter - roundUp(numWritesToSched/numLocalWritesPerSched) + 1
|
|
# - can't move a write past the load it depends on
|
|
# as a simplification, don't move writes past any loads
|
|
if startIter < lastLoadIter:
|
|
startIter = lastLoadIter
|
|
return startIter
|
|
|
|
def getReadsToWait(writer, kernel):
|
|
lwAcnt = len(list(writer.codes.localWriteA.items())) + \
|
|
countDSStoreB192(writer.codes.localWriteA)
|
|
lwBcnt = len(list(writer.codes.localWriteB.items())) + \
|
|
countDSStoreB192(writer.codes.localWriteB)
|
|
lwMXAcnt = len(list(writer.codes.localWriteMXSA.items()))
|
|
lwMXBcnt = len(list(writer.codes.localWriteMXSB.items()))
|
|
readsToWait = lwAcnt + lwBcnt + lwMXAcnt + lwMXBcnt
|
|
readsToWaitNGLL = readsToWait
|
|
return readsToWait, readsToWaitNGLL
|
|
|
|
def schedLocalWrite(writer, kernel, numLocalWriteModPerIter, numLocalWritesPerSched, localWriteEndIter, \
|
|
itemsGRToSchedLater, itemsLWToSched, startIter, readsToWait, readsToWaitNGLL, \
|
|
firstIter, lastLc, isNGLL, startIterItem = None):
|
|
# schedule here
|
|
localwriteCnt = 0
|
|
globalReadInstOffset = 0
|
|
additionalIndexList = {}
|
|
skip = 0
|
|
if itemsLWToSched and itemsLWToSched[0][0] == 0 and itemsLWToSched[0][1] is None:
|
|
itemsLWToSched.pop(0) # remove the dummy item
|
|
|
|
itemsLWToSchedIndexLast = 0
|
|
for u in range(startIter, localWriteEndIter+1):
|
|
itemsLWToSchedLength = itemsLWToSched[-1][0] if itemsLWToSched else 0
|
|
if u == localWriteEndIter:
|
|
itemPerIter = itemsLWToSchedLength # schedule all remaining activity
|
|
else:
|
|
itemPerIter = itemsLWToSchedIndexLast + numLocalWriteModPerIter - 1
|
|
# if localwrite is not multiple of numLocalWriteModPerIter, fill last iteration first.
|
|
# make sure numLocalWriteModPerIter is enough to schedule localwrite
|
|
# TODO: if numLocalWriteModPerIter is not enough to schedule localwrite, need smarter way to distribute localWrite
|
|
if u == startIter and startIterItem:
|
|
itemPerIter = startIterItem - 1
|
|
# Convert timeline index to map index
|
|
if itemsLWToSched:
|
|
foundIndex = -1
|
|
skipInsert = False
|
|
for index, [itemIndex, _] in enumerate(itemsLWToSched):
|
|
if itemPerIter == itemIndex:
|
|
skipInsert = True
|
|
foundIndex = index
|
|
break
|
|
if itemPerIter < itemIndex:
|
|
foundIndex = index
|
|
break
|
|
if not skipInsert: # Insert scheduling point if needed
|
|
if foundIndex != -1:
|
|
itemsLWToSched.insert(foundIndex, [itemPerIter, None])
|
|
itemPerIter = foundIndex + 1
|
|
else:
|
|
itemsLWToSched.append([itemPerIter, None])
|
|
itemPerIter = len(itemsLWToSched)
|
|
itemsLWToSchedLength = itemPerIter
|
|
else:
|
|
itemPerIter = foundIndex + 1
|
|
|
|
perIterLocalWriteCodeCounter = 0
|
|
perIterLocalWriteCodeNGLLCounter = 0
|
|
|
|
itemsLWToSchedLengthLeft = itemsLWToSchedLength - itemsLWToSchedIndexLast
|
|
itemsLWToSchedIndexForSplitDS = 0
|
|
for itemsLWToSchedIndex, item in itemsLWToSched[:itemPerIter]:
|
|
tmp = itemsLWToSchedIndexLast + numLocalWriteModPerIter
|
|
for gapIndex in range(itemsLWToSchedIndexLast, itemsLWToSchedIndex + 1):
|
|
imodList = []
|
|
imodNGLLList = []
|
|
# Use a module to ensure these pieces stay together in the sub-iter scheduler
|
|
if gapIndex == itemsLWToSchedIndex:
|
|
if item:
|
|
writesPerItem = countLocalWrite(item)
|
|
if kernel["ProblemType"]["Sparse"] and not writesPerItem:
|
|
writesPerItem = item.name.startswith("MetadataWrite") and countVMovB32(item)
|
|
if writesPerItem:
|
|
writesPerItem = countLocalWrite(item)
|
|
if kernel["ProblemType"]["Sparse"] and not writesPerItem:
|
|
writesPerItem = item.name.startswith("MetadataWrite") and countVMovB32(item)
|
|
# Split into several dsStore32
|
|
syncEndExpandedNumIndex = itemsLWToSchedIndexForSplitDS
|
|
|
|
if writer.states.numMfmaPerIter and u == (writer.states.lwEndMfmaIndex // writer.states.numMfmaPerIter):
|
|
syncEndExpandedNumIndex = numLocalWriteModPerIter
|
|
syncEndExpandedNumIndex *= ((writer.states.syncPlrMfmaIndex % writer.states.numMfmaPerIter) / writer.states.numMfmaPerIter)
|
|
syncEndExpandedNumIndex = roundUp(syncEndExpandedNumIndex)
|
|
|
|
itemNew, numItemNew, globalReadInstOffset = splitDSInstructionIntoSmaller(writer, kernel, item, numLocalWritesPerSched, syncEndExpandedNumIndex, itemsLWToSchedIndexForSplitDS) if writer.do["AutoSplitDsWrite"] else (None, 0, 0)
|
|
if itemsLWToSchedIndexForSplitDS + globalReadInstOffset <= itemsLWToSchedLengthLeft:
|
|
additionalIndexList.clear()
|
|
for i in range(numItemNew):
|
|
additionalIndexList[int(i * numLocalWritesPerSched + itemsLWToSchedIndexForSplitDS)] = itemNew[i]
|
|
else:
|
|
globalReadInstOffset = 0
|
|
|
|
imodList.append(TextBlock("/* sched write - iter %u writesPerItem=%u */\n"%(u,writesPerItem)))
|
|
imodNGLLList.append(TextBlock("/* sched write - iter %u writesPerItem=%u */\n"%(u,writesPerItem)))
|
|
# if writesPerItem>1 this indicates multiple LocalWrites in the same module
|
|
# this happens in some transpose cases. Here the first write needs to wait
|
|
# for the associated global read to finish, then the remaining writes can flow
|
|
# TODO - can schedule these writes across iters, should figure this out above
|
|
# A ds_store_b128 + ds_store_b64 pair (counted as one B192-equivalent by
|
|
# countDSStoreB192) corresponds to 2 physical VMEM loads. Each extra physical
|
|
# load beyond the first must also decrement readsToWait so that the generated
|
|
# s_wait_loadcnt value correctly tracks the outstanding load count.
|
|
numPhysLoads = 1 + countDSStoreB192(item)
|
|
readsToWait = readsToWait - numPhysLoads
|
|
readsToWaitNGLL = readsToWaitNGLL - numPhysLoads
|
|
imodList.append(SWaitCnt(vlcnt=readsToWait, \
|
|
comment="wait for global read before writing to local"))
|
|
imodNGLLList.append(SWaitCnt(vlcnt=readsToWaitNGLL, \
|
|
comment="wait for global read before writing to local"))
|
|
# PK and StoreCUnroll is removed so you cannot find any HolderContainer in s_waitcnt
|
|
if kernel["PrefetchGlobalRead"]>=2:
|
|
hasHolder, wcList = hasHolderInWaitCnt(item)
|
|
if hasHolder:
|
|
readsToWaitAdjust = readsToWait
|
|
if kernel["NoLdsWriteCode"]:
|
|
# DirectToLds for both A and B case, use the number of global read for both A and B as vmcnt (only for PGR=1)
|
|
readsToWaitAdjust = len(list(writer.codes.globalReadA.middle.items())) + \
|
|
len(list(writer.codes.globalReadMXSA.middle.items())) + \
|
|
len(list(writer.codes.globalReadMXSB.middle.items())) + \
|
|
len(list(writer.codes.globalReadB.middle.items()))
|
|
for wc in wcList:
|
|
replaceHolder(wc, (readsToWaitAdjust))
|
|
if itemsLWToSchedIndexForSplitDS in additionalIndexList:
|
|
imodList.append(additionalIndexList[itemsLWToSchedIndexForSplitDS])
|
|
additionalIndexList.pop(itemsLWToSchedIndexForSplitDS)
|
|
elif gapIndex < itemsLWToSchedIndex or (not item):
|
|
pass
|
|
else:
|
|
imodList.append(item)
|
|
# schedule global instruction that need to be scheduled later
|
|
numGlobalReadA = kernel["NumLoadsPerpendicularA"] * kernel["NumLoadsCoalescedA"]
|
|
numGlobalReadB = kernel["NumLoadsPerpendicularB"] * kernel["NumLoadsCoalescedB"]
|
|
dtvReadNum = numGlobalReadA if kernel["DirectToVgprA"] else numGlobalReadB
|
|
totalNumGR = numGlobalReadA + numGlobalReadB
|
|
nondtvReadNum = totalNumGR - dtvReadNum
|
|
|
|
readCntA = 1
|
|
readCntB = 1
|
|
readCntA = 2 if kernel["DirectToVgprA"] and kernel["reorderGRInstForDTVA"] and \
|
|
kernel["NumLoadsCoalescedA"] % 2 == 0 else 1
|
|
readCntB = 2 if kernel["DirectToVgprB"] and kernel["reorderGRInstForDTVB"] and \
|
|
kernel["NumLoadsCoalescedB"] % 2 == 0 else 1
|
|
|
|
if kernel["DirectToVgprA"]: # In loop, load A first
|
|
readCnt = readCntA if (len(itemsGRToSchedLater) > nondtvReadNum) or isNGLL else readCntB
|
|
elif kernel["DirectToVgprB"]: # In loop, load B first
|
|
readCnt = readCntB if (len(itemsGRToSchedLater) > nondtvReadNum) or isNGLL else readCntA
|
|
else: # not kernel["DirectToVgprA"] and not kernel["DirectToVgprB"]
|
|
readCnt = 1
|
|
|
|
if localwriteCnt % PRECISION == ((numLocalWritesPerSched % PRECISION) + globalReadInstOffset):
|
|
if not skip:
|
|
globalReadInstOffset = 0
|
|
reads = 0
|
|
while itemsGRToSchedLater:
|
|
itemGR = itemsGRToSchedLater[0]
|
|
readsInc = countGlobalRead(itemGR)
|
|
reads = reads + readsInc
|
|
if reads > readCnt * readsInc:
|
|
break
|
|
if kernel["ExpertSchedulingMode"] > 0:
|
|
imodList.append(SWaitAlu(vm_vsrc=0, comment="wait for local read to vgpr complete"))
|
|
# PK and StoreCUnroll is removed so you cannot find any HolderContainer in s_waitcnt
|
|
hasHolder, wcList = hasHolderInWaitCnt(itemGR)
|
|
if hasHolder:
|
|
for wc in wcList:
|
|
replaceHolder(wc, (readsToWait))
|
|
imodList.append(itemGR)
|
|
else:
|
|
imodList.append(itemGR)
|
|
readsToWait = readsToWait + readsInc # GR instruction increments vlcnt
|
|
itemsGRToSchedLater.pop(0)
|
|
|
|
if readCnt == 2:
|
|
skip = skip ^ 1
|
|
else:
|
|
skip = 0
|
|
localwriteCnt += 1
|
|
if gapIndex < itemsLWToSchedIndex or (not item):
|
|
pass
|
|
else:
|
|
imodNGLLList.append(deepcopy(item))
|
|
|
|
perIterLocalWriteCodeCounter += 1
|
|
perIterLocalWriteCodeNGLLCounter += 1
|
|
if imodList:
|
|
imod = Module("LocalWriteMod%u"%u)
|
|
imod.addItems(imodList)
|
|
writer.codes.perIterLocalWrite[u][0].append(perIterLocalWriteCodeCounter)
|
|
writer.codes.perIterLocalWrite[u][1].add(imod)
|
|
if lastLc:
|
|
# local write code for NGLL should be updated at the last lc
|
|
# in init acc opt case, the last inner loop generated is not for the last lc.
|
|
# in that case, local write code for NGLL is not as expected.
|
|
if imodNGLLList:
|
|
imodNGLL = Module("LocalWriteModNGLL%u"%u)
|
|
imodNGLL.addItems(imodNGLLList)
|
|
writer.codes.perIterLocalWriteCodeNGLL[u][0].append(perIterLocalWriteCodeNGLLCounter)
|
|
writer.codes.perIterLocalWriteCodeNGLL[u][1].add(imodNGLL)
|
|
itemsLWToSchedIndexForSplitDS += 1
|
|
itemsLWToSchedIndexLast = itemsLWToSchedIndex + 1
|
|
if writer.codes.perIterLocalWrite[u][0] and writer.codes.perIterLocalWrite[u][0][-1] != perIterLocalWriteCodeCounter:
|
|
writer.codes.perIterLocalWrite[u][0].append(perIterLocalWriteCodeCounter)
|
|
if lastLc and writer.codes.perIterLocalWriteCodeNGLL[u][0] and writer.codes.perIterLocalWriteCodeNGLL[u][0][-1] != perIterLocalWriteCodeNGLLCounter:
|
|
writer.codes.perIterLocalWriteCodeNGLL[u][0].append(perIterLocalWriteCodeNGLLCounter)
|
|
itemsLWToSched = itemsLWToSched[itemPerIter:]
|
|
|
|
# should never run out of items to schedule
|
|
assert not itemsLWToSched # should have scheduled everthing already
|
|
|
|
#For the sparse case, GR and LW are not in paired.
|
|
#Hence, we must add all the remaining GRs into imod at the end.
|
|
if kernel["ProblemType"]["Sparse"]:
|
|
while itemsGRToSchedLater:
|
|
itemGR = itemsGRToSchedLater[0]
|
|
imod.add(itemGR)
|
|
itemsGRToSchedLater.pop(0)
|
|
|
|
def splitDSInstructionIntoSmaller(writer, kernel, item, numLocalWritesPerSched, lenOfItems, currentModIdx):
|
|
if not item:
|
|
return None, 0, 0
|
|
if countDSStoreB128(item) != 1 or countInstruction(item) != 1:
|
|
# only support one b128
|
|
return None, 0, 0
|
|
|
|
itemList = item.flatitems()
|
|
instruction = next(filter(lambda inst: isinstance(inst, DSStoreB128), itemList))
|
|
miLatency = writer.states.miLatency
|
|
div = 1
|
|
dsOffset = 0
|
|
|
|
LocalWriteX = DSStoreB128
|
|
if DSStoreB128.issueLatency() < (miLatency - 1):
|
|
# no need to split
|
|
return None, 0, 0
|
|
elif DSStoreB64.issueLatency() < (miLatency - 1):
|
|
LocalWriteX = DSStoreB64
|
|
dsOffset = 8
|
|
div = 2
|
|
elif DSStoreB32.issueLatency() < (miLatency - 1):
|
|
LocalWriteX = DSStoreB32
|
|
dsOffset = 4
|
|
div = 4
|
|
else:
|
|
# miLatency is not enough
|
|
return None, 0, 0
|
|
|
|
if numLocalWritesPerSched * div >= PRECISION:
|
|
# no enough mfma to split
|
|
return None, 0, 0
|
|
|
|
if (currentModIdx + numLocalWritesPerSched * div) >= lenOfItems:
|
|
# no enough modules to schedule
|
|
return None, 0, 0
|
|
|
|
# LW b32 4-way bank conflict latency ~ 108 cycles
|
|
# round up with quad-cycle
|
|
finalLWCycles = roundUp(108 / 4)
|
|
|
|
# How many mfmas between 2 LWs
|
|
# the MFMA buffer must be larger then the latency
|
|
numMfmaBetweenLW = PRECISION // numLocalWritesPerSched
|
|
mfmaBuffer = numMfmaBetweenLW * miLatency
|
|
if mfmaBuffer <= finalLWCycles:
|
|
# no enough cycles between LWs
|
|
return None, 0, 0
|
|
|
|
extraSched = roundUp(finalLWCycles / miLatency)
|
|
if (currentModIdx + numLocalWritesPerSched * (div - 1 + extraSched)) >= lenOfItems:
|
|
# no enough cycles before barrier
|
|
return None, 0, 0
|
|
|
|
addr = instruction.getParams()[0]
|
|
srcr = instruction.getParams()[1]
|
|
offs = instruction.getParams()[2]
|
|
ds = instruction.ds
|
|
writeInst = []
|
|
for d in range(div):
|
|
ds1 = DSModifiers(na=1, offset=ds.offset + dsOffset * d)
|
|
r1 = deepcopy(srcr)
|
|
r1.regNum //= div
|
|
r1.regName.addOffset(4 // div * d)
|
|
writeInst.append(LocalWriteX(dstAddr=addr, src=r1, ds=ds1, comment=instruction.comment + " splitted"))
|
|
|
|
print2(f"Split ds_write_b128 to 4xds_write_b32 for {str(instruction)}")
|
|
|
|
return writeInst, len(writeInst), numLocalWritesPerSched * (div - 1)
|
|
|
|
|
|
################################################################################
|
|
################################################################################
|
|
###
|
|
### Helper function
|
|
###
|
|
################################################################################
|
|
################################################################################
|
|
|
|
def hasHolderInWaitCnt(module: Item):
|
|
wcList = []
|
|
hasHolder = False
|
|
if isinstance(module, Module):
|
|
for item in module.items():
|
|
tmpHasHolder, tmpList = hasHolderInWaitCnt(item)
|
|
hasHolder = hasHolder or tmpHasHolder
|
|
wcList.extend(tmpList)
|
|
elif isinstance(module, SWaitCnt):
|
|
wcList.append(module)
|
|
if isinstance(module.dscnt, HolderContainer) or \
|
|
isinstance(module.kmcnt, HolderContainer) or \
|
|
isinstance(module.vlcnt, HolderContainer) or \
|
|
isinstance(module.vscnt, HolderContainer):
|
|
hasHolder = True
|
|
return hasHolder, wcList
|