numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

numba_cuda/numba/cuda/tests/cudapy/test_libdevice.py

@@ -0,0 +1,194 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+from numba.cuda import types
+from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
+from numba import cuda
+from numba.cuda import libdevice, compile_ptx
+from numba.cuda.libdevicefuncs import functions, create_signature
+
+
+def use_sincos(s, c, x):
+    i = cuda.grid(1)
+
+    if i < len(x):
+        sr, cr = libdevice.sincos(x[i])
+        s[i] = sr
+        c[i] = cr
+
+
+def use_frexp(frac, exp, x):
+    i = cuda.grid(1)
+
+    if i < len(x):
+        fracr, expr = libdevice.frexp(x[i])
+        frac[i] = fracr
+        exp[i] = expr
+
+
+def use_sad(r, x, y, z):
+    i = cuda.grid(1)
+
+    if i < len(x):
+        r[i] = libdevice.sad(x[i], y[i], z[i])
+
+
+@skip_on_cudasim("Libdevice functions are not supported on cudasim")
+class TestLibdevice(CUDATestCase):
+    """
+    Some tests of libdevice function wrappers that check the returned values.
+
+    These are mainly to check that the generation of the implementations
+    results in correct typing and lowering for each type of function return
+    (e.g. scalar return, UniTuple return, Tuple return, etc.).
+    """
+
+    def test_sincos(self):
+        # Tests return of a UniTuple from a libdevice function
+        arr = np.arange(100, dtype=np.float64)
+        sres = np.zeros_like(arr)
+        cres = np.zeros_like(arr)
+
+        cufunc = cuda.jit(use_sincos)
+        cufunc[4, 32](sres, cres, arr)
+
+        np.testing.assert_allclose(np.cos(arr), cres)
+        np.testing.assert_allclose(np.sin(arr), sres)
+
+    def test_frexp(self):
+        # Test return of a Tuple from a libdevice function
+        arr = np.linspace(start=1.0, stop=10.0, num=100, dtype=np.float64)
+        fracres = np.zeros_like(arr)
+        expres = np.zeros(shape=arr.shape, dtype=np.int32)
+
+        cufunc = cuda.jit(use_frexp)
+        cufunc[4, 32](fracres, expres, arr)
+
+        frac_expect, exp_expect = np.frexp(arr)
+
+        np.testing.assert_array_equal(frac_expect, fracres)
+        np.testing.assert_array_equal(exp_expect, expres)
+
+    def test_sad(self):
+        # Test return of a scalar from a libdevice function
+        x = np.arange(0, 200, 2)
+        y = np.arange(50, 150)
+        z = np.arange(15, 115)
+        r = np.zeros_like(x)
+
+        cufunc = cuda.jit(use_sad)
+        cufunc[4, 32](r, x, y, z)
+
+        np.testing.assert_array_equal(np.abs(x - y) + z, r)
+
+
+# A template for generating tests of compiling calls to libdevice functions.
+# The purpose of the call and assignment of the return variables is to ensure
+# the actual function implementations are not thrown away resulting in a PTX
+# implementation that only contains the ret instruction - this may hide certain
+# errors.
+function_template = """\
+from numba.cuda import libdevice
+
+def pyfunc(%(pyargs)s):
+    ret = libdevice.%(func)s(%(funcargs)s)
+    %(retvars)s = ret
+"""
+
+
+def make_test_call(libname):
+    """
+    Generates a test function for each libdevice function.
+    """
+
+    def _test_call_functions(self):
+        # Strip off '__nv_' from libdevice name to get Python name
+        apiname = libname[5:]
+        apifunc = getattr(libdevice, apiname)  # noqa: F841
+        retty, args = functions[libname]
+        sig = create_signature(retty, args)
+
+        # Construct arguments to the libdevice function. These are all
+        # non-pointer arguments to the underlying bitcode function.
+        funcargs = ", ".join(
+            ["a%d" % i for i, arg in enumerate(args) if not arg.is_ptr]
+        )
+
+        # Arguments to the Python function (`pyfunc` in the template above) are
+        # the arguments to the libdevice function, plus as many extra arguments
+        # as there are in the return type of the libdevice function - one for
+        # scalar-valued returns, or the length of the tuple for tuple-valued
+        # returns.
+        if isinstance(sig.return_type, (types.Tuple, types.UniTuple)):
+            # Start with the parameters for the return values
+            pyargs = ", ".join(["r%d" % i for i in range(len(sig.return_type))])
+            # Add the parameters for the argument values
+            pyargs += ", " + funcargs
+            # Generate the unpacking of the return value from the libdevice
+            # function into the Python function return values (`r0`, `r1`,
+            # etc.).
+            retvars = ", ".join(
+                ["r%d[0]" % i for i in range(len(sig.return_type))]
+            )
+        else:
+            # Scalar return is a more straightforward case
+            pyargs = "r0, " + funcargs
+            retvars = "r0[0]"
+
+        # Create the string containing the function to compile
+        d = {
+            "func": apiname,
+            "pyargs": pyargs,
+            "funcargs": funcargs,
+            "retvars": retvars,
+        }
+        code = function_template % d
+
+        # Convert the string to a Python function
+        locals = {}
+        exec(code, globals(), locals)
+        pyfunc = locals["pyfunc"]
+
+        # Compute the signature for compilation. This mirrors the creation of
+        # arguments to the Python function above.
+        pyargs = [arg.ty for arg in args if not arg.is_ptr]
+        if isinstance(sig.return_type, (types.Tuple, types.UniTuple)):
+            pyreturns = [ret[::1] for ret in sig.return_type]
+            pyargs = pyreturns + pyargs
+        else:
+            pyargs.insert(0, sig.return_type[::1])
+
+        pyargs = tuple(pyargs)
+        ptx, resty = compile_ptx(pyfunc, pyargs)
+
+        # If the function body was discarded by optimization (therefore making
+        # the test a bit weak), there won't be any loading of parameters -
+        # ensure that a load from parameters occurs somewhere in the PTX
+        self.assertIn("ld.param", ptx)
+
+        # Returning the result (through a passed-in array) should also require
+        # a store to global memory, so check for at least one of those too.
+        self.assertIn("st.global", ptx)
+
+    return _test_call_functions
+
+
+@skip_on_cudasim("Compilation to PTX is not supported on cudasim")
+class TestLibdeviceCompilation(unittest.TestCase):
+    """
+    Class for holding all tests of compiling calls to libdevice functions. We
+    generate the actual tests in this class (as opposed to using subTest and
+    one test within this class) because there are a lot of tests, and it makes
+    the test suite appear frozen to test them all as subTests in one test.
+    """
+
+
+for libname in functions:
+    setattr(
+        TestLibdeviceCompilation, "test_%s" % libname, make_test_call(libname)
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_lineinfo.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numba.cuda import float32, int32
+from numba.cuda.core.errors import NumbaInvalidConfigWarning
+import pytest
+
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+import re
+import unittest
+
+
+@skip_on_cudasim("Simulator does not produce lineinfo")
+class TestCudaLineInfo(CUDATestCase):
+    def _loc_directive_regex(self):
+        # This is used in several tests
+
+        pat = (
+            r"\.loc"  # .loc directive beginning
+            r"\s+[0-9]+"  # whitespace then file index
+            r"\s+[0-9]+"  # whitespace then line number
+            r"\s+[0-9]+"  # whitespace then column position
+        )
+        return re.compile(pat)
+
+    def _check(self, fn, sig, expect):
+        fn.compile(sig)
+        llvm = fn.inspect_llvm(sig)
+        ptx = fn.inspect_asm(sig)
+        assertfn = self.assertIsNotNone if expect else self.assertIsNone
+
+        # DICompileUnit debug info metadata should all be of the
+        # DebugDirectivesOnly kind, and not the FullDebug kind
+        pat = (
+            r"!DICompileUnit\(.*"  # Opening of DICompileUnit metadata. Since
+            # the order of attributes is not
+            # guaranteed, we need to match arbitrarily
+            # afterwards.
+            r"emissionKind:\s+"  # The emissionKind attribute followed by
+            # whitespace.
+            r"DebugDirectivesOnly"  # The correct emissionKind.
+        )
+        match = re.compile(pat).search(llvm)
+        assertfn(match, msg=ptx)
+
+        pat = (
+            r"!DICompileUnit\(.*"  # Same as the pattern above, but for the
+            r"emissionKind:\s+"  # incorrect FullDebug emissionKind.
+            r"FullDebug"  #
+        )
+        match = re.compile(pat).search(llvm)
+        self.assertIsNone(match, msg=ptx)
+
+        # The name of this file should be present in the line mapping
+        # if lineinfo was propagated through correctly.
+        pat = (
+            r"\.file"  # .file directive beginning
+            r"\s+[0-9]+\s+"  # file number surrounded by whitespace
+            r'".*test_lineinfo.py"'  # filename in quotes, ignoring full path
+        )
+        match = re.compile(pat).search(ptx)
+        assertfn(match, msg=ptx)
+
+        # .loc directives should be present in the ptx
+        self._loc_directive_regex().search(ptx)
+        assertfn(match, msg=ptx)
+
+        # Debug info sections should not be present when only lineinfo is
+        # generated
+        pat = (
+            r"\.section\s+"  # .section directive beginning
+            r"\.debug_info"  # Section named ".debug_info"
+        )
+        match = re.compile(pat).search(ptx)
+        self.assertIsNone(match, msg=ptx)
+
+    def test_no_lineinfo_in_asm(self):
+        @cuda.jit(lineinfo=False)
+        def foo(x):
+            x[0] = 1
+
+        self._check(foo, sig=(int32[:],), expect=False)
+
+    def test_lineinfo_in_asm(self):
+        @cuda.jit(lineinfo=True)
+        def foo(x):
+            x[0] = 1
+
+        self._check(foo, sig=(int32[:],), expect=True)
+
+    def test_lineinfo_maintains_error_model(self):
+        sig = (float32[::1], float32[::1])
+
+        @cuda.jit(sig, lineinfo=True)
+        def divide_kernel(x, y):
+            x[0] /= y[0]
+
+        llvm = divide_kernel.inspect_llvm(sig)
+
+        # When the error model is Python, the device function returns 1 to
+        # signal an exception (e.g. divide by zero) has occurred. When the
+        # error model is the default NumPy one (as it should be when only
+        # lineinfo is enabled) the device function always returns 0.
+        self.assertNotIn("ret i32 1", llvm)
+
+    def test_no_lineinfo_in_device_function(self):
+        # Ensure that no lineinfo is generated in device functions by default.
+        @cuda.jit
+        def callee(x):
+            x[0] += 1
+
+        @cuda.jit
+        def caller(x):
+            x[0] = 1
+            callee(x)
+
+        sig = (int32[:],)
+        self._check(caller, sig=sig, expect=False)
+
+    def test_lineinfo_in_device_function(self):
+        # First we define a device function / kernel pair and run the usual
+        # checks on the generated LLVM and PTX.
+
+        @cuda.jit(lineinfo=True)
+        def callee(x):
+            x[0] += 1
+
+        @cuda.jit(lineinfo=True)
+        def caller(x):
+            x[0] = 1
+            callee(x)
+
+        sig = (int32[:],)
+        self._check(caller, sig=sig, expect=True)
+
+        # Now we can check the PTX of the device function specifically.
+
+        ptx = caller.inspect_asm(sig)
+        ptxlines = ptx.splitlines()
+
+        # Check that there is no device function in the PTX
+
+        # A line beginning with ".weak .func" that identifies a device function
+        devfn_start = re.compile(r"^\.weak\s+\.func")
+
+        for line in ptxlines:
+            if devfn_start.match(line) is not None:
+                self.fail(f"Found device function in PTX:\n\n{ptx}")
+
+        # Scan for .loc directives that refer to an inlined device function
+
+        loc_directive = self._loc_directive_regex()
+        found = False
+
+        for line in ptxlines:
+            if loc_directive.search(line) is not None:
+                if "inlined_at" in line:
+                    found = True
+                    break
+
+        if not found:
+            self.fail(
+                f"No .loc directive with inlined_at info foundin:\n\n{ptx}"
+            )
+
+        # We also inspect the LLVM to ensure that there's debug info for each
+        # subprogram (function). A lightweight way to check this is to ensure
+        # that we have as many DISubprograms as we expect.
+
+        llvm = caller.inspect_llvm(sig)
+        subprograms = 0
+        for line in llvm.splitlines():
+            if "distinct !DISubprogram" in line:
+                subprograms += 1
+
+        # One DISubprogram for each of:
+        # - The caller
+        # - The callee
+        expected_subprograms = 2
+
+        self.assertEqual(
+            subprograms,
+            expected_subprograms,
+            f'"Expected {expected_subprograms} DISubprograms; '
+            f"got {subprograms}",
+        )
+
+    def test_debug_and_lineinfo_warning(self):
+        with pytest.warns(
+            NumbaInvalidConfigWarning,
+            match="debug and lineinfo are mutually exclusive",
+        ):
+            # We pass opt=False to prevent the warning about opt and debug
+            # occurring as well
+            @cuda.jit(debug=True, lineinfo=True, opt=False)
+            def f():
+                pass
+
+    def test_lineinfo_with_compile_internal(self):
+        # Calling a function implemented using compile_internal should not
+        # enable full debug info generation. See Numba-CUDA Issue #271,
+        # https://github.com/NVIDIA/numba-cuda/issues/271
+
+        @cuda.jit("void(complex128[::1], complex128[::1])", lineinfo=True)
+        def complex_abs_use(r, x):
+            r[0] = abs(x[0])
+
+        cc = cuda.get_current_device().compute_capability
+        ov = complex_abs_use.overloads[complex_abs_use.signatures[0]]
+        ptx = ov.inspect_asm(cc)
+
+        target = ".target sm_%s%s" % cc
+        target_debug = f"{target}, debug"
+        self.assertIn(target, ptx)
+        self.assertNotIn(target_debug, ptx)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_localmem.py

@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda import int32, complex128, void
+from numba.cuda import types
+from numba.cuda import HAS_NUMBA
+
+if HAS_NUMBA:
+    from numba.core.errors import TypingError
+else:
+    from numba.cuda.core.errors import TypingError
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+from .extensions_usecases import struct_model_type, MyStruct
+
+
+def culocal(A, B):
+    C = cuda.local.array(1000, dtype=int32)
+    for i in range(C.shape[0]):
+        C[i] = A[i]
+    for i in range(C.shape[0]):
+        B[i] = C[i]
+
+
+def culocalcomplex(A, B):
+    C = cuda.local.array(100, dtype=complex128)
+    for i in range(C.shape[0]):
+        C[i] = A[i]
+    for i in range(C.shape[0]):
+        B[i] = C[i]
+
+
+def culocal1tuple(A, B):
+    C = cuda.local.array((5,), dtype=int32)
+    for i in range(C.shape[0]):
+        C[i] = A[i]
+    for i in range(C.shape[0]):
+        B[i] = C[i]
+
+
+@skip_on_cudasim("PTX inspection not available in cudasim")
+class TestCudaLocalMem(CUDATestCase):
+    def test_local_array(self):
+        sig = (int32[:], int32[:])
+        jculocal = cuda.jit(sig)(culocal)
+        self.assertTrue(".local" in jculocal.inspect_asm(sig))
+        A = np.arange(1000, dtype="int32")
+        B = np.zeros_like(A)
+        jculocal[1, 1](A, B)
+        self.assertTrue(np.all(A == B))
+
+    def test_local_array_1_tuple(self):
+        """Ensure that local arrays can be constructed with 1-tuple shape"""
+        jculocal = cuda.jit("void(int32[:], int32[:])")(culocal1tuple)
+        # Don't check if .local is in the ptx because the optimizer
+        # may reduce it to registers.
+        A = np.arange(5, dtype="int32")
+        B = np.zeros_like(A)
+        jculocal[1, 1](A, B)
+        self.assertTrue(np.all(A == B))
+
+    def test_local_array_complex(self):
+        sig = "void(complex128[:], complex128[:])"
+        jculocalcomplex = cuda.jit(sig)(culocalcomplex)
+        A = (np.arange(100, dtype="complex128") - 1) / 2j
+        B = np.zeros_like(A)
+        jculocalcomplex[1, 1](A, B)
+        self.assertTrue(np.all(A == B))
+
+    def check_dtype(self, f, dtype):
+        # Find the typing of the dtype argument to cuda.local.array
+        annotation = next(iter(f.overloads.values()))._type_annotation
+        l_dtype = annotation.typemap["l"].dtype
+        # Ensure that the typing is correct
+        self.assertEqual(l_dtype, dtype)
+
+    @skip_on_cudasim("Can't check typing in simulator")
+    def test_numba_dtype(self):
+        # Check that Numba types can be used as the dtype of a local array
+        @cuda.jit(void(int32[::1]))
+        def f(x):
+            l = cuda.local.array(10, dtype=int32)
+            l[0] = x[0]
+            x[0] = l[0]
+
+        self.check_dtype(f, int32)
+
+    @skip_on_cudasim("Can't check typing in simulator")
+    def test_numpy_dtype(self):
+        # Check that NumPy types can be used as the dtype of a local array
+        @cuda.jit(void(int32[::1]))
+        def f(x):
+            l = cuda.local.array(10, dtype=np.int32)
+            l[0] = x[0]
+            x[0] = l[0]
+
+        self.check_dtype(f, int32)
+
+    @skip_on_cudasim("Can't check typing in simulator")
+    def test_string_dtype(self):
+        # Check that strings can be used to specify the dtype of a local array
+        @cuda.jit(void(int32[::1]))
+        def f(x):
+            l = cuda.local.array(10, dtype="int32")
+            l[0] = x[0]
+            x[0] = l[0]
+
+        self.check_dtype(f, int32)
+
+    @skip_on_cudasim("Can't check typing in simulator")
+    def test_invalid_string_dtype(self):
+        # Check that strings of invalid dtypes cause a typing error
+        re = ".*Invalid NumPy dtype specified: 'int33'.*"
+        with self.assertRaisesRegex(TypingError, re):
+
+            @cuda.jit(void(int32[::1]))
+            def f(x):
+                l = cuda.local.array(10, dtype="int33")
+                l[0] = x[0]
+                x[0] = l[0]
+
+    def test_type_with_struct_data_model(self):
+        @cuda.jit(void(struct_model_type[::1]))
+        def f(x):
+            l = cuda.local.array(10, dtype=struct_model_type)
+            l[0] = x[0]
+            x[0] = l[0]
+
+        self.check_dtype(f, struct_model_type)
+
+    def test_struct_model_type_arr(self):
+        @cuda.jit(void(int32[::1], int32[::1]))
+        def f(outx, outy):
+            # Test creation
+            arr = cuda.local.array(10, dtype=struct_model_type)
+            # Test set to arr
+            for i in range(len(arr)):
+                obj = MyStruct(int32(i), int32(i * 2))
+                arr[i] = obj
+            # Test get from arr
+            for i in range(len(arr)):
+                outx[i] = arr[i].x
+                outy[i] = arr[i].y
+
+        arrx = np.array((10,), dtype="int32")
+        arry = np.array((10,), dtype="int32")
+
+        f[1, 1](arrx, arry)
+
+        for i, x in enumerate(arrx):
+            self.assertEqual(x, i)
+        for i, y in enumerate(arry):
+            self.assertEqual(y, i * 2)
+
+    def _check_local_array_size_fp16(self, shape, expected, ty):
+        @cuda.jit
+        def s(a):
+            arr = cuda.local.array(shape, dtype=ty)
+            a[0] = arr.size
+
+        result = np.zeros(1, dtype=np.float16)
+        s[1, 1](result)
+        self.assertEqual(result[0], expected)
+
+    def test_issue_fp16_support(self):
+        self._check_local_array_size_fp16(2, 2, types.float16)
+        self._check_local_array_size_fp16(2, 2, np.float16)
+
+
+if __name__ == "__main__":
+    unittest.main()