Alloc-free pushforward with ForwardDiff (#227)

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -8,7 +8,7 @@ using DifferentiationInterface:
     HessianExtras,
     JacobianExtras,
     NoDerivativeExtras,
-    NoPushforwardExtras
+    PushforwardExtras
 using ForwardDiff.DiffResults: DiffResults, DiffResult, GradientResult, MutableDiffResult
 using ForwardDiff:
     Chunk,

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -1,16 +1,62 @@
 ## Pushforward
 
-DI.prepare_pushforward(f, ::AutoForwardDiff, x, dx) = NoPushforwardExtras()
+struct ForwardDiffOneArgPushforwardExtras{T,X} <: PushforwardExtras
+    xdual_tmp::X
+end
 
-function DI.value_and_pushforward(f, ::AutoForwardDiff, x, dx, ::NoPushforwardExtras)
+function DI.prepare_pushforward(f, ::AutoForwardDiff, x, dx)
     T = tag_type(f, x)
-    xdual = make_dual(T, x, dx)
-    ydual = f(xdual)
+    xdual_tmp = make_dual(T, x, dx)
+    return ForwardDiffOneArgPushforwardExtras{T,typeof(xdual_tmp)}(xdual_tmp)
+end
+
+function compute_ydual_onearg(
+    f, x, dx, extras::ForwardDiffOneArgPushforwardExtras{T}
+) where {T}
+    (; xdual_tmp) = extras
+    xdual_tmp = if x isa Number
+        make_dual(T, x, dx)
+    else
+        make_dual!(T, xdual_tmp, x, dx)
+    end
+    ydual = f(xdual_tmp)
+    return ydual
+end
+
+function DI.value_and_pushforward(
+    f, ::AutoForwardDiff, x, dx, extras::ForwardDiffOneArgPushforwardExtras{T}
+) where {T}
+    ydual = compute_ydual_onearg(f, x, dx, extras)
     y = myvalue(T, ydual)
     new_dy = myderivative(T, ydual)
     return y, new_dy
 end
 
+function DI.pushforward(
+    f, ::AutoForwardDiff, x, dx, extras::ForwardDiffOneArgPushforwardExtras{T}
+) where {T}
+    ydual = compute_ydual_onearg(f, x, dx, extras)
+    new_dy = myderivative(T, ydual)
+    return new_dy
+end
+
+function DI.value_and_pushforward!(
+    f, dy, ::AutoForwardDiff, x, dx, extras::ForwardDiffOneArgPushforwardExtras{T}
+) where {T}
+    ydual = compute_ydual_onearg(f, x, dx, extras)
+    y = myvalue(T, ydual)
+    myderivative!(T, dy, ydual)
+    return y, dy
+end
+
+function DI.pushforward!(
+    f, dy, ::AutoForwardDiff, x, dx, extras::ForwardDiffOneArgPushforwardExtras{T}
+) where {T}
+    ydual = compute_ydual_onearg(f, x, dx, extras)
+    myderivative!(T, dy, ydual)
+    return dy
+end
+
 ## Gradient
 
 struct ForwardDiffGradientExtras{C} <: GradientExtras

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl

@@ -1,15 +1,66 @@
-DI.prepare_pushforward(f!, y, ::AutoForwardDiff, x, dx) = NoPushforwardExtras()
+## Pushforward
 
-function DI.value_and_pushforward(f!, y, ::AutoForwardDiff, x, dx, ::NoPushforwardExtras)
+struct ForwardDiffTwoArgPushforwardExtras{T,X,Y} <: PushforwardExtras
+    xdual_tmp::X
+    ydual_tmp::Y
+end
+
+function DI.prepare_pushforward(f!, y, ::AutoForwardDiff, x, dx)
     T = tag_type(f!, x)
-    xdual = make_dual(T, x, dx)
-    ydual = make_dual(T, y, similar(y))
-    f!(ydual, xdual)
-    y = myvalue!(T, y, ydual)
-    dy = myderivative(T, ydual)
+    xdual_tmp = make_dual(T, x, dx)
+    ydual_tmp = make_dual(T, y, similar(y))
+    return ForwardDiffTwoArgPushforwardExtras{T,typeof(xdual_tmp),typeof(ydual_tmp)}(
+        xdual_tmp, ydual_tmp
+    )
+end
+
+function compute_ydual_twoarg(
+    f!, y, x, dx, extras::ForwardDiffTwoArgPushforwardExtras{T}
+) where {T}
+    (; xdual_tmp, ydual_tmp) = extras
+    xdual_tmp = if x isa Number
+        make_dual(T, x, dx)
+    else
+        make_dual!(T, xdual_tmp, x, dx)
+    end
+    f!(ydual_tmp, xdual_tmp)
+    return ydual_tmp
+end
+
+function DI.value_and_pushforward(
+    f!, y, ::AutoForwardDiff, x, dx, extras::ForwardDiffTwoArgPushforwardExtras{T}
+) where {T}
+    ydual_tmp = compute_ydual_twoarg(f!, y, x, dx, extras)
+    myvalue!(T, y, ydual_tmp)
+    dy = myderivative(T, ydual_tmp)
     return y, dy
 end
 
+function DI.pushforward(
+    f!, y, ::AutoForwardDiff, x, dx, extras::ForwardDiffTwoArgPushforwardExtras{T}
+) where {T}
+    ydual_tmp = compute_ydual_twoarg(f!, y, x, dx, extras)
+    dy = myderivative(T, ydual_tmp)
+    return dy
+end
+
+function DI.value_and_pushforward!(
+    f!, y, dy, ::AutoForwardDiff, x, dx, extras::ForwardDiffTwoArgPushforwardExtras{T}
+) where {T}
+    ydual_tmp = compute_ydual_twoarg(f!, y, x, dx, extras)
+    myvalue!(T, y, ydual_tmp)
+    myderivative!(T, dy, ydual_tmp)
+    return y, dy
+end
+
+function DI.pushforward!(
+    f!, y, dy, ::AutoForwardDiff, x, dx, extras::ForwardDiffTwoArgPushforwardExtras{T}
+) where {T}
+    ydual_tmp = compute_ydual_twoarg(f!, y, x, dx, extras)
+    myderivative!(T, dy, ydual_tmp)
+    return dy
+end
+
 ## Derivative
 
 struct ForwardDiffTwoArgDerivativeExtras{C} <: DerivativeExtras

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -1,13 +1,14 @@
-
 choose_chunk(::AutoForwardDiff{nothing}, x) = Chunk(x)
 choose_chunk(::AutoForwardDiff{C}, x) where {C} = Chunk{C}()
 
 tag_type(::F, x::Number) where {F} = Tag{F,typeof(x)}
 tag_type(::F, x::AbstractArray) where {F} = Tag{F,eltype(x)}
 
-make_dual(::Type{T}, x::Number, dx::Number) where {T} = Dual{T}(x, dx)
+make_dual(::Type{T}, x::Number, dx) where {T} = Dual{T}(x, dx)
 make_dual(::Type{T}, x::AbstractArray, dx) where {T} = Dual{T}.(x, dx)
 
+make_dual!(::Type{T}, xdual, x::AbstractArray, dx) where {T} = xdual .= Dual{T}.(x, dx)
+
 myvalue(::Type{T}, ydual::Number) where {T} = value(T, ydual)
 myvalue(::Type{T}, ydual::AbstractArray) where {T} = value.(T, ydual)
 
@@ -16,6 +17,6 @@ myvalue!(::Type{T}, y::AbstractArray, ydual) where {T} = y .= value.(T, ydual)
 myderivative(::Type{T}, ydual::Number) where {T} = extract_derivative(T, ydual)
 myderivative(::Type{T}, ydual::AbstractArray) where {T} = extract_derivative(T, ydual)
 
-function myderivative!(::Type{T}, dy::AbstractArray, ydual::AbstractArray) where {T}
+function myderivative!(::Type{T}, dy, ydual::AbstractArray) where {T}
     return extract_derivative!(T, dy, ydual)
 end