numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,231 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numba.cuda.testing import CUDATestCase
+import numpy as np
+import sys
+
+
+class UseCase:
+    """
+    Provide a way to call a kernel as if it were a function.
+
+    This allows the CUDA cache tests to closely match the CPU cache tests, and
+    also to support calling cache use cases as njitted functions. The class
+    wraps a function that takes an array for the return value and arguments,
+    and provides an interface that accepts arguments, launches the kernel
+    appropriately, and returns the stored return value.
+
+    The return type is inferred from the type of the first argument, unless it
+    is explicitly overridden by the ``retty`` kwarg.
+    """
+
+    def __init__(self, func, retty=None):
+        self._func = func
+        self._retty = retty
+
+    def __call__(self, *args):
+        array_args = [np.asarray(arg) for arg in args]
+        if self._retty:
+            array_return = np.ndarray((), dtype=self._retty)
+        else:
+            array_return = np.zeros_like(array_args[0])
+
+        self._call(array_return, *array_args)
+        return array_return[()]
+
+    @property
+    def func(self):
+        return self._func
+
+
+class CUDAUseCase(UseCase):
+    def _call(self, ret, *args):
+        self._func[1, 1](ret, *args)
+
+
+@cuda.jit(cache=True)
+def add_usecase_kernel(r, x, y):
+    r[()] = x[()] + y[()] + Z
+
+
+@cuda.jit(cache=False)
+def add_nocache_usecase_kernel(r, x, y):
+    r[()] = x[()] + y[()] + Z
+
+
+add_usecase = CUDAUseCase(add_usecase_kernel)
+add_nocache_usecase = CUDAUseCase(add_nocache_usecase_kernel)
+
+Z = 1
+
+
+# Inner / outer cached / uncached cases
+
+
+@cuda.jit(cache=True)
+def inner(x, y):
+    return x + y + Z
+
+
+@cuda.jit(cache=True)
+def outer_kernel(r, x, y):
+    r[()] = inner(-y[()], x[()])
+
+
+@cuda.jit(cache=False)
+def outer_uncached_kernel(r, x, y):
+    r[()] = inner(-y[()], x[()])
+
+
+outer = CUDAUseCase(outer_kernel)
+outer_uncached = CUDAUseCase(outer_uncached_kernel)
+
+
+# Exercise returning a record instance.  This used to hardcode the dtype
+# pointer's value in the bitcode.
+
+packed_record_type = np.dtype([("a", np.int8), ("b", np.float64)])
+aligned_record_type = np.dtype([("a", np.int8), ("b", np.float64)], align=True)
+
+packed_arr = np.empty(2, dtype=packed_record_type)
+for i in range(packed_arr.size):
+    packed_arr[i]["a"] = i + 1
+    packed_arr[i]["b"] = i + 42.5
+
+aligned_arr = np.array(packed_arr, dtype=aligned_record_type)
+
+
+@cuda.jit(cache=True)
+def record_return(r, ary, i):
+    r[()] = ary[i]
+
+
+record_return_packed = CUDAUseCase(record_return, retty=packed_record_type)
+record_return_aligned = CUDAUseCase(record_return, retty=aligned_record_type)
+
+
+# Closure test cases
+
+
+def make_closure(x):
+    @cuda.jit(cache=True)
+    def closure(r, y):
+        r[()] = x + y[()]
+
+    return CUDAUseCase(closure)
+
+
+closure1 = make_closure(3)
+closure2 = make_closure(5)
+closure3 = make_closure(7)
+closure4 = make_closure(9)
+
+
+# Ambiguous / renamed functions
+
+
+@cuda.jit(cache=True)
+def ambiguous_function(r, x):
+    r[()] = x[()] + 2
+
+
+renamed_function1 = CUDAUseCase(ambiguous_function)
+
+
+@cuda.jit(cache=True)
+def ambiguous_function(r, x):
+    r[()] = x[()] + 6
+
+
+renamed_function2 = CUDAUseCase(ambiguous_function)
+
+
+@cuda.jit(cache=True)
+def many_locals():
+    aa = cuda.local.array((1, 1), np.float64)
+    ab = cuda.local.array((1, 1), np.float64)
+    ac = cuda.local.array((1, 1), np.float64)
+    ad = cuda.local.array((1, 1), np.float64)
+    ae = cuda.local.array((1, 1), np.float64)
+    af = cuda.local.array((1, 1), np.float64)
+    ag = cuda.local.array((1, 1), np.float64)
+    ah = cuda.local.array((1, 1), np.float64)
+    ai = cuda.local.array((1, 1), np.float64)
+    aj = cuda.local.array((1, 1), np.float64)
+    ak = cuda.local.array((1, 1), np.float64)
+    al = cuda.local.array((1, 1), np.float64)
+    am = cuda.local.array((1, 1), np.float64)
+    an = cuda.local.array((1, 1), np.float64)
+    ao = cuda.local.array((1, 1), np.float64)
+    ap = cuda.local.array((1, 1), np.float64)
+    ar = cuda.local.array((1, 1), np.float64)
+    at = cuda.local.array((1, 1), np.float64)
+    au = cuda.local.array((1, 1), np.float64)
+    av = cuda.local.array((1, 1), np.float64)
+    aw = cuda.local.array((1, 1), np.float64)
+    ax = cuda.local.array((1, 1), np.float64)
+    ay = cuda.local.array((1, 1), np.float64)
+    az = cuda.local.array((1, 1), np.float64)
+
+    aa[:] = 0
+    ab[:] = 0
+    ac[:] = 0
+    ad[:] = 0
+    ae[:] = 0
+    af[:] = 0
+    ag[:] = 0
+    ah[:] = 0
+    ai[:] = 0
+    aj[:] = 0
+    ak[:] = 0
+    al[:] = 0
+    am[:] = 0
+    an[:] = 0
+    ao[:] = 0
+    ap[:] = 0
+    ar[:] = 0
+    at[:] = 0
+    au[:] = 0
+    av[:] = 0
+    aw[:] = 0
+    ax[:] = 0
+    ay[:] = 0
+    az[:] = 0
+
+
+# Simple use case for multiprocessing test
+
+
+@cuda.jit(cache=True)
+def simple_usecase_kernel(r, x):
+    r[()] = x[()]
+
+
+simple_usecase_caller = CUDAUseCase(simple_usecase_kernel)
+
+
+class _TestModule(CUDATestCase):
+    """
+    Tests for functionality of this module's functions.
+    Note this does not define any "test_*" method, instead check_module()
+    should be called by hand.
+    """
+
+    def check_module(self, mod):
+        self.assertPreciseEqual(mod.add_usecase(2, 3), 6)
+        self.assertPreciseEqual(mod.outer_uncached(3, 2), 2)
+        self.assertPreciseEqual(mod.outer(3, 2), 2)
+
+        packed_rec = mod.record_return_packed(mod.packed_arr, 1)
+        self.assertPreciseEqual(tuple(packed_rec), (2, 43.5))
+        aligned_rec = mod.record_return_aligned(mod.aligned_arr, 1)
+        self.assertPreciseEqual(tuple(aligned_rec), (2, 43.5))
+
+        mod.simple_usecase_caller(2)
+
+
+def self_test():
+    mod = sys.modules[__name__]
+    _TestModule().check_module(mod)

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

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+
+from numba import cuda
+from numba.cuda import HAS_NUMBA
+from numba.cuda.testing import CUDATestCase, skip_on_standalone_numba_cuda
+from numba.cuda.tests.cudapy.cache_usecases import CUDAUseCase, UseCase
+
+
+class CPUUseCase(UseCase):
+    def _call(self, ret, *args):
+        self._func(ret, *args)
+
+
+# Using the same function as a cached CPU and CUDA-jitted function
+
+
+def target_shared_assign(r, x):
+    r[()] = x[()]
+
+
+assign_cuda_kernel = cuda.jit(cache=True)(target_shared_assign)
+assign_cuda = CUDAUseCase(assign_cuda_kernel)
+if HAS_NUMBA:
+    from numba import njit
+
+    assign_cpu_jitted = njit(cache=True)(target_shared_assign)
+    assign_cpu = CPUUseCase(assign_cpu_jitted)
+
+
+@skip_on_standalone_numba_cuda
+class _TestModule(CUDATestCase):
+    """
+    Tests for functionality of this module's functions.
+    Note this does not define any "test_*" method, instead check_module()
+    should be called by hand.
+    """
+
+    def check_module(self, mod):
+        self.assertPreciseEqual(mod.assign_cpu(5), 5)
+        self.assertPreciseEqual(mod.assign_cpu(5.5), 5.5)
+        self.assertPreciseEqual(mod.assign_cuda(5), 5)
+        self.assertPreciseEqual(mod.assign_cuda(5.5), 5.5)
+
+
+def self_test():
+    mod = sys.modules[__name__]
+    _TestModule().check_module(mod)

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

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numba.cuda.testing import CUDATestCase
+import sys
+
+from numba.cuda.tests.cudapy.cache_usecases import CUDAUseCase
+
+
+# Usecase with cooperative groups
+
+
+@cuda.jit(cache=True)
+def cg_usecase_kernel(r, x):
+    grid = cuda.cg.this_grid()
+    grid.sync()
+
+
+cg_usecase = CUDAUseCase(cg_usecase_kernel)
+
+
+class _TestModule(CUDATestCase):
+    """
+    Tests for functionality of this module's functions.
+    Note this does not define any "test_*" method, instead check_module()
+    should be called by hand.
+    """
+
+    def check_module(self, mod):
+        mod.cg_usecase(0)
+
+
+def self_test():
+    mod = sys.modules[__name__]
+    _TestModule().check_module(mod)

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

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import cmath
+
+
+def div_usecase(x, y):
+    return x / y
+
+
+def real_usecase(x):
+    return x.real
+
+
+def imag_usecase(x):
+    return x.imag
+
+
+def conjugate_usecase(x):
+    return x.conjugate()
+
+
+def acos_usecase(x):
+    return cmath.acos(x)
+
+
+def cos_usecase(x):
+    return cmath.cos(x)
+
+
+def asin_usecase(x):
+    return cmath.asin(x)
+
+
+def sin_usecase(x):
+    return cmath.sin(x)
+
+
+def atan_usecase(x):
+    return cmath.atan(x)
+
+
+def tan_usecase(x):
+    return cmath.tan(x)
+
+
+def acosh_usecase(x):
+    return cmath.acosh(x)
+
+
+def cosh_usecase(x):
+    return cmath.cosh(x)
+
+
+def asinh_usecase(x):
+    return cmath.asinh(x)
+
+
+def sinh_usecase(x):
+    return cmath.sinh(x)
+
+
+def atanh_usecase(x):
+    return cmath.atanh(x)
+
+
+def tanh_usecase(x):
+    return cmath.tanh(x)
+
+
+def exp_usecase(x):
+    return cmath.exp(x)
+
+
+def isfinite_usecase(x):
+    return cmath.isfinite(x)
+
+
+def isinf_usecase(x):
+    return cmath.isinf(x)
+
+
+def isnan_usecase(x):
+    return cmath.isnan(x)
+
+
+def log_usecase(x):
+    return cmath.log(x)
+
+
+def log_base_usecase(x, base):
+    return cmath.log(x, base)
+
+
+def log10_usecase(x):
+    return cmath.log10(x)
+
+
+def phase_usecase(x):
+    return cmath.phase(x)
+
+
+def polar_usecase(x):
+    return cmath.polar(x)
+
+
+def polar_as_complex_usecase(x):
+    return complex(*cmath.polar(x))
+
+
+def rect_usecase(r, phi):
+    return cmath.rect(r, phi)
+
+
+def sqrt_usecase(x):
+    return cmath.sqrt(x)

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

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from enum import Enum, IntEnum
+
+
+class Color(Enum):
+    red = 1
+    green = 2
+    blue = 3
+
+
+class Shake(Enum):
+    vanilla = 7
+    chocolate = 4
+    cookies = 9
+    # Same as Color.blue
+    mint = 3
+
+
+class Planet(Enum):
+    MERCURY = (3.303e23, 2.4397e6)
+    VENUS = (4.869e24, 6.0518e6)
+    EARTH = (5.976e24, 6.37814e6)
+    MARS = (6.421e23, 3.3972e6)
+    JUPITER = (1.9e27, 7.1492e7)
+    SATURN = (5.688e26, 6.0268e7)
+    URANUS = (8.686e25, 2.5559e7)
+    NEPTUNE = (1.024e26, 2.4746e7)
+
+
+class HeterogeneousEnum(Enum):
+    red = 1.0
+    green = 2.0
+    blue = 3j
+
+
+class Shape(IntEnum):
+    # Same as Color.green
+    circle = 2
+    # Same as RequestError.internal_error
+    square = 500
+
+
+class RequestError(IntEnum):
+    dummy = 2
+    not_found = 404
+    internal_error = 500
+
+
+class IntEnumWithNegatives(IntEnum):
+    # Used for testing of hash, need to make sure -1 -> -2 to comply with CPy
+    one = 1
+    two = 2
+    too = 2
+    three = 3
+    negone = -1
+    negtwo = -2
+    negthree = -3

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

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda import types
+from numba.cuda.core import config
+
+
+class MyStruct:
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+
+
+class StructModelType(types.Type):
+    def __init__(self):
+        super().__init__(name="TestStructModelType")
+
+
+struct_model_type = StructModelType()
+
+
+if not config.ENABLE_CUDASIM:
+    from numba.cuda import int32
+    from numba.cuda.extending import (
+        core_models,
+        typeof_impl,
+        type_callable,
+    )
+    from numba.cuda.extending import (
+        register_model,
+        make_attribute_wrapper,
+    )
+    from numba.cuda.cudaimpl import lower
+    from numba.cuda import cgutils
+
+    @typeof_impl.register(MyStruct)
+    def typeof_teststruct(val, c):
+        return struct_model_type
+
+    @register_model(StructModelType)
+    class TestStructModel(core_models.StructModel):
+        def __init__(self, dmm, fe_type):
+            members = [("x", int32), ("y", int32)]
+            super().__init__(dmm, fe_type, members)
+
+    make_attribute_wrapper(StructModelType, "x", "x")
+    make_attribute_wrapper(StructModelType, "y", "y")
+
+    @type_callable(MyStruct)
+    def type_test_struct(context):
+        def typer(x, y):
+            if isinstance(x, types.Integer) and isinstance(y, types.Integer):
+                return struct_model_type
+
+        return typer
+
+    @lower(MyStruct, types.Integer, types.Integer)
+    def lower_test_type_ctor(context, builder, sig, args):
+        obj = cgutils.create_struct_proxy(struct_model_type)(context, builder)
+        obj.x = args[0]
+        obj.y = args[1]
+        return obj._getvalue()

