numba-cuda 0.21.1__cp313-cp313-win_amd64.whl → 0.23.0__cp313-cp313-win_amd64.whl

numba_cuda/VERSION

@@ -1 +1 @@
-0.21.1
+0.23.0

numba_cuda/numba/cuda/api.py

@@ -21,6 +21,7 @@ current_context = devices.get_context
 gpus = devices.gpus
 
 
+@require_context
 def from_cuda_array_interface(desc, owner=None, sync=True):
     """Create a DeviceNDArray from a cuda-array-interface description.
     The ``owner`` is the owner of the underlying memory.
@@ -47,7 +48,9 @@ def from_cuda_array_interface(desc, owner=None, sync=True):
 
     cudevptr_class = driver.binding.CUdeviceptr
     devptr = cudevptr_class(desc["data"][0])
-    data = driver.MemoryPointer(devptr, size=size, owner=owner)
+    data = driver.MemoryPointer(
+        current_context(), devptr, size=size, owner=owner
+    )
     stream_ptr = desc.get("stream", None)
     if stream_ptr is not None:
         stream = external_stream(stream_ptr)

numba_cuda/numba/cuda/cext/_dispatcher.cpp

@@ -12,7 +12,6 @@
 #include "frameobject.h"
 #include "traceback.h"
 #include "typeconv.hpp"
-#include "_devicearray.h"
 
 /*
  * Notes on the C_TRACE macro:
@@ -940,37 +939,6 @@ CLEANUP:
     return retval;
 }
 
-static int
-import_devicearray(void)
-{
-    PyObject *devicearray = PyImport_ImportModule(NUMBA_DEVICEARRAY_IMPORT_NAME);
-    if (devicearray == NULL) {
-        return -1;
-    }
-
-    PyObject *d = PyModule_GetDict(devicearray);
-    if (d == NULL) {
-      Py_DECREF(devicearray);
-      return -1;
-    }
-
-    PyObject *key = PyUnicode_FromString("_DEVICEARRAY_API");
-    PyObject *c_api = PyDict_GetItemWithError(d, key);
-    int retcode = 0;
-    if (PyCapsule_IsValid(c_api, NUMBA_DEVICEARRAY_IMPORT_NAME "._DEVICEARRAY_API")) {
-      DeviceArray_API = (void**)PyCapsule_GetPointer(c_api, NUMBA_DEVICEARRAY_IMPORT_NAME "._DEVICEARRAY_API");
-      if (DeviceArray_API == NULL) {
-        retcode = -1;
-      }
-    } else {
-      retcode = -1;
-    }
-
-    Py_DECREF(key);
-    Py_DECREF(devicearray);
-    return retcode;
-}
-
 static PyMethodDef Dispatcher_methods[] = {
     { "_clear", (PyCFunction)Dispatcher_clear, METH_NOARGS, NULL },
     { "_insert", (PyCFunction)Dispatcher_Insert, METH_VARARGS | METH_KEYWORDS,
@@ -1076,12 +1044,6 @@ static PyMethodDef ext_methods[] = {
 
 
 MOD_INIT(_dispatcher) {
-    if (import_devicearray() < 0) {
-      PyErr_Print();
-      PyErr_SetString(PyExc_ImportError, NUMBA_DEVICEARRAY_IMPORT_NAME " failed to import");
-      return MOD_ERROR_VAL;
-    }
-
     PyObject *m;
     MOD_DEF(m, "_dispatcher", "No docs", ext_methods)
     if (m == NULL)

numba_cuda/numba/cuda/cext/_typeof.cpp

@@ -9,7 +9,6 @@
 
 #include "_typeof.h"
 #include "_hashtable.h"
-#include "_devicearray.h"
 #include "pyerrors.h"
 
 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
@@ -56,9 +55,6 @@ static PyObject *str_typeof_pyval = NULL;
 static PyObject *str_value = NULL;
 static PyObject *str_numba_type = NULL;
 
-/* CUDA device array API */
-void **DeviceArray_API;
-
 /*
  * Type fingerprint computation.
  */
@@ -857,109 +853,6 @@ int typecode_arrayscalar(PyObject *dispatcher, PyObject* aryscalar) {
     return BASIC_TYPECODES[typecode];
 }
 
-static
-int typecode_devicendarray(PyObject *dispatcher, PyObject *ary)
-{
-    int typecode;
-    int dtype;
-    int ndim;
-    int layout = 0;
-    PyObject *ndim_obj = nullptr;
-    PyObject *num_obj = nullptr;
-    PyObject *dtype_obj = nullptr;
-    int dtype_num = 0;
-
-    PyObject* flags = PyObject_GetAttrString(ary, "flags");
-    if (flags == NULL)
-    {
-        PyErr_Clear();
-        goto FALLBACK;
-    }
-
-    if (PyDict_GetItemString(flags, "C_CONTIGUOUS") == Py_True) {
-        layout = 1;
-    } else if (PyDict_GetItemString(flags, "F_CONTIGUOUS") == Py_True) {
-        layout = 2;
-    }
-
-    Py_DECREF(flags);
-
-    ndim_obj = PyObject_GetAttrString(ary, "ndim");
-    if (ndim_obj == NULL) {
-        /* If there's no ndim, try to proceed by clearing the error and using the
-         * fallback. */
-        PyErr_Clear();
-        goto FALLBACK;
-    }
-
-    ndim = PyLong_AsLong(ndim_obj);
-    Py_DECREF(ndim_obj);
-
-    if (PyErr_Occurred()) {
-        /* ndim wasn't an integer for some reason - unlikely to happen, but try
-         * the fallback. */
-        PyErr_Clear();
-        goto FALLBACK;
-    }
-
-    if (ndim <= 0 || ndim > N_NDIM)
-        goto FALLBACK;
-
-    dtype_obj = PyObject_GetAttrString(ary, "dtype");
-    if (dtype_obj == NULL) {
-        /* No dtype: try the fallback. */
-        PyErr_Clear();
-        goto FALLBACK;
-    }
-
-    num_obj = PyObject_GetAttrString(dtype_obj, "num");
-    Py_DECREF(dtype_obj);
-
-    if (num_obj == NULL) {
-        /* This strange dtype has no num - try the fallback. */
-        PyErr_Clear();
-        goto FALLBACK;
-    }
-
-    dtype_num = PyLong_AsLong(num_obj);
-    Py_DECREF(num_obj);
-
-    if (PyErr_Occurred()) {
-        /* num wasn't an integer for some reason - unlikely to happen, but try
-         * the fallback. */
-        PyErr_Clear();
-        goto FALLBACK;
-    }
-
-    dtype = dtype_num_to_typecode(dtype_num);
-    if (dtype == -1) {
-        /* Not a dtype we have in the global lookup table. */
-        goto FALLBACK;
-    }
-
-    /* Fast path, using direct table lookup */
-    assert(layout < N_LAYOUT);
-    assert(ndim <= N_NDIM);
-    assert(dtype < N_DTYPES);
-    typecode = cached_arycode[ndim - 1][layout][dtype];
-
-    if (typecode == -1) {
-        /* First use of this table entry, so it requires populating */
-        typecode = typecode_fallback_keep_ref(dispatcher, (PyObject*)ary);
-        cached_arycode[ndim - 1][layout][dtype] = typecode;
-    }
-
-    return typecode;
-
-FALLBACK:
-    /* Slower path, for non-trivial array types. At present this always uses
-       the fingerprinting to get the typecode. Future optimization might
-       implement a cache, but this would require some fast equivalent of
-       PyArray_DESCR for a device array. */
-
-    return typecode_using_fingerprint(dispatcher, (PyObject *) ary);
-}
-
 extern "C" int
 typeof_typecode(PyObject *dispatcher, PyObject *val)
 {
@@ -994,10 +887,6 @@ typeof_typecode(PyObject *dispatcher, PyObject *val)
     else if (tyobj == &PyArray_Type) {
         return typecode_ndarray(dispatcher, (PyArrayObject*)val);
     }
-    /* Subtype of CUDA device array */
-    else if (PyType_IsSubtype(tyobj, &DeviceArrayType)) {
-        return typecode_devicendarray(dispatcher, val);
-    }
     /* Subtypes of Array handling */
     else if (PyType_IsSubtype(tyobj, &PyArray_Type)) {
         /* By default, Numba will treat all numpy.ndarray subtypes as if they

numba_cuda/numba/cuda/codegen.py

@@ -12,6 +12,7 @@ from numba.cuda.cudadrv.linkable_code import LinkableCode
 from numba.cuda.memory_management.nrt import NRT_LIBRARY
 
 import os
+import pickle
 import subprocess
 import tempfile
 
@@ -189,6 +190,11 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
 
         self.use_cooperative = False
 
+        # Objects that need to be kept alive for the lifetime of the
+        # kernels or device functions generated by this code library,
+        # e.g., device arrays captured from global scope.
+        self.referenced_objects = {}
+
     @property
     def llvm_strs(self):
         if self._llvm_strs is None:
@@ -206,6 +212,9 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         return device.compute_capability
 
     def get_asm_str(self, cc=None):
+        return "\n".join(self.get_asm_strs(cc=cc))
+
+    def get_asm_strs(self, cc=None):
         cc = self._ensure_cc(cc)
 
         ptxes = self._ptx_cache.get(cc, None)
@@ -218,21 +227,25 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
 
         irs = self.llvm_strs
 
-        ptx = nvvm.compile_ir(irs, **options)
+        if "g" in options:
+            ptxes = [nvvm.compile_ir(ir, **options) for ir in irs]
+        else:
+            ptxes = [nvvm.compile_ir(irs, **options)]
 
         # Sometimes the result from NVVM contains trailing whitespace and
         # nulls, which we strip so that the assembly dump looks a little
         # tidier.
-        ptx = ptx.decode().strip("\x00").strip()
+        ptxes = [ptx.decode().strip("\x00").strip() for ptx in ptxes]
 
         if config.DUMP_ASSEMBLY:
             print(("ASSEMBLY %s" % self._name).center(80, "-"))
-            print(ptx)
+            for ptx in ptxes:
+                print(ptx)
             print("=" * 80)
 
-        self._ptx_cache[cc] = ptx
+        self._ptx_cache[cc] = ptxes
 
-        return ptx
+        return ptxes
 
     def get_lto_ptx(self, cc=None):
         """
