WaterLily using ForwardDiff and Enzyme through DifferentiationInterface #355
Description
Hey! I have been playing with this package to implement AD for WaterLily.jl. The ForwardDiff backend works as expected, and it provides the same result as using ForwardDiff without DI, only if compilation is performed through the appropriate Dual type. For example, the following MWE works
using WaterLily
using StaticArrays
using DifferentiationInterface
import ForwardDiff
function make_sim(θ;L=32,U=1,Re=100,mem=Array)
function map(x,_)
s,c = sincos(θ)
SA[c -s; s c]*(x-SA[L,L])
end
function sdf(ξ,_)
p = ξ-SA[0,clamp(ξ[1],-L/2,L/2)]
√(p'*p)-2
end
Simulation((2L,2L),(U,0),L,ν=U*L/Re,body=AutoBody(sdf,map),T=typeof(θ),mem=mem)
end
function dsim(θ)
sim = make_sim(θ)
sim_step!(sim)
sum(sim.flow.p)
end
function main()
FT = Float64
θ = FT(π/36)
# FD
θ1 = θ*0.999
θ2 = θ*1.001
println("FD value and gradient\n",
dsim(θ),",",(dsim(θ2)-dsim(θ1))/(θ2-θ1)
)
# ForwardDiff
T = typeof(ForwardDiff.Tag(dsim,FT))
θAD = ForwardDiff.Dual{T}(θ,one(FT))
println("\nForwardDiff value and gradient\n", dsim(θAD))
# ForwardDiff via DifferentiationInterface
println("\nDifferentiationInterface ForwardDiff value and gradient\n",
value_and_gradient(dsim, AutoForwardDiff(), θ)
)
end
main()but commenting out the ForwardDiff lines
# T = typeof(ForwardDiff.Tag(dsim,FT))
# θAD = ForwardDiff.Dual{T}(θ,one(FT))
# println("\nForwardDiff value and gradient\n", dsim(θAD))makes the DI call to error because of the ordering of Dual tags.
Furthermore, switching to Enzyme with
value_and_gradient(dsim, AutoEnzyme(), θ)throws an error (with a rather large stacktrace...). I am not sure how to specify forward or reverse mode when using the AutoEnzyme() backend, so it could be that by default tries to do backward mode and we have some incompatibility with that. Any help would be appreciated, thanks! :)