numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/testing.py

@@ -0,0 +1,347 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import os
+import platform
+import shutil
+import pytest
+from datetime import datetime
+from numba.cuda.utils import PYVERSION
+from numba.cuda.cuda_paths import get_conda_ctk_libdir
+from numba.cuda.cudadrv import driver, devices, libs
+from numba.cuda.dispatcher import CUDADispatcher
+from numba.cuda import config
+from numba.cuda.tests.support import TestCase
+from pathlib import Path
+
+from typing import Iterable, Union
+from io import StringIO
+import unittest
+import numpy as np
+from numba.cuda import HAS_NUMBA
+
+if PYVERSION >= (3, 10):
+    from filecheck.matcher import Matcher
+    from filecheck.options import Options
+    from filecheck.parser import Parser, pattern_for_opts
+    from filecheck.finput import FInput
+
+numba_cuda_dir = Path(__file__).parent
+test_data_dir = numba_cuda_dir / "tests" / "data"
+
+
+@pytest.mark.usefixtures("initialize_from_pytest_config")
+class CUDATestCase(TestCase):
+    """
+    For tests that use a CUDA device.
+
+    Methods assertFileCheckAsm and assertFileCheckLLVM will inspect a
+    CUDADispatcher and assert that the compilation artifacts match the
+    FileCheck checks given in the kernel's docstring.
+
+    Method assertFileCheckMatches can be used to assert that a given string
+    matches FileCheck checks, and is not specific to CUDADispatcher.
+    """
+
+    FLOAT16_RTOL = np.finfo(np.float16).eps
+
+    def setUp(self):
+        self._low_occupancy_warnings = config.CUDA_LOW_OCCUPANCY_WARNINGS
+        self._warn_on_implicit_copy = config.CUDA_WARN_ON_IMPLICIT_COPY
+
+        # Disable warnings about low gpu utilization in the test suite
+        config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
+        # Disable warnings about host arrays in the test suite
+        config.CUDA_WARN_ON_IMPLICIT_COPY = 0
+
+    def tearDown(self):
+        config.CUDA_LOW_OCCUPANCY_WARNINGS = self._low_occupancy_warnings
+        config.CUDA_WARN_ON_IMPLICIT_COPY = self._warn_on_implicit_copy
+
+    Signature = Union[tuple[type, ...], None]
+
+    def _getIRContents(
+        self,
+        ir_result: Union[dict[Signature, str], str],
+        signature: Union[Signature, None] = None,
+    ) -> Iterable[str]:
+        if isinstance(ir_result, str):
+            assert signature is None, (
+                "Cannot use signature because the kernel was only compiled for one signature"
+            )
+            return [ir_result]
+
+        if signature is None:
+            return list(ir_result.values())
+
+        return [ir_result[signature]]
+
+    def assertFileCheckAsm(
+        self,
+        ir_producer: CUDADispatcher,
+        signature: Union[tuple[type, ...], None] = None,
+        check_prefixes: tuple[str] = ("ASM",),
+        **extra_filecheck_options,
+    ) -> None:
+        """
+        Assert that the assembly output of the given CUDADispatcher matches
+        the FileCheck checks given in the kernel's docstring.
+        """
+        ir_contents = self._getIRContents(ir_producer.inspect_asm(), signature)
+        assert ir_contents, "No assembly output found for the given signature."
+        assert ir_producer.__doc__ is not None, (
+            "Kernel docstring is required. To pass checks explicitly, use assertFileCheckMatches."
+        )
+        check_patterns = ir_producer.__doc__
+        for ir_content in ir_contents:
+            self.assertFileCheckMatches(
+                ir_content,
+                check_patterns=check_patterns,
+                check_prefixes=check_prefixes,
+                **extra_filecheck_options,
+            )
+
+    def assertFileCheckLLVM(
+        self,
+        ir_producer: CUDADispatcher,
+        signature: Union[tuple[type, ...], None] = None,
+        check_prefixes: tuple[str] = ("LLVM",),
+        **extra_filecheck_options,
+    ) -> None:
+        """
+        Assert that the LLVM IR output of the given CUDADispatcher matches
+        the FileCheck checks given in the kernel's docstring.
+        """
+        ir_contents = self._getIRContents(ir_producer.inspect_llvm(), signature)
+        assert ir_contents, "No LLVM IR output found for the given signature."
+        assert ir_producer.__doc__ is not None, (
+            "Kernel docstring is required. To pass checks explicitly, use assertFileCheckMatches."
+        )
+        check_patterns = ir_producer.__doc__
+        for ir_content in ir_contents:
+            assert ir_content, (
+                "LLVM IR content is empty for the given signature."
+            )
+            self.assertFileCheckMatches(
+                ir_content,
+                check_patterns=check_patterns,
+                check_prefixes=check_prefixes,
+                **extra_filecheck_options,
+            )
+
+    def assertFileCheckMatches(
+        self,
+        ir_content: str,
+        check_patterns: str,
+        check_prefixes: tuple[str] = ("CHECK",),
+        **extra_filecheck_options,
+    ) -> None:
+        """
+        Assert that the given string matches the passed FileCheck checks.
+
+        Args:
+            ir_content: The string to check against.
+            check_patterns: The FileCheck checks to use.
+            check_prefixes: The prefixes to use for the FileCheck checks.
+            extra_filecheck_options: Extra options to pass to FileCheck.
+        """
+        if PYVERSION < (3, 10):
+            self.skipTest("FileCheck requires Python 3.10 or later")
+        opts = Options(
+            match_filename="-",
+            check_prefixes=list(check_prefixes),
+            **extra_filecheck_options,
+        )
+        input_file = FInput(fname="-", content=ir_content)
+        parser = Parser(opts, StringIO(check_patterns), *pattern_for_opts(opts))
+        matcher = Matcher(opts, input_file, parser)
+        matcher.stderr = StringIO()
+        result = matcher.run()
+        if result != 0:
+            if self._dump_failed_filechecks:
+                dump_directory = Path(
+                    datetime.now().strftime("numba-ir-%Y_%m_%d_%H_%M_%S")
+                )
+                if not dump_directory.exists():
+                    dump_directory.mkdir(parents=True, exist_ok=True)
+                base_path = self.id().replace(".", "_")
+                ir_dump = dump_directory / Path(base_path).with_suffix(".ll")
+                checks_dump = dump_directory / Path(base_path).with_suffix(
+                    ".checks"
+                )
+                with (
+                    open(ir_dump, "w") as ir_file,
+                    open(checks_dump, "w") as checks_file,
+                ):
+                    _ = ir_file.write(ir_content + "\n")
+                    _ = checks_file.write(check_patterns)
+                    dump_instructions = f"Reproduce with:\n\nfilecheck --check-prefixes={','.join(check_prefixes)} {checks_dump} --input-file {ir_dump}"
+            else:
+                dump_instructions = "Rerun with --dump-failed-filechecks to generate a reproducer."
+
+            self.fail(
+                f"FileCheck failed:\n{matcher.stderr.getvalue()}\n\n"
+                + dump_instructions
+            )
+
+
+def skip_on_cudasim(reason):
+    """Skip this test if running on the CUDA simulator"""
+    assert isinstance(reason, str)
+    return unittest.skipIf(config.ENABLE_CUDASIM, reason)
+
+
+skip_on_standalone_numba_cuda = unittest.skipUnless(
+    HAS_NUMBA, "requires base numba install"
+)
+
+
+def skip_unless_cudasim(reason):
+    """Skip this test if running on CUDA hardware"""
+    assert isinstance(reason, str)
+    return unittest.skipUnless(config.ENABLE_CUDASIM, reason)
+
+
+def skip_unless_conda_cudatoolkit(reason):
+    """Skip test if the CUDA toolkit was not installed by Conda"""
+    assert isinstance(reason, str)
+    return unittest.skipUnless(get_conda_ctk_libdir() is not None, reason)
+
+
+def skip_if_external_memmgr(reason):
+    """Skip test if an EMM Plugin is in use"""
+    assert isinstance(reason, str)
+    return unittest.skipIf(config.CUDA_MEMORY_MANAGER != "default", reason)
+
+
+def skip_under_cuda_memcheck(reason):
+    assert isinstance(reason, str)
+    return unittest.skipIf(os.environ.get("CUDA_MEMCHECK") is not None, reason)
+
+
+def skip_without_nvdisasm(reason):
+    assert isinstance(reason, str)
+    nvdisasm_path = shutil.which("nvdisasm")
+    return unittest.skipIf(nvdisasm_path is None, reason)
+
+
+def skip_with_nvdisasm(reason):
+    assert isinstance(reason, str)
+    nvdisasm_path = shutil.which("nvdisasm")
+    return unittest.skipIf(nvdisasm_path is not None, reason)
+
+
+def skip_on_arm(reason):
+    assert isinstance(reason, str)
+    cpu = platform.processor()
+    is_arm = cpu.startswith("arm") or cpu.startswith("aarch")
+    return unittest.skipIf(is_arm, reason)
+
+
+def skip_on_wsl2(reason):
+    """Skip test when running under WSL2.
+
+    Detection is based on the kernel release string, which typically contains
+    "microsoft-standard-WSL2" on WSL2 systems.
+    """
+    assert isinstance(reason, str)
+    rel = platform.release().lower()
+    is_wsl2 = ("microsoft-standard-wsl2" in rel) or ("wsl2" in rel)
+    return unittest.skipIf(is_wsl2, reason)
+
+
+def skip_if_cuda_includes_missing(fn):
+    # Skip when cuda.h is not available - generally this should indicate
+    # whether the CUDA includes are available or not
+    reason = "CUDA include dir not available on this system"
+    try:
+        cuda_include_path = libs.get_cuda_include_dir()
+    except FileNotFoundError:
+        return unittest.skip(reason)(fn)
+    cuda_h = os.path.join(cuda_include_path, "cuda.h")
+    cuda_h_file = os.path.exists(cuda_h) and os.path.isfile(cuda_h)
+    return unittest.skipUnless(cuda_h_file, reason)(fn)
+
+
+def skip_if_curand_kernel_missing(fn):
+    reason = "curand_kernel.h not available on this system"
+    try:
+        cuda_include_path = libs.get_cuda_include_dir()
+    except FileNotFoundError:
+        return unittest.skip(reason)(fn)
+    curand_kernel_h = os.path.join(cuda_include_path, "curand_kernel.h")
+    curand_kernel_h_file = os.path.exists(curand_kernel_h) and os.path.isfile(
+        curand_kernel_h
+    )
+    return unittest.skipUnless(curand_kernel_h_file, reason)(fn)
+
+
+def skip_if_mvc_enabled(reason):
+    """Skip a test if Minor Version Compatibility is enabled"""
+    assert isinstance(reason, str)
+    return unittest.skipIf(
+        config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY, reason
+    )
+
+
+def cc_X_or_above(major, minor):
+    if not config.ENABLE_CUDASIM:
+        cc = devices.get_context().device.compute_capability
+        return cc >= (major, minor)
+    else:
+        return True
+
+
+def skip_unless_cc_50(fn):
+    return unittest.skipUnless(cc_X_or_above(5, 0), "requires cc >= 5.0")(fn)
+
+
+def skip_unless_cc_53(fn):
+    return unittest.skipUnless(cc_X_or_above(5, 3), "requires cc >= 5.3")(fn)
+
+
+def skip_unless_cc_60(fn):
+    return unittest.skipUnless(cc_X_or_above(6, 0), "requires cc >= 6.0")(fn)
+
+
+def skip_unless_cc_75(fn):
+    return unittest.skipUnless(cc_X_or_above(7, 5), "requires cc >= 7.5")(fn)
+
+
+def xfail_unless_cudasim(fn):
+    if config.ENABLE_CUDASIM:
+        return fn
+    else:
+        return unittest.expectedFailure(fn)
+
+
+def cudadevrt_missing():
+    if config.ENABLE_CUDASIM:
+        return False
+    try:
+        path = libs.get_cudalib("cudadevrt", static=True)
+        libs.check_static_lib(path)
+    except FileNotFoundError:
+        return True
+    return False
+
+
+def skip_if_cudadevrt_missing(fn):
+    return unittest.skipIf(cudadevrt_missing(), "cudadevrt missing")(fn)
+
+
+def skip_if_nvjitlink_missing(reason):
+    assert isinstance(reason, str)
+    return unittest.skipIf(not driver._have_nvjitlink(), reason)
+
+
+class ForeignArray(object):
+    """
+    Class for emulating an array coming from another library through the CUDA
+    Array interface. This just hides a DeviceNDArray so that it doesn't look
+    like a DeviceNDArray.
+    """
+
+    def __init__(self, arr):
+        self._arr = arr
+        self.__cuda_array_interface__ = arr.__cuda_array_interface__