@@ -247,7 +260,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
 
         cc = self._ensure_cc(cc)
 
-        linker = driver._Linker.new(
+        linker = driver._Linker(
             max_registers=self._max_registers,
             cc=cc,
             additional_flags=["-ptx"],
@@ -284,8 +297,9 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             ltoir = self.get_ltoir(cc=cc)
             linker.add_ltoir(ltoir)
         else:
-            ptx = self.get_asm_str(cc=cc)
-            linker.add_ptx(ptx.encode())
+            ptxes = self.get_asm_strs(cc=cc)
+            for ptx in ptxes:
+                linker.add_ptx(ptx.encode())
 
         for path in self._linking_files:
             linker.add_file_guess_ext(path, ignore_nonlto)
@@ -308,7 +322,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             print(ptx)
             print("=" * 80)
 
-        linker = driver._Linker.new(
+        linker = driver._Linker(
             max_registers=self._max_registers, cc=cc, lto=self._lto
         )
         self._link_all(linker, cc, ignore_nonlto=False)
@@ -377,6 +391,9 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         self._setup_functions.extend(library._setup_functions)
         self._teardown_functions.extend(library._teardown_functions)
         self.use_cooperative |= library.use_cooperative
+        self.referenced_objects.update(
+            getattr(library, "referenced_objects", {})
+        )
 
     def add_linking_file(self, path_or_obj):
         if isinstance(path_or_obj, LinkableCode):
@@ -432,7 +449,10 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             for mod in library.modules:
                 for fn in mod.functions:
                     if not fn.is_declaration:
-                        fn.linkage = "linkonce_odr"
+                        if "g" in self._nvvm_options:
+                            fn.linkage = "weak_odr"
+                        else:
+                            fn.linkage = "linkonce_odr"
 
         self._finalized = True
 
@@ -442,6 +462,18 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         but loaded functions are discarded. They are recreated when needed
         after deserialization.
         """
+        # Check for captured device arrays that cannot be safely cached.
+        if self.referenced_objects:
+            if any(
+                getattr(obj, "__cuda_array_interface__", None) is not None
+                for obj in self.referenced_objects.values()
+            ):
+                raise pickle.PicklingError(
+                    "Cannot serialize kernels or device functions referencing "
+                    "global device arrays. Pass the array(s) as arguments "
+                    "to the kernel instead."
+                )
+
         nrt = False
         if self._linking_files:
             if (

numba_cuda/numba/cuda/compiler.py

@@ -1023,10 +1023,9 @@ def compile_all(
     )
 
     if lto:
-        code = lib.get_ltoir(cc=cc)
+        codes = [lib.get_ltoir(cc=cc)]
     else:
-        code = lib.get_asm_str(cc=cc)
-    codes = [code]
+        codes = lib.get_asm_strs(cc=cc)
 
     # linking_files
     is_ltoir = output == "ltoir"
@@ -1241,7 +1240,14 @@ def compile(
     if lto:
         code = lib.get_ltoir(cc=cc)
     else:
-        code = lib.get_asm_str(cc=cc)
+        codes = lib.get_asm_strs(cc=cc)
+        if len(codes) == 1:
+            code = codes[0]
+        else:
+            raise RuntimeError(
+                "Compiling this function results in multiple "
+                "PTX files. Use compile_all() instead"
+            )
     return code, resty
 
 

numba_cuda/numba/cuda/core/analysis.py

@@ -38,13 +38,16 @@ def compute_use_defs(blocks):
                 func = ir_extension_usedefs[type(stmt)]
                 func(stmt, use_set, def_set)
                 continue
-            if isinstance(stmt, ir.Assign):
-                if isinstance(stmt.value, ir.Inst):
+            if isinstance(stmt, ir.assign_types):
+                if isinstance(stmt.value, ir.inst_types):
                     rhs_set = set(var.name for var in stmt.value.list_vars())
-                elif isinstance(stmt.value, ir.Var):
+                elif isinstance(stmt.value, ir.var_types):
                     rhs_set = set([stmt.value.name])
-                elif isinstance(
-                    stmt.value, (ir.Arg, ir.Const, ir.Global, ir.FreeVar)
+                elif (
+                    isinstance(stmt.value, ir.arg_types)
+                    or isinstance(stmt.value, ir.const_types)
+                    or isinstance(stmt.value, ir.global_types)
+                    or isinstance(stmt.value, ir.freevar_types)
                 ):
                     rhs_set = ()
                 else:
@@ -326,7 +329,7 @@ def rewrite_semantic_constants(func_ir, called_args):
         if getattr(val, "op", None) == "getattr":
             if val.attr == "ndim":
                 arg_def = guard(get_definition, func_ir, val.value)
-                if isinstance(arg_def, ir.Arg):
+                if isinstance(arg_def, ir.arg_types):
                     argty = called_args[arg_def.index]
                     if isinstance(argty, types.Array):
                         rewrite_statement(func_ir, stmt, argty.ndim)
@@ -337,17 +340,17 @@ def rewrite_semantic_constants(func_ir, called_args):
             func = guard(get_definition, func_ir, val.func)
             if (
                 func is not None
-                and isinstance(func, ir.Global)
+                and isinstance(func, ir.global_types)
                 and getattr(func, "value", None) is len
             ):
                 (arg,) = val.args
                 arg_def = guard(get_definition, func_ir, arg)
-                if isinstance(arg_def, ir.Arg):
+                if isinstance(arg_def, ir.arg_types):
                     argty = called_args[arg_def.index]
                     if isinstance(argty, types.BaseTuple):
                         rewrite_statement(func_ir, stmt, argty.count)
                 elif (
-                    isinstance(arg_def, ir.Expr)
+                    isinstance(arg_def, ir.expr_types)
                     and arg_def.op == "typed_getitem"
                 ):
                     argty = arg_def.dtype
@@ -358,9 +361,9 @@ def rewrite_semantic_constants(func_ir, called_args):
 
     for blk in func_ir.blocks.values():
         for stmt in blk.body:
-            if isinstance(stmt, ir.Assign):
+            if isinstance(stmt, ir.assign_types):
                 val = stmt.value
-                if isinstance(val, ir.Expr):
+                if isinstance(val, ir.expr_types):
                     rewrite_array_ndim(val, func_ir, called_args)
                     rewrite_tuple_len(val, func_ir, called_args)
 
@@ -391,7 +394,7 @@ def find_literally_calls(func_ir, argtypes):
     for blk in func_ir.blocks.values():
         for assign in blk.find_exprs(op="call"):
             var = ir_utils.guard(ir_utils.get_definition, func_ir, assign.func)
-            if isinstance(var, (ir.Global, ir.FreeVar)):
+            if isinstance(var, ir.global_types + ir.freevar_types):
                 fnobj = var.value
             else:
                 fnobj = ir_utils.guard(
@@ -401,7 +404,7 @@ def find_literally_calls(func_ir, argtypes):
                 # Found
                 [arg] = assign.args
                 defarg = func_ir.get_definition(arg)
-                if isinstance(defarg, ir.Arg):
+                if isinstance(defarg, ir.arg_types):
                     argindex = defarg.index
                     marked_args.add(argindex)
                     first_loc.setdefault(argindex, assign.loc)
@@ -473,14 +476,14 @@ def dead_branch_prune(func_ir, called_args):
         branches = []
         for blk in func_ir.blocks.values():
             branch_or_jump = blk.body[-1]
-            if isinstance(branch_or_jump, ir.Branch):
+            if isinstance(branch_or_jump, ir.branch_types):
                 branch = branch_or_jump
                 pred = guard(get_definition, func_ir, branch.cond.name)
                 if pred is not None and getattr(pred, "op", None) == "call":
                     function = guard(get_definition, func_ir, pred.func)
                     if (
                         function is not None
-                        and isinstance(function, ir.Global)
+                        and isinstance(function, ir.global_types)
                         and function.value is bool
                     ):
                         condition = guard(get_definition, func_ir, pred.args[0])
@@ -539,7 +542,9 @@ def dead_branch_prune(func_ir, called_args):
         try:
             # Just to prevent accidents, whilst already guarded, ensure this
             # is an ir.Const
-            if not isinstance(pred, (ir.Const, ir.FreeVar, ir.Global)):
+            if not isinstance(
+                pred, ir.const_types + ir.freevar_types + ir.global_types
+            ):
                 raise TypeError("Expected constant Numba IR node")
             take_truebr = bool(pred.value)
         except TypeError:
@@ -584,8 +589,11 @@ def dead_branch_prune(func_ir, called_args):
     phi2asgn = dict()
     for lbl, blk in func_ir.blocks.items():
         for stmt in blk.body:
-            if isinstance(stmt, ir.Assign):
-                if isinstance(stmt.value, ir.Expr) and stmt.value.op == "phi":
+            if isinstance(stmt, ir.assign_types):
+                if (
+                    isinstance(stmt.value, ir.expr_types)
+                    and stmt.value.op == "phi"
+                ):
                     phi2lbl[stmt.value] = lbl
                     phi2asgn[stmt.value] = stmt
 
@@ -599,12 +607,12 @@ def dead_branch_prune(func_ir, called_args):
 
     for branch, condition, blk in branch_info:
         const_conds = []
-        if isinstance(condition, ir.Expr) and condition.op == "binop":
+        if isinstance(condition, ir.expr_types) and condition.op == "binop":
             prune = prune_by_value
             for arg in [condition.lhs, condition.rhs]:
                 resolved_const = Unknown()
                 arg_def = guard(get_definition, func_ir, arg)
-                if isinstance(arg_def, ir.Arg):
+                if isinstance(arg_def, ir.arg_types):
                     # it's an e.g. literal argument to the function
                     resolved_const = resolve_input_arg_const(arg_def.index)
                     prune = prune_by_type
@@ -668,7 +676,7 @@ def dead_branch_prune(func_ir, called_args):
     for _, cond, blk in branch_info:
         if cond in deadcond:
             for x in blk.body:
-                if isinstance(x, ir.Assign) and x.value is cond:
+                if isinstance(x, ir.assign_types) and x.value is cond:
                     # rewrite the condition as a true/false bit
                     nullified_info = nullified_conditions[deadcond.index(cond)]
                     # only do a rewrite of conditions, predicates need to retain

numba_cuda/numba/cuda/core/annotations/type_annotations.py

@@ -94,16 +94,16 @@ class TypeAnnotation(object):
             for inst in blk.body:
                 lineno = inst.loc.line
 
-                if isinstance(inst, ir.Assign):
+                if isinstance(inst, ir.assign_types):
                     if found_lifted_loop:
                         atype = "XXX Lifted Loop XXX"
                         found_lifted_loop = False
                     elif (
-                        isinstance(inst.value, ir.Expr)
+                        isinstance(inst.value, ir.expr_types)
                         and inst.value.op == "call"
                     ):
                         atype = self.calltypes[inst.value]
-                    elif isinstance(inst.value, ir.Const) and isinstance(
+                    elif isinstance(inst.value, ir.const_types) and isinstance(
                         inst.value.value, LiftedLoop
                     ):
                         atype = "XXX Lifted Loop XXX"
@@ -113,7 +113,7 @@ class TypeAnnotation(object):
                         atype = self.typemap.get(inst.target.name, "<missing>")
 
                     aline = "%s = %s  :: %s" % (inst.target, inst.value, atype)
-                elif isinstance(inst, ir.SetItem):
+                elif isinstance(inst, ir.setitem_types):
                     atype = self.calltypes[inst]
                     aline = "%s  :: %s" % (inst, atype)
                 else:

numba_cuda/numba/cuda/core/base.py

@@ -933,7 +933,12 @@ class BaseContext(object):
         If *caching* evaluates True, the function keeps the compiled function
         for reuse in *.cached_internal_func*.
         """
-        cache_key = (impl.__code__, sig, type(self.error_model))
+        cache_key = (
+            impl.__code__,
+            sig,
+            type(self.error_model),
+            self.enable_nrt,
+        )
         if not caching:
             cached = None
         else:

numba_cuda/numba/cuda/core/consts.py

@@ -68,7 +68,7 @@ class ConstantInference(object):
         try:
             const = defn.infer_constant()
         except ConstantInferenceError:
-            if isinstance(defn, ir.Expr):
+            if isinstance(defn, ir.expr_types):
                 return self._infer_expr(defn)
             self._fail(defn)
         return const