numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/typing/ctypes_utils.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Support for typing ctypes function pointers.
+"""
+
+import ctypes
+import sys
+
+from numba.cuda import types
+from numba.cuda.typing import templates
+
+_FROM_CTYPES = {
+    ctypes.c_bool: types.boolean,
+    ctypes.c_int8: types.int8,
+    ctypes.c_int16: types.int16,
+    ctypes.c_int32: types.int32,
+    ctypes.c_int64: types.int64,
+    ctypes.c_uint8: types.uint8,
+    ctypes.c_uint16: types.uint16,
+    ctypes.c_uint32: types.uint32,
+    ctypes.c_uint64: types.uint64,
+    ctypes.c_float: types.float32,
+    ctypes.c_double: types.float64,
+    ctypes.c_void_p: types.voidptr,
+    ctypes.py_object: types.ffi_forced_object,
+}
+
+
+_TO_CTYPES = {v: k for (k, v) in _FROM_CTYPES.items()}
+
+
+def from_ctypes(ctypeobj):
+    """
+    Convert the given ctypes type to a Numba type.
+    """
+    if ctypeobj is None:
+        # Special case for the restype of void-returning functions
+        return types.none
+
+    assert isinstance(ctypeobj, type), ctypeobj
+
+    def _convert_internal(ctypeobj):
+        # Recursive helper
+        if issubclass(ctypeobj, ctypes._Pointer):
+            valuety = _convert_internal(ctypeobj._type_)
+            if valuety is not None:
+                return types.CPointer(valuety)
+        else:
+            return _FROM_CTYPES.get(ctypeobj)
+
+    ty = _convert_internal(ctypeobj)
+    if ty is None:
+        raise TypeError("Unsupported ctypes type: %s" % ctypeobj)
+    return ty
+
+
+def to_ctypes(ty):
+    """
+    Convert the given Numba type to a ctypes type.
+    """
+    assert isinstance(ty, types.Type), ty
+
+    if ty is types.none:
+        # Special case for the restype of void-returning functions
+        return None
+
+    def _convert_internal(ty):
+        if isinstance(ty, types.CPointer):
+            return ctypes.POINTER(_convert_internal(ty.dtype))
+        else:
+            return _TO_CTYPES.get(ty)
+
+    ctypeobj = _convert_internal(ty)
+    if ctypeobj is None:
+        raise TypeError("Cannot convert Numba type '%s' to ctypes type" % (ty,))
+    return ctypeobj
+
+
+def is_ctypes_funcptr(obj):
+    try:
+        # Is it something of which we can get the address
+        ctypes.cast(obj, ctypes.c_void_p)
+    except ctypes.ArgumentError:
+        return False
+    else:
+        # Does it define argtypes and restype
+        return hasattr(obj, "argtypes") and hasattr(obj, "restype")
+
+
+def get_pointer(ctypes_func):
+    """
+    Get a pointer to the underlying function for a ctypes function as an
+    integer.
+    """
+    return ctypes.cast(ctypes_func, ctypes.c_void_p).value
+
+
+def make_function_type(cfnptr):
+    """
+    Return a Numba type for the given ctypes function pointer.
+    """
+    if cfnptr.argtypes is None:
+        raise TypeError(
+            "ctypes function %r doesn't define its argument types; "
+            "consider setting the `argtypes` attribute" % (cfnptr.__name__,)
+        )
+    cargs = [from_ctypes(a) for a in cfnptr.argtypes]
+    cret = from_ctypes(cfnptr.restype)
+    # void* return type is a int/long on 32 bit platforms and an int on 64 bit
+    # platforms, explicit conversion to a int64 should match.
+    if cret == types.voidptr:
+        cret = types.uintp
+    if sys.platform == "win32" and not cfnptr._flags_ & ctypes._FUNCFLAG_CDECL:
+        # 'stdcall' calling convention under Windows
+        cconv = "x86_stdcallcc"
+    else:
+        # Default C calling convention
+        cconv = None
+
+    sig = templates.signature(cret, *cargs)
+    return types.ExternalFunctionPointer(
+        sig, cconv=cconv, get_pointer=get_pointer
+    )

numba_cuda/numba/cuda/typing/dictdecl.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+This implements the typing template for `dict()`.
+"""
+
+from numba.cuda import types
+from numba.cuda.core import errors
+from .templates import (
+    AbstractTemplate,
+    Registry,
+    signature,
+)
+
+registry = Registry()
+infer = registry.register
+infer_global = registry.register_global
+infer_getattr = registry.register_attr
+
+
+_message_dict_support = """
+Unsupported use of `dict()` with positional or keyword argument(s). \
+The only supported uses are `dict()` or `dict(iterable)`.
+""".strip()
+
+
+@infer_global(dict)
+class DictBuiltin(AbstractTemplate):
+    def generic(self, args, kws):
+        if kws:
+            raise errors.TypingError(_message_dict_support)
+        if args:
+            (iterable,) = args
+            if isinstance(iterable, types.IterableType):
+                dtype = iterable.iterator_type.yield_type
+                if isinstance(dtype, types.UniTuple):
+                    length = dtype.count
+                    if length != 2:
+                        msg = (
+                            "dictionary update sequence element has length "
+                            f"{length}; 2 is required"
+                        )
+                        raise errors.TypingError(msg)
+                    k = v = dtype.key[0]
+                elif isinstance(dtype, types.Tuple):
+                    k, v = dtype.key
+                else:
+                    raise errors.TypingError(_message_dict_support)
+
+                # dict key must be hashable
+                if not isinstance(k, types.Hashable):
+                    msg = f"Unhashable type: {k}"
+                    raise errors.TypingError(msg)
+
+                return signature(types.DictType(k, v), iterable)
+            else:
+                msg = (
+                    "Non-iterable args used in dict(iterable) "
+                    f"constructor. Got 'dict({args[0]})'"
+                )
+                raise errors.TypingError(msg)
+        return signature(types.DictType(types.undefined, types.undefined))

