perf: allocate Enzyme shadow memory during preparation

DifferentiationInterface/Project.toml

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

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -1,91 +1,104 @@
 ## Pushforward
 
+struct EnzymeOneArgPushforwardPrep{SIG,DF,DC} <: DI.PushforwardPrep{SIG}
+    _sig::Val{SIG}
+    df::DF
+    context_shadows::DC
+end
+
 function DI.prepare_pushforward_nokwarg(
     strict::Val,
     f::F,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::NTuple,
+    tx::NTuple{B},
     contexts::Vararg{DI.Context,C};
-) where {F,C}
+) where {F,C,B}
     _sig = DI.signature(f, backend, x, tx, contexts...; strict)
-    return DI.NoPushforwardPrep(_sig)
+    df = function_shadow(f, backend, Val(B))
+    mode = forward_withprimal(backend)
+    context_shadows = make_context_shadows(backend, mode, Val(B), contexts...)
+    return EnzymeOneArgPushforwardPrep(_sig, df, context_shadows)
 end
 
 function DI.value_and_pushforward(
     f::F,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeOneArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{1},
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
+    (; df, context_shadows) = prep
     mode = forward_withprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode)
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(1))
     dx = only(tx)
     x_and_dx = Duplicated(x, dx)
-    annotated_contexts = translate(backend, mode, Val(1), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(1))
     dy, y = autodiff(mode, f_and_df, x_and_dx, annotated_contexts...)
     return y, (dy,)
 end
 
 function DI.value_and_pushforward(
     f::F,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeOneArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{B},
     contexts::Vararg{DI.Context,C},
 ) where {F,B,C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
+    (; df, context_shadows) = prep
     mode = forward_withprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode, Val(B))
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
     x_and_tx = BatchDuplicated(x, tx)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     ty, y = autodiff(mode, f_and_df, x_and_tx, annotated_contexts...)
     return y, values(ty)
 end
 
 function DI.pushforward(
     f::F,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeOneArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{1},
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
+    (; df, context_shadows) = prep
     mode = forward_noprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode)
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(1))
     dx = only(tx)
     x_and_dx = Duplicated(x, dx)
