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

numba_cuda/numba/cuda/core/consts.py

@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+from types import ModuleType
+
+import weakref
+
+from numba.core.errors import ConstantInferenceError, NumbaError
+from numba.core import ir
+
+
+class ConstantInference(object):
+    """
+    A constant inference engine for a given interpreter.
+    Inference inspects the IR to try and compute a compile-time constant for
+    a variable.
+
+    This shouldn't be used directly, instead call Interpreter.infer_constant().
+    """
+
+    def __init__(self, func_ir):
+        # Avoid cyclic references as some user-visible objects may be
+        # held alive in the cache
+        self._func_ir = weakref.proxy(func_ir)
+        self._cache = {}
+
+    def infer_constant(self, name, loc=None):
+        """
+        Infer a constant value for the given variable *name*.
+        If no value can be inferred, numba.errors.ConstantInferenceError
+        is raised.
+        """
+        if name not in self._cache:
+            try:
+                self._cache[name] = (True, self._do_infer(name))
+            except ConstantInferenceError as exc:
+                # Store the exception args only, to avoid keeping
+                # a whole traceback alive.
+                self._cache[name] = (False, (exc.__class__, exc.args))
+        success, val = self._cache[name]
+        if success:
+            return val
+        else:
+            exc, args = val
+            if issubclass(exc, NumbaError):
+                raise exc(*args, loc=loc)
+            else:
+                raise exc(*args)
+
+    def _fail(self, val):
+        # The location here is set to None because `val` is the ir.Var name
+        # and not the actual offending use of the var. When this is raised it is
+        # caught in the flow control of `infer_constant` and the class and args
+        # (the message) are captured and then raised again but with the location
+        # set to the expression that caused the constant inference error.
+        raise ConstantInferenceError(
+            "Constant inference not possible for: %s" % (val,), loc=None
+        )
+
+    def _do_infer(self, name):
+        if not isinstance(name, str):
+            raise TypeError("infer_constant() called with non-str %r" % (name,))
+        try:
+            defn = self._func_ir.get_definition(name)
+        except KeyError:
+            raise ConstantInferenceError(
+                "no single definition for %r" % (name,)
+            )
+        try:
+            const = defn.infer_constant()
+        except ConstantInferenceError:
+            if isinstance(defn, ir.Expr):
+                return self._infer_expr(defn)
+            self._fail(defn)
+        return const
+
+    def _infer_expr(self, expr):
+        # Infer an expression: handle supported cases
+        if expr.op == "call":
+            func = self.infer_constant(expr.func.name, loc=expr.loc)
+            return self._infer_call(func, expr)
+        elif expr.op == "getattr":
+            value = self.infer_constant(expr.value.name, loc=expr.loc)
+            return self._infer_getattr(value, expr)
+        elif expr.op == "build_list":
+            return [
+                self.infer_constant(i.name, loc=expr.loc) for i in expr.items
+            ]
+        elif expr.op == "build_tuple":
+            return tuple(
+                self.infer_constant(i.name, loc=expr.loc) for i in expr.items
+            )
+        self._fail(expr)
+
+    def _infer_call(self, func, expr):
+        if expr.kws or expr.vararg:
+            self._fail(expr)
+        # Check supported callables
+        _slice = func in (slice,)
+        _exc = isinstance(func, type) and issubclass(func, BaseException)
+        if _slice or _exc:
+            args = [
+                self.infer_constant(a.name, loc=expr.loc) for a in expr.args
+            ]
+            if _slice:
+                return func(*args)
+            elif _exc:
+                # If the exception class is user defined it may implement a ctor
+                # that does not pass the args to the super. Therefore return the
+                # raw class and the args so this can be instantiated at the call
+                # site in the way the user source expects it to be.
+                return func, args
+            else:
+                assert 0, "Unreachable"
+
+        self._fail(expr)
+
+    def _infer_getattr(self, value, expr):
+        if isinstance(value, (ModuleType, type)):
+            # Allow looking up a constant on a class or module
+            try:
+                return getattr(value, expr.attr)
+            except AttributeError:
+                pass
+        self._fail(expr)

numba_cuda/numba/cuda/core/cpu.py

