perf: reduce allocations with Enzyme for in-place functions (#707)

Author: gdalle

.github/workflows/Test.yml

@@ -25,7 +25,7 @@ jobs:
       actions: write
       contents: read
     strategy:
-      fail-fast: false  # TODO: toggle
+      fail-fast: true  # TODO: toggle
       matrix:
         version:
           - "1.10"

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.35"
+version = "0.6.36"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/docs/src/explanation/backends.md

@@ -67,7 +67,7 @@ Moreover, each context type is supported by a specific subset of backends:
 | `AutoFiniteDifferences`    | ✅                  | ✅               |
 | `AutoForwardDiff`          | ✅                  | ✅               |
 | `AutoGTPSA`                | ✅                  | ❌               |
-| `AutoMooncake`             | ✅                  | ❌               |
+| `AutoMooncake`             | ✅                  | ✅               |
 | `AutoPolyesterForwardDiff` | ✅                  | ✅               |
 | `AutoReverseDiff`          | ✅                  | ❌               |
 | `AutoSymbolics`            | ✅                  | ❌               |

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -67,15 +67,22 @@ end
 function DI.value_and_pushforward!(
     f!::F,
     y,
-    ty::NTuple,
-    prep::DI.NoPushforwardPrep,
+    ty::NTuple{B},
+    ::DI.NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::NTuple,
+    tx::NTuple{B},
     contexts::Vararg{DI.Context,C},
-) where {F,C}
-    y, new_ty = DI.value_and_pushforward(f!, y, prep, backend, x, tx, contexts...)
-    foreach(copyto!, ty, new_ty)
+) where {F,B,C}
+    mode = forward_noprimal(backend)
+    f!_and_df! = get_f_and_df(f!, backend, mode, Val(B))
+    tx_sametype = map(Fix1(convert, typeof(x)), tx)
+    ty_sametype = map(Fix1(convert, typeof(y)), ty)
+    x_and_tx = BatchDuplicated(x, tx_sametype)
+    y_and_ty = BatchDuplicated(y, ty_sametype)
+    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    autodiff(mode, f!_and_df!, Const, y_and_ty, x_and_tx, annotated_contexts...)
+    foreach(copyto_if_different_addresses!, ty, ty_sametype)
     return y, ty
 end
 
