Specialize scalar graphs in numba backend - #2310
Draft
ricardoV94 wants to merge 10 commits into
Draft
Conversation
ricardoV94
commented
Jul 27, 2026
ricardoV94
commented
Member
- Add Numba scalarization rewriter (CAReduce producer, Join consumer)
- Numba: accept ScalarType FusedElemwise inputs via constant_inputs
- Allow OpFromGraph to take a pre-built frozen inner fgraph
- Numba: scalarize scalar-index Subtensor to ScalarSubtensor
- Numba: scalarize mixed Elemwise via ScalarizedElemwise
- Numba: scalarize multi-output size-1 Composites
- Numba: squeeze the fused full-reduction epilogue instead of ravel
- Numba: hand fully-reduced FusedElemwise outputs out as scalars
- Numba: keep full-reduction accumulators on the stack
- Numba: squeeze partial fused reductions with as_strided, not reshape
A self-contained graph rewriter that propagates ScalarType so values that are logically single scalars are not boxed into 0-d arrays, each of which costs one heap allocation per call. Producer and consumer are decided independently and matched: one toposort collects the values a consumer would take as a scalar, realises each producible one as tensor_from_scalar(scalar) (type-preserving, shared), then lets consumers bypass the box; an unbypassed box DCEs away or survives as a single shared re-box, so it is never a loss. Covered so far: CAReduce (producer) and Join (consumer). Scalar ops are OpFromGraphs wrapping the faithful tensor graph so non-Numba backends run them via perform; Numba gets an allocation-free dispatch: - ScalarCAReduce reuses the CAReduce reducer with scalar_out=True (returns the accumulator register instead of np.array). - ScalarJoin (and the 1-d Join dispatch) allocate-and-write element-wise, admitting scalar entries directly; this matches np.concatenate for arrays and can take a scalar, which np.concatenate cannot. - Subtensor gains scalar_out (return z vs np.asarray(z)) as groundwork for the ScalarSubtensor producer. tests/benchmarks/test_scalar_ops.py adds a logp_dlogp-shaped validation graph (scalar params + one vector param, reduced to a scalar logp, gradient joined) checked against a numpy value and finite-difference gradient.
Codegen support for handing a loop-invariant FusedElemwise input over as a scalar instead of boxing it in a 0-d array. A ScalarType input carries one value for the whole iteration space, so it goes through _vectorized's constant_inputs (prepended once, never reloaded per iteration) rather than being materialised as an array. - store_core_outputs permutes its signature to receive const inputs first while calling the core op in the original order (no inner-fgraph reorder). - _build_reduce_impl_src splits outer_inputs into the const tuple and the array tuple by static index. - numba_funcify_FusedElemwise computes const_positions (ScalarType inputs) and remaps input_bc_patterns / indexed_inputs / inplace_pattern onto the array-only positions. Routes through the generated-source impl when there are const inputs (the inline closure cannot split by static index). Inert until a rewrite produces ScalarType FusedElemwise inputs: with none, const_positions is empty and every path is unchanged. Fused-elemwise, elemwise and reduction Numba suites green.
Add a `replace=` argument to `construct_nominal_fgraph` that folds an extra substitution into the single nominalizing clone, and an `fgraph=` argument to `OpFromGraph` that accepts an already-nominalized frozen inner graph. Together these let a caller bind a real (possibly shared, multi-client) value at `make_node` -- e.g. rewrite a tensor input into `tensor_from_scalar` of a scalar -- without a second clone or the graph_replace identity trap. Both default to the existing behaviour.
Restructure the scalarize pass as a single forward boundary walk and add a Subtensor producer rule: a scalar-index `Subtensor` (0-d output) becomes a `ScalarSubtensor` OpFromGraph that returns the register value directly instead of `np.asarray`-boxing it, distributed to each consumer that can swallow a scalar. The Numba dispatch reuses the default Subtensor codegen with `scalar_out=True`.
A size-1 or larger Elemwise fed by a mix of scalar-behind and genuine-array inputs becomes a ScalarizedElemwise OpFromGraph that carries the scalar inputs as ScalarType (loop-invariant constant_inputs) while keeping an array output, so their 0-d boxes vanish. numba_funcify_Elemwise is shared between Elemwise and ScalarizedElemwise. FuseElemwise anchors on it and, once it finds a fusion, inlines it back to the tensor_from_scalar form so the existing swallow carries the scalars into the fused loop as const -- and that swallow pulls loop-invariant TensorFromScalar FusedElemwise inputs into const_positions via the OpFromGraph fgraph= path.
A Composite whose outputs are all size-1 is a pure scalar computation regardless of how many outputs it has, so it never needs its scalar inputs boxed. Extend the Elemwise rule to fire on multi-output nodes when every output is broadcastable: all-scalar inputs apply the scalar op and box each output once (_emit_scalar_elemwise); a mix keeps ScalarizedElemwise. This catches the final logp+dlogp accumulation Composite, whose scalar inputs were the last surviving boxes.
The fused full-reduction post-step boxed its single accumulated value with `np.array(raw.ravel()[0])` -- `ravel()` can force a contiguous copy of the strided keepdims buffer. Emit the same strided 0-d view a DimShuffle squeeze uses (`as_strided(raw, shape=(), strides=())`) instead: a view, no copy, cast only when the accumulator dtype differs from the output.
A full reduction yields a single value, but a fused reduced output was emitted as a 0-d array -- one allocation per reduction purely to carry it to a consuming scalar. FuseElemwise now hands every full-reduction fused output out as ScalarType (scalar_from_tensor) and immediately re-boxes it with tensor_from_scalar at the replacement, so the graph stays valid and net-unchanged. The local_scalar_tensor_scalar / local_tensor_scalar_tensor rewrites (registered right after fusion) then cancel the boundary wherever a consumer takes the scalar; where one wants the tensor the single box stays, so it is never a loss. wrap_scalar_careduce does the same for a full CAReduce fusion did not absorb whose only clients are ScalarFromTensor. The Numba epilogue returns the value via .item() for a ScalarType output.
make_outputs allocated every vectorized output with _empty_nd_impl. For a fully-reduced ScalarType output the batch shape is all ones, so that is an NRT heap allocation, an identity memset and a free -- to carry a single value the epilogue immediately reads back out. Put those buffers on the stack instead (_make_stack_array via alloca_once): it costs nothing and, not escaping, LLVM promotes it to a register. Only safe now that these outputs are ScalarType. The reduce spec gains a per-output on_stack flag, set when every axis is reduced and the output is scalar; partial reductions and 0-d tensor outputs keep the heap buffer. Other _vectorized callers pass an empty reduce spec and are unaffected.
The full-reduction epilogue already squeezes its 0-d result with as_strided; the
partial-reduction post-step still went through {src}.reshape(...), which compiles
to a runtime numba_attempt_nocopy_reshape call and only works on contiguous input.
Squeeze the size-1 reduced axes with the same as_strided stride arithmetic
DimShuffle uses instead -- no runtime call, valid on non-contiguous buffers.
Performance-neutral (removes a call and refcounts, no allocations).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.