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

numba_cuda/numba/cuda/lowering.py

@@ -10,18 +10,14 @@ from llvmlite import ir as llvm_ir
 
 from numba.core import (
     typing,
-    utils,
     types,
     ir,
-    debuginfo,
-    funcdesc,
     generators,
-    config,
-    cgutils,
     removerefctpass,
-    targetconfig,
 )
-from numba.cuda.core import ir_utils
+from numba.cuda import debuginfo, cgutils, utils
+from numba.cuda.core import ir_utils, targetconfig, funcdesc, config
+
 from numba.core.errors import (
     LoweringError,
     new_error_context,
@@ -30,8 +26,8 @@ from numba.core.errors import (
     UnsupportedError,
     NumbaDebugInfoWarning,
 )
-from numba.core.funcdesc import default_mangler
-from numba.core.environment import Environment
+from numba.cuda.core.funcdesc import default_mangler
+from numba.cuda.core.environment import Environment
 from numba.core.analysis import compute_use_defs, must_use_alloca
 from numba.misc.firstlinefinder import get_func_body_first_lineno
 from numba import version_info
@@ -466,7 +462,7 @@ class Lower(BaseLower):
         self._blk_local_varmap = {}
 
     def pre_block(self, block):
-        from numba.core.unsafe import eh
+        from numba.cuda.core.unsafe import eh
 
         super(Lower, self).pre_block(block)
         self._cur_ir_block = block
@@ -1029,9 +1025,6 @@ class Lower(BaseLower):
         elif isinstance(fnty, types.RecursiveCall):
             res = self._lower_call_RecursiveCall(fnty, expr, signature)
 
-        elif isinstance(fnty, types.FunctionType):
-            res = self._lower_call_FunctionType(fnty, expr, signature)
-
         else:
             res = self._lower_call_normal(fnty, expr, signature)
 
@@ -1052,7 +1045,7 @@ class Lower(BaseLower):
         )
 
     def _lower_call_ObjModeDispatcher(self, fnty, expr, signature):
-        from numba.core.pythonapi import ObjModeUtils
+        from numba.cuda.core.pythonapi import ObjModeUtils
 
         self.init_pyapi()
         # Acquire the GIL
@@ -1229,136 +1222,6 @@ class Lower(BaseLower):
             )
         return res
 