@@ -0,0 +1,370 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import platform
+
+import llvmlite.binding as ll
+from llvmlite import ir
+
+from numba import _dynfunc
+from numba.core.callwrapper import PyCallWrapper
+from numba.core.base import BaseContext
+from numba.core import (
+    utils,
+    types,
+    config,
+    cgutils,
+    callconv,
+    codegen,
+    externals,
+    fastmathpass,
+    intrinsics,
+)
+from numba.core.options import TargetOptions, include_default_options
+from numba.core.runtime import rtsys
+from numba.core.compiler_lock import global_compiler_lock
+import numba.core.entrypoints
+
+# Re-export these options, they are used from the cpu module throughout the code
+# base.
+from numba.cuda.core.options import (
+    ParallelOptions,  # noqa F401
+    FastMathOptions,  # noqa F401
+    InlineOptions,  # noqa F401
+)  # noqa F401
+from numba.np import ufunc_db
+
+# Keep those structures in sync with _dynfunc.c.
+
+
+class ClosureBody(cgutils.Structure):
+    _fields = [("env", types.pyobject)]
+
+
+class EnvBody(cgutils.Structure):
+    _fields = [
+        ("globals", types.pyobject),
+        ("consts", types.pyobject),
+    ]
+
+
+class CPUContext(BaseContext):
+    """
+    Changes BaseContext calling convention
+    """
+
+    allow_dynamic_globals = True
+
+    def __init__(self, typingctx, target="cpu"):
+        super().__init__(typingctx, target)
+
+    # Overrides
+    def create_module(self, name):
+        return self._internal_codegen._create_empty_module(name)
+
+    @global_compiler_lock
+    def init(self):
+        self.is32bit = utils.MACHINE_BITS == 32
+        self._internal_codegen = codegen.JITCPUCodegen("numba.exec")
+
+        # Add ARM ABI functions from libgcc_s
+        if platform.machine() == "armv7l":
+            ll.load_library_permanently("libgcc_s.so.1")
+
+        # Map external C functions.
+        externals.c_math_functions.install(self)
+
+    def load_additional_registries(self):
+        # Only initialize the NRT once something is about to be compiled. The
+        # "initialized" state doesn't need to be threadsafe, there's a lock
+        # around the internal compilation and the rtsys.initialize call can be
+        # made multiple times, worse case init just gets called a bit more often
+        # than optimal.
+        rtsys.initialize(self)
+
+        # Add implementations that work via import
+        from numba.cpython import (
+            enumimpl,  # noqa F401
+            iterators,  # noqa F401
+            rangeobj,  # noqa F401
+            tupleobj,  # noqa F401
+        )  # noqa F401
+        from numba.core import optional, inline_closurecall  # noqa F401
+        from numba.misc import gdb_hook, literal  # noqa F401
+        from numba.np import linalg, arraymath, arrayobj  # noqa F401
+        from numba.np.random import generator_core, generator_methods  # noqa F401
+        from numba.np.polynomial import polynomial_core, polynomial_functions  # noqa F401
+        from numba.typed import typeddict, dictimpl  # noqa F401
+        from numba.typed import typedlist, listobject  # noqa F401
+        from numba.experimental import jitclass, function_type  # noqa F401
+        from numba.np import npdatetime  # noqa F401
+
+        # Add target specific implementations
+        from numba.np import npyimpl
+        from numba.cpython import cmathimpl, mathimpl, printimpl, randomimpl
+        from numba.misc import cffiimpl
+        from numba.experimental.jitclass.base import (
+            ClassBuilder as jitclassimpl,
+        )
+
+        self.install_registry(cmathimpl.registry)
+        self.install_registry(cffiimpl.registry)
+        self.install_registry(mathimpl.registry)
+        self.install_registry(npyimpl.registry)
+        self.install_registry(printimpl.registry)
+        self.install_registry(randomimpl.registry)
+        self.install_registry(jitclassimpl.class_impl_registry)
+
+        # load 3rd party extensions
+        numba.core.entrypoints.init_all()
+
+        # fix for #8940
+        from numba.np.unsafe import ndarray  # noqa F401
+
+    @property
+    def target_data(self):
+        return self._internal_codegen.target_data
+
+    def with_aot_codegen(self, name, **aot_options):
+        aot_codegen = codegen.AOTCPUCodegen(name, **aot_options)
+        return self.subtarget(_internal_codegen=aot_codegen, aot_mode=True)
+
+    def codegen(self):
+        return self._internal_codegen
+
+    @property
+    def call_conv(self):
+        return callconv.CPUCallConv(self)
+
+    def get_env_body(self, builder, envptr):
+        """
+        From the given *envptr* (a pointer to a _dynfunc.Environment object),
+        get a EnvBody allowing structured access to environment fields.
+        """
+        body_ptr = cgutils.pointer_add(
+            builder, envptr, _dynfunc._impl_info["offsetof_env_body"]
+        )
+        return EnvBody(self, builder, ref=body_ptr, cast_ref=True)
+
+    def get_env_manager(self, builder, return_pyobject=False):
+        envgv = self.declare_env_global(
+            builder.module, self.get_env_name(self.fndesc)
+        )
+        envarg = builder.load(envgv)
+        pyapi = self.get_python_api(builder)
+        pyapi.emit_environment_sentry(
+            envarg,
+            return_pyobject=return_pyobject,
+            debug_msg=self.fndesc.env_name,
+        )
+        env_body = self.get_env_body(builder, envarg)
+        return pyapi.get_env_manager(self.environment, env_body, envarg)
+
+    def get_generator_state(self, builder, genptr, return_type):
+        """
+        From the given *genptr* (a pointer to a _dynfunc.Generator object),
+        get a pointer to its state area.
+        """
+        return cgutils.pointer_add(
+            builder,
+            genptr,
+            _dynfunc._impl_info["offsetof_generator_state"],
+            return_type=return_type,
+        )
+
+    def build_list(self, builder, list_type, items):
+        """
+        Build a list from the Numba *list_type* and its initial *items*.
+        """
+        from numba.cpython import listobj
+
+        return listobj.build_list(self, builder, list_type, items)
+
+    def build_set(self, builder, set_type, items):
+        """
+        Build a set from the Numba *set_type* and its initial *items*.
+        """
+        from numba.cpython import setobj
+
+        return setobj.build_set(self, builder, set_type, items)
+
+    def build_map(self, builder, dict_type, item_types, items):
+        from numba.typed import dictobject
+
+        return dictobject.build_map(self, builder, dict_type, item_types, items)
+
+    def post_lowering(self, mod, library):
+        if self.fastmath:
+            fastmathpass.rewrite_module(mod, self.fastmath)
+
+        if self.is32bit:
+            # 32-bit machine needs to replace all 64-bit div/rem to avoid
+            # calls to compiler-rt
+            intrinsics.fix_divmod(mod)
+
+        library.add_linking_library(rtsys.library)
+
+    def create_cpython_wrapper(
+        self, library, fndesc, env, call_helper, release_gil=False
+    ):
+        wrapper_module = self.create_module("wrapper")
+        fnty = self.call_conv.get_function_type(fndesc.restype, fndesc.argtypes)
+        wrapper_callee = ir.Function(
+            wrapper_module, fnty, fndesc.llvm_func_name
+        )
+        builder = PyCallWrapper(
+            self,
+            wrapper_module,
+            wrapper_callee,
+            fndesc,
+            env,
+            call_helper=call_helper,
+            release_gil=release_gil,
+        )
+        builder.build()
+        library.add_ir_module(wrapper_module)
+
+    def create_cfunc_wrapper(self, library, fndesc, env, call_helper):
+        wrapper_module = self.create_module("cfunc_wrapper")
+        fnty = self.call_conv.get_function_type(fndesc.restype, fndesc.argtypes)
+        wrapper_callee = ir.Function(
+            wrapper_module, fnty, fndesc.llvm_func_name
+        )
+
+        ll_argtypes = [self.get_value_type(ty) for ty in fndesc.argtypes]
+        ll_return_type = self.get_value_type(fndesc.restype)
+        wrapty = ir.FunctionType(ll_return_type, ll_argtypes)
+        wrapfn = ir.Function(
+            wrapper_module, wrapty, fndesc.llvm_cfunc_wrapper_name
+        )
+        builder = ir.IRBuilder(wrapfn.append_basic_block("entry"))
+
+        status, out = self.call_conv.call_function(
+            builder,
+            wrapper_callee,
+            fndesc.restype,
+            fndesc.argtypes,
+            wrapfn.args,
+            attrs=("noinline",),
+        )
+
+        with builder.if_then(status.is_error, likely=False):
+            # If (and only if) an error occurred, acquire the GIL
+            # and use the interpreter to write out the exception.
+            pyapi = self.get_python_api(builder)
+            gil_state = pyapi.gil_ensure()
+            self.call_conv.raise_error(builder, pyapi, status)
+            cstr = self.insert_const_string(builder.module, repr(self))
+            strobj = pyapi.string_from_string(cstr)
+            pyapi.err_write_unraisable(strobj)
+            pyapi.decref(strobj)
+            pyapi.gil_release(gil_state)
+
+        builder.ret(out)
+        library.add_ir_module(wrapper_module)
+
+    def get_executable(self, library, fndesc, env):
+        """
+        Returns
+        -------
+        (cfunc, fnptr)
+
+        - cfunc
+            callable function (Can be None)
+        - fnptr
+            callable function address
+        - env
+            an execution environment (from _dynfunc)
+        """
+        # Code generation
+        fnptr = library.get_pointer_to_function(
+            fndesc.llvm_cpython_wrapper_name
+        )
+
+        # Note: we avoid reusing the original docstring to avoid encoding
+        # issues on Python 2, see issue #1908
+        doc = "compiled wrapper for %r" % (fndesc.qualname,)
+        cfunc = _dynfunc.make_function(
+            fndesc.lookup_module(),
+            fndesc.qualname.split(".")[-1],
+            doc,
+            fnptr,
+            env,
+            # objects to keepalive with the function
+            (library,),
+        )
+        library.codegen.set_env(self.get_env_name(fndesc), env)
+        return cfunc
+
+    def calc_array_sizeof(self, ndim):
+        """
+        Calculate the size of an array struct on the CPU target
+        """
+        aryty = types.Array(types.int32, ndim, "A")
+        return self.get_abi_sizeof(self.get_value_type(aryty))
+
+    # Overrides
+    def get_ufunc_info(self, ufunc_key):
+        return ufunc_db.get_ufunc_info(ufunc_key)
+
+
+# ----------------------------------------------------------------------------
+# TargetOptions
+
+_options_mixin = include_default_options(
+    "nopython",
+    "forceobj",
+    "looplift",
+    "_nrt",
+    "debug",
+    "boundscheck",
+    "nogil",
+    "no_rewrites",
+    "no_cpython_wrapper",
+    "no_cfunc_wrapper",
+    "parallel",
+    "fastmath",
+    "error_model",
+    "inline",
+    "forceinline",
+    "_dbg_extend_lifetimes",
+    "_dbg_optnone",
+)
+
+
+class CPUTargetOptions(_options_mixin, TargetOptions):
+    def finalize(self, flags, options):
+        if not flags.is_set("enable_pyobject"):
+            flags.enable_pyobject = True
+
+        if not flags.is_set("enable_looplift"):
+            flags.enable_looplift = True
+
+        flags.inherit_if_not_set("nrt", default=True)
+
+        if not flags.is_set("debuginfo"):
+            flags.debuginfo = config.DEBUGINFO_DEFAULT
+
+        if not flags.is_set("dbg_extend_lifetimes"):
+            if flags.debuginfo:
+                # auto turn on extend-lifetimes if debuginfo is on and
+                # dbg_extend_lifetimes is not set
+                flags.dbg_extend_lifetimes = True
+            else:
+                # set flag using env-var config
+                flags.dbg_extend_lifetimes = config.EXTEND_VARIABLE_LIFETIMES
+
+        if not flags.is_set("boundscheck"):
+            flags.boundscheck = flags.debuginfo
+
+        flags.enable_pyobject_looplift = True
+
+        flags.inherit_if_not_set("fastmath")
+
+        flags.inherit_if_not_set("error_model", default="python")
+
+        flags.inherit_if_not_set("forceinline")
+
+        if flags.forceinline:
+            # forceinline turns off optnone, just like clang.
+            flags.dbg_optnone = False

