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

numba_cuda/numba/cuda/core/analysis.py

@@ -0,0 +1,319 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from collections import namedtuple
+from numba import types
+from numba.core import ir
+from numba.cuda.core import consts
+from numba.core.analysis import compute_cfg_from_blocks
+
+
+# Used to describe a nullified condition in dead branch pruning
+nullified = namedtuple("nullified", "condition, taken_br, rewrite_stmt")
+
+
+def dead_branch_prune(func_ir, called_args):
+    """
+    Removes dead branches based on constant inference from function args.
+    This directly mutates the IR.
+
+    func_ir is the IR
+    called_args are the actual arguments with which the function is called
+    """
+    from numba.cuda.core.ir_utils import (
+        get_definition,
+        guard,
+        find_const,
+        GuardException,
+    )
+
+    DEBUG = 0
+
+    def find_branches(func_ir):
+        # find *all* branches
+        branches = []
+        for blk in func_ir.blocks.values():
+            branch_or_jump = blk.body[-1]
+            if isinstance(branch_or_jump, ir.Branch):
+                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 function.value is bool
+                    ):
+                        condition = guard(get_definition, func_ir, pred.args[0])
+                        if condition is not None:
+                            branches.append((branch, condition, blk))
+        return branches
+
+    def do_prune(take_truebr, blk):
+        keep = branch.truebr if take_truebr else branch.falsebr
+        # replace the branch with a direct jump
+        jmp = ir.Jump(keep, loc=branch.loc)
+        blk.body[-1] = jmp
+        return 1 if keep == branch.truebr else 0
+
+    def prune_by_type(branch, condition, blk, *conds):
+        # this prunes a given branch and fixes up the IR
+        # at least one needs to be a NoneType
+        lhs_cond, rhs_cond = conds
+        lhs_none = isinstance(lhs_cond, types.NoneType)
+        rhs_none = isinstance(rhs_cond, types.NoneType)
+        if lhs_none or rhs_none:
+            try:
+                take_truebr = condition.fn(lhs_cond, rhs_cond)
+            except Exception:
+                return False, None
+            if DEBUG > 0:
+                kill = branch.falsebr if take_truebr else branch.truebr
+                print(
+                    "Pruning %s" % kill,
+                    branch,
+                    lhs_cond,
+                    rhs_cond,
+                    condition.fn,
+                )
+            taken = do_prune(take_truebr, blk)
+            return True, taken
+        return False, None
+
+    def prune_by_value(branch, condition, blk, *conds):
+        lhs_cond, rhs_cond = conds
+        try:
+            take_truebr = condition.fn(lhs_cond, rhs_cond)
+        except Exception:
+            return False, None
+        if DEBUG > 0:
+            kill = branch.falsebr if take_truebr else branch.truebr
+            print("Pruning %s" % kill, branch, lhs_cond, rhs_cond, condition.fn)
+        do_prune(take_truebr, blk)
+        # It is not safe to rewrite the predicate to a nominal value based on
+        # which branch is taken, the rewritten const predicate needs to
+        # hold the actual computed const value as something else may refer to
+        # it!
+        return True, take_truebr
+
+    def prune_by_predicate(branch, pred, blk):
+        try:
+            # Just to prevent accidents, whilst already guarded, ensure this
+            # is an ir.Const
+            if not isinstance(pred, (ir.Const, ir.FreeVar, ir.Global)):
+                raise TypeError("Expected constant Numba IR node")
+            take_truebr = bool(pred.value)
+        except TypeError:
+            return False, None
+        if DEBUG > 0:
+            kill = branch.falsebr if take_truebr else branch.truebr
+            print("Pruning %s" % kill, branch, pred)
+        taken = do_prune(take_truebr, blk)
+        return True, taken
+
+    class Unknown(object):
+        pass
+
+    def resolve_input_arg_const(input_arg_idx):
+        """
+        Resolves an input arg to a constant (if possible)
+        """
+        input_arg_ty = called_args[input_arg_idx]
+
+        # comparing to None?
+        if isinstance(input_arg_ty, types.NoneType):
+            return input_arg_ty
+
+        # is it a kwarg default
+        if isinstance(input_arg_ty, types.Omitted):
+            val = input_arg_ty.value
+            if isinstance(val, types.NoneType):
+                return val
+            elif val is None:
+                return types.NoneType("none")
+
+        # literal type, return the type itself so comparisons like `x == None`
+        # still work as e.g. x = types.int64 will never be None/NoneType so
+        # the branch can still be pruned
+        return getattr(input_arg_ty, "literal_type", Unknown())
+
+    if DEBUG > 1:
+        print("before".center(80, "-"))
+        print(func_ir.dump())
+
+    phi2lbl = dict()
+    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":
+                    phi2lbl[stmt.value] = lbl
+                    phi2asgn[stmt.value] = stmt
+
+    # This looks for branches where:
+    # at least one arg of the condition is in input args and const
+    # at least one an arg of the condition is a const
+    # if the condition is met it will replace the branch with a jump
+    branch_info = find_branches(func_ir)
+    # stores conditions that have no impact post prune
+    nullified_conditions = []
+
+    for branch, condition, blk in branch_info:
+        const_conds = []
+        if isinstance(condition, ir.Expr) 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):
+                    # it's an e.g. literal argument to the function
+                    resolved_const = resolve_input_arg_const(arg_def.index)
+                    prune = prune_by_type
+                else:
+                    # it's some const argument to the function, cannot use guard
+                    # here as the const itself may be None
+                    try:
+                        resolved_const = find_const(func_ir, arg)
+                        if resolved_const is None:
+                            resolved_const = types.NoneType("none")
+                    except GuardException:
+                        pass
+
+                if not isinstance(resolved_const, Unknown):
+                    const_conds.append(resolved_const)
+
+            # lhs/rhs are consts
+            if len(const_conds) == 2:
+                # prune the branch, switch the branch for an unconditional jump
+                prune_stat, taken = prune(branch, condition, blk, *const_conds)
+                if prune_stat:
+                    # add the condition to the list of nullified conditions
+                    nullified_conditions.append(
+                        nullified(condition, taken, True)
+                    )
+        else:
+            # see if this is a branch on a constant value predicate
+            resolved_const = Unknown()
+            try:
+                pred_call = get_definition(func_ir, branch.cond)
+                resolved_const = find_const(func_ir, pred_call.args[0])
+                if resolved_const is None:
+                    resolved_const = types.NoneType("none")
+            except GuardException:
+                pass
+
+            if not isinstance(resolved_const, Unknown):
+                prune_stat, taken = prune_by_predicate(branch, condition, blk)
+                if prune_stat:
+                    # add the condition to the list of nullified conditions
+                    nullified_conditions.append(
+                        nullified(condition, taken, False)
+                    )
+
+    # 'ERE BE DRAGONS...
+    # It is the evaluation of the condition expression that often trips up type
+    # inference, so ideally it would be removed as it is effectively rendered
+    # dead by the unconditional jump if a branch was pruned. However, there may
+    # be references to the condition that exist in multiple places (e.g. dels)
+    # and we cannot run DCE here as typing has not taken place to give enough
+    # information to run DCE safely. Upshot of all this is the condition gets
+    # rewritten below into a benign const that typing will be happy with and DCE
+    # can remove it and its reference post typing when it is safe to do so
+    # (if desired). It is required that the const is assigned a value that
+    # indicates the branch taken as its mutated value would be read in the case
+    # of object mode fall back in place of the condition itself. For
+    # completeness the func_ir._definitions and ._consts are also updated to
+    # make the IR state self consistent.
+
+    deadcond = [x.condition for x in nullified_conditions]
+    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:
+                    # 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
+                    # their value as they may be used later.
+                    if nullified_info.rewrite_stmt:
+                        branch_bit = nullified_info.taken_br
+                        x.value = ir.Const(branch_bit, loc=x.loc)
+                        # update the specific definition to the new const
+                        defns = func_ir._definitions[x.target.name]
+                        repl_idx = defns.index(cond)
+                        defns[repl_idx] = x.value
+
+    # Check post dominators of dead nodes from in the original CFG for use of
+    # vars that are being removed in the dead blocks which might be referred to
+    # by phi nodes.
+    #
+    # Multiple things to fix up:
+    #
+    # 1. Cases like:
+    #
+    # A        A
+    # |\       |
+    # | B  --> B
+    # |/       |
+    # C        C
+    #
+    # i.e. the branch is dead but the block is still alive. In this case CFG
+    # simplification will fuse A-B-C and any phi in C can be updated as an
+    # direct assignment from the last assigned version in the dominators of the
+    # fused block.
+    #
+    # 2. Cases like:
+    #
+    #   A        A
+    #  / \       |
+    # B   C  --> B
+    #  \ /       |
+    #   D        D
+    #
+    # i.e. the block C is dead. In this case the phis in D need updating to
+    # reflect the collapse of the phi condition. This should result in a direct
+    # assignment of the surviving version in B to the LHS of the phi in D.
+
+    new_cfg = compute_cfg_from_blocks(func_ir.blocks)
+    dead_blocks = new_cfg.dead_nodes()
+
+    # for all phis that are still in live blocks.
+    for phi, lbl in phi2lbl.items():
+        if lbl in dead_blocks:
+            continue
+        new_incoming = [x[0] for x in new_cfg.predecessors(lbl)]
+        if set(new_incoming) != set(phi.incoming_blocks):
+            # Something has changed in the CFG...
+            if len(new_incoming) == 1:
+                # There's now just one incoming. Replace the PHI node by a
+                # direct assignment
+                idx = phi.incoming_blocks.index(new_incoming[0])
+                phi2asgn[phi].value = phi.incoming_values[idx]
+            else:
+                # There's more than one incoming still, then look through the
+                # incoming and remove dead
+                ic_val_tmp = []
+                ic_blk_tmp = []
+                for ic_val, ic_blk in zip(
+                    phi.incoming_values, phi.incoming_blocks
+                ):
+                    if ic_blk in dead_blocks:
+                        continue
+                    else:
+                        ic_val_tmp.append(ic_val)
+                        ic_blk_tmp.append(ic_blk)
+                phi.incoming_values.clear()
+                phi.incoming_values.extend(ic_val_tmp)
+                phi.incoming_blocks.clear()
+                phi.incoming_blocks.extend(ic_blk_tmp)
+
+    # Remove dead blocks, this is safe as it relies on the CFG only.
+    for dead in dead_blocks:
+        del func_ir.blocks[dead]
+
+    # if conditions were nullified then consts were rewritten, update
+    if nullified_conditions:
+        func_ir._consts = consts.ConstantInference(func_ir)
+
+    if DEBUG > 1:
+        print("after".center(80, "-"))
+        print(func_ir.dump())

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

