numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/core/registry.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda import utils
+
+
+class DelayedRegistry(utils.UniqueDict):
+    """
+    A unique dictionary but with deferred initialisation of the values.
+
+    Attributes
+    ----------
+    ondemand:
+
+        A dictionary of key -> value, where value is executed
+        the first time it is is used.  It is used for part of a deferred
+        initialization strategy.
+    """
+
+    def __init__(self, *args, **kws):
+        self.ondemand = utils.UniqueDict()
+        self.key_type = kws.pop("key_type", None)
+        self.value_type = kws.pop("value_type", None)
+        self._type_check = self.key_type or self.value_type
+        super(DelayedRegistry, self).__init__(*args, **kws)
+
+    def __getitem__(self, item):
+        if item in self.ondemand:
+            self[item] = self.ondemand[item]()
+            del self.ondemand[item]
+        return super(DelayedRegistry, self).__getitem__(item)
+
+    def __setitem__(self, key, value):
+        if self._type_check:
+
+            def check(x, ty_x):
+                if isinstance(ty_x, type):
+                    assert ty_x in x.__mro__, (x, ty_x)
+                else:
+                    assert isinstance(x, ty_x), (x, ty_x)
+
+            if self.key_type is not None:
+                check(key, self.key_type)
+            if self.value_type is not None:
+                check(value, self.value_type)
+        return super(DelayedRegistry, self).__setitem__(key, value)

numba_cuda/numba/cuda/core/removerefctpass.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Implement a rewrite pass on a LLVM module to remove unnecessary
+refcount operations.
+"""
+
+from llvmlite.ir.transforms import CallVisitor
+
+from numba.cuda import types
+
+
+class _MarkNrtCallVisitor(CallVisitor):
+    """
+    A pass to mark all NRT_incref and NRT_decref.
+    """
+
+    def __init__(self):
+        self.marked = set()
+
+    def visit_Call(self, instr):
+        if getattr(instr.callee, "name", "") in _accepted_nrtfns:
+            self.marked.add(instr)
+
+
+def _rewrite_function(function):
+    # Mark NRT usage
+    markpass = _MarkNrtCallVisitor()
+    markpass.visit_Function(function)
+    # Remove NRT usage
+    for bb in function.basic_blocks:
+        for inst in list(bb.instructions):
+            if inst in markpass.marked:
+                bb.instructions.remove(inst)
+
+
+_accepted_nrtfns = "NRT_incref", "NRT_decref"
+
+
+def _legalize(module, dmm, fndesc):
+    """
+    Legalize the code in the module.
+    Returns True if the module is legal for the rewrite pass that removes
+    unnecessary refcounts.
+    """
+
+    def valid_output(ty):
+        """
+        Valid output are any type that does not need refcount
+        """
+        model = dmm[ty]
+        return not model.contains_nrt_meminfo()
+
+    def valid_input(ty):
+        """
+        Valid input are any type that does not need refcount except Array.
+        """
+        return valid_output(ty) or isinstance(ty, types.Array)
+
+    # Ensure no reference to function marked as
+    # "numba_args_may_always_need_nrt"
+    try:
+        nmd = module.get_named_metadata("numba_args_may_always_need_nrt")
+    except KeyError:
+        # Nothing marked
+        pass
+    else:
+        # Has functions marked as "numba_args_may_always_need_nrt"
+        if len(nmd.operands) > 0:
+            # The pass is illegal for this compilation unit.
+            return False
+
+    # More legalization base on function type
+    argtypes = fndesc.argtypes
+    restype = fndesc.restype
+    calltypes = fndesc.calltypes
+
+    # Legalize function arguments
+    for argty in argtypes:
+        if not valid_input(argty):
+            return False
+
+    # Legalize function return
+    if not valid_output(restype):
+        return False
+
+    # Legalize all called functions
+    for callty in calltypes.values():
+        if callty is not None and not valid_output(callty.return_type):
+            return False
+
+    # Ensure no allocation
+    for fn in module.functions:
+        if fn.name.startswith("NRT_"):
+            if fn.name not in _accepted_nrtfns:
+                return False
+
+    return True
+
+
+def remove_unnecessary_nrt_usage(function, context, fndesc):
+    """
+    Remove unnecessary NRT incref/decref in the given LLVM function.
+    It uses highlevel type info to determine if the function does not need NRT.
+    Such a function does not:
+
+    - return array object(s);
+    - take arguments that need refcounting except array;
+    - call function(s) that return refcounted object.
+
+    In effect, the function will not capture or create references that extend
+    the lifetime of any refcounted objects beyond the lifetime of the function.
+
+    The rewrite is performed in place.
+    If rewrite has happened, this function returns True, otherwise, it returns False.
+    """
+    dmm = context.data_model_manager
+    if _legalize(function.module, dmm, fndesc):
+        _rewrite_function(function)
+        return True
+    else:
+        return False

numba_cuda/numba/cuda/core/rewrites/__init__.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+A subpackage hosting Numba IR rewrite passes.
+"""
+
+from .registry import register_rewrite, rewrite_registry, Rewrite
+
+# Register various built-in rewrite passes
+from numba.cuda.core.rewrites import (
+    static_getitem,
+    static_raise,
+    static_binop,
+    ir_print,
+)
+
+__all__ = (
+    "static_getitem",
+    "static_raise",
+    "static_binop",
+    "ir_print",
+    "register_rewrite",
+    "rewrite_registry",
+    "Rewrite",
+)

