Numba: rework the vectorize intrinsic interface to fix wide-node compile blowup

[velochy]

Problem

Nodes with more than ~30 inputs (wide fused kernels from additive models, wide scalar Composites from gradient fusion) compiled quadratically in numba, in both time and memory — a 62 GB / killed-by-OOM real-model report in pymc-devs/nutpie#339. Two dispatch-level mechanisms were responsible. Star-args boundaries type as one wide tuple, whose LLVM lowering is O(n²) in the arity: the _vectorized intrinsic's inputs tuple and the Elemwise/FusedElemwise overload targets and impls. And the per-node metadata (broadcast patterns, dtypes, index specs) was passed to the intrinsic as base64-pickled literal-string arguments, which numba's type inference re-hashes and re-decodes on every fixpoint sweep — profiling showed this dominating compile time outright on wide graphs.

Changes

All boundaries now use named parameters, with the outer inputs as trailing individual arguments of a per-node intrinsic instance that closes over its decoded metadata — only runtime values remain in typed signatures. Graphs are untouched. Pairs with #2361 (or numba/numba#10782 once released) for the >30-argument call sites in generated fgraph functions; without it results are unchanged but wide calls still pay the bytecode-chain cost.

Results

Interleaved with main in a single session on an otherwise-idle 1-core box, cold caches (absolute times on this box are load-sensitive, so only same-session interleaved comparisons are quoted; measured together with #2361 and #2362, since the wide-node bottlenecks compose and any single one left unfixed caps the win):

	main	this stack
n=80 additive repro	364.2 / 365.3 s, 2632 MB	105.1 / 105.0 s, 935 / 929 MB
nutpie#339 sentinel, lognormal 18/24	196.5 s, 3.62 GB	88.8 s, 0.76 GB
same, gamma 18/24	137.8 s, 1.17 GB	89.4 s, 0.84 GB

That is 3.5× compile and 2.8× peak RSS on the repro, and on the sentinel it removes the pathology the issue is actually about: the lognormal/gamma disparity goes from 1.4× time / 3.1× memory to 1.0× / 0.9×. logp/grad parity exact on the sentinel; full numba test battery green with cold caches.

[ricardoV94]

See #1971 although I found it to be a wash in the pymc-model catalogue, although the motivating issue shows it can matter.

[ricardoV94]

please no global level comments like this.

[velochy]

Folded into the rewrite docstrings; the registration position is now explained at the optdb.register call instead.

[ricardoV94]

dumbest ever rewrite code I've seen. why is Join not in tracks?

[velochy]

Fixed — tracks=[Join]. The None tracking was left over from an earlier iteration, no excuse.

[ricardoV94]

similar to the wanting to associated add/mul correctly, you also want to think how you associated join, as every join require a full copy of all the inputs. You can easily create bad graphs where you are copying 90% of the contents multiple times to add a scalar at a time. If anything I'd suggesting thinking about a serial set_subtensor rewrite on empty. See #2014 which was more interested in reducing allocations/perfromance, not on what you're at here obviously

[velochy]

You're right that the tree re-copies earlier chunks at every level. Replaced with your suggestion: serial set_subtensor into one preallocated buffer — each input written exactly once, and the chain collapses to inplace updates of a single allocation after the inplace pass (a 50-input join compiles to one AllocEmpty + 50 inplace IncSubtensor). It also removes the wide-tuple np.concatenate callee, which was the other half of the Join cost. #2014's allocation-reduction angle would compose with this but isn't needed for the compile-time goal here.

[ricardoV94]

why between the two?

[velochy]

It must run after the last Add-canonicalization — local_add_canonizer re-flattens the tree, so registering earlier ping-pongs — and before inplace_elemwise (50.5) so the tree/chain nodes get inplace treatment; the gather-fusion db (100) then only sees narrow nodes. Now stated at the registration site. The "after AddDestroyHandler (49.5)" phrasing in the old comment was incidental, not a dependency.

[ricardoV94]

didn't you fix this by attacking numba interpretation directly in #2361 ?

For the join, can't we have a better numba dispatch instead. For instance I implemented one that allowed working with scalars directly for improving perf of scalar ops in #2310

https://github.com/ricardoV94/pytensor/blob/cb8e0e0c33192d2370a3e40931ab6a856011f4b2/pytensor/link/numba/dispatch/tensor_basic.py

[velochy]

#2361 fixes only the call sites (the caller-side prefix-tuple chain). The callee cost is separate: an implementation typed with a wide tuple (np.concatenate's
tuple argument, fused_elemwise_fn's star-args) lowers each element access with the full tuple type spelled out, so it stays quadratic with #2361 alone. Measured on the n=80 additive repro (nutpie, cold caches): main 480 s / 2.6 GB, #2361 alone 360 s / 1.7 GB, this PR alone 83 s / 0.97 GB, both 67 s / 0.97 GB — the two attack different quadratics and this one carries most of the win.

On Join you're right that a dispatch-level fix is possible: a generated per-arity impl with individual parameters (the #2310 style) would be linear with the same copy count, and would avoid the node-count blowup the serial set_subtensor chain causes (~2× nodes on graphs with a wide join). The rewrite's advantage is only that it reuses the existing inplace machinery instead of adding codegen. Happy to replace the Join rewrite with a numba dispatch along the #2310 lines if you prefer that shape — the Add/Mul split can't move to dispatch the same way, since it's the fused kernel's own signature that must stay narrow.

[ricardoV94]

Join is a better representation then nested set subtensor and the inplace doesn't matter. In either case you allocated an empty at the start and then write the values without copying intermediates. A custom join dispatch would never do this while nested set subtensor is at the mercy of the optimizer to get right.

The add/mul split I'm not sold on that yet. In fact we want to flatten more associative commutative Ops like And, Or...

You're also targeting this but I don't see what's specific about it. Fusion would gladly take the same number of args on a Composite and the problem would be the same. What's unique about add or mul other than it showed in your model?

If you want to split it should be in a form that has at least computational motivation like the PR I mentioned.

Otherwise if it's really the scalar add or mul you can tweak the dispatch. Right now it does (a + b + ... c + d), you can consider (a + b + ...) + (... + c + d) if it helps or even a kahman pairing or something like that as it doesn't change the meaning of the scalar add.

[velochy]

Reworked per review: the graph rewrites are gone, everything is now at the dispatch layer, and it turned out better than the rewrites on every measurement. Two smaller pieces are split off for separate review: #2362 (Join dispatch with named parameters and slice writes — your suggestion) and #2363 (pairwise Add/Mul — also your suggestion). This PR keeps the core piece: the vectorize intrinsic interface rework. #2361 remains the call-site companion.

Two findings worth flagging from the profiling that got here. First, splitting wide Add/Mul turned out unnecessary: after fixing the boundaries, the residual cost was not LLVM on big kernels but numba type inference re-hashing and re-unpickling the encoded-literal metadata arguments on every fixpoint sweep — moving the metadata into per-node intrinsic closures fixed it, and dispatch-only now beats the old rewrite branch (n=80 repro: 77 s / 0.84 GB vs 67 s / 0.97 GB; nutpie#339 sentinel lognormal: 58 s / 0.70 GB vs 61 s / 0.77 GB; main: 480 s / 2.6 GB and 200 s / 3.6 GB). Second, the generic caching wrapper in dispatch/basic.py wrapped every funcified op in def op(*args), which re-tupled every node call regardless of what the op's own dispatch did — worth knowing about for any future wide-signature work.

[ricardoV94]

One can't avoid string codegen in numba no matter how hard we try hein

[ricardoV94]

Are you sure about this? Doesn't numba just passthrough without deconstructing anything? Does it depend on numba version? And above 30 arguments wouldn't matter either way right? Because it still triggers CPython tuple thing?

[velochy]

You were right — dropped, basic.py is no longer touched by this PR.

Measurements, since the questions are worth answering properly. Pass-through: at runtime yes, nothing survives — post-optimization IR is 0.884 MB (star) vs 0.910 MB (named) for an 80-argument call chain. At compile time no: the caller has to emit the tuple (an insertvalue chain, 132 KB of IR) and the wrapper to unpack it (extractvalue, 115 KB), so pre-opt IR is 3.39 MB vs 2.67 MB — LLVM parses and then deletes it. Version: not version-dependent, byte-identical IR sizes on numba 0.58.1, 0.61.2 and 0.65.1. >30 arguments: exactly your point, and it cuts deeper than 'no difference' — the wrapper's own call re-triggers the CPython path, so unpatched the named form is worse (5.50 MB vs 3.64 MB at 40 args, paying the chain twice instead of once); it only comes out ahead with #2361 (2.67 vs 3.39 MB at 80).

And on the real model none of it matters: interleaved A/B/A/B on the n=80 repro, cold caches, named 106.9 / 106.5 s vs star 106.9 / 106.7 s, RSS within 1.5%. A 21% pre-opt IR reduction on one wrapper doesn't survive to whole-model compile time, so the special case wasn't earning the kwargs.get("node") branch.

[velochy]

Real Margus here: I'm re-running on my real models to measure the diff with them, as AI's comment on "on the real model" still talks about a toy example. But yes, if that also comes out a wash, it makes sense to close this PR as 2361+2362 do all of the real work.

[velochy]

Also corrected the Results section: the earlier '77 s / 0.84 GB' figure did not reproduce. This box is single-core and shared with a desktop session, so absolute compile times move with ambient load, and that number came from a different session than the main baseline it was compared against — my mistake. Re-measured interleaved with main in one session: n=80 repro 364 s / 2.6 GB → 105 s / 0.93 GB (3.5× / 2.8×), nutpie#339 sentinel lognormal 197 s / 3.62 GB → 89 s / 0.76 GB, and the lognormal-vs-gamma disparity that issue reports goes from 1.4× time / 3.1× memory to parity. Smaller headline than before, same conclusion.

The A/B ablations quoted in this thread were all interleaved within a session, so they are unaffected.

[velochy]

Real Margus: yep, no real performance win, it seems. Everything within noise on this PR, assuming #2361 and #2362 are both present. Feel free to close, unless you want to salvage some of it for cleaner code purposes.

[ricardoV94]

I thought this PR was trying an alternative to #2361 that didn't require the hacking of numba

[velochy]

I thought this PR was trying an alternative to #2361 that didn't require the hacking of numba

No not really:

configuration	compile	peak RSS
main	364.9 s	2696 MB
#2362 alone	158.9 / 152.3 s	2658 MB
#2354 + #2362	186.5 / 185.2 s	2692 MB
#2361 + #2362	106.6 / 105.8 s	929 / 942 MB
all three	105.1 s	935 MB

Also: I understood you were mostly ok with #2361 at least until they merge the fix themselves?

[ricardoV94]

I'm not ok with #2361, not against it either, rather apprehensive. Would really like to get a positive reaction from the numba guys first. It touches internals of numba, including outside of pytensor use.

[ricardoV94]

But now I'm also confused about #2362, the argument names thing is the same thing you were doing here no? If it doesn't help here why does it help there? You sure the speedup is not just the slice assignemnt instead of the np.concatenate call?