-    def _lower_call_FunctionType(self, fnty, expr, signature):
-        self.debug_print("# calling first-class function type")
-        sig = types.unliteral(signature)
-        if not fnty.check_signature(signature):
-            # value dependent polymorphism?
-            raise UnsupportedError(
-                f"mismatch of function types:"
-                f" expected {fnty} but got {types.FunctionType(sig)}"
-            )
-        argvals = self.fold_call_args(
-            fnty,
-            sig,
-            expr.args,
-            expr.vararg,
-            expr.kws,
-        )
-        return self.__call_first_class_function_pointer(
-            fnty.ftype,
-            expr.func.name,
-            sig,
-            argvals,
-        )
-
-    def __call_first_class_function_pointer(self, ftype, fname, sig, argvals):
-        """
-        Calls a first-class function pointer.
-
-        This function is responsible for calling a first-class function pointer,
-        which can either be a JIT-compiled function or a Python function. It
-        determines if a JIT address is available, and if so, calls the function
-        using the JIT address. Otherwise, it calls the function using a function
-        pointer obtained from the `__get_first_class_function_pointer` method.
-
-        Args:
-            ftype: The type of the function.
-            fname: The name of the function.
-            sig: The signature of the function.
-            argvals: The argument values to pass to the function.
-
-        Returns:
-            The result of calling the function.
-        """
-        context = self.context
-        builder = self.builder
-        # Determine if jit address is available
-        fstruct = self.loadvar(fname)
-        struct = cgutils.create_struct_proxy(self.typeof(fname))(
-            context, builder, value=fstruct
-        )
-        jit_addr = struct.jit_addr
-        jit_addr.name = f"jit_addr_of_{fname}"
-
-        ctx = context
-        res_slot = cgutils.alloca_once(
-            builder, ctx.get_value_type(sig.return_type)
-        )
-
-        if_jit_addr_is_null = builder.if_else(
-            cgutils.is_null(builder, jit_addr), likely=False
-        )
-        with if_jit_addr_is_null as (then, orelse):
-            with then:
-                func_ptr = self.__get_first_class_function_pointer(
-                    ftype, fname, sig
-                )
-                res = builder.call(func_ptr, argvals)
-                builder.store(res, res_slot)
-
-            with orelse:
-                llty = ctx.call_conv.get_function_type(
-                    sig.return_type, sig.args
-                ).as_pointer()
-                func_ptr = builder.bitcast(jit_addr, llty)
-                # call
-                status, res = ctx.call_conv.call_function(
-                    builder, func_ptr, sig.return_type, sig.args, argvals
-                )
-                with cgutils.if_unlikely(builder, status.is_error):
-                    context.call_conv.return_status_propagate(builder, status)
-                builder.store(res, res_slot)
-        return builder.load(res_slot)
-
-    def __get_first_class_function_pointer(self, ftype, fname, sig):
-        from numba.experimental.function_type import lower_get_wrapper_address
-
-        llty = self.context.get_value_type(ftype)
-        fstruct = self.loadvar(fname)
-        addr = self.builder.extract_value(
-            fstruct, 0, name="addr_of_%s" % (fname)
-        )
-
-        fptr = cgutils.alloca_once(
-            self.builder, llty, name="fptr_of_%s" % (fname)
-        )
-        with self.builder.if_else(
-            cgutils.is_null(self.builder, addr), likely=False
-        ) as (then, orelse):
-            with then:
-                self.init_pyapi()
-                # Acquire the GIL
-                gil_state = self.pyapi.gil_ensure()
-                pyaddr = self.builder.extract_value(
-                    fstruct, 1, name="pyaddr_of_%s" % (fname)
-                )
-                # try to recover the function address, see
-                # test_zero_address BadToGood example in
-                # test_function_type.py
-                addr1 = lower_get_wrapper_address(
-                    self.context,
-                    self.builder,
-                    pyaddr,
-                    sig,
-                    failure_mode="ignore",
-                )
-                with self.builder.if_then(
-                    cgutils.is_null(self.builder, addr1), likely=False
-                ):
-                    self.return_exception(
-                        RuntimeError,
-                        exc_args=(f"{ftype} function address is null",),
-                        loc=self.loc,
-                    )
-                addr2 = self.pyapi.long_as_voidptr(addr1)
-                self.builder.store(self.builder.bitcast(addr2, llty), fptr)
-                self.pyapi.decref(addr1)
-                self.pyapi.gil_release(gil_state)
-            with orelse:
-                self.builder.store(self.builder.bitcast(addr, llty), fptr)
-        return self.builder.load(fptr)
-
     def _lower_call_normal(self, fnty, expr, signature):
         # Normal function resolution
         self.debug_print("# calling normal function: {0}".format(fnty))

numba_cuda/numba/cuda/mathimpl.py

@@ -4,11 +4,12 @@
 import math
 import operator
 from llvmlite import ir
-from numba.core import types, typing, targetconfig
+from numba.core import types, typing
 from numba.cuda import cgutils
 from numba.core.imputils import Registry
 from numba.types import float32, float64, int64, uint64
 from numba.cuda import libdevice
+from numba.cuda.core import targetconfig
 
 registry = Registry()
 lower = registry.lower

numba_cuda/numba/cuda/memory_management/nrt.py

@@ -6,7 +6,10 @@ import os
 from functools import wraps
 import numpy as np
 