numba_cuda/numba/cuda/core/rewrites/ir_print.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.core import errors
+from numba.cuda.core import ir
+from numba.cuda.core.rewrites import register_rewrite, Rewrite
+
+
+@register_rewrite("before-inference")
+class RewritePrintCalls(Rewrite):
+    """
+    Rewrite calls to the print() global function to dedicated IR print() nodes.
+    """
+
+    def match(self, func_ir, block, typemap, calltypes):
+        self.prints = prints = {}
+        self.block = block
+        # Find all assignments with a right-hand print() call
+        for inst in block.find_insts(ir.Assign):
+            if isinstance(inst.value, ir.Expr) and inst.value.op == "call":
+                expr = inst.value
+                try:
+                    callee = func_ir.infer_constant(expr.func)
+                except errors.ConstantInferenceError:
+                    continue
+                if callee is print:
+                    if expr.kws:
+                        # Only positional args are supported
+                        msg = (
+                            "Numba's print() function implementation does not "
+                            "support keyword arguments."
+                        )
+                        raise errors.UnsupportedError(msg, inst.loc)
+                    prints[inst] = expr
+        return len(prints) > 0
+
+    def apply(self):
+        """
+        Rewrite `var = call <print function>(...)` as a sequence of
+        `print(...)` and `var = const(None)`.
+        """
+        new_block = self.block.copy()
+        new_block.clear()
+        for inst in self.block.body:
+            if inst in self.prints:
+                expr = self.prints[inst]
+                print_node = ir.Print(
+                    args=expr.args, vararg=expr.vararg, loc=expr.loc
+                )
+                new_block.append(print_node)
+                assign_node = ir.Assign(
+                    value=ir.Const(None, loc=expr.loc),
+                    target=inst.target,
+                    loc=inst.loc,
+                )
+                new_block.append(assign_node)
+            else:
+                new_block.append(inst)
+        return new_block
+
+
+@register_rewrite("before-inference")
+class DetectConstPrintArguments(Rewrite):
+    """
+    Detect and store constant arguments to print() nodes.
+    """
+
+    def match(self, func_ir, block, typemap, calltypes):
+        self.consts = consts = {}
+        self.block = block
+        for inst in block.find_insts(ir.Print):
+            if inst.consts:
+                # Already rewritten
+                continue
+            for idx, var in enumerate(inst.args):
+                try:
+                    const = func_ir.infer_constant(var)
+                except errors.ConstantInferenceError:
+                    continue
+                consts.setdefault(inst, {})[idx] = const
+
+        return len(consts) > 0
+
+    def apply(self):
+        """
+        Store detected constant arguments on their nodes.
+        """
+        for inst in self.block.body:
+            if inst in self.consts:
+                inst.consts = self.consts[inst]
+        return self.block

