Numba: work around quadratic lowering of >30-argument calls

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

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
+<https://github.com/numba/numba/pull/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

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)