Fuse Elemwise graphs that have multiple outputs and clients

[ricardoV94]

Continuation of aesara-devs/aesara#1242, as I've been blocked on that repo

More context can be found in aesara-devs/aesara#1237

Clarifying what is the idea behind allowing Composites with multiple clients. The following example:

import pytensor as aesara
import pytensor.tensor as as at

x = at.vector("x")
y = at.exp(x/5)
w = y + 1
z = y * 2

f1 = aesara.function([x], [z, w])
aesara.dprint(f1)

Before this PR it produced this graph:

Elemwise{Mul}[(0, 1)] [id A] 2
 |TensorConstant{(1,) of 2.0} [id B]
 |Elemwise{Composite{exp((i0 * i1))}} [id C] 0
   |TensorConstant{(1,) of 0.2} [id D]
   |x [id E]
Elemwise{add,no_inplace} [id F] 1
 |TensorConstant{(1,) of 1.0} [id G]
 |Elemwise{Composite{exp((i0 * i1))}} [id C] 0

After:

Elemwise{Composite{(i2 * exp((i0 * i1))), (i3 + exp((i0 * i1)))}}.0 [id A] 0
 |TensorConstant{(1,) of 0.2} [id B]
 |x [id C]
 |TensorConstant{(1,) of 2.0} [id D]
 |TensorConstant{(1,) of 1.0} [id E]
Elemwise{Composite{(i2 * exp((i0 * i1))), (i3 + exp((i0 * i1)))}}.1 [id A] 0

The expected savings come from having to iterate only once over the data vector x, vs having to iterate 3 times before, once over x and twice over exp(i0 * i1).

Another example is the following graph:

x = at.dvector("x")
mu = at.dvector("mu")
logp = (- ((x - mu) **2) / 2)
grad = at.grad(logp.sum(), x)
f2 = aesara.function([mu, x], [logp, grad])
aesara.dprint(f2)

Before:

Elemwise{Composite{(i0 * sqr(i1))}}[(0, 1)] [id A] 2
 |TensorConstant{(1,) of -0.5} [id B]
 |Elemwise{sub,no_inplace} [id C] 0
   |x [id D]
   |mu [id E]
Elemwise{neg,no_inplace} [id F] 1
 |Elemwise{sub,no_inplace} [id C] 0

After:

Elemwise{Composite{(-(i0 - i1)), (i2 * sqr((i0 - i1)))}}.1 [id A] 0
 |x [id B]
 |mu [id C]
 |TensorConstant{(1,) of -0.5} [id D]
Elemwise{Composite{(-(i0 - i1)), (i2 * sqr((i0 - i1)))}}.0 [id A] 0

Again we replace 3 loops by 1.

For more detail, here is the C code of the Composite scalar in the first function:

// function 1
{
npy_float64 V13_tmp1;
// i0 * i1
V13_tmp1 = V11_i * V9_i;

npy_float64 V13_tmp2;
// exp(i0 * i1)
V13_tmp2 = exp((npy_float64)V13_tmp1);

// First output
// i2 * exp(i0 * i1)
V1_i = V7_i * V13_tmp2;

// Second output
// i3 + exp(i0 * i1)
V3_i = V5_i + V13_tmp2;
}

// function 2
{
npy_float64 V11_tmp1;
// i0 - i1
V11_tmp1 = V9_i - V7_i;

// First output
// - (i0 - i1)
V1_i = -V11_tmp1;

npy_float64 V11_tmp2;
// sqr(i0 - i1)
V11_tmp2 = V11_tmp1 * V11_tmp1;

// Second output
// i2 * sqr(i0 - i1)
V3_i = V5_i * V11_tmp2;
}

You can see that it avoids recomputing the same sub-expressions in both cases. It does store more intermediate variables than needed, but hopefully the compiler will be able to get rid of those.

This would be further improved by inplacing the scalars as is being pursued in #107

TODO:

• Wait for Numba to allow multi-output Elemwise, or figure out how to disable rewrite in that backend
• Figure out what to do with inplace optimization when there are multiple outputs
  • Deferred as a separate issue: Implement specialized inplace rewriter for Elemwise Composite Ops #138