numba_cuda/numba/cuda/core/rewrites/registry.py

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from collections import defaultdict
+
+from numba.cuda import config
+
+
+class Rewrite(object):
+    """Defines the abstract base class for Numba rewrites."""
+
+    def __init__(self, state=None):
+        """Constructor for the Rewrite class."""
+        pass
+
+    def match(self, func_ir, block, typemap, calltypes) -> bool:
+        """Overload this method to check an IR block for matching terms in the
+        rewrite.
+        """
+        return False
+
+    def apply(self):
+        """Overload this method to return a rewritten IR basic block when a
+        match has been found.
+        """
+        raise NotImplementedError("Abstract Rewrite.apply() called!")
+
+
+class RewriteRegistry(object):
+    """Defines a registry for Numba rewrites."""
+
+    _kinds = frozenset(["before-inference", "after-inference"])
+
+    def __init__(self):
+        """Constructor for the rewrite registry.  Initializes the rewrites
+        member to an empty list.
+        """
+        self.rewrites = defaultdict(list)
+
+    def register(self, kind):
+        """
+        Decorator adding a subclass of Rewrite to the registry for
+        the given *kind*.
+        """
+        if kind not in self._kinds:
+            raise KeyError("invalid kind %r" % (kind,))
+
+        def do_register(rewrite_cls):
+            if not issubclass(rewrite_cls, Rewrite):
+                raise TypeError(
+                    "{0} is not a subclass of Rewrite".format(rewrite_cls)
+                )
+            self.rewrites[kind].append(rewrite_cls)
+            return rewrite_cls
+
+        return do_register
+
+    def apply(self, kind, state):
+        """Given a pipeline and a dictionary of basic blocks, exhaustively
+        attempt to apply all registered rewrites to all basic blocks.
+        """
+        assert kind in self._kinds
+        blocks = state.func_ir.blocks
+        old_blocks = blocks.copy()
+        for rewrite_cls in self.rewrites[kind]:
+            # Exhaustively apply a rewrite until it stops matching.
+            rewrite = rewrite_cls(state)
+            work_list = list(blocks.items())
+            while work_list:
+                key, block = work_list.pop()
+                matches = rewrite.match(
+                    state.func_ir, block, state.typemap, state.calltypes
+                )
+                if matches:
+                    if config.DEBUG or config.DUMP_IR:
+                        print("_" * 70)
+                        print("REWRITING (%s):" % rewrite_cls.__name__)
+                        block.dump()
+                        print("_" * 60)
+                    new_block = rewrite.apply()
+                    blocks[key] = new_block
+                    work_list.append((key, new_block))
+                    if config.DEBUG or config.DUMP_IR:
+                        new_block.dump()
+                        print("_" * 70)
+        # If any blocks were changed, perform a sanity check.
+        for key, block in blocks.items():
+            if block != old_blocks[key]:
+                block.verify()
+
+        # Some passes, e.g. _inline_const_arraycall are known to occasionally
+        # do invalid things WRT ir.Del, others, e.g. RewriteArrayExprs do valid
+        # things with ir.Del, but the placement is not optimal. The lines below
+        # fix-up the IR so that ref counts are valid and optimally placed,
+        # see #4093 for context. This has to be run here opposed to in
+        # apply() as the CFG needs computing so full IR is needed.
+        from numba.cuda.core import postproc
+
+        post_proc = postproc.PostProcessor(state.func_ir)
+        post_proc.run()
+
+
+rewrite_registry = RewriteRegistry()
+register_rewrite = rewrite_registry.register

numba_cuda/numba/cuda/core/rewrites/static_binop.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.core import errors
+from numba.cuda.core import ir
+from numba.cuda.core.rewrites import register_rewrite, Rewrite
+
+
+@register_rewrite("before-inference")
+class DetectStaticBinops(Rewrite):
+    """
+    Detect constant arguments to select binops.
+    """
+
+    # Those operators can benefit from a constant-inferred argument
+    rhs_operators = {"**"}
+
+    def match(self, func_ir, block, typemap, calltypes):
+        self.static_lhs = {}
+        self.static_rhs = {}
+        self.block = block
+        # Find binop expressions with a constant lhs or rhs
+        for expr in block.find_exprs(op="binop"):
+            try:
+                if (
+                    expr.fn in self.rhs_operators
+                    and expr.static_rhs is ir.UNDEFINED
+                ):
+                    self.static_rhs[expr] = func_ir.infer_constant(expr.rhs)
+            except errors.ConstantInferenceError:
+                continue
+
+        return len(self.static_lhs) > 0 or len(self.static_rhs) > 0
+
+    def apply(self):
+        """
+        Store constant arguments that were detected in match().
+        """
+        for expr, rhs in self.static_rhs.items():
+            expr.static_rhs = rhs
+        return self.block

numba_cuda/numba/cuda/core/rewrites/static_getitem.py