numba_cuda/numba/cuda/typing/enumdecl.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Typing for enums.
+"""
+
+import operator
+from numba.cuda import types
+from numba.cuda.typing.templates import (
+    AbstractTemplate,
+    AttributeTemplate,
+    signature,
+    Registry,
+)
+
+registry = Registry()
+infer = registry.register
+infer_global = registry.register_global
+infer_getattr = registry.register_attr
+
+
+@infer_getattr
+class EnumAttribute(AttributeTemplate):
+    key = types.EnumMember
+
+    def resolve_value(self, ty):
+        return ty.dtype
+
+
+@infer_getattr
+class EnumClassAttribute(AttributeTemplate):
+    key = types.EnumClass
+
+    def generic_resolve(self, ty, attr):
+        """
+        Resolve attributes of an enum class as enum members.
+        """
+        if attr in ty.instance_class.__members__:
+            return ty.member_type
+
+
+@infer
+class EnumClassStaticGetItem(AbstractTemplate):
+    key = "static_getitem"
+
+    def generic(self, args, kws):
+        enum, idx = args
+        if (
+            isinstance(enum, types.EnumClass)
+            and idx in enum.instance_class.__members__
+        ):
+            return signature(enum.member_type, *args)
+
+
+class EnumCompare(AbstractTemplate):
+    def generic(self, args, kws):
+        [lhs, rhs] = args
+        if (
+            isinstance(lhs, types.EnumMember)
+            and isinstance(rhs, types.EnumMember)
+            and lhs == rhs
+        ):
+            return signature(types.boolean, lhs, rhs)
+
+
+@infer_global(operator.eq)
+class EnumEq(EnumCompare):
+    pass
+
+
+@infer_global(operator.ne)
+class EnumNe(EnumCompare):
+    pass

numba_cuda/numba/cuda/typing/listdecl.py

@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import operator
+from numba.cuda import types
+from .templates import (
+    AbstractTemplate,
+    AttributeTemplate,
+    Registry,
+    signature,
+    bound_function,
+)
+# Ensure list is typed as a collection as well
+
+
+registry = Registry()
+infer = registry.register
+infer_global = registry.register_global
+infer_getattr = registry.register_attr
+
+
+@infer_global(list)
+class ListBuiltin(AbstractTemplate):
+    def generic(self, args, kws):
+        assert not kws
+        if args:
+            (iterable,) = args
+            if isinstance(iterable, types.IterableType):
+                dtype = iterable.iterator_type.yield_type
+                return signature(types.List(dtype), iterable)
+        else:
+            return signature(types.List(types.undefined))
+
+
+@infer_getattr
+class ListAttribute(AttributeTemplate):
+    key = types.List
+
+    # NOTE: some of these should be Sequence / MutableSequence methods
+
+    @bound_function("list.append")
+    def resolve_append(self, list, args, kws):
+        (item,) = args
+        assert not kws
+        unified = self.context.unify_pairs(list.dtype, item)
+        if unified is not None:
+            sig = signature(types.none, unified)
+            sig = sig.replace(recvr=list.copy(dtype=unified))
+            return sig
+
+    @bound_function("list.clear")
+    def resolve_clear(self, list, args, kws):
+        assert not args
+        assert not kws
+        return signature(types.none)
+
+    @bound_function("list.extend")
+    def resolve_extend(self, list, args, kws):
+        (iterable,) = args
+        assert not kws
+        if not isinstance(iterable, types.IterableType):
+            return
+
+        dtype = iterable.iterator_type.yield_type
+        unified = self.context.unify_pairs(list.dtype, dtype)
+        if unified is not None:
+            sig = signature(types.none, iterable)
+            sig = sig.replace(recvr=list.copy(dtype=unified))
+            return sig
+
+    @bound_function("list.insert")
+    def resolve_insert(self, list, args, kws):
+        idx, item = args
+        assert not kws
+        if isinstance(idx, types.Integer):
+            unified = self.context.unify_pairs(list.dtype, item)
+            if unified is not None:
+                sig = signature(types.none, types.intp, unified)
+                sig = sig.replace(recvr=list.copy(dtype=unified))
+                return sig
+
+    @bound_function("list.pop")
+    def resolve_pop(self, list, args, kws):
+        assert not kws
+        if not args:
+            return signature(list.dtype)
+        else:
+            (idx,) = args
+            if isinstance(idx, types.Integer):
+                return signature(list.dtype, types.intp)
+
+
+@infer_global(operator.add)
+class AddList(AbstractTemplate):
+    def generic(self, args, kws):
+        if len(args) == 2:
+            a, b = args
+            if isinstance(a, types.List) and isinstance(b, types.List):
+                unified = self.context.unify_pairs(a, b)
+                if unified is not None:
+                    return signature(unified, a, b)
+
+
+@infer_global(operator.iadd)
+class InplaceAddList(AbstractTemplate):
+    def generic(self, args, kws):
+        if len(args) == 2:
+            a, b = args
+            if isinstance(a, types.List) and isinstance(b, types.List):
+                if self.context.can_convert(b.dtype, a.dtype):
+                    return signature(a, a, b)
+
+
+@infer_global(operator.mul)
+class MulList(AbstractTemplate):
+    # key = operator.mul
+
+    def generic(self, args, kws):
+        a, b = args
+        if isinstance(a, types.List) and isinstance(b, types.Integer):
+            return signature(a, a, types.intp)
+        elif isinstance(a, types.Integer) and isinstance(b, types.List):
+            return signature(b, types.intp, b)
+
+
+@infer_global(operator.imul)
+class InplaceMulList(MulList):
+    pass
+    # key = operator.imul
+
+
+class ListCompare(AbstractTemplate):
+    def generic(self, args, kws):
+        [lhs, rhs] = args
+        if isinstance(lhs, types.List) and isinstance(rhs, types.List):
+            # Check element-wise comparability
+            res = self.context.resolve_function_type(
+                self.key, (lhs.dtype, rhs.dtype), {}
+            )
+            if res is not None:
+                return signature(types.boolean, lhs, rhs)
+
+
+@infer_global(operator.eq)
+class ListEq(ListCompare):
+    pass
+    # key = operator.eq

numba_cuda/numba/cuda/typing/mathdecl.py

@@ -0,0 +1,158 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+import math
+from numba.cuda import types
+from numba.cuda.typing.templates import ConcreteTemplate, signature, Registry
+
+registry = Registry()
+infer_global = registry.register_global
+
+
+@infer_global(math.exp)
+@infer_global(math.expm1)
+@infer_global(math.fabs)
+@infer_global(math.sqrt)
+@infer_global(math.log)
+@infer_global(math.log1p)
+@infer_global(math.log10)
+@infer_global(math.log2)
+@infer_global(math.sin)
+@infer_global(math.cos)
+@infer_global(math.tan)
+@infer_global(math.sinh)
+@infer_global(math.cosh)
+@infer_global(math.tanh)
+@infer_global(math.asin)
+@infer_global(math.acos)
+@infer_global(math.atan)
+@infer_global(math.asinh)
+@infer_global(math.acosh)
+@infer_global(math.atanh)
+@infer_global(math.degrees)
+@infer_global(math.radians)
+@infer_global(math.erf)
+@infer_global(math.erfc)
+@infer_global(math.gamma)
+@infer_global(math.lgamma)
+class Math_unary(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.int64),
+        signature(types.float64, types.uint64),
+        signature(types.float32, types.float32),
+        signature(types.float64, types.float64),
+    ]
+
+
+if sys.version_info >= (3, 11):
+    Math_unary = infer_global(math.exp2)(Math_unary)
+
+
+@infer_global(math.atan2)
+class Math_atan2(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.int64, types.int64),
+        signature(types.float64, types.uint64, types.uint64),
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.trunc)
+class Math_converter(ConcreteTemplate):
+    cases = [
+        signature(types.intp, types.intp),
+        signature(types.int64, types.int64),
+        signature(types.uint64, types.uint64),
+        signature(types.int64, types.float32),
+        signature(types.int64, types.float64),
+    ]
+
+
+@infer_global(math.floor)
+@infer_global(math.ceil)
+class Math_floor_ceil(Math_converter):
+    pass
+
+
+@infer_global(math.copysign)
+class Math_copysign(ConcreteTemplate):
+    cases = [
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.hypot)
+class Math_hypot(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.int64, types.int64),
+        signature(types.float64, types.uint64, types.uint64),
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.nextafter)
+class Math_nextafter(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.float64, types.float64),
+        signature(types.float32, types.float32, types.float32),
+    ]
+
+
+@infer_global(math.isinf)
+@infer_global(math.isnan)
+class Math_predicate(ConcreteTemplate):
+    cases = [
+        signature(types.boolean, types.int64),
+        signature(types.boolean, types.uint64),
+        signature(types.boolean, types.float32),
+        signature(types.boolean, types.float64),
+    ]
+
+
+@infer_global(math.isfinite)
+class Math_isfinite(Math_predicate):
+    pass
+
+
+@infer_global(math.pow)
+class Math_pow(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.float64, types.int64),
+        signature(types.float64, types.float64, types.uint64),
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.gcd)
+class Math_gcd(ConcreteTemplate):
+    cases = [
+        signature(types.int64, types.int64, types.int64),
+        signature(types.int32, types.int32, types.int32),
+        signature(types.int16, types.int16, types.int16),
+        signature(types.int8, types.int8, types.int8),
+        signature(types.uint64, types.uint64, types.uint64),
+        signature(types.uint32, types.uint32, types.uint32),
+        signature(types.uint16, types.uint16, types.uint16),
+        signature(types.uint8, types.uint8, types.uint8),
+    ]
+
+
+@infer_global(math.frexp)
+class Math_frexp(ConcreteTemplate):
+    cases = [
+        signature(types.Tuple((types.float64, types.intc)), types.float64),
+        signature(types.Tuple((types.float32, types.intc)), types.float32),
+    ]
+
+
+@infer_global(math.ldexp)
+class Math_ldexp(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.float64, types.intc),
+        signature(types.float32, types.float32, types.intc),
+    ]