Numba: back-port compilation speedup for large tuple/ number of arguments

[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 8-core/16-thread box (Ryzen 7 PRO 4750U, BLAS/OMP threads pinned to 1), 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.

[velochy]

Dropped the second commit (inlining the collapsed tuple into direct call arguments). Measured on the n=80 repro it is worth 0.2% — 324.2 / 323.6 s without it vs 322.7 / 323.4 s with it — so it was carrying risk for no gain, including the corresponding semantic change upstream. What is left is the collapse itself plus the alias forwarding that fixes your kwargs repro, which is the part that actually matters. Same drop applied to numba/numba#10782, where it also removes the only behaviour change and the edit to numba's own test suite.

[velochy]

@ricardoV94 numba PR seems to have maintainer blessing now:
numba/numba#10782 (review)

[ricardoV94]

Bot summary of my nudging:

Could we use a faithful backport of the merged Numba fix here? Copying the upstream peephole implementation and monkey-patching it should be straightforward, subject to checking compatibility with our supported Numba versions. This custom pass has broader scope and still leaves quadratic tuple chains after starred unpacking. For temporary code, I would prefer the upstream implementation with a couple of focused regression tests, removed together with the patch.

[velochy]

Done — this is now a faithful backport rather than a custom pass. peep_hole_list_to_tuple is copied verbatim from numba master (the #10782 merge, 2026-08-27); the only edit is a leading underscore on the name. I added it to ruff's exclude list in pyproject.toml so it stays byte-identical and diffable against upstream — it otherwise trips T201/UP031/PERF102 on numba's own debug prints and %-formatting, and reformatting it would defeat the point.

Your two technical objections both go away with the upstream implementation: it coalesces runs of appends around extend, so f(x0, ..., x34, *t) no longer leaves a quadratic chain after the starred unpack (added a regression test asserting the number of adds is bounded by the extends rather than the argument count), and its scope is exactly numba's peephole rather than any binop-add of two single-use tuples anywhere in the IR.

On the version question: the patch installs only when the running numba predates the fix, detected by looking for the coalescing in inspect.getsource, so it self-disables rather than needing a version bound. Exercised the vendored function on numba 0.58.1, 0.61.2, 0.65.1 and 0.66.0 — spanning our whole supported range — with the >30-positional case, your >30-positional-plus-kwarg repro, the starred-unpack case and plain tuple unpacking, all correct on each; and on a numba built from master it correctly does not install and results stay correct via numba's own copy. Five focused tests, all deleted along with the module when our minimum numba includes the fix.

Also rebased onto main — the earlier red CI was numpy 2.5 drift fixed by #2376, unrelated to this patch.

[ricardoV94]

Can we use numba versioning check to include or not the patch, instead of checking the code. we know that if larger than today's release it isn't needed, before or equal yes? we should likely do the same to pointer_add (gotta check when was it released).

[velochy]

Switched to version checks on both, using numba.version_info.short (it parses dev versions correctly, e.g. a master build reports (0, 68)).

The release boundaries, checked against numba's tags rather than guessed: the list_to_tuple fix merged 2026-08-27, after 0.67.0 was tagged on 2026-08-10, so it ships in 0.68 — patch applies on < (0, 68). The pointer_add fix (4b8f84539, 2026-07-15) is contained in the 0.67.0 tag, so that patch now applies on < (0, 67) — good catch, it was unconditional before and would have kept overriding a numba that no longer needs it.

Verified the gates at the boundaries by loading both modules against four interpreters: 0.58.1 both patched, 0.66.0 both patched, 0.67.0 pointer_add off / list_to_tuple on, master build both off. Tests still green.