Enzyme.gradient allocates on SVector

[gdalle]

Hi!
As you know, @ExpandingMan and I are looking to optimize performance for StaticArrays. Forward mode works splendidly, but reverse mode still makes one allocation during the gradient call:

using StaticArrays, Enzyme, BenchmarkTools
f(x) = sum(abs2, x);
x = SVector(1.0, 2.0);
@btime Enzyme.gradient(Enzyme.Reverse, f, $x)  # 8.999 ns (1 allocation: 32 bytes)
@btime Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active($x))  # 4.218 ns (0 allocations: 0 bytes)

I found it surprising because Enzyme guesses the right activity for SVector:

Enzyme.guess_activity(typeof(x), Enzyme.Reverse)  # Active{SVector{2, Float64}}

The allocation happens on the following line:

Enzyme.jl/src/Enzyme.jl

Line 1708 in 42ecd12

Ref(make_zero($arg))

From what I understand, the generated function Enzyme.gradient puts a Ref there to treat every argument as (Mixed)Duplicated. This means that all gradient results are stored in the passed arguments:

Enzyme.jl/src/Enzyme.jl

Line 1741 in 42ecd12

(; derivs = ($(resargs...),), val = res[2])

Otherwise, you would have to recover some gradients from the result and others from the arguments, which is understandably tricky.
Do you think there is an easy fix in Enzyme? Otherwise, since DI only has one differentiated argument, I assume it will be rather straightfoward to call Enzyme.autodiff directly inside DI.gradient and recover allocation-free behavior.

Related:

• Enzyme gradient on StaticArrays allocates JuliaDiff/DifferentiationInterface.jl#583

[wsmoses]

sure, PR welcome!

[gdalle]

Sure! I'll try to handle this case correctly in DI first, because it still errors at the moment. Once I have a handle on the single-argument solution, I'll try to tamper with the generated function to do the same for multiple arguments.

[wsmoses]

bump @gdalle