numba-cuda 0.19.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

numba_cuda/numba/cuda/types.py

@@ -1,4 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from numba.core import types
+from numba.core.typeconv import Conversion
 
 
 class Dim3(types.Type):
@@ -38,3 +42,58 @@ class CUDADispatcher(types.Dispatcher):
     # 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(types.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, types.Float):
+            return Conversion.unsafe
+
+        elif isinstance(other, types.Integer):
+            if other.bitwidth == 8:
+                return Conversion.safe
+            else:
+                return Conversion.unsafe
+
+    def can_convert_to(self, typingctx, other):
+        if isinstance(other, types.Float):
+            if other.bitwidth >= 32:
+                return Conversion.safe
+            else:
+                return Conversion.unsafe
+        elif isinstance(other, types.Integer):
+            return Conversion.unsafe
+
+    def unify(self, typingctx, other):
+        if isinstance(other, (types.Float, types.Integer)):
+            return typingctx.unify_pairs(self, other)
+
+
+bfloat16 = Bfloat16()

numba_cuda/numba/cuda/typing/__init__.py

@@ -1,8 +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"]
+__all__ = [
+    "signature",
+    "make_concrete_template",
+    "Signature",
+    "fold_arguments",
+    "BaseContext",
+    "Context",
+]

numba_cuda/numba/cuda/typing/cffi_utils.py

@@ -0,0 +1,55 @@
+# 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 numba.core import types
+from numba.core.errors import TypingError
+from numba.cuda.typing import templates
+
+try:
+    import cffi
+
+    ffi = cffi.FFI()
+except ImportError:
+    ffi = None
+
+SUPPORTED = ffi is not None
+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)