derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
1311 lines
61 KiB
Python
1311 lines
61 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.
|
|
################################################################################
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock
|
|
|
|
from Tensile.Components.CustomSchedule import hasCustomSchedule, ScheduleInfo
|
|
from Tensile.Components.CMSValidator import isValid, SchedulePosition, ValidatorPass
|
|
from Tensile.Common import IsaVersion
|
|
|
|
# Helper to create a mock data type
|
|
def _mock_dtype(is_16bit=False, is_8bit=False, num_bytes=4):
|
|
mock = MagicMock()
|
|
mock.isHalf.return_value = is_16bit
|
|
mock.isBFloat16.return_value = False # Assuming isHalf is enough for is16bit
|
|
mock.isInt8.return_value = is_8bit
|
|
mock.is8bitFloat.return_value = False # Assuming isInt8 is enough for is8bit
|
|
mock.numBytes.return_value = num_bytes
|
|
return mock
|
|
|
|
# Base kernel configuration factory
|
|
def create_base_kernel():
|
|
kernel = {
|
|
"UseCustomMainLoopSchedule": True,
|
|
"EnableMatrixInstruction": True,
|
|
"UnrollLoopSwapGlobalReadOrder": False,
|
|
"ISA": IsaVersion(9,5,0),
|
|
"WavefrontSize": 64,
|
|
"ProblemType": {
|
|
"DataType": _mock_dtype(),
|
|
"DataTypeA": _mock_dtype(),
|
|
"DataTypeB": _mock_dtype(),
|
|
"TransposeA": False,
|
|
"TransposeB": False,
|
|
"TLUA": True,
|
|
"TLUB": False,
|
|
},
|
|
"MacroTile0": 0, "MacroTile1": 0, "DepthU": 64,
|
|
"PrefetchGlobalRead": 0, "PrefetchLocalRead": 0, "DirectToLds": 1, "DtlPlusLdsBuf": False,
|
|
"GlobalReadVectorWidthA": 0, "GlobalReadVectorWidthB": 0,
|
|
"LocalReadVectorWidthA": 0, "LocalReadVectorWidthB": 0,
|
|
"WaveSeparateGlobalReadA": 0,
|
|
"WaveSeparateGlobalReadB": 0,
|
|
"Use64bShadowLimit" : 1,
|
|
"MatrixInstruction": [16,16,32,1],
|
|
"MIWaveGroup": [],
|
|
"LDSTrInst": False,
|
|
"TransposeLDS": 0,
|
|
"ForceUnrollSubIter": False,
|
|
"SwapGlobalReadOrder": False, # For asserting it gets set
|
|
"UsePLRPack": 0, # For asserting it gets set
|
|
"UseF32XEmulation": False,
|
|
"VectorWidthA": 1,
|
|
"VectorWidthB": 1,
|
|
"MIWaveTileA": 2,
|
|
"MIWaveTileB": 2,
|
|
}
|
|
return kernel
|
|
|
|
def update_kernel(kernel, updates):
|
|
"""Update kernel dict, auto-deriving TLUA/TLUB from TransposeA/TransposeB.
|
|
|
|
Args:
|
|
kernel: kernel dict to modify in-place
|
|
updates: dict mirroring the kernel structure. "ProblemType" key (if present)
|
|
is applied via kernel["ProblemType"].update(); all other keys are
|
|
applied via kernel.update(). If TransposeA or TransposeB appear in
|
|
the ProblemType sub-dict, TLUA and TLUB are auto-derived.
|
|
"""
|
|
if "ProblemType" in updates:
|
|
pt = updates.pop("ProblemType")
|
|
if "TransposeA" in pt or "TransposeB" in pt:
|
|
transA = pt.get("TransposeA", kernel["ProblemType"]["TransposeA"])
|
|
transB = pt.get("TransposeB", kernel["ProblemType"]["TransposeB"])
|
|
pt["TLUA"] = not transA
|
|
pt["TLUB"] = transB
|
|
kernel["ProblemType"].update(pt)
|
|
kernel.update(updates)
|
|
|
|
class TestCustomScheduleBF16:
|
|
@staticmethod
|
|
def get_num_mfma(kernel):
|
|
numMfma = (kernel["MIWaveTileA"] * kernel["MIWaveTileB"] *
|
|
kernel["DepthU"] / kernel["MatrixInstruction"][2] # two sub-iterations due to DepthU=64
|
|
)
|
|
return numMfma
|
|
|
|
def test_no_custom_schedule(self):
|
|
"""Test that a kernel that doesn't match any condition returns False."""
|
|
kernel = create_base_kernel()
|
|
# An empty kernel should not have a custom schedule
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert not has_schedule
|
|
assert schedule_info is None
|
|
|
|
@pytest.mark.parametrize("transA, transB", [(True, False), (False, True), (False, False), (True, True)])
|
|
def test_schedule_256x256x64_16bit(self, transA, transB):
|
|
"""Tests the 256x256x64 16-bit schedule."""
|
|
TN = transA and not transB
|
|
NT = not transA and transB
|
|
NN = not transA and not transB
|
|
TT = transA and transB
|
|
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2], "TransposeLDS": 0 if NT else 1, "MIWaveTileA": 8, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 128
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
if TN:
|
|
assert 'PackA0' not in schedule_info.optSchedule
|
|
assert not kernel["UsePLRPack"]
|
|
elif NT:
|
|
assert 'PackA0' in schedule_info.optSchedule
|
|
assert kernel["UsePLRPack"]
|
|
elif NN:
|
|
assert not kernel["SwapGlobalReadOrder"]
|
|
assert 'PackA0' in schedule_info.optSchedule
|
|
assert 'PackB0' not in schedule_info.optSchedule
|
|
assert kernel["UsePLRPack"]
|
|
elif TT:
|
|
assert kernel["SwapGlobalReadOrder"]
|
|
assert 'PackA0' not in schedule_info.optSchedule
|
|
assert 'PackB0' in schedule_info.optSchedule
|
|
assert kernel["UsePLRPack"]
|
|
|
|
@pytest.mark.parametrize("force_unroll_sub_iter", [True, False])
|
|
def test_schedule_256x256x128_8bit_TN(self, force_unroll_sub_iter: bool):
|
|
"""Tests the 256x256x128 8-bit TNschedule."""
|
|
|
|
kernel = create_base_kernel()
|
|
dtype_8bit = _mock_dtype(is_8bit=True, num_bytes=1)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_8bit, "DataTypeA": dtype_8bit, "DataTypeB": dtype_8bit,
|
|
"TransposeA": True, "TransposeB": False,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 256, "DepthU": 128,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 0,
|
|
"GlobalReadVectorWidthA": 16, "GlobalReadVectorWidthB": 16, "LocalReadVectorWidthA": 16, "LocalReadVectorWidthB": 16,
|
|
"MatrixInstruction": [16,16,128,1], "MIWaveGroup": [2,2], "TransposeLDS": 1, "MIWaveTileA": 8, "MIWaveTileB": 8, "ForceUnrollSubIter": force_unroll_sub_iter,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 1
|
|
assert schedule_info.numMfma == 64
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, dtl_plus_lds_buf", [
|
|
( True, False, True, 1, None),
|
|
( False, True, True, 0, 1),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_256x96x64_16bit(self, transA, transB, lds_tr_inst, tr_lds, dtl_plus_lds_buf):
|
|
"""Tests the 256x96x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 96, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds,
|
|
"MIWaveTileA": 8, "MIWaveTileB": 3,
|
|
})
|
|
|
|
if dtl_plus_lds_buf is not None:
|
|
kernel.update({"DtlPlusLdsBuf": dtl_plus_lds_buf })
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 48
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
# TN case supports both LDSTrInst=True and LDSTrInst=False
|
|
( True, False, False, 1),
|
|
( True, False, True, 1),
|
|
( False, True, True, 0),
|
|
( False, False, False, 1)
|
|
# fmt: on
|
|
])
|
|
def test_schedule_96x256x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 96x256x64 16-bit schedules."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 96, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds,
|
|
"MIWaveTileA": 3, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numMfma == 48
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize("transA, transB", [(False, False), (False, True), (True, False)])
|
|
def test_schedule_192x256x64_16bit(self, transA, transB):
|
|
"""Tests the 192x256x64 16-bit NN schedule."""
|
|
NN = not transA and not transB
|
|
NT = not transA and transB
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 192, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": NN, "TransposeLDS": 0 if NT else 1, "MIWaveTileA": 6, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 96
|
|
if NN:
|
|
assert kernel["SwapGlobalReadOrder"]
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, False, 1),
|
|
( False, False, True, 1),
|
|
( False, True, True, 0)
|
|
# fmt: on
|
|
])
|
|
def test_schedule_256x192x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
|
|
"""Tests the 256x192x64 16-bit TN schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 192, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 8, "MIWaveTileB": 6,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 96
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize("transA, transB", [(True, False), (False, False), (False, True)])
|
|
def test_schedule_160x256x64_16bit(self, transA, transB):
|
|
"""Tests the 160x256x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 160, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": not transA, "TransposeLDS": not transB, "MIWaveTileA": 5, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 80
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize("transA, transB", [(False, False), (False, True), (True, False)])
|
|
def test_schedule_256x160x64_16bit(self, transA, transB):
|
|
"""Tests the 256x160x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 160, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
# Match CustomSchedule.py predicates:
|
|
# - NN/NT useLDSTr=True, TN useLDSTr=False -> useLDSTr == (not TransposeA)
|
|
# - TLDS==1 for NN/TN, TLDS==0 for NT -> TLDS == (not TransposeB)
|
|
"LDSTrInst": not transA, "TransposeLDS": not transB, "MIWaveTileA": 8, "MIWaveTileB": 5,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 80
|
|
# SwapGlobalReadOrder is set for NN/NT branches, not required for TN.
|
|
if not (transA and (not transB)):
|
|
assert kernel["SwapGlobalReadOrder"]
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize("transA, transB", [(True, False), (False, True), (False, False)])
|
|
def test_schedule_256x240x64_16bit(self, transA, transB):
|
|
"""Tests the 256x240x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 240, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 2, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [4,1],
|
|
"LDSTrInst": True, "TransposeLDS": 1 if transA or not (transA or transB) else 0, "MIWaveTileA": 4, "MIWaveTileB": 15,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 1
|
|
assert schedule_info.numMfma == 120
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize("transA, transB", [(True, False), (False, False)])
|
|
def test_schedule_256x208x64_16bit(self, transA, transB):
|
|
"""Tests the 256x208x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 208, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 2, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [4,1],
|
|
"LDSTrInst": not transA , "TransposeLDS": 1, "MIWaveTileA": 4, "MIWaveTileB": 13,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 1
|
|
assert schedule_info.numMfma == 104
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, False, 1),
|
|
( False, True, True, 0),
|
|
( False, False, True, 1)
|
|
# fmt: on
|
|
])
|
|
def test_schedule_224x256x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 224x256x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 224, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 7, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 112
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( False, False, True, 1),
|
|
( True, False, False, 1),
|
|
( False, True, True, 0)
|
|
# fmt: on
|
|
])
|
|
def test_schedule_192x320x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 192x320x64 16-bit NN schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 192, "MacroTile1": 320, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 6, "MIWaveTileB": 10,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 1
|
|
assert schedule_info.numMfma == 120
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, mt0, mt1", [
|
|
( True, False, False, 1, 224, 128), # TN
|
|
( False, True, True, 0, 224, 128), # NT
|
|
( False, False, True, 1, 224, 128), # NN
|
|
( False, True, True, 0, 128, 224), # NT
|
|
])
|
|
def test_schedule_224x128x64_128x224x64_16bit(self, transA, transB, lds_tr_inst, tr_lds, mt0, mt1):
|
|
"""
|
|
Tests the 224x128x64 16-bit schedules (TN/NT/NN).
|
|
Tests the 128x224x64 16-bit schedule (NT).
|
|
"""
|
|
du = 64
|
|
mi = [16,16,32,1]
|
|
mi_wave_group = [2, 2]
|
|
mi_wave_tile = (mt0 // (mi[0] * mi_wave_group[0]), mt1 // (mi[1] * mi_wave_group[1]))
|
|
NT = int(not transA and transB)
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": mt0, "MacroTile1": mt1, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": mi, "MIWaveGroup": mi_wave_group,
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": mi_wave_tile[0], "MIWaveTileB": mi_wave_tile[1],
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == (2 if NT else 1)
|
|
assert schedule_info.numMfma == 56
|
|
assert bool(kernel.get("SwapGlobalReadOrder", 0)) == NT
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, False, 1),
|
|
( False, False, True, 1),
|
|
( False, True, True, 0),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_256x224x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 256x224x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 256, "MacroTile1": 224, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 8, "MIWaveTileB": 7,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 112
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( False, False, True, 1), #NN
|
|
( False, True, True, 0), #NT
|
|
( True, False, False, 1) #TN
|
|
# fmt: on
|
|
])
|
|
def test_schedule_320x192x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 320x192x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 320, "MacroTile1": 192, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 10, "MIWaveTileB": 6,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 120
|
|
assert kernel["SwapGlobalReadOrder"]
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, False, 1),
|
|
( False, True, True, 0),
|
|
( True, False, True, 1),
|
|
( False, True, False, 0),
|
|
( False, False, True, 1)
|
|
# fmt: on
|
|
])
|
|
def test_schedule_240x256x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 240x256x64 16-bit schedule."""
|
|
NT = not transA and transB
|
|
TN = transA and not transB
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 240, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 2, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [1,4],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 15, "MIWaveTileB": 4,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == (1 if NT else 2)
|
|
assert schedule_info.numMfma == 120
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, False, 1),
|
|
( False, False, True, 1),
|
|
( False, True, True, 0)
|
|
# fmt: on
|
|
])
|
|
def test_schedule_208x256x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 208x256x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 208, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 2, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [1,4],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 13, "MIWaveTileB": 4,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 1
|
|
assert schedule_info.numMfma == 104
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, False, 1),
|
|
( False, False, True, 1),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_128x224x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 208x256x64 16-bit schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 128, "MacroTile1": 224, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 4, "MIWaveTileB": 7,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 56
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, mt0, mt1, code_paths", [
|
|
( True, False, False, 1, 128, 192, 1),
|
|
( True, False, False, 1, 192, 128, 2),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_128x192x64_192x128x64_16bit(self, transA, transB, lds_tr_inst, tr_lds, mt0, mt1, code_paths):
|
|
"""Tests the 128x192x64 and 192x128x64 BF16 tiles."""
|
|
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
du = 64
|
|
mi = [16,16,32,1]
|
|
mi_wave_group = [2, 2]
|
|
mi_wave_tile = (mt0 // (mi[0] * mi_wave_group[0]), mt1 // (mi[1] * mi_wave_group[1]))
|
|
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": True, "TransposeB": False,
|
|
},
|
|
"MacroTile0": mt0, "MacroTile1": mt1, "DepthU": du,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": mi, "MIWaveGroup": mi_wave_group,
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": mi_wave_tile[0], "MIWaveTileB": mi_wave_tile[1],
|
|
})
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == code_paths
|
|
assert schedule_info.numMfma == TestCustomScheduleBF16.get_num_mfma(kernel)
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
def test_schedule_128x256x64_16bit(self):
|
|
"""Tests the 128x256x64 NN schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": False, "TransposeB": False,
|
|
},
|
|
"MacroTile0": 128, "MacroTile1": 256, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": True, "TransposeLDS": 1, "MIWaveTileA": 4, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 64
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, True, 1),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_352x192x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 352x192x64 16-bit TN schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 352, "MacroTile1": 192, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 11, "MIWaveTileB": 6,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == TestCustomScheduleBF16.get_num_mfma(kernel)
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds", [
|
|
( True, False, True, 1),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_224x320x64_16bit(self, transA, transB, lds_tr_inst, tr_lds):
|
|
"""Tests the 224x320x64 16-bit TN schedule."""
|
|
kernel = create_base_kernel()
|
|
dtype_16bit = _mock_dtype(is_16bit=True, num_bytes=2)
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"DataType": dtype_16bit, "DataTypeA": dtype_16bit, "DataTypeB": dtype_16bit,
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"MacroTile0": 224, "MacroTile1": 320, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 8, "GlobalReadVectorWidthB": 8, "LocalReadVectorWidthA": 8, "LocalReadVectorWidthB": 8,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 7, "MIWaveTileB": 10,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 140
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
class TestCustomScheduleTF32:
|
|
@staticmethod
|
|
def get_num_mfma(kernel):
|
|
numMfma = (kernel["MIWaveTileA"] * kernel["MIWaveTileB"] *
|
|
3 * # tf32 emulated with 3 bf16
|
|
kernel["DepthU"] / kernel["MatrixInstruction"][2] # two sub-iterations due to DepthU=64
|
|
)
|
|
return numMfma
|
|
|
|
@pytest.mark.parametrize("transA, transB, vwa", [
|
|
(True, False, 1),
|
|
(False, False, 1),
|
|
])
|
|
def test_schedule_192x256x32_TF32(self, transA, transB, vwa):
|
|
"""Tests the 192x256x32 TF32 schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": True,
|
|
"MacroTile0": 192, "MacroTile1": 256, "DepthU": 32,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": False, "TransposeLDS": 1, "MIWaveTileA": 6, "MIWaveTileB": 8,
|
|
})
|
|
if vwa is not None:
|
|
kernel.update({"VectorWidthA": vwa})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 144
|
|
assert kernel["UsePLRPack"]
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
def test_schedule_128x192x32_TF32(self):
|
|
"""Tests the 128x192x32 TF32 TN schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": True, "TransposeB": False,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": True,
|
|
"MacroTile0": 128, "MacroTile1": 192, "DepthU": 32,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": False, "TransposeLDS": 1, "MIWaveTileA": 4, "MIWaveTileB": 6,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 72
|
|
assert kernel["UsePLRPack"]
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
def test_schedule_192x128x32_TF32(self):
|
|
"""Tests the 192x128x32 TF32 TN schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": True, "TransposeB": False,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": True,
|
|
"MacroTile0": 192, "MacroTile1": 128, "DepthU": 32,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": True, "TransposeLDS": 1, "MIWaveTileA": 6, "MIWaveTileB": 4,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 72
|
|
assert kernel["UsePLRPack"]
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, vwa, vwb", [
|
|
( True, False, False, 1, None, None),
|
|
( False, True, False, 0, 4, 4),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_256x256x32_TF32(self,transA, transB, lds_tr_inst, tr_lds, vwa, vwb):
|
|
"""Tests the 256x256x32 TF32 TN schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": True,
|
|
"MacroTile0": 256, "MacroTile1": 256, "DepthU": 32,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 8, "MIWaveTileB": 8,
|
|
})
|
|
|
|
if vwa is not None:
|
|
kernel.update({"VectorWidthA": vwa})
|
|
if vwb is not None:
|
|
kernel.update({"VectorWidthB": vwb})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 192
|
|
assert kernel["UsePLRPack"]
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, vwa", [
|
|
( True, False, 1),
|
|
( False, False, 1),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_256x192x32_TF32(self, transA, transB, vwa):
|
|
"""Tests the 256x192x32 TF32 TN schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": True,
|
|
"MacroTile0": 256, "MacroTile1": 192, "DepthU": 32,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": False, "TransposeLDS": 1, "MIWaveTileA": 8, "MIWaveTileB": 6,
|
|
})
|
|
if vwa is not None:
|
|
kernel.update({"VectorWidthA": vwa})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 144
|
|
assert kernel["UsePLRPack"]
|
|
if transA == False and transB == False:
|
|
assert kernel["UseMFMAF32XEmulation"]
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
def test_schedule_128x256x32_TF32(self):
|
|
"""Tests the 128x256x32 TF32 TN schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": True, "TransposeB": False,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": True,
|
|
"MacroTile0": 128, "MacroTile1": 256, "DepthU": 32,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"DirectToLds": True,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16, 16, 32, 1], "MIWaveGroup": [2, 2],
|
|
"LDSTrInst": False, "TransposeLDS": 1, "MIWaveTileA": 4, "MIWaveTileB": 8,
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == 96
|
|
assert kernel["UsePLRPack"]
|
|
valid, message = isValid(schedule_info, {"kernel": kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, plr, vwa, vwb, mi , ncp", [
|
|
( True, False, False, 1, 1, None, None, [16,16,32,1], 1),
|
|
( True, False, False, 1, 1, None, None, [32,32,16,1], 1),
|
|
( False, False, False, 1, 1, 2, None, [32,32,16,1], 2),
|
|
( False, False, True, 1, 1, 2, None, [32,32,16,1], 2),
|
|
( False, True, True, 0, 1, 2, 2, [32,32,16,1], 2),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_128x128x32(self, transA, transB, lds_tr_inst, tr_lds, plr, vwa, vwb, mi, ncp):
|
|
"""Tests the 128x128x32 TF32 schedule."""
|
|
kernel = create_base_kernel()
|
|
macro_tile = (128,128,32)
|
|
mi_wave_group = (2,2)
|
|
mi_wave_tile = (macro_tile[0] // (mi[0] * mi_wave_group[0]), macro_tile[1] // (mi[1] * mi_wave_group[1]))
|
|
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"ForceUnrollSubIter": (macro_tile[2] == mi[2]), # production sets True when DU == matrixInstK (Solution.py:1442)
|
|
"MacroTile0": macro_tile[0], "MacroTile1": macro_tile[1], "DepthU": macro_tile[2],
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": plr,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": mi, "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": mi_wave_tile[0], "MIWaveTileB": mi_wave_tile[1],
|
|
})
|
|
if vwa is not None:
|
|
kernel.update({"VectorWidthA": vwa})
|
|
if vwb is not None:
|
|
kernel.update({"VectorWidthB": vwb})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == ncp
|
|
numMfma = (mi_wave_tile[0] * mi_wave_tile[1] *
|
|
3 * # tf32 emulated with 3 bf16
|
|
(1 if mi[0] == 16 else 2) # two sub-iterations with mi32
|
|
)
|
|
assert schedule_info.numMfma == numMfma
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, vwa", [
|
|
( True, False, False, 1, None),
|
|
( False, False, True, 1, 4), # NN doesn't depend on lds_tr_inst, so check for both values
|
|
( False, False, False, 1, 4),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_128x128x64(self, transA, transB, lds_tr_inst, tr_lds, vwa):
|
|
"""Tests the 128x128x64 TF32 TN schedule."""
|
|
kernel = create_base_kernel()
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"MacroTile0": 128, "MacroTile1": 128, "DepthU": 64,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": [16,16,32,1], "MIWaveGroup": [2,2],
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": 4, "MIWaveTileB": 4,
|
|
})
|
|
if vwa is not None:
|
|
kernel.update({"VectorWidthA": vwa})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
numMfma = ((kernel["MacroTile0"] // kernel["MIWaveGroup"][0] // kernel["MatrixInstruction"][0]) *
|
|
(kernel["MacroTile1"] // kernel["MIWaveGroup"][1] // kernel["MatrixInstruction"][1]) *
|
|
3 * # tf32 emulated with 3 bf16
|
|
2 # two sub-iterations due to DepthU=64
|
|
)
|
|
assert schedule_info.numMfma == numMfma
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, mt0, mt1", [
|
|
( True, False, False, 1, 128, 160),
|
|
( False, False, True, 1, 160, 128),
|
|
( True, False, False, 1, 160, 128),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_128x160x64_160x128x64(self, transA, transB, lds_tr_inst, tr_lds, mt0, mt1):
|
|
"""Tests the 128x160x64, 160x128x64 TF32 TN schedule and 160x128x64 TF32 NN."""
|
|
|
|
kernel = create_base_kernel()
|
|
|
|
du = 64
|
|
mi = [16,16,32,1]
|
|
mi_wave_group = [2, 2]
|
|
mi_wave_tile = (mt0 // (mi[0] * mi_wave_group[0]), mt1 // (mi[1] * mi_wave_group[1]))
|
|
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"MacroTile0": mt0, "MacroTile1": mt1, "DepthU": du,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": mi, "MIWaveGroup": mi_wave_group,
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": mi_wave_tile[0], "MIWaveTileB": mi_wave_tile[1],
|
|
})
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 2
|
|
assert schedule_info.numMfma == TestCustomScheduleTF32.get_num_mfma(kernel)
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
@pytest.mark.parametrize(
|
|
# fmt: off
|
|
"transA, transB, lds_tr_inst, tr_lds, mt0, mt1", [
|
|
( True, False, False, 1, 64, 128),
|
|
( True, False, False, 1, 128, 64),
|
|
# fmt: on
|
|
])
|
|
def test_schedule_64x128x64_128x64x64(self, transA, transB, lds_tr_inst, tr_lds, mt0, mt1):
|
|
"""Tests the 64x128x64 & 128x64x64 TF32 TN schedules."""
|
|
kernel = create_base_kernel()
|
|
du = 64
|
|
mi = [16,16,32,1]
|
|
mi_wave_group = [2, 2]
|
|
mi_wave_tile = (mt0 // (mi[0] * mi_wave_group[0]), mt1 // (mi[1] * mi_wave_group[1]))
|
|
|
|
update_kernel(kernel, {
|
|
"ProblemType": {
|
|
"TransposeA": transA, "TransposeB": transB,
|
|
},
|
|
"UseF32XEmulation": True, "UseDirect32XEmulation": True,
|
|
"MacroTile0": mt0, "MacroTile1": mt1, "DepthU": du,
|
|
"PrefetchGlobalRead": 2, "PrefetchLocalRead": 1,
|
|
"GlobalReadVectorWidthA": 4, "GlobalReadVectorWidthB": 4, "LocalReadVectorWidthA": 4, "LocalReadVectorWidthB": 4,
|
|
"MatrixInstruction": mi, "MIWaveGroup": mi_wave_group,
|
|
"LDSTrInst": lds_tr_inst, "TransposeLDS": tr_lds, "MIWaveTileA": mi_wave_tile[0], "MIWaveTileB": mi_wave_tile[1],
|
|
})
|
|
|
|
has_schedule, schedule_info = hasCustomSchedule(kernel)
|
|
assert has_schedule
|
|
assert isinstance(schedule_info, ScheduleInfo)
|
|
assert schedule_info.numCodePaths == 1
|
|
numMfma = (mi_wave_tile[0] * mi_wave_tile[1] *
|
|
3 * # tf32 emulated with 3 bf16
|
|
2 # two sub-iterations due to DepthU=64
|
|
)
|
|
assert schedule_info.numMfma == numMfma
|
|
valid, message = isValid(schedule_info, {"kernel" : kernel})
|
|
assert valid, message
|
|
|
|
class TestCustomScheduleValidation:
|
|
def test_disable_single_pass(self):
|
|
"""Disabling a single pass allows an otherwise-invalid schedule to pass that check."""
|
|
kernel = create_base_kernel()
|
|
invalid_schedule = {"P": [[3, 2, 1]]}
|
|
|
|
# Without disabling, the non-ascending schedule fails on ascending order.
|
|
si = ScheduleInfo(1, None, invalid_schedule, None, None, None, None)
|
|
status, message = isValid(si, {"kernel": kernel})
|
|
assert status == False
|
|
assert "Non-descending-order" in message
|
|
|
|
# Disabling VERIFY_ASCENDING_ORDER skips that check.
|
|
# Remaining structural and timeline passes are also disabled because this
|
|
# minimal schedule lacks the keys/data they require.
|
|
si = ScheduleInfo(1, 0, invalid_schedule, None, None, None, None)
|
|
for p in ValidatorPass:
|
|
if p != ValidatorPass.VERIFY_CORRECT_NUMBER_OF_INSTRUCTIONS:
|
|
si.disableValidationPass(p, reason="Not relevant to this test")
|
|
status, message = isValid(si, {"kernel": kernel})
|
|
# The ascending-order error is gone; the schedule passes the remaining enabled check.
|
|
assert "Non-descending-order" not in message
|
|
assert status == True
|
|
|
|
def test_disable_validation_pass_reason_required(self):
|
|
"""disableValidationPass requires a non-empty reason string and a valid ValidatorPass enum member."""
|
|
si = ScheduleInfo(1, None, {}, None, None, None, None)
|
|
|
|
with pytest.raises(ValueError):
|
|
si.disableValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER, reason="")
|
|
|
|
with pytest.raises(ValueError):
|
|
si.disableValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER, reason=" ")
|
|
|
|
with pytest.raises(TypeError):
|
|
si.disableValidationPass("not_an_enum", reason="some reason")
|
|
|
|
with pytest.raises(TypeError):
|
|
si.disableValidationPass(42, reason="some reason")
|
|
|
|
def test_disable_multiple_validation_passes(self):
|
|
"""Multiple validation passes can be disabled independently."""
|
|
invalid_schedule = {"P": [[3, 2, 1]]}
|
|
si = ScheduleInfo(1, None, invalid_schedule, None, None, None, None)
|
|
si.disableValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER, reason="Reason A")
|
|
si.disableValidationPass(ValidatorPass.VERIFY_CORRECT_NUMBER_OF_INSTRUCTIONS, reason="Reason B")
|
|
|
|
assert si.reasonForDisablingValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER) == "Reason A"
|
|
assert si.reasonForDisablingValidationPass(ValidatorPass.VERIFY_CORRECT_NUMBER_OF_INSTRUCTIONS) == "Reason B"
|
|
# Passes not disabled return None.
|
|
assert si.reasonForDisablingValidationPass(ValidatorPass.VERIFY_SCC_OVERLAP) is None
|
|
|
|
def test_reason_for_disabling_validation_pass(self):
|
|
"""reasonForDisablingValidationPass returns the reason for disabled passes, None for enabled ones,
|
|
and raises TypeError for non-enum arguments."""
|
|
si = ScheduleInfo(1, None, {}, None, None, None, None)
|
|
|
|
# Nothing disabled yet.
|
|
assert si.reasonForDisablingValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER) is None
|
|
|
|
si.disableValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER, reason="test reason")
|
|
assert si.reasonForDisablingValidationPass(ValidatorPass.VERIFY_ASCENDING_ORDER) == "test reason"
|
|
|
|
# Non-enum argument raises TypeError.
|
|
with pytest.raises(TypeError):
|
|
si.reasonForDisablingValidationPass("not_an_enum")
|
|
|
|
with pytest.raises(TypeError):
|
|
si.reasonForDisablingValidationPass(42)
|
|
|
|
def test_disable_validation_all_passes(self):
|
|
"""disableValidation disables all passes with the given reason."""
|
|
si = ScheduleInfo(1, None, {}, None, None, None, None)
|
|
si.disableValidation("entire schedule unsupported")
|
|
|
|
for p in ValidatorPass:
|
|
assert si.reasonForDisablingValidationPass(p) == "entire schedule unsupported"
|
|
|
|
def test_disable_validation_isvalid_early_return(self):
|
|
"""isValid returns (True, '') with a single warning when all passes are disabled."""
|
|
kernel = create_base_kernel()
|
|
si = ScheduleInfo(1, 0, {}, None, None, None, None)
|
|
si.disableValidation("not supported yet")
|
|
status, message = isValid(si, {"kernel": kernel})
|
|
assert status == True
|
|
assert message == ""
|
|
|
|
class TestSchedulePositionOrdering:
|
|
"""Test that SchedulePosition comparison handles vmfma_index=-1 correctly.
|
|
|
|
vmfma_index=-1 is a wrap-around position: it represents instructions
|
|
scheduled between iterations (after the last VMFMA of the previous
|
|
iteration, before the first VMFMA of the current iteration). With
|
|
explicit loop_index tracking, -1 naturally sorts before 0 within the
|
|
same loop via integer comparison.
|
|
"""
|
|
|
|
def test_neg1_after_last_vmfma_same_loop(self):
|
|
"""vmfma=-1 in loop 0 must be < vmfma=0 in loop 0."""
|
|
a = SchedulePosition(loop_index=0, vmfma_index=-1, sub_index=0)
|
|
b = SchedulePosition(loop_index=0, vmfma_index=0, sub_index=0)
|
|
assert a < b
|
|
assert b > a
|
|
|
|
def test_neg1_loop1_after_last_vmfma_loop0(self):
|
|
"""vmfma=-1 in loop 1 must be > vmfma=num_vmfma-1 in loop 0."""
|
|
a = SchedulePosition(loop_index=1, vmfma_index=-1, sub_index=0)
|
|
b = SchedulePosition(loop_index=0, vmfma_index=2, sub_index=0)
|
|
assert a > b
|
|
assert b < a
|
|
|
|
def test_neg1_before_next_loop_vmfma0(self):
|
|
"""vmfma=-1 in loop 0 must be < vmfma=0 in loop 1."""
|
|
a = SchedulePosition(loop_index=0, vmfma_index=-1, sub_index=0)
|
|
b = SchedulePosition(loop_index=1, vmfma_index=0, sub_index=0)
|
|
assert a < b
|
|
assert b > a
|
|
|
|
def test_equal_positions(self):
|
|
"""Identical positions must be equal."""
|
|
a = SchedulePosition(loop_index=0, vmfma_index=3, sub_index=1)
|
|
b = SchedulePosition(loop_index=0, vmfma_index=3, sub_index=1)
|
|
assert a == b
|
|
assert not (a != b)
|
|
assert not (a < b)
|
|
assert not (a > b)
|
|
assert a <= b
|
|
assert a >= b
|
|
|
|
def test_not_equal_different_sub_index(self):
|
|
"""Same (loop, vmfma) but different sub_index must not be equal."""
|
|
a = SchedulePosition(loop_index=0, vmfma_index=3, sub_index=0)
|
|
b = SchedulePosition(loop_index=0, vmfma_index=3, sub_index=1)
|
|
assert a != b
|
|
assert a < b
|
|
|
|
def test_sub_index_ordering(self):
|
|
"""Within same (loop, vmfma), sub_index determines order."""
|
|
a = SchedulePosition(loop_index=0, vmfma_index=2, sub_index=0)
|
|
b = SchedulePosition(loop_index=0, vmfma_index=2, sub_index=5)
|
|
c = SchedulePosition(loop_index=0, vmfma_index=2, sub_index=10)
|
|
assert a < b < c
|
|
assert c > b > a
|