ReverseDiff compile mode completed

ext/OptimizationReverseDiffExt.jl

@@ -7,29 +7,19 @@ import Optimization.ADTypes: AutoReverseDiff
 isdefined(Base, :get_extension) ? (using ReverseDiff, ReverseDiff.ForwardDiff) :
 (using ..ReverseDiff, ..ReverseDiff.ForwardDiff)
 
+struct OptimizationReverseDiffTag end
+
 function Optimization.instantiate_function(f, x, adtype::AutoReverseDiff,
     p = SciMLBase.NullParameters(),
     num_cons = 0)
     _f = (θ, args...) -> first(f.f(θ, p, args...))
 
     if f.grad === nothing
-
         if adtype.compile
-            _tape = ReverseDiff.GradientTape((x,)) do θ
-                _f(θ)
-            end
+            _tape = ReverseDiff.GradientTape(_f, x)
             tape = ReverseDiff.compile(_tape)
-
             grad = function (res, θ, args...)
-                tθ = ReverseDiff.input_hook(tape)[1]
-                output = ReverseDiff.output_hook(tape)
-                ReverseDiff.unseed!(tθ) # clear any "leftover" derivatives from previous calls
-                ReverseDiff.value!(tθ, θ)
-                ReverseDiff.forward_pass!(tape)
-                ReverseDiff.increment_deriv!(output, one(eltype(θ)))
-                ReverseDiff.reverse_pass!(tape)
-                copyto!(res, ReverseDiff.deriv(tθ))
-                nothing
+                ReverseDiff.gradient!(res, tape, θ)
             end
         else
             cfg = ReverseDiff.GradientConfig(x)
@@ -40,8 +30,23 @@ function Optimization.instantiate_function(f, x, adtype::AutoReverseDiff,
     end
 
     if f.hess === nothing
-        hess = function (res, θ, args...)
-            ReverseDiff.hessian!(res, x -> _f(x, args...), θ)
+        if adtype.compile
+            T = ForwardDiff.Tag(OptimizationReverseDiffTag(),eltype(x))
+            xdual = ForwardDiff.Dual{typeof(T),eltype(x),length(x)}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), length(x))...,))))
+            h_tape = ReverseDiff.GradientTape(_f, xdual)
+            htape = ReverseDiff.compile(h_tape)
+            function g(θ)
+                res1 = zeros(eltype(θ), length(θ))
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            jaccfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk(x), T)
+            hess = function (res, θ, args...)
+                ForwardDiff.jacobian!(res, g, θ, jaccfg, Val{false}())
+            end
+        else
+            hess = function (res, θ, args...)
+                ReverseDiff.hessian!(res, x -> _f(x, args...), θ)
+            end
         end
     else
         hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
@@ -66,19 +71,43 @@ function Optimization.instantiate_function(f, x, adtype::AutoReverseDiff,
     end
 
     if cons !== nothing && f.cons_j === nothing
-        cjconfig = ReverseDiff.JacobianConfig(x)
-        cons_j = function (J, θ)
-            ReverseDiff.jacobian!(J, cons_oop, θ, cjconfig)
+        if adtype.compile
+            _jac_tape = ReverseDiff.JacobianTape(cons_oop, x)
+            jac_tape = ReverseDiff.compile(_jac_tape)
+            cons_j = function (J, θ)
+                ReverseDiff.jacobian!(J, jac_tape, θ)
+            end
+        else
+            cjconfig = ReverseDiff.JacobianConfig(x)
+            cons_j = function (J, θ)
+                ReverseDiff.jacobian!(J, cons_oop, θ, cjconfig)
+            end
         end
     else
         cons_j = (J, θ) -> f.cons_j(J, θ, p)
     end
 
     if cons !== nothing && f.cons_h === nothing
         fncs = [(x) -> cons_oop(x)[i] for i in 1:num_cons]
-        cons_h = function (res, θ)
-            for i in 1:num_cons
-                ReverseDiff.hessian!(res[i], fncs[i], θ)
+        if adtype.compile
+            consh_tapes = ReverseDiff.GradientTape.(fncs, Ref(xdual))
+            conshtapes = ReverseDiff.compile.(consh_tapes)
+            function grad_cons(θ, htape)
+                res1 = zeros(eltype(θ), length(θ))
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            gs = [x -> grad_cons(x, conshtapes[i]) for i in 1:num_cons]
+            jaccfgs = [ForwardDiff.JacobianConfig(gs[i], x, ForwardDiff.Chunk(x), T) for i in 1:num_cons]
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    ForwardDiff.jacobian!(res[i], gs[i], θ, jaccfgs[i], Val{false}())
+                end
+            end
+        else
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    ReverseDiff.hessian!(res[i], fncs[i], θ)
+                end
             end
         end
     else
@@ -103,14 +132,38 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
     _f = (θ, args...) -> first(f.f(θ, cache.p, args...))
 
     if f.grad === nothing
-        grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ)
+        if adtype.compile
+            _tape = ReverseDiff.GradientTape(_f, cache.u0)
+            tape = ReverseDiff.compile(_tape)
+            grad = function (res, θ, args...)
+                ReverseDiff.gradient!(res, tape, θ)
+            end
+        else
+            cfg = ReverseDiff.GradientConfig(cache.u0)
+            grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ, cfg)
+        end
     else
         grad = (G, θ, args...) -> f.grad(G, θ, cache.p, args...)
     end
 
     if f.hess === nothing
