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

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,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.core import errors, 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.core 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.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,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.core import errors, 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,187 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.core import errors, types, ir
+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

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

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.core import errors, consts, ir
+from numba.cuda.core.rewrites import register_rewrite, Rewrite
+
+
+@register_rewrite("before-inference")
+class RewriteConstRaises(Rewrite):
+    """
+    Rewrite IR statements of the kind `raise(value)`
+    where `value` is the result of instantiating an exception with
+    constant arguments
+    into `static_raise(exception_type, constant args)`.
+
+    This allows lowering in nopython mode, where one can't instantiate
+    exception instances from runtime data.
+    """
+
+    def _is_exception_type(self, const):
+        return isinstance(const, type) and issubclass(const, Exception)
+
+    def _break_constant(self, const, loc):
+        """
+        Break down constant exception.
+        """
+        if isinstance(const, tuple):  # it's a tuple(exception class, args)
+            if not self._is_exception_type(const[0]):
+                msg = "Encountered unsupported exception constant %r"
+                raise errors.UnsupportedError(msg % (const[0],), loc)
+            return const[0], tuple(const[1])
+        elif self._is_exception_type(const):
+            return const, None
+        else:
+            if isinstance(const, str):
+                msg = (
+                    "Directly raising a string constant as an exception is "
+                    "not supported."
+                )
+            else:
+                msg = "Encountered unsupported constant type used for exception"
+            raise errors.UnsupportedError(msg, loc)
+
+    def _try_infer_constant(self, func_ir, inst):
+        try:
+            return func_ir.infer_constant(inst.exception)
+        except consts.ConstantInferenceError:
+            # not a static exception
+            return None
+
+    def match(self, func_ir, block, typemap, calltypes):
+        self.raises = raises = {}
+        self.tryraises = tryraises = {}
+        self.block = block
+        # Detect all raise statements and find which ones can be
+        # rewritten
+        for inst in block.find_insts((ir.Raise, ir.TryRaise)):
+            if inst.exception is None:
+                # re-reraise
+                exc_type, exc_args = None, None
+            else:
+                # raise <something> => find the definition site for <something>
+                const = self._try_infer_constant(func_ir, inst)
+
+                # failure to infer constant indicates this isn't a static
+                # exception
+                if const is None:
+                    continue
+
+                loc = inst.exception.loc
+                exc_type, exc_args = self._break_constant(const, loc)
+
+            if isinstance(inst, ir.Raise):
+                raises[inst] = exc_type, exc_args
+            elif isinstance(inst, ir.TryRaise):
+                tryraises[inst] = exc_type, exc_args
+            else:
+                raise ValueError("unexpected: {}".format(type(inst)))
+        return (len(raises) + len(tryraises)) > 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.raises:
+                exc_type, exc_args = self.raises[inst]
+                new_inst = ir.StaticRaise(exc_type, exc_args, inst.loc)
+                new_block.append(new_inst)
+            elif inst in self.tryraises:
+                exc_type, exc_args = self.tryraises[inst]
+                new_inst = ir.StaticTryRaise(exc_type, exc_args, inst.loc)
+                new_block.append(new_inst)
+            else:
+                new_block.append(inst)
+        return new_block