From 9da1e5b7235775060bc89add76454014933309ff Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Mon, 17 Aug 2026 17:34:55 +0300 Subject: [PATCH] Numba: patch quadratic lowering of >30-item calls and tuple displays MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit numba's peep_hole_list_to_tuple rewrites CPython's >30-item LIST_APPEND bytecode (STACK_USE_GUIDELINE) into a tuple of every prefix length, which types and lowers as O(n^2) LLVM IR plus NRT refcount churn — generated fgraph functions make such calls routinely. Collapse the chain into a single build_tuple after the peephole runs. Fixed upstream in numba/numba#10782; this patch covers already-released numba versions, in the same mold as _patch_pointer_add. Co-Authored-By: Claude Fable 5 --- pytensor/link/numba/dispatch/__init__.py | 3 + .../numba/dispatch/_patch_list_to_tuple.py | 99 +++++++++++++++++++ tests/link/numba/test_patch_list_to_tuple.py | 93 +++++++++++++++++ 3 files changed, 195 insertions(+) create mode 100644 pytensor/link/numba/dispatch/_patch_list_to_tuple.py create mode 100644 tests/link/numba/test_patch_list_to_tuple.py diff --git a/pytensor/link/numba/dispatch/__init__.py b/pytensor/link/numba/dispatch/__init__.py index e8c4aab8db..1dd5bb43bb 100644 --- a/pytensor/link/numba/dispatch/__init__.py +++ b/pytensor/link/numba/dispatch/__init__.py @@ -2,6 +2,9 @@ # Patch numba's cgutils.pointer_add before any dispatch codegen runs import pytensor.link.numba.dispatch._patch_pointer_add +# Patch numba's list_to_tuple peephole so >30-argument calls don't lower quadratically +import pytensor.link.numba.dispatch._patch_list_to_tuple + from pytensor.link.numba.dispatch.basic import numba_funcify, numba_typify # Load dispatch specializations diff --git a/pytensor/link/numba/dispatch/_patch_list_to_tuple.py b/pytensor/link/numba/dispatch/_patch_list_to_tuple.py new file mode 100644 index 0000000000..4b71520e30 --- /dev/null +++ b/pytensor/link/numba/dispatch/_patch_list_to_tuple.py @@ -0,0 +1,99 @@ +"""Patch ``numba.core.interpreter.peep_hole_list_to_tuple`` to emit one +``build_tuple`` instead of a chain of tuple concatenations. + +Numba turns >30-item calls and tuple displays (CPython's +``STACK_USE_GUIDELINE`` bytecode) into an IR tuple of every prefix length, +which types and lowers as O(n^2) LLVM IR — generated fgraph functions make +such calls routinely. Fixed upstream in `numba#10782 +`_; imported for its side effect, +drop this module once a numba release ships the fix. +""" + +import operator +from collections import Counter + +from numba.core import interpreter, ir + + +_peep_hole_list_to_tuple_orig = interpreter.peep_hole_list_to_tuple + + +def _collapse_tuple_chains(func_ir): + # Number of *uses* of each variable (definition sites excluded) + uses = Counter() + for blk in func_ir.blocks.values(): + for stmt in blk.body: + vars_ = ( + stmt.value.list_vars() + if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) + else stmt.list_vars() + ) + for var in vars_: + uses[var.name] += 1 + + defs = func_ir._definitions + for blk in func_ir.blocks.values(): + # var name -> (body index, items) of single-def build_tuple assignments + tuples: dict = {} + body = blk.body + changed = False + for idx, stmt in enumerate(body): + if not isinstance(stmt, ir.Assign): + continue + if isinstance(stmt.value, ir.Var): + # Forward the tuple through single-use aliases: numba's + # CALL_FUNCTION_EX peephole requires a call's vararg to be + # defined by a build_tuple directly when kwargs are present + name = stmt.value.name + if name in tuples and uses[name] == 1: + t_idx, t_items = tuples.pop(name) + new_expr = ir.Expr.build_tuple(t_items, stmt.loc) + defs[name].clear() + defs[stmt.target.name].remove(stmt.value) + defs[stmt.target.name].append(new_expr) + stmt = ir.Assign(new_expr, stmt.target, stmt.loc) + body[t_idx] = None + body[idx] = stmt + changed = True + if len(defs[stmt.target.name]) == 1: + tuples[stmt.target.name] = (idx, list(t_items)) + continue + if not isinstance(stmt.value, ir.Expr): + continue + expr = stmt.value + if ( + expr.op == "binop" + and expr.fn is operator.add + and expr.lhs.name in tuples + and expr.rhs.name in tuples + and uses[expr.lhs.name] == 1 + and uses[expr.rhs.name] == 1 + ): + l_idx, l_items = tuples.pop(expr.lhs.name) + r_idx, r_items = tuples.pop(expr.rhs.name) + new_expr = ir.Expr.build_tuple(l_items + r_items, expr.loc) + defs[expr.lhs.name].clear() + defs[expr.rhs.name].clear() + defs[stmt.target.name].remove(expr) + defs[stmt.target.name].append(new_expr) + stmt = ir.Assign(new_expr, stmt.target, stmt.loc) + body[l_idx] = None + body[r_idx] = None + body[idx] = stmt + expr = new_expr + changed = True + if expr.op == "build_tuple" and len(defs.get(stmt.target.name, ())) == 1: + tuples[stmt.target.name] = (idx, list(expr.items)) + if changed: + new_body = [s for s in body if s is not None] + body.clear() + body.extend(new_body) + return func_ir + + +def _peep_hole_list_to_tuple_flat(func_ir): + return _collapse_tuple_chains(_peep_hole_list_to_tuple_orig(func_ir)) + + +# ``Interpreter.interpret`` resolves the peephole through the module at call time +interpreter.peep_hole_list_to_tuple = _peep_hole_list_to_tuple_flat diff --git a/tests/link/numba/test_patch_list_to_tuple.py b/tests/link/numba/test_patch_list_to_tuple.py new file mode 100644 index 0000000000..48ec747afa --- /dev/null +++ b/tests/link/numba/test_patch_list_to_tuple.py @@ -0,0 +1,93 @@ +import operator + +import pytest + + +numba = pytest.importorskip("numba") + +import pytensor.link.numba.dispatch # noqa: F401 (installs the patch) + + +def _interpret(fn): + from numba.core import bytecode, interpreter + + func_id = bytecode.FunctionIdentity.from_function(fn) + interp = interpreter.Interpreter(func_id) + return interp.interpret(bytecode.ByteCode(func_id)) + + +def _make_wide_caller(n_args): + args_sig = ", ".join(f"a{i}" for i in range(n_args)) + src = f"def caller(callee, {args_sig}):\n return callee({args_sig})\n" + glb: dict = {} + exec(src, glb) + return glb["caller"] + + +def test_wide_call_collapses_to_single_build_tuple(): + """A >30-argument call leaves one build_tuple, not a chain of prefix tuples.""" + from numba.core import ir + + func_ir = _interpret(_make_wide_caller(40)) + tuple_adds = [ + stmt + for blk in func_ir.blocks.values() + for stmt in blk.body + if isinstance(stmt, ir.Assign) + and isinstance(stmt.value, ir.Expr) + and stmt.value.op == "binop" + and stmt.value.fn is operator.add + ] + assert not tuple_adds + widths = [ + len(stmt.value.items) + for blk in func_ir.blocks.values() + for stmt in blk.body + if isinstance(stmt, ir.Assign) + and isinstance(stmt.value, ir.Expr) + and stmt.value.op == "build_tuple" + ] + assert widths == [40] + + +def test_wide_call_computes_correctly(): + @numba.njit + def callee(*args): + total = 0.0 + for a in args: + total += a + return total + + caller = numba.njit(_make_wide_caller(35)) + vals = [float(i) for i in range(35)] + assert caller(callee, *vals) == sum(vals) + + +def test_wide_call_with_kwarg(): + """>30 positional args plus a keyword argument: the collapsed tuple must + reach the call's vararg directly or numba's kwargs peephole rejects it.""" + n_args = 40 + params = ", ".join(f"a{i}" for i in range(n_args)) + glb: dict = {} + exec( + f"def callee({params}, k):\n" + f" return a0 + k\n" + f"def caller({params}):\n" + f" return jitted_callee({params}, k=1.0)\n", + glb, + ) + glb["jitted_callee"] = numba.njit(glb["callee"]) + caller = numba.njit(glb["caller"]) + assert caller(*[float(i) for i in range(n_args)]) == 1.0 + + +def test_tuple_unpacking_still_works(): + """Chains interleaved with genuine unpacking still compute correctly.""" + + @numba.njit + def spread(a, b): + t = (*a, 1.0, *b, 2.0) + return len(t), t[len(a)], t[-1] + + n, mid, last = spread((3.0, 4.0), (5.0,)) + assert (n, mid, last) == (5, 1.0, 2.0)