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

numba_cuda/numba/cuda/types/cuda_abstract.py

@@ -0,0 +1,533 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from abc import ABCMeta, abstractmethod, abstractproperty
+from typing import Dict as ptDict, Type as ptType
+import itertools
+import weakref
+from functools import cached_property
+
+import numpy as np
+
+from numba.cuda.utils import get_hashable_key
+
+# Types are added to a global registry (_typecache) in order to assign
+# them unique integer codes for fast matching in _dispatcher.c.
+# However, we also want types to be disposable, therefore we ensure
+# each type is interned as a weak reference, so that it lives only as
+# long as necessary to keep a stable type code.
+# NOTE: some types can still be made immortal elsewhere (for example
+# in _dispatcher.c's internal caches).
+_typecodes = itertools.count()
+
+
+def _autoincr():
+    n = next(_typecodes)
+    # 4 billion types should be enough, right?
+    assert n < 2**32, "Limited to 4 billion types"
+    return n
+
+
+_typecache: ptDict[weakref.ref, weakref.ref] = {}
+
+
+def _on_type_disposal(wr, _pop=_typecache.pop):
+    _pop(wr, None)
+
+
+class _TypeMetaclass(ABCMeta):
+    """
+    A metaclass that will intern instances after they are created.
+    This is done by first creating a new instance (including calling
+    __init__, which sets up the required attributes for equality
+    and hashing), then looking it up in the _typecache registry.
+    """
+
+    def __init__(cls, name, bases, orig_vars):
+        # __init__ is hooked to mark whether a Type class being defined is a
+        # Numba internal type (one which is defined somewhere under the `numba`
+        # module) or an external type (one which is defined elsewhere, for
+        # example a user defined type).
+        super(_TypeMetaclass, cls).__init__(name, bases, orig_vars)
+        root = (cls.__module__.split("."))[0]
+        cls._is_internal = root == "numba"
+
+    def _intern(cls, inst):
+        # Try to intern the created instance
+        wr = weakref.ref(inst, _on_type_disposal)
+        orig = _typecache.get(wr)
+        orig = orig and orig()
+        if orig is not None:
+            return orig
+        else:
+            inst._code = _autoincr()
+            _typecache[wr] = wr
+            return inst
+
+    def __call__(cls, *args, **kwargs):
+        """
+        Instantiate *cls* (a Type subclass, presumably) and intern it.
+        If an interned instance already exists, it is returned, otherwise
+        the new instance is returned.
+        """
+        inst = type.__call__(cls, *args, **kwargs)
+        return cls._intern(inst)
+
+
+def _type_reconstructor(reconstructor, reconstructor_args, state):
+    """
+    Rebuild function for unpickling types.
+    """
+    obj = reconstructor(*reconstructor_args)
+    if state:
+        obj.__dict__.update(state)
+    return type(obj)._intern(obj)
+
+
+class Type(metaclass=_TypeMetaclass):
+    """
+    The base class for all Numba types.
+    It is essential that proper equality comparison is implemented.  The
+    default implementation uses the "key" property (overridable in subclasses)
+    for both comparison and hashing, to ensure sane behaviour.
+    """
+
+    mutable = False
+    # Rather the type is reflected at the python<->nopython boundary
+    reflected = False
+
+    def __init__(self, name):
+        self.name = name
+
+    @property
+    def key(self):
+        """
+        A property used for __eq__, __ne__ and __hash__.  Can be overridden
+        in subclasses.
+        """
+        return self.name
+
+    @property
+    def mangling_args(self):
+        """
+        Returns `(basename, args)` where `basename` is the name of the type
+        and `args` is a sequence of parameters of the type.
+
+        Subclass should override to specialize the behavior.
+        By default, this returns `(self.name, ())`.
+        """
+        return self.name, ()
+
+    def __repr__(self):
+        return self.name
+
+    def __str__(self):
+        return self.name
+
+    def __hash__(self):
+        return hash(self.key)
+
+    def __eq__(self, other):
+        return self.__class__ is other.__class__ and self.key == other.key
+
+    def __ne__(self, other):
+        return not (self == other)
+
+    def __reduce__(self):
+        reconstructor, args, state = super(Type, self).__reduce__()
+        return (_type_reconstructor, (reconstructor, args, state))
+
+    def unify(self, typingctx, other):
+        """
+        Try to unify this type with the *other*.  A third type must
+        be returned, or None if unification is not possible.
+        Only override this if the coercion logic cannot be expressed
+        as simple casting rules.
+        """
+        return None
+
+    def can_convert_to(self, typingctx, other):
+        """
+        Check whether this type can be converted to the *other*.
+        If successful, must return a string describing the conversion, e.g.
+        "exact", "promote", "unsafe", "safe"; otherwise None is returned.
+        """
+        return None
+
+    def can_convert_from(self, typingctx, other):
+        """
+        Similar to *can_convert_to*, but in reverse.  Only needed if
+        the type provides conversion from other types.
+        """
+        return None
+
+    def is_precise(self):
+        """
+        Whether this type is precise, i.e. can be part of a successful
+        type inference.  Default implementation returns True.
+        """
+        return True
+
+    def augment(self, other):
+        """
+        Augment this type with the *other*.  Return the augmented type,
+        or None if not supported.
+        """
+        return None
+
+    # User-facing helpers.  These are not part of the core Type API but
+    # are provided so that users can write e.g. `numba.boolean(1.5)`
+    # (returns True) or `types.int32(types.int32[:])` (returns something
+    # usable as a function signature).
+
+    def __call__(self, *args):
+        from numba.cuda.typing import signature
+
+        if len(args) == 1 and not isinstance(args[0], Type):
+            return self.cast_python_value(args[0])
+        return signature(
+            self,  # return_type
+            *args,
+        )
+
+    def __getitem__(self, args):
+        """
+        Return an array of this type.
+        """
+        from numba.cuda.types import Array
+
+        ndim, layout = self._determine_array_spec(args)
+        return Array(dtype=self, ndim=ndim, layout=layout)
+
+    def _determine_array_spec(self, args):
+        # XXX non-contiguous by default, even for 1d arrays,
+        # doesn't sound very intuitive
+        def validate_slice(s):
+            return isinstance(s, slice) and s.start is None and s.stop is None
+
+        if isinstance(args, (tuple, list)) and all(map(validate_slice, args)):
+            ndim = len(args)
+            if args[0].step == 1:
+                layout = "F"
+            elif args[-1].step == 1:
+                layout = "C"
+            else:
+                layout = "A"
+        elif validate_slice(args):
+            ndim = 1
+            if args.step == 1:
+                layout = "C"
+            else:
+                layout = "A"
+        else:
+            # Raise a KeyError to not be handled by collection constructors (e.g. list).
+            raise KeyError(
+                f"Can only index numba types with slices with no start or stop, got {args}."
+            )
+
+        return ndim, layout
+
+    def cast_python_value(self, args):
+        raise NotImplementedError
+
+    @property
+    def is_internal(self):
+        """Returns True if this class is an internally defined Numba type by
+        virtue of the module in which it is instantiated, False else."""
+        return self._is_internal
+
+    def dump(self, tab=""):
+        print(
+            f"{tab}DUMP {type(self).__name__}[code={self._code}, name={self.name}]"
+        )
+
+
+# XXX we should distinguish between Dummy (no meaningful
+# representation, e.g. None or a builtin function) and Opaque (has a
+# meaningful representation, e.g. ExternalFunctionPointer)
+
+
+class Dummy(Type):
+    """
+    Base class for types that do not really have a representation and are
+    compatible with a void*.
+    """
+
+
+class Hashable(Type):
+    """
+    Base class for hashable types.
+    """
+
+
+class Number(Hashable):
+    """
+    Base class for number types.
+    """
+
+    def unify(self, typingctx, other):
+        """
+        Unify the two number types using Numpy's rules.
+        """
+        from numba.cuda.np import numpy_support
+
+        if isinstance(other, Number):
+            # XXX: this can produce unsafe conversions,
+            # e.g. would unify {int64, uint64} to float64
+            a = numpy_support.as_dtype(self)
+            b = numpy_support.as_dtype(other)
+            sel = np.promote_types(a, b)
+            return numpy_support.from_dtype(sel)
+
+
+class Callable(Type):
+    """
+    Base class for callables.
+    """
+
+    @abstractmethod
+    def get_call_type(self, context, args, kws):
+        """
+        Using the typing *context*, resolve the callable's signature for
+        the given arguments.  A signature object is returned, or None.
+        """
+
+    @abstractmethod
+    def get_call_signatures(self):
+        """
+        Returns a tuple of (list of signatures, parameterized)
+        """
+
+    @abstractmethod
+    def get_impl_key(self, sig):
+        """
+        Returns the impl key for the given signature
+        """
+
+
+class DTypeSpec(Type):
+    """
+    Base class for types usable as "dtype" arguments to various Numpy APIs
+    (e.g. np.empty()).
+    """
+
+    @abstractproperty
+    def dtype(self):
+        """
+        The actual dtype denoted by this dtype spec (a Type instance).
+        """
+
+
+class IterableType(Type):
+    """
+    Base class for iterable types.
+    """
+
+    @abstractproperty
+    def iterator_type(self):
+        """
+        The iterator type obtained when calling iter() (explicitly or implicitly).
+        """
+
+
+class Sized(Type):
+    """
+    Base class for objects that support len()
+    """
+
+
+class ConstSized(Sized):
+    """
+    For types that have a constant size
+    """
+
+    @abstractmethod
+    def __len__(self):
+        pass
+
+
+class IteratorType(IterableType):
+    """
+    Base class for all iterator types.
+    Derived classes should implement the *yield_type* attribute.
+    """
+
+    def __init__(self, name, **kwargs):
+        super(IteratorType, self).__init__(name, **kwargs)
+
+    @abstractproperty
+    def yield_type(self):
+        """
+        The type of values yielded by the iterator.
+        """
+
+    # This is a property to avoid recursivity (for pickling)
+
+    @property
+    def iterator_type(self):
+        return self
+
+
+class Container(Sized, IterableType):
+    """
+    Base class for container types.
+    """
+
+
+class Sequence(Container):
+    """
+    Base class for 1d sequence types.  Instances should have the *dtype*
+    attribute.
+    """
+
+
+class MutableSequence(Sequence):
+    """
+    Base class for 1d mutable sequence types.  Instances should have the
+    *dtype* attribute.
+    """
+
+    mutable = True
+
+
+class ArrayCompatible(Type):
+    """
+    Type class for Numpy array-compatible objects (typically, objects
+    exposing an __array__ method).
+    Derived classes should implement the *as_array* attribute.
+    """
+
+    # If overridden by a subclass, it should also implement typing
+    # for '__array_wrap__' with arguments (input, formal result).
+    array_priority = 0.0
+
+    @abstractproperty
+    def as_array(self):
+        """
+        The equivalent array type, for operations supporting array-compatible
+        objects (such as ufuncs).
+        """
+
+    # For compatibility with types.Array
+
+    @cached_property
+    def ndim(self):
+        return self.as_array.ndim
+
+    @cached_property
+    def layout(self):
+        return self.as_array.layout
+
+    @cached_property
+    def dtype(self):
+        return self.as_array.dtype
+
+
+class Literal(Type):
+    """Base class for Literal types.
+    Literal types contain the original Python value in the type.
+
+    A literal type should always be constructed from the `literal(val)`
+    function.
+    """
+
+    # *ctor_map* is a dictionary mapping Python types to Literal subclasses
+    # for constructing a numba type for a given Python type.
+    # It is used in `literal(val)` function.
+    # To add new Literal subclass, register a new mapping to this dict.
+    ctor_map: ptDict[type, ptType["Literal"]] = {}
+
+    # *_literal_type_cache* is used to cache the numba type of the given value.
+    _literal_type_cache = None
+
+    def __init__(self, value):
+        if type(self) is Literal:
+            raise TypeError(
+                "Cannot be constructed directly. "
+                "Use `numba.cuda.types.literal(value)` instead",
+            )
+        self._literal_init(value)
+        fmt = "Literal[{}]({})"
+        super(Literal, self).__init__(fmt.format(type(value).__name__, value))
+
+    def _literal_init(self, value):
+        self._literal_value = value
+        # We want to support constants of non-hashable values, therefore
+        # fall back on the value's id() if necessary.
+        self._key = get_hashable_key(value)
+
+    @property
+    def literal_value(self):
+        return self._literal_value
+
+    @property
+    def literal_type(self):
+        if self._literal_type_cache is None:
+            from numba.cuda import typing
+
+            ctx = typing.Context()
+            try:
+                res = ctx.resolve_value_type(self.literal_value)
+            except ValueError as e:
+                if "Int value is too large" in str(e):
+                    # If a string literal cannot create an IntegerLiteral
+                    # because of overflow we generate this message.
+                    msg = f"Cannot create literal type. {str(e)}"
+                    raise TypeError(msg)
+                # Not all literal types have a literal_value that can be
+                # resolved to a type, for example, LiteralStrKeyDict has a
+                # literal_value that is a python dict for which there's no
+                # `typeof` support.
+                msg = "{} has no attribute 'literal_type'".format(self)
+                raise AttributeError(msg)
+            self._literal_type_cache = res
+
+        return self._literal_type_cache
+
+
+class TypeRef(Dummy):
+    """Reference to a type.
+
+    Used when a type is passed as a value.
+    """
+
+    def __init__(self, instance_type):
+        self.instance_type = instance_type
+        super(TypeRef, self).__init__("typeref[{}]".format(self.instance_type))
+
+    @property
+    def key(self):
+        return self.instance_type
+
+
+class InitialValue(object):
+    """
+    Used as a mixin for a type will potentially have an initial value that will
+    be carried in the .initial_value attribute.
+    """
+
+    def __init__(self, initial_value):
+        self._initial_value = initial_value
+
+    @property
+    def initial_value(self):
+        return self._initial_value
+
+
+class Poison(Type):
+    """
+    This is the "bottom" type in the type system. It won't unify and it's
+    unliteral version is Poison of itself. It's advisable for debugging purposes
+    to call the constructor with the type that's being poisoned (for whatever
+    reason) but this isn't strictly required.
+    """
+
+    def __init__(self, ty):
+        self.ty = ty
+        super(Poison, self).__init__(name="Poison<%s>" % ty)
+
+    def __unliteral__(self):
+        return Poison(self)
+
+    def unify(self, typingctx, other):
+        return None

