Vendor in generators and removerefctpass for CUDA-specific changes

Author: VijayKandiah

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")