numba_cuda/numba/cuda/tests/__init__.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from fnmatch import fnmatch
+import unittest
+from numba import cuda
+from os.path import dirname, isfile, join, normpath, relpath, splitext
+
+import os
+import sys
+import traceback
+
+
+# Copied and modified from numba/testing/__init__.py, to handle the difference
+# between the top dirs for Numba and the CUDA target
+def load_testsuite(loader, dir):
+    """Find tests in 'dir'."""
+    top_level_dir = dirname(dirname(dirname(dirname(__file__))))
+    try:
+        suite = unittest.TestSuite()
+        files = []
+        for f in os.listdir(dir):
+            path = join(dir, f)
+            if isfile(path) and fnmatch(f, "test_*.py"):
+                files.append(f)
+            elif isfile(join(path, "__init__.py")):
+                suite.addTests(
+                    loader.discover(path, top_level_dir=top_level_dir)
+                )
+        for f in files:
+            # turn 'f' into a filename relative to the toplevel dir and
+            # translate it to a module name. This differs from the
+            # implementation in Numba, because the toplevel dir is the
+            # numba_cuda module location, not the numba one.
+            f = relpath(join(dir, f), top_level_dir)
+            f = splitext(normpath(f.replace(os.path.sep, ".")))[0]
+            suite.addTests(loader.loadTestsFromName(f))
+        return suite
+    except Exception:
+        traceback.print_exc(file=sys.stderr)
+        sys.exit(-1)
+
+
+def load_tests(loader, tests, pattern):
+    suite = unittest.TestSuite()
+    this_dir = dirname(__file__)
+    suite.addTests(load_testsuite(loader, join(this_dir, "nocuda")))
+    if cuda.is_available():
+        suite.addTests(load_testsuite(loader, join(this_dir, "cudasim")))
+        gpus = cuda.list_devices()
+        if gpus and gpus[0].compute_capability >= (2, 0):
+            suite.addTests(load_testsuite(loader, join(this_dir, "cudadrv")))
+            suite.addTests(load_testsuite(loader, join(this_dir, "cudapy")))
+            suite.addTests(load_testsuite(loader, join(this_dir, "nrt")))
+            suite.addTests(
+                load_testsuite(loader, join(this_dir, "doc_examples"))
+            )
+        else:
+            print("skipped CUDA tests because GPU CC < 2.0")
+    else:
+        print("skipped CUDA tests")
+    return suite