-from numba import cuda, config
+
+from numba import cuda, types
+from numba.cuda import config
+
 from numba.core.runtime.nrt import _nrt_mstats
 from numba.cuda.cudadrv.driver import (
     _Linker,
@@ -17,24 +20,11 @@ from numba.cuda.cudadrv.driver import (
 )
 from numba.cuda.cudadrv import devices
 from numba.cuda.api import get_current_device
-from numba.cuda.utils import _readenv, cached_file_read
+from numba.cuda.utils import cached_file_read
 from numba.cuda.cudadrv.linkable_code import CUSource
+from numba.cuda.typing.templates import signature
 
-
-# Check environment variable or config for NRT statistics enablement
-NRT_STATS = _readenv("NUMBA_CUDA_NRT_STATS", bool, False) or getattr(
-    config, "NUMBA_CUDA_NRT_STATS", False
-)
-if not hasattr(config, "NUMBA_CUDA_NRT_STATS"):
-    config.CUDA_NRT_STATS = NRT_STATS
-
-
-# Check environment variable or config for NRT enablement
-ENABLE_NRT = _readenv("NUMBA_CUDA_ENABLE_NRT", bool, False) or getattr(
-    config, "NUMBA_CUDA_ENABLE_NRT", False
-)
-if not hasattr(config, "NUMBA_CUDA_ENABLE_NRT"):
-    config.CUDA_ENABLE_NRT = ENABLE_NRT
+from numba.core.extending import intrinsic, overload_classmethod
 
 
 def get_include():
@@ -42,6 +32,34 @@ def get_include():
     return os.path.dirname(os.path.abspath(__file__))
 
 
+# Provide an implementation of Array._allocate() for the CUDA target (used
+# internally by Numba when generating the allocation of an array)
+
+
+@intrinsic
+def intrin_alloc(typingctx, allocsize, align):
+    """Intrinsic to call into the allocator for Array"""
+
+    def codegen(context, builder, signature, args):
+        allocsize, align = args
+        meminfo = context.nrt.meminfo_alloc_aligned(builder, allocsize, align)
+        return meminfo
+
+    mip = types.MemInfoPointer(types.voidptr)  # return untyped pointer
+    sig = signature(mip, allocsize, align)
+    return sig, codegen
+
+
+@overload_classmethod(types.Array, "_allocate", target="CUDA")
+def _ol_array_allocate(cls, allocsize, align):
+    """Implements a Numba-only CUDA-target classmethod on the array type."""
+
+    def impl(cls, allocsize, align):
+        return intrin_alloc(allocsize, align)
+
+    return impl
+
+
 # Protect method to ensure NRT memory allocation and initialization
 def _alloc_init_guard(method):
     """
@@ -69,10 +87,18 @@ class _Runtime:
 
     def __init__(self):
         """Initialize memsys module and variable"""
+        self._reset()
+
+    def _reset(self):
+        """Reset to the uninitialized state"""
         self._memsys_module = None
         self._memsys = None
         self._initialized = False
 
+    def close(self):
+        """Close and reset"""
+        self._reset()
+
     def _compile_memsys_module(self):
         """
         Compile memsys.cu and create a module from it in the current context

numba_cuda/numba/cuda/misc/findlib.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+import os
+import re
+
+
+def get_lib_dirs():
+    """
+    Anaconda specific
+    """
+    if sys.platform == "win32":
+        # CUDA 12 puts in "bin" directory, whereas CUDA 13 puts in "bin\x64" directory
+        dirnames = [
+            os.path.join("Library", "bin"),
+            os.path.join("Library", "bin", "x64"),
+            os.path.join("Library", "nvvm", "bin"),
+            os.path.join("Library", "nvvm", "bin", "x64"),
+        ]
+    else:
+        dirnames = [
+            "lib",
+        ]
+    libdirs = [os.path.join(sys.prefix, x) for x in dirnames]
+    return libdirs
+
+
+DLLNAMEMAP = {
+    "linux": r"lib%(name)s\.so\.%(ver)s$",
+    "linux2": r"lib%(name)s\.so\.%(ver)s$",
+    "linux-static": r"lib%(name)s\.a$",
+    "darwin": r"lib%(name)s\.%(ver)s\.dylib$",
+    "win32": r"%(name)s%(ver)s\.dll$",
+    "win32-static": r"%(name)s\.lib$",
+    "bsd": r"lib%(name)s\.so\.%(ver)s$",
+}
+
+RE_VER = r"[0-9]*([_\.][0-9]+)*"
+
+
+def find_lib(libname, libdir=None, platform=None, static=False):
+    platform = platform or sys.platform
+    platform = "bsd" if "bsd" in platform else platform
+    if static:
+        platform = f"{platform}-static"
+    if platform not in DLLNAMEMAP:
+        # Return empty list if platform name is undefined.
+        # Not all platforms define their static library paths.
+        return []
+    pat = DLLNAMEMAP[platform] % {"name": libname, "ver": RE_VER}
+    regex = re.compile(pat)
+    return find_file(regex, libdir)
+
+
+def find_file(pat, libdir=None):
+    if libdir is None:
+        libdirs = get_lib_dirs()
+    elif isinstance(libdir, str):
+        libdirs = [
+            libdir,
+        ]
+    else:
+        libdirs = list(libdir)
+    files = []
+    for ldir in libdirs:
+        try:
+            entries = os.listdir(ldir)
+        except FileNotFoundError:
+            continue
+        candidates = [
+            os.path.join(ldir, ent) for ent in entries if pat.match(ent)
+        ]
+        files.extend([c for c in candidates if os.path.isfile(c)])
+    return files

numba_cuda/numba/cuda/models.py

@@ -6,9 +6,10 @@ import functools
 from llvmlite import ir
 
 from numba.core.datamodel.registry import DataModelManager, register
+from numba.core.datamodel import PrimitiveModel
 from numba.core.extending import models
 from numba.core import types
-from numba.cuda.types import Dim3, GridGroup, CUDADispatcher
+from numba.cuda.types import Dim3, GridGroup, CUDADispatcher, Bfloat16
 
 
 cuda_data_manager = DataModelManager()
@@ -45,3 +46,10 @@ class FloatModel(models.PrimitiveModel):
 
 
 register_model(CUDADispatcher)(models.OpaqueModel)
+
+
+@register_model(Bfloat16)
+class _model___nv_bfloat16(PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        be_type = ir.IntType(16)
+        super(_model___nv_bfloat16, self).__init__(dmm, fe_type, be_type)

numba_cuda/numba/cuda/np/npdatetime_helpers.py

@@ -0,0 +1,217 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Helper functions for np.timedelta64 and np.datetime64.
+For now, multiples-of-units (for example timedeltas expressed in tens
+of seconds) are not supported.
+"""
+
+import numpy as np
+
+
+DATETIME_UNITS = {
+    "Y": 0,  # Years
+    "M": 1,  # Months
+    "W": 2,  # Weeks
+    # Yes, there's a gap here
+    "D": 4,  # Days
+    "h": 5,  # Hours
+    "m": 6,  # Minutes
+    "s": 7,  # Seconds
+    "ms": 8,  # Milliseconds
+    "us": 9,  # Microseconds
+    "ns": 10,  # Nanoseconds
+    "ps": 11,  # Picoseconds
+    "fs": 12,  # Femtoseconds
+    "as": 13,  # Attoseconds
+    "": 14,  # "generic", i.e. unit-less
+}
+
+NAT = np.timedelta64("nat").astype(np.int64)
+
+# NOTE: numpy has several inconsistent functions for timedelta casting:
+# - can_cast_timedelta64_{metadata,units}() disallows "safe" casting
+#   to and from generic units
+# - cast_timedelta_to_timedelta() allows casting from (but not to)
+#   generic units
+# - compute_datetime_metadata_greatest_common_divisor() allows casting from
+#   generic units (used for promotion)
+
+
+def same_kind(src, dest):
+    """
+    Whether the *src* and *dest* units are of the same kind.
+    """
+    return (DATETIME_UNITS[src] < 5) == (DATETIME_UNITS[dest] < 5)
+
+
+def can_cast_timedelta_units(src, dest):
+    # Mimic NumPy's "safe" casting and promotion
+    # `dest` must be more precise than `src` and they must be compatible
+    # for conversion.
+    # XXX should we switch to enforcing "same-kind" for Numpy 1.10+ ?
+    src = DATETIME_UNITS[src]
+    dest = DATETIME_UNITS[dest]
+    if src == dest:
+        return True
+    if src == 14:
+        return True
+    if src > dest:
+        return False
+    if dest == 14:
+        # unit-less timedelta64 is not compatible with anything else
+        return False
+    if src <= 1 and dest > 1:
+        # Cannot convert between months or years and other units
+        return False
+    return True
+
+
+# Exact conversion factors from one unit to the immediately more precise one
+_factors = {
+    0: (1, 12),  # Years -> Months
+    2: (4, 7),  # Weeks -> Days
+    4: (5, 24),  # Days -> Hours
+    5: (6, 60),  # Hours -> Minutes
+    6: (7, 60),  # Minutes -> Seconds
+    7: (8, 1000),
+    8: (9, 1000),
+    9: (10, 1000),
+    10: (11, 1000),
+    11: (12, 1000),
+    12: (13, 1000),
+}
+
+
+def _get_conversion_multiplier(big_unit_code, small_unit_code):
+    """
+    Return an integer multiplier allowing to convert from *big_unit_code*
+    to *small_unit_code*.
+    None is returned if the conversion is not possible through a
+    simple integer multiplication.
+    """
+    # Mimics get_datetime_units_factor() in NumPy's datetime.c,
+    # with a twist to allow no-op conversion from generic units.
+    if big_unit_code == 14:
+        return 1
+    c = big_unit_code
+    factor = 1
+    while c < small_unit_code:
+        try:
+            c, mult = _factors[c]
+        except KeyError:
+            # No possible conversion
+            return None
+        factor *= mult
+    if c == small_unit_code:
+        return factor
+    else:
+        return None
+
+
+def get_timedelta_conversion_factor(src_unit, dest_unit):
+    """
+    Return an integer multiplier allowing to convert from timedeltas
+    of *src_unit* to *dest_unit*.
+    """
+    return _get_conversion_multiplier(
+        DATETIME_UNITS[src_unit], DATETIME_UNITS[dest_unit]
+    )
+
+
+def get_datetime_timedelta_conversion(datetime_unit, timedelta_unit):
+    """
+    Compute a possible conversion for combining *datetime_unit* and
+    *timedelta_unit* (presumably for adding or subtracting).
+    Return (result unit, integer datetime multiplier, integer timedelta
+    multiplier). RuntimeError is raised if the combination is impossible.
+    """
+    # XXX now unused (I don't know where / how Numpy uses this)
+    dt_unit_code = DATETIME_UNITS[datetime_unit]
+    td_unit_code = DATETIME_UNITS[timedelta_unit]
+    if td_unit_code == 14 or dt_unit_code == 14:
+        return datetime_unit, 1, 1
+    if td_unit_code < 2 and dt_unit_code >= 2:
+        # Cannot combine Y or M timedelta64 with a finer-grained datetime64
+        raise RuntimeError(
+            "cannot combine datetime64(%r) and timedelta64(%r)"
+            % (datetime_unit, timedelta_unit)
+        )
+    dt_factor, td_factor = 1, 1
+
+    # If years or months, the datetime unit is first scaled to weeks or days,
+    # then conversion continues below.  This is the same algorithm as used
+    # in Numpy's get_datetime_conversion_factor() (src/multiarray/datetime.c):
+    # """Conversions between years/months and other units use
+    # the factor averaged over the 400 year leap year cycle."""
+    if dt_unit_code == 0:
+        if td_unit_code >= 4:
+            dt_factor = 97 + 400 * 365
+            td_factor = 400
+            dt_unit_code = 4
+        elif td_unit_code == 2:
+            dt_factor = 97 + 400 * 365
+            td_factor = 400 * 7
+            dt_unit_code = 2
+    elif dt_unit_code == 1:
+        if td_unit_code >= 4:
+            dt_factor = 97 + 400 * 365
+            td_factor = 400 * 12
+            dt_unit_code = 4
+        elif td_unit_code == 2:
+            dt_factor = 97 + 400 * 365
+            td_factor = 400 * 12 * 7
+            dt_unit_code = 2
+
+    if td_unit_code >= dt_unit_code:
+        factor = _get_conversion_multiplier(dt_unit_code, td_unit_code)
+        assert factor is not None, (dt_unit_code, td_unit_code)
+        return timedelta_unit, dt_factor * factor, td_factor
+    else:
+        factor = _get_conversion_multiplier(td_unit_code, dt_unit_code)
+        assert factor is not None, (dt_unit_code, td_unit_code)
+        return datetime_unit, dt_factor, td_factor * factor
+
+
+def combine_datetime_timedelta_units(datetime_unit, timedelta_unit):
+    """
+    Return the unit result of combining *datetime_unit* with *timedelta_unit*
+    (e.g. by adding or subtracting).  None is returned if combining
+    those units is forbidden.
+    """
+    dt_unit_code = DATETIME_UNITS[datetime_unit]
+    td_unit_code = DATETIME_UNITS[timedelta_unit]
+    if dt_unit_code == 14:
+        return timedelta_unit
+    elif td_unit_code == 14:
+        return datetime_unit
+    if td_unit_code < 2 and dt_unit_code >= 2:
+        return None
+    if dt_unit_code > td_unit_code:
+        return datetime_unit
+    else:
+        return timedelta_unit
+
+
+def get_best_unit(unit_a, unit_b):
+    """
+    Get the best (i.e. finer-grained) of two units.
+    """
+    a = DATETIME_UNITS[unit_a]
+    b = DATETIME_UNITS[unit_b]
+    if a == 14:
+        return unit_b
+    if b == 14:
+        return unit_a
+    if b > a:
+        return unit_b
+    return unit_a
+
+
+def datetime_minimum(a, b):
+    pass
+
+
+def datetime_maximum(a, b):
+    pass