Numba: fuse AdvancedSubtensor->Elemwise->AdvancedIncSubtensor

[ricardoV94]

Summary

Introduce IndexedElemwise, an OpFromGraph that wraps AdvancedSubtensor + Elemwise + AdvancedIncSubtensor subgraphs so the Numba backend can generate a single loop with indirect indexing, avoiding materializing AdvancedSubtensor input arrays, and writing directly on the output buffer, doing the job of AdvancedIncSubtensor in the same loop, without having to loop again through the intermediate elemwise output.

Supports arbitrary-axis indexing, ND index arrays, multi-index groups, broadcast indices, repeated accumulation (inc/set modes), and multi-client outputs via Composite identity duplication.

Also included:

• Extend FusionOptimizer to merge independent sibling subgraphs that share inputs but have no producer-consumer edge (e.g. f(x) and g(x))

Motivation

In hierarchical models with mu = beta[group_idx] * x + ..., the logp+gradient graph combines indexed reads and indexed updates in the same Elemwise (the forward expands group-level parameters via advanced subtensor, and the gradient accumulates back into the source via advanced inc subtensor).

A simple example

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import get_mode

numba_mode = get_mode("NUMBA")
numba_mode_before = numba_mode.excluding("fuse_indexed_elemwise")

x = pt.vector("x")
idx = pt.vector("idx", dtype=int)
value = pt.vector("value")

y = pt.zeros(100)
out = ((x[idx] - value) ** 2).sum()
grad_wrt_x = pt.grad(out, x)
fn_before = pytensor.function([x, value, idx], [out, grad_wrt_x], mode=numba_mode_before, trust_input=True)
fn_before.dprint(print_op_info=True, print_destroy_map=True)
# Sum{axes=None} [id A] 5
#  └─ Composite{...}.0 [id B] d={0: [0]} 1
#     ├─ AdvancedSubtensor1 [id C] 0
#     │  ├─ x [id D]
#     │  └─ idx [id E]
#     └─ value [id F]
# AdvancedIncSubtensor1{inplace,inc} [id G] d={0: [0]} 4
#  ├─ Alloc [id H] 3
#  │  ├─ [0.] [id I]
#  │  └─ Shape_i{0} [id J] 2
#  │     └─ x [id D]
#  ├─ Composite{...}.1 [id B] d={0: [0]} 1
#  │  └─ ···
#  └─ idx [id E]

# Inner graphs:

# Composite{...} [id B] d={0: [0]}
#  ← sqr [id K] 'o0'
#     └─ sub [id L] 't5'
#        ├─ i0 [id M]
#        └─ i1 [id N]
#  ← mul [id O] 'o1'
#     ├─ 2.0 [id P]
#     └─ sub [id L] 't5'
#        └─ ···

fn = pytensor.function([x, value, idx], [out, grad_wrt_x], mode=numba_mode, trust_input=True)
fn.dprint(print_op_info=True, print_destroy_map=True)

# Sum{axes=None} [id A] 3
#  └─ IndexedElemwise{Composite{...}}.0 [id B] d={1: [3]} 2
#     ├─ x [id C] (indexed read (idx_0))
#     ├─ value [id D]
#     ├─ idx [id E] (idx_0)
#     └─ Alloc [id F] 1 (buf_0)
#        ├─ [0.] [id G]
#        └─ Shape_i{0} [id H] 0
#           └─ x [id C]
# IndexedElemwise{Composite{...}}.1 [id B] d={1: [3]} 2 (indexed inc (buf_0, idx_0))
#  └─ ···

# Inner graphs:

# IndexedElemwise{Composite{...}} [id B] d={1: [3]}
#  ← Composite{...}.0 [id I]
#     ├─ AdvancedSubtensor1 [id J]
#     │  ├─ *0-<Vector(float64, shape=(?,))> [id K]
#     │  └─ *2-<Vector(int64, shape=(?,))> [id L]
#     └─ *1-<Vector(float64, shape=(?,))> [id M]
#  ← AdvancedIncSubtensor1{inplace,inc} [id N] d={0: [0]}
#     ├─ *3-<Vector(float64, shape=(?,))> [id O]
#     ├─ Composite{...}.1 [id I]
#     │  └─ ···
#     └─ *2-<Vector(int64, shape=(?,))> [id L]

# Composite{...} [id I]
#  ← sqr [id P] 'o0'
#     └─ sub [id Q] 't0'
#        ├─ i0 [id R]
#        └─ i1 [id S]
#  ← mul [id T] 'o1'
#     ├─ 2.0 [id U]
#     └─ sub [id Q] 't0'
#        └─ ···

x_test = np.arange(15, dtype="float64")
idx_test = np.random.randint(15, size=(10_000,))
value_test = np.random.normal(size=idx_test.shape)

logp_before, dlogp_before = fn_before(x_test, value_test, idx_test)
logp, dlogp = fn(x_test, value_test, idx_test)
np.testing.assert_allclose(logp_before, logp)
np.testing.assert_allclose(dlogp_before, dlogp)

%timeit fn_before(x_test, value_test, idx_test)  # 29.4 μs ± 2.57 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit fn(x_test, value_test, idx_test)  # 13.8 μs ± 136 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

Next step would be to also fuse the sum directly on the elemwise, so we end up with a single loop over the data.

[ricardoV94]

Here is a benchmark experiment with a subset of models where the rewrite fires: https://ricardov94.github.io/pymc-model-catalogue/experiments.html#base=index_elemwise_fusion_curated_base&compare=index_elemwise_fusion_curated

Nice gains in eval time for quite some models

[ricardoV94]

TODO: This is not equivalent, only in +=, but most of the times is what we want. Anyway better to use a custom fn in the store_core_outputs and only overload that

[ricardoV94]

store_inplace, inc_inplace.

[ricardoV94]

todo: check didn't we get rid of inc/set

[jessegrabowski]

we went over this in a call, Ricardo explained it to me and it seems good.