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
5677 lines
252 KiB
Python
5677 lines
252 KiB
Python
################################################################################
|
|
#
|
|
# Copyright (C) 2025 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 dataclasses import dataclass, field, asdict
|
|
from rocisa.code import KernelBody, Label, Macro, Module, RegSet, SrdUpperValue, \
|
|
StructuredModule, TextBlock, ValueEndif, ValueIf, ValueElseIf, ValueSet, SignatureBase
|
|
from rocisa.container import vgpr, sgpr, SMEMModifiers, replaceHolder, EXEC,\
|
|
VOP3PModifiers, ContinuousRegister
|
|
from rocisa.instruction import BufferLoadB128, BufferLoadB32, BufferLoadB64, \
|
|
BufferLoadD16B16, BufferLoadD16U8, DSLoad2B32, DSLoad2B64, DSLoadB128, \
|
|
DSLoadB32, DSLoadB64, DSLoadB64TrB16, DSLoadInstruction, DSLoadU16, \
|
|
DSLoadU8, DSStore2B32, DSStore2B64, DSStoreB128, DSStoreB16, DSStoreB256, \
|
|
DSStoreB32, DSStoreB64, DSStoreB8, DSStoreInstruction, FlatLoadB128, FlatLoadB32, \
|
|
FlatLoadB64, FlatStoreB128, FlatStoreB32, FlatStoreB64, Instruction, \
|
|
MFMAInstruction, SBarrier, SBranch, SCBranchSCC0, SCBranchSCC1, SCmpLeU32, \
|
|
SMovB32, SMFMAInstruction, SNop, SSetPrior, SSetRegIMM32B32, SSubU32, SWaitCnt, SWaitAlu, \
|
|
SLongBranchPositive, VFmaMixF32, VMadMixF32, VMovB32
|
|
from rocisa.instruction import SAddU32, SAddCU32, SCmpEQU32, SCSelectB32, SSubBU32
|
|
from Tensile.Common import IsaVersion
|
|
from Tensile.Common.Utilities import printWarning
|
|
from Tensile.Utilities.Decorators.Shared import CallableGuard
|
|
|
|
from copy import deepcopy
|
|
from typing import Callable, Optional, Union, Tuple
|
|
from enum import Enum, auto
|
|
import Tensile.Components.CMSValidator as cmsv
|
|
from typing import Callable
|
|
from itertools import product
|
|
|
|
# Enum to distinguish between different schedule matching outcomes
|
|
class ScheduleMatchStatus(Enum):
|
|
FOUND = auto() # Schedule found and supported
|
|
NO_MATCH = auto() # Criteria don't match, continue searching
|
|
UNSUPPORTED_VARIANT = auto() # Criteria match but variant unsupported, stop searching
|
|
|
|
# Global registry for schedule functions
|
|
_SCHEDULE_REGISTRY = []
|
|
|
|
_SCHEDULE_METADATA: list["CMSKernelInfo"] = []
|
|
|
|
# Map dtype predicate functions to human-readable names
|
|
_DTYPE_PREDICATE_NAMES: dict[Callable, str] = {}
|
|
|
|
def _register_dtype_name(func: Callable, name: str) -> Callable:
|
|
"""Helper to register a dtype predicate name mapping."""
|
|
_DTYPE_PREDICATE_NAMES[func] = name
|
|
return func
|
|
|
|
@dataclass(frozen=True)
|
|
class CMSKernelInfo:
|
|
"""
|
|
Metadata about registered CMS kernels
|
|
Contains the minimum combination of parameters needed to use the CMS kernel.
|
|
Important Note:
|
|
If you are adding new parameters to this list (of params use in CMS kernels), please make sure those names match Tensile names.
|
|
These names will be used by caller/tuning codes to set correct parameter/values.
|
|
"""
|
|
name: str
|
|
dtype: str
|
|
MacroTile0: int
|
|
MacroTile1: int
|
|
DepthU: int
|
|
PrefetchGlobalRead: int
|
|
PrefetchLocalRead: int
|
|
DirectToLds: bool
|
|
DtlPlusLdsBuf: int
|
|
WaveSeparateGlobalReadA: int
|
|
WaveSeparateGlobalReadB: int
|
|
GlobalReadVectorWidthA: int
|
|
GlobalReadVectorWidthB: int
|
|
LocalReadVectorWidth: int
|
|
MatrixInstruction: list[int]
|
|
MIWaveGroup: list[int]
|
|
LDSTrInst: bool
|
|
TransposeLDS: int
|
|
TransposeA: bool
|
|
TransposeB: bool
|
|
|
|
def matches(self, dtype: Optional[str] = None, layout: Optional[str] = None) -> bool:
|
|
"""Check if this kernel info matches the given dtype and/or layout filter.
|
|
|
|
Args:
|
|
dtype: Data type filter string (e.g. "16bit", "8bit", "TF32"), or None for any.
|
|
layout: Layout filter string (e.g. "TN", "NT", "NN", "TT"), or None for any.
|
|
|
|
Returns:
|
|
True if the kernel matches all provided filters.
|
|
"""
|
|
if dtype is not None and self.dtype.lower() != dtype.lower():
|
|
return False
|
|
if layout is not None:
|
|
layout = layout.upper()
|
|
if self.TransposeA != (layout[0] == "T") or self.TransposeB != (layout[1] == "T"):
|
|
return False
|
|
return True
|
|
|
|
|
|
|
|
@dataclass
|
|
class SyncSchedule:
|
|
schedule: list[tuple[int, Union[SWaitCnt, SBarrier]]] = field(default_factory=list)
|
|
|
|
def add(self, idx: int, dscnt: int = -1, vlcnt: int = -1, vscnt: int = -1, comment: str = "", barrier: bool = False, barrier_idx: Optional[int] = None, barrier_comment: str = ""):
|
|
""" Add a SWaitCnt (and optionally a SBarrier) to the schedule at the given index.
|
|
Args:
|
|
idx: The index at which to add the SWaitCnt.
|
|
dscnt: The dscnt value for the SWaitCnt.
|
|
vlcnt: The vlcnt value for the SWaitCnt.
|
|
vscnt: The vscnt value for the SWaitCnt.
|
|
comment: An optional comment for the SWaitCnt.
|
|
barrier: If True, also add a SBarrier.
|
|
barrier_idx: The index at which to add the SBarrier. If None, uses idx.
|
|
barrier_comment: An optional comment for the SBarrier.
|
|
|
|
Example:
|
|
wait.add(2, dscnt=3) adds SWaitCnt at index 2 with dscnt=3
|
|
wait.add(5, dscnt=0, sbarrier=True) adds SWaitCnt at index 5 with dscnt=0 and a SBarrier at the same index
|
|
wait.add(5, dscnt=0, sbarrier=True, barrier_idx=6) adds SWaitCnt at index 5 with dscnt=0 and a SBarrier at index 6
|
|
"""
|
|
self.schedule.append( (idx, SWaitCnt(dscnt=dscnt, vlcnt=vlcnt, vscnt=vscnt, comment=comment)) )
|
|
if barrier:
|
|
barrier_idx = barrier_idx if barrier_idx is not None else idx
|
|
self.schedule.append( (barrier_idx, SBarrier(comment=barrier_comment)) )
|
|
|
|
def get_indicies(self):
|
|
return [item[0] for item in self.schedule]
|
|
def get_code(self):
|
|
return [item[1] for item in self.schedule]
|
|
|
|
def create_range(min_val: int, num: int, max_val: int = -1, step: int = 1, repeat: int = 2) -> list[int]:
|
|
"""
|
|
Generate a list where each value in range(min_val, min_val+num, step) is repeated 'repeat' times.
|
|
Value is clamped to max_val
|
|
|
|
Args:
|
|
min_val: Starting value (inclusive)
|
|
num: Number of values
|
|
step: Step between values
|
|
max_val: Maximum value (clamp)
|
|
repeat: Number of times to repeat each value
|
|
|
|
Example:
|
|
create_range(100, 5,200, 1, 2) => [100, 100, 101, 101, 102, 102, 103, 103, 104, 104]
|
|
create_range(0, 5, 10, 2, 3) => [0, 0, 0, 2, 2, 2, 4, 4, 4, 6, 6, 6, 8, 8, 8]
|
|
create_range(0, 5, 6, 2, 3) => [0, 0, 0, 2, 2, 2, 4, 4, 4, 6, 6, 6, 6, 6, 6]
|
|
"""
|
|
if max_val == -1:
|
|
max_val = min_val + step*num
|
|
return [min(val, max_val) for val in range(min_val, min_val + step*num, step) for _ in range(repeat)]
|
|
|
|
def inflight(lst, index):
|
|
"""
|
|
Return number of inflight loads in a given list of instructions at a specified index
|
|
"""
|
|
return sum(val < (index) for val in lst)
|
|
|
|
def duplicate_list_items(input_list: list, repeat_count: int, step: int = 0) -> list:
|
|
"""
|
|
Duplicate each item in input_list repeat_count times. Optionally duplicate with a step
|
|
|
|
Example:
|
|
duplicate_list_items([1, 2, 3], 3) => [1,1,1, 2,2,2, 3,3,3]
|
|
duplicate_list_items([1, 2, 3], 3, 1) => [1,2,3, 2,3,4, 3,4,5]
|
|
"""
|
|
return [item + step * j for item in input_list for j in range(repeat_count)]
|
|
|
|
def count_items(input_list: list[int], sv: Optional[int] = None, ev: Optional[int] = None):
|
|
"""
|
|
Count how many items in the list are between start value `sv` (inclusive) and end value `ev` (exclusive)
|
|
|
|
Example:
|
|
count_items([1,2,3,4,5], sv=2, ev=5) => 3 (2,3,4)
|
|
count_items([1,2,3,4,5], sv=3) => 3 (3,4,5)
|
|
count_items([1,2,3,4,5], ev=4) => 3 (1,2,3)
|
|
"""
|
|
count = 0
|
|
sv = sv if sv is not None else input_list[0]
|
|
ev = ev if ev is not None else input_list[-1]
|
|
for item in input_list:
|
|
if sv <= item < ev:
|
|
count += 1
|
|
return count
|
|
|
|
def switch_A_B_schedule(optSchedule):
|
|
# Swap A and B entries in the schedule
|
|
# Only replace A/B if it's the last or second-last character
|
|
swappedSchedule = dict()
|
|
for key, value in optSchedule.items():
|
|
# Check if A or B is in the last or second-last position
|
|
if len(key) >= 1 and key[-1] in ('A', 'B'):
|
|
# Last character is A or B
|
|
new_key = key[:-1] + ('B' if key[-1] == 'A' else 'A')
|
|
elif len(key) >= 2 and key[-2] in ('A', 'B'):
|
|
# Second-last character is A or B
|
|
new_key = key[:-2] + ('B' if key[-2] == 'A' else 'A') + key[-1]
|
|
else:
|
|
# No A or B in last or second-last position, keep unchanged
|
|
new_key = key
|
|
swappedSchedule[new_key] = value
|
|
return swappedSchedule
|
|
|
|
class ScheduleInfo:
|
|
def __init__(
|
|
self,
|
|
numCodePaths: int,
|
|
numMfma: int,
|
|
optSchedule: dict[str, list[list[int]]],
|
|
syncCode: list[Union[SWaitCnt, SBarrier]],
|
|
nglshift: int,
|
|
nllshift: int,
|
|
nllZeroDscnt: bool = False,
|
|
mfmaReorder = [],
|
|
snopCode: list[SNop] = [],
|
|
):
|
|
self.numCodePaths = numCodePaths
|
|
self.numMfma = numMfma
|
|
self.optSchedule = optSchedule
|
|
self.syncCode = syncCode
|
|
self.nglshift = nglshift # vmcnt shift for noglobalload loop
|
|
self.nllshift = nllshift # vmcnt shift for nolocalload loop
|
|
self.nllZeroDscnt = nllZeroDscnt
|
|
self.mfmaReorder = mfmaReorder
|
|
self.snopCode = snopCode
|
|
self._disabledPasses: dict[cmsv.ValidatorPass, str] = {}
|
|
|
|
def disableValidationPass(self, pass_id: cmsv.ValidatorPass, reason: str) -> None:
|
|
"""Disable a specific validator pass for this schedule.
|
|
|
|
Args:
|
|
pass_id: The ValidatorPass enum member to disable.
|
|
reason: Mandatory explanation of why this pass is being disabled.
|
|
"""
|
|
if not isinstance(pass_id, cmsv.ValidatorPass):
|
|
raise TypeError(f"pass_id must be a ValidatorPass enum member, got {type(pass_id).__name__}")
|
|
if not isinstance(reason, str) or not reason.strip():
|
|
raise ValueError("Reason for disabling pass must be a non-empty string")
|
|
self._disabledPasses[pass_id] = reason
|
|
|
|
def disableValidation(self, reason: str) -> None:
|
|
"""Disable all validator passes for this schedule."""
|
|
for pass_id in cmsv.ValidatorPass:
|
|
self.disableValidationPass(pass_id, reason)
|
|
|
|
def reasonForDisablingValidationPass(self, pass_id: cmsv.ValidatorPass) -> Optional[str]:
|
|
"""Return the reason this pass was disabled, or None if it is enabled.
|
|
|
|
Raises TypeError if pass_id is not a ValidatorPass enum member.
|
|
"""
|
|
if not isinstance(pass_id, cmsv.ValidatorPass):
|
|
raise TypeError(f"pass_id must be a ValidatorPass enum member, got {type(pass_id).__name__}")
|
|
return self._disabledPasses.get(pass_id)
|
|
|
|
def pretty_print(self):
|
|
print("{")
|
|
keys = list(self.optSchedule.keys())
|
|
maxKeyLen = max(len(k) for k in keys) if keys else 0
|
|
for i, k in enumerate(keys):
|
|
v = self.optSchedule[k]
|
|
comma = "," if i < len(keys) - 1 else ""
|
|
pad = " " * (maxKeyLen - len(k))
|
|
if len(v) == 1:
|
|
print(f" '{k}':{pad} [{v[0]}]{comma}")
|
|
else:
|
|
# Align continuation rows after the opening bracket
|
|
bracketCol = 8 + maxKeyLen
|
|
indent = " " * (bracketCol + 1)
|
|
print(f" '{k}':{pad} [")
|
|
for j, row in enumerate(v):
|
|
row_comma = "," if j < len(v) - 1 else ""
|
|
print(f"{indent}{row}{row_comma}")
|
|
print(f"{' ' * bracketCol}]{comma}")
|
|
print("}")
|
|
|
|
if snops := self.optSchedule.get('SNOP', []):
|
|
print("---- SNOP code ----")
|
|
for idx, code in zip(snops[0], self.snopCode):
|
|
print(f"{idx:>2}: {str(code).strip()}")
|
|
|
|
if syncs := self.optSchedule.get('SYNC', []):
|
|
print("---- SYNC code ----")
|
|
for idx, code in zip(syncs[0], self.syncCode):
|
|
print(f"{idx:>2}: {str(code).strip()}")
|
|
|
|
def removeComments(module):
|
|
retModule = Module()
|
|
for i in module.flatitems():
|
|
if type(i) != TextBlock and not isinstance(i, (SCBranchSCC1, SNop)):
|
|
retModule.add(i)
|
|
return retModule.flatitems()
|
|
|
|
|
|
def customMainLoopSchedule(writer, kernel, tensorParametersA, tensorParametersB, \
|
|
globalReadIncACode, globalReadIncBCode, \
|
|
LRCodeA, PackCodeA, LRCodeB, PackCodeB,\
|
|
LRSwapA, LRSwapB, \
|
|
globalReadA, globalReadB, \
|
|
LWSwapA, LWSwapB, \
|
|
mfmaCode, loopCounterCode, \
|
|
nta=0, ntb=0, \
|
|
):
|
|
strNta = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTA%s"%nta
|
|
strNtb = "" if kernel["AdaptiveGemmNTAB"] == 0 else "_NTB%s"%ntb
|
|
module = Module()
|
|
|
|
globalReadIncACode = removeComments(globalReadIncACode)
|
|
globalReadIncBCode = removeComments(globalReadIncBCode)
|
|
numLoopIter = kernel["LoopIters"]
|
|
ph = -2 # placeholder index
|
|
|
|
if numLoopIter > 1:
|
|
for uIdx in range(0, numLoopIter):
|
|
LRCodeA[uIdx] = removeComments(LRCodeA[uIdx])
|
|
LRCodeB[uIdx] = removeComments(LRCodeB[uIdx])
|
|
PackCodeA[uIdx] = removeComments(PackCodeA[uIdx])
|
|
PackCodeB[uIdx] = removeComments(PackCodeB[uIdx])
|
|
else: # numIterPLR = 0 case
|
|
numLoopIter = 2
|
|
# Split instruction stream for numiterPLR=0
|
|
# First half is done in subiter 1, second half is done in subiter 0
|
|
def splitForPLR(oldModule):
|
|
newList0 = []
|
|
newList1 = []
|
|
numInst = len(oldModule.flatitems())
|
|
for ii in range(numInst):
|
|
if ii < numInst // 2:
|
|
newList1.append(oldModule.flatitems()[ii])
|
|
else:
|
|
newList0.append(oldModule.flatitems()[ii])
|
|
return [newList0, newList1]
|
|
|
|
LRCodeA = splitForPLR(LRCodeA[0])
|
|
LRCodeB = splitForPLR(LRCodeB[0])
|
|
PackCodeA = splitForPLR(PackCodeA[0])
|
|
PackCodeB = splitForPLR(PackCodeB[0])
|
|
|
|
|
|
LRSwapA = removeComments(LRSwapA)
|
|
LRSwapB = removeComments(LRSwapB)
|
|
globalReadA = removeComments(globalReadA)
|
|
globalReadB = removeComments(globalReadB)
|
|
localWriteA = removeComments(writer.codes.localWriteA)
|
|
localWriteB = removeComments(writer.codes.localWriteB)
|
|
LWSwapA = removeComments(LWSwapA)
|
|
LWSwapB = removeComments(LWSwapB)
|
|
loopCounterCode = removeComments(loopCounterCode)
|
|
mfmaCode = removeComments(mfmaCode)
|
|
|
|
_, opt1 = hasCustomSchedule(kernel)
|
|
numCodePath = opt1.numCodePaths
|
|
assert opt1.numMfma == len(mfmaCode)
|
|
|
|
|
|
for _, indexList in opt1.optSchedule.items():
|
|
assert len(indexList) <= opt1.numCodePaths
|
|
|
|
if len(opt1.mfmaReorder) > 0:
|
|
mfmaCode = [mfmaCode[x] for x in opt1.mfmaReorder]
|
|
|
|
idMap = {
|
|
'GRIncA' : globalReadIncACode,
|
|
'GRIncB' : globalReadIncBCode,
|
|
'GRA' : globalReadA,
|
|
'GRB' : globalReadB,
|
|
'LWA' : localWriteA,
|
|
'LWB' : localWriteB,
|
|
'LRSA' : LRSwapA,
|
|
'LRSB' : LRSwapB,
|
|
'LWSA' : LWSwapA,
|
|
'LWSB' : LWSwapB,
|
|
'LCC' : loopCounterCode,
|
|
}
|
|
|
|
|
|
for uIdx in range(0, numLoopIter):
|
|
idMap["LRA%u"%uIdx] = LRCodeA[uIdx]
|
|
idMap["LRB%u"%uIdx] = LRCodeB[uIdx]
|
|
idMap["PackA%u"%uIdx] = PackCodeA[uIdx]
|
|
idMap["PackB%u"%uIdx] = PackCodeB[uIdx]
|
|
|
|
idMap['SYNC'] = opt1.syncCode
|
|
idMap['SNOP'] = opt1.snopCode
|
|
|
|
status, message = cmsv.isValid(opt1, {'kernel' : kernel, "idMap": idMap})
|
|
# create the case str (TN, NT, TT, or NN)
|
|
if isTN(kernel):
|
|
case_str = "TN"
|
|
elif isNT(kernel):
|
|
case_str = "NT"
|
|
elif isTT(kernel):
|
|
case_str = "TT"
|
|
elif isNN(kernel):
|
|
case_str = "NN"
|
|
else:
|
|
case_str = "Unknown"
|
|
assert status is True, f"CMS validation failed for kernel {kernel['MacroTile0']}x{kernel['MacroTile1']}x{kernel['DepthU']} {case_str}: {message}"
|
|
|
|
InstStreams = {key: [stream, idMap[key]] for key, stream in opt1.optSchedule.items()}
|
|
|
|
macro = Macro("MAINLOOP%s%s"%(strNta, strNtb), ["ID", "useGR=1", "usePLR=1", "useGRInc=1", "useLoop=1"])
|
|
|
|
lastIter = numLoopIter - 1
|
|
|
|
for miIndex in range(-1, len(mfmaCode)):
|
|
if miIndex >= 0:
|
|
macro.addComment0("mfmaIndex:%u"%(miIndex))
|
|
macro.add(mfmaCode[miIndex])
|
|
|
|
def scheduleInst(indexList, instructionList):
|
|
ret = [None]*len(indexList)
|
|
totalNumInst = len(instructionList)
|
|
for i in range(len(indexList)):
|
|
if indexList[i]: # For specific codepath
|
|
cc = 0
|
|
# Add slower, but allow reordering of instructions
|
|
while miIndex in indexList[i]:
|
|
ind = indexList[i].index(miIndex)
|
|
if cc == 0:
|
|
ret[i] = Module()
|
|
ret[i].add(instructionList[ind])
|
|
cc += 1
|
|
indexList[i][ind] = ph
|
|
if ret.count(None) == len(ret):
|
|
return [None]
|
|
else:
|
|
return ret
|
|
|
|
ToSched = {k: scheduleInst(stream[0], stream[1]) for k, stream in InstStreams.items()}
|
|
|
|
def nllvmcntHandling(inst, shift0, shift1):
|
|
if isinstance(inst, SWaitCnt) and (inst.vlcnt != -1 or (inst.dscnt != -1 and opt1.nllZeroDscnt)):
|
|
macro.add(ValueIf("\\useGR == 1 && \\usePLR == 1")) # in main loop
|
|
macro.addComment0("vmcnt used in main loop")
|
|
macro.add(inst)
|
|
macro.add(ValueElseIf("\\useGR == 0 && \\usePLR == 1")) # in NGL
|
|
instModified = deepcopy(inst)
|
|
if inst.vlcnt != -1:
|
|
macro.addComment0("vmcnt used in ngl, applying %u shift"%shift0)
|
|
instModified.vlcnt = max(0, instModified.vlcnt - shift0)
|
|
macro.add(instModified)
|
|
macro.add(ValueElseIf("\\useGR == 0 && \\usePLR == 0")) # in NLL
|
|
instModified = deepcopy(inst)
|
|
if inst.vlcnt != -1:
|
|
macro.addComment0("vmcnt used in nll, applying %u shift"%shift1)
|
|
instModified.vlcnt = max(0, instModified.vlcnt - shift1)
|
|
if (inst.dscnt != -1 and opt1.nllZeroDscnt):
|
|
macro.addComment0("setting dscnt = 0 for NLL")
|
|
instModified.dscnt = 0
|
|
macro.add(instModified)
|
|
macro.add(ValueEndif())
|
|
else:
|
|
macro.add(inst)
|
|
|
|
def get_macro_guard(key):
|
|
"""Determine the macro guard for a given instruction key."""
|
|
if key in ['GRIncA', 'GRIncB']:
|
|
return "\\useGRInc == 1"
|
|
elif key in ['GRA', 'GRB', 'LWSA', 'LWSB']:
|
|
return "\\useGR == 1"
|
|
elif key in ['LRA%u' % lastIter, 'LRB%u' % lastIter, 'LRSA', 'LRSB']:
|
|
return "\\usePLR == 1"
|
|
elif key in ['LCC']:
|
|
return "\\useLoop == 1"
|
|
return ""
|
|
|
|
def emit_instructions(instModule, macroGuard: str):
|
|
"""Emit instructions from a module with optional macro guard."""
|
|
if instModule is not None:
|
|
for inst in instModule.flatitems():
|
|
if isinstance(inst, SWaitCnt):
|
|
nllvmcntHandling(inst, opt1.nglshift, opt1.nllshift)
|
|
else:
|
|
if macroGuard:
|
|
macro.add(ValueIf(macroGuard))
|
|
macro.add(inst)
|
|
if macroGuard:
|
|
macro.add(ValueEndif(comment="EndIf %s" % macroGuard))
|
|
|
|
for k, ts in ToSched.items():
|
|
macroGuard = get_macro_guard(k)
|
|
|
|
if len(ts) == 1:
|
|
emit_instructions(ts[0], macroGuard)
|
|
elif len(ts) == numCodePath:
|
|
# Multi codepath - emit inside ID conditionals
|
|
for codepath in range(numCodePath):
|
|
if codepath == 0:
|
|
macro.add(ValueIf("\\ID == %u" % codepath))
|
|
else:
|
|
macro.add(ValueElseIf("\\ID == %u" % codepath))
|
|
emit_instructions(ts[codepath], macroGuard)
|
|
macro.add(ValueEndif(comment="EndIf \\ID checks"))
|
|
else:
|
|
raise ValueError(f"Invalid number of instructions for {k}: {len(ts)}")
|
|
|
|
module.add(macro)
|
|
return module, numCodePath
|
|
|
|
|
|
def hasCustomSchedule(kernel):
|
|
"""
|
|
Trampoline function that checks if a custom schedule is available.
|
|
Iterates through registered schedule functions and returns the first match.
|
|
"""
|
|
if not kernel["UseCustomMainLoopSchedule"]:
|
|
return False, None
|
|
if not kernel["EnableMatrixInstruction"]:
|
|
return False, None
|
|
if not kernel["ISA"] == IsaVersion(9,5,0):
|
|
return False, None
|
|
if isMixed(kernel):
|
|
return False, None
|
|
|
|
useLDSTr = kernel["LDSTrInst"]
|
|
TLDS = kernel["TransposeLDS"]
|
|
|
|
for schedule_func in _SCHEDULE_REGISTRY:
|
|
status, schedule = schedule_func(kernel, useLDSTr, TLDS)
|
|
if status == ScheduleMatchStatus.FOUND:
|
|
return True, schedule
|
|
elif status == ScheduleMatchStatus.UNSUPPORTED_VARIANT:
|
|
# Criteria matched but variant unsupported - stop searching
|
|
return False, None
|
|
# status == NO_MATCH: continue to next schedule
|
|
|
|
return False, None
|
|
|
|
|
|
def query_cms_kernels(dtype: Optional[str] = None, layout: Optional[str] = None) -> list[dict]:
|
|
"""Query for available CMS kernels matching the given data type and/or layout.
|
|
|
|
This function searches the CMS kernel registry and returns the minimum
|
|
combination of parameters needed for each matching CMS kernel.
|
|
|
|
Args:
|
|
dtype: Data type filter (case-insensitive).
|
|
Accepted values: "16bit", "8bit", "TF32", or None for all.
|
|
layout: Layout / transpose e.g. ("TN", "NT", "NN", "TT", or None for all)
|
|
|
|
|
|
Returns:
|
|
A list of dicts, each containing the minimum parameter combination
|
|
needed for a matching CMS kernel. Each dict includes the minimal parameters/values combinations needed for using each CMS kernel.
|
|
|
|
"""
|
|
results = []
|
|
for info in _SCHEDULE_METADATA:
|
|
if info.matches(dtype=dtype, layout=layout):
|
|
results.append(asdict(info))
|
|
return results
|
|
|
|
|
|
def get_cms_kernel_info_objects(dtype: Optional[str] = None, layout: Optional[str] = None) -> list[CMSKernelInfo]:
|
|
"""Query for available CMS kernels and return CMSKernelInfo objects.
|
|
|
|
Same filtering as :func:`query_cms_kernels` but returns the raw
|
|
``CMSKernelInfo`` dataclass instances instead of dicts.
|
|
|
|
Args:
|
|
dtype: Data type filter (case-insensitive), or None for all.
|
|
layout: Layout filter (case-insensitive), or None for all.
|
|
|
|
Returns:
|
|
A list of CMSKernelInfo objects matching the filters.
|
|
"""
|
|
return [info for info in _SCHEDULE_METADATA if info.matches(dtype=dtype, layout=layout)]
|
|
|
|
|
|
def get_available_dtypes() -> set[str]:
|
|
"""Return set of all data type strings that have at least one CMS kernel."""
|
|
return {info.dtype for info in _SCHEDULE_METADATA}
|
|
|
|
|
|
def get_available_layouts(dtype: Optional[str] = None) -> set[str]:
|
|
"""Return a set of all layout strings available for the given data type.
|
|
|
|
Args:
|
|
dtype: Optional data type filter, or None for all data types.
|
|
|
|
Returns:
|
|
Sorted list of unique layout strings (e.g. ["NN", "NT", "TN", "TT"]).
|
|
"""
|
|
def as_str(transpose: bool) -> str:
|
|
return "T" if transpose else "N"
|
|
|
|
layouts: set[str] = set()
|
|
for info in _SCHEDULE_METADATA:
|
|
if dtype is None or info.dtype.lower() == dtype.lower():
|
|
layouts.add(as_str(info.TransposeA) + as_str(info.TransposeB))
|
|
return layouts
|
|
|
|
@CallableGuard
|
|
def isNN(kernel):
|
|
return not kernel["ProblemType"]["TransposeA"] and not kernel["ProblemType"]["TransposeB"]
|
|
|
|
@CallableGuard
|
|
def isNT(kernel):
|
|
return not kernel["ProblemType"]["TransposeA"] and kernel["ProblemType"]["TransposeB"]
|
|
|
|
@CallableGuard
|
|
def isTT(kernel):
|
|
return kernel["ProblemType"]["TransposeA"] and kernel["ProblemType"]["TransposeB"]
|
|
|
|
@CallableGuard
|
|
def isTN(kernel):
|
|
return kernel["ProblemType"]["TransposeA"] and not kernel["ProblemType"]["TransposeB"]
|
|
|
|
@CallableGuard
|
|
def is16bit(kernel):
|
|
return kernel["ProblemType"]["DataType"].isHalf() or kernel["ProblemType"]["DataType"].isBFloat16()
|
|
_register_dtype_name(is16bit, "16bit")
|
|
|
|
@CallableGuard
|
|
def is8bit(kernel):
|
|
return kernel["ProblemType"]["DataType"].isInt8() or kernel["ProblemType"]["DataType"].is8bitFloat()
|
|
_register_dtype_name(is8bit, "8bit")
|
|
|
|
@CallableGuard
|
|
def isMixed(kernel):
|
|
return kernel["ProblemType"]["DataTypeA"].numBytes() != kernel["ProblemType"]["DataTypeB"].numBytes()
|
|
|
|
@CallableGuard
|
|
def isTF32(kernel):
|
|
return kernel["UseF32XEmulation"]
|
|
_register_dtype_name(isTF32, "TF32")
|
|
|
|
@dataclass(frozen=True)
|
|
class TileConfig:
|
|
macro_tile_size_0: int
|
|
macro_tile_size_1: int
|
|
depth_u: int
|
|
prefetch_global_read: int
|
|
prefetch_local_read: int
|
|
direct_to_lds: int
|
|
dtl_plus_lds_buf: bool
|
|
wave_separate_global_read_a: int
|
|
wave_separate_global_read_b: int
|
|
|
|
class _ProbeDataType:
|
|
"""Minimal DataType stub for layout probing at registration time."""
|
|
def isHalf(self): return False
|
|
def isBFloat16(self): return False
|
|
def isInt8(self): return False
|
|
def is8bitFloat(self): return False
|
|
def numBytes(self): return 2
|
|
|
|
class RegisterSchedule:
|
|
"""
|
|
Decorator that registers a schedule function with its matching criteria.
|
|
The function is wrapped with logic that checks if the kernel matches the criteria.
|
|
Supported layouts are auto-detected by probing the inner function at registration time.
|
|
|
|
Usage:
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 96, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x96x64_16bit(kernel, useLDSTr, TLDS):
|
|
...
|
|
"""
|
|
|
|
def __init__(self, tile_config: TileConfig, dtype_predicate: Callable, vector_widths: list[int], matrix_inst: list[int], mfma_wave_group: list[int]):
|
|
"""
|
|
Initialize the registration decorator with matching criteria.
|
|
|
|
Args:
|
|
tile_config: TileConfig object
|
|
dtype_predicate: Callable that takes kernel and returns True if dtype matches
|
|
vector_widths: List of [GRVWA, GRVWB, LRVW]
|
|
matrix_inst: List [M, N, K, B] for MI
|
|
mfma_wave_group: List [rows, cols] for MIWG
|
|
"""
|
|
self.tile_config = tile_config
|
|
self.dtype_predicate = dtype_predicate
|
|
self.vector_widths = vector_widths
|
|
self.matrix_inst = matrix_inst
|
|
self.mfma_wave_group = mfma_wave_group
|
|
|
|
def _make_probe_kernel(self, transA: bool, transB: bool, useLDSTr: bool, TLDS: int, vectorWidthA: int, vectorWidthB: int) -> dict:
|
|
"""Build a synthetic kernel dict for probing layout support."""
|
|
tc = self.tile_config
|
|
mi = self.matrix_inst
|
|
miwg = self.mfma_wave_group
|
|
probe_dtype = _ProbeDataType()
|
|
return {
|
|
"ProblemType": {
|
|
"DataType": probe_dtype,
|
|
"DataTypeA": probe_dtype,
|
|
"DataTypeB": probe_dtype,
|
|
"TransposeA": transA,
|
|
"TransposeB": transB,
|
|
},
|
|
"MacroTile0": tc.macro_tile_size_0,
|
|
"MacroTile1": tc.macro_tile_size_1,
|
|
"DepthU": tc.depth_u,
|
|
"PrefetchGlobalRead": tc.prefetch_global_read,
|
|
"PrefetchLocalRead": tc.prefetch_local_read,
|
|
"DirectToLds": tc.direct_to_lds,
|
|
"WaveSeparateGlobalReadA": tc.wave_separate_global_read_a,
|
|
"WaveSeparateGlobalReadB": tc.wave_separate_global_read_b,
|
|
"GlobalReadVectorWidthA": self.vector_widths[0],
|
|
"GlobalReadVectorWidthB": self.vector_widths[1],
|
|
"VectorWidthA": vectorWidthA,
|
|
"VectorWidthB": vectorWidthB,
|
|
"LocalReadVectorWidth": self.vector_widths[2],
|
|
"MatrixInstruction": list(self.matrix_inst),
|
|
"MIWaveGroup": list(self.mfma_wave_group),
|
|
"LDSTrInst": useLDSTr,
|
|
"TransposeLDS": TLDS,
|
|
"MIWaveTileA": tc.macro_tile_size_0 // (mi[0] * miwg[0]),
|
|
"MIWaveTileB": tc.macro_tile_size_1 // (mi[1] * miwg[1]),
|
|
# Standard flags that inner functions may read/write
|
|
"UseCustomMainLoopSchedule": True,
|
|
"EnableMatrixInstruction": True,
|
|
"UnrollLoopSwapGlobalReadOrder": False,
|
|
"ISA": IsaVersion(9, 5, 0),
|
|
"WavefrontSize": 64,
|
|
"Use64bShadowLimit": 1,
|
|
"ForceUnrollSubIter": False,
|
|
"SwapGlobalReadOrder": 0,
|
|
"UsePLRPack": 0,
|
|
"UseF32XEmulation": False,
|
|
"UseDirect32XEmulation": False,
|
|
"MfmaInitCVgprs": False,
|
|
}
|
|
|
|
def _detect_supported_layouts(self, func: Callable) -> list[Tuple[bool, bool, bool, int]]:
|
|
"""Probe the inner function to discover which layouts it actually handles."""
|
|
def as_str(transpose: bool) -> str:
|
|
return "T" if transpose else "N"
|
|
|
|
valid_vector_widths = [1, 2, 3, 4, 6, 8]
|
|
detected = set()
|
|
for transA, transB in product([True, False], repeat=2):
|
|
for useLDSTr, TLDS in product([True, False], [1, 0]):
|
|
for vwA, vwB in product(valid_vector_widths, repeat=2):
|
|
probe = self._make_probe_kernel(transA, transB, useLDSTr, TLDS, vwA, vwB)
|
|
try:
|
|
found, _ = func(probe, useLDSTr, TLDS)
|
|
if found:
|
|
detected_info_tuple = (transA, transB, useLDSTr, TLDS)
|
|
detected.add(detected_info_tuple)
|
|
except (ValueError, KeyError) as e:
|
|
layout = as_str(transA) + as_str(transB)
|
|
printWarning(
|
|
f"Layout probe failed for func '{func.__name__}' "
|
|
f"with layout={layout}, useLDSTr={useLDSTr}, TLDS={TLDS}, "
|
|
f"VectorWidthA={vwA}, VectorWidthB={vwB}\n"
|
|
f" Kernel: {probe['MacroTile0']}x{probe['MacroTile1']}x{probe['DepthU']} {layout}\n"
|
|
f" Error: {e}"
|
|
)
|
|
continue
|
|
|
|
return list(detected)
|
|
|
|
def __call__(self, func: Callable) -> Callable:
|
|
"""Wrap the function with matching logic and register it."""
|
|
def wrapped_func(kernel: dict, useLDSTr: bool, TLDS: int) -> tuple[ScheduleMatchStatus, Optional[ScheduleInfo]]:
|
|
# TODO: Currently ULSGRO not checked for in CMS, disabled for now
|
|
if kernel["UnrollLoopSwapGlobalReadOrder"]:
|
|
return ScheduleMatchStatus.NO_MATCH, None
|
|
|
|
if not self.dtype_predicate(kernel):
|
|
return ScheduleMatchStatus.NO_MATCH, None
|
|
|
|
MT0, MT1, DU = kernel["MacroTile0"], kernel["MacroTile1"], kernel["DepthU"]
|
|
PGR, PLR, DTL, DPLB = kernel["PrefetchGlobalRead"], kernel["PrefetchLocalRead"], kernel["DirectToLds"], kernel["DtlPlusLdsBuf"]
|
|
WSGRA, WSGRB = kernel["WaveSeparateGlobalReadA"], kernel["WaveSeparateGlobalReadB"]
|
|
kernel_tile_config = TileConfig(MT0, MT1, DU, PGR, PLR, DTL, DPLB, WSGRA, WSGRB)
|
|
if self.tile_config != kernel_tile_config:
|
|
return ScheduleMatchStatus.NO_MATCH, None
|
|
|
|
GRVWA, GRVWB = kernel["GlobalReadVectorWidthA"], kernel["GlobalReadVectorWidthB"]
|
|
LRVWA, LRVWB = kernel["LocalReadVectorWidthA"], kernel["LocalReadVectorWidthB"]
|
|
kernel_vector_widths = [GRVWA, GRVWB, LRVWA, LRVWB]
|
|
# WA: if need to support different LRVW for A and B, add a new parameter to vector_widths
|
|
extended_vector_widths = self.vector_widths + [self.vector_widths[2]]
|
|
if extended_vector_widths != kernel_vector_widths:
|
|
return ScheduleMatchStatus.NO_MATCH, None
|
|
|
|
if self.matrix_inst != kernel["MatrixInstruction"]:
|
|
return ScheduleMatchStatus.NO_MATCH, None
|
|
|
|
if self.mfma_wave_group != kernel["MIWaveGroup"]:
|
|
return ScheduleMatchStatus.NO_MATCH, None
|
|
|
|
# All wrapper criteria matched - call inner function
|
|
match, schedule = func(kernel, useLDSTr, TLDS)
|
|
|
|
if match:
|
|
return ScheduleMatchStatus.FOUND, schedule
|
|
# Inner function returned False - variant unsupported, stop searching
|
|
return ScheduleMatchStatus.UNSUPPORTED_VARIANT, None
|
|
|
|
_SCHEDULE_REGISTRY.append(wrapped_func)
|
|
|
|
# Auto-detect supported layouts by probing the inner function
|
|
detected_infos = self._detect_supported_layouts(func)
|
|
|
|
# Store metadata for query API
|
|
dtype_name = _DTYPE_PREDICATE_NAMES.get(self.dtype_predicate, str(self.dtype_predicate))
|
|
tc = self.tile_config
|
|
for detected_info in detected_infos:
|
|
_transA, _transB, _useLDSTr, _TLDS = detected_info
|
|
_SCHEDULE_METADATA.append(CMSKernelInfo(
|
|
name=func.__name__,
|
|
dtype=dtype_name,
|
|
TransposeA=_transA,
|
|
TransposeB=_transB,
|
|
MacroTile0=tc.macro_tile_size_0,
|
|
MacroTile1=tc.macro_tile_size_1,
|
|
DepthU=tc.depth_u,
|
|
PrefetchGlobalRead=tc.prefetch_global_read,
|
|
PrefetchLocalRead=tc.prefetch_local_read,
|
|
DirectToLds=tc.direct_to_lds,
|
|
DtlPlusLdsBuf=tc.dtl_plus_lds_buf,
|
|
WaveSeparateGlobalReadA=tc.wave_separate_global_read_a,
|
|
WaveSeparateGlobalReadB=tc.wave_separate_global_read_b,
|
|
GlobalReadVectorWidthA=self.vector_widths[0],
|
|
GlobalReadVectorWidthB=self.vector_widths[1],
|
|
LocalReadVectorWidth=self.vector_widths[2],
|
|
MatrixInstruction=list(self.matrix_inst),
|
|
MIWaveGroup=list(self.mfma_wave_group),
|
|
LDSTrInst=_useLDSTr,
|
|
TransposeLDS=_TLDS,
|
|
))
|
|
|
|
# Return original function unchanged (so it can still be called directly)
|
|
return func
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 96, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x96x64_16bit(kernel, useLDSTr, TLDS):
|
|
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
|
|
if isTN(kernel) and TLDS == 1:
|
|
|
|
nglshift = nllshift = 11
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Finish all LRA1 and 1/3 LRB1"),
|
|
7, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="Finish 2/3 LRB1"),
|
|
|
|
15, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="All LRB1 and LRA0 done"),
|
|
15, SBarrier(comment=""),
|
|
|
|
23, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="1/3 LRB0 done"),
|
|
|
|
29, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="All LRB0 done"),
|
|
29, SBarrier(comment=""),
|
|
|
|
35, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="All GRA launched, 3 prev GRB."),
|
|
35, SBarrier(comment=""),
|
|
|
|
42, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Only global reads for this iter"),
|
|
42, SBarrier(comment="")]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA' : [[1,1,2,2,3,3,3,4,4]],
|
|
'GRIncB' : [[5,5,6,6,6,7,7,8,8]],
|
|
'LRA0' : [[1,2,3,4,5,6, 8,10],
|
|
[1,2,3,4,5,6, 9,11]],
|
|
'LRB0' : [[12,16,18],
|
|
[13,17,19]],
|
|
'GRB' : [[36,36,38,38,40,40],
|
|
[37,37,39,39,41,41]],
|
|
'GRA' : [[16,16,18,18,20,20,22,22,24,24,26,26,28,28,30,30],
|
|
[17,17,19,19,21,21,23,23,25,25,27,27,29,29,31,31]],
|
|
'LRA1' : [[36,37,38,39,40,41,42,43]],
|
|
'LRB1' : [[44,45,46]],
|
|
'LRSA' : [[30]], # this must come before next reads of A X0 - so the LRA1
|
|
'LRSB' : [[31]], # this must come before next reads of A X0 - so the LRB1
|
|
'LWSA' : [[32]], # swap after last gr a
|
|
'LWSB' : [[42]], # swap after last gr b
|
|
'LCC' : [[47, 47]],
|
|
}
|
|
elif isNN(kernel) and useLDSTr and TLDS == 1:
|
|
|
|
nglshift = nllshift = 11
|
|
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for LRB1 in prev iteration"),
|
|
|
|
7, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for prior 5 LRA0"),
|
|
20, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="All LRA0 is launched"),
|
|
20, SBarrier(comment=""),
|
|
|
|
21, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="All LRB0 launched"),
|
|
21, SBarrier(comment=""),
|
|
|
|
36, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="All GRA launched"),
|
|
36, SBarrier(comment=""),
|
|
|
|
43, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="All GRB launched"),
|
|
43, SBarrier(comment=""),
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
|
|
'GRIncA' : [[1,1,1, 2,2,2, 3,3,3]],
|
|
'GRIncB' : [[4,4,4, 5,5,5, 6,6,6]],
|
|
|
|
'LRA0' : [[1, 3,3, 5,5, 7,7, 9,9, 11,11, 13,13, 15,15, 17],
|
|
[2, 4,4, 6,6, 8,8, 10,10, 12,12, 14,14, 16,16, 18]],
|
|
'LRB0' : [[13, 15, 17],
|
|
[14, 16, 18]],
|
|
|
|
'GRA' : [[21,21, 23,23, 25,25, 27,27, 29,29, 31,31, 33,33, 35,35],
|
|
[20,20, 22,22, 24,24, 26,26, 28,28, 30,30, 32,32, 34,34]],
|
|
'GRB' : [[37,37, 39,39, 41,41],
|
|
[38,38, 40,40, 42,42]],
|
|
|
|
'LRSA' : [[30]],
|
|
'LRSB' : [[31]],
|
|
|
|
'LWSA' : [[36]],
|
|
'LWSB' : [[43]],
|
|
|
|
'LRA1' : [[36,36, 37,37, 38,38, 39,39, 40,40, 41,41, 42,42, 43,43]],
|
|
'LRB1' : [[43, 44, 45]],
|
|
|
|
'LCC' : [[47, 47]],
|
|
}
|
|
else:
|
|
return False, None
|
|
|
|
numMfma = 48
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 96, 64, 2, 1, 1, True, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x96x64_16bit_DPLB(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
|
|
if isNT(kernel) and useLDSTr and TLDS == 0:
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt= 4, vlcnt=-1, vscnt=-1, comment="wait for all LRA1 and one LRB1"),
|
|
7, SWaitCnt(dscnt= 6, vlcnt=-1, vscnt=-1, comment="wait the rest of LRB1"),
|
|
23, SWaitCnt(dscnt= 0, vlcnt=-1, vscnt=-1, comment="wait for all LR0 before starting 2nd sub-iteration"),
|
|
27, SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="wait for GRAs before starting LRA1"),
|
|
27, SBarrier(comment=""),
|
|
41, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="wait for GRBs before starting LRB1"),
|
|
41, SBarrier(comment=""),
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncTable[::2]],
|
|
|
|
'GRIncA': [[0,1,1, 1,2,3, 3,3,4],
|
|
[0,0,0, 1,2,2, 2,3,4]],
|
|
'GRIncB': [[18,18,18, 21,21,21, 22,22,22]],
|
|
|
|
'LRA0': [[0,0, 2,2, 4,4, 6,6, 8,8, 10,10, 12,12, 14,14],
|
|
[1,1, 3,3, 5,5, 7,7, 9,9, 11,11, 13,13, 15,15]],
|
|
'LRB0': [[9,11, 13,15, 16,16],
|
|
[10,12, 14,16, 17,17]],
|
|
|
|
'GRA': [[5,5, 5,6, 7, 9, 11,11, 15,16, 19,20, 24,25, 26,27],
|
|
[4,5, 6,6, 7,10, 11,12, 15,16, 19,20, 24,25, 26,28]],
|
|
'GRB': [[31,31, 35,35, 39,39],
|
|
[32,32, 36,36, 40,40]],
|
|
|
|
'LRA1': [[28,28, 30,30, 32,32, 34,34, 36,36, 37,37, 38,38, 40,40],
|
|
[29,29, 31,31, 33,33, 35,35, 37,37, 39,39, 41,41, 42,42]],
|
|
'LRB1': [[42,42, 43,43, 45,45],
|
|
[43,43, 44,44, 46,46]],
|
|
|
|
'LRSA': [[26,26,26,27]],
|
|
'LRSB': [[27]],
|
|
'LWSA': [[44,44,44],
|
|
[44,45,45]],
|
|
'LWSB': [[]],
|
|
'LCC': [[46,46],
|
|
[45,46]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 11
|
|
else:
|
|
return False, None
|
|
|
|
numMfma = 48
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(192, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_192x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isNN(kernel) and useLDSTr and TLDS==1:
|
|
# TODO: This schedule can be improved when BC are resolved for MT192
|
|
# Note: A/B Global read orders are swapped
|
|
# i.e. GRA contains GR for B
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
optSchedule = {
|
|
'SYNC' : [[12,13, 47,48,49,50,51, 52,53, 56,56, 95]],
|
|
'GRIncB' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncA' : [[42,42,43,43,44,44,45,45,46]],
|
|
'LRB0' : [[0,0,1,1,2,2,6,8],
|
|
[3,3,4,4,5,5,7,9]],
|
|
# These local reads have BC
|
|
'LRA0' : [[10, 15,17,19,21,23, 25,27,29,33,37,39],
|
|
[11, 14,16,18,20,22, 24,26,28,32,36,38]],
|
|
'GRA' : [[14,14, 16,16, 18,18, 20,20, 22,22, 34,34,36,36,38,38],
|
|
[15,15, 17,17, 19,19, 21,21, 23,23, 35,35,37,37,39,39]],
|
|
'GRB' : [[54,54, 56,56, 58,58, 60,60, 62,62, 64,64],
|
|
[55,55, 57,57, 59,59, 61,61, 63,63, 65,65]],
|
|
'LRSA' : [[40]],
|
|
'LRSB' : [[40]],
|
|
'LWSB' : [[41]], # For B
|
|
'LWSA' : [[66]], # For A
|
|
'LRB1' : [[57,57,59,59,61,61,63,65],
|
|
[58,58,60,60,62,62,64,64]],
|
|
'LRA1' : [[67,71,73,75,77,79,81,85,87,89,91,93],
|
|
[68,72,74,76,78,80,82,86,88,90,92,94]],
|
|
'LCC' : [[95, 95]],
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=10, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=8, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=9, vscnt=-1, comment="Wait for GRA & GRB to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 & LRB1 to complete"),]
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
elif isTN(kernel) and not useLDSTr and TLDS == 1:
|
|
#index and code pair
|
|
syncTable = [-1, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="for LRB1-0"),
|
|
5, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="for LRB1"),
|
|
14, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="for LRA0 complete"),
|
|
14, SBarrier(comment="for GRA start"),
|
|
46, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="for LRB0"),
|
|
46, SBarrier(comment="for GRB start"),
|
|
50, SWaitCnt(dscnt=-1, vlcnt=14+1, vscnt=-1, comment="for LRA1"),
|
|
50, SBarrier(comment="for LRA1 start"),
|
|
65, SWaitCnt(dscnt=-1, vlcnt=6+5, vscnt=-1, comment="for LRB0"),
|
|
65, SBarrier(comment="for LRB1 start"),]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[6,6,7,7,8,8,9,9,9]],
|
|
'GRIncB': [[33,34,35,36,37,38,39,40,41]],
|
|
|
|
'LRA0' : [[0, 1, 2, 3, 4, 5],
|
|
[-1, 0, 1, 2, 3, 4]],
|
|
'LRB0' : [[7, 9, 11, 13, 15, 17, 19, 21],
|
|
[8, 10, 12, 13, 16, 18, 20, 22]],
|
|
'GRA' : [[14,14, 16,16, 18,18, 20,20, 25,25, 31,31],
|
|
[15,15, 17,17, 19,19, 21,21, 26,26, 32,32]],
|
|
|
|
'GRB' : [[46,46, 50,50, 54,54, 58,58, 62,62, 66,66, 70,70, 76,76],
|
|
[47,47, 51,51, 55,55, 59,59, 63,63, 67,67, 71,71, 77,77]],
|
|
'LRA1' : [[50, 52, 56, 58, 60, 62],
|
|
[51, 53, 57, 59, 61, 63]],
|
|
'LRB1' : [[65, 67, 69, 71, 73, 75, 77, 79],
|
|
[66, 68, 70, 72, 74, 76, 78, 80]],
|
|
|
|
'LRSA' : [[47]],
|
|
'LRSB' : [[47]],
|
|
'LWSA' : [[47]],
|
|
'LWSB' : [[80]],
|
|
'LCC' : [[95, 95]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
elif isNT(kernel) and not useLDSTr and TLDS == 0:
|
|
optSchedule = {
|
|
'SYNC' : [[-1, 25, 25, 46, 46, 55, 55, 72, 72]],
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
'LRA0' : [[0, 0, 2, 2, 4, 4, 6, 6, 8, 8, 10, 10, 12, 12, 14, 14, 16, 16, 18, 18, 20, 20, 22, 22],
|
|
[1, 1, 3, 3, 5, 5, 7, 7, 9, 9, 11, 11, 13, 13, 15, 15, 17, 17, 19, 19, 21, 21, 23, 23]],
|
|
'LRB0' : [[24, 24, 26, 26, 28, 28, 30 ,30],
|
|
[25, 25, 27, 27, 29, 29, 31, 31]],
|
|
'GRA' : [[25, 25, 27, 27, 29, 29, 31, 31, 33, 33, 35, 35],
|
|
[26, 26, 28, 28, 30, 30, 32, 32, 34, 34, 36, 36]],
|
|
'GRB' : [[47, 47, 49, 49, 51, 51, 53, 53, 64, 64, 66, 66, 68, 68, 70, 70],
|
|
[48, 48, 50, 50, 52, 52, 54, 54, 65, 65, 67, 67, 69, 69, 71, 71]],
|
|
'LRA1' : [[55, 55, 57, 57, 59, 59, 61, 61, 63, 63, 65, 65, 67, 67, 69, 69, 87, 87, 89, 89, 91, 91, 93, 93],
|
|
[56, 56, 58, 58, 60, 60, 62, 62, 64, 64, 66, 66, 68, 68, 70, 70, 88, 88, 90, 90, 92, 92, 94, 94]],
|
|
'LRB1' : [[72, 74, 76, 78, 80, 82, 84, 86],
|
|
[73, 75, 77, 79, 81, 83, 85, 87]],
|
|
'LRSA' : [[37]],
|
|
'LRSB' : [[37]],
|
|
'LWSA' : [[71]],
|
|
'LWSB' : [[71]],
|
|
'PackB1': [[-1, -1, -1, -1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5]],
|
|
'PackA1': [[-1, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2]],
|
|
'PackB0': [[47, 47, 47, 47, 50, 50, 50, 50, 50, 50, 51, 51, 51, 51, 51, 51, 51, 51, 52, 52, 52, 52, 52, 52, 52, 52, 53, 53, 53, 53, 53, 53]],
|
|
'PackA0': [[47, 47, 47, 47, 47, 47, 48, 48, 48, 48, 48, 48, 48, 48, 49, 49, 49, 49, 49, 49, 49, 49, 50, 50]],
|
|
'LCC' : [[95, 95]],
|
|
}
|
|
|
|
syncCode = [SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 to complete") ,
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=-1, vlcnt=14+4, vscnt=-1, comment="Wait for global reads to complete") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=-1, vlcnt=14, vscnt=-1, comment="Wait for global reads to complete") ,
|
|
SBarrier(comment="") ,
|
|
]
|
|
nglshift = nllshift = 14
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 96
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 192, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x192x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 96
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and not useLDSTr and TLDS == 1:
|
|
#index and code pair
|
|
syncTable = [-1, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="wait for LRB1-0"),
|
|
7, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="wait for LRB1"),
|
|
10, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="wait for LRA0"),
|
|
10, SBarrier(comment="for GRA"),
|
|
|
|
47, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0-0"),
|
|
50, SWaitCnt(dscnt=-1, vlcnt=14, vscnt=-1, comment="for previous GRA"),
|
|
50, SBarrier(comment="for GRA"),
|
|
|
|
70, SWaitCnt(dscnt=-1, vlcnt=12, vscnt=-1, comment="for previous GRB"),
|
|
70, SBarrier(comment="for GRB"),
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB': [[37,37,38,38,39,39,40,40,41]],
|
|
'LRA0': [[-1, 0, 1, 2, 3, 4, 5, 6],
|
|
[0, 1, 2, 3, 4, 5, 6, 7]],
|
|
#interleave LRB0 , GRA
|
|
'LRB0': [[7, 9, 11, 13, 15, 17],
|
|
[8, 10, 12, 14, 16, 18]],
|
|
'GRA': [[10,10, 12,12, 14,14, 16,16, 20,20, 31,31, 33,33, 35,35],
|
|
[11,11, 13,13, 15,15, 17,17, 21,21, 32,32, 34,34, 36,36]],
|
|
#interleave GRB, LRB1
|
|
'GRB': [[51,51, 55,55, 59,59, 63,63, 83,83, 85,85],
|
|
[52,52, 56,56, 60,60, 64,64, 84,84, 86,86]],
|
|
'LRA1': [[50, 52, 57, 60, 62, 64, 66, 68],
|
|
[51, 53, 58, 61, 63, 65, 67, 69]],
|
|
|
|
'LRB1': [[70, 72, 74, 76, 78, 79],
|
|
[71, 73, 75, 77, 79, 80]],
|
|
'LRSA': [[20]],
|
|
'LRSB': [[64]],
|
|
'LWSA': [[41]],
|
|
'LWSB': [[90]],
|
|
'LCC' : [[95, 95]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
#index and code pair
|
|
syncTable = [-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="for LRB1"),
|
|
29, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="wait for LRB0. For code path 0, this is actually wait for LRB0 + 1/16 LRA0"),
|
|
29, SBarrier(comment="for GRA"),
|
|
47, SWaitCnt(dscnt=0, vlcnt=14, vscnt=-1, comment="wait for previous GRB and LRA0"),
|
|
47, SBarrier(comment="for GRB"),
|
|
70, SWaitCnt(dscnt=-1, vlcnt=14-3, vscnt=-1, comment="wait for previous GRA"),
|
|
70, SBarrier(comment="for GRB"),
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[18, 19,20,21,22,23,24,25,26]],
|
|
'GRIncB': [[9,10,11,12,13,14,15,16,17]],
|
|
|
|
'LRB0': [[-1,1, 3,5, 7,9, 11,13, 15,17, 19,21],
|
|
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]],
|
|
'LRA0': [[23, 24, 25, 26, 27, 28, 29, 30,31, 32,33, 34,35, 36,37, 38],
|
|
[24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39]],
|
|
'GRA': [[29,29, 31,31, 33,33, 38,38, 40,40, 41,41],
|
|
[30, 30, 32, 32, 34, 34, 39, 39, 41, 41, 42, 42]],
|
|
|
|
'GRB': [[57,57, 59,59, 61,61, 63,63, 65,65, 70,70, 75,75, 80,80],
|
|
[58, 58, 60, 60, 62, 62, 64, 64, 66, 66, 71, 71, 76, 76, 81, 81]],
|
|
'LRB1': [[47,48, 50,52, 54,56, 58,60, 62,64, 66,68],
|
|
[48, 49, 51, 53, 55, 57, 59, 61, 63, 65, 67, 69]],
|
|
'LRA1': [[70,71, 72,73, 74,75, 76,77, 78,79, 80,81, 82,83, 84,85],
|
|
[71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86]],
|
|
|
|
'LRSB': [[39]],
|
|
'LRSA': [[46]],
|
|
'LWSB': [[78]],
|
|
'LWSA': [[95]],
|
|
'LCC' : [[95, 95]],}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNN(kernel) and useLDSTr and TLDS == 1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
#index and code pair
|
|
syncTable = [-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA1"),
|
|
15, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="wait for LRB0"),
|
|
15, SBarrier(comment=""),
|
|
46, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
51, SWaitCnt(dscnt=-1, vlcnt=14, vscnt=-1, comment="wait for previous set of global reads"),
|
|
51, SBarrier(comment=""),
|
|
63, SWaitCnt(dscnt=-1, vlcnt=14-4, vscnt=-1, comment="wait for previous set of global reads"),
|
|
63, SBarrier(comment=""),
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[35,36,37,38,39,40,41,42,43]],
|
|
'GRIncB': [[0,1,2,3,4,5,6,7,8]],
|
|
|
|
'LRB0': [[-1, 0, 1, 2, 3, 4],
|
|
[0, 1, 2, 3, 4, 5]],
|
|
'LRA0': [[6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 25, 27, 29, 31, 33, 35],
|
|
[7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 26, 28, 30, 32, 34, 36]],
|
|
'GRA': [[15,15, 17,17, 27,27, 29,29, 31,31, 33,33],
|
|
[16, 16, 18, 18, 28, 28, 30, 30, 32, 32, 34, 34]],
|
|
|
|
'GRB': [[51,51, 53,53, 55,55, 57,57, 67,67, 69,69, 71,71, 73,73],
|
|
[52,52, 54,54, 56,56, 58,58, 68,68, 70,70, 72,72, 74,74]],
|
|
'LRB1': [[51, 53, 55, 57, 59, 61],
|
|
[52, 54, 56, 58, 60, 62]],
|
|
'LRA1': [[63, 65, 67, 69, 71, 73, 75, 77, 79, 81, 83, 85, 87, 89, 91, 93],
|
|
[64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94]],
|
|
|
|
'LRSB': [[14]],
|
|
'LRSA': [[45]],
|
|
'LWSB': [[94]],
|
|
'LWSA': [[94]],
|
|
'LCC' : [[95, 95]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 256, 128, 2, 0, 1, False, 0, 0),
|
|
dtype_predicate=is8bit,
|
|
vector_widths=[16, 16, 16],
|
|
matrix_inst=[16, 16, 128, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x256x128_8bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
|
|
plr = 3 if kernel["ForceUnrollSubIter"] else 1
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and TLDS == 1:
|
|
optSchedule = {
|
|
'SYNC' : [[6,7, 20,21, 46,47, 61]],
|
|
'GRIncA' : [[0,1,2,3,4,4,4,4,4]],
|
|
'GRIncB' : [[5,5,5,5,5,6,6,6,6]],
|
|
'LRA0' : [[0,0, 1,1, 2,2, 3,3]],
|
|
'GRA' : [[8,8,9,9,10,10,11,11,12,12, 23,23,24,24,25,25]],
|
|
'LRB0' : [[13,13,14,14,15,15,16,16]],
|
|
'LRA%u'%plr : [[48,48,49,49,50,50,51,51]],
|
|
'LRB%u'%plr : [[52,52,54,54,55,55,56,56]],
|
|
'GRB' : [[26,26,27,27, 39,39,40,40,41,41,42,42,43,43, 53,53]],
|
|
'LCC' : [[60, 60]],
|
|
'LRSA' : [[17]],
|
|
'LRSB' : [[17]],
|
|
'LWSA' : [[57]],
|
|
'LWSB' : [[57]],
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0/LRB0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0/LRB0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=15, vscnt=-1, comment="Wait for GRA to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for PLR to complete")]
|
|
nglshift = nllshift = 16
|
|
else:
|
|
return False, None
|
|
|
|
numMfma = 64
|
|
# B0A0, B0A1, B1A0, B1A1
|
|
mfmaReorder = []
|
|
kernel["MfmaInitCVgprs"] = True
|
|
if not kernel["ForceUnrollSubIter"]:
|
|
mfmaReorder = [0,1,2,3, 8,9,10,11, 16,17,18,19, 24,25,26,27,
|
|
4,5,6,7, 12,13,14,15, 20,21,22,23, 28,29,30,31,
|
|
32,33,34,35, 40,41,42,43, 48,49,50,51, 56,57,58,59,
|
|
36,37,38,39, 44,45,46,47, 52,53,54,55, 60,61,62,63]
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift, mfmaReorder)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and TLDS == 1:
|
|
optSchedule = {
|
|
'SYNC' : [[19,20, 50,51, 67,68, 104, 105, 127]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[9,10,11,12,13,14,15,16,17]],
|
|
'LRA0' : [[0,2,4,6,8,10,12,14],
|
|
[1,3,5,7,9,11,13,15]],
|
|
'LRB0' : [[24,27,30,33,36,38,40,42],
|
|
[22,25,28,31,34,37,39,41]],
|
|
'GRA' : [[21,22, 23,25, 26,28, 29,31, 32,34, 35,52, 53,55, 56,58],
|
|
[21,23, 24,26, 27,29, 30,32, 33,35, 36,53, 54,56, 57,59]],
|
|
'GRB' : [[59,61, 62,64, 65,85, 86,87, 88,89, 94,96, 98,100, 102,124],
|
|
[60,62, 63,65, 66,84, 85,86, 87,88, 93,95, 97,99, 103,123]],
|
|
'LRA1' : [[69,71,73,75,77,79,81,83],
|
|
[70,72,74,76,78,80,82,90]],
|
|
'LRB1' : [[106,108,110,112,114,116,118,120],
|
|
[107,109,111,113,115,117,119,121]],
|
|
'LRSA' : [[16]],
|
|
'LRSB' : [[83]],
|
|
'LWSA' : [[125]],
|
|
'LWSB' : [[125]],
|
|
'LCC' : [[126, 126]],
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=(2 + 8 + 8), vscnt=-1, comment="Wait for previous GRA to completely"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=15, vscnt=-1, comment="Wait for previous GRB to completely"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 and 3/8 LRB1 to complete")]
|
|
nglshift = nllshift = 16
|
|
elif isNT(kernel) and not useLDSTr and TLDS == 0:
|
|
kernel["UsePLRPack"] = 1
|
|
|
|
optSchedule = {
|
|
'SYNC' : [[12,13, 36,44, 56,59, 66,68, 73,92]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[28,29,30,31,32,33,34,35,36]],
|
|
'LRA0' : [[0,0,2,2,4,4,6,6],
|
|
[1,1,3,3,5,5,7,7]],
|
|
'LRB0' : [[8,8,10,10,15,15,18,18],
|
|
[9,9,11,11,14,14,17,17]],
|
|
'GRA' : [[14,14, 17,17, 20,20, 23,23, 26,26, 45,45, 48,48, 51,51],
|
|
[15,15, 18,18, 21,21, 24,24, 27,27, 46,46, 49,49, 52,52]],
|
|
'GRB' : [[54,54, 57,57, 87,87,90,90,93,93,96,96,99,99, 123,123],
|
|
[55,55, 58,58, 88,88,91,91,94,94,97,97,100,100, 124,124]],
|
|
'LRA1' : [[60,60,62,62,64,64,66,66],
|
|
[61,61,63,63,65,65,67,67]],
|
|
'LRB1' : [[69,69,71,71,73,73,75,75],
|
|
[70,70,72,72,74,74,76,76]],
|
|
'LRSA' : [[59]],
|
|
'LRSB' : [[59]],
|
|
'LWSA' : [[125]],
|
|
'LWSB' : [[125]],
|
|
'LCC' : [[126, 126]],
|
|
'PackA0' : [[16,16, 19,19, 21,21, 22,22, 24,24, 25,25, 27,27, 28,28, 29,29, 30,30, 31,31, 32,32, 33,33, 34,34, 35,35, 36,36],
|
|
[16,16, 19,19, 20,20, 22,22, 23,23, 25,25, 26,26, 28,28, 29,29, 30,30, 31,31, 32,32, 33,33, 34,34, 35,35, 36,36]],
|
|
'PackB0' : [[37,37, 38,38, 39,39, 40,40, 41,41, 42,42, 43,43, 46,46, 47,47, 49,49, 50,50, 52,52, 53,53, 55,55, 56,56, 58,58],
|
|
[37,37, 38,38, 39,39, 40,40, 41,41, 42,42, 43,43, 45,45, 47,47, 48,48, 50,50, 51,51, 53,53, 54,54, 56,56, 57,57]],
|
|
'PackA1' : [[74,74, 76,76, 77,77, 78,78, 79,79, 80,80, 81,81, 82,82, 83,83, 84,84, 85,85, 86,86, 88,88, 89,89, 91,91, 92,92],
|
|
[75,75, 77,77, 78,78, 79,79, 80,80, 81,81, 82,82, 83,83, 84,84, 85,85, 86,86, 87,87, 89,89, 90,90, 92,92, 93,93]],
|
|
'PackB1' : [[94,94, 95,95, 97,97, 98,98, 100,100, 101,101, 102,102, 103,103, 104,104, 105,105, 106,106, 107,107, 108,108, 109,109, 110,110, 111,111],
|
|
[95,95, 96,96, 98,98, 99,99, 101,101, 102,102, 103,103, 104,104, 105,105, 106,106, 107,107, 108,108, 109,109, 110,110, 111,111, 112,112]],
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=17, vscnt=-1, comment="Wait for GRA to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=9, vscnt=-1, comment="Wait for GRB to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB1 to complete")]
|
|
nglshift = nllshift = 16
|
|
elif (isNN(kernel) or isTT(kernel)) and not useLDSTr and TLDS == 1:
|
|
kernel["UsePLRPack"] = 1
|
|
|
|
optSchedule = {
|
|
'SYNC' : [[8, 12,13, 36,44, 56,59, 66,68, 74, 127]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[28,29,30,31,32,33,34,35,36]],
|
|
'LRA0' : [[0,0,2,2,4,4,6,6],
|
|
[1,1,3,3,5,5,7,7]],
|
|
'LRB0' : [[9,11, 15,18,21,24,27,30],
|
|
[10,12, 14,17,20,23,26,29]],
|
|
'GRA' : [[14,14, 17,17, 20,20, 23,23, 26,26, 45,45, 48,48, 51,51],
|
|
[15,15, 18,18, 21,21, 24,24, 27,27, 46,46, 49,49, 52,52]],
|
|
'GRB' : [[54,54, 57,57, 87,87,90,90,93,93,96,96,99,99, 123,123],
|
|
[55,55, 58,58, 88,88,91,91,94,94,97,97,100,100, 124,124]],
|
|
'LRA1' : [[60,60,62,62,64,64,66,66],
|
|
[61,61,63,63,65,65,67,67]],
|
|
'LRB1' : [[68,70,72,74,76,78,80,82],
|
|
[69,71,73,75,77,79,81,83]],
|
|
'LRSA' : [[59]],
|
|
'LRSB' : [[59]],
|
|
'LWSA' : [[125]],
|
|
'LWSB' : [[125]],
|
|
'LCC' : [[126, 126]],
|
|
'PackA0' : [[8,8, 16,16, 19,19, 22,22, 25,25, 28,28, 29,29, 31,31, 32,32, 33,33, 34,34, 35,35, 36,36, 37,37, 38,38, 39,39]],
|
|
'PackA1' : [[75,75, 77,77, 79,79, 81,81, 83,83, 84,84, 85,85, 86,86, 88,88, 89,89, 91,91, 92,92, 94,94, 95,95, 97,97, 98,98],
|
|
[74,74, 76,76, 78,78, 80,80, 82,82, 84,84, 85,85, 86,86, 87,87, 89,89, 90,90, 92,92, 93,93, 95,95, 96,96, 98,98]],
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 first half to complete"),
|
|
SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=17, vscnt=-1, comment="Wait for GRA to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=9, vscnt=-1, comment="Wait for GRB to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB1 to complete")]
|
|
if isTT(kernel):
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
|
|
optSchedule['GRIncA'], optSchedule['GRIncB'] = optSchedule['GRIncB'], optSchedule['GRIncA']
|
|
optSchedule['LRA0'], optSchedule['LRB0'] = optSchedule['LRB0'], optSchedule['LRA0']
|
|
optSchedule['LRA1'], optSchedule['LRB1'] = optSchedule['LRB1'], optSchedule['LRA1']
|
|
optSchedule['PackB0'] = optSchedule['PackA0']
|
|
optSchedule['PackB1'] = optSchedule['PackA1']
|
|
del optSchedule['PackA0'], optSchedule['PackA1']
|
|
nglshift = nllshift = 16
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 128
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(160, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_160x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 80
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and TLDS==1:
|
|
optSchedule = {
|
|
|
|
'SYNC' : [[-1, 4, 13,13, 38,39, 42,43, 70,70]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[29,30,31,32,33,34,35,36,37]],
|
|
|
|
'LRA0' : [[0,2,3,4,5]], ## -2 is place holder
|
|
|
|
'LRB0' : [[13,15,18,21,24,26,28,30], ## After LRA0, we can mix LRB0 and GRA
|
|
[14,16,19,22,25,27,29,31]],
|
|
## GRA should start after LRA0 is done.
|
|
'GRA' : [[11,14, 17,17, 20,20, 23,23, 26,27],
|
|
[12,15, 18,18, 21,21, 24,24, 27,28]],
|
|
|
|
## GRB should start after LRB0 is done
|
|
'GRB' : [[40,40, 43,43, 46,46, 49,49, 59,59, 62,62, 65,65, 67,68], # m0 inc is part of GRA/GRB
|
|
[41,41, 44,44, 47,47, 57,57, 60,60, 63,63, 66,66, 68,69]],
|
|
'LRA1' : [[44, 47, 53, 58, 63],
|
|
[45, 48, 54, 59, 64]],
|
|
|
|
#After GRB is done.
|
|
'LRB1' : [[70,71,72,73,75,76,77,78]],
|
|
|
|
'LRSA' : [[33]], # after LRA0 and before LRA1
|
|
'LRSB' : [[33]], # after LRB0 and before LRB2
|
|
'LWSA' : [[74]], # For A
|
|
'LWSB' : [[76]],
|
|
|
|
'LCC' : [[79, 79]],
|
|
}
|
|
# note: syncCode needs to be
|
|
syncCode = [SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="Wait for necessary prior LRA1/LRB1 before starting main loop"),
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 for the remaining main loop"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete to start GRA"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete to start GRB"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=(13 + 1), vscnt=-1, comment="Wait for GRA to complete to start LRA1"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=13, vscnt=-1, comment="Wait for GRB to complete to start LRB1"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 13 # vmcnt shift for ngl and nll
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNN(kernel) and useLDSTr and TLDS==1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
optSchedule = {
|
|
'SYNC' : [[-1,
|
|
12, 12, # Wait for B
|
|
24, 24, # wait LRB0.
|
|
41, 41,
|
|
61, 61 # wait GRA.
|
|
]],
|
|
# Addr. update (be done before GRA/GRB).
|
|
'GRIncA' : [[5,6,7,8,9,10,11,12,13]],
|
|
'GRIncB' : [[0,0,1,1,2,2,3,3,4]],
|
|
# Current iteration.
|
|
'LRA0' : [[8,9,10,11,12,13,14,15,16,17]],
|
|
'LRB0' : [[0,1,2,3,4,5,6,7]],
|
|
# Buffer loads.
|
|
'GRB' : [[51,51, 55,55, 59,61, 76,77, 78,78]],
|
|
'GRA' : [[11,12, 16,16, 20,20, 24,24, 28, 28, 32,32, 36, 36, 40, 40]],
|
|
# Prefetch next iteration.
|
|
'LRA1' : [[62,63,64,65,66,67,68,69,70,71]],
|
|
'LRB1' : [[41,42,43,44,45,46,47,49]],
|
|
'LRSA' : [[39]],
|
|
'LRSB' : [[39]],
|
|
'LWSA' : [[60]],
|
|
'LWSB' : [[60]],
|
|
'LCC' : [[79, 79]], # Loop control.
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 LRB1"),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRB0"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=13, vscnt=-1, comment="Wait for previous GRA(B) to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="Wait for previous GRB(A) to complete"),
|
|
SBarrier(comment="")]
|
|
nglshift = nllshift = 13
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNT(kernel) and useLDSTr and TLDS==0:
|
|
optSchedule = {
|
|
'SYNC': [[-1,17,17,57,57]],
|
|
'GRA': [[16,17,20,20,24,24,28,28,31,31]],
|
|
'GRB': [[35,35,39,39,68,68,70,70,71,71,76,76,77,77,78,78]],
|
|
'GRIncA': [[0,0,1,1,2,2,3,3,4]],
|
|
'GRIncB': [[4,5,5,13,13,13,14,14,14]],
|
|
'LCC': [[79,79]],
|
|
'LRA0': [[0,1,1,2,2,3,3,4,4,5]],
|
|
'LRA1': [[58,60,62,64,66,67,67,68,68,69]],
|
|
'LRB0': [[0,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12]],
|
|
'LRB1': [[59,61,63,65,67,71,72,72,73,73,74,74,75,75,76,76]],
|
|
'LRSA': [[38]],
|
|
'LRSB': [[38]],
|
|
'LWSA': [[61]],
|
|
'LWSB': [[61]]
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0 for iteration == 0"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=7, vscnt=-1, comment="wait for previous set of global reads"),
|
|
SBarrier(comment="")
|
|
]
|
|
nglshift = nllshift = 13
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(96, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_96x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
kernel["MfmaInitCVgprs"] = True
|
|
|
|
numMfma = 48
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 11
|
|
|
|
if isTN(kernel) and TLDS==1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
|
|
syncTable = [
|
|
7, SWaitCnt(dscnt=8, vlcnt=-1, vscnt=-1, comment="Finish all LRA1s and LRB1s"),
|
|
|
|
9, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="All LRB0 done"),
|
|
9, SBarrier(comment=""),
|
|
|
|
23, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="All LRA0 done"),
|
|
|
|
35, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Wait for prev GRBs"),
|
|
35, SBarrier(comment=""),
|
|
|
|
43, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Wait for prev GRA"),
|
|
43, SBarrier(comment=""),
|
|
|
|
47, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Finish all LRB1 and 1/3 LRA1"),
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
'GRIncA' : [[0,0,0, 2,2,2, 3, 4,4],
|
|
[-1,-1,-1, 1,1,1, 3, 3,3]],
|
|
'GRIncB' : [[4, 5,5,5, 6,6,6, 7,7],
|
|
[4, 5,5,5, 6,6,6, 7,7]],
|
|
|
|
'LRB0' : [[-1,-1,-1, 1,1,1, 3,3],
|
|
[0,0,0, 2,2,2, 4,4]],
|
|
'LRSB' : [[8], [9]],
|
|
|
|
'SYNC' : [syncTable[::2]],
|
|
|
|
# Actually loads GRB
|
|
'GRA' : [[8,9, 11,11, 13,13, 15,15, 20,20, 22,22, 24,24, 26,26],
|
|
[8,10, 12,12, 14,14, 16,16, 21,21, 23,23, 25,25, 27,27]],
|
|
'LWSB' : [[32]],
|
|
|
|
'LRA0' : [[17, 17, 17],
|
|
[15, 15, 15]],
|
|
'LRSA' : [[19]],
|
|
|
|
# Actually loads GRA
|
|
'GRB' : [[36,36, 38,38, 40,40],
|
|
[37,37, 39,39, 41,41]],
|
|
'LWSA' : [[45]],
|
|
|
|
'LRB1' : [[35,35, 37,37, 39,39, 41,41],
|
|
[36,36, 38,38, 40,40, 42,42]],
|
|
'LRA1' : [[43,44,46],
|
|
[43,45,46]],
|
|
|
|
'LCC' : [[47, 47]],
|
|
}
|
|
|
|
# Reorder to basically create the 256x96 case
|
|
mfmaReorder = [
|
|
0, 3, 6, 9, 12, 15, 18, 21,
|
|
1, 4, 7, 10, 13, 16, 19, 22,
|
|
2, 5, 8, 11, 14, 17, 20, 23,
|
|
# Second half
|
|
24, 27, 30, 33, 36, 39, 42, 45,
|
|
25, 28, 31, 34, 37, 40, 43, 46,
|
|
26, 29, 32, 35, 38, 41, 44, 47
|
|
]
|
|
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift, mfmaReorder=mfmaReorder)
|
|
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
# A: MIWaveTileA=3 => 2*3 = 6 local/global reads
|
|
# B: MIWaveTileB=8 => 2*8 = 16 local/global reads
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=12, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LR/LW for iteration == 0"),
|
|
5, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LR/LW for iteration == 0"),
|
|
12, SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="wait for LRA0 to complete before GRA"),
|
|
12, SBarrier(comment="barrier after LRA0, before GRA"),
|
|
23, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0 to complete before GRB"),
|
|
23, SBarrier(comment="barrier after LRB0, before GRB"),
|
|
|
|
25, SWaitCnt(dscnt=-1, vlcnt=11+1, vscnt=-1, comment="wait for global reads before next-tile LR"),
|
|
25, SBarrier(comment="final barrier before LRA1/LRB1"),
|
|
36, SWaitCnt(dscnt=-1, vlcnt=11-2, vscnt=-1, comment="wait for global reads before next-tile LR"),
|
|
36, SBarrier(comment="final barrier before LRA1/LRB1"),
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
|
|
# Address increments
|
|
'GRIncA' : [[0, 0, 1, 1, 2, 2, 3, 3, 4]],
|
|
'GRIncB' : [[4, 5, 5, 6, 6, 7, 7, 8, 8]],
|
|
|
|
# Current-iteration local reads
|
|
'LRA0' : [[0, 1, 2, 3, 4, 5],
|
|
[1, 2, 3, 4, 4, 6]],
|
|
# Ensure LRB0 completes early enough for upcoming MFMA consumers.
|
|
'LRB0' : [[6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
|
|
[7, 8, 9, 10, 11, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]],
|
|
|
|
# Global reads (DirectToLds). Issue after LRA0 barrier, while finishing LRB0.
|
|
# Do not interleave GRA between GRIncB indices (SCC hazard).
|
|
'GRA' : [[12, 12, 18, 18, 21, 21],
|
|
[13, 13, 19, 19, 22, 22]],
|
|
# Start GRB after the LRB0 barrier. Use stride=2 starting at 24.
|
|
'GRB' : [[24, 24, 26, 26, 28, 28, 30, 30, 32, 32, 34, 34, 38, 38, 40, 40],
|
|
[24, 24, 27, 27, 29, 29, 31, 31, 33, 33, 35, 35, 37, 37, 39, 39]],
|
|
|
|
# Next-iteration local reads
|
|
'LRA1' : [[25, 26, 27, 28, 29, 30],
|
|
[26, 27, 28, 29, 30, 31]],
|
|
# Keep within the mfma window; repeated indices are allowed (multiple instructions scheduled at same mfma slot).
|
|
'LRB1' : [[36, 37, 37, 38, 38, 39, 39, 40, 40, 41, 41, 42, 43, 45, 46, 47],
|
|
[37, 38, 38, 39, 39, 40, 40, 41, 41, 42, 42, 43, 44, 46, 47, 47]],
|
|
|
|
# Epilogue-related
|
|
'LRSA' : [[22]],
|
|
'LRSB' : [[23]],
|
|
'LWSA' : [[43]],
|
|
'LWSB' : [[44]],
|
|
'LCC' : [[45, 46]],
|
|
}
|
|
|
|
syncCode = syncTable[1::2]
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNN(kernel) and not useLDSTr and TLDS==1:
|
|
snops = []
|
|
syncs = SyncSchedule()
|
|
syncs.add(-1, dscnt=4, comment="Wait for prior local read local write")
|
|
syncs.add(2, dscnt=3, comment="Wait for prior local read local write")
|
|
syncs.add(8, dscnt=6, comment="Wait for partial LRA0")
|
|
syncs.add(16, dscnt=0, barrier=True, comment="Wait for LRA0 to complete before starting LRB0+GRA")
|
|
syncs.add(23, dscnt=0, vlcnt=3, barrier=True, comment="Wait for LRB0+GRA")
|
|
|
|
snopIdxs = [1, 25]
|
|
snops = [[x, SNop(1, comment="")] for x in snopIdxs]
|
|
|
|
lra0 = [0,0,1,1,2,2,3,3,4,5,5,6,6,7,7,8,8,9,10,11,12,13,14,15]
|
|
|
|
# Issue LRB0+GRA after LRA0 completes.
|
|
lrb0 = [16, 17, 18, 19, 20, 21, 22, 22]
|
|
grA = [18, 18, 20, 20, 21, 21]
|
|
|
|
# Issue LRA1+GRB after LRB0+GRA complete.
|
|
lra1 = [24,24, 25,25, 26,26, 27,27, 28, 29,29, 30,30, 31,31, 32,32, 33,34,35,36,37,38,39]
|
|
grB = [24,24,26,26,28,28,30,30,32,32,36,36,38,38,40,40]
|
|
lrb1 = [39,40,41,42,43,44,45,46]
|
|
|
|
packA1 = [
|
|
-1,-1,-1,-1,-1,-1,
|
|
0,0,0,0,
|
|
1,1,
|
|
]
|
|
packA0 = [
|
|
23,23,23,23,23,23,
|
|
24,24,24,24,
|
|
25,25,
|
|
]
|
|
|
|
# GRIncs should be ordered AFTER LRs.
|
|
grIncA = [0,1,2,3,4,5,6,7,8]
|
|
grIncB = [9,10,11,12,13,14,15,16,17]
|
|
|
|
lwsa = [46]
|
|
lwsb = [46]
|
|
lrsa = [22]
|
|
lrsb = [22]
|
|
num_gr = (len(grA) + len(grB)) // 2
|
|
optSchedule = {
|
|
'SYNC' : [syncs.get_indicies()],
|
|
'LRA0' : [lra0],
|
|
'LRA1' : [lra1],
|
|
'PackA0' : [packA0],
|
|
'PackA1' : [packA1],
|
|
'LRB0' : [lrb0],
|
|
'LRB1' : [lrb1],
|
|
'GRIncA' : [grIncA],
|
|
'GRIncB' : [grIncB],
|
|
'GRA' : [grA],
|
|
'GRB' : [grB],
|
|
'LRSA' : [lrsa],
|
|
'LRSB' : [lrsb],
|
|
'LWSA' : [lwsa],
|
|
'LWSB' : [lwsb],
|
|
'LCC' : [[47, 47]],
|
|
}
|
|
nllshift = nglshift = num_gr
|
|
if snops:
|
|
optSchedule['SNOP'] = [[s[0] for s in snops]]
|
|
snopCode = [s[1] for s in snops]
|
|
|
|
opt1 = ScheduleInfo(1, 48, optSchedule=optSchedule, syncCode=syncs.get_code(), nglshift=nglshift, nllshift=nllshift, snopCode=snopCode)
|
|
else:
|
|
return False, None
|
|
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 160, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x160x64_16bit(kernel, useLDSTr, TLDS):
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
numMfma = 80
|
|
if isNN(kernel) and useLDSTr and TLDS==1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
optSchedule = {
|
|
'SYNC' : [[-1,
|
|
12,12, # Wait for LRB0
|
|
24, 24,# Wait LRA0.
|
|
41,41,
|
|
61, 61 # Wait previous GR.
|
|
]],
|
|
# Addr. update (be done before GRA/GRB).
|
|
'GRIncA' : [[21,21,21,22,22,22,23,23,23]],
|
|
'GRIncB' : [[0,1,2,3,4,5,6,7,8]],
|
|
# Current iteration.
|
|
'LRA0' : [[5,5,7,7,9,9,11,11,13,13,15,15,17,18,19,20]],
|
|
'LRB0' : [[0,0,1,2,3]],
|
|
# Buffer loads.
|
|
'GRB' : [[30,30, 33,33, 36,36, 52,52, 56,56, 60,61, 76,77, 78,78]],
|
|
'GRA' : [[11,12, 16,16, 20,20, 25,25, 26, 28]],
|
|
# Prefetch next iteration.
|
|
'LRA1' : [[62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77]],
|
|
'LRB1' : [[41,42,43,44,45]],
|
|
'LRSA' : [[39]],
|
|
'LRSB' : [[39]],
|
|
'LWSA' : [[60]],
|
|
'LWSB' : [[60]],
|
|
'LCC' : [[79, 79]], # Loop control
|
|
}
|
|
syncCode = [SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA1 LRB1"),
|
|
SWaitCnt(dscnt=8, vlcnt=-1, vscnt=-1, comment="Wait for LRB0"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for previous GRB to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=(13+8-5), vscnt=-1, comment="Wait for previous GRB to complete"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=(13+10-13), vscnt=-1, comment="Wait for previous GRA to complete"),
|
|
SBarrier(comment="")]
|
|
nglshift = nllshift = 13
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isTN(kernel) and (not useLDSTr) and TLDS==1:
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1 (partial) before starting main loop"),
|
|
4, SWaitCnt(dscnt=0+2, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1 (complete) for the remaining main loop"),
|
|
14, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete to start GRB"),
|
|
14, SBarrier(comment=""),
|
|
# Must be dscnt=0 here: validator requires proving all LRA0 are complete
|
|
# before the first GRA is issued (vmfma_index window [31,41)).
|
|
39, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 / ensure LRA0 complete before starting GRB/GRA"),
|
|
39, SBarrier(comment=""),
|
|
45, SWaitCnt(dscnt=-1, vlcnt=13+2, vscnt=-1, comment="Wait for GRB to complete before LRB1"),
|
|
45, SBarrier(comment=""),
|
|
69, SWaitCnt(dscnt=-1, vlcnt=13, vscnt=-1, comment="Wait for GRA to complete before LRA1"),
|
|
69, SBarrier(comment=""),
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA' : [[29,30,31,32,33,34,35,36,37]],
|
|
'GRIncB' : [[0,1,2,3,4,5,6,7,8]],
|
|
|
|
# Current iteration.
|
|
'LRB0' : [[0,2,3,4,5],
|
|
[1,3,4,5,6]],
|
|
'LRA0' : [[13,15,18,21,24,26,28,30],
|
|
[13,16,19,22,25,27,29,31]],
|
|
|
|
# GRB must not start before the SYNC at idx 15 (LRB0 completion).
|
|
'GRB' : [[14,14, 17,17, 20,20, 23,23, 26,26],
|
|
[15,15, 18,18, 21,21, 24,24, 27,27]],
|
|
# Buffer loads.
|
|
'GRA' : [[40,40, 43,43, 46,46, 49,49, 59,59, 62,62, 65,65, 67,67],
|
|
[41,41, 44,44, 47,47, 57,57, 60,60, 63,63, 66,66, 68,68]],
|
|
# Prefetch next iteration.
|
|
# Need 5 local reads for B (MIWaveTileB=5).
|
|
'LRB1' : [[45,46,47,48,49],
|
|
[46,47,48,49,50]],
|
|
# Need 8 local reads for A (MIWaveTileA=8) in each code path.
|
|
# Path1 LRA1 must be earlier than path0 (validator requirement).
|
|
'LRA1' : [[69, 70, 71, 72, 73, 74, 75, 76],
|
|
[70, 71, 72, 73, 74, 75, 76, 77]],
|
|
'LRSA' : [[32]],
|
|
'LRSB' : [[33]],
|
|
'LWSA' : [[74]],
|
|
'LWSB' : [[76]],
|
|
'LCC' : [[77, 78]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 13
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNT(kernel) and useLDSTr and TLDS==0:
|
|
nglshift = nllshift = 0
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
optSchedule = {
|
|
'SYNC': [[-1,17,17,57,57]],
|
|
'GRA': [[16,17,20,20,24,24,28,28,31,31]],
|
|
'GRB': [[35,35,39,39,68,68,70,70,71,71,76,76,77,77,78,78]],
|
|
'GRIncA': [[0,0,1,1,2,2,3,3,4]],
|
|
'GRIncB': [[4,5,5,13,13,13,14,14,14]],
|
|
'LCC': [[79,79]],
|
|
'LRA0': [[0,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12]],
|
|
'LRB0': [[0,1,1,2,2,3,3,4,4,5]],
|
|
'LRA1': [[59,61,63,65,67,71,72,72,73,73,74,74,75,75,76,76]],
|
|
'LRB1': [[58,60,62,64,66,67,67,68,68,69]],
|
|
'LRSA': [[38]],
|
|
'LRSB': [[38]],
|
|
'LWSA': [[61]],
|
|
'LWSB': [[61]]
|
|
}
|
|
syncCode = [
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0 for iteration == 0"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=7, vscnt=-1, comment="wait for previous set of global reads"),
|
|
SBarrier(comment="")
|
|
]
|
|
nglshift = nllshift = 13
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 240, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 2, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[4, 1]
|
|
)
|
|
def _get_schedule_256x240x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0
|
|
if isTN(kernel) and TLDS==1:
|
|
optSchedule = {
|
|
'GRIncA': [[0, 0, 1, 1, 2, 2, 3, 3, 4]],
|
|
'GRIncB': [[30, 30, 31, 31, 32, 32, 33, 33, 34]],
|
|
'LRA0': [[0, 1, 1, 2]],
|
|
'LRB0': [[3, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29]],
|
|
'GRA': [[5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]],
|
|
'GRB': [[35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 89, 90, 90, 91, 91]],
|
|
'LRA1': [[93, 94, 95, 96]],
|
|
'LRB1': [[97, 98, 99, 100, 102, 104, 106, 108, 110, 112, 114, 116, 116, 116, 116]],
|
|
'LRSA': [[59]],
|
|
'LRSB': [[59]],
|
|
'LWSA': [[91]],
|
|
'LWSB': [[91]],
|
|
'LCC': [[119, 119]],
|
|
'SYNC': [[-1, -1, 4, 4, 33, 33, 92, 92]],
|
|
}
|
|
nglshift = 38
|
|
nllshift = 38
|
|
syncCode = [
|
|
SBarrier(comment="wavefront sync at loop start"),
|
|
SWaitCnt(dscnt=13, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LR/LW"),
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="wait for LRA0 to complete before GRA DirectToLds"),
|
|
SBarrier(comment="barrier after LRA0 (idx 3), before GRA starts (idx 5)"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0 to complete before GRB DirectToLds"),
|
|
SBarrier(comment="barrier after LRB0 (idx 29), before GRB starts (idx 35)"),
|
|
SWaitCnt(dscnt=-1, vlcnt=38, vscnt=-1, comment="wait for global reads before using data"),
|
|
SBarrier(comment="earlier final barrier to reduce idle time"),
|
|
]
|
|
elif isNT(kernel) and TLDS==0 and useLDSTr:
|
|
optSchedule = {
|
|
'LRA0': [[0, 1, 1, 2, 2, 3, 3, 4]],
|
|
'LRB0': [[0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]],
|
|
'LRA1': [[98, 99, 99, 100, 100, 101, 101, 102]],
|
|
'LRB1': [[98, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 119, 119, 119, 119, 119, 119, 119, 119, 119, 119, 119, 119]],
|
|
'GRA': [[8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]],
|
|
'GRB': [[40, 41, 43, 44, 46, 47, 49, 50, 52, 53, 55, 56, 58, 59, 61, 62, 64, 65, 67, 68, 70, 71, 73, 74, 76, 77, 79, 80, 82, 83, 85, 86, 88, 89, 91, 92, 94, 95, 97, 98, 100, 101, 103, 104, 106, 107, 109, 110, 112, 113, 115, 116, 118, 119, 119, 119, 119, 119, 119, 119]],
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
'LCC': [[119, 119]],
|
|
'LRSA': [[57]],
|
|
'LRSB': [[57]],
|
|
'LWSA': [[96]],
|
|
'LWSB': [[96]],
|
|
'SYNC': [[-1, 6, 6, 38, 38, 96, 96]],
|
|
}
|
|
nglshift = 38
|
|
nllshift = 38
|
|
syncCode = [
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LR/LW"),
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="wait for LRA0 to complete before GRA DirectToLds"),
|
|
SBarrier(comment="barrier after LRA0 (idx 4), before GRA starts (idx 8)"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0 to complete before GRB DirectToLds"),
|
|
SBarrier(comment="barrier after LRB0 (idx 33), before GRB starts (idx 40)"),
|
|
SWaitCnt(dscnt=-1, vlcnt=27, vscnt=-1, comment="wait for 54 global reads before idx 96 (16 GRA + 38 GRB). vlcnt = 38 - 11 = 27"),
|
|
SBarrier(comment="barrier at idx 96 - before LRA1/LRB1 start at 98"),
|
|
]
|
|
elif isNN(kernel) and TLDS==1 and useLDSTr:
|
|
optSchedule = {
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
'LRA0': [[0, 1, 1, 2, 2, 3, 3, 4]],
|
|
'LRB0': [[1, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]],
|
|
'GRA': [[8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]],
|
|
'GRB': [[26, 27, 29, 30, 32, 33, 35, 36, 38, 39, 41, 42, 44, 45, 47, 48, 50, 51, 53, 54, 56, 57, 59, 60, 62, 63, 65, 66, 68, 69, 71, 72, 74, 75, 77, 78, 80, 81, 83, 84, 86, 87, 89, 90, 92, 93, 95, 96, 98, 99, 101, 102, 104, 105, 107, 108, 110, 111, 113, 114]],
|
|
'LRA1': [[93, 95, 95, 96, 96, 97, 97, 98]],
|
|
'LRB1': [[94, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112]],
|
|
'LRSA': [[57]],
|
|
'LRSB': [[57]],
|
|
'LWSA': [[91]],
|
|
'LWSB': [[91]],
|
|
'LCC': [[119, 119]],
|
|
'SYNC': [[-1, 6, 6, 26, 26, 90, 90]],
|
|
}
|
|
nglshift = 38
|
|
nllshift = 38
|
|
syncCode = [
|
|
SWaitCnt(dscnt=13, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LR/LW"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA0 to complete before GRA DirectToLds"),
|
|
SBarrier(comment="barrier after LRA0 (idx 4), before GRA starts (idx 8)"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0 to complete before GRB DirectToLds"),
|
|
SBarrier(comment="barrier after LRB0 (idx 19), before GRB starts (idx 26)"),
|
|
SWaitCnt(dscnt=-1, vlcnt=30, vscnt=-1, comment="wait for 59 global reads before idx 90 (16 GRA + 43 GRB). vlcnt = 38 - 8 = 30"),
|
|
SBarrier(comment="barrier at idx 90 - before LRA1/LRB1 start at 93/94"),
|
|
]
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 120 # Must match actual MFMA count for 256x240x64 tile
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 208, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 2, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[4, 1]
|
|
)
|
|
def _get_schedule_256x208x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0
|
|
if isTN(kernel) and TLDS==1:
|
|
optSchedule = {
|
|
'SYNC': [[-1, 3, 23, 23, 35, 35, 81, 81]],
|
|
'LRA0': [[0, 1, 2, 4]],
|
|
'LRB0': [[5, 6, 8, 9, 11, 14, 16, 19, 21, 24, 26, 29, 31]],
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
'GRA': [[23, 23, 23, 23, 25, 25, 28, 28, 29, 29, 31, 31, 33, 33, 35, 35]],
|
|
'GRB': [[36, 36, 37, 37, 39, 39, 41, 41, 42, 42, 44, 44, 47, 47, 48, 48, 50, 50, 52, 52, 54, 54, 55, 55, 57, 57, 59, 59, 60, 60, 62, 62, 64, 64, 66, 66, 67, 67, 69, 69, 71, 71, 73, 73, 74, 74, 76, 76, 78, 78, 79, 79]],
|
|
'LRSA': [[50]],
|
|
'LRSB': [[50]],
|
|
'LWSA': [[79]],
|
|
'LWSB': [[80]],
|
|
'LRA1': [[82, 84, 87, 90]],
|
|
'LRB1': [[83, 85, 86, 88, 89, 91, 92, 93, 94, 95, 96, 97, 98]],
|
|
'LCC': [[98, 98]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=8, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LRB1-4"),
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LRB1-remaining"),
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="wait for LRA0 to complete before GRA DirectToLds"),
|
|
SBarrier(comment="barrier after LRA0, before GRA starts"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0 to complete before GRB DirectToLds"),
|
|
SBarrier(comment="barrier after LRB0, before GRB starts"),
|
|
SWaitCnt(dscnt=-1, vlcnt=34, vscnt=-1, comment="wait for all GRA/GRB before next-tile LDS reads"),
|
|
SBarrier(comment="final barrier before LRA1/LRB1"),
|
|
]
|
|
|
|
nglshift = nllshift = 34
|
|
|
|
elif isNN(kernel) and useLDSTr and TLDS==1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
nglshift = nllshift = 0
|
|
|
|
optSchedule = {
|
|
# last index of producer <SYNC> first index of consumer
|
|
# SYNC[0] = -1 to align all waves at the start of the loop
|
|
# A fence at 23, annd B fence at 36, final vmem fence at 81
|
|
'SYNC': [[-1, 3, 7, 11, 23, 36, 36, 81, 81]],
|
|
|
|
# Avoid interleaving of LRA0 and LRB0
|
|
# LRA0: tightly packed at the beginning
|
|
'LRA0': [[0, 1, 2, 3, 4, 5, 6, 7]],
|
|
# LRB0 scheduled after the A fence, overlapping with GRA/GRB
|
|
'LRB0': [[24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 35]],
|
|
|
|
# Address increments for GR
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
|
|
# first stream (A side in NN+swapped)
|
|
'GRB': [[23, 23, 24, 24,
|
|
26, 26, 28, 28,
|
|
29, 29, 31, 31,
|
|
33, 33, 35, 35]],
|
|
|
|
# second stream (B side in NN+swapped)
|
|
'GRA': [[36, 36, 38, 38, 40, 40, 41, 41,
|
|
43, 43, 45, 45, 47, 47, 48, 48,
|
|
50, 50, 52, 52, 54, 54, 55, 55,
|
|
57, 57, 59, 59, 60, 60, 62, 62,
|
|
64, 64, 66, 66, 67, 67, 69, 69,
|
|
71, 71, 73, 73, 74, 74, 76, 76,
|
|
78, 78, 79, 79]],
|
|
|
|
# from epilogue in the default schedule
|
|
# these are not updated in the updated schedule
|
|
'LRSA': [[50]],
|
|
'LRSB': [[50]],
|
|
'LWSA': [[80]],
|
|
'LWSB': [[80]],
|
|
'LRA1': [[82, 84, 84, 85, 85, 86, 86, 87]], # 8
|
|
'LRB1': [[83, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99]], # 13
|
|
'LCC': [[100, 100]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=12, vlcnt=-1, vscnt=-1, comment="wait for prior iteration LR/LW"),
|
|
SWaitCnt(dscnt=11, vlcnt=-1, vscnt=-1, comment="ensure all previous LRA1/LRB1 done before early MFMA use"),
|
|
SWaitCnt(dscnt=10, vlcnt=-1, vscnt=-1, comment="ensure all previous LRA1/LRB1 done before early MFMA use"),
|
|
|
|
|
|
# A fence: all LRA0 are done before DTL writes from the first GR stream startign at 23 (swapped)
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA0 to complete before first DTL writes"),
|
|
# barrier after LRA0, before first global DTL phase at 23
|
|
SBarrier(comment="barrier after LRA0 , before GR at 23"),
|
|
|
|
# B fence : all LRB0 are done before second DTL stream
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0 to complete before second DTL writes"),
|
|
# barrier after LRB0 before second global stream at 36
|
|
SBarrier(comment="barrier after LRB0, before GR at 36"),
|
|
|
|
# final vmem fence: ensure all GRA/GRB are done before next tile LDS reads LRA1/LRB1
|
|
SWaitCnt(dscnt=-1, vlcnt=34, vscnt=-1, comment="wait for all GRA/GRB before next-tile LDS reads"),
|
|
# final barrier : all waves; make next-tile LDS visible to LRA1/LRB1
|
|
SBarrier(comment="final barrier before LRA1/LRB1 (at 83)"),
|
|
]
|
|
|
|
nglshift = nllshift = 34
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 104
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(192, 128, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_192x128x64_16bit(kernel, useLDSTr, TLDS):
|
|
"""192x128x64 TN schedule (BF16/FP16)."""
|
|
kernel["MfmaInitCVgprs"] = True
|
|
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
|
|
# 192 = 16 * MIWaveTileA(6) * MIWaveGroup0(2)
|
|
# 128 = 16 * MIWaveTileB(4) * MIWaveGroup1(2)
|
|
if isTN(kernel) and (not useLDSTr) and TLDS == 1:
|
|
numMfma = 2 * kernel["MIWaveTileA"] * kernel["MIWaveTileB"]
|
|
# Number of global reads per iter (A:6, B:4) = 10
|
|
nglshift = nllshift = 10
|
|
|
|
# Use syncTable format (idx, wait/barrier, idx, wait/barrier, ...)
|
|
syncTable = [
|
|
# Loop start: must guarantee prior-iteration LR1 completion for the next iteration's early MFMA use.
|
|
-1, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 before starting main loop"),
|
|
5, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 for the remaining main loop"),
|
|
|
|
14, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete to start GRA"),
|
|
14, SBarrier(comment=""),
|
|
|
|
23, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete to start GRB"),
|
|
23, SBarrier(comment=""),
|
|
|
|
# GR -> SWait(vmcnt=0) -> SBarrier -> LR1 ordering (global-read to LDS must be visible before LR1)
|
|
28, SWaitCnt(dscnt=-1, vlcnt=10+1, vscnt=-1, comment="Wait for all GRs to complete before starting LR1"),
|
|
28, SBarrier(comment=""),
|
|
|
|
44, SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="Wait for all GRs to complete before starting LRB1"),
|
|
44, SBarrier(comment=""),
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[9,9,10,10,11,11,12,13,13]],
|
|
|
|
'LRA0' : [[0,1,2,3,4,5]],
|
|
'LRB0' : [[6,9,12,15]],
|
|
|
|
'GRA' : [[14,14, 16,16, 18,18, 20,20, 22,22, 24,24],
|
|
[15,15, 17,17, 19,19, 21,21, 23,23, 25,25]],
|
|
'GRB' : [[26,26, 30,30, 37,37, 43,43]],
|
|
|
|
# Prefetch next iteration
|
|
'LRA1' : [[28,30,32,34,36,38],
|
|
[29,31,33,35,37,39]],
|
|
'LRB1' : [[44,45,46,47]],
|
|
|
|
'LRSA' : [[21]],
|
|
'LRSB' : [[21]],
|
|
'LWSA' : [[46]],
|
|
'LWSB' : [[46]],
|
|
'LCC' : [[47, 47]],
|
|
}
|
|
|
|
syncCode = syncTable[1::2]
|
|
else:
|
|
return False, None
|
|
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(224, 128, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_224x128x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
numMfma = 56
|
|
numCodePaths = 1
|
|
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
nglshift = nllshift = 11 # vmcnt shift for ngl and nll
|
|
optSchedule = {
|
|
'SYNC': [[-1, 6, 14, 14, 27,27, 47, 47]],
|
|
'LRA0': [[0,1, 2,3,4,5,5]],
|
|
'GRIncA': [[0, 0, 1, 1, 2,2 , 3,3, 4]],
|
|
'LRB0': [[9, 11,13, 19]],
|
|
'GRIncB': [[ 6,6,7,7,8,8,9,9,10]],
|
|
'GRA': [[14, 14, 16,16,18,18,20,20,23,23, 26,26, 27, 27]],
|
|
'LRSA': [[26]],
|
|
'LRSB': [[26]],
|
|
'GRB': [[33,34,36,38,38,42,42,46]],
|
|
'LWSA': [[54]],
|
|
'LWSB': [[54]],
|
|
'LRA1': [[30,35,44, 45, 46, 48,51]],
|
|
'LRB1': [[47,52,54,55]],
|
|
'LCC': [[55, 55]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="wait for all of LRA1 and the first instance of LRB1"),
|
|
SWaitCnt(dscnt=8, vlcnt=-1, vscnt=-1, comment="wait for the second instance of LRB1"),
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="wait for all LRA0 to complete so GRA could begin. Makes sure LRB1 is completed so no need for a barrier at 21"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=10, vscnt=-1, comment="wait for all LR. All of previous GRB (4) and current GRA (6), total of 10 can be outstanding"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Outstanding LR are all LRA so no need to wait. All of GR from previous iteration must be done."),
|
|
SBarrier(comment="")
|
|
]
|
|
elif isNN(kernel) and useLDSTr and TLDS==1:
|
|
optSchedule = {
|
|
|
|
'SYNC' : [[-1,3, 16,16, 27, 35,35, 48,48]],
|
|
'GRIncA' : [[1,2,3,4,4,5,5,6,6]],
|
|
'GRIncB' : [[7,7,8,8,9,9,10,11,11]],
|
|
|
|
'LRA0' : [[0,1,1,2,2,3,4,5,6,7,8,9,10,11]],
|
|
'LRB0' : [[17,18,19,20]], ## After LRA0, we can mix LRB0 and GRA
|
|
|
|
## GRA should start after LRA0 is done.
|
|
'GRA' : [[15,16, 19,19, 22,22, 25,25, 28,28, 31,31, 34,34]],
|
|
## GRB should start after LRB0 is done.
|
|
'GRB' : [[42,42, 45,45, 48,48, 51,51]],
|
|
|
|
#After previous GRA is done.
|
|
'LRA1' : [[35,36,37,38,39,40,41,42,43,44,45,46,47,48]],
|
|
#After previous GRB is done.
|
|
'LRB1' : [[49,50, 52,53]],
|
|
|
|
'LRSA' : [[24]], # after LRA0 and before LRA1
|
|
'LRSB' : [[24]], # after LRB0 and before LRB2
|
|
'LWSA' : [[53]], # For A
|
|
'LWSB' : [[54]],
|
|
|
|
'LCC' : [[54, 55]],
|
|
}
|
|
# note: syncCode needs to be
|
|
syncCode = [SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for necessary prior LRA1/LRB1 before starting main loop"),
|
|
SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 for the remaining main loop"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete to start GRA"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Wait for previous GRA to complete to start LRA1"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=9, vscnt=-1, comment="Wait for previous GRB to complete to start LRB1"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 11 # vmcnt shift for ngl and nll
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
# Global read scheduling:
|
|
# Each GR has two instructions (addr update + buffer_load), so we list them explicitly as
|
|
# two adjacent MFMA indices per GR.
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
numCodePaths = 2
|
|
|
|
syncTable = [
|
|
# Loop start:
|
|
# - LRB1 waits previous-iter GRB
|
|
# - LRA1 waits previous-iter GRA
|
|
-1, SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 before starting main loop"),
|
|
|
|
# After early MFMAs (keep prior-iter LR fully fenced)
|
|
3, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 for the remaining main loop"),
|
|
|
|
# GRB must wait for LRB0 (interleave LRA0 + GRB safely)
|
|
15, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete to start GRB"),
|
|
15, SBarrier(comment=""),
|
|
|
|
# GRA must wait for LRA0; LRB1 can be interleaved with GRA after this fence
|
|
27, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0/GRB to complete to start GRA/LRB1"),
|
|
27, SBarrier(comment=""),
|
|
|
|
# Mid-loop global-read safety fence (GR-to-LDS)
|
|
30, SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Mid-loop fence (wait for outstanding GR-to-LDS)"),
|
|
30, SBarrier(comment=""),
|
|
|
|
# Ensure all GR-to-LDS are complete before LRA1 (next-iter A LDS reads)
|
|
43, SWaitCnt(dscnt=-1, vlcnt=11-2, vscnt=-1, comment="Wait for all GR-to-LDS to complete before LRA1"),
|
|
43, SBarrier(comment=""),
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
# Swap A/B increments
|
|
'GRIncB' : [[0, 0, 1, 1, 2, 2, 4, 4, 5]],
|
|
'GRIncA' : [[5, 6, 6, 7, 7, 8, 8, 9, 9]],
|
|
|
|
'LRB0' : [[0, 1, 2, 3, 4, 5, 6, 7],
|
|
[1, 2, 3, 4, 5, 6, 7, 8]],
|
|
'LRA0' : [[9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22],
|
|
[10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]],
|
|
|
|
'GRA' : [[14,15, 17,18, 20,21, 23,24],
|
|
[15,16, 18,19, 21,22, 24,25]],
|
|
'GRB' : [[28,29, 31,32, 34,35, 37,38, 40,41, 43,44, 46,47],
|
|
[29,30, 32,33, 35,36, 38,39, 41,42, 44,45, 45,46]],
|
|
|
|
'LRB1' : [[30, 31, 32, 33, 34, 35, 36, 37],
|
|
[31, 32, 33, 34, 35, 36, 37, 38]],
|
|
'LRA1' : [[43, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55]],
|
|
|
|
'LRSA' : [[24]],
|
|
'LRSB' : [[25]],
|
|
'LWSA' : [[48]],
|
|
'LWSB' : [[49]],
|
|
'LCC' : [[53, 54]],
|
|
}
|
|
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 11
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(numCodePaths, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(224, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_224x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
if isTN(kernel) and TLDS==1:
|
|
optSchedule = {
|
|
'SYNC' : [[ -1, 18, 18, 51, 51, 90, 90]],
|
|
'GRIncA' : [[ 1, 1, 3, 3, 5, 5, 7, 7, 9],
|
|
[ 0, 0, 2, 2, 4, 4, 6, 6, 8]],
|
|
'GRIncB' : [[ 9, 11, 11, 13, 13, 15, 15, 17, 17],
|
|
[ 8, 10, 10, 12, 12, 14, 14, 16, 16]],
|
|
'LRA0' : [[ 0, 2, 4, 6, 8, 10, 12],
|
|
[ 1, 3, 5, 7, 9, 11, 13]],
|
|
|
|
'LRB0' : [[ 14, 16, 19, 21, 23, 25, 27, 29],
|
|
[ 15, 17, 20, 22, 24, 26, 28, 30]],
|
|
'GRA' : [[ 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 46, 47, 48, 49],
|
|
[ 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 47, 48, 49, 50]],
|
|
|
|
'LRA1' : [[ 56, 58, 77, 79, 81, 83, 85],
|
|
[ 57, 59, 78, 80, 82, 84, 86]],
|
|
'GRB' : [[ 52, 53, 54, 55, 56, 57, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86],
|
|
[ 53, 54, 55, 56, 57, 58, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87]],
|
|
|
|
'LRB1' : [[ 91, 93, 95, 97, 99, 101, 103, 105],
|
|
[ 92, 94, 96, 98, 100, 102, 104, 106]],
|
|
'LRSA' : [[ 50], [52]],
|
|
'LRSB' : [[ 50], [52]],
|
|
'LWSA' : [[108]],
|
|
'LWSB' : [[109]],
|
|
'LCC' : [[110, 111]]
|
|
}
|
|
syncCode = [
|
|
SWaitCnt(dscnt= 0, vlcnt=-1, vscnt=-1, comment="Wait for LRBs"),
|
|
SWaitCnt(dscnt= 2, vlcnt=-1, vscnt=-1, comment="Wait for LRAs"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt= 0, vlcnt=15, vscnt=-1, comment="Wait for LRBs and previous set of GRs"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=15, vscnt=-1, comment="Wait for previous set of GRs"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 15
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
optSchedule = {
|
|
'SYNC' : [[-1, 21, 21, 51, 51, 79, 79]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[9,10,11,12,13,14,15,16,17]],
|
|
|
|
'LRA0' : [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]],
|
|
'LRB0' : [[14, 17, 20, 23, 26, 29, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50]],
|
|
|
|
'GRA' : [[22,22, 26,26, 30,30, 34,34, 38,38, 42,42, 46,46]],
|
|
'GRB' : [[52,53, 56,57, 60,61, 64,65, 68,69, 71,72, 74,75, 77,78]],
|
|
|
|
'LRA1' : [[79,80, 81,82, 83,84, 85,86, 87,88, 89,90, 91,92]],
|
|
'LRB1' : [[93,94, 95,96, 97,98, 99,100, 101,102, 103,104, 105,106, 107,108]],
|
|
'LRSA' : [[54]],
|
|
'LRSB' : [[54]],
|
|
'LWSA' : [[91]],
|
|
'LWSB' : [[91]],
|
|
'LCC' : [[111, 111]]
|
|
}
|
|
syncCode = [
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0 for iteration == 0"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=15, vscnt=-1, comment="wait for previous set of global reads"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 15
|
|
elif isNN(kernel) and useLDSTr and TLDS == 1:
|
|
optSchedule = {
|
|
'SYNC': [[ -1,6,
|
|
23,23,
|
|
55,55,
|
|
94,
|
|
]],
|
|
'GRIncA': [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB': [[9,10,11,12,13,14,15,16,17]],
|
|
'LRA0': [[0,1,2,3,4,5,6,7,8,9,10,11,12,13]],
|
|
'GRA': [[23,23,27,27,32,32,37,37,42,42,46,46,51,51],
|
|
[25,25,29,29,34,34,39,39,44,44,48,48,53,53]],
|
|
'LRB0': [[23,27,32,37,42,46,51,54],
|
|
[25,29,34,39,44,48,53,54]],
|
|
|
|
'LWSA': [[87],[86]],
|
|
'LWSB': [[90],[89]],
|
|
'LRA1': [[57,57,61,61,65,65,75,75,81,81,87,90,96,96],
|
|
[55,55,59,59,63,63,69,69,79,79,86,89,94,94]],
|
|
'GRB': [[56,56, 60,60, 70,70, 80,80, 90,92, 100,100, 106,106, 109,109],
|
|
[58,58, 62,62, 72,72, 82,82, 91,93, 101,101, 107,107, 110,110]],
|
|
'LRB1': [[91,92,97,100,102,103,106,109]],
|
|
'LRSA': [[54]],
|
|
'LRSB': [[54]],
|
|
'LCC': [[109,110]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="wait for 8-5 local reads // oldleft=8, completed=3"),
|
|
SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="wait for 11-6 local reads // oldleft=5, new=6, completed=5"),
|
|
SWaitCnt(dscnt=0, vlcnt=8, vscnt=-1, comment="wait for prior global reads and local reads // oldleft=6, new=8, completed=14"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=7, vscnt=-1, comment="wait for prior global reads and local reads // oldleft=0, new=8, completed=8"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="wait for 14-5 local reads // oldleft=0, new=14, completed=9"),
|
|
]
|
|
|
|
nglshift = nllshift = 15
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 112
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(192, 320, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_192x320x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 120
|
|
nllZeroDscnt = False
|
|
syncs = SyncSchedule()
|
|
gr_inc_step = 0
|
|
|
|
if isNN(kernel) and useLDSTr and TLDS==1:
|
|
syncs.add(-1, dscnt=9, comment="wait for all LRA1 and one item from LRB1 before starting the sub-iteration")
|
|
lra0 = [0,1,2,3,4,6,8,10,12,14,16,18]
|
|
|
|
syncs.add(5, dscnt=5, comment="wait for the rest of LRB1 to complete")
|
|
grinca = [7,7,7,9,9,9,11,11,11]
|
|
grincb = [13,13,13,15,15,15,17,17,17]
|
|
lrb0 = [20,22,24,25,27,29,31,33,35,37]
|
|
|
|
syncs.add(26, dscnt=4, barrier=True, comment="wait for all LRA0 to complete before GRA start")
|
|
gra = [28,30,32,36,39,42] # one index for two instructions
|
|
syncs.add(44, dscnt=0, barrier=True, comment="wait for all LRB0 to complete before GRB start")
|
|
grb = [53,58,63,67,72,77,82,86,91,96] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
lrsa = [58]
|
|
lrsb = [58]
|
|
|
|
syncs.add(71, vlcnt=10, barrier=True, comment="wait for previous set of global reads")
|
|
lra1 = [72,74,76,78,80,82,84,87,90,92,98,100]
|
|
lrb1 = [99,106,107,108,109,110,111,112,113,114]
|
|
lwsa = [95]
|
|
lwsb = [95]
|
|
|
|
elif isTN(kernel) and not useLDSTr and TLDS==1:
|
|
syncs.add(-1, dscnt=9, comment="wait for all LRA1 and one item from LRB1 before starting the sub-iteration")
|
|
lra0 = [0,1,2,3,4,6]
|
|
|
|
syncs.add(5, dscnt=5, comment="wait for the rest of LRB1 to complete")
|
|
grinca = [0,1,2,3,4,5,6,7,8]
|
|
grincb = [9,10,12,13,14,15,16,17,18]
|
|
lrb0 = [8,10,12,14,16,18,20,22,24,27]
|
|
|
|
syncs.add(26, dscnt=9, barrier=True, comment="wait for all LRA0 to complete before GRA start")
|
|
gra = [28,30,32,36,39,42] # one index for two instructions
|
|
syncs.add(44, dscnt=0, barrier=True, comment="wait for all LRB0 to complete before GRB start")
|
|
grb = [53,58,63,67,72,77,82,86,91,96] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
lrsa = [58]
|
|
lrsb = [58]
|
|
syncs.add(71, vlcnt=10, barrier=True, comment="wait for previous set of global reads")
|
|
|
|
lra1 = [72,76,80,84,90,98]
|
|
lrb1 = [99,106,107,108,109,110,111,112,113,114]
|
|
lwsa = [95]
|
|
lwsb = [95]
|
|
|
|
gr_inc_step = 1
|
|
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
lra0 = [0,1,3,5,7, 9,10,12,14,16, 18,19] # 12 loads
|
|
lrb0 = [21,23,25,27,28, 30,32,34,36,37, 39,41,43,45,46, 48,50,52,54,55] # 20 loads
|
|
# need two LRB1 items because a single LRB read gets only half of the data needed for MFMA
|
|
# note, max dscnt value is 15, so in this case 19 will be maxed at 15, thus we will wait more than needed :(
|
|
syncs.add(-1, dscnt=len(lrb0)-2, comment="wait for all LRA1 and two items from LRB1 before starting the sub-iteration")
|
|
|
|
i = 5 # next LRB1 is needed at index 6, so insert wait at 5
|
|
syncs.add(i, dscnt=count_items(lra0,ev=i), comment="wait for the rest of LRB1 to complete")
|
|
grinca = [0,1,2,3,4,5,6,7,8]
|
|
grincb = [9,10,12,13,14,15,16,17,18]
|
|
|
|
i = 26
|
|
syncs.add(i, dscnt=count_items(lrb0,ev=i), barrier=True, comment="wait for all LRA0 to complete before GRA start")
|
|
gra = [28,30,32,36,39,42] # one index for two instructions
|
|
|
|
lrsa = [57]
|
|
lrsb = [58]
|
|
|
|
# This wait serves dual purpose
|
|
syncs.add(59, dscnt=len(lrb0)-2, vlcnt=len(gra), barrier=True,
|
|
comment="wait for the first LRB0 to complete to start 2nd batch of MFMAs and also make GRs from the previous iteration is done before LRA1 starts")
|
|
lra1 = [61,62,63,64,65, 66,67,69,71,73, 75,76] # 12 loads
|
|
|
|
i = 65 # next LRB0 is needed at index 66, so insert wait at 65
|
|
syncs.add(i, dscnt=count_items(lra1,ev=i), barrier=True,
|
|
comment="wait for the rest of LRB0 to complete across all waves before GRB start")
|
|
grb = [66,70,75,79,84, 88,93,97,102,106] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
lwsa = [68]
|
|
lwsb = [77]
|
|
lrb1 = [78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110,112,114,116] # 20 loads
|
|
|
|
gr_inc_step = 1
|
|
nllZeroDscnt = True
|
|
else:
|
|
return False, None
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'LRA0': [lra0],
|
|
'GRIncA': [grinca],
|
|
'LRB0': [lrb0],
|
|
'GRIncB': [grincb],
|
|
# Note: each GRA/GRB item corresponds to two instructions (addr increment and read). So duplicate each item twice.
|
|
'GRA': [duplicate_list_items(gra, 2, gr_inc_step)],
|
|
'GRB': [duplicate_list_items(grb, 2, gr_inc_step)],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'LRA1': [lra1],
|
|
'LRB1': [lrb1],
|
|
'LCC': [[numMfma-2, numMfma-1]],
|
|
}
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
kernel["SwapGlobalReadOrder"] = 0
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift, nllZeroDscnt)
|
|
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 224, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x224x64_16bit(kernel, useLDSTr, TLDS):
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
if isTN(kernel) and TLDS==1:
|
|
optSchedule = {
|
|
'SYNC' : [[ -1, 18, 18, 51, 51, 90, 90]],
|
|
'GRIncA' : [[ 1, 1, 3, 3, 5, 5, 7, 7, 9],
|
|
[ 0, 0, 2, 2, 4, 4, 6, 6, 8]],
|
|
'GRIncB' : [[ 9, 11, 11, 13, 13, 15, 15, 17, 17],
|
|
[ 8, 10, 10, 12, 12, 14, 14, 16, 16]],
|
|
'LRA0' : [[ 0, 2, 4, 6, 8, 10, 12, 14],
|
|
[ 1, 3, 5, 7, 9, 11, 13, 15]],
|
|
# schduling GRIncA/B and LRA0 as follow,
|
|
# SIMD 0 | ... | MFMA | GRInc | GRInc | MFMA | LDS Load | MFMA | GRInc | GRInc | MFMA | ...
|
|
# SIMD 1 | ... | MFMA | LDS Load | MFMA | GRInc | GRInc | MFMA | LDS Load | MFMA | ...
|
|
|
|
'LRB0' : [[ 16, 19, 21, 23, 25, 27, 29],
|
|
[ 17, 20, 22, 24, 26, 28, 30]],
|
|
'GRA' : [[ 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 46, 47, 48, 49, 52, 53],
|
|
[ 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 47, 48, 49, 50, 53, 54]],
|
|
|
|
'LRA1' : [[ 52, 56, 58, 77, 79, 81, 83, 85],
|
|
[ 53, 57, 59, 78, 80, 82, 84, 86]],
|
|
'GRB' : [[ 54, 55, 56, 57, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86],
|
|
[ 55, 56, 57, 58, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87]],
|
|
|
|
'LRB1' : [[ 91, 93, 95, 97, 99, 101, 103],
|
|
[ 92, 94, 96, 98, 100, 102, 104]],
|
|
'LRSA' : [[ 50], [52]],
|
|
'LRSB' : [[ 50], [52]],
|
|
'LWSA' : [[108]],
|
|
'LWSB' : [[109]],
|
|
'LCC' : [[110, 111]]
|
|
}
|
|
syncCode = [
|
|
SWaitCnt(dscnt= 0, vlcnt=-1, vscnt=-1, comment="Wait for LRBs"),
|
|
SWaitCnt(dscnt= 1, vlcnt=-1, vscnt=-1, comment="Wait for LRAs"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt= 0, vlcnt=14, vscnt=-1, comment="Wait for LRBs and previous set of GRAs"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=15, vscnt=-1, comment="Wait for previous set of GRBs"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 15
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
optSchedule = {
|
|
'SYNC': [[-1,6,
|
|
21,21,55,55,60,60]],
|
|
'GRIncA': [[0,1,2,3,4,5,6,7,8]],
|
|
'LRA0': [[0,0,2,2,4,4,6,6,8,8,10,10,12,12,14,14],
|
|
[1,1,3,3,5,5,7,7,9,9,11,11,13,13,15,15]],
|
|
|
|
'GRIncB': [[9,10,11,12,13,14,15,16,17]],
|
|
|
|
'GRA': [[21,22, 25,26, 30,31, 35,36, 40,41, 45,46, 50,51, 53,54]],
|
|
'LRB0': [[21,22, 25,26, 30,31, 35,36, 40,41, 45,46, 50,51]],
|
|
|
|
'GRB': [[61,61, 63,63, 65,65, 79,79, 85,85, 95,95, 100,100],
|
|
[62,62, 64,64, 66,66, 80,80, 91,91, 96,96, 101,101]],
|
|
'LWSA': [[93],[99]],
|
|
'LWSB': [[91],[87]],
|
|
'LRA1': [[61,61, 63,63, 65,65, 79,79, 85,85, 93,93, 100,100, 104,104],
|
|
[62,62, 64,64, 66,66, 80,80, 91,91, 96,96, 101,101, 106,106]],
|
|
'LRB1': [[91,91,95,95,98,98,100,100,110,110,110,111,111,111],
|
|
[87,87,94,94,99,99,101,101,103,103,105,105,107,107]],
|
|
|
|
'LRSA': [[54]],
|
|
'LRSB': [[54]],
|
|
'LCC': [[110,111]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="wait for prior local read"),
|
|
SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="wait for prior local read"),
|
|
SWaitCnt(dscnt=0, vlcnt=7, vscnt=-1, comment="wait for previous set of global reads and Local Reads"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=8, vscnt=-1, comment="wait for previous set of global reads"),
|
|
SBarrier(comment="")
|
|
]
|
|
nglshift = nllshift = 15
|
|
|
|
elif isNN(kernel) and useLDSTr and TLDS == 1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
#index and code pair
|
|
syncTable = [-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA1"),
|
|
17, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="wait for LRB0"),
|
|
17, SBarrier(comment=""),
|
|
# 46, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
51, SWaitCnt(dscnt=0, vlcnt=15, vscnt=-1, comment="wait for previous set of global reads"),
|
|
51, SBarrier(comment=""),
|
|
65, SWaitCnt(dscnt=-1, vlcnt=15-4, vscnt=-1, comment="wait for previous set of global reads"),
|
|
65, SBarrier(comment=""),
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[37,38,39,40,41,42,43,44,45]],
|
|
'GRIncB': [[0,1,2,3,4,5,6,7,8]],
|
|
|
|
'LRB0': [[-1, 0, 1, 2, 3, 4, 5],
|
|
[ 0, 1, 2, 3, 4, 5, 6]],
|
|
'LRA0': [[8, 10, 12, 14, 16, 18, 20, 22, 24, 25, 27, 29, 31, 33, 35, 37],
|
|
[9, 11, 13, 15, 17, 19, 21, 23, 25, 26, 28, 30, 32, 34, 36, 38]],
|
|
'GRA': [[17,17, 19,19, 26,26, 28,28, 30,30, 32,32, 34,34],
|
|
[18,18, 20,20, 27,27, 29,29, 31,31, 33,33, 35,35]],
|
|
|
|
'GRB': [[52,52, 54,54, 56,56, 58,58, 66,66, 68,68, 70,70, 72,72],
|
|
[53,53, 55,55, 57,57, 59,59, 67,67, 69,69, 71,71, 73,73]],
|
|
'LRB1': [[51, 53, 55, 57, 59, 61, 63],
|
|
[52, 54, 56, 58, 60, 62, 64]],
|
|
'LRA1': [[65, 67, 69, 71, 73, 75, 77, 79, 81, 83, 85, 87, 89, 91, 93, 95],
|
|
[66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96]],
|
|
|
|
'LRSB': [[14]],
|
|
'LRSA': [[45]],
|
|
'LWSB': [[97]],
|
|
'LWSA': [[97]],
|
|
'LCC' : [[110, 110]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 15 # vmcnt shift for ngl and nll
|
|
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 112
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(320, 192, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_320x192x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
|
|
if isNN(kernel) and useLDSTr and TLDS == 1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA1 "),
|
|
19, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="before DirectToLds load, ensure LRB0 have finished"),
|
|
19, SBarrier(comment=""),
|
|
52, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA0 finish"),
|
|
52, SBarrier(comment=""),
|
|
53, SWaitCnt(dscnt=-1, vlcnt=16+1, vscnt=-1, comment="wait for previous GRB finish"),
|
|
53, SBarrier(comment=""),
|
|
71, SWaitCnt(dscnt=-1, vlcnt=16-8, vscnt=-1, comment="wait for previous GRA finish"),
|
|
71, SBarrier(comment=""),
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncTable[::2]],
|
|
'GRIncA': [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB': [[9,10,11,12,13,14,16,16,16]],
|
|
|
|
'LRB0': [[0, 1, 2, 3, 4, 5],
|
|
[1, 2, 3, 4, 5, 6]],
|
|
'LRA0': [[6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44],
|
|
[7, 8, 11, 13, 15, 17, 18, 21, 23, 25, 27, 29, 31, 33, 35, 37, 39, 41, 43, 45]],
|
|
'GRA': [[19,19, 21,21, 28,28, 31,31, 33,33, 41,41],
|
|
[20,20, 22,22, 29,29, 32,32, 34,34, 42,42]],
|
|
'GRB': [[52,52, 64,64, 74,74, 78,78, 83,83, 88,88, 93,93, 98,98, 105,105, 109,109],
|
|
[52,52, 65,65, 75,75, 79,79, 84,84, 89,89, 94,94, 99,99, 106,106, 110,110]],
|
|
'LRB1': [[53, 55, 57, 59, 63, 67],
|
|
[54, 56, 58, 60, 64, 68]],
|
|
'LRA1': [[71,73, 75,77, 79,81, 83,85, 87,89, 91,93, 95,97, 99,101, 103,105, 107,109],
|
|
[72,74, 76,78, 80,82, 84,86, 88,90, 92,94, 96,98, 100,102, 104,106, 108,110]],
|
|
'LRSB': [[16]],
|
|
'LRSA': [[48]],
|
|
'LWSB': [[48]],
|
|
'LWSA': [[112]],
|
|
'LCC': [[119, 119]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 16
|
|
elif isTN(kernel) and TLDS==1:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
# Note: A/B Global read orders are swapped
|
|
# i.e. GRA contains GR for B
|
|
optSchedule = {
|
|
'SYNC' : [[-1, 4, 16, 16, 50, 50, 54, 55, 84, 85]],
|
|
'GRIncA': [[0, 1, 2, 3, 4, 5, 6, 7, 8]],
|
|
'GRIncB': [[9, 10, 11, 12, 13, 14, 16, 16, 16]],
|
|
'LRB0' : [[0, 2, 4, 6, 8, 10],
|
|
[1, 3, 5, 7, 9, 11]],
|
|
'LRA0' : [[12, 14, 24, 26, 28, 30, 32, 34, 36, 38],
|
|
[13, 15, 25, 27, 29, 31, 33, 35, 37, 39]],
|
|
'GRA' : [[16, 16, 18, 18, 20, 20, 22, 22, 46, 46, 48, 48],
|
|
[17, 17, 19, 19, 21, 21, 23, 23, 47, 47, 49, 49]],
|
|
'GRB' : [[50, 50, 52, 52, 74, 75, 78, 78, 80, 80, 82, 82, 106, 106, 108, 108, 110, 110, 112, 112],
|
|
[51, 51, 53, 53, 76, 77, 79, 79, 81, 81, 83, 83, 107, 107, 109, 109, 111, 111, 113, 113]],
|
|
'LRB1' : [[56, 58, 60, 62, 64, 66],
|
|
[57, 59, 61, 63, 65, 67]],
|
|
'LRA1' : [[86, 88, 90, 92, 94, 96, 98, 100, 102, 104],
|
|
[87, 89, 91, 93, 95, 97, 99, 101, 103, 105]],
|
|
'LRSA' : [[44]],
|
|
'LRSB' : [[45]],
|
|
'LWSA' : [[115]],
|
|
'LWSB' : [[117]],
|
|
'LCC' : [[119, 119]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for prior local read. Relax a bit to dscnt=4 to reduce latency") ,
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for all prior local read requested as the input for MFMA") ,
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for prior LRB0 before GRA (GR for MatrixB). Skip LRA0*2") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA0 before GRB (GR for MatrixA)") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=-1, vlcnt=18, vscnt=-1, comment="Wait for GRA (GR for MatrixB) from previous iteration before LRB1. Skip GRB*10 from last iter and GRA*6 + GRB*2 from this iter.") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=-1, vlcnt=12, vscnt=-1, comment="Wait for GRB (GR for MatrixA) from previous iteration before LRA1. Skip GRA*6 + GRB*6 from this iter.") ,
|
|
SBarrier(comment="") ,
|
|
]
|
|
nglshift = nllshift = 16
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
kernel["SwapGlobalReadOrder"] = 1
|
|
# Note: A/B Global read orders are swapped
|
|
# i.e. GRA contains GR for B
|
|
optSchedule = {
|
|
'SYNC' : [[-1, 7, 17, 17, 49, 49, 59, 59]],
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
'LRB0' : [[0, 0, 2, 2, 4, 4, 6, 6, 8, 8, 10, 10],
|
|
[1, 1, 3, 3, 5, 5, 7, 7, 9, 9, 11, 11]],
|
|
'LRA0' : [[11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 31, 33, 33, 35, 37, 39, 41, 43, 45],
|
|
[12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 36, 38, 38, 40, 42, 44, 46]],
|
|
'GRA' : [[18, 18, 20, 20, 22, 22, 24, 24, 26, 26, 28, 28],
|
|
[19, 19, 21, 21, 23, 23, 25, 25, 27, 27, 29, 29]],
|
|
'GRB' : [[49, 49, 51, 51, 53, 53, 55, 55, 57, 57, 89, 89, 91, 91, 93, 93, 95, 95, 97, 97],
|
|
[50, 50, 52, 52, 54, 54, 56, 56, 58, 58, 90, 90, 92, 92, 94, 94, 96, 96, 98, 98]],
|
|
'LRB1' : [[60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82],
|
|
[61, 63, 65, 67, 69, 71, 73, 75, 77, 79, 81, 83]],
|
|
'LRA1' : [[85, 87, 89, 91, 93, 95, 97, 99, 101, 103, 103, 105, 105, 107, 107, 109, 111, 113, 115, 117],
|
|
[86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 106, 108, 108, 110, 110, 112, 114, 116, 118],],
|
|
'LRSA' : [[58]],
|
|
'LRSB' : [[58]],
|
|
'LWSA' : [[99]],
|
|
'LWSB' : [[99]],
|
|
'LCC' : [[119, 119]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for prior local read. Relax a bit to dscnt=4 to reduce latency") ,
|
|
SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for remaining LRA1.") ,
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for all LRB0 prior to LRA0*3") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for prior local read") ,
|
|
SBarrier(comment="") ,
|
|
SWaitCnt(dscnt=-1, vlcnt=11, vscnt=-1, comment="Wait for prior GRA*6 + GRB*5 = 11 global reads") ,
|
|
SBarrier(comment="") ,
|
|
]
|
|
nglshift = nllshift = 16
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 120
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(352, 192, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_352x192x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 132
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
|
|
if isTN(kernel) and TLDS==1:
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=5 newLW=0 newLR=5 for iteration == 0"),
|
|
14, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="wait for 5 LRA0s to complete"),
|
|
25, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="wait for next 4 LRA0s to complete"),
|
|
|
|
34, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRA0s to complete"),
|
|
34, SBarrier(comment="Barrier before GRA"),
|
|
|
|
46, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="wait for 3 LRB0s to complete"),
|
|
|
|
87, SWaitCnt(dscnt=0, vlcnt=17, vscnt=-1, comment="wait for previous GRA to complete"),
|
|
87, SBarrier(comment="Barrier before LRA1"),
|
|
118, SWaitCnt(dscnt=-1, vlcnt=12, vscnt=-1, comment="wait for previous GRB to complete"),
|
|
118, SBarrier(comment="Barrier before LRB1"),
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncTable[::2]], # 6
|
|
|
|
'GRIncA': [[0, 0, 1, 1, 2, 2, 3, 3, 4]], # 9
|
|
'GRIncB': [[4, 5, 5, 6, 6, 7, 7, 8, 8]], # 9
|
|
|
|
'LRA0': [[0, 2, 4, 6, 8 , 10, 14, 18, 20, 24, 26],
|
|
[1, 3, 5, 7, 9, 11, 15, 17, 21, 25, 27]], # 11
|
|
|
|
'LRB0': [[31, 33, 36, 39, 41, 44],
|
|
[32, 34, 37, 40, 42, 45]], # 6
|
|
|
|
'GRA': [[35, 35, 40, 40, 46, 46, 51, 51, 56, 56, 61, 61, 67, 67, 72, 72, 77, 77, 82, 82, 86, 86]], # 22
|
|
'GRB': [[88, 88, 93, 93, 98, 98, 103, 103, 109, 109, 114, 114]], # 12
|
|
|
|
'LRSA': [[64]], # 1
|
|
'LRSB': [[64]], # 1
|
|
|
|
'LWSA': [[109]], # 1
|
|
'LWSB': [[109]], # 1
|
|
'LRA1': [[87, 90, 92, 94, 96, 98, 100, 102, 104, 106, 114]], # 11
|
|
'LRB1': [[118, 120, 122, 124, 127, 130]], # 6
|
|
|
|
'LCC': [[numMfma-3, numMfma-3]], # 2
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = len(optSchedule["GRA"][0])/2 + len(optSchedule["GRB"][0])/2
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(240, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[2, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[1, 4]
|
|
)
|
|
def _get_schedule_240x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
if isTN(kernel) and TLDS==1:
|
|
kernel["SwapGlobalReadOrder"] = 0
|
|
optSchedule = {
|
|
'SYNC': [[-1,
|
|
14,
|
|
26,26,
|
|
59,
|
|
69,69,
|
|
98,98]],
|
|
'LRA0': [[0,2,3,4,5,6,8,10,12,14, 16,18,20,22,24]],
|
|
'GRIncA': [[0,0,0,1,1,1,2,2,2]],
|
|
'GRIncB': [[3,3,3,4,4,4,5,5,5]],
|
|
'LRB0': [[28,30,36,38]],
|
|
'GRA': [[26,26,27,27,29,29,31,31,33,33,35,35,37,37,39,39,41,41,42,42,44,44,46,46,48,48,50,50,52,52,54,54,56,56,58,58,59,59,61,61,63,63,65,65,67,67,69,69,71,71,73,73,75,75,76,76,78,78,80,80]],
|
|
'GRB': [[82,82,84,84,86,86,88,88,90,90,92,92,94,94,96,96],
|
|
[81,81,83,83,85,85,87,87,91,91,93,93,95,95,97,97]],
|
|
'LRSA': [[58]],
|
|
'LRSB': [[58]],
|
|
'LWSA': [[98]],
|
|
'LWSB': [[98]],
|
|
'LRA1': [[70,72,74,76,77,79,80,82,84,86,88,90,92,94,96]],
|
|
'LRB1': [[99,114,115,116]],
|
|
'LCC': [[119, 119]]
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=3 newLW=0 newLR=3 for iteration == 0"),
|
|
SWaitCnt(dscnt=9, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=3 newLW=0 newLR=3 for iteration == 0"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRA0"),
|
|
SBarrier(comment=""),
|
|
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LRB0"),
|
|
|
|
SWaitCnt(dscnt=-1, vlcnt=23+8, vscnt=-1, comment="wait for previous set of GRA"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=38, vscnt=-1, comment="wait for previous set of GRB"),
|
|
SBarrier(comment="")
|
|
]
|
|
numMfma = 120
|
|
nglshift = nllshift = len(optSchedule["GRA"][0])/2 + len(optSchedule["GRB"][0])/2
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNT(kernel) and TLDS==0:
|
|
optSchedule = {
|
|
'SYNC': [[-1,24,24,59,59]],
|
|
'LRA0': [[0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13,14,14,15]],
|
|
'LRB0': [[24,24,26,26,28,28,30,30]],
|
|
'GRA': [[24,24,25,25,27,27,29,29,31,31,33,33,35,35,37,37,39,39,41,41,43,43,45,45,47,47,49,49,50,50,52,52,54,54,56,56,58,58,60,60,61,61,62,62,63,63,64,64, 65,65,66,66,67,67,68,68,69,69,70,70]],
|
|
'GRB': [[75,75,76,76,77,77,78,78, 89,89,91,91,93,93,95,95]],
|
|
'LRA1': [[60,60,61,61,62,62,63,63,64,64, 65,65,66,66,67,67,68,68,69,69,70,70, 75,75,76,76,77,77,78,78]],
|
|
'LRB1': [[89,89,91,91,93,93,95,95]],
|
|
'GRIncA': [[0,0,0,1,1,1,2,2,2]],
|
|
'GRIncB': [[3,3,3,4,4,4,5,5,5]],
|
|
'LRSA': [[58]],
|
|
'LRSB': [[58]],
|
|
'LWSA': [[95]],
|
|
'LWSB': [[95]],
|
|
'LCC': [[119,119]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0 for iteration == 0"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=19, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0, wait for previous set of global reads"),
|
|
SBarrier(comment="")
|
|
]
|
|
numMfma = 120
|
|
nglshift = nllshift = len(optSchedule["GRA"][0])/2 + len(optSchedule["GRB"][0])/2
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
elif isNN(kernel) and useLDSTr and TLDS==1:
|
|
optSchedule = {
|
|
'SYNC': [[-1,
|
|
26,26,
|
|
59,59
|
|
]],
|
|
'LRA0': [[0,2,2,3,3,4,4,5,5,6,6,7,7,9,9,11,11,13,13,15,15,17,17,19,19,21,21,23,23,24],
|
|
[0,2,2,3,3,4,4,5,5,6,6,7,7,8,8,10,10,12,12,14,14,16,16,18,18,20,20,22,22,25]],
|
|
'LRB0': [[26,27,28,29]],
|
|
'GRA': [[26,26,27,27,29,29,31,31,33,33,35,35,37,37,39,39,41,41,42,42,44,44,46,46,48,48,50,50,52,52,54,54,56,56,58,58,59,59,61,61,63,63,65,65,67,67,69,69,71,71,73,73,75,75,76,76,78,78,80,80]],
|
|
'GRB': [[82,82,84,84,86,86,88,88,90,90,92,92,93,93,96,96],
|
|
[83,83,85,85,87,87,89,89,91,91,94,94,99,99,103,103]],
|
|
'LRA1': [[59,59,61,61,63,63,65,65,67,67,69,69,71,71,73,73,75,75,76,76,78,78,80,80,82,82,84,84, 86,86]],
|
|
'LRB1': [[88,90,92,94]],
|
|
'LRSA': [[58]],
|
|
'LRSB': [[58]],
|
|
'LWSA': [[95]],
|
|
'LWSB': [[95]],
|
|
'GRIncA': [[1,1,1,17,17,17,18,18,18]],
|
|
'GRIncB': [[19,19,19,20,20,20,21,21,21]],
|
|
'LCC': [[119,119]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=3 newLW=0 newLR=3 for iteration == 0"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=18, vscnt=-1, comment="wait for prior local read local write old=0, new=0 newLW=0 newLR=0"),
|
|
SBarrier(comment=""),
|
|
]
|
|
numMfma = 120
|
|
nglshift = nllshift = len(optSchedule["GRA"][0])/2 + len(optSchedule["GRB"][0])/2
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(208, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[2, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[1, 4]
|
|
)
|
|
def _get_schedule_208x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 104
|
|
syncs = SyncSchedule()
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
syncs.add(-1, dscnt=3, comment="wait for all LRA1 and one item from LRB1 before starting the sub-iteration")
|
|
lra0 = [0,1,2,3,5,7,9,11,13,15,17,19,20]
|
|
|
|
syncs.add(12, dscnt=8, comment="wait for the rest of LRB1 to complete")
|
|
grinca = [4,4,4,10,10,10,14,14,14]
|
|
lrb0 = [22,24,26,28]
|
|
grincb = [16,16,16,18,18,18,21,21,21]
|
|
|
|
syncs.add(29, dscnt=4, barrier=True, comment="wait for all LRA0 to complete before GRA start")
|
|
# one index for two instructions
|
|
gra = [30,31,32,33,34,35,36,37,38,39,41,43,45,46,48,50,52,53,55,57,59,60,62,64,66,67]
|
|
|
|
syncs.add(40, dscnt=0, barrier=True, comment="wait for all LRB0 to complete before GRB start")
|
|
# one index for two instructions
|
|
grb = [69,71,73,77,81,85,89,93]
|
|
num_gr = len(gra) + len(grb)
|
|
lrsa = [50]
|
|
lrsb = [50]
|
|
lwsa = [70]
|
|
lwsb = [70]
|
|
|
|
syncs.add(72, vlcnt=num_gr-6, barrier=True, comment="wait for previous set of global reads")
|
|
lra1 = [73,74,75,76,78,82,84,86,88,90,92,94,96]
|
|
lrb1 = [80,98,99,100]
|
|
|
|
elif isNN(kernel) and useLDSTr and TLDS==1:
|
|
syncs.add(-1, dscnt=3, comment="wait for all LRA1 and one item from LRB1 before starting the sub-iteration")
|
|
grinca = [8,9,10,11,13,14,15,16,17]
|
|
grincb = [19,20,21,22,23,24,25,26,27]
|
|
|
|
syncs.add(12, dscnt=12, comment="wait for the rest of LRB1 to complete")
|
|
lra0 = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25] # 26 loads
|
|
|
|
syncs.add(30, dscnt=2, barrier=True, comment="wait for all LRA0 to complete before GRA start")
|
|
lrb0 = [26,28,30,33]
|
|
|
|
syncs.add(38, dscnt=0, barrier=True, comment="wait for all LRB0 to complete before GRB start")
|
|
# one index for two instructions
|
|
gra = [31,32,34,36,37,39,41,43,45,46,48,50,52,53,55,57,59,60,62,64,66,67,68,69,70,71] # 26 loads
|
|
grb = [73,74,81,83,87,91,92,93]
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
syncs.add(72, vlcnt=num_gr-8, barrier=True, comment="wait for previous set of global reads")
|
|
lra1 = [73,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99] # 26 loads
|
|
lrb1 = [74,100,101,102]
|
|
|
|
lrsa = [49]
|
|
lrsb = [51]
|
|
lwsa = [84]
|
|
lwsb = [85]
|
|
|
|
elif isNT(kernel) and useLDSTr and TLDS==0:
|
|
syncs.add(-1, dscnt=3, comment="wait for all LRA1 and one item from LRB1 before starting the sub-iteration")
|
|
grinca = [8,9,10,11,13,14,15,16,17]
|
|
grincb = [19,20,21,22,23,24,25,26,27]
|
|
|
|
syncs.add(12, dscnt=12, comment="wait for the rest of LRB1 to complete")
|
|
lra0 = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25] # 26 loads
|
|
lrb0 = [27,29,31,33,35,38,40,42] # 8 loads
|
|
syncs.add(30, dscnt=2, barrier=True, comment="wait for all LRA0 to complete before GRA start")
|
|
|
|
syncs.add(49, dscnt=0, barrier=True, comment="wait for all LRB0 to complete before GRB start")
|
|
# one index for two instructions
|
|
gra = [31,32,34,36,37,39,41,43,45,46,48,50,52,53,55,57,59,60,62,64,66,67,68,69,70,71] # 26 loads
|
|
grb = [73,74,81,83,87,91,92,93] # 8 loads
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
# 8 GRBs from previous iteration + 16 GRAs from current iteration (indices 31-57) can be still pending
|
|
syncs.add(58, vlcnt=(8+16), barrier=True, comment="wait for previous set of GRA")
|
|
lra1 = [59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84] # 26 loads
|
|
|
|
# 20 GRAs (indices 31-64) from current iteration can be still pending
|
|
syncs.add(65,vlcnt=20, barrier=True, comment="wait for previous set of GRB")
|
|
lrb1 = [66,85,86,88,90,92,94,96] # 8 loads
|
|
|
|
lrsa = [49]
|
|
lrsb = [51]
|
|
lwsa = [84]
|
|
lwsb = [85]
|
|
else:
|
|
return False, None
|
|
|
|
def duplicate_list_items(input_list, repeat_count):
|
|
"""Example: duplicate_list_items([1, 2, 3], 3) => [1,1,1, 2,2,2, 3,3,3]"""
|
|
return [item for item in input_list for _ in range(repeat_count)]
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'LRA0': [lra0],
|
|
'GRIncA': [grinca],
|
|
'LRB0': [lrb0],
|
|
'GRIncB': [grincb],
|
|
# Note: each GRA/GRB item corresponds to two instructions (addr increment and read). So duplicate each item twice.
|
|
'GRA': [duplicate_list_items(gra, 2)],
|
|
'GRB': [duplicate_list_items(grb, 2)],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'LRA1': [lra1],
|
|
'LRB1': [lrb1],
|
|
'LCC': [[numMfma-2, numMfma-1]],
|
|
}
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
kernel["SwapGlobalReadOrder"] = 0
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 224, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x224x64_16bit(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
optSchedule = {
|
|
|
|
'SYNC' : [[-1,3, 10,10, 26,27, 45,45]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[22,22,24,24,25,25,26,27,27]],
|
|
|
|
'LRA0' : [[0,2,3,4]], ## -2 is place holder
|
|
|
|
'LRB0' : [[8,11,13,15,17,19,21], ## After LRA0, we can mix LRB0 and GRA
|
|
[9,12,14,16,18,20,22]],
|
|
## GRA should start after LRA0 is done.
|
|
'GRA' : [[10,11, 14,14, 17,17, 20,20],
|
|
[11,12, 15,15, 18,18, 21,21]],
|
|
|
|
## GRB should start after LRB0 is done
|
|
'GRB' : [[28,28, 31,31, 34,34, 37,37, 40,40, 43,43, 46,46], # m0 inc is part of GRA/GRB
|
|
[29,29, 32,32, 35,35, 38,38, 41,41, 44,44, 47,47]],
|
|
'LRA1' : [[29, 32, 35, 38],
|
|
[30, 33, 36, 39]],
|
|
|
|
#After GRB is done.
|
|
'LRB1' : [[45,46,47,48,49,50,51]],
|
|
|
|
'LRSA' : [[23]], # after LRA0 and before LRA1
|
|
'LRSB' : [[23]], # after LRB0 and before LRB2
|
|
'LWSA' : [[54]], # For A
|
|
'LWSB' : [[54]],
|
|
|
|
'LCC' : [[55, 55]],
|
|
}
|
|
# note: syncCode needs to be
|
|
syncCode = [SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for necessary prior LRA1/LRB1 before starting main loop"),
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 for the remaining main loop"),
|
|
SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete to start GRA"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=11, vscnt=-1, comment="Wait for LRB0/GRA to complete to start GRB/LRA1"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="Wait for GRB to complete to start LRB1"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 11 # vmcnt shift for ngl and nll
|
|
elif isNN(kernel) and useLDSTr and TLDS==1:
|
|
optSchedule = {
|
|
|
|
'SYNC' : [[-1,3, 12,12, 26,27, 45,45]],
|
|
'GRIncA' : [[0,1,2,3,4,5,6,7,8]],
|
|
'GRIncB' : [[22,22,24,24,25,25,26,27,27]],
|
|
|
|
'LRA0' : [[0,0,2,2,3,3,5,6]], ## -2 is place holder
|
|
|
|
'LRB0' : [[8,10,13,15,17,19,21], ## After LRA0, we can mix LRB0 and GRA
|
|
[9,11,14,16,18,20,22]],
|
|
## GRA should start after LRA0 is done.
|
|
'GRA' : [[10,12, 14,14, 17,17, 20,20],
|
|
[11,13, 15,15, 18,18, 21,21]],
|
|
|
|
## GRB should start after LRB0 is done
|
|
'GRB' : [[28,28, 31,31, 34,34, 37,37, 40,40, 43,43, 46,46], # m0 inc is part of GRA/GRB
|
|
[29,29, 32,32, 35,35, 38,38, 41,41, 44,44, 47,47]],
|
|
'LRA1' : [[29,30, 32,33, 35,36, 38,39],
|
|
[30,31, 33,34, 36,37, 39,40]],
|
|
|
|
#After GRB is done.
|
|
'LRB1' : [[45,46,47,48,49,50,51]],
|
|
|
|
'LRSA' : [[23]], # after LRA0 and before LRA1
|
|
'LRSB' : [[23]], # after LRB0 and before LRB2
|
|
'LWSA' : [[54]], # For A
|
|
'LWSB' : [[54]],
|
|
|
|
'LCC' : [[55, 55]],
|
|
}
|
|
# note: syncCode needs to be
|
|
syncCode = [SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for necessary prior LRA1/LRB1 before starting main loop"),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1/LRB1 for the remaining main loop"),
|
|
SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete to start GRA"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=0, vlcnt=11, vscnt=-1, comment="Wait for LRB0/GRA to complete to start GRB/LRA1"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="Wait for GRB to complete to start LRB1"),
|
|
SBarrier(comment=""),
|
|
]
|
|
nglshift = nllshift = 11 # vmcnt shift for ngl and nll
|
|
elif isNT(kernel) and useLDSTr and TLDS == 0:
|
|
valid, opt = _get_schedule_224x128x64_16bit(kernel, useLDSTr, TLDS)
|
|
if not valid:
|
|
return False, None
|
|
optSchedule = switch_A_B_schedule(opt.optSchedule)
|
|
return True, ScheduleInfo(opt.numCodePaths, opt.numMfma, optSchedule, opt.syncCode, opt.nglshift, opt.nllshift)
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 56
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 192, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x192x64_16bit(kernel, useLDSTr, TLDS):
|
|
"""128x192x64 TN schedule (BF16/FP16)."""
|
|
kernel["MfmaInitCVgprs"] = True
|
|
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
numMfma = 2 * kernel["MIWaveTileA"] * kernel["MIWaveTileB"] # 48
|
|
|
|
syncs = SyncSchedule()
|
|
gr_inc_step = 0
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
grinca = [0,1,2, 3,4,5, 6,7,7]
|
|
grincb = [7,8,9, 9,9,10, 10,10,11]
|
|
|
|
syncs.add(-1, dscnt=5, comment="wait for all LRA1 and one item from LRB1 before starting the sub-iteration")
|
|
lra0 = [0,1,2,3]
|
|
syncs.add( 3, dscnt=3, comment="wait for the rest of LRB1 to complete")
|
|
lrb0 = [ 4,5,6,8,11, 14]
|
|
|
|
syncs.add( 12, dscnt=5, barrier=True, comment="wait for LRA0 before GRA start")
|
|
gra = [ 13,15,17,19] # one index for two instructions
|
|
|
|
syncs.add( 21, dscnt=0, vlcnt=4+6, barrier=True, comment="wait for LRB0 before GRB start + wait for previous GRAs before LRA1")
|
|
grb = [ 21,24, 27,31,34,37] # one index for two instructions
|
|
lra1 = [ 22,25,26,29]
|
|
syncs.add( 35, vlcnt=4+5, barrier=True, comment="wait for previous GRBs to complete before LRB1")
|
|
lrb1 = [ 35,38,40,42,44,46]
|
|
|
|
num_gr = len(gra) + len(grb)
|
|
lrsa = [18]
|
|
lrsb = [18]
|
|
lwsa = [30]
|
|
lwsb = [30]
|
|
|
|
else:
|
|
return False, None
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'LRA0': [lra0],
|
|
'GRIncA': [grinca],
|
|
'LRB0': [lrb0],
|
|
'GRIncB': [grincb],
|
|
# Note: each GRA/GRB item corresponds to two instructions (addr increment and read). So duplicate each item twice.
|
|
'GRA': [duplicate_list_items(gra, 2, gr_inc_step)],
|
|
'GRB': [duplicate_list_items(grb, 2, gr_inc_step)],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'LRA1': [lra1],
|
|
'LRB1': [lrb1],
|
|
'LCC': [[numMfma-3, numMfma-2]],
|
|
}
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
opt1 = ScheduleInfo(1, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 192, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x192x32_TF32(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isNN(kernel) and not useLDSTr and TLDS==1:
|
|
# TODO: Add NN schedule in upcoming PR
|
|
return False, None
|
|
elif isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = False
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
syncTable = [
|
|
5, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Before PackA0. Wait for all LRA0. Skip 1*LRB0.") ,
|
|
17, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Before PackB0. Wait for all prior LRB0 for PackB0.") ,
|
|
17, SBarrier(comment="GRA") ,
|
|
32, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Before GRB. Wait for all prior LRB0.") ,
|
|
32, SBarrier(comment="GRB") ,
|
|
35, SWaitCnt(dscnt=-1, vlcnt=6, vscnt=-1, comment="Before LRB3. Wait for GRB from previous iter. Skip 4*GRA + 2*GRB") ,
|
|
35, SBarrier(comment="LRB") ,
|
|
44, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Before PackB3. Wait for all prior LRB3.") ,
|
|
53, SWaitCnt(dscnt=0, vlcnt=10, vscnt=-1, comment="Before LRA3. Wait for GRA from previous iter. Skip 4*GRA + 6*GRB") ,
|
|
53, SBarrier(comment="LRA") ,
|
|
63, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Before PackA3. Wait for all prior LRA3.") ,
|
|
]
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[0, 0, 0, 1, 1, 1, 2, 2, 2]],
|
|
'GRIncB': [[3, 3, 3, 4, 4, 4, 5, 5, 5]],
|
|
'LRA0' : [[0, 1, 2, 3]],
|
|
'LRB0' : [[4, 6, 8, 10, 12, 14],
|
|
[4, 7, 9, 11, 13, 15]],
|
|
'PackA0': [create_range(5,12,17, 1, 4)],
|
|
'PackB0': [create_range(18,18,34, 1, 4)],
|
|
'GRA' : [[17,17, 18,18, 19,19, 20,20]],
|
|
'GRB' : [[33, 33, 34, 34, 41, 42, 43, 44, 51, 51, 52, 52]],
|
|
'LRB3' : [[36, 37, 38, 39, 40, 41]],
|
|
'LRA3' : [[53, 55, 57, 59],
|
|
[54, 56, 58, 60]],
|
|
'PackB3': [create_range(44,9,52, 1, 8)],
|
|
'PackA3': [create_range(63,8,71, 1, 6)],
|
|
'LRSA' : [[16]],
|
|
'LRSB' : [[16]],
|
|
'LWSA' : [[61]],
|
|
'LWSB' : [[62]],
|
|
'LCC' : [[71, 71]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 10
|
|
elif isNT(kernel) and useLDSTr and TLDS==0:
|
|
# TODO: Add NT schedule in upcoming PR
|
|
return False, None
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 72
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(192, 256, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_192x256x32_TF32(kernel, useLDSTr, TLDS):
|
|
numMfma = 144
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
mfmaReorder = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
|
|
# Used the following constrains to create schedule
|
|
# - LRA0 + PACKA0 needs to be done before 1/4 MFMAs
|
|
# - LBR0 + PACKB0 needs to be done before 2/4 MFMAs
|
|
# - LRB3 + PACKB3 needs to start after 2/4 MFMAs
|
|
# - LRA3 + PACKA3 needs to start after 3/4 MFMAs
|
|
|
|
# LRA0 + GRIncA
|
|
lra0 = create_range(min_val = 0, num = 6, step = 1, repeat = 1)
|
|
grIncA = create_range(min_val = max(lra0)+1, num = 3, step = 1, repeat = 3)
|
|
# Hide LRA0 latency behind GRIncA
|
|
waitLRA0 = max(grIncA)+5
|
|
startPACKA0 = waitLRA0
|
|
|
|
# Reordering of packA instructions.
|
|
# 4 CVT + 2 4x4x4_16B MFMAs + 4 CVTs
|
|
# we interleave the 3 blocks together to avoid :
|
|
# - having a 5 state wait after each 4x4x4_16B MFMA
|
|
# - having extra latency when switching between MFMA types
|
|
packAOffset = [
|
|
0, 0, 1, 1,
|
|
6, 6,
|
|
7, 7, 8, 8,
|
|
|
|
2, 2, 3, 3,
|
|
6, 6,
|
|
9, 9, 10, 11,
|
|
|
|
4, 4, 5, 5,
|
|
6, 6,
|
|
12, 12, 13, 13,
|
|
]
|
|
|
|
|
|
packA0 = [x + startPACKA0 for x in packAOffset]
|
|
packA0Done = max(packA0)
|
|
|
|
# Sanity check
|
|
assert packA0Done < numMfma//4
|
|
|
|
# LRB0 + GRIncB
|
|
lrb0 = create_range(min_val = max(packA0)+1, num = 8, step = 1, repeat = 1)
|
|
grIncB = create_range(max(lrb0)+1,3,max(lrb0)+4,1,3)
|
|
waitLRB0 = max(grIncB)+6
|
|
startPACKB0 = waitLRB0
|
|
packBOffset = [
|
|
0, 0, 1, 1,
|
|
8, 8,
|
|
9, 9, 10, 11,
|
|
|
|
2, 2, 3, 3,
|
|
8, 8,
|
|
12, 12, 13, 13,
|
|
|
|
4, 4, 5, 5,
|
|
8, 8,
|
|
14, 14, 15, 15,
|
|
|
|
6, 6, 7, 7,
|
|
8, 8,
|
|
16, 16, 17, 17,
|
|
]
|
|
|
|
packB0 = [x + startPACKB0 for x in packBOffset]
|
|
|
|
# GRA
|
|
grA = [create_range(min_val = max(packB0)+1, num = 6, step = 2,repeat = 2),
|
|
create_range(min_val = max(packB0)+2, num = 6, step = 2,repeat = 2)]
|
|
|
|
halfMFMA = numMfma//2
|
|
assert max(packB0) < halfMFMA
|
|
|
|
# LR3
|
|
startLRB3 = halfMFMA
|
|
lrb3 = create_range(min_val = startLRB3, num = 2, step = 1, repeat = 2)
|
|
lrb3 += create_range(min_val = max(lrb3)+6,num = 2, step = 1, repeat = 2)
|
|
|
|
# GRB (split in two blocks)
|
|
grB = create_range(min_val = max(lrb3)+1,num = 4,step = 2, repeat = 2)
|
|
waitLRB3 = max(grB)+1
|
|
grB += create_range(min_val = max(grB)+47,num = 4,step = 2, repeat = 2)
|
|
|
|
# PackB3 (starts after 1st GRB block)
|
|
packB3 = [x + waitLRB3 for x in packBOffset]
|
|
|
|
# LRA3 + PACKA3
|
|
startLRA3 = (3*numMfma)//4 # Can't start before 3/4 MFMAs
|
|
lra3 = create_range(min_val = startLRA3,num=6,step=1,repeat=1)
|
|
waitLRA3 = max(lra3) + 8
|
|
packA3 = [x + waitLRA3 for x in packAOffset]
|
|
|
|
syncTable = [
|
|
waitLRA0, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
waitLRB0, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
|
|
max(packB0)+1, SBarrier(comment="Barrier before GRA&GRB"),
|
|
|
|
startLRB3-1,SWaitCnt(dscnt=-1, vlcnt=5, vscnt=-1, comment="Wait for previous GRA&B"),
|
|
startLRB3-1,SBarrier(comment=""),
|
|
|
|
waitLRB3,SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB3 to complete"),
|
|
waitLRA3, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
'SYNC': [syncTable[::2]],
|
|
|
|
'GRIncA': [grIncA],
|
|
'GRIncB': [grIncB],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0' : [packA0],
|
|
'PackB0' : [packB0],
|
|
|
|
'GRA': [*grA],
|
|
'GRB': [grB],
|
|
'LRSA': [[max(grIncB)+1]],
|
|
'LRSB': [[max(grIncB)+2]],
|
|
'LWSA': [[142]],
|
|
'LWSB': [[142]],
|
|
'LCC': [[143, 143]],
|
|
'LRA3': [lra3],
|
|
'LRB3': [lrb3],
|
|
'PackB3' : [packB3],
|
|
'PackA3' : [packA3],
|
|
}
|
|
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
elif isNN(kernel) and TLDS==1 and kernel["VectorWidthA"] == 1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
|
|
numLrReadA = 24
|
|
numLrReadB = 8
|
|
|
|
# A is 24 instructions as current codegen can't generate ds_read_b128 in NN case.
|
|
# Instead of reading A first, this schedule re-orders the mfma instructions and reads B first (less instructions)
|
|
# Before :
|
|
# B0 - A0
|
|
# B0 - A1
|
|
# B1 - A0
|
|
# B1 - A1
|
|
# Now :
|
|
# B0 - A0
|
|
# B1 - A0
|
|
# B0 - A1
|
|
# B1 - A1
|
|
|
|
# mfma Reordering
|
|
mfmaReorder = [i for i in range(0,numMfma//4)] + [i for i in range(numMfma//2,3*numMfma//4)]+[i for i in range(numMfma//4,numMfma//2)]+[i for i in range(3*numMfma//4,numMfma)]
|
|
|
|
# Interleaving of LBR0 and GRINCB to hide LRB0 latency
|
|
lrb0 = create_range(min_val = 0, num = 6, step = 1, repeat = 1)
|
|
grIncB = create_range(min_val = max(lrb0)+1, num = 3, step = 1, repeat = 3)
|
|
lrb0 += create_range(min_val = max(grIncB)+1, num = 2, step = 1, repeat = 1)
|
|
grIncA = create_range(min_val = max(lrb0)+1, num = 3, step = 1, repeat = 3)
|
|
waitLRB0 = max(grIncA)+4
|
|
|
|
# PackB0 using mfma4x4x4_16b
|
|
startPACKB0 = waitLRB0
|
|
packBOffset = [
|
|
0, 0, 1, 1,
|
|
8, 8,
|
|
9, 9, 10, 10,
|
|
|
|
2, 2, 3, 3,
|
|
8, 8,
|
|
11, 11, 12, 12,
|
|
|
|
4, 4, 5, 5,
|
|
8, 8,
|
|
13, 13, 14, 14,
|
|
|
|
6, 6, 7, 7,
|
|
8, 8,
|
|
15, 15, 16, 16,
|
|
]
|
|
|
|
packB0 = [x + startPACKB0 for x in packBOffset]
|
|
packB0Done = max(packB0)
|
|
|
|
# Sanity check
|
|
assert packB0Done < numMfma//4
|
|
|
|
# GRB (1st block) interleaved with LRA0
|
|
grB = [create_range(min_val = packB0Done+2,num = 4,step = 4, repeat = 2),
|
|
create_range(min_val = packB0Done+1,num = 4,step = 4, repeat = 2)]
|
|
|
|
# LRA0
|
|
lra0 = [create_range(min_val = max(packB0)+1, num = numLrReadA // 2, step = 2, repeat = 2),
|
|
create_range(min_val = max(packB0)+2, num = numLrReadA // 2, step = 2, repeat = 2)]
|
|
|
|
# PackA0
|
|
waitLRA0 = max(lra0[1])+2
|
|
startPACKA0 = waitLRA0
|
|
packAOffset = [
|
|
0, 0, 1, 1,
|
|
6, 6,
|
|
7, 7, 8, 8,
|
|
|
|
2, 2, 3, 3,
|
|
6, 6,
|
|
9, 9, 10, 10,
|
|
|
|
4, 4, 5, 5,
|
|
6, 6,
|
|
11, 11, 12, 12,
|
|
]
|
|
|
|
packA0 = [x + startPACKA0 for x in packAOffset]
|
|
|
|
halfMFMA = numMfma//2
|
|
assert max(packA0) < halfMFMA
|
|
|
|
# LRA3 interleaved with GRB (2nd half)
|
|
startLRA3 = halfMFMA
|
|
lra3 = [create_range(min_val = startLRA3, num = numLrReadA // 2, step = 2, repeat = 2),
|
|
create_range(min_val = startLRA3+1, num = numLrReadA // 2, step = 2, repeat = 2)]
|
|
grB[0] += create_range(min_val = startLRA3+1,num = 4,step = 2, repeat = 2)
|
|
grB[1] += create_range(min_val = startLRA3,num = 4,step = 2, repeat = 2)
|
|
waitLRA3 = max(lra3[0])+4
|
|
# LRB3 + PACKA3 & PACKB3
|
|
startLRB3 = (3*numMfma)//4 - 4 # Starts 4 indexes before 3/4 MFMAs to accommodate LRB3 latency
|
|
lrb3 = create_range(min_val = startLRB3,num=numLrReadB - 2,step=1,repeat=1)
|
|
grA = [create_range(min_val = min(lrb3)+1, num = 8, step = 1,repeat = 1),
|
|
create_range(min_val = min(lrb3)+3, num = 8, step = 1,repeat = 1)]
|
|
lrb3 += create_range(min_val = max(lrb3)+3,num=2,step=1,repeat=1)
|
|
waitLRB3 = max(lrb3) + 6
|
|
|
|
# Grouping segment of 4x4x4_16B MFMAs together for PackB3 & PackA3 (reduce MFMA type switching cost)
|
|
packB3 = [x + waitLRB3 for x in packBOffset]
|
|
start_4x4x4 = packB3[4] # 5th index is start of 4x4x4_16B MFMA for PackB3
|
|
packA3 = [
|
|
*create_range(min_val = waitLRA3, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4,start_4x4x4,
|
|
*create_range(min_val = max(packB3)+1, num = 2, step = 1, repeat = 2),
|
|
|
|
*create_range(min_val = waitLRA3+2, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4,start_4x4x4,
|
|
*create_range(min_val = max(packB3)+3, num = 2, step = 1, repeat = 2),
|
|
|
|
*create_range(min_val = waitLRA3+4, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4,start_4x4x4,
|
|
*create_range(min_val = max(packB3)+5, num = 2, step = 1, repeat = 2),
|
|
]
|
|
|
|
# GRA 2nd half
|
|
grA[0] += create_range(min_val = max(packB3)+1, num = 2, step = 1,repeat = 2)
|
|
grA[1] += create_range(min_val = max(packB3)+1, num = 2, step = 1,repeat = 2)
|
|
|
|
syncTable = [
|
|
waitLRB0, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for 4/8 LRB0 to complete"),
|
|
waitLRB0+4, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRB0 to complete"),
|
|
waitLRB0+4, SBarrier(comment="Barrier before GRB"), #Barrier can be after CVT
|
|
|
|
# dscnt has a max value of 15
|
|
waitLRA0, SWaitCnt(dscnt=min(15,numLrReadA-4), vlcnt=-1, vscnt=-1, comment="Wait for 4 LRA0 to complete"),
|
|
waitLRA0+1, SWaitCnt(dscnt=min(15,numLrReadA-8), vlcnt=-1, vscnt=-1, comment="Wait for 8 LRA0 to complete"),
|
|
waitLRA0+2, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
|
|
startLRA3-1,SWaitCnt(dscnt=-1, vlcnt=4, vscnt=-1, comment="Wait for previous GRA&B"),
|
|
startLRA3-1,SBarrier(comment="Sync before GRA, LRA3 & LRB3"),
|
|
|
|
# incremental wait on LRA3
|
|
waitLRA3, SWaitCnt(dscnt=min(15,numLrReadA-4), vlcnt=-1, vscnt=-1, comment="Wait for 4 LRA3 to complete"),
|
|
waitLRA3+1, SWaitCnt(dscnt=min(15,numLrReadA-8), vlcnt=-1, vscnt=-1, comment="Wait for 8 LRA3 to complete"),
|
|
waitLRA3+2, SWaitCnt(dscnt=min(15,numLrReadA-12), vlcnt=-1, vscnt=-1, comment="Wait for 12 LRA3 to complete"),
|
|
waitLRA3+3, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
|
|
# incremental wait on LRB3
|
|
waitLRB3, SWaitCnt(dscnt=(numLrReadB-1), vlcnt=-1, vscnt=-1, comment="Wait for 1st LRB3 to complete"),
|
|
waitLRB3+1, SWaitCnt(dscnt=(numLrReadB-2), vlcnt=-1, vscnt=-1, comment="Wait for 2nd LRB3 to complete"),
|
|
waitLRB3+2, SWaitCnt(dscnt=(numLrReadB-3), vlcnt=-1, vscnt=-1, comment="Wait for 3rd LRB3 to complete"),
|
|
waitLRB3+3, SWaitCnt(dscnt=(numLrReadB-4), vlcnt=-1, vscnt=-1, comment="Wait for 4th LRB3 to complete"),
|
|
waitLRB3+4, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRB3 to complete"),
|
|
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
|
|
'SYNC': [syncTable[::2]],
|
|
|
|
'GRIncA': [grIncA],
|
|
'GRIncB': [grIncB],
|
|
'LRA0': [*lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0' : [packA0],
|
|
'PackB0' : [packB0],
|
|
'GRA': [*grA],
|
|
'GRB': [*grB],
|
|
'LRSA': [[max(lra0[1])+1]],
|
|
'LRSB': [[max(lra0[1])+1]],
|
|
'LWSA': [[numMfma-2]],
|
|
'LWSB': [[numMfma-2]],
|
|
'LCC': [[numMfma-1, numMfma-1]],
|
|
'LRA3': [*lra3],
|
|
'LRB3': [lrb3],
|
|
'PackB3' : [packB3],
|
|
'PackA3' : [packA3],
|
|
|
|
}
|
|
|
|
nglshift = nllshift = 14
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift, mfmaReorder=mfmaReorder)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 192, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x192x32_TF32(kernel, useLDSTr, TLDS):
|
|
numMfma = 144
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and not useLDSTr and TLDS == 1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = False
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
numPackInstr = 24
|
|
numPackIndices = numPackInstr // 2 # Assign 2 pack instructions per mfma index
|
|
|
|
# LRA0 + PACKA0 - done before 1/4 MFMAs - index 36
|
|
lrA0 = [0,0, 1,1, 2,2, 3,3]
|
|
waitLRA0 = max(lrA0) + 2
|
|
startPACKA0 = waitLRA0
|
|
packA0 = create_range(startPACKA0, (len(lrA0)//2)*numPackIndices, numMfma//4-1)
|
|
|
|
# LBR0 + PACKB0 - done before 2/4 MFMAs - index 72
|
|
lrB0 = [7,7, 11,11, 15,15]
|
|
waitLRB0 = max(lrB0) + 2
|
|
startPACKB0 = max(waitLRB0,max(packA0)) # Starts after waitLRB0 and packA0
|
|
packB0 = create_range(startPACKB0, (len(lrB0)//2)*numPackIndices, numMfma//2-1)
|
|
|
|
# LBR3 + PACKB3 - start after 2/4 MFMAs - index 72
|
|
halfMFMA = numMfma//2
|
|
startLRB3 = halfMFMA
|
|
lrB3 = create_range(startLRB3, 1, numMfma-1)
|
|
lrB3 += create_range(max(lrB3)+6, 2, numMfma-1)
|
|
waitLRB3 = startLRB3 + 6
|
|
packB3 = create_range(waitLRB3, (len(lrB3)//2)*numPackIndices, numMfma-1)
|
|
|
|
# LRA3 + PACKA3 - start after 3/4 MFMAs - index 108
|
|
startLRA3 = (3*numMfma)//4
|
|
lrA3 = create_range(startLRA3, 4, numMfma-1)
|
|
waitLRA3 = startLRA3 + 5
|
|
packA3 = create_range(waitLRA3, (len(lrA3)//2)*numPackIndices, numMfma-1)
|
|
|
|
syncTable = [
|
|
waitLRA0, SWaitCnt(dscnt=inflight(lrA0, waitLRA0)-2, vlcnt=-1, vscnt=-1, comment="wait for 1st 2 LRA0 to complete"),
|
|
waitLRA0+numPackIndices, SWaitCnt(dscnt=inflight(lrB0, waitLRA0+numPackIndices), vlcnt=-1, vscnt=-1, comment="wait for all LRA0 to complete"),
|
|
|
|
waitLRB0, SWaitCnt(dscnt=inflight(lrB0, waitLRB0)-2, vlcnt=-1, vscnt=-1, comment="wait for 1st 2 LRB0 to complete"),
|
|
waitLRB0+numPackIndices, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRB0 to complete"),
|
|
waitLRB0+numPackIndices, SBarrier(comment="Barrier before GRA&GRB"),
|
|
|
|
startLRB3-1, SWaitCnt(dscnt=-1, vlcnt=6, vscnt=-1, comment="Wait for prev GRA&GRB"),
|
|
startLRB3-1, SBarrier(comment=""),
|
|
|
|
waitLRB3,SWaitCnt(dscnt=inflight(lrB3, waitLRB3)-2, vlcnt=-1, vscnt=-1, comment="Wait for 1st 2 LRB3 to complete"),
|
|
waitLRB3+numPackIndices,SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRB3 to complete"),
|
|
|
|
waitLRA3, SWaitCnt(dscnt=inflight(lrA3,waitLRA3)-2, vlcnt=-1, vscnt=-1, comment="Wait for 1st 2 LRA3 to complete"),
|
|
waitLRA3+numPackIndices, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRA3 to complete")
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
|
|
'GRIncA' : [[0,0,0, 1,1,1, 2,2,2]],
|
|
'GRIncB' : [[3,3,3, 4,4,4, 5,5,5]],
|
|
|
|
'LRA0' : [lrA0],
|
|
'PackA0' : [packA0],
|
|
|
|
'LRB0' : [lrB0],
|
|
'PackB0' : [packB0],
|
|
|
|
'GRA' : [[72,72, 74,74, 76,76, 100,100, 102,102, 104,104, 106,106, 108,108],
|
|
[73,73, 75,75, 77,77, 101,101, 103,103, 105,105, 107,107, 109,109]],
|
|
'GRB' : [[38,38, 40,40, 42,42, 44,44, 46,46, 48,48],
|
|
[39,39, 41,41, 43,43, 45,45, 47,47, 49,49]],
|
|
|
|
'LRA3' : [lrA3],
|
|
'PackA3' : [packA3],
|
|
|
|
'LRB3' : [lrB3],
|
|
'PackB3' : [packB3],
|
|
|
|
'LRSA' : [[35]],
|
|
'LRSB' : [[35]],
|
|
|
|
'LWSA' : [[107]],
|
|
'LWSB' : [[107]],
|
|
|
|
'LCC' : [[143, 143]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 14 # vmcnt shift for ngl and nll
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
|
|
elif isNN(kernel) and TLDS==1 and kernel["VectorWidthA"] == 1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
|
|
numLrReadA = 32
|
|
numLrReadB = 6
|
|
|
|
# mfma Reordering
|
|
mfmaReorder = [i for i in range(0,numMfma//4)] + [i for i in range(numMfma//2,3*numMfma//4)]+[i for i in range(numMfma//4,numMfma//2)]+[i for i in range(3*numMfma//4,numMfma)]
|
|
|
|
# Interleave LBR0 and GRINCB
|
|
lrb0 = create_range(min_val = 0, num = 6, step = 1, repeat = 1)
|
|
grIncB = create_range(min_val = 0, num = 6, step = 1, repeat = 1)
|
|
grIncB += create_range(min_val = max(lrb0)+1, num = 1, step = 1, repeat = 3)
|
|
|
|
# Interleave GRINCA and PACKB0
|
|
grIncA = create_range(min_val = max(grIncB)+1, num = 2, step = 1, repeat = 3)
|
|
waitLRB0 = max(grIncA)
|
|
grIncA += create_range(min_val = max(grIncA)+2, num = 3, step = 1, repeat = 1)
|
|
|
|
startPACKB0 = waitLRB0
|
|
packBOffset = [
|
|
0, 0, 1, 1,
|
|
6, 6,
|
|
7, 7, 8, 8,
|
|
|
|
2, 2, 3, 3,
|
|
6, 6,
|
|
9, 9, 10, 10,
|
|
|
|
4, 4, 5, 5,
|
|
6, 6,
|
|
11, 11, 12, 12,
|
|
]
|
|
|
|
packB0 = [x + startPACKB0 for x in packBOffset]
|
|
packB0Done = max(packB0)
|
|
|
|
# Sanity check
|
|
assert packB0Done < numMfma//4, f"packB0Done {packB0Done} >= numMfma//4 {numMfma//4}"
|
|
|
|
# GRB (1st block) interleaved with LRA0
|
|
grB = [create_range(min_val = packB0Done+2,num = 3,step = 4, repeat = 2),
|
|
create_range(min_val = packB0Done+1,num = 3,step = 4, repeat = 2)]
|
|
|
|
# LRA0
|
|
lra0 = [create_range(min_val = max(packB0)+1, num = numLrReadA // 2, step = 2, repeat = 2),
|
|
create_range(min_val = max(packB0)+2, num = numLrReadA // 2, step = 2, repeat = 2)]
|
|
|
|
# PackA0
|
|
waitLRA0 = max(lra0[1])+1
|
|
startPACKA0 = waitLRA0
|
|
|
|
packAOffset = [
|
|
0, 0, 1, 1,
|
|
8, 8,
|
|
9, 9, 10, 10,
|
|
|
|
2, 2, 3, 3,
|
|
8, 8,
|
|
11, 11, 12, 12,
|
|
|
|
4, 4, 5, 5,
|
|
8, 8,
|
|
13, 13, 14, 14,
|
|
|
|
6, 6, 7, 7,
|
|
8, 8,
|
|
15, 15, 16, 16,
|
|
]
|
|
packA0 = [x + startPACKA0 for x in packAOffset]
|
|
|
|
halfMFMA = numMfma//2
|
|
assert max(packA0) < halfMFMA, f"max(packA0) {max(packA0)} >= halfMFMA {halfMFMA}"
|
|
|
|
# LRA3 interleaved with GRB (2nd half)
|
|
startLRA3 = halfMFMA
|
|
lra3 = [create_range(min_val = startLRA3+1, num = numLrReadA // 2, step = 2, repeat = 2),
|
|
create_range(min_val = startLRA3, num = numLrReadA // 2, step = 2, repeat = 2)]
|
|
|
|
# M0 update before barrier to prevent bad interleaving between the 2 codepaths
|
|
grB[0] += [startLRA3-2,startLRA3]
|
|
grB[1] += [startLRA3-2,startLRA3+1]
|
|
|
|
grB[0] += create_range(min_val = startLRA3+2,num = 2,step = 2, repeat = 2)
|
|
grB[1] += create_range(min_val = startLRA3+3,num = 2,step = 2, repeat = 2)
|
|
|
|
# LRB3 + PACKA3 & PACKB3
|
|
startLRB3 = (3*numMfma)//4 - 4 # Starts 4 indexes before 3/4 MFMAs to accommodate LRB3 latency
|
|
lrb3 = create_range(min_val = startLRB3,num=numLrReadB - 2,step=1,repeat=1)
|
|
|
|
grA = [create_range(min_val = max(lra3[0])+1, num = 8, step = 1,repeat = 1),
|
|
create_range(min_val = max(lra3[1])+1, num = 8, step = 1,repeat = 1)]
|
|
lrb3 += create_range(min_val = max(lrb3)+3,num=2,step=1,repeat=1)
|
|
|
|
waitLRB3 = max(lrb3) + 9
|
|
|
|
# Grouping segment of 4x4x4_16B MFMAs together for PackB3 & PackA3 (reduce MFMA type switching cost)
|
|
packB3 = [x + waitLRB3 for x in packBOffset]
|
|
start_4x4x4 = packB3[4] # 5th index is start of 4x4x4_16B MFMA for PackB3
|
|
waitLRA3 = max(lrb3)+1
|
|
packA3 = [
|
|
*create_range(min_val = waitLRA3, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4, start_4x4x4,
|
|
*create_range(min_val = max(packB3)+1, num = 2, step = 1, repeat = 2),
|
|
|
|
*create_range(min_val = waitLRA3+2, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4, start_4x4x4,
|
|
*create_range(min_val = max(packB3)+3, num = 2, step = 1, repeat = 2),
|
|
|
|
*create_range(min_val = waitLRA3+4, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4, start_4x4x4,
|
|
*create_range(min_val = max(packB3)+5, num = 2, step = 1, repeat = 2),
|
|
|
|
*create_range(min_val = waitLRA3+6, num = 2, step = 1, repeat = 2),
|
|
start_4x4x4, start_4x4x4,
|
|
*create_range(min_val = max(packB3)+7, num = 2, step = 1, repeat = 2),
|
|
]
|
|
|
|
# GRA 2nd half
|
|
grA[0] += create_range(min_val = max(packB3)+1, num = 4, step = 1,repeat = 2)
|
|
grA[1] += create_range(min_val = max(packB3)+1, num = 4, step = 1,repeat = 2)
|
|
|
|
syncTable = [
|
|
waitLRB0, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for 1/6 LRB0 to complete"),
|
|
waitLRB0+1, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for 2/6 LRB0 to complete"),
|
|
waitLRB0+2, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for 3/6 LRB0 to complete"),
|
|
waitLRB0+3, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for 4/6 LRB0 to complete"),
|
|
waitLRB0+4, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for 5/6 LRB0 to complete"),
|
|
waitLRB0+5, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for 6/6 LRB0 to complete"),
|
|
waitLRB0+5, SBarrier(comment="Barrier before GRB"),
|
|
|
|
# incremental wait on LRA0
|
|
waitLRA0, SWaitCnt(dscnt=min(15,numLrReadA-4), vlcnt=-1, vscnt=-1, comment="Wait for 4 LRA0 to complete"),
|
|
waitLRA0+4, SWaitCnt(dscnt=numLrReadA-20, vlcnt=-1, vscnt=-1, comment="Wait for 20 LRA0 to complete"),
|
|
waitLRA0+5, SWaitCnt(dscnt=numLrReadA-24, vlcnt=-1, vscnt=-1, comment="Wait for 24 LRA0 to complete"),
|
|
waitLRA0+6, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
|
|
startLRA3-1,SWaitCnt(dscnt=-1, vlcnt=3, vscnt=-1, comment="Wait for previous GRA&B"),
|
|
startLRA3-1,SBarrier(comment="Sync before GRA, LRA3 & LRB3"),
|
|
|
|
# incremental wait on LRA3 & LRB3
|
|
waitLRA3, SWaitCnt(dscnt=15, vlcnt=-1, vscnt=-1, comment="Wait for 17/32 LRA3 to complete"),
|
|
waitLRA3+4, SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
waitLRA3+7, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB3 to complete"),
|
|
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
|
|
'SYNC': [syncTable[::2]],
|
|
|
|
'GRIncA': [grIncA],
|
|
'GRIncB': [grIncB],
|
|
'LRA0': [*lra0],
|
|
'LRB0': [lrb0],
|
|
'LRSA': [[packA0[4]]],#Use slot between MFMA 16x16x32 & 4x4x4 for LRSA
|
|
'LRSB': [[packA0[4]]],
|
|
'PackA0' : [packA0],
|
|
'PackB0' : [packB0],
|
|
'GRA': [*grA],
|
|
'GRB': [*grB],
|
|
'LWSA': [[numMfma-2]],
|
|
'LWSB': [[numMfma-2]],
|
|
'LCC': [[numMfma-1, numMfma-1]],
|
|
'LRA3': [*lra3],
|
|
'LRB3': [lrb3],
|
|
'PackB3' : [packB3],
|
|
'PackA3' : [packA3],
|
|
|
|
}
|
|
|
|
nglshift = nllshift = 14
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift, mfmaReorder=mfmaReorder)
|
|
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 256, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x256x32_TF32(kernel, useLDSTr, TLDS):
|
|
numMfma = 192
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
# This schedule follows similar pattern as 192x256x32 TF32 schedule
|
|
|
|
# LRA0 + GRIncA
|
|
lra0 = create_range(min_val = 0, num = 4, step = 1, repeat = 1)
|
|
lra0 += create_range(min_val = max(lra0)+4, num = 4, step = 1, repeat = 1)
|
|
|
|
grIncA = create_range(min_val = max(lra0)+1, num = 3, step = 1, repeat = 3)
|
|
waitLRA0 = max(grIncA)+5
|
|
startPACKA0 = waitLRA0
|
|
|
|
# Use a common packOffset re-ordering for both A and B
|
|
packOffset = [
|
|
0, 0, 1, 1,
|
|
8, 8,
|
|
9, 9, 10, 11,
|
|
|
|
2, 2, 3, 3,
|
|
8, 8,
|
|
12, 12, 13, 13,
|
|
|
|
4, 4, 5, 5,
|
|
8, 8,
|
|
14, 14, 15, 15,
|
|
|
|
6, 6, 7, 7,
|
|
8, 8,
|
|
16, 16, 17, 17,
|
|
]
|
|
|
|
|
|
packA0 = [x + startPACKA0 for x in packOffset]
|
|
packA0Done = max(packA0)
|
|
|
|
# Sanity check
|
|
assert packA0Done < numMfma//4
|
|
|
|
# LRB0 + GRIncB (Split LRB0 into two halves to hide latency)
|
|
lrb0 = create_range(min_val = max(packA0)+1, num = 4, step = 1, repeat = 1)
|
|
lrb0 += create_range(min_val = max(lrb0)+4, num = 4, step = 1, repeat = 1)
|
|
grIncB = create_range(max(lrb0)+1,3,max(lrb0)+4,1,3)
|
|
waitLRB0 = max(grIncB)+6
|
|
startPACKB0 = waitLRB0
|
|
|
|
packB0 = [x + startPACKB0 for x in packOffset]
|
|
|
|
# GRA - 1st half (4 reads)
|
|
grA = create_range(min_val = max(packB0)+1, num = 4, step = 2,repeat = 2)
|
|
|
|
halfMFMA = numMfma//2
|
|
assert max(packB0) < halfMFMA
|
|
|
|
# LR3
|
|
startLRB3 = halfMFMA
|
|
lrb3 = create_range(min_val = startLRB3, num = 4, step = 1, repeat = 1)
|
|
lrb3 += create_range(min_val = max(lrb3)+4, num = 4, step = 1, repeat = 1)
|
|
|
|
# GRA - 2nd half (4 reads)
|
|
grA += create_range(min_val = max(lrb3)+1, num = 4, step = 2,repeat = 2)
|
|
waitLRB3 = max(grA)+1
|
|
|
|
# PackB3
|
|
packB3 = [x + waitLRB3 for x in packOffset]
|
|
|
|
# GRB - 1st half (4 reads)
|
|
grB = create_range(min_val = max(packB3)+1,num = 4,step = 2, repeat = 2)
|
|
|
|
# LRA3 + PACKA3
|
|
startLRA3 = (3*numMfma)//4
|
|
lra3 = create_range(min_val = startLRA3,num=4,step=1,repeat=1)
|
|
lra3 += create_range(min_val = max(lra3)+4,num=4,step=1,repeat=1)
|
|
waitLRA3 = max(lra3) + 10
|
|
packA3 = [x + waitLRA3 for x in packOffset]
|
|
|
|
# GRB - 2nd half (4 reads)
|
|
grB += create_range(min_val = max(packA3)+1,num = 4,step = 2, repeat = 2)
|
|
|
|
syncTable = [
|
|
waitLRA0, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
waitLRB0, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
|
|
max(packB0)+1, SBarrier(comment="Barrier before GRA&GRB"),
|
|
|
|
startLRB3-1,SWaitCnt(dscnt=-1, vlcnt=4, vscnt=-1, comment="Wait for previous GRA&B"),
|
|
startLRB3-1,SBarrier(comment=""),
|
|
|
|
waitLRB3,SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB3 to complete"),
|
|
waitLRA3, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
|
|
'SYNC': [syncTable[::2]],
|
|
|
|
'GRIncA': [grIncA],
|
|
'GRIncB': [grIncB],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0' : [packA0],
|
|
'PackB0' : [packB0],
|
|
|
|
'GRA': [grA],
|
|
'GRB': [grB],
|
|
'LRSA': [[max(grIncB)+1]],
|
|
'LRSB': [[max(grIncB)+2]],
|
|
'LWSA': [[numMfma-2]],
|
|
'LWSB': [[numMfma-2]],
|
|
'LCC': [[numMfma-1, numMfma-1]],
|
|
'LRA3': [lra3],
|
|
'LRB3': [lrb3],
|
|
'PackB3' : [packB3],
|
|
'PackA3' : [packA3],
|
|
|
|
}
|
|
|
|
nglshift = nllshift = 16 # vmcnt shift for ngl and nll
|
|
elif isNT(kernel) and not useLDSTr and TLDS==0 and kernel["VectorWidthA"] == 4 and kernel["VectorWidthB"] == 4:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
swap_idx = [1,2,3, # depend on DS1
|
|
7,8, # depend on DS2
|
|
11, # depend on DS3
|
|
4,5,6, # depend on DS5
|
|
9,10, # depend on DS6
|
|
12, # depend on DS7
|
|
]
|
|
optSchedule = {
|
|
'SYNC': [[9, 14, 38, 43, 73, 95, 95, 106, 111, 157, 162]],
|
|
'GRIncA': [[0,1,2,3, 5,6,7,8, 11]],
|
|
'GRIncB': [[38, 39, 40, 41, 45, 46, 47, 50, 51]],
|
|
'LRA0': [[0,1,2,3, 5,6,7,8]],
|
|
'LRB0': [[28, 30, 31, 32, 34, 35, 36, 37]],
|
|
|
|
'PackA0': [[*[x + 8 for x in swap_idx],
|
|
20,20,21,21, 28,28, 29,29,30,30,
|
|
22,22,23,23, 28,28, 31,31,32,32,
|
|
24,24,25,25, 28,28, 33,33,34,34,
|
|
26,26,27,27, 28,28, 34,35,36,36]],
|
|
'PackB0': [[*[x + 37 for x in swap_idx],
|
|
55,55,56,56, 63,63, 64,64,65,65,
|
|
57,57,58,58, 63,63, 66,66,67,67,
|
|
59,59,60,60, 63,63, 68,68,69,69,
|
|
61,61,62,62, 63,63, 70,70,71,71]],
|
|
'LRSA': [[50]],
|
|
'LRSB': [[51]],
|
|
'LRB3': [[96,96,98,98, 103,103,105,105],
|
|
[97,97,99,99, 102,102,104,104]],
|
|
'GRA': [[75,75,80,80,85,85,90,90, 107,107,109,109,111,111,113,113],
|
|
[77,77,82,82,87,87,92,92, 108,108,110,110,112,112,114,114]],
|
|
'LRA3': [[138, 138, 144, 144, 150, 150, 154, 154],
|
|
[139, 139, 146, 146, 152, 152, 156, 156]],
|
|
'GRB': [[130,130,135,135,140,140,145,145, 165,165,170,170,175,175,180,180],
|
|
[132,132,137,137,142,142,147,147, 167,167,172,172,177,177,182,182]],
|
|
|
|
'PackB3': [[*[x + 105 for x in swap_idx],
|
|
119,119,120,120, 128,128, 129,129,130,130,
|
|
121,121,122,122, 128,128, 131,131,132,132,
|
|
123,123,124,124, 128,128, 133,133,134,134,
|
|
125,125,126,126, 128,128, 135,135,136,136]],
|
|
'PackA3': [[*[x + 156 for x in swap_idx],
|
|
169,169,170,170, 177,177, 178,178,179,179,
|
|
171,171,172,172, 177,177, 180,180,181,181,
|
|
173,173,174,174, 177,177, 182,182,183,183,
|
|
175,175,176,176, 177,177, 184,184,185,185]],
|
|
'LWSA': [[187]],
|
|
'LWSB': [[188]],
|
|
'LCC': [[189, 190]],
|
|
}
|
|
|
|
syncCode = [
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0 to complete"),
|
|
SBarrier(comment="Barrier before GRA&GRB"),
|
|
SWaitCnt(dscnt=-1, vlcnt=4, vscnt=-1, comment="Wait for previous GRA&B"),
|
|
SBarrier(comment=""),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for LRB3 to complete"),
|
|
SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB3 to complete"),
|
|
]
|
|
nglshift = nllshift = 16 # vmcnt shift for ngl and nll
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(192, 128, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_192x128x32_TF32(kernel, useLDSTr, TLDS):
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
if isTN(kernel) and useLDSTr and TLDS==1:
|
|
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = False
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
# Used the following constrains to create schedule
|
|
# - LRA0 + PACKA0 needs to be done before 1/4 MFMAs - index 18
|
|
# - LBR0 + PACKB0 needs to be done before 2/4 MFMAs - index 36
|
|
# - LRB3 + PACKB3 needs to start after 2/4 MFMAs - index 36
|
|
# - LRA3 + PACKA3 needs to start after 3/4 MFMAs - index 54
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment=""),
|
|
|
|
6, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="wait for first 2 LRA0"),
|
|
11, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRA0"),
|
|
|
|
18, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="wait for first 2 LRB0"),
|
|
23, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRB0"),
|
|
23, SBarrier(comment="Barrier before GRA&GRB"),
|
|
|
|
44, SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="Wait for prev GRA&GRB"),
|
|
44, SBarrier(comment=""),
|
|
|
|
49, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="wait for first 2 LRB3"),
|
|
56, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRB3"),
|
|
|
|
|
|
57, SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="Wait for prev GRA&GRB"),
|
|
57, SBarrier(comment=""),
|
|
|
|
61, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="wait for first 2 LRA3"),
|
|
66, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRA3"),
|
|
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA': [[0,0,0,1,1,1,2,2,2]],
|
|
'GRIncB': [[3,3,3,4,4,4,5,5,5]],
|
|
# - LRA0 + PACKA0 needs to be done before 1/4 MFMAs - index 18
|
|
'LRA0': [[0,1,2,3,4,5]],
|
|
'PackA0' : [[7]*4 + [8]*10 + [9]*10 + [12]*4 + [13]*20 + [14]*8 + [15]*8 + [16]*8],
|
|
# - LBR0 + PACKB0 needs to be done before 2/4 MFMAs - index 36
|
|
'LRB0': [[12,13,14,15]],
|
|
# First two LRB0 need to be done at 18, all LRB0 done by 23
|
|
'PackB0' : [create_range(19,4,22, repeat=3) + create_range(24,12,35, repeat=3) ],
|
|
'GRB': [[36,36,38,38,40,40,42,42],
|
|
[37,37,39,39,41,41,43,43]],
|
|
'GRA': [[45,45,47,47,49,49,51,51,53,53,55,55],
|
|
[46,46,48,48,50,50,52,52,54,54,56,56]],
|
|
'LRSA': [[36]],
|
|
'LRSB': [[36]],
|
|
'LWSA': [[52]],
|
|
'LWSB': [[52]],
|
|
'LCC': [[71, 71]],
|
|
# - LRB3 + PACKB3 needs to start after 2/4 MFMAs - index 36
|
|
'LRB3': [[45, 46, 47, 48]],
|
|
# First two LRB3 need to be done by 43, all LRB3 done at 48
|
|
'PackB3' : [create_range(50,12,71, repeat=4)],
|
|
# - LRA3 + PACKA3 needs to start after 3/4 MFMAs - index 54
|
|
'LRA3': [[58,59,60,62,63,64]],
|
|
# First two LRA3 need to be done by 61. All LRA0 done by 66.
|
|
'PackA3' : [create_range(62,12,71, repeat=6)],
|
|
|
|
}
|
|
|
|
nglshift = nllshift = 10 # vmcnt shift for ngl and nll
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
numMfma = 72
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 128, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x128x32_TF32(kernel, useLDSTr, TLDS):
|
|
n_mfma = 4 * 4 * 3 # 128 MT0 / 2 WT0 / 16 mfma dim * 128/2/16 * 3 bf16 MFMAs per tf32 mfma
|
|
|
|
optSchedule = dict()
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
syncs = SyncSchedule()
|
|
syncCode = []
|
|
snops: list[tuple[int, SNop]] = []
|
|
snopCode = []
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
|
|
lra0 = [0,0,1,1]
|
|
syncs.add( 3, dscnt=2, comment="Wait for the first 2 LRA0 to complete before pack")
|
|
syncs.add( 5, dscnt=0, comment="Wait for the rest LRA0 to complete before pack")
|
|
pack_a0 = [3,3,4,4, 7,7, 8,8,9,9, 5,5,6,6, 7,7, 10,10,11,11]
|
|
pack_b0 = [12,12,13,13, 16,16, 17,17,18,18, 14,14,15,15, 16,16, 19,19,20,20]
|
|
|
|
lrb0 = [6,6,7,7]
|
|
syncs.add(11, dscnt=0, comment="Wait for LRB0 to complete before pack",
|
|
barrier=True, barrier_comment="Wait for all waves to finish LRs before GRs")
|
|
grinca = [0,1,2, 2,2,2, 2,4,5]
|
|
grincb = [6,8,9, 10,11,12, 13,14,15]
|
|
lrsa = [13]
|
|
lrsb = [14]
|
|
lwsa = [45]
|
|
lwsb = [45]
|
|
|
|
gra = [15,17, 18,19, 20,21, 25,26]
|
|
grb = [27,28, 31,33, 36,37, 39,40]
|
|
num_gr = (len(gra[1::2]) + len(grb[1::2]))
|
|
|
|
gr_wait = 23
|
|
v = count_items(gra[1::2]+grb[1::2], ev=gr_wait)
|
|
syncs.add(gr_wait, vlcnt=v, barrier=True, comment = "Wait for previous GRA&B")
|
|
|
|
lrb3 = [24,24,25,25]
|
|
syncs.add( 28, dscnt=2, comment="Wait for the first 2 LRB3 to complete")
|
|
syncs.add( 30, dscnt=0, comment="Wait for the rest LRB3 to complete")
|
|
pack_b3 = [28,28,29,29, 32,32, 33,33,34,34, 30,30,31,31, 32,32, 35,35,36,36]
|
|
|
|
lra3 = [36,36,37,37]
|
|
syncs.add(39, dscnt=2, comment="Wait for the first 2 LRA3 to complete")
|
|
syncs.add(41, dscnt=0, comment="Wait for the rest LRA3 to complete")
|
|
pack_a3 = [39,39,40,40, 43,43, 44,44,45,45, 41,41,42,42, 43,43, 46,46,47,47]
|
|
|
|
else:
|
|
return False, None
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'GRIncA': [grinca],
|
|
'GRIncB': [grincb],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0': [pack_a0],
|
|
'PackB0': [pack_b0],
|
|
'GRA': [gra],
|
|
'GRB': [grb],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'LRA3': [lra3],
|
|
'LRB3': [lrb3],
|
|
'PackB3': [pack_b3],
|
|
'PackA3': [pack_a3],
|
|
'LCC': [[n_mfma-2, n_mfma-1]],
|
|
}
|
|
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
if snops:
|
|
optSchedule['SNOP'] = [ [s[0] for s in snops] ]
|
|
snopCode = [s[1] for s in snops]
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
kernel["UsePLRPack"] = 1
|
|
opt1 = ScheduleInfo(1, n_mfma, optSchedule, syncCode, nglshift, nllshift, snopCode=snopCode)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 128, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[32, 32, 16, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x128x32_TF32_plr1(kernel, useLDSTr, TLDS):
|
|
n_mfma = 128//2//32 * 128//2//32 * 3 * 2 # 128 MT0 / 2 WT0 / 32 mfma dim * 128/2/32 * 3 bf16 MFMAs per tf32 mfma * 2 PLR=1
|
|
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
syncs = SyncSchedule()
|
|
gr_inc_step = 0
|
|
num_code_paths = 1
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
lra0 = [0,1,2,3]
|
|
lrb0 = [ 4,5,6,7]
|
|
# wait then read
|
|
syncs.add( 4, dscnt=2, comment="wait for the first 2 LRAs before packing")
|
|
syncs.add( 5, dscnt=1, comment="wait for the rest of LRAs before packing them")
|
|
pack_a0 = [ 4,4,4,4, 6,6, 7,7,7,7,
|
|
5,5,5,5, 6,6, 8,8,8,8]
|
|
# because of GR starting at 10, we need barrier at 9, will use that for sync too.
|
|
syncs.add( 9, dscnt=0, comment="wait for LRBs before the packing them",
|
|
barrier=True, barrier_comment="make sure all LRs are done before starting GR")
|
|
pack_b0= [ 9,9,9,9, 10,10, 11,11,11,11,
|
|
9,9,9,9, 10,10, 11,11,11,11]
|
|
|
|
grinca = [0,0,0,1,1,1,2,2,2]
|
|
grincb = [3,3,3,6,6,6,6,6,6]
|
|
lrsa = [10]
|
|
lrsb = [10]
|
|
|
|
gra = [ 10,10,11,11] # one index for two instructions
|
|
grb = [ 13,13,14,14] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
syncs.add( 12, vlcnt=8, barrier=True, comment="wait for the previous GRAs")
|
|
|
|
lra1 = [ 12,12,13,14] # twice on 12 since we are waiting for GRA anyway at 12
|
|
lrb1 = [ 15,16,16,17]
|
|
# wait then read
|
|
syncs.add( 15, dscnt=2, vlcnt=8, comment="wait for the first 2 LRAs before packing. Also wait for GRBs",
|
|
barrier=True, barrier_comment="make sure GRBs are done before starting LRBs" )
|
|
syncs.add( 17, dscnt=3, comment="wait for the rest of LRAs before packing them")
|
|
pack_a1 = [ 16,16,16,16, 20,20, 21,21,21,21,
|
|
17,17,17,17, 20,20, 21,21,21,21]
|
|
syncs.add( 18, dscnt=2, comment="wait for 2 LRBs before the packing them")
|
|
syncs.add( 19, dscnt=0, comment="wait for the rest of LRBs before the packing them")
|
|
pack_b1= [ 18,18,18,18, 20,20, 22,22,22,22,
|
|
19,19,19,19, 20,20, 22,22,23,23]
|
|
lwsa = [ 20] # use delay before mfma4x4x4
|
|
lwsb = [ 20]
|
|
|
|
elif isNN(kernel) and TLDS==1 and kernel["VectorWidthA"] == 2:
|
|
lra0 = [0,0,0,0,
|
|
1,1,1,1]
|
|
lrb0 = [ 3, 4,6,6]
|
|
# wait then read
|
|
syncs.add( 3, dscnt=4, comment="wait for the first 2x2 LRAs before packing")
|
|
syncs.add( 4, dscnt=1, comment="wait for the rest of LRAs")
|
|
pack_a0 = [ 3,3,4,4, # swap instructions, must come after LR and before other packs
|
|
4,5,5,5, 6,6, 7,7,7,7,
|
|
5,5,6,6, 6,6, 8,8,8,8]
|
|
# because of GR starting at 10, we need barrier at 9, will use that for sync too.
|
|
syncs.add( 9, dscnt=0, comment="wait for LRBs before the packing them",
|
|
barrier=True, barrier_comment="make sure all LRs are done before starting GR")
|
|
pack_b0= [ 10,10,10,10, 10,10, 11,11,11,11,
|
|
9,9,9,9, 10,10, 11,11,11,11]
|
|
grinca = [0,0,1, 1,2,2, 2,2,2]
|
|
grincb = [2,2,6, 7,7,7, 8,8,8]
|
|
lrsa = [10]
|
|
lrsb = [10]
|
|
|
|
num_code_paths = 2
|
|
gra = [ 9,9, 11,11]
|
|
gra2 = [ 10,10,11,11]
|
|
grb = [ 13, 14,14,17] # one index for two instructions
|
|
grb2 = [ 13, 15,15,17] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
syncs.add( 12, vlcnt=8, barrier=True, comment="wait for the previous GRAs")
|
|
|
|
lra1 = [ 12,12,12,12,
|
|
13,13,13,13]
|
|
syncs.add( 14, vlcnt=4+1, barrier=True, barrier_comment="make sure GRBs are done before starting LRBs" )
|
|
lrb1 = [ 14,15,16,16]
|
|
syncs.add( 15, dscnt=1, comment="wait for LRAs")
|
|
pack_a1 =[ 15,15,16,16, # swap instructions, must come after LR and before other packs
|
|
17,17,17,17, 20,20, 21,21,21,21,
|
|
18,18,18,18, 20,20, 21,21,21,21]
|
|
syncs.add( 19, dscnt=2, comment="wait for the first 2 LRBs before the packing them")
|
|
syncs.add( 20, dscnt=0, comment="wait for the rest of LRBs")
|
|
pack_b1= [ 19,19,19,19, 20,20, 22,22,22,22,
|
|
20,20,20,20, 20,20, 22,22,22,22]
|
|
lwsa = [ 20] # use delay before mfma4x4x4
|
|
lwsb = [ 20]
|
|
|
|
elif isNT(kernel) and useLDSTr and TLDS==0 and kernel["VectorWidthA"] == 2 and kernel["VectorWidthB"] == 2:
|
|
lra0 = [0,0,0,0,
|
|
1,1,1,1]
|
|
lrb0 = [ 3,3,4,4,
|
|
6,6,6,6]
|
|
# wait then read
|
|
syncs.add( 3, dscnt=4, comment="wait for the first 2x2 LRAs before packing")
|
|
syncs.add( 4, dscnt=2, comment="wait for the rest of LRAs")
|
|
pack_a0 = [ 3,3,4,4, # swap instructions, must come after LR and before other packs
|
|
4,5,5,5, 6,6, 7,7,7,7,
|
|
5,5,6,6, 6,6, 8,8,8,8]
|
|
# because of GR starting at 10, we need barrier at 9, will use that for sync too.
|
|
syncs.add( 9, dscnt=0, comment="wait for LRBs",
|
|
barrier=True, barrier_comment="make sure all LRs are done before starting GR")
|
|
pack_b0= [ 9,9, 9,9, # swap instructions, must come after LR and before other packs
|
|
10,10,10,10, 10,10, 11,11,11,11,
|
|
9,9,9,9, 10,10, 11,11,11,11]
|
|
grinca = [0,0,1, 1,2,2, 2,2,2]
|
|
grincb = [2,2,6, 7,7,7, 8,8,8]
|
|
lrsa = [10]
|
|
lrsb = [10]
|
|
|
|
num_code_paths = 2
|
|
gra = [ 9,9, 11,11]
|
|
gra2 = [ 10,10,11,11]
|
|
grb = [ 13, 14,14,17] # one index for two instructions
|
|
grb2 = [ 13, 15,15,17] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
syncs.add( 12, vlcnt=8, barrier=True, comment="wait for the previous GRAs")
|
|
|
|
lra1 = [ 12,12,12,12,
|
|
13,13,13,13]
|
|
syncs.add( 14, vlcnt=4+1, barrier=True, barrier_comment="make sure GRBs are done before starting LRBs" )
|
|
lrb1 = [ 14,14,15,15,
|
|
16,16,16,16]
|
|
syncs.add( 15, dscnt=2, comment="wait for LRAs")
|
|
pack_a1 =[ 15,15,16,16, # swap instructions, must come after LR and before other packs
|
|
17,17,17,17, 20,20, 21,21,21,21,
|
|
18,18,18,18, 20,20, 21,21,21,21]
|
|
syncs.add( 19, dscnt=0, comment="wait for LRBs")
|
|
pack_b1= [ 19,19,19,19, # swap instructions, must come after LR and before other packs
|
|
19,19,19,19, 20,20, 22,22,22,22,
|
|
20,20,20,20, 20,20, 22,22,22,22]
|
|
lwsa = [ 20] # use delay before mfma4x4x4
|
|
lwsb = [ 20]
|
|
|
|
else:
|
|
return False, None
|
|
|
|
final_gra = [duplicate_list_items(gra, 2, gr_inc_step)]
|
|
if num_code_paths == 2:
|
|
final_gra += [duplicate_list_items(gra2, 2, gr_inc_step)]
|
|
|
|
final_grb = [duplicate_list_items(grb, 2, gr_inc_step)]
|
|
if num_code_paths == 2:
|
|
final_grb += [duplicate_list_items(grb2, 2, gr_inc_step)]
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'GRIncA': [grinca],
|
|
'GRIncB': [grincb],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'GRA': final_gra,
|
|
'GRB': final_grb,
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'PackA0': [pack_a0],
|
|
'PackB0': [pack_b0],
|
|
'LRA1': [lra1],
|
|
'LRB1': [lrb1],
|
|
'PackB1': [pack_b1],
|
|
'PackA1': [pack_a1],
|
|
'LCC': [[n_mfma-1, n_mfma-1]],
|
|
}
|
|
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
opt1 = ScheduleInfo(num_code_paths, n_mfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 128, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x128x64_TF32(kernel, useLDSTr, TLDS):
|
|
n_mfma = 96
|
|
optSchedule = dict()
|
|
nglshift = nllshift = 0
|
|
|
|
syncs = SyncSchedule()
|
|
syncCode = []
|
|
gr_inc_step = 1
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
offset=[0,0,1,1, 8,8, 9, 9,10,10,
|
|
2,2,3,3, 8,8, 11,11,12,12,
|
|
4,4,5,5, 8,8, 13,13,14,14,
|
|
6,6,7,7, 8,8, 15,15,16,16]
|
|
|
|
lra0 = [ 0,1,2,3,4,5,6,7]
|
|
lrb0 = [ 8,9,10,11,12,13,14,15]
|
|
# wait then read
|
|
syncs.add( 10, dscnt=6, comment="wait for the first 2 LRAs before packing")
|
|
syncs.add( 14, dscnt=6, comment="wait for the rest of LRAs before packing them")
|
|
pack_a0= [ i+11 for i in offset] # last at 27
|
|
# because of GR starting at 22, we need barrier at 21, will use that for sync too.
|
|
syncs.add( 21, dscnt=0, comment="wait for LRBs before the packing them",
|
|
barrier=True, barrier_comment="barrier before GR")
|
|
pack_b0= [ i+28 for i in offset] # last at 44
|
|
|
|
grinca = [0,0,1,1,2,2,3,3,4]
|
|
grincb = [5,5,6,6,7,7,8,8,9]
|
|
lrsa = [45]
|
|
lrsb = [45]
|
|
lwsa = [72]
|
|
lwsb = [72]
|
|
|
|
gra = [ 22,25,29,33, 37,41,45,49] # one index for two instructions
|
|
grb = [ 53,57,61, 64,69,75,79,84] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
syncs.add( 48, vlcnt=7, barrier=True, comment="wait for the previous GRs")
|
|
|
|
lra1 =[[ 48,49,50,51,52,53,54,55]]
|
|
lrb1 = [ 56,57, 59,60,61,62,63,64]
|
|
syncs.add( 58, dscnt=6, comment="wait for the first two LRAs before packing")
|
|
syncs.add( 63, dscnt=6, comment="wait for the rest of LRAs before packing them")
|
|
pack_a1 = [ i+59 for i in offset] # last at 75
|
|
syncs.add( 76, dscnt=0, comment="wait for LRBs before the packing them")
|
|
pack_b1 =[ i+77 for i in offset] # last at 93
|
|
|
|
elif isNN(kernel) and TLDS==1 and kernel["VectorWidthA"] == 4:
|
|
offset=[0,0,1,1, 8,8, 9, 9,10,10,
|
|
2,2,3,3, 8,8, 11,11,12,12,
|
|
4,4,5,5, 8,8, 13,13,14,14,
|
|
6,6,7,7, 8,8, 15,15,16,16]
|
|
|
|
lra0 = [ 0,0,2,2,4,4,6,6]
|
|
lrb0 = [ 11,11,13,13,15,15,17,17]
|
|
# wait then read
|
|
syncs.add( 6, dscnt=2, comment="wait for the first 4 LRAs before swapping/packing")
|
|
syncs.add( 9, dscnt=0, comment="wait for the rest of LRAs before swapping/packing them")
|
|
pack_a0= [ 7,7,7, 9,9,9, 8,8, 10,10, 8, 10] # swap instructions
|
|
pack_a0+=[ i+11 for i in offset] # last at 27
|
|
# because of GR starting at 22, we need barrier at 21, will use that for sync too.
|
|
syncs.add( 21, dscnt=0, comment="wait for LRBs before the packing them",
|
|
barrier=True, barrier_comment="barrier before GR")
|
|
pack_b0= [ i+28 for i in offset] # last at 44
|
|
|
|
grinca = [0,1,1,1,2,3,3,3,4]
|
|
grincb = [5,5,5,7,8,9,10,19,19]
|
|
lrsa = [45]
|
|
lrsb = [45]
|
|
lwsa = [72]
|
|
lwsb = [72]
|
|
|
|
gra = [ 22,25,29,33, 37,41,45,49] # one index for two instructions
|
|
grb = [ 53,57,61, 64,69,75,79,84] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
syncs.add( 48, vlcnt=7, barrier=True, comment="wait for the previous GRs")
|
|
lra1 =[[ 48,48,50,50,52,52,54,54],
|
|
[49,49,51,51,53,53,54,54]]
|
|
lrb1 = [ 59,59, 61,61,63,63,65,65]
|
|
syncs.add( 54, dscnt=2, comment="wait for the first two LRAs before packing")
|
|
syncs.add( 57, dscnt=0, comment="wait for the rest of LRAs before packing them")
|
|
pack_a1 =[ 55,55,55, 57,57,57, 55,55, 57,57, 56, 58] # swap instructions
|
|
pack_a1+= [ i+59 for i in offset] # last at 75
|
|
syncs.add( 76, dscnt=0, comment="wait for LRBs before the packing them")
|
|
pack_b1 =[ i+77 for i in offset] # last at 93
|
|
|
|
else:
|
|
return False, None
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'GRIncA': [grinca],
|
|
'GRIncB': [grincb],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'GRA': [duplicate_list_items(gra, 2, gr_inc_step),
|
|
duplicate_list_items([i+1 for i in gra], 2, gr_inc_step)],
|
|
'GRB': [duplicate_list_items(grb, 2, gr_inc_step),
|
|
duplicate_list_items([i+1 for i in grb], 2, gr_inc_step)],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'PackA0': [pack_a0],
|
|
'PackB0': [pack_b0],
|
|
'LRA1': lra1,
|
|
'LRB1': [lrb1],
|
|
'PackB1': [pack_b1],
|
|
'PackA1': [pack_a1],
|
|
'LCC': [[n_mfma-2, n_mfma-2]],
|
|
}
|
|
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
kernel["UsePLRPack"] = 1
|
|
opt1 = ScheduleInfo(2, n_mfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 256, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x256x32_TF32(kernel, useLDSTr, TLDS):
|
|
numMfma = 96
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
mfmaReorder = []
|
|
nglshift = nllshift = 0
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UsePLRPack"] = 1
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
|
|
# LRA0 + GRIncA
|
|
lra0 = create_range(min_val = 0, num = 4, step = 1, repeat = 1)
|
|
grIncA = create_range(min_val = max(lra0)+1, num = 3, step = 1, repeat = 3)
|
|
|
|
waitLRA0 = max(grIncA)+2
|
|
startPACKA0 = waitLRA0
|
|
|
|
packAOffset = [
|
|
0, 0, 1, 1,
|
|
4, 4,
|
|
5, 5, 6, 6,
|
|
|
|
2, 2, 3, 3,
|
|
4, 4,
|
|
7, 7, 8, 8,
|
|
]
|
|
packA0 = [x + startPACKA0 for x in packAOffset]
|
|
packA0Done = max(packA0)
|
|
|
|
# Sanity check
|
|
assert packA0Done < numMfma//4 , f"packA0Done {packA0Done} >= {numMfma//4}"
|
|
|
|
# 1st part of LRB0 + GRIncB
|
|
lrb0 = create_range(min_val = max(packA0)+1, num = 6, step = 1, repeat = 1)
|
|
grIncB = create_range(min_val = max(packA0)+1, num = 4, step = 1, repeat= 2)
|
|
grIncB += [max(grIncB)+1]
|
|
|
|
# GRA + 2nd part of LBR0
|
|
grA = create_range(min_val = max(lrb0)+1, num = 4, step = 2,repeat = 2)
|
|
lrb0 += create_range(min_val = max(lrb0)+4, num = 2, step = 1, repeat = 1)
|
|
|
|
waitLRB0 = max(lrb0)+2
|
|
startPACKB0 = waitLRB0
|
|
|
|
packBOffset = [
|
|
0, 0, 1, 1,
|
|
8, 8,
|
|
9, 9, 10, 10,
|
|
|
|
2, 2, 3, 3,
|
|
8, 8,
|
|
11, 11, 12, 12,
|
|
|
|
4, 4, 5, 5,
|
|
8, 8,
|
|
13, 13, 14, 14,
|
|
|
|
6, 6, 7, 7,
|
|
8, 8,
|
|
15, 15, 16, 16,
|
|
]
|
|
packB0 = [x + startPACKB0 for x in packBOffset]
|
|
|
|
halfMFMA = numMfma//2
|
|
assert max(packB0) < halfMFMA, f"max(packB0) {max(packB0)} >= halfMFMA {halfMFMA}"
|
|
|
|
# LR3
|
|
startLRB3 = halfMFMA
|
|
# Interleave GRA and LBR3
|
|
lrb3 = [create_range(min_val = startLRB3, num = 3, step = 2, repeat = 2),
|
|
create_range(min_val = startLRB3+1, num = 3, step = 2, repeat = 2)]
|
|
|
|
# Splitting LRB3 to avoid LDS Issue latency
|
|
lrb3[0]+= create_range(min_val = max(lrb3[0])+5, num = 1, step = 2, repeat = 2)
|
|
lrb3[1]+= create_range(min_val = max(lrb3[1])+5, num = 1, step = 2, repeat = 2)
|
|
|
|
grB = [create_range(min_val = startLRB3+1,num = 4,step = 2, repeat = 2),
|
|
create_range(min_val = startLRB3,num = 4,step = 2, repeat = 2)]
|
|
|
|
waitLRB3 = max(lrb3[1])+2
|
|
|
|
# Use different PackBOffset to shift last 5 CVTs iterations after GRB/LRA3
|
|
packB3Offset = [
|
|
0, 0, 1, 1,
|
|
8, 8,
|
|
9, 9, 10, 10,
|
|
|
|
2, 2, 3, 3,
|
|
8, 8,
|
|
11, 11, 19, 19,
|
|
|
|
4, 4, 5, 5,
|
|
8, 8,
|
|
20, 20, 21, 21,
|
|
|
|
6, 6, 7, 7,
|
|
8, 8,
|
|
22, 22, 23, 23,
|
|
]
|
|
|
|
# PackB3
|
|
packB3 = [x + waitLRB3 for x in packB3Offset]
|
|
|
|
startLRA3 = (3*numMfma)//4
|
|
# GRB + LRA3 (interleaved)
|
|
grB[0] += create_range(min_val = startLRA3,num = 4,step = 2, repeat = 2)
|
|
grB[1] += create_range(min_val = startLRA3+1,num = 4,step = 2, repeat = 2)
|
|
|
|
lra3 = [create_range(min_val = startLRA3+1,num=4,step=2,repeat=1),
|
|
create_range(min_val = startLRA3,num=4,step=2,repeat=1)]
|
|
|
|
waitLRA3 = max(lra3[0]) + 6
|
|
packA3 = [x + waitLRA3 for x in packAOffset]
|
|
|
|
syncTable = [
|
|
-1, SBarrier(comment="Sync codepath"),
|
|
waitLRA0, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for 1st LRA0 to complete"),
|
|
waitLRA0+1, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for 2nd LRA0 to complete"),
|
|
waitLRA0+2, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for 3rd LRA0 to complete"),
|
|
waitLRA0+3, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRA0 to complete"),
|
|
min(grIncB), SBarrier(comment="Barrier before GRA"),
|
|
|
|
waitLRB0, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="Wait for 1/8 LRB0 to complete"),
|
|
waitLRB0+1, SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for 2/8 LRB0 to complete"),
|
|
waitLRB0+2, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for 3/8 LRB0 to complete"),
|
|
waitLRB0+3, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for 4/8 LRB0 to complete"),
|
|
waitLRB0+4, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for 5/8 LRB0 to complete"),
|
|
waitLRB0+5, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for 6/8 LRB0 to complete"),
|
|
waitLRB0+6, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for 7/8 LRB0 to complete"),
|
|
waitLRB0+7, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for 8/8 LRB0 to complete"),
|
|
|
|
startLRB3-1, SWaitCnt(dscnt=-1, vlcnt=4, vscnt=-1, comment="Wait for previous GRA&B"),
|
|
startLRB3-1, SBarrier(comment="Barrier before GRB and before LRBA3/LBRB3"),
|
|
|
|
waitLRB3, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="Wait for 1/8 LRB3 to complete"),
|
|
waitLRB3+1, SWaitCnt(dscnt=6, vlcnt=-1, vscnt=-1, comment="Wait for 2/8 LRB3 to complete"),
|
|
waitLRB3+2, SWaitCnt(dscnt=5, vlcnt=-1, vscnt=-1, comment="Wait for 3/8 LRB3 to complete"),
|
|
waitLRB3+3, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for 4/8 LRB3 to complete"),
|
|
waitLRB3+4, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for 5/8 LRB3 to complete"),
|
|
waitLRB3+5, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for 6/8 LRB3 to complete"),
|
|
waitLRB3+6, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for 7/8 LRB3 to complete"),
|
|
waitLRB3+7, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for 8/8 LRB3 to complete"),
|
|
|
|
waitLRA3, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for 1/4 LRA3 to complete"),
|
|
waitLRA3+1, SWaitCnt(dscnt=2, vlcnt=-1, vscnt=-1, comment="Wait for 2/4 LRA3 to complete"),
|
|
waitLRA3+2, SWaitCnt(dscnt=1, vlcnt=-1, vscnt=-1, comment="Wait for 3/4 LRA3 to complete"),
|
|
waitLRA3+3, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for 4/4 LRA3 to complete"),
|
|
]
|
|
|
|
syncCode = syncTable[1::2]
|
|
optSchedule = {
|
|
'SYNC': [syncTable[::2]],
|
|
|
|
'GRIncA': [grIncA],
|
|
'GRIncB': [grIncB],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0' : [packA0],
|
|
'PackB0' : [packB0],
|
|
|
|
'GRA': [grA],
|
|
'GRB': [*grB],
|
|
'LRSA': [[max(lrb0)+1]],
|
|
'LRSB': [[max(lrb0)+1]],
|
|
'LWSA': [[numMfma-2]],
|
|
'LWSB': [[numMfma-2]],
|
|
'LCC': [[numMfma-1, numMfma-1]],
|
|
'LRA3': [*lra3],
|
|
'LRB3': [*lrb3],
|
|
'PackB3' : [packB3],
|
|
'PackA3' : [packA3],
|
|
}
|
|
nglshift = nllshift = 12 # vmcnt shift for ngl and nll
|
|
elif isNN(kernel) and TLDS==1:
|
|
return False, None
|
|
# kernel["UsePLRPack"] = 1
|
|
# kernel["UseMFMAF32XEmulation"] = True
|
|
# kernel["UseDot2F32XEmulation"] = False
|
|
|
|
# numLrReadA = 16
|
|
# numLrReadB = 8
|
|
# # mfma Reordering
|
|
# mfmaReorder = [i for i in range(0, numMfma//4)] + [i for i in range(numMfma//2, 3*numMfma//4)] + [i for i in range(numMfma//4, numMfma//2)] + [i for i in range(3*numMfma//4, numMfma)]
|
|
|
|
# # LBR0
|
|
# lrb0 = create_range(min_val = 0, num = numLrReadB//2, step = 1, repeat = 2)
|
|
# grIncB = create_range(min_val = 0, num = 3, step = 1, repeat = 3)
|
|
# grIncA = create_range(min_val = max(grIncB)+1, num = 3, step = 1, repeat = 3)
|
|
# waitLRB0 = max(lrb0)
|
|
# # PackB0 using mfma4x4x4_16b
|
|
# startPACKB0 = waitLRB0 + 4
|
|
# packBOffset = [
|
|
# 0, 0, 1, 1,
|
|
# 8, 8,
|
|
# 9, 9, 10, 10,
|
|
|
|
# 2, 2, 3, 3,
|
|
# 8, 8,
|
|
# 11, 11, 12, 12,
|
|
|
|
# 4, 4, 5, 5,
|
|
# 8, 8,
|
|
# 13, 13, 14, 14,
|
|
|
|
# 6, 6, 7, 7,
|
|
# 8, 8,
|
|
# 15, 15, 16, 16,
|
|
# ]
|
|
# packB0 = [x + startPACKB0 for x in packBOffset]
|
|
# packB0Done = max(packB0)
|
|
# assert packB0Done < numMfma//4
|
|
|
|
# # LRA0
|
|
# lra0 = create_range(min_val = waitLRB0+4, num = numLrReadA, step = 1, repeat = 1)
|
|
# grA = create_range(min_val = packB0Done+6, num = 8, step = 2, repeat = 1)
|
|
# waitLRA0 = max(lra0)
|
|
# startPACKA0 = waitLRA0 + 2
|
|
# packAOffset = [
|
|
# 0, 0, 1, 1,
|
|
# 4, 4,
|
|
# 5, 5, 6, 6,
|
|
|
|
# 2, 2, 3, 3,
|
|
# 4, 4,
|
|
# 7, 7, 8, 8,
|
|
# ]
|
|
# packA0 = [x + startPACKA0 for x in packAOffset]
|
|
# halfMFMA = numMfma//2
|
|
# assert max(packA0) < halfMFMA
|
|
|
|
# # LRA3
|
|
# startLRA3 = halfMFMA
|
|
# lra3 = create_range(min_val = startLRA3, num = numLrReadA, step = 1, repeat = 1)
|
|
# grB = create_range(min_val = startLRA3+1, num = 4, step = 2, repeat = 2)
|
|
# grB += create_range(min_val = max(lra3)+1, num = 4, step = 2, repeat = 2)
|
|
# waitLRA3 = max(lra3)
|
|
# startPACKA3 = waitLRA3 + 4
|
|
|
|
# # LRB3
|
|
# startLRB3 = (3*numMfma)//4 - 4 # Starts 4 indexes before 3/4 MFMAs to accommodate LRB3 latency
|
|
# lrb3 = create_range(min_val = startLRB3, num = numLrReadB//2, step = 1, repeat = 2)
|
|
# waitLRB3 = max(lrb3)
|
|
# startPACKB3 = waitLRB3 + 4
|
|
|
|
# # Grouping segment of 4x4x4_16B MFMAs together for PackB3 & PackA3 (reduce MFMA type switching cost)
|
|
# packB3 = [x + startPACKB3 for x in packBOffset]
|
|
# start_4x4x4 = packB3[4] # 5th index is start of 4x4x4_16B MFMA for PackB3
|
|
# packA3 = [
|
|
# *create_range(min_val = startPACKA3, num = 2, step = 1, repeat = 2),
|
|
# start_4x4x4, start_4x4x4,
|
|
# *create_range(min_val = max(packB3)+1, num = 2, step = 1, repeat = 2),
|
|
|
|
# *create_range(min_val = startPACKA3+2, num = 2, step = 1, repeat = 2),
|
|
# start_4x4x4, start_4x4x4,
|
|
# *create_range(min_val = max(packB3)+3, num = 2, step = 1, repeat = 2),
|
|
# ]
|
|
|
|
# syncTable = [
|
|
# waitLRB0, SWaitCnt(dscnt=min(15,numLrReadA-4), vlcnt=-1, vscnt=-1, comment="Wait for 4 LRA0 to complete"),
|
|
# waitLRB0+4, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 to complete"),
|
|
# waitLRB0+4, SBarrier(comment=""),
|
|
|
|
# waitLRA0, SWaitCnt(dscnt=4, vlcnt=-1, vscnt=-1, comment="Wait for 4/8 LRB0 to complete"),
|
|
# waitLRA0+4, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRB0 to complete"),
|
|
# waitLRA0+4, SBarrier(comment=""),
|
|
|
|
# startLRA3-1, SWaitCnt(dscnt=-1, vlcnt=6, vscnt=-1, comment="Wait for previous GRA&GRB"),
|
|
# startLRA3-1, SBarrier(comment=""),
|
|
|
|
# waitLRA3, SWaitCnt(dscnt=(numLrReadB-1), vlcnt=-1, vscnt=-1, comment="Wait for 1st LRB3 to complete"),
|
|
# waitLRA3+4, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRB3 to complete"),
|
|
|
|
# startLRB3-1, SWaitCnt(dscnt=-1, vlcnt=6, vscnt=-1, comment="Wait for previous GRA&GRB"),
|
|
# startLRB3-1, SBarrier(comment=""),
|
|
|
|
# waitLRB3, SWaitCnt(dscnt=min(15,numLrReadA-4), vlcnt=-1, vscnt=-1, comment="Wait for 4 LRA3 to complete"),
|
|
# waitLRB3+4, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA3 to complete"),
|
|
# ]
|
|
|
|
# optSchedule = {
|
|
# 'SYNC': [syncTable[::2]],
|
|
# 'GRIncA': [grIncA],
|
|
# 'GRIncB': [grIncB],
|
|
# 'LRA0': [lra0],
|
|
# 'LRB0': [lrb0],
|
|
# 'PackA0' : [packA0],
|
|
# 'PackB0' : [packB0],
|
|
# 'GRA': [grA],
|
|
# 'GRB': [grB],
|
|
# 'LRSA': [[max(lra0)+1]],
|
|
# 'LRSB': [[max(lra0)+1]],
|
|
# 'LWSA': [[max(lrb3)+1]],
|
|
# 'LWSB': [[max(lrb3)+1]],
|
|
# 'LCC': [[numMfma-1, numMfma-1]],
|
|
# 'LRA3': [lra3],
|
|
# 'LRB3': [lrb3],
|
|
# 'PackA3' : [packA3],
|
|
# 'PackB3' : [packB3],
|
|
# }
|
|
# syncCode = syncTable[1::2]
|
|
# nglshift = nllshift = 12
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift, mfmaReorder=mfmaReorder)
|
|
return True, opt1
|
|
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 160, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x160x64_TF32(kernel, useLDSTr, TLDS):
|
|
n_mfma = 120
|
|
optSchedule = dict()
|
|
nglshift = nllshift = 0
|
|
|
|
syncs = SyncSchedule()
|
|
syncCode = []
|
|
gr_inc_step = 0
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UsePLRPack"] = 1
|
|
|
|
grinca = [0,0,1,1,2,2,3,3,4]
|
|
grincb = [4,5,5,6,6,7,7,8,8]
|
|
lrsa = [58]
|
|
lrsb = [59]
|
|
lwsa = [118]
|
|
lwsb = [118]
|
|
|
|
pack_a = [0,0,1,1, 8,8, 9,9,10,10,
|
|
2,2,3,3, 8,8, 11,11,12,12,
|
|
4,4,5,5, 8,8, 13,13,14,14,
|
|
6,6,7,7, 8,8, 15,15,16,16
|
|
]
|
|
pack_b = [0,0,1,1, 10,10, 11,11,12,12,
|
|
2,2,3,3, 10,10, 13,13,14,14,
|
|
4,4,5,5, 10,10, 15,15,16,16,
|
|
6,6,7,7, 10,10, 17,17,18,18,
|
|
8,8,9,9, 10,10, 19,19,20,20
|
|
]
|
|
lra0 = [0,1,2,3,4,5,6,7]
|
|
syncs.add( 12, dscnt=4, barrier=True, comment="wait for LRA0 before pack to complete + barrier for GRA")
|
|
pack_a0 = [ i+13 for i in pack_a] ## last element = 13 + 16 = 29
|
|
|
|
lrb0 = [ 8,9,10,11, 13,14,15,16, 18,19]
|
|
syncs.add( 24, dscnt=0, comment="wait for LRB0 before pack to complete")
|
|
pack_b0 = [ i+30 for i in pack_b] ## last element = 30 + 20 = 50
|
|
|
|
gra = [ 17,22,27,32, 42,47,52,57] # one index for two instructions
|
|
grb = [ 67,71,75,79, 89,93,97,101, 112,116] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
syncs.add( 59, vlcnt=8, barrier=True, comment="wait for previous set of global reads + barrier for GRB")
|
|
|
|
lra1 = [60,61,62,63,64,65,66,67]
|
|
syncs.add( 72, dscnt=4, comment="wait for LRA1 before pack to complete")
|
|
pack_a1 = [ i+73 for i in pack_a] ## last element = 73 + 16 = 89
|
|
|
|
lrb1 = [ 68,69,70,71, 73,74,75,76, 78,79]
|
|
syncs.add( 85, dscnt=0, comment="wait for LRB1 before pack to complete")
|
|
pack_b1 = [ i+90 for i in pack_b] ## last element = 90 + 20 = 110
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'GRIncA': [grinca],
|
|
'GRIncB': [grincb],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0': [pack_a0],
|
|
'PackB0': [pack_b0],
|
|
'GRA': [duplicate_list_items(gra, 2, gr_inc_step),
|
|
duplicate_list_items([x+1 for x in gra], 2, gr_inc_step)],
|
|
'GRB': [duplicate_list_items(grb, 2, gr_inc_step),
|
|
duplicate_list_items([x+1 for x in grb], 2, gr_inc_step)],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'LRA1': [lra1],
|
|
'LRB1': [lrb1],
|
|
'PackB1': [pack_b1],
|
|
'PackA1': [pack_a1],
|
|
'LCC': [[n_mfma-2, n_mfma-1]],
|
|
}
|
|
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(2, n_mfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(256, 128, 32, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_256x128x32_TF32(kernel, useLDSTr, TLDS):
|
|
numMfma = 96
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
|
|
if isTN(kernel) and useLDSTr and TLDS==1:
|
|
kernel["UseMFMAF32XEmulation"] = False
|
|
kernel["UseDot2F32XEmulation"] = False
|
|
kernel["UsePLRPack"] = 1
|
|
numPackInstr = 24
|
|
numPackIndices = numPackInstr // 2 # Assign 2 pack instructions per mfma index
|
|
|
|
# LRA0 + PACKA0 - done before 1/4 MFMAs - index 24
|
|
lrA0 = [0,0, 1,1, 2,2, 3,3]
|
|
waitLRA0 = max(lrA0) + 2
|
|
startPACKA0 = waitLRA0
|
|
packA0 = create_range(startPACKA0, (len(lrA0)//2)*numPackIndices, numMfma//4-1)
|
|
|
|
# LBR0 + PACKB0 - done before 2/4 MFMAs - index 48
|
|
lrB0 = [7,7, 15,15]
|
|
waitLRB0 = max(lrB0) + 2
|
|
startPACKB0 = max(waitLRB0,max(packA0)) # Starts after waitLRB0 and packA0
|
|
packB0 = create_range(startPACKB0, (len(lrB0)//2)*numPackIndices, numMfma//2-1)
|
|
|
|
# LRB3 + PACKB3 - start after 2/4 MFMAs - index 48
|
|
halfMFMA = numMfma//2
|
|
startLRB3 = halfMFMA
|
|
lrB3 = create_range(startLRB3, 1, numMfma-1)
|
|
lrB3 += create_range(max(lrB3)+6, 1, numMfma-1)
|
|
waitLRB3 = startLRB3 + 4
|
|
packB3 = create_range(waitLRB3, (len(lrB3)//2)*numPackIndices, numMfma-1)
|
|
|
|
# LRA3 + PACKA3 - start after 3/4 MFMAs - index 72
|
|
startLRA3 = (3*numMfma)//4
|
|
lrA3 = create_range(startLRA3, 4, numMfma-1)
|
|
waitLRA3 = startLRA3 + 4
|
|
packA3 = create_range(waitLRA3, (len(lrA3)//2)*numPackIndices, numMfma-1)
|
|
|
|
syncTable = [
|
|
waitLRA0, SWaitCnt(dscnt=inflight(lrA0, waitLRA0)-2, vlcnt=-1, vscnt=-1, comment="wait for 1st 2 LRA0 to complete"),
|
|
waitLRA0+numPackIndices, SWaitCnt(dscnt=inflight(lrA0, waitLRA0+numPackIndices), vlcnt=-1, vscnt=-1, comment="wait for all LRA0 to complete"),
|
|
|
|
waitLRB0, SWaitCnt(dscnt=inflight(lrB0, waitLRB0)-2, vlcnt=-1, vscnt=-1, comment="wait for 1st 2 LRB0 to complete"),
|
|
waitLRB0+numPackIndices, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for all LRB0 to complete"),
|
|
waitLRB0+numPackIndices, SBarrier(comment="Barrier before GRA&GRB"),
|
|
|
|
startLRB3-1, SWaitCnt(dscnt=-1, vlcnt=6, vscnt=-1, comment="Wait for prev GRA&GRB"),
|
|
startLRB3-1, SBarrier(comment=""),
|
|
|
|
waitLRB3,SWaitCnt(dscnt=inflight(lrB3, waitLRB3)-2, vlcnt=-1, vscnt=-1, comment="Wait for 1st 2 LRB3 to complete"),
|
|
waitLRB3+numPackIndices,SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRB3 to complete"),
|
|
|
|
startLRA3, SWaitCnt(dscnt=-1, vlcnt=6, vscnt=-1, comment="Wait for prev GRA&GRB"),
|
|
startLRA3, SBarrier(comment=""),
|
|
|
|
waitLRA3, SWaitCnt(dscnt=inflight(lrA3,waitLRA3)-2, vlcnt=-1, vscnt=-1, comment="Wait for 1st 2 LRA3 to complete"),
|
|
waitLRA3+numPackIndices, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for all LRA3 to complete")
|
|
]
|
|
|
|
optSchedule = {
|
|
'SYNC' : [syncTable[::2]],
|
|
'GRIncA' : [[0,0,0, 1,1,1, 2,2,2]],
|
|
'GRIncB' : [[3,3,3, 4,4,4, 5,5,5]],
|
|
|
|
'LRA0' : [lrA0],
|
|
'PackA0' : [packA0],
|
|
'LRB0' : [lrB0],
|
|
'PackB0' : [packB0],
|
|
|
|
'GRA': [[48, 48, 50, 50, 52, 52, 54, 54, 66, 66, 68, 68, 70, 70, 72, 72]],
|
|
'GRB': [[30, 32, 34, 36, 40, 42, 44, 46]],
|
|
|
|
'LRA3' : [lrA3],
|
|
'PackA3' : [packA3],
|
|
'LRB3' : [lrB3],
|
|
'PackB3' : [packB3],
|
|
|
|
'LRSA': [[22]],
|
|
'LRSB': [[22]],
|
|
'LWSA': [[70]],
|
|
'LWSB': [[70]],
|
|
'LCC': [[95, 95]],
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 12 # vmcnt shift for ngl and nll
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(64, 128, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_64x128x64_TF32(kernel, useLDSTr, TLDS):
|
|
n_mfma = 48
|
|
optSchedule = dict()
|
|
nglshift = nllshift = 0
|
|
|
|
syncs = SyncSchedule()
|
|
syncCode = []
|
|
gr_inc_step = 0
|
|
|
|
if isTN(kernel) and not useLDSTr and TLDS==1:
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UsePLRPack"] = 1
|
|
|
|
grinca = [0,0,1,1,2,2,3,3,4]
|
|
grincb = [4,5,5,6,6,7,7,8,8]
|
|
lrsa = [23]
|
|
lrsb = [23]
|
|
lwsa = [47]
|
|
lwsb = [47]
|
|
|
|
pack_a_offset = [0,0,0,0, 2,2, 3,3,3,3,
|
|
1,1,1,1, 2,2, 4,4,4,4
|
|
]
|
|
pack_b_offset = [0,0,0,0, 4,4, 5,5,5,5,
|
|
1,1,1,1, 4,4, 6,6,6,6,
|
|
2,2,2,2, 4,4, 7,7,7,7,
|
|
3,3,3,3, 4,4, 8,8,8,8
|
|
]
|
|
lra0 = [0,1,2,3]
|
|
syncs.add( 8, dscnt=5, comment="wait for necessary LRA0 before pack to start")
|
|
syncs.add( 10, dscnt=4, barrier=True, comment="wait for remaining LRA0 before pack to complete + barrier for GRA")
|
|
pack_a0 = [ i+9 for i in pack_a_offset] ## last index = 9 + 4 = 13
|
|
|
|
lrb0 = [ 4,5, 7, 9,10,11,12,13]
|
|
syncs.add( 14, dscnt=4, comment="wait for necessary LRB0 before pack to start")
|
|
syncs.add( 16, dscnt=0, comment="wait for remaining LRB0 before pack to complete")
|
|
pack_b0 = [ i+14 for i in pack_b_offset] ## last index = 14 + 8 = 22
|
|
|
|
gra = [ 10,13,17,21] # one index for two instructions
|
|
grb = [ 24,28,32,35,38,41,43,45] # one index for two instructions
|
|
num_gr = len(gra) + len(grb)
|
|
|
|
syncs.add( 24, vlcnt=4, barrier=True, comment="wait for previous set of global reads + barrier for GRB")
|
|
|
|
lra1 = [24,25,26,27]
|
|
syncs.add( 32, dscnt=5, comment="wait for necessary LRA1 before pack to start")
|
|
syncs.add( 34, dscnt=4, comment="wait for remaining LRA1 before pack to complete")
|
|
pack_a1 = [ i+33 for i in pack_a_offset] ## last index = 33 + 4 = 37
|
|
|
|
lrb1 = [ 28,29, 31,32, 34,35,36,37]
|
|
syncs.add( 38, dscnt=4, comment="wait for necessary LRB1 before pack to start")
|
|
syncs.add( 40, dscnt=0, comment="wait for remaining LRB1 before pack to complete")
|
|
pack_b1 = [ i+38 for i in pack_b_offset] ## last index = 38 + 8 = 46
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
'GRIncA': [grinca],
|
|
'GRIncB': [grincb],
|
|
'LRA0': [lra0],
|
|
'LRB0': [lrb0],
|
|
'PackA0': [pack_a0],
|
|
'PackB0': [pack_b0],
|
|
'GRA': [duplicate_list_items(gra, 2, gr_inc_step)],
|
|
'GRB': [duplicate_list_items(grb, 2, gr_inc_step)],
|
|
'LRSA': [lrsa],
|
|
'LRSB': [lrsb],
|
|
'LWSA': [lwsa],
|
|
'LWSB': [lwsb],
|
|
'LRA1': [lra1],
|
|
'LRB1': [lrb1],
|
|
'PackB1': [pack_b1],
|
|
'PackA1': [pack_a1],
|
|
'LCC': [[n_mfma-1, n_mfma-1]],
|
|
}
|
|
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = num_gr
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
opt1 = ScheduleInfo(1, n_mfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|
|
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 64, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x64x64_TF32(kernel, useLDSTr, TLDS):
|
|
valid, opt = _get_schedule_64x128x64_TF32(kernel, useLDSTr, TLDS)
|
|
if not valid:
|
|
return False, None
|
|
|
|
optSchedule = switch_A_B_schedule(opt.optSchedule)
|
|
return True, ScheduleInfo(opt.numCodePaths, opt.numMfma, optSchedule, opt.syncCode, opt.nglshift, opt.nllshift)
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(160, 128, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=isTF32,
|
|
vector_widths=[4, 4, 4],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_160x128x64_TF32(kernel, useLDSTr, TLDS):
|
|
n_mfma = 120
|
|
optSchedule = dict()
|
|
nglshift = nllshift = 0
|
|
|
|
syncs = SyncSchedule()
|
|
syncCode = []
|
|
|
|
if isNN(kernel) and useLDSTr and TLDS==1:
|
|
kernel["UseMFMAF32XEmulation"] = True
|
|
kernel["UsePLRPack"] = 1
|
|
syncs.add(11, dscnt=8, comment="wait for LRB0 before pack to complete")
|
|
syncs.add(16, dscnt=8, barrier=True, comment="wait for LRB0 before pack to complete", barrier_comment="barrier for GRA")
|
|
syncs.add(33, dscnt=0, comment="wait for LRA0 before pack to complete")
|
|
syncs.add(59, vlcnt=8, barrier=True, comment="wait for previous set of global reads", barrier_comment="barrier for GRB")
|
|
syncs.add(73, dscnt=8, comment="wait for LRB1 before pack to complete")
|
|
syncs.add(78, dscnt=8, comment="wait for LRB1 before pack to complete")
|
|
syncs.add(94, dscnt=0, comment="wait for LRA1 before pack to complete")
|
|
|
|
optSchedule = {
|
|
'SYNC': [syncs.get_indicies()],
|
|
|
|
'GRIncB': [[0, 1, 2, 3, 4, 6, 7, 8, 9]],
|
|
'LRB0' : [[0,2,
|
|
4,5,
|
|
6,7,
|
|
9,11],
|
|
[0,1,
|
|
3,5,
|
|
6,8,
|
|
10,12]],
|
|
|
|
'GRIncA': [[10, 11, 12, 13, 14, 15, 15, 16, 16]],
|
|
|
|
'PackB0': [[13,13,14,14, 21,21, 22,22,23,23,
|
|
15,15,16,16, 21,21, 24,24,25,25,
|
|
17,17,18,18, 21,21, 26,26,27,27,
|
|
19,19,20,20, 21,21, 28,28,29,29]],
|
|
|
|
'LRA0' : [[0, 0, 2, 2, 4, 4, 5, 5,
|
|
7, 7, 8, 8, 10,10,12,12,
|
|
13,13,13,13,14,14,14,14,
|
|
16,16,16,16,18,18,18,18,
|
|
20,20,20,20,22,24,26,28],
|
|
[0, 0, 1, 1, 3, 3, 5, 5,
|
|
7, 7, 9, 9, 11,11,12,12,
|
|
13,13,13,13,15,15,15,15,
|
|
17,17,17,17,19,19,19,19,
|
|
21,21,21,21,23,23,23,23]],
|
|
|
|
'PackA0': [[30,30,31,31, 40,40, 41,41,42,42,
|
|
32,32,33,33, 40,40, 43,43,44,44,
|
|
34,34,35,35, 40,40, 45,45,46,46,
|
|
36,36,37,37, 40,40, 47,47,48,48,
|
|
38,38,39,39, 40,40, 49,49,50,50]],
|
|
|
|
'GRA' : [[60, 60, 62, 62, 64, 64, 66, 66, 68, 68,
|
|
81, 81, 83, 83, 85, 85, 87, 87, 90, 90]],
|
|
|
|
'GRB' : [[17, 17, 22, 22, 27, 27, 32, 32, 42, 42, 47, 47, 52, 52, 57, 57]],
|
|
|
|
'LRSA' : [[58]], 'LRSB' : [[59]],'LWSA' : [[118]], 'LWSB' : [[118]],
|
|
|
|
'LRB1' : [[59,61,
|
|
63,65,
|
|
67,69,
|
|
71,73],
|
|
[60,62,
|
|
64,66,
|
|
68,70,
|
|
72,74]],
|
|
|
|
'PackB1': [[73,73,74,74, 81,81, 82,82,83,83,
|
|
75,75,76,76, 81,81, 84,84,85,85,
|
|
77,77,78,78, 81,81, 86,86,87,87,
|
|
79,79,80,80, 81,81, 88,88,89,89]],
|
|
|
|
'LRA1' : [[59,59,61,61, 63,63,65,65,
|
|
67,67,69,69, 71,71,73,73,
|
|
74,74,74,74, 75,75,75,75,
|
|
77,77,77,77, 79,79,79,79,
|
|
81,81,81,81, 83,83,83,83],
|
|
[60,60,62,62, 64,64,65,65,
|
|
66,66,68,68, 70,70,72,72,
|
|
74,74,74,74, 76,76,76,76,
|
|
78,78,78,78, 80,80,80,80,
|
|
81,81,81,81, 82,82,82,82]],
|
|
|
|
'PackA1': [[90,90,91,91, 100,100, 101,101,102,102,
|
|
92,92,93,93, 100,100, 103,103,104,104,
|
|
94,94,95,95, 100,100, 105,105,106,106,
|
|
96,96,97,97, 100,100, 107,107,108,108,
|
|
98,98,99,99, 100,100, 109,109,110,110]],
|
|
|
|
'LCC': [[118, 119]]
|
|
}
|
|
|
|
syncCode = syncs.get_code()
|
|
nglshift = nllshift = len(optSchedule["GRA"][0])/2 + len(optSchedule["GRB"][0])/2
|
|
|
|
opt1 = ScheduleInfo(2, n_mfma, optSchedule, syncCode, nglshift, nllshift)
|
|
|
|
elif isTN(kernel) and not useLDSTr and TLDS==1:
|
|
valid, opt = _get_schedule_128x160x64_TF32(kernel, useLDSTr, TLDS)
|
|
if not valid:
|
|
return False, None
|
|
optSchedule = switch_A_B_schedule(opt.optSchedule)
|
|
opt1 = ScheduleInfo(opt.numCodePaths, opt.numMfma, optSchedule, opt.syncCode, opt.nglshift, opt.nllshift)
|
|
|
|
else:
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(128, 256, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_128x256x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 64
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0
|
|
if isNN(kernel) and useLDSTr and TLDS == 1:
|
|
lra0 = [create_range(min_val = 1, num = 4, step = 2, repeat = 2),
|
|
create_range(min_val = 0, num = 4, step = 2, repeat = 2)]
|
|
|
|
GRIncA = [create_range(min_val = 2, num = 3, step = 2, repeat = 3),
|
|
create_range(min_val = 1, num = 3, step = 2, repeat = 3)]
|
|
|
|
waitLRA0 = max(lra0[1])+5
|
|
gra = create_range(min_val = waitLRA0+1, num = 4, step = 2, repeat = 2)
|
|
lrb0 = create_range(min_val = max(gra)+1, num = 8, step = 1, repeat = 1)
|
|
GRIncB = create_range(min_val = max(gra)+1, num = 9, step = 1, repeat = 1)
|
|
|
|
assert max(lrb0) < numMfma // 2, "lrb0 max {} numMfma/2 {}".format(max(lrb0), numMfma//2)
|
|
|
|
startGRB = max(lrb0) + 5
|
|
|
|
assert startGRB < numMfma // 2, "startGRB {} numMfma/2 {}".format(startGRB, numMfma//2)
|
|
grb = create_range(min_val = startGRB, num = 4, step = 2, repeat = 2)
|
|
startLRA1 = max(grb) + 3
|
|
|
|
lra1 = create_range(min_val = startLRA1, num = 8, step = 1, repeat = 1)
|
|
startLRB1 = max(lra1) + 1
|
|
grb += create_range(min_val = startLRB1, num = 4, step = 2, repeat = 2)
|
|
lrb1 = create_range(min_val = startLRB1+1, num = 4, step = 2, repeat = 1)
|
|
lrb1 += create_range(min_val = max(lrb1)+2, num = 4, step = 1, repeat = 1)
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0 & LRB0"),
|
|
waitLRA0, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRA0"),
|
|
waitLRA0, SBarrier(comment=""),
|
|
|
|
startGRB-1, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="Wait for LRB0"),
|
|
startGRB-1, SBarrier(comment=""),
|
|
startLRA1-1, SWaitCnt(dscnt=-1, vlcnt=16, vscnt=-1, comment="wait for previous GRA & GRB"),
|
|
startLRA1-1, SBarrier(comment=""),
|
|
|
|
startLRB1-1, SWaitCnt(dscnt=-1, vlcnt=8, vscnt=-1, comment="wait for previous GRA & GRB"),
|
|
startLRB1-1, SBarrier(comment="")
|
|
]
|
|
|
|
optSchedule = {
|
|
'GRA': [gra],
|
|
'GRB': [grb],
|
|
'GRIncA': [*GRIncA],
|
|
'GRIncB': [GRIncB],
|
|
'LCC': [[numMfma-2,numMfma-2]],
|
|
'LRA0': [*lra0],
|
|
'LRA1': [lra1],
|
|
'LRB0': [lrb0],
|
|
'LRB1': [lrb1],
|
|
'LRSA': [[startGRB-1]],
|
|
'LRSB': [[startGRB-1]],
|
|
'LWSA': [[numMfma-3]],
|
|
'LWSB': [[numMfma-3]],
|
|
'SYNC': [syncTable[::2]],
|
|
}
|
|
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 12
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
else:
|
|
# No matching variant found
|
|
return False, None
|
|
|
|
kernel["MfmaInitCVgprs"] = True
|
|
return True, opt1
|
|
|
|
|
|
@RegisterSchedule(
|
|
tile_config=TileConfig(224, 320, 64, 2, 1, 1, False, 0, 0),
|
|
dtype_predicate=is16bit,
|
|
vector_widths=[8, 8, 8],
|
|
matrix_inst=[16, 16, 32, 1],
|
|
mfma_wave_group=[2, 2]
|
|
)
|
|
def _get_schedule_224x320x64_16bit(kernel, useLDSTr, TLDS):
|
|
numMfma = 140
|
|
optSchedule = dict()
|
|
syncCode = []
|
|
nglshift = nllshift = 0 # vmcnt shift for ngl and nll
|
|
kernel["MfmaInitCVgprs"] = True
|
|
kernel["SwapGlobalReadOrder"] = 0
|
|
|
|
if isTN(kernel) and useLDSTr and TLDS==1:
|
|
syncTable = [
|
|
-1, SWaitCnt(dscnt=9, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write old=0, new=9 newLW=0 newLR=9 for iteration == 0"),
|
|
6, SWaitCnt(dscnt=7, vlcnt=-1, vscnt=-1, comment="wait for prior local read local write"),
|
|
25, SWaitCnt(dscnt=0, vlcnt=-1, vscnt=-1, comment="wait for LR0 before DTL"),
|
|
25, SBarrier(comment=""),
|
|
78, SWaitCnt(dscnt=-1, vlcnt=10, vscnt=-1, comment="wait for prev iter GR before LRA1 and LRB1"),
|
|
78, SBarrier(comment=""),
|
|
]
|
|
optSchedule = {
|
|
'SYNC': [syncTable[::2]],
|
|
'GRIncA': [[0, 1, 2, 3, 4, 5, 6, 7, 8]], # 9
|
|
'GRIncB': [[9, 10, 11, 12, 13, 14, 15, 16, 17]], # 9
|
|
|
|
'LRA0': [[0, 3, 6, 9, 12, 15, 18]], # 7
|
|
'LRB0': [[1, 2, 4, 5, 7, 8, 10, 11, 13, 17]], # 10
|
|
|
|
'GRA': [[25,25, 30,30, 36,36, 42,42, 48,48, 54,54, 60,60]], # 14
|
|
'GRB': [[62,62, 67,67, 72,72, 77,77, 88,88, 94,94, 100,100, 106,106, 112,112, 118,118]], # 20
|
|
|
|
'LRA1': [[79, 81, 82, 83, 84, 85, 86]], # 7
|
|
'LRB1': [[120, 121, 122, 123, 124, 125, 126, 127, 128, 129]], # 10
|
|
|
|
'LRSA': [[66]], # 1
|
|
'LRSB': [[66]], # 1
|
|
'LWSA': [[118]], # 1
|
|
'LWSB': [[118]], # 1
|
|
'LCC': [[138, 138]], # 2
|
|
}
|
|
syncCode = syncTable[1::2]
|
|
nglshift = nllshift = 17
|
|
|
|
else:
|
|
return False, None
|
|
|
|
opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift)
|
|
return True, opt1
|