@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.core import errors
+from numba.cuda.core import ir
+from numba.cuda import types
+from numba.cuda.core.rewrites import register_rewrite, Rewrite
+
+
+@register_rewrite("before-inference")
+class RewriteConstGetitems(Rewrite):
+    """
+    Rewrite IR expressions of the kind `getitem(value=arr, index=$constXX)`
+    where `$constXX` is a known constant as
+    `static_getitem(value=arr, index=<constant value>)`.
+    """
+
+    def match(self, func_ir, block, typemap, calltypes):
+        self.getitems = getitems = {}
+        self.block = block
+        # Detect all getitem expressions and find which ones can be
+        # rewritten
+        for expr in block.find_exprs(op="getitem"):
+            if expr.op == "getitem":
+                try:
+                    const = func_ir.infer_constant(expr.index)
+                except errors.ConstantInferenceError:
+                    continue
+                getitems[expr] = const
+
+        return len(getitems) > 0
+
+    def apply(self):
+        """
+        Rewrite all matching getitems as static_getitems.
+        """
+        new_block = self.block.copy()
+        new_block.clear()
+        for inst in self.block.body:
+            if isinstance(inst, ir.Assign):
+                expr = inst.value
+                if expr in self.getitems:
+                    const = self.getitems[expr]
+                    new_expr = ir.Expr.static_getitem(
+                        value=expr.value,
+                        index=const,
+                        index_var=expr.index,
+                        loc=expr.loc,
+                    )
+                    inst = ir.Assign(
+                        value=new_expr, target=inst.target, loc=inst.loc
+                    )
+            new_block.append(inst)
+        return new_block
+
+
+@register_rewrite("after-inference")
+class RewriteStringLiteralGetitems(Rewrite):
+    """
+    Rewrite IR expressions of the kind `getitem(value=arr, index=$XX)`
+    where `$XX` is a StringLiteral value as
+    `static_getitem(value=arr, index=<literal value>)`.
+    """
+
+    def match(self, func_ir, block, typemap, calltypes):
+        """
+        Detect all getitem expressions and find which ones have
+        string literal indexes
+        """
+        self.getitems = getitems = {}
+        self.block = block
+        self.calltypes = calltypes
+        for expr in block.find_exprs(op="getitem"):
+            if expr.op == "getitem":
+                index_ty = typemap[expr.index.name]
+                if isinstance(index_ty, types.StringLiteral):
+                    getitems[expr] = (expr.index, index_ty.literal_value)
+
+        return len(getitems) > 0
+
+    def apply(self):
+        """
+        Rewrite all matching getitems as static_getitems where the index
+        is the literal value of the string.
+        """
+        new_block = ir.Block(self.block.scope, self.block.loc)
+        for inst in self.block.body:
+            if isinstance(inst, ir.Assign):
+                expr = inst.value
+                if expr in self.getitems:
+                    const, lit_val = self.getitems[expr]
+                    new_expr = ir.Expr.static_getitem(
+                        value=expr.value,
+                        index=lit_val,
+                        index_var=expr.index,
+                        loc=expr.loc,
+                    )
+                    self.calltypes[new_expr] = self.calltypes[expr]
+                    inst = ir.Assign(
+                        value=new_expr, target=inst.target, loc=inst.loc
+                    )
+            new_block.append(inst)
+        return new_block
+
+
+@register_rewrite("after-inference")
+class RewriteStringLiteralSetitems(Rewrite):
+    """
+    Rewrite IR expressions of the kind `setitem(value=arr, index=$XX, value=)`
+    where `$XX` is a StringLiteral value as
+    `static_setitem(value=arr, index=<literal value>, value=)`.
+    """
+
+    def match(self, func_ir, block, typemap, calltypes):
+        """
+        Detect all setitem expressions and find which ones have
+        string literal indexes
+        """
+        self.setitems = setitems = {}
+        self.block = block
+        self.calltypes = calltypes
+        for inst in block.find_insts(ir.SetItem):
+            index_ty = typemap[inst.index.name]
+            if isinstance(index_ty, types.StringLiteral):
+                setitems[inst] = (inst.index, index_ty.literal_value)
+
+        return len(setitems) > 0
+
+    def apply(self):
+        """
+        Rewrite all matching setitems as static_setitems where the index
+        is the literal value of the string.
+        """
+        new_block = ir.Block(self.block.scope, self.block.loc)
+        for inst in self.block.body:
+            if isinstance(inst, ir.SetItem):
+                if inst in self.setitems:
+                    const, lit_val = self.setitems[inst]
+                    new_inst = ir.StaticSetItem(
+                        target=inst.target,
+                        index=lit_val,
+                        index_var=inst.index,
+                        value=inst.value,
+                        loc=inst.loc,
+                    )
+                    self.calltypes[new_inst] = self.calltypes[inst]
+                    inst = new_inst
+            new_block.append(inst)
+        return new_block
+
+
+@register_rewrite("before-inference")
+class RewriteConstSetitems(Rewrite):
+    """
+    Rewrite IR statements of the kind `setitem(target=arr, index=$constXX, ...)`
+    where `$constXX` is a known constant as
+    `static_setitem(target=arr, index=<constant value>, ...)`.
+    """
+
+    def match(self, func_ir, block, typemap, calltypes):
+        self.setitems = setitems = {}
+        self.block = block
+        # Detect all setitem statements and find which ones can be
+        # rewritten
+        for inst in block.find_insts(ir.SetItem):
+            try:
+                const = func_ir.infer_constant(inst.index)
+            except errors.ConstantInferenceError:
+                continue
+            setitems[inst] = const
+
+        return len(setitems) > 0
+
+    def apply(self):
+        """
+        Rewrite all matching setitems as static_setitems.
+        """
+        new_block = self.block.copy()
+        new_block.clear()
+        for inst in self.block.body:
+            if inst in self.setitems:
+                const = self.setitems[inst]
+                new_inst = ir.StaticSetItem(
+                    inst.target, const, inst.index, inst.value, inst.loc
+                )
+                new_block.append(new_inst)
+            else:
+                new_block.append(inst)
+        return new_block