ReverseDiff performance about reversediff.jl HOT 4 CLOSED

using ReverseDiff: GradientTape, gradient!, compile, @forward
using BenchmarkTools: @btime

@forward φ1(r) = r^2 + r^4 + exp(-r)
@forward φ2(r) = -0.1 * r^3

function F(x)
    N = length(x)
    f = φ1(x[2]-x[1])
    for i = 3:N
        f += φ1(x[i]-x[i-1]) + φ2(x[i] - x[i-2])
    end
    return f
end

function Fvec(x)
    N = length(x)
    f = φ1(x[2] - x[1])
    tmp1 = x[3:N]
    tmp2 = x[2:(N-1)]
    tmp3 = x[1:(N-2)]
    f += sum(φ1.(tmp1 - tmp2) + φ2.(tmp1 - tmp3))
    return f
end

N = 40000
f_tape = GradientTape(F, rand(N))
compiled_f_tape = compile(f_tape)
fvec_tape = GradientTape(Fvec, rand(N))
compiled_fvec_tape = compile(fvec_tape)

x = rand(N)
results = zeros(N)

println("Time for F(x):")
@btime F($x)

println("Time for gradient!(results, f_tape, x):")
@btime gradient!($results, $f_tape, $x)

println("Time for gradient!(results, compiled_f_tape, x):")
@btime gradient!($results, $compiled_f_tape, $x)

println("Time for Fvec(x):")
@btime Fvec($x)

println("Time for gradient!(results, fvec_tape, x):")
@btime gradient!($results, $fvec_tape, $x)

println("Time for gradient!(results, compiled_fvec_tape, x):")
@btime gradient!($results, $compiled_fvec_tape, $x)

Time for F(x):
  1.729 ms (0 allocations: 0 bytes)
Time for gradient!(results, f_tape, x):
  35.348 ms (0 allocations: 0 bytes)
Time for gradient!(results, compiled_f_tape, x):
  37.608 ms (0 allocations: 0 bytes)
Time for Fvec(x):
  2.040 ms (16 allocations: 2.44 MiB)
Time for gradient!(results, fvec_tape, x):
  6.098 ms (0 allocations: 0 bytes)
Time for gradient!(results, compiled_fvec_tape, x):
  6.123 ms (0 allocations: 0 bytes)