numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/core/funcdesc.py

@@ -0,0 +1,330 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Function descriptors.
+"""
+
+from collections import defaultdict
+import importlib
+
+from numba.cuda import types
+from numba.cuda import itanium_mangler
+from numba.cuda.utils import _dynamic_modname, _dynamic_module
+
+
+def default_mangler(name, argtypes, *, abi_tags=(), uid=None):
+    return itanium_mangler.mangle(name, argtypes, abi_tags=abi_tags, uid=uid)
+
+
+def qualifying_prefix(modname, qualname):
+    """
+    Returns a new string that is used for the first half of the mangled name.
+    """
+    # XXX choose a different convention for object mode
+    return "{}.{}".format(modname, qualname) if modname else qualname
+
+
+class FunctionDescriptor(object):
+    """
+    Base class for function descriptors: an object used to carry
+    useful metadata about a natively callable function.
+
+    Note that while `FunctionIdentity` denotes a Python function
+    which is being concretely compiled by Numba, `FunctionDescriptor`
+    may be more "abstract".
+    """
+
+    __slots__ = (
+        "native",
+        "modname",
+        "qualname",
+        "doc",
+        "typemap",
+        "calltypes",
+        "args",
+        "kws",
+        "restype",
+        "argtypes",
+        "mangled_name",
+        "unique_name",
+        "env_name",
+        "global_dict",
+        "inline",
+        "noalias",
+        "abi_tags",
+        "uid",
+    )
+
+    def __init__(
+        self,
+        native,
+        modname,
+        qualname,
+        unique_name,
+        doc,
+        typemap,
+        restype,
+        calltypes,
+        args,
+        kws,
+        mangler=None,
+        argtypes=None,
+        inline=False,
+        noalias=False,
+        env_name=None,
+        global_dict=None,
+        abi_tags=(),
+        uid=None,
+    ):
+        self.native = native
+        self.modname = modname
+        self.global_dict = global_dict
+        self.qualname = qualname
+        self.unique_name = unique_name
+        self.doc = doc
+        # XXX typemap and calltypes should be on the compile result,
+        # not the FunctionDescriptor
+        self.typemap = typemap
+        self.calltypes = calltypes
+        self.args = args
+        self.kws = kws
+        self.restype = restype
+        # Argument types
+        if argtypes is not None:
+            assert isinstance(argtypes, tuple), argtypes
+            self.argtypes = argtypes
+        else:
+            # Get argument types from the type inference result
+            # (note the "arg.FOO" convention as used in typeinfer
+            self.argtypes = tuple(self.typemap["arg." + a] for a in args)
+        mangler = default_mangler if mangler is None else mangler
+        # The mangled name *must* be unique, else the wrong function can
+        # be chosen at link time.
+        qualprefix = qualifying_prefix(self.modname, self.qualname)
+        self.uid = uid
+        self.mangled_name = mangler(
+            qualprefix,
+            self.argtypes,
+            abi_tags=abi_tags,
+            uid=uid,
+        )
+        if env_name is None:
+            env_name = mangler(
+                ".NumbaEnv.{}".format(qualprefix),
+                self.argtypes,
+                abi_tags=abi_tags,
+                uid=uid,
+            )
+        self.env_name = env_name
+        self.inline = inline
+        self.noalias = noalias
+        self.abi_tags = abi_tags
+
+    def lookup_globals(self):
+        """
+        Return the global dictionary of the function.
+        It may not match the Module's globals if the function is created
+        dynamically (i.e. exec)
+        """
+        return self.global_dict or self.lookup_module().__dict__
+
+    def lookup_module(self):
+        """
+        Return the module in which this function is supposed to exist.
+        This may be a dummy module if the function was dynamically
+        generated or the module can't be found.
+        """
+        if self.modname == _dynamic_modname:
+            return _dynamic_module
+        else:
+            try:
+                # ensure module exist
+                return importlib.import_module(self.modname)
+            except ImportError:
+                return _dynamic_module
+
+    def lookup_function(self):
+        """
+        Return the original function object described by this object.
+        """
+        return getattr(self.lookup_module(), self.qualname)
+
+    @property
+    def llvm_func_name(self):
+        """
+        The LLVM-registered name for the raw function.
+        """
+        return self.mangled_name
+
+    @property
+    def llvm_cpython_wrapper_name(self):
+        """
+        The LLVM-registered name for a CPython-compatible wrapper of the
+        raw function (i.e. a PyCFunctionWithKeywords).
+        """
+        return itanium_mangler.prepend_namespace(
+            self.mangled_name, ns="cpython"
+        )
+
+    @property
+    def llvm_cfunc_wrapper_name(self):
+        """
+        The LLVM-registered name for a C-compatible wrapper of the
+        raw function.
+        """
+        return "cfunc." + self.mangled_name
+
+    def __repr__(self):
+        return "<function descriptor %r>" % (self.unique_name)
+
+    @classmethod
+    def _get_function_info(cls, func_ir):
+        """
+        Returns
+        -------
+        qualname, unique_name, modname, doc, args, kws, globals
+
+        ``unique_name`` must be a unique name.
+        """
+        func = func_ir.func_id.func
+        qualname = func_ir.func_id.func_qualname
+        # XXX to func_id
+        modname = func.__module__
+        doc = func.__doc__ or ""
+        args = tuple(func_ir.arg_names)
+        kws = ()  # TODO
+        global_dict = None
+
+        if modname is None:
+            # Dynamically generated function.
+            modname = _dynamic_modname
+            # Retain a reference to the dictionary of the function.
+            # This disables caching, serialization and pickling.
+            global_dict = func_ir.func_id.func.__globals__
+
+        unique_name = func_ir.func_id.unique_name
+
+        return qualname, unique_name, modname, doc, args, kws, global_dict
+
+    @classmethod
+    def _from_python_function(
+        cls,
+        func_ir,
+        typemap,
+        restype,
+        calltypes,
+        native,
+        mangler=None,
+        inline=False,
+        noalias=False,
+        abi_tags=(),
+    ):
+        (
+            qualname,
+            unique_name,
+            modname,
+            doc,
+            args,
+            kws,
+            global_dict,
+        ) = cls._get_function_info(func_ir)
+
+        self = cls(
+            native,
+            modname,
+            qualname,
+            unique_name,
+            doc,
+            typemap,
+            restype,
+            calltypes,
+            args,
+            kws,
+            mangler=mangler,
+            inline=inline,
+            noalias=noalias,
+            global_dict=global_dict,
+            abi_tags=abi_tags,
+            uid=func_ir.func_id.unique_id,
+        )
+        return self
+
+
+class PythonFunctionDescriptor(FunctionDescriptor):
+    """
+    A FunctionDescriptor subclass for Numba-compiled functions.
+    """
+
+    __slots__ = ()
+
+    @classmethod
+    def from_specialized_function(
+        cls,
+        func_ir,
+        typemap,
+        restype,
+        calltypes,
+        mangler,
+        inline,
+        noalias,
+        abi_tags,
+    ):
+        """
+        Build a FunctionDescriptor for a given specialization of a Python
+        function (in nopython mode).
+        """
+        return cls._from_python_function(
+            func_ir,
+            typemap,
+            restype,
+            calltypes,
+            native=True,
+            mangler=mangler,
+            inline=inline,
+            noalias=noalias,
+            abi_tags=abi_tags,
+        )
+
+    @classmethod
+    def from_object_mode_function(cls, func_ir):
+        """
+        Build a FunctionDescriptor for an object mode variant of a Python
+        function.
+        """
+        typemap = defaultdict(lambda: types.pyobject)
+        calltypes = typemap.copy()
+        restype = types.pyobject
+        return cls._from_python_function(
+            func_ir, typemap, restype, calltypes, native=False
+        )
+
+
+class ExternalFunctionDescriptor(FunctionDescriptor):
+    """
+    A FunctionDescriptor subclass for opaque external functions
+    (e.g. raw C functions).
+    """
+
+    __slots__ = ()
+
+    def __init__(self, name, restype, argtypes):
+        args = ["arg%d" % i for i in range(len(argtypes))]
+
+        def mangler(a, x, abi_tags, uid=None):
+            return a
+
+        super(ExternalFunctionDescriptor, self).__init__(
+            native=True,
+            modname=None,
+            qualname=name,
+            unique_name=name,
+            doc="",
+            typemap=None,
+            restype=restype,
+            calltypes=None,
+            args=args,
+            kws=None,
+            mangler=mangler,
+            argtypes=argtypes,
+        )

numba_cuda/numba/cuda/core/generators.py

@@ -0,0 +1,387 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Support for lowering generators.
+"""
+
+import llvmlite.ir
+from llvmlite.ir import Constant, IRBuilder
+
+from numba.cuda import types, config, cgutils
+from numba.cuda.core.funcdesc import FunctionDescriptor
+
+
+class GeneratorDescriptor(FunctionDescriptor):
+    """
+    The descriptor for a generator's next function.
+    """
+
+    __slots__ = ()
+
+    @classmethod
+    def from_generator_fndesc(cls, func_ir, fndesc, gentype, mangler):
+        """
+        Build a GeneratorDescriptor for the generator returned by the
+        function described by *fndesc*, with type *gentype*.
+
+        The generator inherits the env_name from the *fndesc*.
+        All emitted functions for the generator shares the same Env.
+        """
+        assert isinstance(gentype, types.Generator)
+        restype = gentype.yield_type
+        args = ["gen"]
+        argtypes = (gentype,)
+        qualname = fndesc.qualname + ".next"
+        unique_name = fndesc.unique_name + ".next"
+        self = cls(
+            fndesc.native,
+            fndesc.modname,
+            qualname,
+            unique_name,
+            fndesc.doc,
+            fndesc.typemap,
+            restype,
+            fndesc.calltypes,
+            args,
+            fndesc.kws,
+            argtypes=argtypes,
+            mangler=mangler,
+            inline=False,
+            env_name=fndesc.env_name,
+        )
+        return self
+
+    @property
+    def llvm_finalizer_name(self):
+        """
+        The LLVM name of the generator's finalizer function
+        (if <generator type>.has_finalizer is true).
+        """
+        return "finalize_" + self.mangled_name
+
+
+class BaseGeneratorLower(object):
+    """
+    Base support class for lowering generators.
+    """
+
+    def __init__(self, lower):
+        self.context = lower.context
+        self.fndesc = lower.fndesc
+        self.library = lower.library
+        self.func_ir = lower.func_ir
+        self.lower = lower
+
+        self.geninfo = lower.generator_info
+        self.gentype = self.get_generator_type()
+        self.gendesc = GeneratorDescriptor.from_generator_fndesc(
+            lower.func_ir, self.fndesc, self.gentype, self.context.mangler
+        )
+        # Helps packing non-omitted arguments into a structure
+        self.arg_packer = self.context.get_data_packer(self.fndesc.argtypes)
+
+        self.resume_blocks = {}
+
+    @property
+    def call_conv(self):
+        return self.lower.call_conv
+
+    def get_args_ptr(self, builder, genptr):
+        return cgutils.gep_inbounds(builder, genptr, 0, 1)
+
+    def get_resume_index_ptr(self, builder, genptr):
+        return cgutils.gep_inbounds(
+            builder, genptr, 0, 0, name="gen.resume_index"
+        )
+
+    def get_state_ptr(self, builder, genptr):
+        return cgutils.gep_inbounds(builder, genptr, 0, 2, name="gen.state")
+
+    def lower_init_func(self, lower):
+        """
+        Lower the generator's initialization function (which will fill up
+        the passed-by-reference generator structure).
+        """
+        lower.setup_function(self.fndesc)
+
+        builder = lower.builder
+
+        # Insert the generator into the target context in order to allow
+        # calling from other Numba-compiled functions.
+        lower.context.insert_generator(
+            self.gentype, self.gendesc, [self.library]
+        )
+
+        # Init argument values
+        lower.extract_function_arguments()
+
+        lower.pre_lower()
+
+        # Initialize the return structure (i.e. the generator structure).
+        retty = self.context.get_return_type(self.gentype)
+        # Structure index #0: the initial resume index (0 == start of generator)
+        resume_index = self.context.get_constant(types.int32, 0)
+        # Structure index #2: the states
+        statesty = retty.elements[2]
+
+        lower.debug_print("# low_init_func incref")
+        # Incref all NRT arguments before storing into generator states
+        if self.context.enable_nrt:
+            for argty, argval in zip(self.fndesc.argtypes, lower.fnargs):
+                self.context.nrt.incref(builder, argty, argval)
+
+        # Filter out omitted arguments
+        argsval = self.arg_packer.as_data(builder, lower.fnargs)
+
+        # Zero initialize states
+        statesval = Constant(statesty, None)
+        gen_struct = cgutils.make_anonymous_struct(
+            builder, [resume_index, argsval, statesval], retty
+        )
+
+        retval = self.box_generator_struct(lower, gen_struct)
+
+        lower.debug_print("# low_init_func before return")
+        self.call_conv.return_value(builder, retval)
+        lower.post_lower()
+
+    def lower_next_func(self, lower):
+        """
+        Lower the generator's next() function (which takes the
+        passed-by-reference generator structure and returns the next
+        yielded value).
+        """
+        lower.setup_function(self.gendesc)
+        lower.debug_print(
+            "# lower_next_func: {0}".format(self.gendesc.unique_name)
+        )
+        assert self.gendesc.argtypes[0] == self.gentype
+        builder = lower.builder
+        function = lower.function
+
+        # Extract argument values and other information from generator struct
+        (genptr,) = self.call_conv.get_arguments(function)
+        self.arg_packer.load_into(
+            builder, self.get_args_ptr(builder, genptr), lower.fnargs
+        )
+
+        self.resume_index_ptr = self.get_resume_index_ptr(builder, genptr)
+        self.gen_state_ptr = self.get_state_ptr(builder, genptr)
+
+        prologue = function.append_basic_block("generator_prologue")
+
+        # Lower the generator's Python code
+        entry_block_tail = lower.lower_function_body()
+
+        # Add block for StopIteration on entry
+        stop_block = function.append_basic_block("stop_iteration")
+        builder.position_at_end(stop_block)
+        self.call_conv.return_stop_iteration(builder)
+
+        # Add prologue switch to resume blocks
+        builder.position_at_end(prologue)
+        # First Python block is also the resume point on first next() call
+        self.resume_blocks[0] = lower.blkmap[lower.firstblk]
+
+        # Create front switch to resume points
+        switch = builder.switch(builder.load(self.resume_index_ptr), stop_block)
+        for index, block in self.resume_blocks.items():
+            switch.add_case(index, block)
+
+        # Close tail of entry block
+        builder.position_at_end(entry_block_tail)
+        builder.branch(prologue)
+
+    def lower_finalize_func(self, lower):
+        """
+        Lower the generator's finalizer.
+        """
+        fnty = llvmlite.ir.FunctionType(
+            llvmlite.ir.VoidType(), [self.context.get_value_type(self.gentype)]
+        )
+        function = cgutils.get_or_insert_function(
+            lower.module, fnty, self.gendesc.llvm_finalizer_name
+        )
+        entry_block = function.append_basic_block("entry")
+        builder = IRBuilder(entry_block)
+
+        genptrty = self.context.get_value_type(self.gentype)
+        genptr = builder.bitcast(function.args[0], genptrty)
+        self.lower_finalize_func_body(builder, genptr)
+
+    def return_from_generator(self, lower):
+        """
+        Emit a StopIteration at generator end and mark the generator exhausted.
+        """
+        indexval = Constant(self.resume_index_ptr.type.pointee, -1)
+        lower.builder.store(indexval, self.resume_index_ptr)
+        self.call_conv.return_stop_iteration(lower.builder)
+
+    def create_resumption_block(self, lower, index):
+        block_name = "generator_resume%d" % (index,)
+        block = lower.function.append_basic_block(block_name)
+        lower.builder.position_at_end(block)
+        self.resume_blocks[index] = block
+
+    def debug_print(self, builder, msg):
+        if config.DEBUG_JIT:
+            self.context.debug_print(builder, "DEBUGJIT: {0}".format(msg))
+
+
+class GeneratorLower(BaseGeneratorLower):
+    """
+    Support class for lowering nopython generators.
+    """
+
+    def get_generator_type(self):
+        return self.fndesc.restype
+
+    def box_generator_struct(self, lower, gen_struct):
+        return gen_struct
+
+    def lower_finalize_func_body(self, builder, genptr):
+        """
+        Lower the body of the generator's finalizer: decref all live
+        state variables.
+        """
+        self.debug_print(builder, "# generator: finalize")
+        if self.context.enable_nrt:
+            # Always dereference all arguments
+            # self.debug_print(builder, "# generator: clear args")
+            args_ptr = self.get_args_ptr(builder, genptr)
+            for ty, val in self.arg_packer.load(builder, args_ptr):
+                self.context.nrt.decref(builder, ty, val)
+
+        self.debug_print(builder, "# generator: finalize end")
+        builder.ret_void()
+
+
+class PyGeneratorLower(BaseGeneratorLower):
+    """
+    Support class for lowering object mode generators.
+    """
+
+    def get_generator_type(self):
+        """
+        Compute the actual generator type (the generator function's return
+        type is simply "pyobject").
+        """
+        return types.Generator(
+            gen_func=self.func_ir.func_id.func,
+            yield_type=types.pyobject,
+            arg_types=(types.pyobject,) * self.func_ir.arg_count,
+            state_types=(types.pyobject,) * len(self.geninfo.state_vars),
+            has_finalizer=True,
+        )
+
+    def box_generator_struct(self, lower, gen_struct):
+        """
+        Box the raw *gen_struct* as a Python object.
+        """
+        gen_ptr = cgutils.alloca_once_value(lower.builder, gen_struct)
+        return lower.pyapi.from_native_generator(
+            gen_ptr, self.gentype, lower.envarg
+        )
+
+    def init_generator_state(self, lower):
+        """
+        NULL-initialize all generator state variables, to avoid spurious
+        decref's on cleanup.
+        """
+        lower.builder.store(
+            Constant(self.gen_state_ptr.type.pointee, None), self.gen_state_ptr
+        )
+
+    def lower_finalize_func_body(self, builder, genptr):
+        """
+        Lower the body of the generator's finalizer: decref all live
+        state variables.
+        """
+        pyapi = self.context.get_python_api(builder)
+        resume_index_ptr = self.get_resume_index_ptr(builder, genptr)
+        resume_index = builder.load(resume_index_ptr)
+        # If resume_index is 0, next() was never called
+        # If resume_index is -1, generator terminated cleanly
+        # (note function arguments are saved in state variables,
+        #  so they don't need a separate cleanup step)
+        need_cleanup = builder.icmp_signed(
+            ">", resume_index, Constant(resume_index.type, 0)
+        )
+
+        with cgutils.if_unlikely(builder, need_cleanup):
+            # Decref all live vars (some may be NULL)
+            gen_state_ptr = self.get_state_ptr(builder, genptr)
+            for state_index in range(len(self.gentype.state_types)):
+                state_slot = cgutils.gep_inbounds(
+                    builder, gen_state_ptr, 0, state_index
+                )
+                ty = self.gentype.state_types[state_index]
+                val = self.context.unpack_value(builder, ty, state_slot)
+                pyapi.decref(val)
+
+        builder.ret_void()
+
+
+class LowerYield(object):
+    """
+    Support class for lowering a particular yield point.
+    """
+
+    def __init__(self, lower, yield_point, live_vars):
+        self.lower = lower
+        self.context = lower.context
+        self.builder = lower.builder
+        self.genlower = lower.genlower
+        self.gentype = self.genlower.gentype
+
+        self.gen_state_ptr = self.genlower.gen_state_ptr
+        self.resume_index_ptr = self.genlower.resume_index_ptr
+        self.yp = yield_point
+        self.inst = self.yp.inst
+        self.live_vars = live_vars
+        self.live_var_indices = [
+            lower.generator_info.state_vars.index(v) for v in live_vars
+        ]
+
+    def lower_yield_suspend(self):
+        self.lower.debug_print("# generator suspend")
+        # Save live vars in state
+        for state_index, name in zip(self.live_var_indices, self.live_vars):
+            state_slot = cgutils.gep_inbounds(
+                self.builder, self.gen_state_ptr, 0, state_index
+            )
+            ty = self.gentype.state_types[state_index]
+            # The yield might be in a loop, in which case the state might
+            # contain a predicate var that branches back to the loop head, in
+            # this case the var is live but in sequential lowering won't have
+            # been alloca'd yet, so do this here.
+            fetype = self.lower.typeof(name)
+            self.lower._alloca_var(name, fetype)
+            val = self.lower.loadvar(name)
+            # IncRef newly stored value
+            if self.context.enable_nrt:
+                self.context.nrt.incref(self.builder, ty, val)
+
+            self.context.pack_value(self.builder, ty, val, state_slot)
+        # Save resume index
+        indexval = Constant(self.resume_index_ptr.type.pointee, self.inst.index)
+        self.builder.store(indexval, self.resume_index_ptr)
+        self.lower.debug_print("# generator suspend end")
+
+    def lower_yield_resume(self):
+        # Emit resumption point
+        self.genlower.create_resumption_block(self.lower, self.inst.index)
+        self.lower.debug_print("# generator resume")
+        # Reload live vars from state
+        for state_index, name in zip(self.live_var_indices, self.live_vars):
+            state_slot = cgutils.gep_inbounds(
+                self.builder, self.gen_state_ptr, 0, state_index
+            )
+            ty = self.gentype.state_types[state_index]
+            val = self.context.unpack_value(self.builder, ty, state_slot)
+            self.lower.storevar(val, name)
+            # Previous storevar is making an extra incref
+            if self.context.enable_nrt:
+                self.context.nrt.decref(self.builder, ty, val)
+        self.lower.debug_print("# generator resume end")