• Refactor the specialized add_mul_fusion
• Do not split fusion ops depending on c-code presence in non-c backends
  • Similarly, get rid of fragile elemwise_max_input_fct
• Improve Op string, it now looks like a monster with many outputs
• Profile expected performance gains! # Will do after 🔄 From Aesara: #1347 and # 1365: " Add CI support for benchmarking" #139

New features

• Fuse Elemwise graphs that have multiple outputs and non-fuseable clients

Bugfixes

• Fix bug when creating Composite with multiple identical outputs

[ricardoV94]

@OriolAbril Any idea why readthedocs is failing?

[OriolAbril]

I suspect github is laggy or partially down. The error was that readthedocs was unable to find the reference to this PR. I restarted the job and now it has been able to get the PR changes. 🤷🏿

[ricardoV94]

Yup, the slow and skip marks seem to be ignored in the CI #5 (comment)

[ricardoV94]

Alright first time the tests pass (minus the ones I disabled :D)

[ricardoV94]

Benchmark, running this snippet after this PR vs main:

import pytensor
import pytensor.tensor as pt
import numpy as np

rng = np.random.default_rng(123)
size = 100_000
x = pytensor.shared(rng.normal(size=size), name="x")
mu = pytensor.shared(rng.normal(size=size), name="mu")

logp = (- ((x - mu) **2) / 2)
grad = pt.grad(logp.sum(), x)

func = pytensor.function([], [logp, grad])
pytensor.dprint(func)
%timeit -n 1000 func()

This PR: 145 µs ± 15.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
Main: 220 µs ± 27.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Exact numbers change quite a bit when I rerun the script, but they don't overlap. The difference grows with size, and becomes way less noticeable (or vanishes) for smaller sizes.

Would appreciate if someone else could replicate the benchmark.

[aseyboldt]

Timings on my desktop machine give a similar pattern.

[michaelosthege]

The main loop looks much better, but the local functions make it still really hard to read, as it is ~450 lines in total.
Would it be feasible to extract these local functions and pass a few more things as args/kwargs if needed?

Your call though

[michaelosthege]

I don't want to sound nitpicky, but you're adding more lines to the parametrize (74) than the test method has to begin with (61).

I mean sure, there were 600 lines of parameterize before, but like... maybe some of these parametrize items deserve their own test method

[ricardoV94]

This is a rewrite that affects 99.9% of all graphs in a very fundamental manner, I don't think it's crazy that we would test so many parametrizations.

Besides autodiff, loop fusion is basically the other thing that PyTensor does that goes beyond vanilla Numpy

[michaelosthege]

I fully support having all these tests. I'm just questioning whether this is the most maintainable way to organize the test code.
For example, the test implementation could be extracted into a function, and some of the most involved parametrizations that have these 10ish levels of indentation could just be their own test methods.

[ricardoV94]

Hmm... We can also use a generator to yield the parametrizations, but I am not sure that's much better? I went through the parametrizations and there is only one that seems too long. I rearranged it to have less nesting:

            (
                (
                    log(
                        ge(
                            assert_op(
                                at_abs(fx),
                                at_all(ge(at_abs(fx), 0)),
                            ),
                            0,
                        )
                    ),
                ),
                (fx,),
                (fxv,),
                4,
                (np.zeros_like(fxv),),
                ("float32",),
            ),

[ricardoV94]

Ready for re-review

[michaelosthege]

thanks for cleaning up the code some more, @ricardoV94

I know that you invested a lot of attention on the details here. If there is anything where you're in doubt and need a careful review let me know. It's just too much for one pass..

[ricardoV94]

thanks for cleaning up the code some more, @ricardoV94

I know that you invested a lot of attention on the details here. If there is anything where you're in doubt and need a careful review let me know. It's just too much for one pass..

Yeah, it's a though one to review. Thanks for the suggestions to make it a tiny bit more readable. I think I will merge it and iterate in future PRs

[ricardoV94]

Actually still have to clean the comments from one function... BRB