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

numba_cuda/numba/cuda/types/cuda_scalars.py

@@ -0,0 +1,280 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import enum
+
+import numpy as np
+
+from .abstract import Dummy, Hashable, Literal, Number, Type
+from functools import total_ordering, cached_property
+from numba.cuda.typeconv import Conversion
+from numba.cuda.np import npdatetime_helpers
+
+
+class Boolean(Hashable):
+    def cast_python_value(self, value):
+        return bool(value)
+
+
+def parse_integer_bitwidth(name):
+    for prefix in ("int", "uint"):
+        if name.startswith(prefix):
+            bitwidth = int(name[len(prefix) :])
+    return bitwidth
+
+
+def parse_integer_signed(name):
+    signed = name.startswith("int")
+    return signed
+
+
+@total_ordering
+class Integer(Number):
+    def __init__(self, name, bitwidth=None, signed=None):
+        super(Integer, self).__init__(name)
+        if bitwidth is None:
+            bitwidth = parse_integer_bitwidth(name)
+        if signed is None:
+            signed = parse_integer_signed(name)
+        self.bitwidth = bitwidth
+        self.signed = signed
+
+    @classmethod
+    def from_bitwidth(cls, bitwidth, signed=True):
+        name = ("int%d" if signed else "uint%d") % bitwidth
+        return cls(name)
+
+    def cast_python_value(self, value):
+        return getattr(np, self.name)(value)
+
+    def __lt__(self, other):
+        if self.__class__ is not other.__class__:
+            return NotImplemented
+        if self.signed != other.signed:
+            return NotImplemented
+        return self.bitwidth < other.bitwidth
+
+    @property
+    def maxval(self):
+        """
+        The maximum value representable by this type.
+        """
+        if self.signed:
+            return (1 << (self.bitwidth - 1)) - 1
+        else:
+            return (1 << self.bitwidth) - 1
+
+    @property
+    def minval(self):
+        """
+        The minimal value representable by this type.
+        """
+        if self.signed:
+            return -(1 << (self.bitwidth - 1))
+        else:
+            return 0
+
+
+class IntegerLiteral(Literal, Integer):
+    def __init__(self, value):
+        self._literal_init(value)
+        name = "Literal[int]({})".format(value)
+        basetype = self.literal_type
+        Integer.__init__(
+            self,
+            name=name,
+            bitwidth=basetype.bitwidth,
+            signed=basetype.signed,
+        )
+
+    def can_convert_to(self, typingctx, other):
+        conv = typingctx.can_convert(self.literal_type, other)
+        if conv is not None:
+            return max(conv, Conversion.promote)
+
+
+Literal.ctor_map[int] = IntegerLiteral
+
+
+class BooleanLiteral(Literal, Boolean):
+    def __init__(self, value):
+        self._literal_init(value)
+        name = "Literal[bool]({})".format(value)
+        Boolean.__init__(self, name=name)
+
+    def can_convert_to(self, typingctx, other):
+        conv = typingctx.can_convert(self.literal_type, other)
+        if conv is not None:
+            return max(conv, Conversion.promote)
+
+
+Literal.ctor_map[bool] = BooleanLiteral
+
+
+@total_ordering
+class Float(Number):
+    def __init__(self, *args, **kws):
+        super(Float, self).__init__(*args, **kws)
+        # Determine bitwidth
+        assert self.name.startswith("float")
+        bitwidth = int(self.name[5:])
+        self.bitwidth = bitwidth
+
+    def cast_python_value(self, value):
+        return getattr(np, self.name)(value)
+
+    def __lt__(self, other):
+        if self.__class__ is not other.__class__:
+            return NotImplemented
+        return self.bitwidth < other.bitwidth
+
+
+@total_ordering
+class Complex(Number):
+    def __init__(self, name, underlying_float, **kwargs):
+        super(Complex, self).__init__(name, **kwargs)
+        self.underlying_float = underlying_float
+        # Determine bitwidth
+        assert self.name.startswith("complex")
+        bitwidth = int(self.name[7:])
+        self.bitwidth = bitwidth
+
+    def cast_python_value(self, value):
+        return getattr(np, self.name)(value)
+
+    def __lt__(self, other):
+        if self.__class__ is not other.__class__:
+            return NotImplemented
+        return self.bitwidth < other.bitwidth
+
+
+class _NPDatetimeBase(Type):
+    """
+    Common base class for np.datetime64 and np.timedelta64.
+    """
+
+    def __init__(self, unit, *args, **kws):
+        name = "%s[%s]" % (self.type_name, unit)
+        self.unit = unit
+        self.unit_code = npdatetime_helpers.DATETIME_UNITS[self.unit]
+        super(_NPDatetimeBase, self).__init__(name, *args, **kws)
+
+    def __lt__(self, other):
+        if self.__class__ is not other.__class__:
+            return NotImplemented
+        # A coarser-grained unit is "smaller", i.e. less precise values
+        # can be represented (but the magnitude of representable values is
+        # also greater...).
+        return self.unit_code < other.unit_code
+
+    def cast_python_value(self, value):
+        cls = getattr(np, self.type_name)
+        if self.unit:
+            return cls(value, self.unit)
+        else:
+            return cls(value)
+
+
+@total_ordering
+class NPTimedelta(_NPDatetimeBase):
+    type_name = "timedelta64"
+
+
+@total_ordering
+class NPDatetime(_NPDatetimeBase):
+    type_name = "datetime64"
+
+
+class EnumClass(Dummy):
+    """
+    Type class for Enum classes.
+    """
+
+    basename = "Enum class"
+
+    def __init__(self, cls, dtype):
+        assert isinstance(cls, type)
+        assert isinstance(dtype, Type)
+        self.instance_class = cls
+        self.dtype = dtype
+        name = "%s<%s>(%s)" % (
+            self.basename,
+            self.dtype,
+            self.instance_class.__name__,
+        )
+        super(EnumClass, self).__init__(name)
+
+    @property
+    def key(self):
+        return self.instance_class, self.dtype
+
+    @cached_property
+    def member_type(self):
+        """
+        The type of this class' members.
+        """
+        return EnumMember(self.instance_class, self.dtype)
+
+
+class IntEnumClass(EnumClass):
+    """
+    Type class for IntEnum classes.
+    """
+
+    basename = "IntEnum class"
+
+    @cached_property
+    def member_type(self):
+        """
+        The type of this class' members.
+        """
+        return IntEnumMember(self.instance_class, self.dtype)
+
+
+class EnumMember(Type):
+    """
+    Type class for Enum members.
+    """
+
+    basename = "Enum"
+    class_type_class = EnumClass
+
+    def __init__(self, cls, dtype):
+        assert isinstance(cls, type)
+        assert isinstance(dtype, Type)
+        self.instance_class = cls
+        self.dtype = dtype
+        name = "%s<%s>(%s)" % (
+            self.basename,
+            self.dtype,
+            self.instance_class.__name__,
+        )
+        super(EnumMember, self).__init__(name)
+
+    @property
+    def key(self):
+        return self.instance_class, self.dtype
+
+    @property
+    def class_type(self):
+        """
+        The type of this member's class.
+        """
+        return self.class_type_class(self.instance_class, self.dtype)
+
+
+class IntEnumMember(EnumMember):
+    """
+    Type class for IntEnum members.
+    """
+
+    basename = "IntEnum"
+    class_type_class = IntEnumClass
+
+    def can_convert_to(self, typingctx, other):
+        """
+        Convert IntEnum members to plain integers.
+        """
+        if issubclass(self.instance_class, enum.IntEnum):
+            conv = typingctx.can_convert(self.dtype, other)
+            return max(conv, Conversion.safe)

