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

numba_cuda/numba/cuda/core/typed_passes.py

@@ -1,9 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 import abc
-import warnings
 from contextlib import contextmanager
-from numba.core import errors, types, funcdesc
-from numba.core.compiler_machinery import LoweringPass
-from llvmlite import binding as llvm
+from collections import defaultdict, namedtuple
+from copy import copy
+import warnings
+
+from numba.cuda.core import typeinfer
+from numba.core import (
+    errors,
+    types,
+    typing,
+    ir,
+    lowering,
+)
+
+from numba.cuda.core.compiler_machinery import (
+    FunctionPass,
+    LoweringPass,
+    AnalysisPass,
+    register_pass,
+)
+from numba.cuda.core.annotations import type_annotations
+from numba.cuda.core.ir_utils import (
+    raise_on_unsupported_feature,
+    warn_deprecated,
+    check_and_legalize_ir,
+    guard,
+    dead_code_elimination,
+    simplify_CFG,
+    get_definition,
+    compute_cfg_from_blocks,
+    is_operator_or_getitem,
+)
+
+from numba.cuda.core import postproc, rewrites, funcdesc, config
+
+
+try:
+    # llvmlite < 0.45
+    from llvmlite.binding import passmanagers
+except ImportError:
+    # llvmlite >= 0.45
+    from llvmlite.binding import newpassmanagers as passmanagers
+
+# Outputs of type inference pass
+_TypingResults = namedtuple(
+    "_TypingResults",
+    [
+        "typemap",
+        "return_type",
+        "calltypes",
+        "typing_errors",
+    ],
+)
 
 
 @contextmanager
@@ -35,6 +86,216 @@ def fallback_context(state, msg):
             raise
 
 
