Improve Enzyme in forward mode with chunk size (#186)

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -7,7 +7,8 @@ using DifferentiationInterface:
     NoGradientExtras,
     NoJacobianExtras,
     NoPullbackExtras,
-    NoPushforwardExtras
+    NoPushforwardExtras,
+    pick_chunksize
 using DocStringExtensions
 using Enzyme:
     Active,
@@ -22,6 +23,7 @@ using Enzyme:
     ReverseMode,
     autodiff,
     autodiff_thunk,
+    chunkedonehot,
     gradient,
     gradient!,
     jacobian,

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -33,56 +33,78 @@ end
 
 ## Gradient
 
-DI.prepare_gradient(f, ::AutoForwardEnzyme, x) = NoGradientExtras()
+struct EnzymeForwardGradientExtras{C,O}
+    shadow::O
+end
+
+function DI.prepare_gradient(f, ::AutoForwardEnzyme, x)
+    C = pick_chunksize(length(x))
+    shadow = chunkedonehot(x, Val(C))
+    return EnzymeForwardGradientExtras{C,typeof(shadow)}(shadow)
+end
 
-function DI.gradient(f, backend::AutoForwardEnzyme, x::AbstractArray, ::NoGradientExtras)
-    return reshape(collect(gradient(backend.mode, f, x)), size(x))
+function DI.gradient(
+    f, backend::AutoForwardEnzyme, x, extras::EnzymeForwardGradientExtras{C}
+) where {C}
+    grad_tup = gradient(backend.mode, f, x, Val{C}(); shadow=extras.shadow)
+    return reshape(collect(grad_tup), size(x))
 end
 
 function DI.value_and_gradient(
-    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoGradientExtras
+    f, backend::AutoForwardEnzyme, x, extras::EnzymeForwardGradientExtras
 )
     return f(x), DI.gradient(f, backend, x, extras)
 end
 
 function DI.gradient!(
-    f, grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoGradientExtras
-)
-    return copyto!(grad, DI.gradient(f, backend, x, extras))
+    f, grad, backend::AutoForwardEnzyme, x, extras::EnzymeForwardGradientExtras{C}
+) where {C}
+    grad_tup = gradient(backend.mode, f, x, Val{C}(); shadow=extras.shadow)
+    return copyto!(grad, grad_tup)
 end
 
 function DI.value_and_gradient!(
-    f, grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoGradientExtras
-)
-    y, new_grad = DI.value_and_gradient(f, backend, x, extras)
-    return y, copyto!(grad, new_grad)
+    f, grad, backend::AutoForwardEnzyme, x, extras::EnzymeForwardGradientExtras{C}
+) where {C}
+    grad_tup = gradient(backend.mode, f, x, Val{C}(); shadow=extras.shadow)
+    return f(x), copyto!(grad, grad_tup)
 end
 
 ## Jacobian
 
-DI.prepare_jacobian(f, ::AutoForwardEnzyme, x) = NoJacobianExtras()
+struct EnzymeForwardOneArgJacobianExtras{C,O}
+    shadow::O
+end
+
+function DI.prepare_jacobian(f, ::AutoForwardEnzyme, x)
+    C = pick_chunksize(length(x))
+    shadow = chunkedonehot(x, Val(C))
+    return EnzymeForwardOneArgJacobianExtras{C,typeof(shadow)}(shadow)
+end
 
-function DI.jacobian(f, backend::AutoForwardEnzyme, x::AbstractArray, ::NoJacobianExtras)
-    jac_wrongshape = jacobian(backend.mode, f, x)
+function DI.jacobian(
+    f, backend::AutoForwardEnzyme, x, extras::EnzymeForwardOneArgJacobianExtras{C}
+) where {C}
+    jac_wrongshape = jacobian(backend.mode, f, x, Val{C}(); shadow=extras.shadow)
     nx = length(x)
     ny = length(jac_wrongshape) ÷ length(x)
     return reshape(jac_wrongshape, ny, nx)
 end
 
 function DI.value_and_jacobian(
-    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoJacobianExtras
+    f, backend::AutoForwardEnzyme, x, extras::EnzymeForwardOneArgJacobianExtras
 )
     return f(x), DI.jacobian(f, backend, x, extras)
 end
 
 function DI.jacobian!(
-    f, jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoJacobianExtras
+    f, jac, backend::AutoForwardEnzyme, x, extras::EnzymeForwardOneArgJacobianExtras
 )
     return copyto!(jac, DI.jacobian(f, backend, x, extras))
 end
 
 function DI.value_and_jacobian!(
-    f, jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoJacobianExtras
+    f, jac, backend::AutoForwardEnzyme, x, extras::EnzymeForwardOneArgJacobianExtras
 )
     y, new_jac = DI.value_and_jacobian(f, backend, x, extras)
     return y, copyto!(jac, new_jac)

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -84,3 +84,63 @@ end
 function DI.value_and_gradient!(f, grad, backend::AutoReverseEnzyme, x, ::NoGradientExtras)
     return DI.value_and_pullback!(f, grad, backend, x, one(eltype(x)), NoPullbackExtras())
 end