numba_cuda/numba/cuda/tests/benchmarks/test_kernel_launch.py

@@ -0,0 +1,119 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import string
+from numba import cuda
+from numba.cuda.core import config
+import numpy as np
+import pytest
+from pytest import param
+
+
+pytestmark = pytest.mark.skipif(
+    condition=config.ENABLE_CUDASIM,
+    reason="no reason to run benchmarks in the simulator",
+)
+
+
+@pytest.mark.parametrize(
+    "array_func",
+    [
+        param(
+            lambda: cuda.device_array(128, dtype=np.float32),
+            id="device_array",
+        ),
+        param(
+            lambda: pytest.importorskip("torch").empty(
+                (128,),
+                dtype=pytest.importorskip("torch").float32,
+                device="cuda:0",
+            ),
+            id="torch",
+        ),
+        param(
+            lambda: pytest.importorskip("cupy").empty(128, dtype=np.float32),
+            id="cupy",
+        ),
+    ],
+)
+def test_one_arg(benchmark, array_func):
+    @cuda.jit("void(float32[:])")
+    def one_arg(arr1):
+        return
+
+    def bench(func, arr):
+        for _ in range(100):
+            func(arr)
+
+    benchmark(bench, one_arg[128, 128], array_func())
+
+
+@pytest.mark.parametrize(
+    "array_func",
+    [
+        param(
+            lambda: [
+                cuda.device_array(128, dtype=np.float32)
+                for _ in range(len(string.ascii_lowercase))
+            ],
+            id="device_array",
+        ),
+        param(
+            lambda: [
+                pytest.importorskip("torch").empty(
+                    (128,),
+                    dtype=pytest.importorskip("torch").float32,
+                    device="cuda:0",
+                )
+                for _ in range(len(string.ascii_lowercase))
+            ],
+            id="torch",
+        ),
+        param(
+            lambda: [
+                pytest.importorskip("cupy").empty(128, dtype=np.float32)
+                for _ in range(len(string.ascii_lowercase))
+            ],
+            id="cupy",
+        ),
+    ],
+)
+def test_many_args(benchmark, array_func):
+    many_arrs = array_func()
+
+    @cuda.jit("void({})".format(", ".join(["float32[:]"] * len(many_arrs))))
+    def many_args(
+        a,
+        b,
+        c,
+        d,
+        e,
+        f,
+        g,
+        h,
+        i,
+        j,
+        k,
+        l,
+        m,
+        n,
+        o,
+        p,
+        q,
+        r,
+        s,
+        t,
+        u,
+        v,
+        w,
+        x,
+        y,
+        z,
+    ):
+        return
+
+    def bench(func, *arrs):
+        for _ in range(100):
+            func(*arrs)
+
+    benchmark(bench, many_args[128, 128], *many_arrs)

