Overhead of function call to get output shape in pushforward-via-pullback

[wsmoses]

using DifferentiationInterface
import Enzyme, Mooncake, ForwardDiff, Zygote, FiniteDiff, BenchmarkTools, Reactant

function f(x)
    A = [x+1 x+2 x+3
         x+4 x+5 x+6
         x+7 x+8 x+9]
    B = [x-1 x-2 x-3
         x-4 x-5 x-6
         x-7 x-8 x-9]
    
    sum(A' * (B .* 1.5))
end

function bench()
    for backend ∈ (AutoEnzyme(mode=Enzyme.Forward),
                   AutoEnzyme(mode=Enzyme.Reverse),
                   AutoMooncake(),
                   AutoMooncakeForward(),
                   AutoForwardDiff(),
                   AutoZygote(),
                   AutoFiniteDiff())
        prep = prepare_derivative(f, backend, 1.0)
        x = Ref(1.0)
        @show backend
        print("  "); @btime derivative(f, $prep, $backend, $x[])
        println()
    end
end

bench()

@btime Enzyme.autodiff(Reverse, f, Active(1.0))

@btime Enzyme.gradient(Reverse, f, 1.0)

[gdalle]

Thanks for opening the issue, I confirm your observations:

Details

using DifferentiationInterface
import Enzyme
using BenchmarkTools

function f(x)
    A = [
        x + 1 x + 2 x + 3
        x + 4 x + 5 x + 6
        x + 7 x + 8 x + 9
    ]
    B = [
        x - 1 x - 2 x - 3
        x - 4 x - 5 x - 6
        x - 7 x - 8 x - 9
    ]

    return sum(A' * (B .* 1.5))
end

back_for = AutoEnzyme(mode = Enzyme.Forward)
back_rev = AutoEnzyme(mode = Enzyme.Reverse)

julia> prep_der_rev = prepare_derivative(f, back_rev, 1.0);

julia> @btime derivative($f, $prep_der_rev, $back_rev, 1.0) seconds = 1;
  6.683 μs (112 allocations: 3.72 KiB)

julia> @btime Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active(1.0)) seconds = 1;
  5.942 μs (84 allocations: 2.72 KiB)

The overhead we're seeing here is due to the fact that as far as DI is concerned, derivative is a fundamentally forward-mode operation. Indeed, it is easily implemented with one pushforward. But when you call derivative in reverse mode, DI decomposes that single pushforward as many pullbacks, one for each output dimension. To do that, it calls the function f once to figure out the shape of the output and dispatch on the correct implementation (see this part of the code).
I don't know how to fuse that one function call with the evaluation of the pullbacks, since I need to know the output space in order to seed those pullbacks. However, I don't think it matters too much because this is a part of the package that only serves as a backup implementation. You see it when you call "the wrong mode" and it is easy to get around.

Here are a few ways to get the same performance with DI as with Enzyme:

• Use forward mode for the derivative (that's what users should do anyway).

julia> prep_der_for = prepare_derivative(f, back_for, 1.0);

julia> @btime derivative($f, $prep_der_for, $back_for, 1.0) seconds = 1;
  1.413 μs (54 allocations: 1.75 KiB)

julia> @btime Enzyme.autodiff(Enzyme.Forward, f, Enzyme.Duplicated(1.0, 1.0)) seconds = 1;
  1.408 μs (54 allocations: 1.75 KiB)

• Use a reverse-mode pullback, where you can specify the shape of the output through the cotangent argument and hence DI doesn't need that extra call to f.

julia> prep_pb_rev = prepare_pullback(f, back_rev, 1.0, (1.0,));

julia> @btime pullback($f, $prep_pb_rev, $back_rev, 1.0, (1.0,)) seconds = 1;
  6.117 μs (84 allocations: 2.72 KiB)

julia> @btime Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active(1.0)) seconds = 1;
  5.950 μs (84 allocations: 2.72 KiB)

• Use a reverse-mode gradient.

julia> prep_grad_rev = prepare_gradient(f, back_rev, 1.0);

julia> @btime gradient($f, $prep_grad_rev, $back_rev, 1.0) seconds = 1;
  5.938 μs (84 allocations: 2.72 KiB)

julia> @btime Enzyme.gradient(Enzyme.Reverse, f, 1.0) seconds = 1;
  5.942 μs (84 allocations: 2.72 KiB)

Is this answer to your satisfaction?

[gdalle]

Actually I just realized I can get that shape info from preparation, so #931 fixes most of the overhead:

julia> prep_der_rev = prepare_derivative(f, back_rev, 1.0);

julia> @btime derivative($f, $prep_der_rev, $back_rev, 1.0) seconds = 1;
  6.042 μs (84 allocations: 2.72 KiB)

julia> @btime Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active(1.0)) seconds = 1;
  5.896 μs (84 allocations: 2.72 KiB)

[wsmoses]

what is the source of the remaining gap?

[gdalle]

It's a 2.5% gap, come on 🤣 I'm sure it can be explained by plenty of things and might even be a benchmarking artifact. Unlike Enzyme, DI doesn't do everything at compile time, so it has strictly more stuff to compute than Enzyme for any given call. But in the grand scheme of things, especially for non-trivial functions, I'd consider this kind of gap to be small enough not to be noticeable.
If you have another MWE where you find a larger gap, feel free to open a new issue, and like for this one, I'll see what I can do!

Details

julia> 6.042 / 5.896
1.0247625508819538

[wsmoses]

it's very much reproducible so I don't think it's noise. It's also surprising becuase I wouldn't understand why it exists. Shouldn't DI just fowrard to Enzyme directly here?

[gdalle]

I just profiled the DI version and all of the visible time is spent inside Enzyme calls, so there's no DI overhead before or after (see flame graph below). I guess the key difference would be that DI.value_and_pullback uses Enzyme.autodiff_thunk¹ instead of Enzyme.autodiff, since in general it has no way of knowing that the seed will be 1.

Details

In the specific case of derivatives with scalar outputs, we do indeed know that the seed will be 1, so we could specialize it further for Enzyme (since it is the only backend for which seed = 1 is faster). However, given that (a) it's a 2.5% performance difference and (b) people should compute derivatives of scalar-input functions with forward mode anyway, does it really matter?
If you strongly care about this, I'll take another look, because in this specific instance it is feasible. But I don't have unlimited free time so it will not be on top of my priority list.

More importantly, users should not expect DI to have absolutely zero overhead. That's the reason behind the big red warning in the README, saying that native Enzyme might be faster. Here we reached a 2.5% difference, so I'd call that a win. My personal take is that users for whom this general level of DI service is not enough should definitely switch to native Enzyme. I hope this is understandable to you?

Footnotes

1. btw I still need help for upstreaming this to Enzyme itself ↩