Enzyme slower than FiniteDiff with `Cache`

[ErikQQY]

The Cache in DifferentiationInterface is really a great feature for in-place mutating arrays inside some preallocated cache stuff, but I found it's slow and allocates a lot when trying the Enzyme backend. I'm not sure if I am using it correctly; here is a simple MWE:

using ConcreteStructs
using DifferentiationInterface
using Enzyme
using FiniteDiff
const DI = DifferentiationInterface
using BenchmarkTools

@concrete struct TESTCache
    res
    u
end

# helper function to transform between vector of arrays and flat vector
@views function recursive_unflatten!(y::Vector{<:AbstractArray}, x::AbstractVector)
    i = 0
    for yᵢ in y
        copyto!(yᵢ, x[(i + 1):(i + length(yᵢ))])
        i += length(yᵢ)
    end
    return y
end
@views function recursive_flatten!(y::AbstractVector, x::Vector{<:AbstractArray})
    i = 0
    for xᵢ in x
        copyto!(y[(i + 1):(i + length(xᵢ))], xᵢ)
        i += length(xᵢ)
    end
    return y
end
recursive_length(x::Vector{<:AbstractArray}) = sum(length, x)

@views function bc(residual, u, p, t)
    residual[1] = u[1] .+ pi / 2
    residual[2] = u[end] .- pi / 2
end

@views function test(du, u, y, residual)
    y_ = recursive_unflatten!(y, u)
    bc(residual, y_, 0, 0)
    recursive_flatten!(du, residual)
end

function f(du, u, c::TESTCache)
    test(du, u, c.u, c.res)
end

x = rand(100)
y = similar(x)
test_cache = TESTCache([rand(2) for _ in 1:50], [rand(2) for _ in 1:50])
backend_finitediff = AutoFiniteDiff()
backend_enzyme = AutoEnzyme(; mode = Enzyme.Forward)
prep_finitediff = DI.prepare_jacobian(f, y, backend_finitediff, x, Cache(test_cache))
@benchmark DI.jacobian(f, y, prep_finitediff, backend_finitediff, x, Cache(test_cache))
prep_enzyme = DI.prepare_jacobian(f, y, backend_enzyme, x, Cache(test_cache))
@benchmark DI.jacobian(f, y, prep_enzyme, backend_enzyme, x, Cache(test_cache))

FiniteDiff

BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  45.334 μs …   5.164 ms  ┊ GC (min … max): 0.00% … 98.88%
 Time  (median):     52.333 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   56.231 μs ± 105.995 μs  ┊ GC (mean ± σ):  6.85% ±  4.87%

  ▅▃▁ ▄▇▄▂▄▇█▇▅▄▂▂▂▂▁▁▁                                        ▂
  ██████████████████████▇▇▆▆▆▆▅▅▄▁▄▅▄▄▃▄▁▄▁▃▄▃▄▁▄▄▁▃▅▄▄▃▁▄▄▄▄▄ █
  45.3 μs       Histogram: log(frequency) by time        87 μs <

 Memory estimate: 94.11 KiB, allocs estimate: 411.

Enzyme

BenchmarkTools.Trial: 5592 samples with 1 evaluation per sample.
 Range (min … max):  584.958 μs … 51.009 ms  ┊ GC (min … max):  0.00% … 97.69%
 Time  (median):     665.562 μs              ┊ GC (median):     0.00%
 Time  (mean ± σ):   890.874 μs ±  1.279 ms  ┊ GC (mean ± σ):  14.73% ± 16.62%

  ▆█▆▃▃▁                                                       ▁
  ███████▇▆▆▄▄▂▄▄▄▄▅▄▃▄▄▄▄▄▄▃▂▃▆▆▆▇▇██▇▇▆▇▆▆▅▆▆▄▆▅▅▄▂▄▄▃▄▄▄▃▃▃ █
  585 μs        Histogram: log(frequency) by time      3.74 ms <

 Memory estimate: 2.24 MiB, allocs estimate: 23486.

