numba-cuda 0.19.1__py3-none-any.whl → 0.20.1__py3-none-any.whl

numba_cuda/numba/cuda/tests/support.py

@@ -21,8 +21,9 @@ from functools import cached_property
 import numpy as np
 
 from numba import types
-from numba.core import errors, config
-from numba.core.typing import cffi_utils
+from numba.core import errors
+from numba.cuda.core import config
+from numba.cuda.typing import cffi_utils
 from numba.cuda.memory_management.nrt import rtsys
 from numba.core.extending import (
     typeof_impl,
@@ -31,7 +32,7 @@ from numba.core.extending import (
     NativeValue,
 )
 from numba.core.datamodel.models import OpaqueModel
-from numba.np import numpy_support
+from numba.cuda.np import numpy_support
 
 
 class EnableNRTStatsMixin(object):
@@ -751,16 +752,55 @@ class TestCase(unittest.TestCase):
 
         return Dummy, DummyType
 
-    def skip_if_no_external_compiler(self):
-        """
-        Call this to ensure the test is skipped if no suitable external compiler
-        is found. This is a method on the TestCase opposed to a stand-alone
-        decorator so as to make it "lazy" via runtime evaluation opposed to
-        running at test-discovery time.
-        """
-        # This is a local import to avoid deprecation warnings being generated
-        # through the use of the numba.pycc module.
-        from numba.pycc.platform import external_compiler_works
 
-        if not external_compiler_works():
-            self.skipTest("No suitable external compiler was found.")
+class MemoryLeak(object):
+    __enable_leak_check = True
+
+    def memory_leak_setup(self):
+        # Clean up any NRT-backed objects hanging in a dead reference cycle
+        gc.collect()
+        self.__init_stats = rtsys.get_allocation_stats()
+
+    def memory_leak_teardown(self):
+        if self.__enable_leak_check:
+            self.assert_no_memory_leak()
+
+    def assert_no_memory_leak(self):
+        old = self.__init_stats
+        new = rtsys.get_allocation_stats()
+        total_alloc = new.alloc - old.alloc
+        total_free = new.free - old.free
+        total_mi_alloc = new.mi_alloc - old.mi_alloc
+        total_mi_free = new.mi_free - old.mi_free
+        self.assertEqual(total_alloc, total_free)
+        self.assertEqual(total_mi_alloc, total_mi_free)
+
+    def disable_leak_check(self):
+        # For per-test use when MemoryLeakMixin is injected into a TestCase
+        self.__enable_leak_check = False
+
+
+class MemoryLeakMixin(EnableNRTStatsMixin, MemoryLeak):
+    def setUp(self):
+        super(MemoryLeakMixin, self).setUp()
+        self.memory_leak_setup()
+
+    def tearDown(self):
+        gc.collect()
+        self.memory_leak_teardown()
+        super(MemoryLeakMixin, self).tearDown()
+
+
+class CheckWarningsMixin(object):
+    @contextlib.contextmanager
+    def check_warnings(self, messages, category=RuntimeWarning):
+        with warnings.catch_warnings(record=True) as catch:
+            warnings.simplefilter("always")
+            yield
+        found = 0
+        for w in catch:
+            for m in messages:
+                if m in str(w.message):
+                    self.assertEqual(w.category, category)
+                    found += 1
+        self.assertEqual(found, len(messages))

numba_cuda/numba/cuda/tests/test_tracing.py

@@ -0,0 +1,200 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from io import StringIO
+import logging
+
+import unittest
+from numba.cuda.core import tracing
+
+logger = logging.getLogger("trace")
+logger.setLevel(logging.INFO)
+
+# Make sure tracing is enabled
+orig_trace = tracing.trace
+tracing.trace = tracing.dotrace
+
+
+class CapturedTrace:
+    """Capture the trace temporarily for validation."""
+
+    def __init__(self):
+        self.buffer = StringIO()
+        self.handler = logging.StreamHandler(self.buffer)
+
+    def __enter__(self):
+        self._handlers = logger.handlers
+        self.buffer = StringIO()
+        logger.handlers = [logging.StreamHandler(self.buffer)]
+
+    def __exit__(self, type, value, traceback):
+        logger.handlers = self._handlers
+
+    def getvalue(self):
+        # Depending on how the tests are run, object names may be
+        # qualified by their containing module.
+        # Remove that to make the trace output independent from the testing mode.
+        log = self.buffer.getvalue()
+        log = log.replace(__name__ + ".", "")
+        return log
+
+
+class Class(object):
+    @tracing.trace
+    @classmethod
+    def class_method(cls):
+        pass
+
+    @tracing.trace
+    @staticmethod
+    def static_method():
+        pass
+
+    __test = None
+
+    def _test_get(self):
+        return self.__test
+
+    def _test_set(self, value):
+        self.__test = value
+
+    test = tracing.trace(property(_test_get, _test_set))
+
+    @tracing.trace
+    def method(self, some, other="value", *args, **kwds):
+        pass
+
+    def __repr__(self):
+        """Generate a deterministic string for testing."""
+        return "<Class instance>"
+
+
+class Class2(object):
+    @classmethod
+    def class_method(cls):
+        pass
+
+    @staticmethod
+    def static_method():
+        pass
+
+    __test = None
+
+    @property
+    def test(self):
+        return self.__test
+
+    @test.setter
+    def test(self, value):
+        self.__test = value
+
+    def method(self):
+        pass
+
+    def __str__(self):
+        return "Test(" + str(self.test) + ")"
+
+    def __repr__(self):
+        """Generate a deterministic string for testing."""
+        return "<Class2 instance>"
+
+
+@tracing.trace
+def test_traced_function():
+    # Test the tracing functionality with fixed values
+    x, y = 5, 5
+    z = True
+
+    a = x + y
+    b = x * y
+    if z:
+        result = a
+    else:
+        result = b
+
+    # The function should return 10 (5 + 5) when z is True
+    assert result == 10
+
+
+class TestTracing(unittest.TestCase):
+    def __init__(self, *args):
+        super(TestTracing, self).__init__(*args)
+
+    def setUp(self):
+        self.capture = CapturedTrace()
+
+    def tearDown(self):
+        del self.capture
+
+    def test_method(self):
+        with self.capture:
+            Class().method("foo", bar="baz")
+        self.assertEqual(
+            self.capture.getvalue(),
+            ">> Class.method(self=<Class instance>, some='foo', other='value', bar='baz')\n"
+            + "<< Class.method\n",
+        )
+
+    def test_class_method(self):
+        with self.capture:
+            Class.class_method()
+        self.assertEqual(
+            self.capture.getvalue(),
+            ">> Class.class_method(cls=<class 'Class'>)\n"
+            + "<< Class.class_method\n",
+        )
+
+    def test_static_method(self):
+        with self.capture:
+            Class.static_method()
+        self.assertEqual(
+            self.capture.getvalue(),
+            ">> static_method()\n" + "<< static_method\n",
+        )
+
+    def test_property(self):
+        with self.capture:
+            test = Class()
+            test.test = 1
+            assert 1 == test.test
+        self.assertEqual(
+            self.capture.getvalue(),
+            ">> Class._test_set(self=<Class instance>, value=1)\n"
+            + "<< Class._test_set\n"
+            + ">> Class._test_get(self=<Class instance>)\n"
+            + "<< Class._test_get -> 1\n",
+        )
+
+    def test_function(self):
+        with self.capture:
+            test_traced_function()
+        # The test function should be traced when called
+        trace_output = self.capture.getvalue()
+        self.assertIn(">> test_traced_function()", trace_output)
+        self.assertIn("<< test_traced_function", trace_output)
+
+    @unittest.skip("recursive decoration not yet implemented")
+    def test_injected(self):
+        with self.capture:
+            tracing.trace(Class2, recursive=True)
+            Class2.class_method()
+            Class2.static_method()
+            test = Class2()
+            test.test = 1
+            assert 1 == test.test
+            test.method()
+
+            self.assertEqual(
+                self.capture.getvalue(),
+                ">> Class2.class_method(cls=<type 'Class2'>)\n"
+                + "<< Class2.class_method\n"
+                ">> static_method()\n"
+                "<< static_method\n",
+            )
+
+
+# Reset tracing to its original value
+tracing.trace = orig_trace
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/types.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-2-Clause
 
 from numba.core import types
+from numba.core.typeconv import Conversion
 
 
 class Dim3(types.Type):
@@ -41,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

@@ -7,5 +7,13 @@ from .templates import (
     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)