numba_cuda/numba/cuda/types/ext_types.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from . import Dispatcher, Number, Integer, Float
+from .abstract import Type
+
+from numba.cuda.typeconv import Conversion
+
+
+class Dim3(Type):
+    """
+    A 3-tuple (x, y, z) representing the position of a block or thread.
+    """
+
+    def __init__(self):
+        super().__init__(name="Dim3")
+
+
+class GridGroup(Type):
+    """
+    The grid of all threads in a cooperative kernel launch.
+    """
+
+    def __init__(self):
+        super().__init__(name="GridGroup")
+
+
+dim3 = Dim3()
+grid_group = GridGroup()
+
+
+class CUDADispatcher(Dispatcher):
+    """The type of CUDA dispatchers"""
+
+    # This type exists (instead of using types.Dispatcher as the type of CUDA
+    # dispatchers) so that we can have an alternative lowering for them to the
+    # lowering of CPU dispatchers - the CPU target lowers all dispatchers as a
+    # constant address, but we need to lower to a dummy value because it's not
+    # generally valid to use the address of CUDA kernels and functions.
+    #
+    # Notes: it may be a bug in the CPU target that it lowers all dispatchers to
+    # a constant address - it should perhaps only lower dispatchers acting as
+    # first-class functions to a constant address. Even if that bug is fixed, it
+    # is still probably a good idea to have a separate type for CUDA
+    # dispatchers, and this type might get other differentiation from the CPU
+    # dispatcher type in future.
+
+
+class Bfloat16(Number):
+    """
+    A bfloat16 type. Has 8 exponent bits and 7 significand bits.
+
+    Conversion rules:
+    Floats:
+    from:
+        fp32, fp64: UNSAFE
+        fp16: UNSAFE (loses precision)
+    to:
+        fp32, fp64: PROMOTE (same exponent, more mantissa)
+        fp16: UNSAFE (loses range)
+
+    Integers:
+    from:
+        int8: SAFE
+        other int: All UNSAFE (bf16 cannot represent all integers in range)
+    to: UNSAFE (loses precision, round to zeros)
+
+    All other conversions are not allowed.
+    """
+
+    def __init__(self):
+        super().__init__(name="__nv_bfloat16")
+
+        self.alignof_ = 2
+        self.bitwidth = 16
+
+    def can_convert_from(self, typingctx, other):
+        if isinstance(other, Float):
+            return Conversion.unsafe
+
+        elif isinstance(other, Integer):
+            if other.bitwidth == 8:
+                return Conversion.safe
+            else:
+                return Conversion.unsafe
+
+    def can_convert_to(self, typingctx, other):
+        if isinstance(other, Float):
+            if other.bitwidth >= 32:
+                return Conversion.safe
+            else:
+                return Conversion.unsafe
+        elif isinstance(other, Integer):
+            return Conversion.unsafe
+
+    def unify(self, typingctx, other):
+        if isinstance(other, (Float, Integer)):
+            return typingctx.unify_pairs(self, other)
+
+
+bfloat16 = Bfloat16()

numba_cuda/numba/cuda/types/function_type.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(),
+    "numba.core.types.function_type",
+    "numba.cuda.types.cuda_function_type",
+)

numba_cuda/numba/cuda/types/functions.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.types.functions", "numba.cuda.types.cuda_functions"
+)

numba_cuda/numba/cuda/types/iterators.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.types.iterators", "numba.cuda.types.cuda_iterators"
+)

numba_cuda/numba/cuda/types/misc.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.types.misc", "numba.cuda.types.cuda_misc"
+)

numba_cuda/numba/cuda/types/npytypes.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.types.npytypes", "numba.cuda.types.cuda_npytypes"
+)

numba_cuda/numba/cuda/types/scalars.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.types.scalars", "numba.cuda.types.cuda_scalars"
+)

numba_cuda/numba/cuda/typing/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from .templates import (
+    signature,
+    make_concrete_template,
+    Signature,
+    fold_arguments,
+)
+from .context import BaseContext, Context
+
+__all__ = [
+    "signature",
+    "make_concrete_template",
+    "Signature",
+    "fold_arguments",
+    "BaseContext",
+    "Context",
+]