+
+## Jacobian
+
+# see https://github.com/EnzymeAD/Enzyme.jl/issues/1391
+
+#=
+
+struct EnzymeReverseOneArgJacobianExtras{C,N} end
+
+function DI.prepare_jacobian(f, ::AutoReverseEnzyme, x)
+    C = pick_chunksize(length(x))
+    y = f(x)
+    N = length(y)
+    return EnzymeReverseOneArgJacobianExtras{C,N}()
+end
+
+function DI.jacobian(
+    f,
+    backend::AutoReverseEnzyme,
+    x::AbstractArray,
+    ::EnzymeReverseOneArgJacobianExtras{C,N},
+) where {C,N}
+    jac_wrongshape = jacobian(backend.mode, f, x, Val{N}(), Val{C}())
+    nx = length(x)
+    ny = length(jac_wrongshape) ÷ length(x)
+    jac_rightshape = reshape(jac_wrongshape, ny, nx)
+    return jac_rightshape
+end
+
+function DI.value_and_jacobian(
+    f,
+    backend::AutoReverseEnzyme,
+    x::AbstractArray,
+    extras::EnzymeReverseOneArgJacobianExtras,
+)
+    return f(x), DI.jacobian(f, backend, x, extras)
+end
+
+function DI.jacobian!(
+    f,
+    jac,
+    backend::AutoReverseEnzyme,
+    x::AbstractArray,
+    extras::EnzymeReverseOneArgJacobianExtras,
+)
+    return copyto!(jac, DI.jacobian(f, backend, x, extras))
+end
+
+function DI.value_and_jacobian!(
+    f,
+    jac,
+    backend::AutoReverseEnzyme,
+    x::AbstractArray,
+    extras::EnzymeReverseOneArgJacobianExtras,
+)
+    y, new_jac = DI.value_and_jacobian(f, backend, x, extras)
+    return y, copyto!(jac, new_jac)
+end
+
+=#

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -57,6 +57,7 @@ include("hessian.jl")
 
 include("check.jl")
 include("sparse.jl")
+include("chunk.jl")
 
 export AutoChainRules,
     AutoDiffractor,

DifferentiationInterface/src/chunk.jl

@@ -0,0 +1,22 @@
+#=
+This heuristic is taken from ForwardDiff.jl.
+Source file: https://github.com/JuliaDiff/ForwardDiff.jl/blob/master/src/prelude.jl
+=#
+
+const DEFAULT_CHUNKSIZE = 8
+
+"""
+    pick_chunksize(input_length)
+
+Pick a reasonable chunk size for chunked derivative evaluation with an input of length `input_length`.
+    
+The result cannot be larger than `DEFAULT_CHUNKSIZE=$DEFAULT_CHUNKSIZE`.
+"""
+function pick_chunksize(input_length::Integer; threshold::Integer=DEFAULT_CHUNKSIZE)
+    if input_length <= threshold
+        return input_length
+    else
+        nchunks = round(Int, input_length / threshold, RoundUp)
+        return round(Int, input_length / nchunks, RoundUp)
+    end
+end

DifferentiationInterface/test/chunk.jl

@@ -0,0 +1,9 @@
+using DifferentiationInterface: pick_chunksize, DEFAULT_CHUNKSIZE
+
+@test pick_chunksize.(1:DEFAULT_CHUNKSIZE) == 1:DEFAULT_CHUNKSIZE
+@test all(
+    pick_chunksize.((DEFAULT_CHUNKSIZE + 1):(5DEFAULT_CHUNKSIZE)) .<= DEFAULT_CHUNKSIZE
+)
+@test all(
+    pick_chunksize.((DEFAULT_CHUNKSIZE + 1):(5DEFAULT_CHUNKSIZE)) .>= DEFAULT_CHUNKSIZE / 2
+)

DifferentiationInterface/test/runtests.jl

@@ -43,5 +43,9 @@ include("test_imports.jl")
         @testset "Weird arrays" begin
             include("weird_arrays.jl")
         end
+
+        @testset "Chunks" begin
+            include("chunk.jl")
+        end
     end
 end;

DifferentiationInterface/test/type_stability.jl

@@ -1,4 +1,6 @@
-type_stable_backends = [AutoForwardDiff(), AutoEnzyme(Enzyme.Reverse)]
+type_stable_backends = [
+    AutoForwardDiff(), AutoEnzyme(Enzyme.Forward), AutoEnzyme(Enzyme.Reverse)
+]
 
 test_differentiation(
     type_stable_backends;