Numba: generate Join with named parameters and slice writes - #2362
Numba: generate Join with named parameters and slice writes#2362velochy wants to merge 1 commit into
Conversation
| post = ", :" * (ndim - ax - 1) | ||
| names = [f"x{i}" for i in range(len(node.inputs))] | ||
|
|
||
| total_src = "\n".join(f" total += {name}.shape[{ax}]" for name in names[1:]) |
There was a problem hiding this comment.
can you use our string codegen helpers? makes it a tiny bit more readable
There was a problem hiding this comment.
Done — rewritten with CODE_TOKEN/build_source_code.
|
|
||
| return join | ||
| cache_version = 3 | ||
| return numba_basic.numba_njit(join_fn), cache_version |
There was a problem hiding this comment.
checkbounds=False iff the function already checks shapes are correct
There was a problem hiding this comment.
Added boundscheck=False — the slice writes are within the output allocated from the validated shapes by construction.
np.concatenate typed with an n-tuple lowers as O(n^2) LLVM IR; one named parameter and one slice-write per input into a preallocated output is linear, with the same single copy per input. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
You were right about the argument names, so I removed them: the signature is back to |
| ) | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("n_terms, axis", [(35, 0), (35, 1), (35, -1)]) |
There was a problem hiding this comment.
are the new tests relevant now? Sounds like no?
2 participants
Join's numba implementation wasnp.concatenateon a*tensorstuple; numba types and lowers an implementation that receives its inputs as one wide tuple quadratically in the input count, and >30-input joins are routine (gradient assembly concatenates one piece per parameter).The generated implementation takes one named parameter per input and writes each into its slice of a preallocated output — linear codegen, same single copy per input. Split off from #2354 per review there (a dispatch-level Join fix instead of a graph rewrite).
Results
Interleaved with
mainin one session on an otherwise-idle 1-core box, cold caches (absolute times here are load-sensitive, so only same-session comparisons are quoted), n=80 additive repro whose gradient assembly is a 60-tensor join:The two are complementary: this one removes the wide-tuple callee (
np.concatenatetyped with a 60-tuple) and carries the compile-time win, while #2361 removes the caller-side bytecode chain and carries the memory win.