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

numba_cuda/numba/cuda/typing/cffi_utils.py

@@ -0,0 +1,219 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Support for CFFI. Allows checking whether objects are CFFI functions and
+obtaining the pointer and numba signature.
+"""
+
+from functools import partial
+import numpy as np
+
+from numba.cuda import types
+from numba.cuda.core.errors import TypingError
+from numba.cuda.typing import templates
+from numba.cuda.np import numpy_support
+
+try:
+    import cffi
+
+    ffi = cffi.FFI()
+except ImportError:
+    ffi = None
+
+SUPPORTED = ffi is not None
+_ool_func_types = {}
+_ool_func_ptr = {}
+_ffi_instances = set()
+
+
+def is_ffi_instance(obj):
+    # Compiled FFI modules have a member, ffi, which is an instance of
+    # CompiledFFI, which behaves similarly to an instance of cffi.FFI. In
+    # order to simplify handling a CompiledFFI object, we treat them as
+    # if they're cffi.FFI instances for typing and lowering purposes.
+    try:
+        return obj in _ffi_instances or isinstance(obj, cffi.FFI)
+    except TypeError:  # Unhashable type possible
+        return False
+
+
+def is_cffi_func(obj):
+    """Check whether the obj is a CFFI function"""
+    try:
+        return ffi.typeof(obj).kind == "function"
+    except TypeError:
+        try:
+            return obj in _ool_func_types
+        except Exception:
+            return False
+
+
+def get_pointer(cffi_func):
+    """
+    Get a pointer to the underlying function for a CFFI function as an
+    integer.
+    """
+    if cffi_func in _ool_func_ptr:
+        return _ool_func_ptr[cffi_func]
+    return int(ffi.cast("uintptr_t", cffi_func))
+
+
+_cached_type_map = None
+
+
+def _type_map():
+    """
+    Lazily compute type map, as calling ffi.typeof() involves costly
+    parsing of C code...
+    """
+    global _cached_type_map
+    if _cached_type_map is None:
+        _cached_type_map = {
+            ffi.typeof("bool"): types.boolean,
+            ffi.typeof("char"): types.char,
+            ffi.typeof("short"): types.short,
+            ffi.typeof("int"): types.intc,
+            ffi.typeof("long"): types.long_,
+            ffi.typeof("long long"): types.longlong,
+            ffi.typeof("unsigned char"): types.uchar,
+            ffi.typeof("unsigned short"): types.ushort,
+            ffi.typeof("unsigned int"): types.uintc,
+            ffi.typeof("unsigned long"): types.ulong,
+            ffi.typeof("unsigned long long"): types.ulonglong,
+            ffi.typeof("int8_t"): types.char,
+            ffi.typeof("uint8_t"): types.uchar,
+            ffi.typeof("int16_t"): types.short,
+            ffi.typeof("uint16_t"): types.ushort,
+            ffi.typeof("int32_t"): types.intc,
+            ffi.typeof("uint32_t"): types.uintc,
+            ffi.typeof("int64_t"): types.longlong,
+            ffi.typeof("uint64_t"): types.ulonglong,
+            ffi.typeof("float"): types.float32,
+            ffi.typeof("double"): types.double,
+            ffi.typeof("ssize_t"): types.intp,
+            ffi.typeof("size_t"): types.uintp,
+            ffi.typeof("void"): types.void,
+        }
+    return _cached_type_map
+
+
+def map_type(cffi_type, use_record_dtype=False):
+    """
+    Map CFFI type to numba type.
+
+    Parameters
+    ----------
+    cffi_type:
+        The CFFI type to be converted.
+    use_record_dtype: bool (default: False)
+        When True, struct types are mapped to a NumPy Record dtype.
+
+    """
+    primed_map_type = partial(map_type, use_record_dtype=use_record_dtype)
+    kind = getattr(cffi_type, "kind", "")
+    if kind == "union":
+        raise TypeError("No support for CFFI union")
+    elif kind == "function":
+        if cffi_type.ellipsis:
+            raise TypeError("vararg function is not supported")
+        restype = primed_map_type(cffi_type.result)
+        argtypes = [primed_map_type(arg) for arg in cffi_type.args]
+        return templates.signature(restype, *argtypes)
+    elif kind == "pointer":
+        pointee = cffi_type.item
+        if pointee.kind == "void":
+            return types.voidptr
+        else:
+            return types.CPointer(primed_map_type(pointee))
+    elif kind == "array":
+        dtype = primed_map_type(cffi_type.item)
+        nelem = cffi_type.length
+        return types.NestedArray(dtype=dtype, shape=(nelem,))
+    elif kind == "struct" and use_record_dtype:
+        return map_struct_to_record_dtype(cffi_type)
+    else:
+        result = _type_map().get(cffi_type)
+        if result is None:
+            raise TypeError(cffi_type)
+        return result
+
+
+def map_struct_to_record_dtype(cffi_type):
+    """Convert a cffi type into a NumPy Record dtype"""
+    fields = {
+        "names": [],
+        "formats": [],
+        "offsets": [],
+        "itemsize": ffi.sizeof(cffi_type),
+    }
+    is_aligned = True
+    for k, v in cffi_type.fields:
+        # guard unsupported values
+        if v.bitshift != -1:
+            msg = "field {!r} has bitshift, this is not supported"
+            raise ValueError(msg.format(k))
+        if v.flags != 0:
+            msg = "field {!r} has flags, this is not supported"
+            raise ValueError(msg.format(k))
+        if v.bitsize != -1:
+            msg = "field {!r} has bitsize, this is not supported"
+            raise ValueError(msg.format(k))
+        dtype = numpy_support.as_dtype(
+            map_type(v.type, use_record_dtype=True),
+        )
+        fields["names"].append(k)
+        fields["formats"].append(dtype)
+        fields["offsets"].append(v.offset)
+        # Check alignment
+        is_aligned &= v.offset % dtype.alignment == 0
+
+    return numpy_support.from_dtype(np.dtype(fields, align=is_aligned))
+
+
+def make_function_type(cffi_func, use_record_dtype=False):
+    """
+    Return a Numba type for the given CFFI function pointer.
+    """
+    cffi_type = _ool_func_types.get(cffi_func) or ffi.typeof(cffi_func)
+    if getattr(cffi_type, "kind", "") == "struct":
+        raise TypeError("No support for CFFI struct values")
+    sig = map_type(cffi_type, use_record_dtype=use_record_dtype)
+    return types.ExternalFunctionPointer(sig, get_pointer=get_pointer)
+
+
+registry = templates.Registry()
+
+
+@registry.register
+class FFI_from_buffer(templates.AbstractTemplate):
+    key = "ffi.from_buffer"
+
+    def generic(self, args, kws):
+        if kws or len(args) != 1:
+            return
+        [ary] = args
+        if not isinstance(ary, types.Buffer):
+            raise TypingError(
+                "from_buffer() expected a buffer object, got %s" % (ary,)
+            )
+        if ary.layout not in ("C", "F"):
+            raise TypingError(
+                "from_buffer() unsupported on non-contiguous buffers (got %s)"
+                % (ary,)
+            )
+        if ary.layout != "C" and ary.ndim > 1:
+            raise TypingError(
+                "from_buffer() only supports multidimensional arrays with C layout (got %s)"
+                % (ary,)
+            )
+        ptr = types.CPointer(ary.dtype)
+        return templates.signature(ptr, ary)
+
+
+@registry.register_attr
+class FFIAttribute(templates.AttributeTemplate):
+    key = types.ffi
+
+    def resolve_from_buffer(self, ffi):
+        return types.BoundFunction(FFI_from_buffer, types.ffi)

numba_cuda/numba/cuda/typing/cmathdecl.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import cmath
+
+from numba.cuda import types
+from numba.cuda.typing.templates import ConcreteTemplate, signature, Registry
+
+registry = Registry()
+infer_global = registry.register_global
+
+# TODO: support non-complex arguments (floats and ints)
+
+
+@infer_global(cmath.acos)
+@infer_global(cmath.asin)
+@infer_global(cmath.asinh)
+@infer_global(cmath.atan)
+@infer_global(cmath.atanh)
+@infer_global(cmath.cos)
+@infer_global(cmath.exp)
+@infer_global(cmath.sin)
+@infer_global(cmath.sqrt)
+@infer_global(cmath.tan)
+class CMath_unary(ConcreteTemplate):
+    cases = [signature(tp, tp) for tp in sorted(types.complex_domain)]
+
+
+@infer_global(cmath.isinf)
+@infer_global(cmath.isnan)
+class CMath_predicate(ConcreteTemplate):
+    cases = [
+        signature(types.boolean, tp) for tp in sorted(types.complex_domain)
+    ]
+
+
+@infer_global(cmath.isfinite)
+class CMath_isfinite(CMath_predicate):
+    pass
+
+
+@infer_global(cmath.log)
+class Cmath_log(ConcreteTemplate):
+    # unary cmath.log()
+    cases = [signature(tp, tp) for tp in sorted(types.complex_domain)]
+    # binary cmath.log()
+    cases += [signature(tp, tp, tp) for tp in sorted(types.complex_domain)]

numba_cuda/numba/cuda/typing/collections.py

@@ -0,0 +1,138 @@
+# 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 errors
+from numba.cuda import utils
+import operator
+from .templates import (
+    AttributeTemplate,
+    AbstractTemplate,
+    Registry,
+    signature,
+    make_callable_template,
+)
+from .builtins import normalize_1d_index
+
+registry = Registry()
+infer_global = registry.register_global
+infer_getattr = registry.register_attr
+
+
+@infer_global(operator.contains)
+class InContainer(AbstractTemplate):
+    key = operator.contains
+
+    def generic(self, args, kws):
+        cont, item = args
+        if isinstance(cont, types.Container):
+            return signature(types.boolean, cont, cont.dtype)
+
+
+@infer_global(len)
+class ContainerLen(AbstractTemplate):
+    def generic(self, args, kws):
+        assert not kws
+        (val,) = args
+        if isinstance(val, (types.Container)):
+            return signature(types.intp, val)
+
+
+@infer_global(operator.truth)
+class SequenceBool(AbstractTemplate):
+    key = operator.truth
+
+    def generic(self, args, kws):
+        assert not kws
+        (val,) = args
+        if isinstance(val, (types.Sequence)):
+            return signature(types.boolean, val)
+
+
+@infer_global(operator.getitem)
+class GetItemSequence(AbstractTemplate):
+    key = operator.getitem
+
+    def generic(self, args, kws):
+        seq, idx = args
+        if isinstance(seq, types.Sequence):
+            idx = normalize_1d_index(idx)
+            if isinstance(idx, types.SliceType):
+                # Slicing a tuple only supported with static_getitem
+                if not isinstance(seq, types.BaseTuple):
+                    return signature(seq, seq, idx)
+            elif isinstance(idx, types.Integer):
+                return signature(seq.dtype, seq, idx)
+
+
+@infer_global(operator.setitem)
+class SetItemSequence(AbstractTemplate):
+    def generic(self, args, kws):
+        seq, idx, value = args
+        if isinstance(seq, types.MutableSequence):
+            idx = normalize_1d_index(idx)
+            if isinstance(idx, types.SliceType):
+                return signature(types.none, seq, idx, seq)
+            elif isinstance(idx, types.Integer):
+                if not self.context.can_convert(value, seq.dtype):
+                    msg = "invalid setitem with value of {} to element of {}"
+                    raise errors.TypingError(
+                        msg.format(types.unliteral(value), seq.dtype)
+                    )
+                return signature(types.none, seq, idx, seq.dtype)
+
+
+@infer_global(operator.delitem)
+class DelItemSequence(AbstractTemplate):
+    def generic(self, args, kws):
+        seq, idx = args
+        if isinstance(seq, types.MutableSequence):
+            idx = normalize_1d_index(idx)
+            return signature(types.none, seq, idx)
+
+
+# --------------------------------------------------------------------------
+# named tuples
+
+
+@infer_getattr
+class NamedTupleAttribute(AttributeTemplate):
+    key = types.BaseNamedTuple
+
+    def resolve___class__(self, tup):
+        return types.NamedTupleClass(tup.instance_class)
+
+    def generic_resolve(self, tup, attr):
+        # Resolution of other attributes
+        try:
+            index = tup.fields.index(attr)
+        except ValueError:
+            return
+        return tup[index]
+
+
+@infer_getattr
+class NamedTupleClassAttribute(AttributeTemplate):
+    key = types.NamedTupleClass
+
+    def resolve___call__(self, classty):
+        """
+        Resolve the named tuple constructor, aka the class's __call__ method.
+        """
+        instance_class = classty.instance_class
+        pysig = utils.pysignature(instance_class)
+
+        def typer(*args, **kws):
+            # Fold keyword args
+            try:
+                bound = pysig.bind(*args, **kws)
+            except TypeError as e:
+                msg = "In '%s': %s" % (instance_class, e)
+                e.args = (msg,)
+                raise
+            assert not bound.kwargs
+            return types.BaseTuple.from_types(bound.args, instance_class)
+
+        # Override the typer's pysig to match the namedtuple constructor's
+        typer.pysig = pysig
+        return types.Function(make_callable_template(self.key, typer))