numba_cuda/numba/cuda/types/cuda_common.py

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Helper classes / mixins for defining types.
+"""
+
+from .abstract import ArrayCompatible, Dummy, IterableType, IteratorType
+from numba.cuda.core.errors import NumbaTypeError, NumbaValueError
+
+
+class Opaque(Dummy):
+    """
+    A type that is a opaque pointer.
+    """
+
+
+class SimpleIterableType(IterableType):
+    def __init__(self, name, iterator_type):
+        self._iterator_type = iterator_type
+        super(SimpleIterableType, self).__init__(name)
+
+    @property
+    def iterator_type(self):
+        return self._iterator_type
+
+
+class SimpleIteratorType(IteratorType):
+    def __init__(self, name, yield_type):
+        self._yield_type = yield_type
+        super(SimpleIteratorType, self).__init__(name)
+
+    @property
+    def yield_type(self):
+        return self._yield_type
+
+
+class Buffer(IterableType, ArrayCompatible):
+    """
+    Type class for objects providing the buffer protocol.
+    Derived classes exist for more specific cases.
+    """
+
+    mutable = True
+    slice_is_copy = False
+    aligned = True
+
+    # CS and FS are not reserved for inner contig but strided
+    LAYOUTS = frozenset(["C", "F", "CS", "FS", "A"])
+
+    def __init__(self, dtype, ndim, layout, readonly=False, name=None):
+        from .misc import unliteral
+
+        if isinstance(dtype, Buffer):
+            msg = (
+                "The dtype of a Buffer type cannot itself be a Buffer type, "
+                "this is unsupported behaviour."
+                "\nThe dtype requested for the unsupported Buffer was: {}."
+            )
+            raise NumbaTypeError(msg.format(dtype))
+        if layout not in self.LAYOUTS:
+            raise NumbaValueError("Invalid layout '%s'" % layout)
+        self.dtype = unliteral(dtype)
+        self.ndim = ndim
+        self.layout = layout
+        if readonly:
+            self.mutable = False
+        if name is None:
+            type_name = self.__class__.__name__.lower()
+            if readonly:
+                type_name = "readonly %s" % type_name
+            name = "%s(%s, %sd, %s)" % (type_name, dtype, ndim, layout)
+        super(Buffer, self).__init__(name)
+
+    @property
+    def iterator_type(self):
+        from .iterators import ArrayIterator
+
+        return ArrayIterator(self)
+
+    @property
+    def as_array(self):
+        return self
+
+    def copy(self, dtype=None, ndim=None, layout=None):
+        if dtype is None:
+            dtype = self.dtype
+        if ndim is None:
+            ndim = self.ndim
+        if layout is None:
+            layout = self.layout
+        return self.__class__(
+            dtype=dtype, ndim=ndim, layout=layout, readonly=not self.mutable
+        )
+
+    @property
+    def key(self):
+        return self.dtype, self.ndim, self.layout, self.mutable
+
+    @property
+    def is_c_contig(self):
+        return self.layout == "C" or (self.ndim <= 1 and self.layout in "CF")
+
+    @property
+    def is_f_contig(self):
+        return self.layout == "F" or (self.ndim <= 1 and self.layout in "CF")
+
+    @property
+    def is_contig(self):
+        return self.layout in "CF"