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

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

@@ -0,0 +1,312 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+
+from collections import namedtuple
+from functools import partial
+from itertools import product
+from numba.cuda import vectorize as cuda_vectorize
+from numba import cuda, vectorize as numba_vectorize
+from numba.cuda.types import int32, float32, float64
+from numba.cuda.cudadrv.driver import CudaAPIError, driver
+from numba.cuda.testing import skip_on_cudasim
+from numba.cuda.testing import CUDATestCase
+import unittest
+
+
+# Signatures to test with - these are all homogeneous in dtype, so the output
+# dtype should match the input dtype - the output should not have been cast
+# upwards, as reported in #8400: https://github.com/numba/numba/issues/8400
+signatures = [
+    int32(int32, int32),
+    float32(float32, float32),
+    float64(float64, float64),
+]
+
+# The order here is chosen such that each subsequent dtype might have been
+# casted to a previously-used dtype. This is unlikely to be an issue for CUDA,
+# but there might be future circumstances in which it becomes relevant, perhaps
+# if it supported Dynamic UFuncs, and we want to ensure that an implementation
+# for a the given dtype is used rather than casting the input upwards.
+dtypes = (np.float64, np.float32, np.int32)
+
+# NumPy ndarray orders
+orders = ("C", "F")
+
+# Input sizes corresponding to operations:
+# - Less than one warp,
+# - Less than one block,
+# - Greater than one block (i.e. many blocks)
+input_sizes = (8, 100, 2**10 + 1)
+
+# Vectorize functions to test
+# cuda.vectorize doesn't need target parameter, numba.vectorize needs target="cuda"
+vectorize_funcs = [cuda_vectorize, partial(numba_vectorize, target="cuda")]
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestCUDAVectorize(CUDATestCase):
+    # Presumably chosen as an odd number unlikely to coincide with the total
+    # thread count, and large enough to ensure a significant number of blocks
+    # are used.
+    N = 1000001
+
+    def test_scalar(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            a = 1.2
+            b = 2.3
+            c = vector_add(a, b)
+            self.assertEqual(c, a + b)
+
+    def test_1d(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            for ty in dtypes:
+                data = np.array(np.random.random(self.N), dtype=ty)
+                expected = np.add(data, data)
+                actual = vector_add(data, data)
+                np.testing.assert_allclose(expected, actual)
+                self.assertEqual(actual.dtype, ty)
+
+    def test_1d_async(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            stream = cuda.stream()
+
+            for ty in dtypes:
+                data = np.array(np.random.random(self.N), dtype=ty)
+                device_data = cuda.to_device(data, stream)
+
+                dresult = vector_add(device_data, device_data, stream=stream)
+                actual = dresult.copy_to_host()
+
+                expected = np.add(data, data)
+
+                np.testing.assert_allclose(expected, actual)
+                self.assertEqual(actual.dtype, ty)
+
+    def test_nd(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            for nd, dtype, order in product(range(1, 8), dtypes, orders):
+                shape = (4,) * nd
+                data = np.random.random(shape).astype(dtype)
+                data2 = np.array(data.T, order=order)
+
+                expected = data + data2
+                actual = vector_add(data, data2)
+                np.testing.assert_allclose(expected, actual)
+                self.assertEqual(actual.dtype, dtype)
+
+    def test_output_arg(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            A = np.arange(10, dtype=np.float32)
+            B = np.arange(10, dtype=np.float32)
+
+            expected = A + B
+            actual = np.empty_like(A)
+            vector_add(A, B, out=actual)
+
+            np.testing.assert_allclose(expected, actual)
+            self.assertEqual(expected.dtype, actual.dtype)
+
+    def test_reduce(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            dtype = np.int32
+
+            for n in input_sizes:
+                x = np.arange(n, dtype=dtype)
+                expected = np.add.reduce(x)
+                actual = vector_add.reduce(x)
+                np.testing.assert_allclose(expected, actual)
+                # np.add.reduce is special-cased to return an int64 for any int
+                # arguments, so we can't compare against its returned dtype when
+                # we're checking the general reduce machinery (which just happens
+                # to be using addition). Instead, compare against the input dtype.
+                self.assertEqual(dtype, actual.dtype)
+
+    def test_reduce_async(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            stream = cuda.stream()
+            dtype = np.int32
+
+            for n in input_sizes:
+                x = np.arange(n, dtype=dtype)
+                expected = np.add.reduce(x)
+                dx = cuda.to_device(x, stream)
+                actual = vector_add.reduce(dx, stream=stream)
+                np.testing.assert_allclose(expected, actual)
+                # Compare against the input dtype as in test_reduce().
+                self.assertEqual(dtype, actual.dtype)
+
+    def test_manual_transfer(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            n = 10
+            x = np.arange(n, dtype=np.int32)
+            dx = cuda.to_device(x)
+            expected = x + x
+            actual = vector_add(x, dx).copy_to_host()
+            np.testing.assert_equal(expected, actual)
+            self.assertEqual(expected.dtype, actual.dtype)
+
+    def test_ufunc_output_2d(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            n = 10
+            x = np.arange(n, dtype=np.int32).reshape(2, 5)
+            dx = cuda.to_device(x)
+            vector_add(dx, dx, out=dx)
+
+            expected = x + x
+            actual = dx.copy_to_host()
+            np.testing.assert_equal(expected, actual)
+            self.assertEqual(expected.dtype, actual.dtype)
+
+    def check_tuple_arg(self, a, b):
+        for vectorize in vectorize_funcs:
+
+            @vectorize(signatures)
+            def vector_add(a, b):
+                return a + b
+
+            r = vector_add(a, b)
+            np.testing.assert_equal(np.asarray(a) + np.asarray(b), r)
+
+    def test_tuple_arg(self):
+        a = (1.0, 2.0, 3.0)
+        b = (4.0, 5.0, 6.0)
+        self.check_tuple_arg(a, b)
+
+    def test_namedtuple_arg(self):
+        Point = namedtuple("Point", ("x", "y", "z"))
+        a = Point(x=1.0, y=2.0, z=3.0)
+        b = Point(x=4.0, y=5.0, z=6.0)
+        self.check_tuple_arg(a, b)
+
+    def test_tuple_of_array_arg(self):
+        arr = np.arange(10, dtype=np.int32)
+        a = (arr, arr + 1)
+        b = (arr + 2, arr + 2)
+        self.check_tuple_arg(a, b)
+
+    def test_tuple_of_namedtuple_arg(self):
+        Point = namedtuple("Point", ("x", "y", "z"))
+        a = (Point(x=1.0, y=2.0, z=3.0), Point(x=1.5, y=2.5, z=3.5))
+        b = (Point(x=4.0, y=5.0, z=6.0), Point(x=4.5, y=5.5, z=6.5))
+        self.check_tuple_arg(a, b)
+
+    def test_namedtuple_of_array_arg(self):
+        xs1 = np.arange(10, dtype=np.int32)
+        ys1 = xs1 + 2
+        xs2 = np.arange(10, dtype=np.int32) * 2
+        ys2 = xs2 + 1
+        Points = namedtuple("Points", ("xs", "ys"))
+        a = Points(xs=xs1, ys=ys1)
+        b = Points(xs=xs2, ys=ys2)
+        self.check_tuple_arg(a, b)
+
+    def test_name_attribute(self):
+        for vectorize in vectorize_funcs:
+
+            @vectorize("f8(f8)")
+            def bar(x):
+                return x**2
+
+            self.assertEqual(bar.__name__, "bar")
+
+    def test_no_transfer_for_device_data(self):
+        for vectorize in vectorize_funcs:
+            # Initialize test data on the device prior to banning host <-> device
+            # transfer
+
+            noise = np.random.randn(1, 3, 64, 64).astype(np.float32)
+            noise = cuda.to_device(noise)
+
+            # A mock of a CUDA function that always raises a CudaAPIError
+
+            def raising_transfer(*args, **kwargs):
+                raise CudaAPIError(999, "Transfer not allowed")
+
+            # Use the mock for transfers between the host and device
+
+            old_HtoD = getattr(driver, "cuMemcpyHtoD", None)
+            old_DtoH = getattr(driver, "cuMemcpyDtoH", None)
+
+            setattr(driver, "cuMemcpyHtoD", raising_transfer)
+            setattr(driver, "cuMemcpyDtoH", raising_transfer)
+
+            # Ensure that the mock functions are working as expected
+
+            with self.assertRaisesRegex(CudaAPIError, "Transfer not allowed"):
+                noise.copy_to_host()
+
+            with self.assertRaisesRegex(CudaAPIError, "Transfer not allowed"):
+                cuda.to_device([1])
+
+            try:
+                # Check that defining and calling a ufunc with data on the device
+                # induces no transfers
+
+                @vectorize(["float32(float32)"])
+                def func(noise):
+                    return noise + 1.0
+
+                func(noise)
+            finally:
+                # Replace our mocks with the original implementations. If there was
+                # no original implementation, simply remove ours.
+
+                if old_HtoD is not None:
+                    setattr(driver, "cuMemcpyHtoD", old_HtoD)
+                else:
+                    del driver.cuMemcpyHtoD
+                if old_DtoH is not None:
+                    setattr(driver, "cuMemcpyDtoH", old_DtoH)
+                else:
+                    del driver.cuMemcpyDtoH
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,23 @@
+# 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 vectorize
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+import unittest
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestVectorizeComplex(CUDATestCase):
+    def test_vectorize_complex(self):
+        @vectorize(["complex128(complex128)"], target="cuda")
+        def vcomp(a):
+            return a * a + 1.0
+
+        A = np.arange(5, dtype=np.complex128)
+        B = vcomp(A)
+        self.assertTrue(np.allclose(A * A + 1.0, B))
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+import math
+
+from numba import cuda
+from numba.cuda import vectorize, int32, uint32, float32, float64
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+from numba.cuda.tests.support import CheckWarningsMixin
+
+import unittest
+
+
+pi = math.pi
+
+
+def sinc(x):
+    if x == 0.0:
+        return 1.0
+    else:
+        return math.sin(x * pi) / (pi * x)
+
+
+def scaled_sinc(x, scale):
+    if x == 0.0:
+        return scale
+    else:
+        return scale * (math.sin(x * pi) / (pi * x))
+
+
+def vector_add(a, b):
+    return a + b
+
+
+class BaseVectorizeDecor(object):
+    target = None
+    wrapper = None
+    funcs = {
+        "func1": sinc,
+        "func2": scaled_sinc,
+        "func3": vector_add,
+    }
+
+    @classmethod
+    def _run_and_compare(cls, func, sig, A, *args, **kwargs):
+        if cls.wrapper is not None:
+            func = cls.wrapper(func)
+        numba_func = vectorize(sig, target=cls.target)(func)
+        numpy_func = np.vectorize(func)
+        result = numba_func(A, *args)
+        gold = numpy_func(A, *args)
+        np.testing.assert_allclose(result, gold, **kwargs)
+
+    def test_1(self):
+        sig = ["float64(float64)", "float32(float32)"]
+        func = self.funcs["func1"]
+        A = np.arange(100, dtype=np.float64)
+        self._run_and_compare(func, sig, A)
+
+    def test_2(self):
+        sig = [float64(float64), float32(float32)]
+        func = self.funcs["func1"]
+        A = np.arange(100, dtype=np.float64)
+        self._run_and_compare(func, sig, A)
+
+    def test_3(self):
+        sig = ["float64(float64, uint32)"]
+        func = self.funcs["func2"]
+        A = np.arange(100, dtype=np.float64)
+        scale = np.uint32(3)
+        self._run_and_compare(func, sig, A, scale, atol=1e-8)
+
+    def test_4(self):
+        sig = [
+            int32(int32, int32),
+            uint32(uint32, uint32),
+            float32(float32, float32),
+            float64(float64, float64),
+        ]
+        func = self.funcs["func3"]
+        A = np.arange(100, dtype=np.float64)
+        self._run_and_compare(func, sig, A, A)
+        A = A.astype(np.float32)
+        self._run_and_compare(func, sig, A, A)
+        A = A.astype(np.int32)
+        self._run_and_compare(func, sig, A, A)
+        A = A.astype(np.uint32)
+        self._run_and_compare(func, sig, A, A)
+
+
+class BaseVectorizeNopythonArg(unittest.TestCase, CheckWarningsMixin):
+    """
+    Test passing the nopython argument to the vectorize decorator.
+    """
+
+    def _test_target_nopython(self, target, warnings, with_sig=True):
+        a = np.array([2.0], dtype=np.float32)
+        b = np.array([3.0], dtype=np.float32)
+        sig = [float32(float32, float32)]
+        args = with_sig and [sig] or []
+        with self.check_warnings(warnings):
+            f = vectorize(*args, target=target, nopython=True)(vector_add)
+            f(a, b)
+
+
+class BaseVectorizeUnrecognizedArg(unittest.TestCase, CheckWarningsMixin):
+    """
+    Test passing an unrecognized argument to the vectorize decorator.
+    """
+
+    def _test_target_unrecognized_arg(self, target, with_sig=True):
+        a = np.array([2.0], dtype=np.float32)
+        b = np.array([3.0], dtype=np.float32)
+        sig = [float32(float32, float32)]
+        args = with_sig and [sig] or []
+        with self.assertRaises(KeyError) as raises:
+            f = vectorize(*args, target=target, nonexistent=2)(vector_add)
+            f(a, b)
+        self.assertIn("Unrecognized options", str(raises.exception))
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestVectorizeDecor(CUDATestCase, BaseVectorizeDecor):
+    """
+    Runs the tests from BaseVectorizeDecor with the CUDA target.
+    """
+
+    target = "cuda"
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestGPUVectorizeBroadcast(CUDATestCase):
+    def test_broadcast(self):
+        a = np.random.randn(100, 3, 1)
+        b = a.transpose(2, 1, 0)
+
+        def fn(a, b):
+            return a - b
+
+        @vectorize(["float64(float64,float64)"], target="cuda")
+        def fngpu(a, b):
+            return a - b
+
+        expect = fn(a, b)
+        got = fngpu(a, b)
+        np.testing.assert_almost_equal(expect, got)
+
+    def test_device_broadcast(self):
+        """
+        Same test as .test_broadcast() but with device array as inputs
+        """
+
+        a = np.random.randn(100, 3, 1)
+        b = a.transpose(2, 1, 0)
+
+        def fn(a, b):
+            return a - b
+
+        @vectorize(["float64(float64,float64)"], target="cuda")
+        def fngpu(a, b):
+            return a - b
+
+        expect = fn(a, b)
+        got = fngpu(cuda.to_device(a), cuda.to_device(b))
+        np.testing.assert_almost_equal(expect, got.copy_to_host())
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestVectorizeNopythonArg(BaseVectorizeNopythonArg, CUDATestCase):
+    def test_target_cuda_nopython(self):
+        warnings = ["nopython kwarg for cuda target is redundant"]
+        self._test_target_nopython("cuda", warnings)
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestVectorizeUnrecognizedArg(BaseVectorizeUnrecognizedArg, CUDATestCase):
+    def test_target_cuda_unrecognized_arg(self):
+        self._test_target_unrecognized_arg("cuda")
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda import vectorize
+from numba import cuda
+from numba.cuda import float32
+import numpy as np
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+import unittest
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestCudaVectorizeDeviceCall(CUDATestCase):
+    def test_cuda_vectorize_device_call(self):
+        @cuda.jit(float32(float32, float32, float32), device=True)
+        def cu_device_fn(x, y, z):
+            return x**y / z
+
+        def cu_ufunc(x, y, z):
+            return cu_device_fn(x, y, z)
+
+        ufunc = vectorize([float32(float32, float32, float32)], target="cuda")(
+            cu_ufunc
+        )
+
+        N = 100
+
+        X = np.array(np.random.sample(N), dtype=np.float32)
+        Y = np.array(np.random.sample(N), dtype=np.float32)
+        Z = np.array(np.random.sample(N), dtype=np.float32) + 0.1
+
+        out = ufunc(X, Y, Z)
+
+        gold = (X**Y) / Z
+
+        self.assertTrue(np.allclose(out, gold))
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,40 @@
+# 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 vectorize
+from numba import cuda
+from numba.cuda import float64
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+import unittest
+
+sig = [float64(float64, float64)]
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestCUDAVectorizeScalarArg(CUDATestCase):
+    def test_vectorize_scalar_arg(self):
+        @vectorize(sig, target="cuda")
+        def vector_add(a, b):
+            return a + b
+
+        A = np.arange(10, dtype=np.float64)
+        dA = cuda.to_device(A)
+        v = vector_add(1.0, dA)
+
+        np.testing.assert_array_almost_equal(
+            v.copy_to_host(), np.arange(1, 11, dtype=np.float64)
+        )
+
+    def test_vectorize_all_scalars(self):
+        @vectorize(sig, target="cuda")
+        def vector_add(a, b):
+            return a + b
+
+        v = vector_add(1.0, 1.0)
+
+        np.testing.assert_almost_equal(2.0, v)
+
+
+if __name__ == "__main__":
+    unittest.main()