[gdalle]

Cache hasn't yet been optimized for Enzyme because no one has asked for it yet, it's on my to-do list. If you want to take a crack at it, you can draw inspiration from what I did in the ForwardDiff case. Essentially, it requires preallocating cache contexts at preparation time, and then reusing them at execution time.
Otherwise I should be able to get to it soon.

[ErikQQY]

Seems I am really good at finding eage cases😆. I will try to implement cache context optimizations for Enzyme following #741

FYI, I am trying to support the Enzyme and Mooncake backend in SciML/BoundaryValueDiffEq.jl#291, but I found out ForwardDiff is always like 2x faster than Enzyme using function_annotation = Enzyme.Duplicated(because of Enzyme cache-mutating related readonly error), and 100x faster than Enzyme using Cache, so I thought maybe it's the Cache for Enzyme is implemented erroneously and it turns out no one is using it😂

[gdalle]

I mean, it's not erroneous, the semantics are correct, it's just the optimization missing 😅
Why can't you use a constant function annotation for AutoEnzyme? Does your function contain writable storage?

[ErikQQY]

Why can't you use a constant function annotation for AutoEnzyme? Does your function contain writable storage?

Does constant function annotation suitable for the scenario where the function involving actively mutating cache storage? That is also the reason I am trying Cache for this.

[gdalle]

No, in that case you need function_annotation=Enzyme.Duplicated indeed.

[gdalle]

Hey there @ErikQQY, got a fix for you! You can take it out for a spin on the branch #782.

I wouldn't do it with your initial example though, because with my new DI additions it seems to generate invalid IR that crashes the REPL.
But if you take a simpler example like the one below, you can see that Enzyme with Caches is now fully non-allocating and much faster!

using BenchmarkTools
using DifferentiationInterface
using Enzyme: Enzyme
using FiniteDiff: FiniteDiff

function f!(y, x, c)
    copyto!(c, x)
    c[1] = c[end] = 0
    for i in 2:(length(c) - 1)
        c[i] = exp(c[i]) - cos(c[i - 1])
    end
    y .= 2 .* c .+ 1
    return nothing
end

x = rand(100);
y, c = similar(x), similar(x);
J = similar(x, 100, 100);
backend1 = AutoEnzyme(; mode=Enzyme.Forward)
backend2 = AutoFiniteDiff()
prep1 = prepare_jacobian(f!, y, backend1, x, Cache(x));
prep2 = prepare_jacobian(f!, y, backend2, x, Cache(x));

Enzyme benchmark:

julia> @benchmark jacobian!($f!, $y, $J, $prep1, $backend1, $x, Cache($c))
BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  19.708 μs … 38.375 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     19.958 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   20.330 μs ±  1.490 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▄█▆▂ ▂     ▁▁                                               ▁
  ███████▆▆▆▅██▆▆▆▇▇▇█▆▆▅▄▅▆▅▅▄▅▄▄▅▅▅▆▅▅▅▅▅▅▅▅▄▄▄▅▆▆▃▅▄▄▃▆▅▆▆ █
  19.7 μs      Histogram: log(frequency) by time      27.8 μs <

 Memory estimate: 0 bytes, allocs estimate: 0.

FiniteDiff benchmark:

julia> @benchmark jacobian!($f!, $y, $J, $prep2, $backend2, $x, Cache($c))
BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  138.209 μs … 260.875 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     138.958 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   141.994 μs ±   6.544 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ██▃▅▁▁▄▄▁▂▁     ▁   ▅▃▁   ▁                                   ▂
  ██████████████▇█████████████▇███▇▇▆▆▇▆▅▆▅▆▆▆▆▆▆▅▆▆▅▅▅▅▅▄▄▅▅▅▅ █
  138 μs        Histogram: log(frequency) by time        169 μs <

 Memory estimate: 0 bytes, allocs estimate: 0