add pytest support for tutorial gemm (#2826)
* add pytest support for tutorial gemm * add license
This commit is contained in:
Linfeng Zheng
@@ -141,15 +141,15 @@ def kernel(
|
||||
# (bM, bN)
|
||||
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
|
||||
thr_mma = tiled_mma.get_slice(0)
|
||||
# (MMA, MMA_M, MMA_K)
|
||||
# (MMA, MMA_M, MMA_K, RestK)
|
||||
tCgA = thr_mma.partition_A(gA)
|
||||
# (MMA, MMA_N, MMA_K)
|
||||
# (MMA, MMA_N, MMA_K, RestK)
|
||||
tCgB = thr_mma.partition_B(gB)
|
||||
# (MMA, MMA_M, MMA_N)
|
||||
tCgC = thr_mma.partition_C(gC)
|
||||
# (MMA, MMA_M, MMA_K)
|
||||
# (MMA, MMA_M, MMA_K, STAGE)
|
||||
tCrA = tiled_mma.make_fragment_A(sA)
|
||||
# (MMA, MMA_N, MMA_K)
|
||||
# (MMA, MMA_N, MMA_K, STAGE)
|
||||
tCrB = tiled_mma.make_fragment_B(sB)
|
||||
# (MMA, MMA_M, MMA_N)
|
||||
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
|
||||
|
||||
@@ -174,15 +174,15 @@ def kernel(
|
||||
# (bM, bN)
|
||||
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
|
||||
thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
|
||||
# (MMA, MMA_M, MMA_K)
|
||||
# (MMA, MMA_M, MMA_K, RestK)
|
||||
tCgA = thr_mma.partition_A(gA)
|
||||
# (MMA, MMA_N, MMA_K)
|
||||
# (MMA, MMA_N, MMA_K, RestK)
|
||||
tCgB = thr_mma.partition_B(gB)
|
||||
# (MMA, MMA_M, MMA_N)
|
||||
tCgC = thr_mma.partition_C(gC)
|
||||
# (MMA, MMA_M, MMA_K)
|
||||
# (MMA, MMA_M, MMA_K, STAGE)
|
||||
tCrA = tiled_mma.make_fragment_A(sA)
|
||||
# (MMA, MMA_N, MMA_K)
|
||||
# (MMA, MMA_N, MMA_K, STAGE)
|
||||
tCrB = tiled_mma.make_fragment_B(sB)
|
||||
# (MMA, MMA_M, MMA_N)
|
||||
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
|
||||
|
||||
148
test/examples/CuTeDSL/conftest.py
Normal file
148
test/examples/CuTeDSL/conftest.py
Normal file
@@ -0,0 +1,148 @@
|
||||
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import logging
|
||||
import hashlib
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
project_root = Path(__file__).resolve().parent.parent.parent.parent
|
||||
example_path = project_root / "examples" / "python" / "CuTeDSL"
|
||||
utils_path = project_root / "test" / "utils"
|
||||
sys.path.append(str(example_path))
|
||||
sys.path.append(str(utils_path))
|
||||
|
||||
# The helper class to prevent modification of sys.path from test files
|
||||
# Only allow modification of sys.path from pytest monkeypatch API calls
|
||||
class ImmutableSysPath(list):
|
||||
mutating_methods = {
|
||||
"append",
|
||||
"extend",
|
||||
"insert",
|
||||
"remove",
|
||||
"pop",
|
||||
"clear",
|
||||
"reverse",
|
||||
"sort",
|
||||
"__setitem__",
|
||||
"__delitem__",
|
||||
}
|
||||
|
||||
for mtd in mutating_methods:
|
||||
def mutating_method(self, *args, mtd=mtd, **kwargs):
|
||||
frame = sys._getframe().f_back
|
||||
if (
|
||||
frame
|
||||
and hasattr(frame, "f_locals")
|
||||
and "__file__" in frame.f_locals
|
||||
and frame.f_locals["__file__"].startswith(str(project_root))
|
||||
):
|
||||
err_msg = (
|
||||
"Modification of sys.path is forbidden in test file! "
|
||||
"Please use pytest monkeypatch.syspath_prepend(...) instead."
|
||||
)
|
||||
raise RuntimeError(err_msg)
|
||||
else:
|
||||
return getattr(super(), mtd)(*args, **kwargs)
|
||||
|
||||
locals()[mtd] = mutating_method
|
||||
|
||||
def __init__(self, initial=None):
|
||||
if initial is None:
|
||||
initial = []
|
||||
super().__init__(initial)
|
||||
|
||||
|
||||
sys.path = ImmutableSysPath(list(sys.path))
|
||||
|
||||
pytest_plugins = ["test_sharding"]
|
||||
|
||||
def pytest_addoption(parser):
|
||||
parser.addoption(
|
||||
"--sample-interval",
|
||||
action="store",
|
||||
type=int,
|
||||
default=4,
|
||||
help="If value x is provided, then 1 / x of random picked tests will be run",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_interval(request):
|
||||
return request.config.getoption("--sample-interval")
|
||||
|
||||
|
||||
# Removes all StreamHandlers from loggers at the end of test session
|
||||
# This prevents errors when atexit-registered functions try to use loggers
|
||||
# whose handlers have already been closed during pytest teardown
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def cleanup_logging_handlers():
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
loggers = [logging.getLogger()] + list(
|
||||
logging.Logger.manager.loggerDict.values()
|
||||
)
|
||||
for logger in loggers:
|
||||
handlers = getattr(logger, "handlers", [])
|
||||
for handler in handlers:
|
||||
if isinstance(handler, logging.StreamHandler):
|
||||
logger.removeHandler(handler)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def torch_sanity_check():
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("GPU is required to run example tests!")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def torch_empty_cache():
|
||||
"""
|
||||
Automatically empty the torch CUDA cache at the end of each test, to reduce risk of OOM errors.
|
||||
"""
|
||||
yield
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def random_seed(request):
|
||||
test_case = request.node.nodeid.split(os.sep)[-1]
|
||||
seed = int(hashlib.md5(test_case.encode("utf-8")).hexdigest(), 16) % (2**32)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.manual_seed(seed)
|
||||
30
test/examples/CuTeDSL/sm_100a/conftest.py
Normal file
30
test/examples/CuTeDSL/sm_100a/conftest.py
Normal file
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
def pytest_configure(config):
|
||||
config.default_SMs[__file__] = "100f"
|
||||
57
test/examples/CuTeDSL/sm_100a/test_tutorial_gemm.py
Normal file
57
test/examples/CuTeDSL/sm_100a/test_tutorial_gemm.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
from blackwell.tutorial_gemm import fp16_gemm_0
|
||||
from blackwell.tutorial_gemm import fp16_gemm_1
|
||||
|
||||
import pytest
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mnk",
|
||||
[(512, 512, 256)],
|
||||
)
|
||||
@pytest.mark.parametrize("tolerance", [1e-01])
|
||||
def test_fp16_gemm_0(
|
||||
mnk: Tuple[int, int, int],
|
||||
tolerance: float,
|
||||
):
|
||||
fp16_gemm_0.run_dense_gemm(mnk, tolerance)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mnk",
|
||||
[(512, 512, 256)],
|
||||
)
|
||||
@pytest.mark.parametrize("tolerance", [1e-01])
|
||||
def test_fp16_gemm_1(
|
||||
mnk: Tuple[int, int, int],
|
||||
tolerance: float,
|
||||
):
|
||||
fp16_gemm_1.run_dense_gemm(mnk, tolerance)
|
||||
59
test/utils/device_info.py
Normal file
59
test/utils/device_info.py
Normal file
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
|
||||
def _get_device_compute_capability():
|
||||
try:
|
||||
import cuda.bindings.driver as drv
|
||||
from cuda.bindings.driver import CUdevice_attribute as dev_attr
|
||||
|
||||
def drv_api(api_name, *args):
|
||||
ret_code, *result = getattr(drv, api_name)(*args)
|
||||
if ret_code:
|
||||
raise ValueError(f"CUDA error: {ret_code}")
|
||||
return result[0] if len(result) == 1 else result
|
||||
|
||||
drv_api("cuInit", 0)
|
||||
device = drv_api("cuDeviceGet", 0)
|
||||
major = drv_api(
|
||||
"cuDeviceGetAttribute",
|
||||
dev_attr.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
|
||||
device,
|
||||
)
|
||||
minor = drv_api(
|
||||
"cuDeviceGetAttribute",
|
||||
dev_attr.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
|
||||
device,
|
||||
)
|
||||
return f"{major}{minor}"
|
||||
except Exception as e:
|
||||
print(f"Failed to get CUDA compute capability: {e}")
|
||||
return None
|
||||
|
||||
|
||||
compute_capability = _get_device_compute_capability()
|
||||
527
test/utils/test_sharding.py
Normal file
527
test/utils/test_sharding.py
Normal file
@@ -0,0 +1,527 @@
|
||||
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
|
||||
import os
|
||||
import math
|
||||
import random
|
||||
from collections import defaultdict, UserDict
|
||||
from pathlib import Path
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from _pytest.assertion.util import running_on_ci
|
||||
|
||||
|
||||
# The dict class to record the default SMs for all subtests under certain directory.
|
||||
class DefaultSMs(UserDict):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.rootdir = os.path.abspath(config.rootdir)
|
||||
|
||||
def __setitem__(self, conftest_file, sms):
|
||||
dir_path = os.path.dirname(os.path.abspath(conftest_file))
|
||||
rel_path = os.path.relpath(dir_path, self.rootdir)
|
||||
if isinstance(sms, (list, set)):
|
||||
sms = [str(sm) for sm in sms]
|
||||
else:
|
||||
sms = [str(sms)]
|
||||
self.data[rel_path] = sms
|
||||
|
||||
def __getitem__(self, test_file):
|
||||
dir_path = os.path.dirname(os.path.abspath(test_file))
|
||||
rel_path = os.path.relpath(dir_path, self.rootdir)
|
||||
sms = None
|
||||
if rel_path in self.data:
|
||||
sms = self.data[rel_path]
|
||||
else:
|
||||
for parent in Path(rel_path).parents:
|
||||
if str(parent) in self.data:
|
||||
sms = self.data[str(parent)]
|
||||
break
|
||||
return sms
|
||||
|
||||
|
||||
def pytest_addoption(parser):
|
||||
group = parser.getgroup(
|
||||
"test_sharding",
|
||||
"Test sharding for pytest. ",
|
||||
)
|
||||
|
||||
group.addoption(
|
||||
"--test-level",
|
||||
type=str,
|
||||
default="L2",
|
||||
choices=["L0", "L1", "L2"],
|
||||
help="Test level to run",
|
||||
)
|
||||
|
||||
group.addoption(
|
||||
"--random-seed", type=int, default=0, help="Seed for random subtest selection"
|
||||
)
|
||||
|
||||
group.addoption(
|
||||
"--target-cc",
|
||||
type=int,
|
||||
default=None,
|
||||
help="The target device compute capability to run on",
|
||||
)
|
||||
|
||||
group.addoption(
|
||||
"--runtime-sm",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The runtime SM for current test session.",
|
||||
)
|
||||
|
||||
group.addoption(
|
||||
"--deselect-not-run",
|
||||
action="store_true",
|
||||
help="Deselect the skipped, invalid, excluded test cases.",
|
||||
)
|
||||
|
||||
group.addoption(
|
||||
"--only-large-case",
|
||||
action="store_true",
|
||||
help="Only select the test cases marked by @pytest.mark.large_case(...). This option implies --deselect-not-run.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True)
|
||||
def pytest_configure(config):
|
||||
seed = config.getoption("--random-seed")
|
||||
random.seed(seed)
|
||||
|
||||
config.default_SMs = DefaultSMs(config)
|
||||
|
||||
# Register custom markers
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"L0(ratio: float=1.0, filter: Callable=lambda: True): mark test to run on test level 0.\n"
|
||||
"The 'ratio' value stands for the probability of each test case being selected from all subtests of a certain test function. "
|
||||
"The valid range of 'ratio' is [0, 1.0]. '0' means explicitly exclude the test case(s). "
|
||||
"'1' means explicitly select the test case(s). The default value is 1. "
|
||||
"If a test case isn't explicitly marked, then it is implicitly tagged with 'pytest.mark.L0(1)'.\n"
|
||||
"The 'filter' argument should be a callable that returns a boolean to indicate "
|
||||
"if the current test case is included(True)/excluded(False). All the valid parameters of the test function "
|
||||
"will be automatically set in its locals, so test writers can use them directly. "
|
||||
"By using the filter function, more flexible test case selections can be applied.\n"
|
||||
"The mark decorator used at lower-level can override the same mark used at higher-level. "
|
||||
"For example, 'pytest.mark.L0(0)' at parameter-level can override 'pytest.mark.L0(ratio=1)' at function-level.\n"
|
||||
"Note that when the mark is used at parameter-level, the ratio value other than 0 or 1 is meaningless.\n"
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"L1(ratio: float=1.0, filter: Callable=lambda: True): same as L0, but for test level 1.",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"L2(ratio: float=1.0, filter: Callable=lambda: True): same as L0, but for test level 2.",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"invalid_case(checker: Callable=lambda: True): marker to indicate certain test cases as invalid, and skip them when pytest collects test cases.\n"
|
||||
"The 'checker' argument should be a callable that returns True or a string to describe the invalid reason if it's an invalid case. "
|
||||
"Otherwise, return False, empty string or None. The test case parameters will be automatically set in the checker's locals, test writers can use them directly. ",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"xfail_case(checker: Callable=lambda: 'expected to fail'): marker to mark a test case as xfail given the parameters of the test function.\n"
|
||||
"The test function parameters will be automatically set in the checker's locals, so test writers can use them directly. "
|
||||
"If it's an xfail case, return a string to describe the xfail reason. Otherwise return None. This is similar to pytest.xfail(reason='xfail reason'). "
|
||||
"If you want to be more specific as to why the test is failing, the checker can return a single exception like TypeError, "
|
||||
"or a tuple of exception along with the regex string that matches the string representation of the exception, i.e. (RuntimeError, 'error string'). "
|
||||
"This is similar to pytest.raises(expected_exception, match='error string').",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"large_case(checker: Callable=lambda: True): marker to tag certain test cases that will occupy a large chunk of memory.\n"
|
||||
"The 'checker' argument should be a callable that returns True if it's a large case. Otherwise, return False. "
|
||||
"The test function parameters will be automatically set in the checker's locals, so test writers can use them directly. "
|
||||
"Note that the large cases are skipped by default. They will only be selected when the command line option '--only-large-case' is specified.",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"arch(sms: List[str]): mark test to run on specific architectures.",
|
||||
)
|
||||
|
||||
config.lvl_mark_values = mark_values(
|
||||
("ratio", 1.0, lambda x: 0 <= x <= 1.0, None),
|
||||
("filter", lambda: True, lambda x: callable(x), None),
|
||||
)
|
||||
config.shard_mark_values = mark_values(
|
||||
("size", None, lambda x: type(x) is int and x > 0, None),
|
||||
)
|
||||
config.invalid_case_mark_values = mark_values(
|
||||
("checker", lambda: True, lambda x: callable(x), None),
|
||||
)
|
||||
config.xfail_case_mark_values = mark_values(
|
||||
("checker", lambda: "expected to fail", lambda x: callable(x), None),
|
||||
)
|
||||
config.large_case_mark_values = mark_values(
|
||||
("checker", lambda: True, lambda x: callable(x), None),
|
||||
)
|
||||
config.arch_mark_values = mark_values(
|
||||
(
|
||||
"sms",
|
||||
None,
|
||||
lambda x: isinstance(x, (list, str, int)),
|
||||
lambda x: [str(x)] if not isinstance(x, list) else [str(sm) for sm in x],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def mark_values(*args_list):
|
||||
def gen_mark_value(mark):
|
||||
assert len(mark.args) <= len(
|
||||
args_list
|
||||
), f"The marker '{mark.name}' specifies more arguments than expected!"
|
||||
for idx, (key, default_value, validator, formatter) in enumerate(args_list):
|
||||
if idx < len(mark.args):
|
||||
value = mark.args[idx]
|
||||
else:
|
||||
value = mark.kwargs.get(key, default_value)
|
||||
assert validator(
|
||||
value
|
||||
), f"The value '{value}' for 'pytest.mark.{mark.name}({key}=...)' is invalid!"
|
||||
if formatter:
|
||||
value = formatter(value)
|
||||
yield value
|
||||
|
||||
supported_args = set(arg[0] for arg in args_list)
|
||||
|
||||
def get_mark_values(mark):
|
||||
assert (
|
||||
set(mark.kwargs.keys()) <= supported_args
|
||||
), f"The marker '{mark.name}' used unknown arguments: {set(mark.kwargs.keys()) - supported_args}.\n"
|
||||
return tuple(gen_mark_value(mark))
|
||||
|
||||
return get_mark_values
|
||||
|
||||
|
||||
def params_in_context(func):
|
||||
if defined_params := inspect.signature(func).parameters.keys():
|
||||
|
||||
def wrapper(full_params):
|
||||
return func(*[full_params[k] for k in defined_params])
|
||||
|
||||
else:
|
||||
|
||||
def wrapper(full_params):
|
||||
context = func.__globals__.copy()
|
||||
context.update(full_params)
|
||||
return eval(func.__code__, context)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
test_level = config.getoption("--test-level")
|
||||
|
||||
if runtime_sm := config.getoption("--runtime-sm"):
|
||||
assert re.match(r"^\d+[af]?$", runtime_sm), f"Invalid runtime SM: {runtime_sm}"
|
||||
compatible_SMs = [runtime_sm]
|
||||
elif target_cc := config.getoption("--target-cc"):
|
||||
if target_cc < 90:
|
||||
compatible_SMs = {f"{target_cc}"}
|
||||
elif target_cc == 90:
|
||||
compatible_SMs = {f"{target_cc}a", f"{target_cc}"}
|
||||
else:
|
||||
compatible_SMs = {f"{target_cc}a", f"{target_cc}f", f"{target_cc}"}
|
||||
else:
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if script_dir not in sys.path:
|
||||
sys.path.append(script_dir)
|
||||
|
||||
from device_info import compute_capability
|
||||
|
||||
if compute_capability:
|
||||
target_cc = int(compute_capability)
|
||||
else:
|
||||
raise SystemError("Failed to get CUDA compute capability!")
|
||||
|
||||
if target_cc < 90:
|
||||
compatible_SMs = {f"{target_cc}"}
|
||||
else:
|
||||
compatible_SMs = {f"{target_cc}a"}
|
||||
|
||||
item_runnable_SMs = defaultdict(list)
|
||||
|
||||
item_groups = defaultdict(
|
||||
lambda: {
|
||||
"ratio": None, # random selection ratio of test function
|
||||
"candidates": [], # candidates of test function
|
||||
"invalid": [], # marked by invalid_case
|
||||
"excluded": [], # marked by skip or L0(ratio=0)
|
||||
"arch_excluded": [], # arch not matched
|
||||
"force_selected": [], # marked by L0(ratio=1.0)
|
||||
"rand_selected": [], # randomly selected
|
||||
"rand_skipped": [], # randomly skipped
|
||||
}
|
||||
)
|
||||
|
||||
def get_closest_level_mark(item_):
|
||||
if test_level and (lvl_ := item_.get_closest_marker(test_level)):
|
||||
ratio_, filter_ = config.lvl_mark_values(lvl_)
|
||||
else:
|
||||
ratio_ = 1.0
|
||||
filter_ = lambda: True
|
||||
return ratio_, filter_
|
||||
|
||||
def get_closest_arch_mark(item_):
|
||||
if hasattr(item_, "callspec"):
|
||||
for mark in item_.callspec.marks:
|
||||
if mark.name == "arch":
|
||||
sms, *_ = config.arch_mark_values(mark)
|
||||
return sms
|
||||
if mark := item_.get_closest_marker("arch"):
|
||||
sms, *_ = config.arch_mark_values(mark)
|
||||
return sms
|
||||
if running_on_ci():
|
||||
return config.default_SMs[item_.fspath]
|
||||
else:
|
||||
return None
|
||||
|
||||
for item in items:
|
||||
func_name = item.nodeid[: len(item.originalname) - len(item.name)]
|
||||
params = item.callspec.params if hasattr(item, "callspec") else {}
|
||||
|
||||
if item.get_closest_marker("skip"):
|
||||
continue
|
||||
|
||||
# check if the test case is valid
|
||||
if invalid_case := item.get_closest_marker("invalid_case"):
|
||||
checker, *_ = config.invalid_case_mark_values(invalid_case)
|
||||
checker_ = params_in_context(checker)
|
||||
if msg := checker_(params):
|
||||
invalid_reason = (
|
||||
msg if isinstance(msg, str) else "invalid params combination"
|
||||
)
|
||||
item_groups[func_name]["invalid"].append((item, invalid_reason))
|
||||
continue
|
||||
|
||||
# check if the test case is compatible with the target arch
|
||||
sms = get_closest_arch_mark(item)
|
||||
if sms is None:
|
||||
# assume compatible if no arch mark found
|
||||
item_runnable_SMs[item.nodeid].extend(compatible_SMs)
|
||||
else:
|
||||
for sm in sms:
|
||||
if sm in compatible_SMs:
|
||||
item_runnable_SMs[item.nodeid].append(sm)
|
||||
if not item_runnable_SMs[item.nodeid]:
|
||||
item_groups[func_name]["arch_excluded"].append(item)
|
||||
continue
|
||||
|
||||
if hasattr(item, "callspec"):
|
||||
param_args_cnt = 0
|
||||
for mark in item.own_markers:
|
||||
if mark.name == "parametrize":
|
||||
param_args_cnt += 1
|
||||
|
||||
# clean up the level marks from parameter-level
|
||||
if param_args_cnt > 0:
|
||||
refined_marks = []
|
||||
for mark in reversed(item.callspec.marks):
|
||||
if mark.name != test_level:
|
||||
refined_marks.append(mark)
|
||||
assert mark is item.own_markers[-1]
|
||||
item.own_markers.pop()
|
||||
item.own_markers.extend(reversed(refined_marks))
|
||||
|
||||
# get the function-level ratio and filter
|
||||
ratio, flt = get_closest_level_mark(item)
|
||||
|
||||
# check the function-level filter
|
||||
flt_ = params_in_context(flt)
|
||||
if not flt_(params):
|
||||
ratio = 0
|
||||
|
||||
# gather the parameter-level selection count and check the malformed usage
|
||||
param_selection_cnt = 0
|
||||
for mark in item.callspec.marks:
|
||||
if test_level and mark.name == test_level:
|
||||
param_ratio, *_ = config.lvl_mark_values(mark)
|
||||
assert param_ratio in [0, 1], (
|
||||
f"The 'ratio' parameter can only be 1/True or 0/False (but now it's {param_ratio}) to denote mandatorily selecting or excluding current test case "
|
||||
f"for testing when the marker '{test_level}' is used at the parameter-level.\n"
|
||||
f'Irregular subtest: "{item.nodeid}"'
|
||||
)
|
||||
if param_ratio == 0:
|
||||
# force deselect the test case
|
||||
ratio = 0
|
||||
else:
|
||||
param_selection_cnt += 1
|
||||
|
||||
assert param_args_cnt >= param_selection_cnt
|
||||
# force select the test case at parameter-level
|
||||
if param_args_cnt != 0 and param_args_cnt == param_selection_cnt:
|
||||
ratio = 1.0
|
||||
else:
|
||||
ratio, _ = get_closest_level_mark(item)
|
||||
|
||||
# check if the test case is large
|
||||
if large_case := item.get_closest_marker("large_case"):
|
||||
checker, *_ = config.large_case_mark_values(large_case)
|
||||
checker_ = params_in_context(checker)
|
||||
if bool(checker_(params)) ^ bool(config.getoption("--only-large-case")):
|
||||
item_groups[func_name]["excluded"].append(item)
|
||||
continue
|
||||
else:
|
||||
if config.getoption("--only-large-case"):
|
||||
item_groups[func_name]["excluded"].append(item)
|
||||
continue
|
||||
|
||||
if ratio > 0:
|
||||
if ratio == 1.0:
|
||||
item_groups[func_name]["force_selected"].append(item)
|
||||
else:
|
||||
item_groups[func_name]["candidates"].append(item)
|
||||
|
||||
if item_groups[func_name]["ratio"] is None:
|
||||
item_groups[func_name]["ratio"] = ratio
|
||||
assert item_groups[func_name]["ratio"] == ratio, (
|
||||
f"The ratio of subtests for {func_name} is not consistent.\n"
|
||||
f'First occurrence: {item_groups[func_name]["candidates"][0].nodeid} -> {item_groups[func_name]["ratio"]}\n'
|
||||
f"Current: {item.nodeid} -> {ratio}"
|
||||
)
|
||||
else:
|
||||
item_groups[func_name]["excluded"].append(item)
|
||||
|
||||
# Randomly select subtests
|
||||
for func_name, test_info in item_groups.items():
|
||||
if test_info["candidates"]:
|
||||
if test_info["ratio"] <= 0 or test_info["ratio"] > 1.0:
|
||||
raise ValueError("Split ratio must be between (0, 1.0]")
|
||||
|
||||
random.shuffle(test_info["candidates"])
|
||||
split_point = math.ceil(test_info["ratio"] * len(test_info["candidates"]))
|
||||
test_info["rand_selected"] += test_info["candidates"][:split_point]
|
||||
test_info["rand_skipped"] += test_info["candidates"][split_point:]
|
||||
|
||||
# Assign markers to each subtest
|
||||
for test_info in item_groups.values():
|
||||
for item, invalid_reason in test_info["invalid"]:
|
||||
item.add_marker(pytest.mark.skip(reason=invalid_reason), append=False)
|
||||
for item in test_info["excluded"]:
|
||||
item.add_marker(pytest.mark.skip(reason=f"non-{test_level}"), append=False)
|
||||
for item in test_info["arch_excluded"]:
|
||||
item.add_marker(pytest.mark.skip(reason=f"arch not matched"), append=False)
|
||||
for item in test_info["rand_skipped"]:
|
||||
item.add_marker(pytest.mark.skip(reason=f"randomly"), append=False)
|
||||
for item in test_info["force_selected"] + test_info["rand_selected"]:
|
||||
item.add_marker(f"{test_level}", append=False)
|
||||
item.add_marker(
|
||||
pytest.mark.arch(sms=item_runnable_SMs[item.nodeid]), append=False
|
||||
)
|
||||
if any(
|
||||
"report_resource_usage" == name
|
||||
for name, _ in config.pluginmanager.list_name_plugin()
|
||||
):
|
||||
# item.add_marker(pytest.mark.report_uss(interval=0.1), append=False)
|
||||
# item.add_marker(pytest.mark.report_duration, append=False)
|
||||
...
|
||||
|
||||
if config.getoption("--deselect-not-run") or config.getoption("--only-large-case"):
|
||||
selected = []
|
||||
deselected = []
|
||||
for item in items:
|
||||
if item.get_closest_marker("skip"):
|
||||
deselected.append(item)
|
||||
else:
|
||||
selected.append(item)
|
||||
items[:] = selected
|
||||
config.hook.pytest_deselected(items=deselected)
|
||||
|
||||
|
||||
@pytest.hookimpl(hookwrapper=True)
|
||||
def pytest_runtest_makereport(item, call):
|
||||
outcome = yield
|
||||
report = outcome.get_result()
|
||||
|
||||
if report.outcome != "skipped" and call.when == "call":
|
||||
if xfail_case := item.get_closest_marker("xfail_case"):
|
||||
params = item.callspec.params if hasattr(item, "callspec") else {}
|
||||
checker, *_ = item.config.xfail_case_mark_values(xfail_case)
|
||||
checker_ = params_in_context(checker)
|
||||
if xfail_info := checker_(params):
|
||||
if report.outcome == "passed":
|
||||
report.wasxfail = "expected to fail, but passed"
|
||||
else:
|
||||
if isinstance(xfail_info, str):
|
||||
report.outcome = "skipped"
|
||||
report.wasxfail = xfail_info
|
||||
elif isinstance(xfail_info, tuple):
|
||||
if isinstance(call.excinfo.value, xfail_info[0]):
|
||||
if re.search(xfail_info[1], str(call.excinfo.value)):
|
||||
report.outcome = "skipped"
|
||||
report.wasxfail = str(call.excinfo.value)
|
||||
else:
|
||||
report.longrepr = f"Expect '{xfail_info[1]}', but got '{call.excinfo.value}'"
|
||||
else:
|
||||
report.longrepr = f"Expect '{xfail_info[0]}', but got '{type(call.excinfo.value)}'"
|
||||
else:
|
||||
if isinstance(call.excinfo.value, xfail_info):
|
||||
report.outcome = "skipped"
|
||||
report.wasxfail = str(call.excinfo.value)
|
||||
else:
|
||||
report.longrepr = f"Expect exception '{xfail_info}', but got '{type(call.excinfo.value)}'"
|
||||
|
||||
|
||||
def pytest_report_collectionfinish(config, start_path, startdir, items):
|
||||
if not config.getoption("--collect-only"):
|
||||
return
|
||||
|
||||
report = []
|
||||
selected_count = 0
|
||||
for item in items:
|
||||
additional_info = ""
|
||||
if skip_mark := item.get_closest_marker("skip"):
|
||||
status = "SKIPPED"
|
||||
additional_info = skip_mark.kwargs.get("reason", "")
|
||||
else:
|
||||
status = "SELECTED"
|
||||
if sms := item.get_closest_marker("arch").kwargs.get("sms"):
|
||||
additional_info += f"SMs {','.join(sms)}"
|
||||
selected_count += 1
|
||||
|
||||
line = f"SHARDING {item.nodeid} ({status}:{additional_info})"
|
||||
report.append(line)
|
||||
|
||||
report.append(
|
||||
f"\n{len(items)} tests collected, {selected_count} tests selected, {len(items) - selected_count} tests skipped."
|
||||
)
|
||||
|
||||
if report:
|
||||
return [
|
||||
"\n============= Sharding Report Start =============\n",
|
||||
*report,
|
||||
"\n============= Sharding Report End =============\n",
|
||||
]
|
||||
Reference in New Issue
Block a user