Enzyme custom rule alloactions

General Usage Performance
, ,
1

Hello,
I got a problem with allocations when using custom rules in Enzyme.
The example below is just a toy example but the same problem arises there.
In total I get 4 alloactions.
2 by the line primal = func.val(y.val, x.val) and
1 by x.dval .+= 2 .* xval .* dret.val and 1 by x.dval .+= 2 .* xval .* y.dval.
If I dont define any custom rule I measure 0 allocations.

Code:

using Enzyme
import Enzyme.EnzymeRules: forward, reverse, augmented_primal
using Enzyme.EnzymeRules


x = [3.0, 1.0]
dx = [1.0, 0.0]
y = [0.0, 0.0]
dy = [0.0, 0.0]
function f(y, x)
    y .= x.^2
    return sum(y)
end


gf(y, x) = f(y, x)^2 # function to differentiate

function EnzymeRules.augmented_primal(config::R, 
    func::Const{F}, ::Type{A},
    y::Duplicated{Vector{Float64}}, x::Duplicated{Vector{Float64}}) where {R<:RevConfigWidth{1},F <: typeof(f),A<:Active}
    primal = func.val(y.val, x.val)
    return AugmentedReturn(primal, nothing, nothing)
end 

function EnzymeRules.reverse(config::RevConfigWidth{1}, func::Const{F}, dret::Active, tape,
    y::Duplicated, x::Duplicated) where F <: typeof(f)
    # println("In custom reverse rule.")
    # retrieve x value, either from original x or from tape if x may have been overwritten.
    xval = x.val
    # accumulate dret into x's shadow. don't assign!
    x.dval .+= 2 .* xval .* dret.val
    # # also accumulate any derivative in y's shadow into x's shadow.
    x.dval .+= 2 .* xval .* y.dval
    make_zero!(y.dval)
    return (nothing, nothing)
end

Enzyme.autodiff(Enzyme.Reverse, gf, Enzyme.Duplicated(y, dy), Enzyme.Duplicated(x, dx))

using Chairmarks
@b Enzyme.autodiff(Enzyme.Reverse, $gf, Enzyme.Duplicated($y, $dy), Enzyme.Duplicated($x, $dx))

Thank you for you help