diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py index 6e720235..341b47d2 100644 --- a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py @@ -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]) diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py index c2f0c820..1d1d98ae 100644 --- a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py @@ -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]) diff --git a/test/examples/CuTeDSL/conftest.py b/test/examples/CuTeDSL/conftest.py new file mode 100644 index 00000000..24df499f --- /dev/null +++ b/test/examples/CuTeDSL/conftest.py @@ -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) diff --git a/test/examples/CuTeDSL/sm_100a/conftest.py b/test/examples/CuTeDSL/sm_100a/conftest.py new file mode 100644 index 00000000..d3c949e5 --- /dev/null +++ b/test/examples/CuTeDSL/sm_100a/conftest.py @@ -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" \ No newline at end of file diff --git a/test/examples/CuTeDSL/sm_100a/test_tutorial_gemm.py b/test/examples/CuTeDSL/sm_100a/test_tutorial_gemm.py new file mode 100644 index 00000000..4a5de35f --- /dev/null +++ b/test/examples/CuTeDSL/sm_100a/test_tutorial_gemm.py @@ -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) \ No newline at end of file diff --git a/test/utils/device_info.py b/test/utils/device_info.py new file mode 100644 index 00000000..1e6f5b4d --- /dev/null +++ b/test/utils/device_info.py @@ -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() diff --git a/test/utils/test_sharding.py b/test/utils/test_sharding.py new file mode 100644 index 00000000..fe31d734 --- /dev/null +++ b/test/utils/test_sharding.py @@ -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", + ]