More batch mode for Enzyme (#495)

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -21,6 +21,31 @@ function seeded_autodiff_thunk(
     end
 end
 
+function batch_seeded_autodiff_thunk(
+    rmode::ReverseModeSplit{ReturnPrimal},
+    dresults::NTuple,
+    f::FA,
+    ::Type{RA},
+    args::Vararg{Annotation,N},
+) where {ReturnPrimal,FA<:Annotation,RA<:Annotation,N}
+    forward, reverse = autodiff_thunk(rmode, FA, RA, typeof.(args)...)
+    tape, result, shadow_results = forward(f, args...)
+    if RA <: Active
+        dresults_righttype = map(Fix1(convert, typeof(result)), dresults)
+        dinputs = only(reverse(f, args..., dresults_righttype, tape))
+    else
+        foreach(shadow_results, dresults) do d0, d
+            d0 .+= d  # use recursive_add here?
+        end
+        dinputs = only(reverse(f, args..., tape))
+    end
+    if ReturnPrimal
+        return (dinputs, result)
+    else
+        return (dinputs,)
+    end
+end
+
 ## Pullback
 
 function DI.prepare_pullback(
@@ -35,24 +60,6 @@ end
 
 ### Out-of-place
 
-function DI.value_and_pullback(
-    f::F,
-    prep::NoPullbackPrep,
-    backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
-    x,
-    ty::Tangents,
-    contexts::Vararg{Context,C},
-) where {F,C}
-    ys_and_dxs = map(ty.d) do dy
-        y, tx = DI.value_and_pullback(f, prep, backend, x, Tangents(dy), contexts...)
-        y, only(tx)
-    end
-    y = first(ys_and_dxs[1])
-    dxs = last.(ys_and_dxs)
-    tx = Tangents(dxs...)
-    return y, tx
-end
-
 function DI.value_and_pullback(
     f::F,
     ::NoPullbackPrep,
@@ -70,6 +77,24 @@ function DI.value_and_pullback(
     return result, Tangents(first(dinputs))
 end
 
+function DI.value_and_pullback(
+    f::F,
+    prep::NoPullbackPrep,
+    backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
+    x::Number,
+    ty::Tangents{B},
+    contexts::Vararg{Context,C},
+) where {F,B,C}
+    # TODO: improve
+    ys_and_dxs = map(ty.d) do dy
+        y, tx = DI.value_and_pullback(f, prep, backend, x, Tangents(dy), contexts...)
+        y, only(tx)
+    end
+    y = first(ys_and_dxs[1])
+    dxs = last.(ys_and_dxs)
+    return y, Tangents(dxs...)
+end
+
 function DI.value_and_pullback(
     f::F,
     ::NoPullbackPrep,
@@ -88,53 +113,37 @@ function DI.value_and_pullback(
     return result, Tangents(dx)
 end
 
-function DI.pullback(
-    f::F,
-    prep::NoPullbackPrep,
-    backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
-    x::Number,
-    ty::Tangents{1},
-    contexts::Vararg{Context,C},
-) where {F,C}
-    return last(DI.value_and_pullback(f, prep, backend, x, ty, contexts...))
-end
-
-### In-place
-
-function DI.value_and_pullback!(
+function DI.value_and_pullback(
     f::F,
-    tx::Tangents,
-    prep::NoPullbackPrep,
+    ::NoPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
-    ty::Tangents,
+    ty::Tangents{B},
     contexts::Vararg{Context,C},
-) where {F,C}
-    ys = map(tx.d, ty.d) do dx, dy
-        y, _ = DI.value_and_pullback!(
-            f, Tangents(dx), prep, backend, x, Tangents(dy), contexts...
-        )
-        y
-    end
-    y = first(ys)
-    return y, tx
+) where {F,B,C}
+    f_and_df = force_annotation(get_f_and_df(f, backend, Val(B)))
+    mode = reverse_mode_split_withprimal(backend)
+    RA = eltype(ty) <: Number ? Active : BatchDuplicated
+    dxs = ntuple(_ -> make_zero(x), Val(B))
+    _, result = batch_seeded_autodiff_thunk(
+        mode, NTuple(ty), f_and_df, RA, BatchDuplicated(x, dxs), map(translate, contexts)...
+    )
+    return result, Tangents(dxs...)
 end
 
-function DI.pullback!(
+function DI.pullback(
     f::F,
-    tx::Tangents,
     prep::NoPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     ty::Tangents,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    for b in eachindex(tx.d, ty.d)
-        DI.pullback!(f, Tangents(tx.d[b]), prep, backend, x, Tangents(ty.d[b]), contexts...)
-    end
-    return tx
+    return last(DI.value_and_pullback(f, prep, backend, x, ty, contexts...))
 end
 
+### In-place
+
 function DI.value_and_pullback!(
     f::F,
     tx::Tangents{1},
@@ -161,13 +170,39 @@ function DI.value_and_pullback!(
     return result, tx
 end
 
+function DI.value_and_pullback!(
+    f::F,
+    tx::Tangents{B},
+    ::NoPullbackPrep,
+    backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
+    x,
+    ty::Tangents{B},
+    contexts::Vararg{Context,C},
+) where {F,B,C}
+    f_and_df = force_annotation(get_f_and_df(f, backend, Val(B)))
+    mode = reverse_mode_split_withprimal(backend)
+    RA = eltype(ty) <: Number ? Active : BatchDuplicated
+    dxs_righttype = map(Fix1(convert, typeof(x)), NTuple(tx))
+    make_zero!(dxs_righttype)
+    _, result = batch_seeded_autodiff_thunk(
+        mode,
+        NTuple(ty),
+        f_and_df,
+        RA,
+        BatchDuplicated(x, dxs_righttype),
+        map(translate, contexts)...,
+    )
+    foreach(copyto!, NTuple(tx), dxs_righttype)
+    return result, tx
+end
+
 function DI.pullback!(
     f::F,
-    tx::Tangents{1},
+    tx::Tangents,
     prep::NoPullbackPrep,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
-    ty::Tangents{1},
+    ty::Tangents,
     contexts::Vararg{Context,C},
 ) where {F,C}
     return last(DI.value_and_pullback!(f, tx, prep, backend, x, ty, contexts...))

DifferentiationInterface/src/utils/tangents.jl

@@ -10,9 +10,12 @@ pick_batchsize(::AbstractADType, dimension::Integer) = 1
 """
     Tangents{B}
 
-Storage for a batch of `B` tangents (behaves like an `NTuple`).
+Storage for a batch of `B` tangents (wrapper around an `NTuple`).
 
-Must be passed as an argument to [`pushforward`](@ref), [`pullback`](@ref) and [`hvp`](@ref), in addition to the input `x`.
+The operators [`pushforward`](@ref), [`pullback`](@ref) and [`hvp`](@ref) require a `Tangents` argument in addition to the input `x`.
+
+The underlying `NTuple` of `t::Tangents` can be retrieved with `NTuple(t)`.
+We also define a few utility functions, as shown below.
 
 # Constructors
 
@@ -27,6 +30,9 @@ julia> using DifferentiationInterface
 julia> t = Tangents(2.0)
 Tangents{1, Float64}((2.0,))
 
+julia> NTuple(t)
+(2.0,)
+
 julia> length(t)
 1
 
@@ -36,6 +42,9 @@ julia> only(t)
 julia> t = Tangents([2.0], [4.0], [6.0])
 Tangents{3, Vector{Float64}}(([2.0], [4.0], [6.0]))
 
+julia> NTuple(t)
+([2.0], [4.0], [6.0])
+
 julia> length(t)
 3
 
@@ -59,25 +68,25 @@ end
 Base.length(::Tangents{B,T}) where {B,T} = B
 Base.eltype(::Tangents{B,T}) where {B,T} = T
 
-Base.only(t::Tangents) = only(t.d)
-Base.getindex(t::Tangents, ind) = t.d[ind]
-Base.firstindex(t::Tangents) = firstindex(t.d)
-Base.lastindex(t::Tangents) = lastindex(t.d)
+Base.NTuple(t::Tangents) = t.d
 
-Base.iterate(t::Tangents) = iterate(t.d)
-Base.iterate(t::Tangents, state) = iterate(t.d, state)
+Base.only(t::Tangents) = only(NTuple(t))
+Base.getindex(t::Tangents, ind) = NTuple(t)[ind]
+Base.firstindex(t::Tangents) = firstindex(NTuple(t))
+Base.lastindex(t::Tangents) = lastindex(NTuple(t))
 
-Base.map(f, t::Tangents) = Tangents(map(f, t.d)...)
+Base.iterate(t::Tangents) = iterate(NTuple(t))
+Base.iterate(t::Tangents, state) = iterate(NTuple(t), state)
 
-Base.:(==)(t1::Tangents{B}, t2::Tangents{B}) where {B} = t1.d == t2.d
+Base.map(f, t::Tangents) = Tangents(map(f, NTuple(t))...)
+
+Base.:(==)(t1::Tangents{B}, t2::Tangents{B}) where {B} = NTuple(t1) == NTuple(t2)
 
 function Base.isapprox(t1::Tangents{B}, t2::Tangents{B}; kwargs...) where {B}
-    return all(isapprox.(t1.d, t2.d; kwargs...))
+    return all(isapprox.(NTuple(t1), NTuple(t2); kwargs...))
 end
 
 function Base.copyto!(t1::Tangents{B}, t2::Tangents{B}) where {B}
-    for b in eachindex(t1.d, t2.d)
-        copyto!(t1.d[b], t2.d[b])
-    end
+    foreach(copyto!, NTuple(t1), NTuple(t2))
     return t1
 end

DifferentiationInterface/test/Internals/tangents.jl

@@ -4,12 +4,14 @@ using Test
 @test_throws ArgumentError Tangents()
 
 t = Tangents([2.0])
+@test NTuple(t) == ([2.0],)
 @test length(t) == 1
 @test eltype(t) == Vector{Float64}
 @test only(t) == [2.0]
 @test copyto!(map(zero, t), t) ≈ t
 
 t = Tangents(2.0, 4.0, 6.0)
+@test NTuple(t) == (2.0, 4.0, 6.0)
 @test length(t) == 3
 @test eltype(t) == Float64
 @test t[begin] == first(t) == 2.0