numba_cuda/numba/cuda/tests/cloudpickle_main_class.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+# Expected to run this module as __main__
+
+
+# Cloudpickle will think this is a dynamic class when this module is __main__
+class Klass:
+    classvar = None

numba_cuda/numba/cuda/tests/core/serialize_usecases.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Separate module with function samples for serialization tests,
+to avoid issues with __main__.
+"""
+
+import math
+from math import sqrt
+import numpy as np
+import numpy.random as nprand
+
+# This does not need a guard, it's already guarded at the import site
+from numba import jit
+
+
+@jit("int32(int32, int32)")
+def add_with_sig(a, b):
+    return a + b
+
+
+@jit
+def add_without_sig(a, b):
+    return a + b
+
+
+@jit(nopython=True)
+def add_nopython(a, b):
+    return a + b
+
+
+@jit(nopython=True)
+def add_nopython_fail(a, b):
+    object()
+    return a + b
+
+
+def closure(a):
+    @jit(nopython=True)
+    def inner(b, c):
+        return a + b + c
+
+    return inner
+
+
+K = 3.0
+
+
+def closure_with_globals(x, **jit_args):
+    @jit(**jit_args)
+    def inner(y):
+        # Exercise a builtin function and a module-level constant
+        k = max(K, K + 1)
+        # Exercise two functions from another module, one accessed with
+        # dotted notation, one imported explicitly.
+        return math.hypot(x, y) + sqrt(k)
+
+    return inner
+
+
+@jit(nopython=True)
+def other_function(x, y):
+    return math.hypot(x, y)
+
+
+@jit(forceobj=True)
+def get_global_objmode(x):
+    return K * x
+
+
+@jit(nopython=True)
+def get_renamed_module(x):
+    nprand.seed(42)
+    return np.cos(x), nprand.random()
+
+
+def closure_calling_other_function(x):
+    @jit(nopython=True)
+    def inner(y, z):
+        return other_function(x, y) + z
+
+    return inner
+
+
+def closure_calling_other_closure(x):
+    @jit(nopython=True)
+    def other_inner(y):
+        return math.hypot(x, y)
+
+    @jit(nopython=True)
+    def inner(y):
+        return other_inner(y) + x
+
+    return inner
+
+
+# A dynamic function calling a builtin function
+def _get_dyn_func(**jit_args):
+    code = """
+        def dyn_func(x):
+            res = 0
+            for i in range(x):
+                res += x
+            return res
+        """
+    ns = {}
+    exec(code.strip(), ns)
+    return jit(**jit_args)(ns["dyn_func"])
+
+
+dyn_func = _get_dyn_func(nopython=True)
+dyn_func_objmode = _get_dyn_func(forceobj=True)