numba_cuda/numba/cuda/core/environment.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import weakref
+import importlib
+
+from numba import _dynfunc
+
+
+class Environment(_dynfunc.Environment):
+    """Stores globals and constant pyobjects for runtime.
+
+    It is often needed to convert b/w nopython objects and pyobjects.
+    """
+
+    __slots__ = ("env_name", "__weakref__")
+    # A weak-value dictionary to store live environment with env_name as the
+    # key.
+    _memo = weakref.WeakValueDictionary()
+
+    @classmethod
+    def from_fndesc(cls, fndesc):
+        try:
+            # Avoid creating new Env
+            return cls._memo[fndesc.env_name]
+        except KeyError:
+            inst = cls(fndesc.lookup_globals())
+            inst.env_name = fndesc.env_name
+            cls._memo[fndesc.env_name] = inst
+            return inst
+
+    def can_cache(self):
+        is_dyn = "__name__" not in self.globals
+        return not is_dyn
+
+    def __reduce__(self):
+        return _rebuild_env, (
+            self.globals.get("__name__"),
+            self.consts,
+            self.env_name,
+        )
+
+    def __del__(self):
+        return
+
+    def __repr__(self):
+        return f"<Environment {self.env_name!r} >"
+
+
+def _rebuild_env(modname, consts, env_name):
+    env = lookup_environment(env_name)
+    if env is not None:
+        return env
+
+    mod = importlib.import_module(modname)
+    env = Environment(mod.__dict__)
+    env.consts[:] = consts
+    env.env_name = env_name
+    # Cache loaded object
+    Environment._memo[env_name] = env
+    return env
+
+
+def lookup_environment(env_name):
+    """Returns the Environment object for the given name;
+    or None if not found
+    """
+    return Environment._memo.get(env_name)