numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""Example: sum each row using guvectorize
+
+See Numpy documentation for detail about gufunc:
+    http://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html
+"""
+
+import numpy as np
+from numba import cuda, guvectorize
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+import unittest
+
+
+@skip_on_cudasim("ufunc API unsupported in the simulator")
+class TestGUFuncScalar(CUDATestCase):
+    def test_gufunc_scalar_output(self):
+        #    function type:
+        #        - has no void return type
+        #        - array argument is one dimension fewer than the source array
+        #        - scalar output is passed as a 1-element array.
+        #
+        #    signature: (n)->()
+        #        - the function takes an array of n-element and output a scalar.
+
+        @guvectorize(["void(int32[:], int32[:])"], "(n)->()", target="cuda")
+        def sum_row(inp, out):
+            tmp = 0.0
+            for i in range(inp.shape[0]):
+                tmp += inp[i]
+            out[0] = tmp
+
+        # inp is (10000, 3)
+        # out is (10000)
+        # The outer (leftmost) dimension must match or numpy broadcasting
+        # is performed. But, broadcasting on CUDA arrays is not supported.
+
+        inp = np.arange(300, dtype=np.int32).reshape(100, 3)
+
+        # invoke on CUDA with manually managed memory
+        out1 = np.empty(100, dtype=inp.dtype)
+        out2 = np.empty(100, dtype=inp.dtype)
+
+        dev_inp = cuda.to_device(inp)  # alloc and copy input data
+        dev_out1 = cuda.to_device(out1, copy=False)  # alloc only
+
+        sum_row(dev_inp, out=dev_out1)  # invoke the gufunc
+        dev_out2 = sum_row(dev_inp)  # invoke the gufunc
+
+        dev_out1.copy_to_host(out1)  # retrieve the result
+        dev_out2.copy_to_host(out2)  # retrieve the result
+
+        # verify result
+        for i in range(inp.shape[0]):
+            self.assertTrue(out1[i] == inp[i].sum())
+            self.assertTrue(out2[i] == inp[i].sum())
+
+    def test_gufunc_scalar_output_bug(self):
+        # Issue 2812: Error due to using input argument types as output argument
+        @guvectorize(["void(int32, int32[:])"], "()->()", target="cuda")
+        def twice(inp, out):
+            out[0] = inp * 2
+
+        self.assertEqual(twice(10), 20)
+        arg = np.arange(10).astype(np.int32)
+        self.assertPreciseEqual(twice(arg), arg * 2)
+
+    def test_gufunc_scalar_input_saxpy(self):
+        @guvectorize(
+            ["void(float32, float32[:], float32[:], float32[:])"],
+            "(),(t),(t)->(t)",
+            target="cuda",
+        )
+        def saxpy(a, x, y, out):
+            for i in range(out.shape[0]):
+                out[i] = a * x[i] + y[i]
+
+        A = np.float32(2)
+        X = np.arange(10, dtype=np.float32).reshape(5, 2)
+        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
+        out = saxpy(A, X, Y)
+
+        for j in range(5):
+            for i in range(2):
+                exp = A * X[j, i] + Y[j, i]
+                self.assertTrue(exp == out[j, i])
+
+        X = np.arange(10, dtype=np.float32)
+        Y = np.arange(10, dtype=np.float32)
+        out = saxpy(A, X, Y)
+
+        for j in range(10):
+            exp = A * X[j] + Y[j]
+            self.assertTrue(exp == out[j], (exp, out[j]))
+
+        A = np.arange(5, dtype=np.float32)
+        X = np.arange(10, dtype=np.float32).reshape(5, 2)
+        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
+        out = saxpy(A, X, Y)
+
+        for j in range(5):
+            for i in range(2):
+                exp = A[j] * X[j, i] + Y[j, i]
+                self.assertTrue(exp == out[j, i], (exp, out[j, i]))
+
+    def test_gufunc_scalar_cast(self):
+        @guvectorize(
+            ["void(int32, int32[:], int32[:])"], "(),(t)->(t)", target="cuda"
+        )
+        def foo(a, b, out):
+            for i in range(b.size):
+                out[i] = a * b[i]
+
+        a = np.int64(2)  # type does not match signature (int32)
+        b = np.arange(10).astype(np.int32)
+        out = foo(a, b)
+        np.testing.assert_equal(out, a * b)
+
+        # test error
+        a = np.array(a)
+        da = cuda.to_device(a)
+        self.assertEqual(da.dtype, np.int64)
+        with self.assertRaises(TypeError) as raises:
+            foo(da, b)
+
+        self.assertIn("does not support .astype()", str(raises.exception))
+
+    def test_gufunc_old_style_scalar_as_array(self):
+        # Example from issue #2579
+        @guvectorize(
+            ["void(int32[:],int32[:],int32[:])"], "(n),()->(n)", target="cuda"
+        )
+        def gufunc(x, y, res):
+            for i in range(x.shape[0]):
+                res[i] = x[i] + y[0]
+
+        # Case 1
+        a = np.array([1, 2, 3, 4], dtype=np.int32)
+        b = np.array([2], dtype=np.int32)
+
+        res = np.zeros(4, dtype=np.int32)
+
+        expected = res.copy()
+        expected = a + b
+
+        gufunc(a, b, out=res)
+
+        np.testing.assert_almost_equal(expected, res)
+
+        # Case 2
+        a = np.array([1, 2, 3, 4] * 2, dtype=np.int32).reshape(2, 4)
+        b = np.array([2, 10], dtype=np.int32)
+
+        res = np.zeros((2, 4), dtype=np.int32)
+
+        expected = res.copy()
+        expected[0] = a[0] + b[0]
+        expected[1] = a[1] + b[1]
+
+        gufunc(a, b, res)
+
+        np.testing.assert_almost_equal(expected, res)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.deviceufunc import GUFuncEngine
+import unittest
+
+
+def template(signature, shapes, expects):
+    gufb = GUFuncEngine.from_signature(signature)
+    sch = gufb.schedule(shapes)
+    for k, v in expects.items():
+        got = getattr(sch, k)
+        if got != v:
+            fmt = "error for %s: got=%s but expect=%s"
+            raise AssertionError(fmt % (k, got, v))
+
+
+class TestGUFuncScheduling(unittest.TestCase):
+    def test_signature_1(self):
+        signature = "(m, n), (n, p) -> (m, p)"
+        shapes = (100, 4, 5), (1, 5, 7)
+        expects = dict(
+            ishapes=[(4, 5), (5, 7)],
+            oshapes=[(4, 7)],
+            loopdims=(100,),
+            pinned=[False, True],
+        )
+        template(signature, shapes, expects)
+
+    def test_signature_2(self):
+        signature = "(m, n), (n, p) -> (m, p)"
+        shapes = (100, 4, 5), (100, 5, 7)
+        expects = dict(
+            ishapes=[(4, 5), (5, 7)],
+            oshapes=[(4, 7)],
+            loopdims=(100,),
+            pinned=[False, False],
+        )
+        template(signature, shapes, expects)
+
+    def test_signature_3(self):
+        signature = "(m, n), (n, p) -> (m, p)"
+        shapes = (12, 34, 4, 5), (12, 34, 5, 7)
+        expects = dict(
+            ishapes=[(4, 5), (5, 7)],
+            oshapes=[(4, 7)],
+            loopdims=(12, 34),
+            pinned=[False, False],
+        )
+        template(signature, shapes, expects)
+
+    def test_signature_4(self):
+        signature = "(m, n), (n, p) -> (m, p)"
+        shapes = (4, 5), (5, 7)
+        expects = dict(
+            ishapes=[(4, 5), (5, 7)],
+            oshapes=[(4, 7)],
+            loopdims=(),
+            pinned=[False, False],
+        )
+        template(signature, shapes, expects)
+
+    def test_signature_5(self):
+        signature = "(a), (a) -> (a)"
+        shapes = (5,), (5,)
+        expects = dict(
+            ishapes=[(5,), (5,)],
+            oshapes=[(5,)],
+            loopdims=(),
+            pinned=[False, False],
+        )
+        template(signature, shapes, expects)
+
+    def test_signature_6(self):
+        signature = "(), () -> ()"
+        shapes = (5,), (5,)
+        expects = dict(
+            ishapes=[(), ()], oshapes=[()], loopdims=(5,), pinned=[False, False]
+        )
+        template(signature, shapes, expects)
+
+    def test_signature_7(self):
+        signature = "(), () -> ()"
+        shapes = (5,), ()
+        expects = dict(
+            ishapes=[(), ()], oshapes=[()], loopdims=(5,), pinned=[False, True]
+        )
+        template(signature, shapes, expects)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,39 @@
+# 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 float32, float64, int32, void
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestCudaIDiv(CUDATestCase):
+    def test_inplace_div(self):
+        @cuda.jit(void(float32[:, :], int32, int32))
+        def div(grid, l_x, l_y):
+            for x in range(l_x):
+                for y in range(l_y):
+                    grid[x, y] /= 2.0
+
+        x = np.ones((2, 2), dtype=np.float32)
+        grid = cuda.to_device(x)
+        div[1, 1](grid, 2, 2)
+        y = grid.copy_to_host()
+        self.assertTrue(np.all(y == 0.5))
+
+    def test_inplace_div_double(self):
+        @cuda.jit(void(float64[:, :], int32, int32))
+        def div_double(grid, l_x, l_y):
+            for x in range(l_x):
+                for y in range(l_y):
+                    grid[x, y] /= 2.0
+
+        x = np.ones((2, 2), dtype=np.float64)
+        grid = cuda.to_device(x)
+        div_double[1, 1](grid, 2, 2)
+        y = grid.copy_to_host()
+        self.assertTrue(np.all(y == 0.5))
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,170 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import re
+import numpy as np
+from numba import cuda
+from numba.cuda import types
+from numba.cuda.testing import (
+    unittest,
+    CUDATestCase,
+    skip_on_cudasim,
+)
+
+
+@skip_on_cudasim("Cudasim does not support inline and forceinline")
+class TestCudaInline(CUDATestCase):
+    def _test_call_inline(self, inline, inline_expected):
+        """Test @cuda.jit(inline=...)"""
+        a = np.ones(2, dtype=np.int32)
+
+        sig = (types.int32[::1],)
+
+        @cuda.jit(inline=inline)
+        def set_zero(a):
+            a[0] = 0
+
+        @cuda.jit(sig)
+        def call_set_zero(a):
+            set_zero(a)
+
+        call_set_zero[1, 2](a)
+
+        expected = np.arange(2, dtype=np.int32)
+        self.assertTrue(np.all(a == expected))
+
+        llvm_ir = call_set_zero.inspect_llvm(sig)
+        pat = r"call [a-zA-Z0-9]* @"
+        match = re.compile(pat).search(llvm_ir)
+
+        if inline_expected:
+            # check that call was inlined
+            self.assertIsNone(match, msg=llvm_ir)
+        else:
+            # check that call was not inlined
+            self.assertIsNotNone(match, msg=llvm_ir)
+
+        # alwaysinline should not be in the IR when the inline kwarg is used
+        self.assertNotIn("alwaysinline", llvm_ir)
+
+    def test_call_inline_always(self):
+        self._test_call_inline("always", True)
+
+    def test_call_inline_never(self):
+        self._test_call_inline("never", False)
+
+    def test_call_inline_true(self):
+        self._test_call_inline(True, True)
+
+    def test_call_inline_false(self):
+        self._test_call_inline(False, False)
+
+    def test_call_inline_costmodel_false(self):
+        def cost_model(expr, caller_info, callee_info):
+            return False
+
+        self._test_call_inline(cost_model, False)
+
+    def test_call_inline_costmodel_true(self):
+        def cost_model(expr, caller_info, callee_info):
+            return True
+
+        self._test_call_inline(cost_model, True)
+
+    def _test_call_forceinline(self, forceinline):
+        """Test @cuda.jit(forceinline=...)"""
+        a = np.ones(2, dtype=np.int32)
+
+        sig = (types.int32[::1],)
+
+        @cuda.jit(forceinline=forceinline)
+        def set_zero(a):
+            a[0] = 0
+
+        @cuda.jit(sig)
+        def call_set_zero(a):
+            set_zero(a)
+
+        call_set_zero[1, 2](a)
+
+        expected = np.arange(2, dtype=np.int32)
+        self.assertTrue(np.all(a == expected))
+
+        llvm_ir = call_set_zero.inspect_llvm(sig)
+        pat = r"call [a-zA-Z0-9]* @"
+        match = re.compile(pat).search(llvm_ir)
+
+        # Check that call was not inlined at the Numba IR level - the call
+        # should still be present in the IR
+        self.assertIsNotNone(match)
+
+        # Check the definition of set_zero - it is a definition where the
+        # name does not include an underscore just before "set_zero", because
+        # that would match the "call_set_zero" definition
+        pat = r"define.*[^_]set_zero.*"
+        match = re.compile(pat).search(llvm_ir)
+        self.assertIsNotNone(match)
+        if forceinline:
+            self.assertIn("alwaysinline", match.group())
+        else:
+            self.assertNotIn("alwaysinline", match.group())
+
+        # The kernel, "call_set_zero", should never have "alwaysinline" set
+        pat = r"define.*call_set_zero.*"
+        match = re.compile(pat).search(llvm_ir)
+        self.assertIsNotNone(match)
+        self.assertNotIn("alwaysinline", match.group())
+
+    def test_call_forceinline_true(self):
+        self._test_call_forceinline(True)
+
+    def test_call_forceinline_false(self):
+        self._test_call_forceinline(False)
+
+    def test_compile_forceinline_ltoir_only(self):
+        def set_zero(a):
+            a[0] = 0
+
+        args = (types.float32[::1],)
+        msg = r"Can only designate forced inlining in LTO-IR"
+        with self.assertRaisesRegex(ValueError, msg):
+            cuda.compile(
+                set_zero,
+                args,
+                device=True,
+                forceinline=True,
+            )
+
+    def _compile_set_zero(self, forceinline):
+        def set_zero(a):
+            a[0] = 0
+
+        args = (types.float32[::1],)
+        ltoir, resty = cuda.compile(
+            set_zero,
+            args,
+            device=True,
+            output="ltoir",
+            forceinline=forceinline,
+        )
+
+        # Sanity check
+        self.assertEqual(resty, types.none)
+
+        return ltoir
+
+    def test_compile_forceinline(self):
+        ltoir_noinline = self._compile_set_zero(False)
+        ltoir_forceinline = self._compile_set_zero(True)
+
+        # As LTO-IR is opaque, the best we can do is check that changing the
+        # flag resulted in a change in the generated LTO-IR in some way.
+        self.assertNotEqual(
+            ltoir_noinline,
+            ltoir_forceinline,
+            "forceinline flag appeared to have no effect on LTO-IR",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()