-        hess = function (res, θ, args...)
-            ReverseDiff.hessian!(res, x -> _f(x, args...), θ)
+        if adtype.compile
+            T = ForwardDiff.Tag(OptimizationReverseDiffTag(),eltype(cache.u0))
+            xdual = ForwardDiff.Dual{typeof(T),eltype(cache.u0),length(cache.u0)}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), length(cache.u0))...,))))
+            h_tape = ReverseDiff.GradientTape(_f, xdual)
+            htape = ReverseDiff.compile(h_tape)
+            function g(θ)
+                res1 = zeros(eltype(θ), length(θ))
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            jaccfg = ForwardDiff.JacobianConfig(g, cache.u0, ForwardDiff.Chunk(cache.u0), T)
+            hess = function (res, θ, args...)
+                ForwardDiff.jacobian!(res, g, θ, jaccfg, Val{false}())
+            end
+        else
+            hess = function (res, θ, args...)
+                ReverseDiff.hessian!(res, x -> _f(x, args...), θ)
+            end
         end
     else
         hess = (H, θ, args...) -> f.hess(H, θ, cache.p, args...)
@@ -135,19 +188,43 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
     end
 
     if cons !== nothing && f.cons_j === nothing
-        cjconfig = ReverseDiff.JacobianConfig(cache.u0)
-        cons_j = function (J, θ)
-            ReverseDiff.jacobian!(J, cons_oop, θ, cjconfig)
+        if adtype.compile
+            _jac_tape = ReverseDiff.JacobianTape(cons_oop, cache.u0)
+            jac_tape = ReverseDiff.compile(_jac_tape)
+            cons_j = function (J, θ)
+                ReverseDiff.jacobian!(J, jac_tape, θ)
+            end
+        else
+            cjconfig = ReverseDiff.JacobianConfig(cache.u0)
+            cons_j = function (J, θ)
+                ReverseDiff.jacobian!(J, cons_oop, θ, cjconfig)
+            end
         end
     else
         cons_j = (J, θ) -> f.cons_j(J, θ, cache.p)
     end
 
     if cons !== nothing && f.cons_h === nothing
         fncs = [(x) -> cons_oop(x)[i] for i in 1:num_cons]
-        cons_h = function (res, θ)
-            for i in 1:num_cons
-                ReverseDiff.hessian!(res[i], fncs[i], θ)
+        if adtype.compile
+            consh_tapes = ReverseDiff.GradientTape.(fncs, Ref(xdual))
+            conshtapes = ReverseDiff.compile.(consh_tapes)
+            function grad_cons(θ, htape)
+                res1 = zeros(eltype(θ), length(θ))
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            gs = [x -> grad_cons(x, conshtapes[i]) for i in 1:num_cons]
+            jaccfgs = [ForwardDiff.JacobianConfig(gs[i], cache.u0, ForwardDiff.Chunk(cache.u0), T) for i in 1:num_cons]
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    ForwardDiff.jacobian!(res[i], gs[i], θ, jaccfgs[i], Val{false}())
+                end
+            end
+        else
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    ReverseDiff.hessian!(res[i], fncs[i], θ)
+                end
             end
         end
     else