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

[velochy]

Problem

Numba's peep_hole_list_to_tuple rewrites the bytecode CPython emits for >30-item calls and tuple displays (STACK_USE_GUIDELINE) into incremental tuple concatenation — an IR tuple of every prefix length — which types and lowers as O(n²) LLVM IR plus NRT refcount churn, with a hard cliff at 31 items (21× module-size jump). Generated fgraph functions make such calls for every >30-input node (MakeVector, Scan, wide fused kernels), so compile time and memory blow up quadratically in node arity.

Fix

_patch_list_to_tuple wraps the peephole, collapses the concatenation chain into a single build_tuple, and inlines that tuple back into direct call arguments (so no aggregate forms at either side of a wide call boundary) — the same patch-until-upstream-releases mold as _patch_pointer_add. The bug is fixed upstream in numba/numba#10782, but numba releases roughly twice a year and pytensor supports already-released versions, so the workaround is warranted until the minimum supported numba ships the fix; the module then just gets deleted.

Results

Interleaved with main in one session on an otherwise-idle 1-core box, cold caches (absolute times here are load-sensitive, so only same-session comparisons are quoted): on the n=80 additive repro this patch alone takes 364.9 s / 2696 MB to 323.3 s / 1689 MB — most of the win is memory, because the caller-side bytecode chain is only one of the two quadratics. Combined with #2362, which removes the wide-tuple callee in Join: 105.9 s / 933 MB. The two are complementary — this PR carries the memory win, #2362 the compile-time one. Pure-numba: a 160-argument jitted call compiles ~3x faster and its caller module is 34x smaller.

[ricardoV94]

Adversarial regression:

_patch_list_to_tuple breaks njit compilation of any call that has >30 positional arguments and at least one keyword argument.

import numba

import pytensor.link.numba.dispatch  # noqa: F401  -- installs the peephole patch

N = 40  # CPython emits CALL_FUNCTION_EX once a call has >30 positional argsparams = ", ".join(f"a{i}" for i in range(N))
ns = {}
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",
    ns,
)
ns["jitted_callee"] = numba.njit(ns["callee"])

print(numba.njit(ns["caller"])(*[float(i) for i in range(N)]))  # expected: 1.0

  File "numba/core/interpreter.py", line 714, in peep_hole_call_function_ex_to_call_function_kw
    args = _call_function_ex_replace_args_large(...)
  File "numba/core/interpreter.py", line 428, in _call_function_ex_replace_args_large
    raise UnsupportedBytecodeError(errmsg)
numba.core.errors.UnsupportedBytecodeError:
CALL_FUNCTION_EX with **kwargs not supported.
...

You need to put the patch in the right location or smth

---

We also need to see if this patch is compatible with all numba versions in our supported range.
I'm weary of patching numba like this by just importing pytensor, although with unique_ids we have a precedent.

[velochy]

Confirmed, thanks — the collapse left the tuple behind a single-use Var alias, and the CALL_FUNCTION_EX peephole only accepts a vararg defined by a build_tuple directly when kwargs are present. Fixed by forwarding the collapsed tuple through such aliases; your repro compiles (prints 1.0) and is added as a regression test. The numba-side PR doesn't have this bug — it writes the tuple into the final variable to begin with, and numba's own test_large_args_small_kws/test_large_args_large_kws cover exactly this shape (they were the failing CI there until yesterday's fix).

On version range: verified standalone against numba 0.58.1 (py3.11), 0.61.2, 0.65.1 and 0.66.0 — i.e. both ends and the middle of our >=0.58,<=0.66 range. The peephole and the concat-chain shape it emits are unchanged across that range, and the collapse only rewrites the exact pattern it matches — if a future numba changes the shape, it finds nothing and is a no-op.

On the wariness: agreed it's not free — this follows the same import-side-effect mold as _patch_pointer_add (numba#10605), and like it, it's self-retiring: once the numba floor includes numba/numba#10782 the module gets deleted.

[ricardoV94]

pointer add is already merged upstream so it's a backport mostly. The unique_ids they acknowledge it's a limitation but are still looking for how to tackle.

Let's give them a few days to review your PR over there to see if they don't have any qualms over the approach.

Maybe there's a less intrusive approach. Your original one of changing our codegen was more limited but less intrusive. Although it touched other backends as they share the logic (easy to fix that though)

Otherwise the gains make sense.