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

[wsmoses]

gentle ping @gdalle

[gdalle]

Essentially this comes down to adding the option for Active inputs here:

Enzyme.jl/src/Enzyme.jl

Lines 1714 to 1720 in 42ecd12

	if $arg isa Enzyme.Const
	$arg
	elseif $act
	MixedDuplicated($arg, $shad)
	else
	Duplicated($arg, $shad)
	end

The variable interpolated as $arg is a boolean defined as follows, which is a bit obscure to me, care to shed some light?

Enzyme.jl/src/Enzyme.jl

Lines 1686 to 1692 in 42ecd12

	!($argidx isa Enzyme.Const) &&
	Compiler.active_reg_inner(
	Core.Typeof($argidx),
	(),
	nothing,
	Val(true),
	) == Compiler.ActiveState #=justActive=#