numba_cuda/numba/cuda/tests/cudapy/jitlink.ptx

@@ -0,0 +1,28 @@
+//
+// Generated by NVIDIA NVVM Compiler
+// Compiler built on Tue Apr  1 03:34:02 2014 (1396341242)
+// Cuda compilation tools, release 6.0, V6.0.1
+//
+
+.version 4.0
+.target sm_20
+.address_size 64
+
+
+.visible .func  (.param .b32 func_retval0) bar(
+	.param .b64 bar_param_0,
+	.param .b32 bar_param_1
+)
+{
+	.reg .s32 	%r<4>;
+	.reg .s64 	%rd<2>;
+
+
+	ld.param.u64 	%rd1, [bar_param_0];
+	ld.param.u32 	%r1, [bar_param_1];
+	shl.b32 	%r2, %r1, 1;
+	st.u32 	[%rd1], %r2;
+	mov.u32 	%r3, 0;
+	st.param.b32	[func_retval0+0], %r3;
+	ret;
+}

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

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+# Python 3 syntax only use cases, used in test_extending.py
+
+# arg name is different, and there's no arg name to match before *args
+
+
+def impl4(z, *args, kw=None):
+    if z > 10:
+        return 1
+    else:
+        return -1
+
+
+# arg name is different but at a detectable location, with *args
+
+
+def impl5(z, b, *args, kw=None):
+    if z > 10:
+        return 1
+    else:
+        return -1
+
+
+def var_positional_impl(a, *star_args_token, kw=None, kw1=12):
+    def impl(a, b, f, kw=None, kw1=12):
+        if a > 10:
+            return 1
+        else:
+            return -1
+
+    return impl

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

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Usecases of recursive functions in the CUDA target, many derived from
+numba/tests/recursion_usecases.py.
+
+Some functions are compiled at import time, hence a separate module.
+"""
+
+from numba import cuda
+
+
+@cuda.jit("i8(i8)", device=True)
+def fib1(n):
+    if n < 2:
+        return n
+    # Note the second call does not use a named argument, unlike the CPU target
+    # usecase
+    return fib1(n - 1) + fib1(n - 2)
+
+
+def make_fib2():
+    @cuda.jit("i8(i8)", device=True)
+    def fib2(n):
+        if n < 2:
+            return n
+        return fib2(n - 1) + fib2(n - 2)
+
+    return fib2
+
+
+fib2 = make_fib2()
+
+
+@cuda.jit
+def type_change_self(x, y):
+    if x > 1 and y > 0:
+        return x + type_change_self(x - y, y)
+    else:
+        return y
+
+
+# Implicit signature
+@cuda.jit(device=True)
+def fib3(n):
+    if n < 2:
+        return n
+
+    return fib3(n - 1) + fib3(n - 2)
+
+
+# Run-away self recursion
+@cuda.jit(device=True)
+def runaway_self(x):
+    return runaway_self(x)
+
+
+@cuda.jit(device=True)
+def raise_self(x):
+    if x == 1:
+        raise ValueError("raise_self")
+    elif x > 0:
+        return raise_self(x - 1)
+    else:
+        return 1
+
+
+@cuda.jit(debug=True, opt=False)
+def raise_self_kernel(x):
+    raise_self(x)
+
+
+def make_optional_return_case(jit=lambda x: x):
+    @jit
+    def foo(x):
+        if x > 5:
+            return x - 1
+        else:
+            return
+
+    @jit
+    def bar(x):
+        out = foo(x)
+        if out is None:
+            return out
+        elif out < 8:
+            return out
+        else:
+            return x * bar(out)
+
+    return bar
+
+
+def make_growing_tuple_case(jit=lambda x: x):
+    # From issue #4387
+    @jit
+    def make_list(n):
+        if n <= 0:
+            return None
+
+        return (n, make_list(n - 1))
+
+    return make_list