+def type_inference_stage(
+    typingctx,
+    targetctx,
+    interp,
+    args,
+    return_type,
+    locals=None,
+    raise_errors=True,
+):
+    if locals is None:
+        locals = {}
+    if len(args) != interp.arg_count:
+        raise TypeError("Mismatch number of argument types")
+    warnings = errors.WarningsFixer(errors.NumbaWarning)
+
+    infer = typeinfer.TypeInferer(typingctx, interp, warnings)
+    callstack_ctx = typingctx.callstack.register(
+        targetctx.target, infer, interp.func_id, args
+    )
+    # Setup two contexts: 1) callstack setup/teardown 2) flush warnings
+    with callstack_ctx, warnings:
+        # Seed argument types
+        for index, (name, ty) in enumerate(zip(interp.arg_names, args)):
+            infer.seed_argument(name, index, ty)
+
+        # Seed return type
+        if return_type is not None:
+            infer.seed_return(return_type)
+
+        # Seed local types
+        for k, v in locals.items():
+            infer.seed_type(k, v)
+
+        infer.build_constraint()
+        # return errors in case of partial typing
+        errs = infer.propagate(raise_errors=raise_errors)
+        typemap, restype, calltypes = infer.unify(raise_errors=raise_errors)
+
+    return _TypingResults(typemap, restype, calltypes, errs)
+
+
+class BaseTypeInference(FunctionPass):
+    _raise_errors = True
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    def run_pass(self, state):
+        """
+        Type inference and legalization
+        """
+        with fallback_context(
+            state,
+            'Function "%s" failed type inference' % (state.func_id.func_name,),
+        ):
+            # Type inference
+            typemap, return_type, calltypes, errs = type_inference_stage(
+                state.typingctx,
+                state.targetctx,
+                state.func_ir,
+                state.args,
+                state.return_type,
+                state.locals,
+                raise_errors=self._raise_errors,
+            )
+            state.typemap = typemap
+            # save errors in case of partial typing
+            state.typing_errors = errs
+            if self._raise_errors:
+                state.return_type = return_type
+            state.calltypes = calltypes
+
+        def legalize_return_type(return_type, interp, targetctx):
+            """
+            Only accept array return type iff it is passed into the function.
+            Reject function object return types if in nopython mode.
+            """
+            if not targetctx.enable_nrt and isinstance(
+                return_type, types.Array
+            ):
+                # Walk IR to discover all arguments and all return statements
+                retstmts = []
+                caststmts = {}
+                argvars = set()
+                for bid, blk in interp.blocks.items():
+                    for inst in blk.body:
+                        if isinstance(inst, ir.Return):
+                            retstmts.append(inst.value.name)
+                        elif isinstance(inst, ir.Assign):
+                            if (
+                                isinstance(inst.value, ir.Expr)
+                                and inst.value.op == "cast"
+                            ):
+                                caststmts[inst.target.name] = inst.value
+                            elif isinstance(inst.value, ir.Arg):
+                                argvars.add(inst.target.name)
+
+                assert retstmts, "No return statements?"
+
+                for var in retstmts:
+                    cast = caststmts.get(var)
+                    if cast is None or cast.value.name not in argvars:
+                        if self._raise_errors:
+                            msg = (
+                                "Only accept returning of array passed into "
+                                "the function as argument"
+                            )
+                            raise errors.NumbaTypeError(msg)
+
+            elif isinstance(return_type, types.Function) or isinstance(
+                return_type, types.Phantom
+            ):
+                if self._raise_errors:
+                    msg = "Can't return function object ({}) in nopython mode"
+                    raise errors.NumbaTypeError(msg.format(return_type))
+
+        with fallback_context(
+            state,
+            'Function "%s" has invalid return type'
+            % (state.func_id.func_name,),
+        ):
+            legalize_return_type(
+                state.return_type, state.func_ir, state.targetctx
+            )
+        return True
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class NopythonTypeInference(BaseTypeInference):
+    _name = "nopython_type_inference"
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class PartialTypeInference(BaseTypeInference):
+    _name = "partial_type_inference"
+    _raise_errors = False
+
+
+@register_pass(mutates_CFG=False, analysis_only=False)
+class AnnotateTypes(AnalysisPass):
+    _name = "annotate_types"
+
+    def __init__(self):
+        AnalysisPass.__init__(self)
+
+    def get_analysis_usage(self, AU):
+        AU.add_required(IRLegalization)
+
+    def run_pass(self, state):
+        """
+        Create type annotation after type inference
+        """
+        func_ir = state.func_ir.copy()
+        state.type_annotation = type_annotations.TypeAnnotation(
+            func_ir=func_ir,
+            typemap=state.typemap,
+            calltypes=state.calltypes,
+            lifted=state.lifted,
+            lifted_from=state.lifted_from,
+            args=state.args,
+            return_type=state.return_type,
+            html_output=config.HTML,
+        )
+
+        if config.ANNOTATE:
+            print("ANNOTATION".center(80, "-"))
+            print(state.type_annotation)
+            print("=" * 80)
+        if config.HTML:
+            with open(config.HTML, "w") as fout:
+                state.type_annotation.html_annotate(fout)
+
+        return False
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class NopythonRewrites(FunctionPass):
+    _name = "nopython_rewrites"
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    def run_pass(self, state):
+        """
+        Perform any intermediate representation rewrites after type
+        inference.
+        """
+        # a bunch of these passes are either making assumptions or rely on some
+        # very picky and slightly bizarre state particularly in relation to
+        # ir.Del presence. To accommodate, ir.Dels are added ahead of running
+        # this pass and stripped at the end.
+
+        # Ensure we have an IR and type information.
+        assert state.func_ir
+        assert isinstance(getattr(state, "typemap", None), dict)
+        assert isinstance(getattr(state, "calltypes", None), dict)
+        msg = (
+            "Internal error in post-inference rewriting "
+            "pass encountered during compilation of "
+            'function "%s"' % (state.func_id.func_name,)
+        )
+
+        pp = postproc.PostProcessor(state.func_ir)
+        pp.run(True)
+        with fallback_context(state, msg):
+            rewrites.rewrite_registry.apply("after-inference", state)
+        pp.remove_dels()
+        return True
+
+
 class BaseNativeLowering(abc.ABC, LoweringPass):
     """The base class for a lowering pass. The lowering functionality must be
     specified in inheriting classes by providing an appropriate lowering class
@@ -68,7 +329,7 @@ class BaseNativeLowering(abc.ABC, LoweringPass):
         calltypes = state.calltypes
         flags = state.flags
         metadata = state.metadata
-        pre_stats = llvm.passmanagers.dump_refprune_stats()
+        pre_stats = passmanagers.dump_refprune_stats()
 
         msg = "Function %s failed at nopython mode lowering" % (
             state.func_id.func_name,
@@ -114,7 +375,7 @@ class BaseNativeLowering(abc.ABC, LoweringPass):
                 call_helper = lower.call_helper
                 del lower
 
-            from numba.core.compiler import _LowerResult  # TODO: move this
+            from numba.cuda.compiler import _LowerResult  # TODO: move this
 
             if flags.no_compile:
                 state["cr"] = _LowerResult(
@@ -131,9 +392,482 @@ class BaseNativeLowering(abc.ABC, LoweringPass):
                 )
 
             # capture pruning stats
-            post_stats = llvm.passmanagers.dump_refprune_stats()
+            post_stats = passmanagers.dump_refprune_stats()
             metadata["prune_stats"] = post_stats - pre_stats
 
             # Save the LLVM pass timings
             metadata["llvm_pass_timings"] = library.recorded_timings
         return True
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class NativeLowering(BaseNativeLowering):
+    """Lowering pass for a native function IR described solely in terms of
+    Numba's standard `numba.core.ir` nodes."""
+
+    _name = "native_lowering"
+
+    @property
+    def lowering_class(self):
+        return lowering.Lower
+
+
+@register_pass(mutates_CFG=False, analysis_only=True)
+class NoPythonSupportedFeatureValidation(AnalysisPass):
+    """NoPython Mode check: Validates the IR to ensure that features in use are
+    in a form that is supported"""
+
+    _name = "nopython_supported_feature_validation"
+
+    def __init__(self):
+        AnalysisPass.__init__(self)
+
+    def run_pass(self, state):
+        raise_on_unsupported_feature(state.func_ir, state.typemap)
+        warn_deprecated(state.func_ir, state.typemap)
+        return False
+
+
+@register_pass(mutates_CFG=False, analysis_only=True)
+class IRLegalization(AnalysisPass):
+    _name = "ir_legalization"
+
+    def __init__(self):
+        AnalysisPass.__init__(self)
+
+    def run_pass(self, state):
+        # NOTE: this function call must go last, it checks and fixes invalid IR!
+        check_and_legalize_ir(state.func_ir, flags=state.flags)
+        return True
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class NoPythonBackend(LoweringPass):
+    _name = "nopython_backend"
+
+    def __init__(self):
+        LoweringPass.__init__(self)
+
+    def run_pass(self, state):
+        """
+        Back-end: Generate LLVM IR from Numba IR, compile to machine code
+        """
+        lowered = state["cr"]
+        signature = typing.signature(state.return_type, *state.args)
+
+        from numba.cuda.compiler import cuda_compile_result
+
+        state.cr = cuda_compile_result(
+            typing_context=state.typingctx,
+            target_context=state.targetctx,
+            entry_point=lowered.cfunc,
+            typing_error=state.status.fail_reason,
+            type_annotation=state.type_annotation,
+            library=state.library,
+            call_helper=lowered.call_helper,
+            signature=signature,
+            objectmode=False,
+            lifted=state.lifted,
+            fndesc=lowered.fndesc,
+            environment=lowered.env,
+            metadata=state.metadata,
+            reload_init=state.reload_init,
+        )
+        return True
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class InlineOverloads(FunctionPass):
+    """
+    This pass will inline a function wrapped by the numba.extending.overload
+    decorator directly into the site of its call depending on the value set in
+    the 'inline' kwarg to the decorator.
+
+    This is a typed pass. CFG simplification and DCE are performed on
+    completion.
+    """
+
+    _name = "inline_overloads"
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    _DEBUG = False
+
+    def run_pass(self, state):
+        """Run inlining of overloads"""
+        if self._DEBUG:
+            print("before overload inline".center(80, "-"))
+            print(state.func_id.unique_name)
+            print(state.func_ir.dump())
+            print("".center(80, "-"))
+        from numba.cuda.core.inline_closurecall import (
+            InlineWorker,
+            callee_ir_validator,
+        )
+
+        inline_worker = InlineWorker(
+            state.typingctx,
+            state.targetctx,
+            state.locals,
+            state.pipeline,
+            state.flags,
+            callee_ir_validator,
+            state.typemap,
+            state.calltypes,
+        )
+        modified = False
+        work_list = list(state.func_ir.blocks.items())
+        # use a work list, look for call sites via `ir.Expr.op == call` and
+        # then pass these to `self._do_work` to make decisions about inlining.
+        while work_list:
+            label, block = work_list.pop()
+            for i, instr in enumerate(block.body):
+                # TO-DO: other statements (setitem)
+                if isinstance(instr, ir.Assign):
+                    expr = instr.value
+                    if isinstance(expr, ir.Expr):
+                        workfn = self._do_work_expr
+
+                        if guard(
+                            workfn,
+                            state,
+                            work_list,
+                            block,
+                            i,
+                            expr,
+                            inline_worker,
+                        ):
+                            modified = True
+                            break  # because block structure changed
+
+        if self._DEBUG:
+            print("after overload inline".center(80, "-"))
+            print(state.func_id.unique_name)
+            print(state.func_ir.dump())
+            print("".center(80, "-"))
+
+        if modified:
+            # Remove dead blocks, this is safe as it relies on the CFG only.
+            cfg = compute_cfg_from_blocks(state.func_ir.blocks)
+            for dead in cfg.dead_nodes():
+                del state.func_ir.blocks[dead]
+            # clean up blocks
+            dead_code_elimination(state.func_ir, typemap=state.typemap)
+            # clean up unconditional branches that appear due to inlined
+            # functions introducing blocks
+            state.func_ir.blocks = simplify_CFG(state.func_ir.blocks)
+
+        if self._DEBUG:
+            print("after overload inline DCE".center(80, "-"))
+            print(state.func_id.unique_name)
+            print(state.func_ir.dump())
+            print("".center(80, "-"))
+        return True
+
+    def _get_attr_info(self, state, expr):
+        recv_type = state.typemap[expr.value.name]
+        recv_type = types.unliteral(recv_type)
+        matched = state.typingctx.find_matching_getattr_template(
+            recv_type,
+            expr.attr,
+        )
+        if not matched:
+            return None
+
+        template = matched["template"]
+        if getattr(template, "is_method", False):
+            # The attribute template is representing a method.
+            # Don't inline the getattr.
+            return None
+
+        templates = [template]
+        sig = typing.signature(matched["return_type"], recv_type)
+        arg_typs = sig.args
+        is_method = False
+
+        return templates, sig, arg_typs, is_method
+
+    def _get_callable_info(self, state, expr):
+        def get_func_type(state, expr):
+            func_ty = None
+            if expr.op == "call":
+                # check this is a known and typed function
+                try:
+                    func_ty = state.typemap[expr.func.name]
+                except KeyError:
+                    # e.g. Calls to CUDA Intrinsic have no mapped type
+                    # so KeyError
+                    return None
+                if not hasattr(func_ty, "get_call_type"):
+                    return None
+
+            elif is_operator_or_getitem(expr):
+                func_ty = state.typingctx.resolve_value_type(expr.fn)
+            else:
+                return None
+
+            return func_ty
+
+        if expr.op == "call":
+            # try and get a definition for the call, this isn't always
+            # possible as it might be a eval(str)/part generated
+            # awaiting update etc. (parfors)
+            to_inline = None
+            try:
+                to_inline = state.func_ir.get_definition(expr.func)
+            except Exception:
+                return None
+
+            # do not handle closure inlining here, another pass deals with that
+            if getattr(to_inline, "op", False) == "make_function":
+                return None
+
+        func_ty = get_func_type(state, expr)
+        if func_ty is None:
+            return None
+
+        sig = state.calltypes[expr]
+        if not sig:
+            return None
+
+        templates, arg_typs, is_method = None, None, False
+        if getattr(func_ty, "template", None) is not None:
+            # @overload_method
+            is_method = True
+            templates = [func_ty.template]
+            arg_typs = (func_ty.template.this,) + sig.args
+        else:
+            # @overload case
+            templates = getattr(func_ty, "templates", None)
+            arg_typs = sig.args
+
+        return templates, sig, arg_typs, is_method
+
+    def _do_work_expr(self, state, work_list, block, i, expr, inline_worker):
+        def select_template(templates, args):
+            if templates is None:
+                return None
+
+            impl = None
+            for template in templates:
+                inline_type = getattr(template, "_inline", None)
+                if inline_type is None:
+                    # inline not defined
+                    continue
+                if args not in template._inline_overloads:
+                    # skip overloads not matching signature
+                    continue
+                if not inline_type.is_never_inline:
+                    try:
+                        impl = template._overload_func(*args)
+                        if impl is None:
+                            raise Exception  # abort for this template
+                        break
+                    except Exception:
+                        continue
+            else:
+                return None
+
+            return template, inline_type, impl
+
+        inlinee_info = None
+        if expr.op == "getattr":
+            inlinee_info = self._get_attr_info(state, expr)
+        else:
+            inlinee_info = self._get_callable_info(state, expr)
+
+        if not inlinee_info:
+            return False
+
+        templates, sig, arg_typs, is_method = inlinee_info
+        inlinee = select_template(templates, arg_typs)
+        if inlinee is None:
+            return False
+        template, inlinee_type, impl = inlinee
+
+        return self._run_inliner(
+            state,
+            inlinee_type,
+            sig,
+            template,
+            arg_typs,
+            expr,
+            i,
+            impl,
+            block,
+            work_list,
+            is_method,
+            inline_worker,
+        )
+
+    def _run_inliner(
+        self,
+        state,
+        inline_type,
+        sig,
+        template,
+        arg_typs,
+        expr,
+        i,
+        impl,
+        block,
+        work_list,
+        is_method,
+        inline_worker,
+    ):
+        do_inline = True
+        if not inline_type.is_always_inline:
+            from numba.core.typing.templates import _inline_info
+
+            caller_inline_info = _inline_info(
+                state.func_ir, state.typemap, state.calltypes, sig
+            )
+
+            # must be a cost-model function, run the function
+            iinfo = template._inline_overloads[arg_typs]["iinfo"]
+            if inline_type.has_cost_model:
+                do_inline = inline_type.value(expr, caller_inline_info, iinfo)
+            else:
+                assert "unreachable"
+
+        if do_inline:
+            if is_method:
+                if not self._add_method_self_arg(state, expr):
+                    return False
+            arg_typs = template._inline_overloads[arg_typs]["folded_args"]
+            iinfo = template._inline_overloads[arg_typs]["iinfo"]
+            freevars = iinfo.func_ir.func_id.func.__code__.co_freevars
+            _, _, _, new_blocks = inline_worker.inline_ir(
+                state.func_ir,
+                block,
+                i,
+                iinfo.func_ir,
+                freevars,
+                arg_typs=arg_typs,
+            )
+            if work_list is not None:
+                for blk in new_blocks:
+                    work_list.append(blk)
+            return True
+        else:
+            return False
+
+    def _add_method_self_arg(self, state, expr):
+        func_def = guard(get_definition, state.func_ir, expr.func)
+        if func_def is None:
+            return False
+        expr.args.insert(0, func_def.value)
+        return True
+
+
+@register_pass(mutates_CFG=False, analysis_only=False)
+class DeadCodeElimination(FunctionPass):
+    """
+    Does dead code elimination
+    """
+
+    _name = "dead_code_elimination"
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    def run_pass(self, state):
+        dead_code_elimination(state.func_ir, state.typemap)
+        return True
+
+
+@register_pass(mutates_CFG=False, analysis_only=False)
+class PreLowerStripPhis(FunctionPass):
+    """Remove phi nodes (ir.Expr.phi) introduced by SSA.
+
+    This is needed before Lowering because the phi nodes in Numba IR do not
+    match the semantics of phi nodes in LLVM IR. In Numba IR, phi nodes may
+    expand into multiple LLVM instructions.
+    """
+
+    _name = "strip_phis"
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    def run_pass(self, state):
+        state.func_ir = self._strip_phi_nodes(state.func_ir)
+
+        # Rerun postprocessor to update metadata
+        post_proc = postproc.PostProcessor(state.func_ir)
+        post_proc.run(emit_dels=False)
+
+        # Ensure we are not in objectmode generator
+        if (
+            state.func_ir.generator_info is not None
+            and state.typemap is not None
+        ):
+            # Rebuild generator type
+            # TODO: move this into PostProcessor
+            gentype = state.return_type
+            state_vars = state.func_ir.generator_info.state_vars
+            state_types = [state.typemap[k] for k in state_vars]
+            state.return_type = types.Generator(
+                gen_func=gentype.gen_func,
+                yield_type=gentype.yield_type,
+                arg_types=gentype.arg_types,
+                state_types=state_types,
+                has_finalizer=gentype.has_finalizer,
+            )
+        return True
+
+    def _strip_phi_nodes(self, func_ir):
+        """Strip Phi nodes from ``func_ir``
+
+        For each phi node, put incoming value to their respective incoming
+        basic-block at possibly the latest position (i.e. after the latest
+        assignment to the corresponding variable).
+        """
+        exporters = defaultdict(list)
+        phis = set()
+        # Find all variables that needs to be exported
+        for label, block in func_ir.blocks.items():
+            for assign in block.find_insts(ir.Assign):
+                if isinstance(assign.value, ir.Expr):
+                    if assign.value.op == "phi":
+                        phis.add(assign)
+                        phi = assign.value
+                        for ib, iv in zip(
+                            phi.incoming_blocks, phi.incoming_values
+                        ):
+                            exporters[ib].append((assign.target, iv))
+
+        # Rewrite the blocks with the new exporting assignments
+        newblocks = {}
+        for label, block in func_ir.blocks.items():
+            newblk = copy(block)
+            newblocks[label] = newblk
+
+            # strip phis
+            newblk.body = [stmt for stmt in block.body if stmt not in phis]
+
+            # insert exporters
+            for target, rhs in exporters[label]:
+                # If RHS is undefined
+                if rhs is ir.UNDEFINED:
+                    # Put in a NULL initializer, set the location to be in what
+                    # will eventually materialize as the prologue.
+                    rhs = ir.Expr.null(loc=func_ir.loc)
+
+                assign = ir.Assign(target=target, value=rhs, loc=rhs.loc)
+                # Insert at the earliest possible location; i.e. after the
+                # last assignment to rhs
+                assignments = [
+                    stmt
+                    for stmt in newblk.find_insts(ir.Assign)
+                    if stmt.target == rhs
+                ]
+                if assignments:
+                    last_assignment = assignments[-1]
+                    newblk.insert_after(assign, last_assignment)
+                else:
+                    newblk.prepend(assign)
+
+        func_ir.blocks = newblocks
+        return func_ir