@@ -89,7 +96,6 @@ function DI.pushforward!(
     tx::NTuple,
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
-    new_ty = DI.pushforward(f!, y, prep, backend, x, tx, contexts...)
-    foreach(copyto!, ty, new_ty)
+    DI.value_and_pushforward!(f!, y, ty, prep, backend, x, tx, contexts...)
     return ty
 end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -49,21 +49,26 @@ end
 
 ## Pullback
 
+struct EnzymeReverseOneArgPullbackPrep{Y} <: DI.PullbackPrep
+    y_example::Y  # useful to create return activity
+end
+
 function DI.prepare_pullback(
     f::F,
     ::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple,
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
-    return DI.NoPullbackPrep()
+    y = f(x, map(DI.unwrap, contexts)...)
+    return EnzymeReverseOneArgPullbackPrep(y)
 end
 
 ### Out-of-place
 
 function DI.value_and_pullback(
     f::F,
-    ::DI.NoPullbackPrep,
+    prep::EnzymeReverseOneArgPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple{1},
@@ -72,7 +77,7 @@ function DI.value_and_pullback(
     mode = reverse_split_withprimal(backend)
     f_and_df = force_annotation(get_f_and_df(f, backend, mode))
     IA = guess_activity(typeof(x), mode)
-    RA = guess_activity(eltype(ty), mode)
+    RA = guess_activity(typeof(prep.y_example), mode)
     dx = make_zero(x)
     annotated_contexts = translate(backend, mode, Val(1), contexts...)
     dinputs, result = seeded_autodiff_thunk(
@@ -88,7 +93,7 @@ end
 
 function DI.value_and_pullback(
     f::F,
-    ::DI.NoPullbackPrep,
+    prep::EnzymeReverseOneArgPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple{B},
@@ -97,7 +102,7 @@ function DI.value_and_pullback(
     mode = reverse_split_withprimal(backend)
     f_and_df = force_annotation(get_f_and_df(f, backend, mode, Val(B)))
     IA = batchify_activity(guess_activity(typeof(x), mode), Val(B))
-    RA = batchify_activity(guess_activity(eltype(ty), mode), Val(B))
+    RA = batchify_activity(guess_activity(typeof(prep.y_example), mode), Val(B))
     tx = ntuple(_ -> make_zero(x), Val(B))
     annotated_contexts = translate(backend, mode, Val(B), contexts...)
     dinputs, result = batch_seeded_autodiff_thunk(
@@ -113,7 +118,7 @@ end
 
 function DI.pullback(
     f::F,
-    prep::DI.NoPullbackPrep,
+    prep::EnzymeReverseOneArgPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple,
@@ -127,51 +132,51 @@ end
 function DI.value_and_pullback!(
     f::F,
     tx::NTuple{1},
-    ::DI.NoPullbackPrep,
+    prep::EnzymeReverseOneArgPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple{1},
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     mode = reverse_split_withprimal(backend)
     f_and_df = force_annotation(get_f_and_df(f, backend, mode))
-    RA = guess_activity(eltype(ty), mode)
+    RA = guess_activity(typeof(prep.y_example), mode)
     dx_righttype = convert(typeof(x), only(tx))
     make_zero!(dx_righttype)
     annotated_contexts = translate(backend, mode, Val(1), contexts...)
     _, result = seeded_autodiff_thunk(
         mode, only(ty), f_and_df, RA, Duplicated(x, dx_righttype), annotated_contexts...
     )
-    only(tx) === dx_righttype || copyto!(only(tx), dx_righttype)
+    copyto_if_different_addresses!(only(tx), dx_righttype)
     return result, tx
 end
 
 function DI.value_and_pullback!(
     f::F,
     tx::NTuple{B},
-    ::DI.NoPullbackPrep,
+    prep::EnzymeReverseOneArgPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple{B},
     contexts::Vararg{DI.Context,C},
 ) where {F,B,C}
     mode = reverse_split_withprimal(backend)
     f_and_df = force_annotation(get_f_and_df(f, backend, mode, Val(B)))
-    RA = batchify_activity(guess_activity(eltype(ty), mode), Val(B))
+    RA = batchify_activity(guess_activity(typeof(prep.y_example), mode), Val(B))
     tx_righttype = map(Fix1(convert, typeof(x)), tx)
     make_zero!(tx_righttype)
     annotated_contexts = translate(backend, mode, Val(B), contexts...)
     _, result = batch_seeded_autodiff_thunk(
         mode, ty, f_and_df, RA, BatchDuplicated(x, tx_righttype), annotated_contexts...
     )
-    foreach(copyto!, tx, tx_righttype)
+    foreach(copyto_if_different_addresses!, tx, tx_righttype)
     return result, tx
 end
 
 function DI.pullback!(
     f::F,
     tx::NTuple,
-    prep::DI.NoPullbackPrep,
+    prep::EnzymeReverseOneArgPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::NTuple,
@@ -265,7 +270,7 @@ function DI.gradient!(
     make_zero!(grad_righttype)
     annotated_contexts = translate(backend, mode, Val(1), contexts...)
     autodiff(mode, f_and_df, Active, Duplicated(x, grad_righttype), annotated_contexts...)
-    grad === grad_righttype || copyto!(grad, grad_righttype)
+    copyto_if_different_addresses!(grad, grad_righttype)
     return grad
 end
 
@@ -295,70 +300,6 @@ function DI.value_and_gradient!(
     _, y = autodiff(
         mode, f_and_df, Active, Duplicated(x, grad_righttype), annotated_contexts...
     )
-    grad === grad_righttype || copyto!(grad, grad_righttype)
+    copyto_if_different_addresses!(grad, grad_righttype)
     return y, grad
 end
-
-## Jacobian
-
-# TODO: does not support static arrays
-
-#=
-struct EnzymeReverseOneArgJacobianPrep{Sy,B} <:DI.JacobianPrep end
-
-function EnzymeReverseOneArgJacobianPrep(::Val{Sy}, ::Val{B}) where {Sy,B}
-    return EnzymeReverseOneArgJacobianPrep{Sy,B}()
-end
-
-function DI.prepare_jacobian(f::F, backend::AutoEnzyme{<:ReverseMode,Nothing}, x) where {F}
-    y = f(x)
-    Sy = size(y)
-    valB = to_val(DI.pick_batchsize(backend, y))
-    return EnzymeReverseOneArgJacobianPrep(Val(Sy), valB)
-end
-
-function DI.jacobian(
-    f::F,
-    ::EnzymeReverseOneArgJacobianPrep{Sy,B},
-    backend::AutoEnzyme{<:ReverseMode,Nothing},
-    x,
-) where {F,Sy,B}
-    derivs = jacobian(reverse_noprimal(backend), f, x; n_outs=Val(Sy), chunk=Val(B))
-    jac_tensor = only(derivs)
-    return maybe_reshape(jac_tensor, prod(Sy), length(x))
-end
-
-function DI.value_and_jacobian(
-    f::F,
-    ::EnzymeReverseOneArgJacobianPrep{Sy,B},
-    backend::AutoEnzyme{<:ReverseMode,Nothing},
-    x,
-) where {F,Sy,B}
-    (; derivs, val) = jacobian(
-        reverse_withprimal(backend), f, x; n_outs=Val(Sy), chunk=Val(B)
-    )
-    jac_tensor = only(derivs)
-    return val, maybe_reshape(jac_tensor, prod(Sy), length(x))
-end
-
-function DI.jacobian!(
-    f::F,
-    jac,
-    prep::EnzymeReverseOneArgJacobianPrep,
-    backend::AutoEnzyme{<:ReverseMode,Nothing},
-    x,
-) where {F}
-    return copyto!(jac, DI.jacobian(f, prep, backend, x))
-end
-
-function DI.value_and_jacobian!(
-    f::F,
-    jac,
-    prep::EnzymeReverseOneArgJacobianPrep,
-    backend::AutoEnzyme{<:ReverseMode,Nothing},
-    x,
-) where {F}
-    y, new_jac = DI.value_and_jacobian(f, prep, backend, x)
-    return y, copyto!(jac, new_jac)
-end
-=#