-    annotated_contexts = translate(backend, mode, Val(1), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(1))
     dy = only(autodiff(mode, f_and_df, x_and_dx, annotated_contexts...))
     return (dy,)
 end
 
 function DI.pushforward(
     f::F,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeOneArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{B},
     contexts::Vararg{DI.Context,C},
 ) where {F,B,C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
+    (; df, context_shadows) = prep
     mode = forward_noprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode, Val(B))
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
     x_and_tx = BatchDuplicated(x, tx)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     ty = only(autodiff(mode, f_and_df, x_and_tx, annotated_contexts...))
     return values(ty)
 end
 
 function DI.value_and_pushforward!(
     f::F,
     ty::NTuple,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeOneArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple,
@@ -101,7 +114,7 @@
 function DI.pushforward!(
     f::F,
     ty::NTuple,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeOneArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple,
@@ -116,10 +129,12 @@
 
 ## Gradient
 
-struct EnzymeForwardGradientPrep{SIG,B,O} <: DI.GradientPrep{SIG}
+struct EnzymeForwardGradientPrep{SIG,B,DF,DC,O} <: DI.GradientPrep{SIG}
     _sig::Val{SIG}
     _valB::Val{B}
-    shadows::O
+    df::DF
+    context_shadows::DC
+    basis_shadows::O
 end
 
 function DI.prepare_gradient_nokwarg(
@@ -131,8 +146,11 @@
 ) where {F,C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
     valB = to_val(DI.pick_batchsize(backend, x))
-    shadows = create_shadows(valB, x)
-    return EnzymeForwardGradientPrep(_sig, valB, shadows)
+    df = function_shadow(f, backend, valB)
+    mode = forward_withprimal(backend)
+    context_shadows = make_context_shadows(backend, mode, valB, contexts...)
+    basis_shadows = create_shadows(valB, x)
+    return EnzymeForwardGradientPrep(_sig, valB, df, context_shadows, basis_shadows)
 end
 
 function DI.gradient(
@@ -143,11 +161,12 @@
     contexts::Vararg{DI.Constant,C},
 ) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
+    (; df, context_shadows, basis_shadows) = prep
     mode = forward_noprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     derivs = gradient(
-        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=prep.shadows
+        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=basis_shadows
     )
     return first(derivs)
 end
@@ -160,11 +179,12 @@
     contexts::Vararg{DI.Constant,C},
 ) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
+    (; df, context_shadows, basis_shadows) = prep
     mode = forward_withprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     (; derivs, val) = gradient(
-        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=prep.shadows
+        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=basis_shadows
     )
     return val, first(derivs)
 end
@@ -196,10 +216,12 @@
 
 ## Jacobian
 
-struct EnzymeForwardOneArgJacobianPrep{SIG,B,O} <: DI.JacobianPrep{SIG}
+struct EnzymeForwardOneArgJacobianPrep{SIG,B,DF,DC,O} <: DI.JacobianPrep{SIG}
     _sig::Val{SIG}
     _valB::Val{B}
-    shadows::O
+    df::DF
+    context_shadows::DC
+    basis_shadows::O
     output_length::Int
 end
 
@@ -213,8 +235,13 @@
     _sig = DI.signature(f, backend, x, contexts...; strict)
     y = f(x, map(DI.unwrap, contexts)...)
     valB = to_val(DI.pick_batchsize(backend, x))
-    shadows = create_shadows(valB, x)
-    return EnzymeForwardOneArgJacobianPrep(_sig, valB, shadows, length(y))
+    mode = forward_withprimal(backend)
+    df = function_shadow(f, backend, valB)
+    context_shadows = make_context_shadows(backend, mode, valB, contexts...)
+    basis_shadows = create_shadows(valB, x)
+    return EnzymeForwardOneArgJacobianPrep(
+        _sig, valB, df, context_shadows, basis_shadows, length(y)
+    )
 end
 
 function DI.jacobian(
@@ -225,14 +252,15 @@
     contexts::Vararg{DI.Constant,C},
 ) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
+    (; df, context_shadows, basis_shadows, output_length) = prep
     mode = forward_noprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     derivs = jacobian(
-        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=prep.shadows
+        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=basis_shadows
     )
     jac_tensor = first(derivs)
-    return maybe_reshape(jac_tensor, prep.output_length, length(x))
+    return maybe_reshape(jac_tensor, output_length, length(x))
 end
 
 function DI.value_and_jacobian(
@@ -243,14 +271,15 @@
     contexts::Vararg{DI.Constant,C},
 ) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
+    (; df, context_shadows, basis_shadows, output_length) = prep
     mode = forward_withprimal(backend)
-    f_and_df = get_f_and_df(f, backend, mode)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     (; derivs, val) = jacobian(
-        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=prep.shadows
+        mode, f_and_df, x, annotated_contexts...; chunk=Val(B), shadows=basis_shadows
     )
     jac_tensor = first(derivs)
-    return val, maybe_reshape(jac_tensor, prep.output_length, length(x))
+    return val, maybe_reshape(jac_tensor, output_length, length(x))
 end
 
 function DI.jacobian!(

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -1,63 +1,74 @@
 ## Pushforward
 
+struct EnzymeTwoArgPushforwardPrep{SIG,DF,DC} <: DI.PushforwardPrep{SIG}
+    _sig::Val{SIG}
+    df!::DF
+    context_shadows::DC
+end
+
 function DI.prepare_pushforward_nokwarg(
     strict::Val,
     f!::F,
     y,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::NTuple,
+    tx::NTuple{B},
     contexts::Vararg{DI.Context,C};
-) where {F,C}
+) where {F,B,C}
     _sig = DI.signature(f!, y, backend, x, tx, contexts...; strict)
-    return DI.NoPushforwardPrep(_sig)
+    df! = function_shadow(f!, backend, Val(B))
+    mode = forward_noprimal(backend)
+    context_shadows = make_context_shadows(backend, mode, Val(B), contexts...)
+    return EnzymeTwoArgPushforwardPrep(_sig, df!, context_shadows)
 end
 
 function DI.value_and_pushforward(
     f!::F,
     y,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeTwoArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{1},
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
+    (; df!, context_shadows) = prep
     mode = forward_noprimal(backend)
-    f!_and_df! = get_f_and_df(f!, backend, mode)
+    f!_and_df! = get_f_and_df_prepared!(df!, f!, backend, Val(1))
     dx = only(tx)
     dy = make_zero(y)
     x_and_dx = Duplicated(x, dx)
     y_and_dy = Duplicated(y, dy)
-    annotated_contexts = translate(backend, mode, Val(1), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(1))
     autodiff(mode, f!_and_df!, Const, y_and_dy, x_and_dx, annotated_contexts...)
     return y, (dy,)
 end
 
 function DI.value_and_pushforward(
     f!::F,
     y,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeTwoArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{B},
     contexts::Vararg{DI.Context,C},
 ) where {F,B,C}
     DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
+    (; df!, context_shadows) = prep
     mode = forward_noprimal(backend)
-    f!_and_df! = get_f_and_df(f!, backend, mode, Val(B))
+    f!_and_df! = get_f_and_df_prepared!(df!, f!, backend, Val(B))
     ty = ntuple(_ -> make_zero(y), Val(B))
     x_and_tx = BatchDuplicated(x, tx)
     y_and_ty = BatchDuplicated(y, ty)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     autodiff(mode, f!_and_df!, Const, y_and_ty, x_and_tx, annotated_contexts...)
     return y, ty
 end
 
 function DI.pushforward(
     f!::F,
     y,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeTwoArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple,
@@ -72,18 +83,19 @@
     f!::F,
     y,
     ty::NTuple{B},
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeTwoArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple{B},
     contexts::Vararg{DI.Context,C},
 ) where {F,B,C}
     DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
+    (; df!, context_shadows) = prep
     mode = forward_noprimal(backend)
-    f!_and_df! = get_f_and_df(f!, backend, mode, Val(B))
+    f!_and_df! = get_f_and_df_prepared!(df!, f!, backend, Val(B))
     x_and_tx = BatchDuplicated(x, tx)
     y_and_ty = BatchDuplicated(y, ty)
-    annotated_contexts = translate(backend, mode, Val(B), contexts...)
+    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
     autodiff(mode, f!_and_df!, Const, y_and_ty, x_and_tx, annotated_contexts...)
     return y, ty
 end
@@ -92,7 +104,7 @@
     f!::F,
     y,
     ty::NTuple,
-    prep::DI.NoPushforwardPrep,
+    prep::EnzymeTwoArgPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
     tx::NTuple,