@@ -0,0 +1,304 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from collections import defaultdict, OrderedDict
+from collections.abc import Mapping
+from contextlib import closing
+import copy
+import inspect
+import os
+import re
+import textwrap
+from io import StringIO
+
+import numba.core.dispatcher
+from numba.core import ir
+
+
+class SourceLines(Mapping):
+    def __init__(self, func):
+        try:
+            lines, startno = inspect.getsourcelines(func)
+        except OSError:
+            self.lines = ()
+            self.startno = 0
+        else:
+            self.lines = textwrap.dedent("".join(lines)).splitlines()
+            self.startno = startno
+
+    def __getitem__(self, lineno):
+        try:
+            return self.lines[lineno - self.startno].rstrip()
+        except IndexError:
+            return ""
+
+    def __iter__(self):
+        return iter((self.startno + i) for i in range(len(self.lines)))
+
+    def __len__(self):
+        return len(self.lines)
+
+    @property
+    def avail(self):
+        return bool(self.lines)
+
+
+class TypeAnnotation(object):
+    # func_data dict stores annotation data for all functions that are
+    # compiled. We store the data in the TypeAnnotation class since a new
+    # TypeAnnotation instance is created for each function that is compiled.
+    # For every function that is compiled, we add the type annotation data to
+    # this dict and write the html annotation file to disk (rewrite the html
+    # file for every function since we don't know if this is the last function
+    # to be compiled).
+    func_data = OrderedDict()
+
+    def __init__(
+        self,
+        func_ir,
+        typemap,
+        calltypes,
+        lifted,
+        lifted_from,
+        args,
+        return_type,
+        html_output=None,
+    ):
+        self.func_id = func_ir.func_id
+        self.blocks = func_ir.blocks
+        self.typemap = typemap
+        self.calltypes = calltypes
+        self.filename = func_ir.loc.filename
+        self.linenum = str(func_ir.loc.line)
+        self.signature = str(args) + " -> " + str(return_type)
+
+        # lifted loop information
+        self.lifted = lifted
+        self.num_lifted_loops = len(lifted)
+
+        # If this is a lifted loop function that is being compiled, lifted_from
+        # points to annotation data from function that this loop lifted function
+        # was lifted from. This is used to stick lifted loop annotations back
+        # into original function.
+        self.lifted_from = lifted_from
+
+    def prepare_annotations(self):
+        # Prepare annotations
+        groupedinst = defaultdict(list)
+        found_lifted_loop = False
+        # for blkid, blk in self.blocks.items():
+        for blkid in sorted(self.blocks.keys()):
+            blk = self.blocks[blkid]
+            groupedinst[blk.loc.line].append("label %s" % blkid)
+            for inst in blk.body:
+                lineno = inst.loc.line
+
+                if isinstance(inst, ir.Assign):
+                    if found_lifted_loop:
+                        atype = "XXX Lifted Loop XXX"
+                        found_lifted_loop = False
+                    elif (
+                        isinstance(inst.value, ir.Expr)
+                        and inst.value.op == "call"
+                    ):
+                        atype = self.calltypes[inst.value]
+                    elif isinstance(inst.value, ir.Const) and isinstance(
+                        inst.value.value, numba.core.dispatcher.LiftedLoop
+                    ):
+                        atype = "XXX Lifted Loop XXX"
+                        found_lifted_loop = True
+                    else:
+                        # TODO: fix parfor lowering so that typemap is valid.
+                        atype = self.typemap.get(inst.target.name, "<missing>")
+
+                    aline = "%s = %s  :: %s" % (inst.target, inst.value, atype)
+                elif isinstance(inst, ir.SetItem):
+                    atype = self.calltypes[inst]
+                    aline = "%s  :: %s" % (inst, atype)
+                else:
+                    aline = "%s" % inst
+                groupedinst[lineno].append("  %s" % aline)
+        return groupedinst
+
+    def annotate(self):
+        source = SourceLines(self.func_id.func)
+        # if not source.avail:
+        #     return "Source code unavailable"
+
+        groupedinst = self.prepare_annotations()
+
+        # Format annotations
+        io = StringIO()
+        with closing(io):
+            if source.avail:
+                print("# File: %s" % self.filename, file=io)
+                for num in source:
+                    srcline = source[num]
+                    ind = _getindent(srcline)
+                    print("%s# --- LINE %d --- " % (ind, num), file=io)
+                    for inst in groupedinst[num]:
+                        print("%s# %s" % (ind, inst), file=io)
+                    print(file=io)
+                    print(srcline, file=io)
+                    print(file=io)
+                if self.lifted:
+                    print("# The function contains lifted loops", file=io)
+                    for loop in self.lifted:
+                        print(
+                            "# Loop at line %d" % loop.get_source_location(),
+                            file=io,
+                        )
+                        print(
+                            "# Has %d overloads" % len(loop.overloads), file=io
+                        )
+                        for cres in loop.overloads.values():
+                            print(cres.type_annotation, file=io)
+            else:
+                print("# Source code unavailable", file=io)
+                for num in groupedinst:
+                    for inst in groupedinst[num]:
+                        print("%s" % (inst,), file=io)
+                    print(file=io)
+
+            return io.getvalue()
+
+    def html_annotate(self, outfile):
+        # ensure that annotation information is assembled
+        self.annotate_raw()
+        # make a deep copy ahead of the pending mutations
+        func_data = copy.deepcopy(self.func_data)
+
+        key = "python_indent"
+        for this_func in func_data.values():
+            if key in this_func:
+                idents = {}
+                for line, amount in this_func[key].items():
+                    idents[line] = "&nbsp;" * amount
+                this_func[key] = idents
+
+        key = "ir_indent"
+        for this_func in func_data.values():
+            if key in this_func:
+                idents = {}
+                for line, ir_id in this_func[key].items():
+                    idents[line] = ["&nbsp;" * amount for amount in ir_id]
+                this_func[key] = idents
+
+        try:
+            from jinja2 import Template
+        except ImportError:
+            raise ImportError("please install the 'jinja2' package")
+
+        root = os.path.join(os.path.dirname(__file__))
+        template_filename = os.path.join(root, "template.html")
+        with open(template_filename, "r") as template:
+            html = template.read()
+
+        template = Template(html)
+        rendered = template.render(func_data=func_data)
+        outfile.write(rendered)
+
+    def annotate_raw(self):
+        """
+        This returns "raw" annotation information i.e. it has no output format
+        specific markup included.
+        """
+        python_source = SourceLines(self.func_id.func)
+        ir_lines = self.prepare_annotations()
+        line_nums = [num for num in python_source]
+        lifted_lines = [l.get_source_location() for l in self.lifted]
+
+        def add_ir_line(func_data, line):
+            line_str = line.strip()
+            line_type = ""
+            if line_str.endswith("pyobject"):
+                line_str = line_str.replace("pyobject", "")
+                line_type = "pyobject"
+            func_data["ir_lines"][num].append((line_str, line_type))
+            indent_len = len(_getindent(line))
+            func_data["ir_indent"][num].append(indent_len)
+
+        func_key = (
+            self.func_id.filename + ":" + str(self.func_id.firstlineno + 1),
+            self.signature,
+        )
+        if (
+            self.lifted_from is not None
+            and self.lifted_from[1]["num_lifted_loops"] > 0
+        ):
+            # This is a lifted loop function that is being compiled. Get the
+            # numba ir for lines in loop function to use for annotating
+            # original python function that the loop was lifted from.
+            func_data = self.lifted_from[1]
+            for num in line_nums:
+                if num not in ir_lines.keys():
+                    continue
+                func_data["ir_lines"][num] = []
+                func_data["ir_indent"][num] = []
+                for line in ir_lines[num]:
+                    add_ir_line(func_data, line)
+                    if line.strip().endswith("pyobject"):
+                        func_data["python_tags"][num] = "object_tag"
+                        # If any pyobject line is found, make sure original python
+                        # line that was marked as a lifted loop start line is tagged
+                        # as an object line instead. Lifted loop start lines should
+                        # only be marked as lifted loop lines if the lifted loop
+                        # was successfully compiled in nopython mode.
+                        func_data["python_tags"][self.lifted_from[0]] = (
+                            "object_tag"
+                        )
+
+            # We're done with this lifted loop, so decrement lifted loop counter.
+            # When lifted loop counter hits zero, that means we're ready to write
+            # out annotations to html file.
+            self.lifted_from[1]["num_lifted_loops"] -= 1
+
+        elif func_key not in TypeAnnotation.func_data.keys():
+            TypeAnnotation.func_data[func_key] = {}
+            func_data = TypeAnnotation.func_data[func_key]
+
+            for i, loop in enumerate(self.lifted):
+                # Make sure that when we process each lifted loop function later,
+                # we'll know where it originally came from.
+                loop.lifted_from = (lifted_lines[i], func_data)
+            func_data["num_lifted_loops"] = self.num_lifted_loops
+
+            func_data["filename"] = self.filename
+            func_data["funcname"] = self.func_id.func_name
+            func_data["python_lines"] = []
+            func_data["python_indent"] = {}
+            func_data["python_tags"] = {}
+            func_data["ir_lines"] = {}
+            func_data["ir_indent"] = {}
+
+            for num in line_nums:
+                func_data["python_lines"].append(
+                    (num, python_source[num].strip())
+                )
+                indent_len = len(_getindent(python_source[num]))
+                func_data["python_indent"][num] = indent_len
+                func_data["python_tags"][num] = ""
+                func_data["ir_lines"][num] = []
+                func_data["ir_indent"][num] = []
+
+                for line in ir_lines[num]:
+                    add_ir_line(func_data, line)
+                    if num in lifted_lines:
+                        func_data["python_tags"][num] = "lifted_tag"
+                    elif line.strip().endswith("pyobject"):
+                        func_data["python_tags"][num] = "object_tag"
+        return self.func_data
+
+    def __str__(self):
+        return self.annotate()
+
+
+re_longest_white_prefix = re.compile(r"^\s*")
+
+
+def _getindent(text):
+    m = re_longest_white_prefix.match(text)
+    if not m:
+        return ""
+    else:
+        return " " * len(m.group(0))