Numba: fix quadratic compile time/memory for graphs with wide (>30-input) nodes - #2354
Numba: fix quadratic compile time/memory for graphs with wide (>30-input) nodes#2354velochy wants to merge 5 commits into
Conversation
Calls (and tuple displays) with >30 items compile to LIST_APPEND + CALL_FUNCTION_EX bytecode, which numba lowers as incremental tuple concatenation: O(n^2) LLVM IR with a large constant per wide call. - fgraph_to_python passes >30-arg calls as <=30-sized tuples through a chunked njit wrapper (safety net for unsplittable wide nodes e.g. Scan) - numba-only rewrites split wide Add/Mul (by prospective fused arity) and Join into balanced trees before inplace_elemwise, so gather-fusion kernels stay <=30 outer inputs and sibling nodes share numba signatures Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
mypy flagged 'None not callable': max_call_args and chunked_call_wrapper_fn were coupled but independently optional. The wrapper factory now enables the behaviour and max_call_args is a plain int with a default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
See #1971 although I found it to be a wash in the pymc-model catalogue, although the motivating issue shows it can matter. |
| # --- Wide-arity splitting (numba) ------------------------------------------- | ||
| # Calls (and tuple displays) with >30 items compile to LIST_APPEND + | ||
| # CALL_FUNCTION_EX bytecode, which numba lowers quadratically in the argument | ||
| # count. Split wide associative Elemwise and Join nodes into balanced trees so | ||
| # that no call — including the idx arrays gather-fusion may later add to an | ||
| # elemwise — crosses that threshold. Small chunks are deliberate: sibling tree | ||
| # nodes share a numba signature, so they compile once and hit the cache. |
There was a problem hiding this comment.
please no global level comments like this.
There was a problem hiding this comment.
Folded into the rewrite docstrings; the registration position is now explained at the optdb.register call instead.
| from pytensor.tensor.basic import Join, join | ||
|
|
||
| if not isinstance(node.op, Join): | ||
| return None |
There was a problem hiding this comment.
dumbest ever rewrite code I've seen. why is Join not in tracks?
There was a problem hiding this comment.
Fixed — tracks=[Join]. The None tracking was left over from an earlier iteration, no excuse.
| if len(tensors) <= _JOIN_SPLIT_THRESHOLD: | ||
| return None | ||
| axis = node.op.axis | ||
| new_out = _tree_reduce(lambda items: join(axis, *items), tensors, _JOIN_CHUNK) |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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.
| "split_wide_add_mul", | ||
| dfs_rewriter(split_wide_add_mul), | ||
| "numba", | ||
| # after AddDestroyHandler (49.5), before inplace_elemwise (50.5) |
There was a problem hiding this comment.
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.
… comment - Fix split_wide_add_mul non-termination on multi-index gathers: each such gathered input adds more than one idx array to the prospective fused arity, so the fixed half-chunk could leave chunks over the threshold and -- when no chunking reduces the node -- _tree_reduce rebuilt an equal-arity clone that the walking rewriter revisited forever. The chunk now scales with the measured extra arity and the rewrite declines when it cannot reduce the node (the chunked-call codegen keeps such call sites legal). A 142-RV production model spent its entire 40+ minute compile in this loop (~150k firings); it now compiles. - Replace the tree-of-Joins split with a serial set_subtensor rewrite into one preallocated buffer (review suggestion): a Join tree re-copies earlier chunks' contents at every level, while the chain writes each input exactly once and collapses into inplace updates of a single buffer after the inplace pass (verified: 50-input join compiles to one AllocEmpty + 50 inplace IncSubtensor). Also gives the rewrite proper tracks ([Join]) instead of tracking None. - Fold the section banner comment into the rewrite docstrings and explain the registration position where it is chosen.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2 participants
Problem
Numba compile time and peak RSS grow quadratically in the arity of graph nodes,
with a hard cliff at 31 inputs. Additive models of the form
mu = sum_i params_i[idx_i]hit this hard: parameters, data and FLOPs all growlinearly with the number of terms n, but canonicalization flattens the sum into
one n-ary
Add, gather-fusion doubles its outer arity, and gradient assemblyJoins n pieces — so compile cost explodes while the model itself stays cheap.Measured on the minimal repro (n independent
(8, 20)Normal blocks gatheredinto a softmax likelihood; nutpie
compile_pymc_model, cold caches, 1-core box):A production PyMC model with a few hundred additive terms reached 29 GB of
compile-time RSS.
Root cause
CPython emits
BUILD_LIST+LIST_APPEND×m +CALL_INTRINSIC_1+CALL_FUNCTION_EXbytecode for any call — or tuple display — with more than 30items (
STACK_USE_GUIDELINE). Numba lowers this pattern as incremental tupleconcatenation: it materializes a tuple of every prefix length 1..m, each step a
full LLVM aggregate copy (7 IR fields per array, with the complete tuple type
spelled out per line) plus NRT incref/decref churn. One wide call therefore
produces O(m²) LLVM IR text with a large constant; LLVM parse + optimization of
that text dominates compile time (~70% of wall time is spent under the LLVM
lock) and the in-memory module accounts for the RSS blowup.
Pure-numba demonstration (no pytensor): a jitted caller passing m array
arguments to a jitted callee produces a 0.071 MB module at m=30 and a 1.478 MB
module at m=31 (21× discontinuity), reaching 37.6 MB at m=160.
Pre-optimization IR of the generated
numba_funcified_fgraphfor the repro:11.3 MB at n=20 → 52.7 MB at n=40 (4.65× per doubling). Post-optimization IR is
near-linear — refprune/DCE clean it up, but the cost has been paid by then.
Both call sites and callees are affected: the fgraph function's wide calls, and
*argsimplementations typed with an m-tuple (np.concatenateinJoin,fused_elemwise_fn), whose lowering unrolls with whole-tuple aggregate copies.Changes
Two coordinated pieces; each keeps every generated call and tuple display at
or below 30 items (
MAX_CALL_ARGSinlink/utils.py).1.
fgraph_to_python: chunked calls for wide nodes (link/utils.py)Generated calls with more than 30 arguments pass their arguments as ≤30-sized
tuples. The chunk layout is computed by a shared helper
(
call_arg_chunk_sizes); past 900 arguments the chunks grow instead, so thecall itself stays under the threshold (mildly superlinear, never quadratic).
The chunk tuples go through a small njit wrapper (built by the numba dispatch)
that concatenates and star-calls the real function. This is the generic safety
net: it handles wide nodes that cannot be split at graph level —
Scan, wideComposites, etc. (a real catalogue model,bayesian_var_hierarchical, has a108-input node that only this path covers).
2. Split wide
Add/Mul; rewrite wideJoinas serialset_subtensor(tensor/rewriting/fused_elemwise.py)Two numba-only rewrites, triggered only when the threshold would actually be
crossed:
Add/MulElemwise nodes split into a balanced tree when theirprospective fused arity — inputs plus the idx arrays gather-fusion would
later absorb — exceeds 28. The chunk size scales with the measured extra
arity (multi-index gathers contribute more than one idx array per input),
and the rewrite declines when no chunking can reduce the node — rebuilding
an equal-arity clone would make the walking rewriter revisit it forever, a
non-termination a 142-RV production model actually hit (its entire 40+
minute "compile" was ~150k firings of this rewrite; it now compiles). Such
irreducible nodes are left wide: fusion declines oversized kernels and the
chunked-call codegen keeps the call site legal.
Joinwith more than 30 tensors is rewritten as serialset_subtensorinto one preallocated buffer (per review: a tree of narrower
Joins wouldre-copy earlier chunks' contents at every level). Each input's contents are
written exactly once, and after the inplace pass the chain collapses into
destructive updates of the single buffer — a 50-input join compiles to one
AllocEmpty+ 50 inplaceIncSubtensor. This also removes the wide-tuplenp.concatenatecallee from such graphs entirely.2. Split wide
Add/MulandJoininto trees (tensor/rewriting/fused_elemwise.py)Two numba-only rewrites split wide nodes into balanced trees, triggered only
when the threshold would actually be crossed:
Add/MulElemwise nodes split when their prospective fused arity —inputs plus the idx arrays gather-fusion would later absorb — exceeds 28;
chunk 14 with gathers present, 28 without, so
FusedElemwisekernels stay≤30 outer inputs.
Joinsplits when it has more than 30 tensors (the axis is an op parameter,so a 30-tensor join call is exactly at the bytecode threshold); chunk 16.
Small chunks are deliberate: sibling tree nodes share a numba signature, so
they compile once and every other sibling hits the compilation cache. An
alternative implementation was measured and rejected — a bespoke generated
wide-
Join(chunk-tuple parameters, unrolled copies into a preallocatedoutput, single data copy) produced strictly worse compile times than the tree
(63.2s vs 47.7s on the n=80 repro) because its one big single-use body enjoys
no signature reuse, while the tree's second data copy was not measurable in
per-call runtime.
Both rewrites register at optdb position ~50: after
AddDestroyHandler(49.5),before
inplace_elemwise(50.5) — so tree nodes still receive inplacetreatment — and well before the gather-fusion database (100).
Results
Repro (nutpie, cold caches; same 1-core box):
(Same-day interleaved base/patch runs confirm the n=80 figures: 372s vs
47-48s.)
Numeric parity: logp diff 5e-10, gradient rel. diff 3e-12 at n=40 (float64
reassociation from the add tree).
On current
main(4ab7498c0, i.e. including the three numba commitsmerged 08-14), same repro, interleaved before/after, 2 rounds — the pathology is
unaffected by those commits and the patch removes it:
A/B across the pymc-model-catalogue (200 models, their own
build_logp_fn,mode=NUMBA, cold caches, 3 reps each in default and frozen form, 1678 runs):
compile, RSS and runtime alike in both forms.
bayesian_var_hierarchical(46-way join, 32-ary adds, a 108-input nodeonly the chunked wrapper can handle): compile 192.3s → 154.5s (−20%),
peak RSS 2059 → 1444 MB (−30%), runtime 168.5 → 175.3 µs (+4%; the
frozen form of the same model instead shows −2.8%).
CFA_SEM_indirect: compile unchanged, peak RSS −4.5%, runtime withinnoise. Its 17-way join is deliberately not split — below the threshold.
sr18_missing_data_primates,longitudinal_external_polynomial_gender:compile/RSS −0.1..−2.6%, runtime within noise.
Interpretation of the runtime column: on this 1-core box, models whose graphs
are byte-identical between the two variants still show up to ±9% (default form)
swings run-to-run, so only the
bayesian_var_hierarchical+4% (tight repeats,non-overlapping ranges) is above the noise. It buys a 20%/30% compile/memory
reduction on a model that takes >3 minutes to compile.
Caveats / open questions
chunked codegen changes generated source without changing the fgraph key, so
warm on-disk caches from before the patch would be stale.
function still grows ~n^1.6 (dominant beyond n≈300; 351s of 594s at n=640).
That term is redundant-LLVM-work-shaped and is the territory of
Prevent compilation when
no_cpython_wrapperis set and restructure linking IR modules numba/numba#9566 (deferred codegen / linking restructure); a pytensor-sidealternative is segmenting the generated fgraph function, not attempted here.
gathered/plain inputs; per-chunk arity accounting could squeeze out slightly
larger fused kernels.
LIST_APPEND-chain lowering deserves an upstreamissue with the pure-numba repro (collapsing the chain into one
build_tuplewould fix all frameworks); happy to file it alongside this PR.
Tests
Existing suite, run on
origin/mainand on this branch with cold caches:tests/link/numba/{test_basic,test_tensor_basic,test_elemwise,test_fused_elemwise,test_scan}.py,tests/link/test_utils.py,tests/tensor/rewriting/{test_elemwise,test_numba}.py— 480 passed, 32 xfailed on both, identical sets.
New tests to add with this PR:
axis, ragged sizes along the join axis, non-contiguous inputs), narrow
fallback, gradient through a wide join.
the prospective fused arity is ≤28, split when above, and idx arrays counted
for nodes feeding gather-fusion.
dtypes), asserting results match the unsplit graph.
Related
inline="always"for trivial ops): orthogonal —measured no effect on this pathology. Note numba: inline="always" to speedup trivial op compilation #2111 as written crashes on >30-arg
n-ary
AddunderFusedElemwise("Calling a closure with *args isunsupported"); the ≤30-arg guard it applies to
makevectorneeds to beextended to
numba_funcify_Add/Mul/genericScalarOp.no_cpython_wrapperis set and restructure linking IR modules numba/numba#9566: complementary; addresses redundant per-function LLVMcodegen, which matches the residual ~n^1.6 term left after this fix.
🤖 Generated with Claude Code