diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index bb96ff54f..821820830 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -95,6 +95,7 @@ jobs: - DifferentiateWith # - Diffractor - Enzyme + - EnzymeReactant - FastDifferentiation - FiniteDiff - FiniteDifferences diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index a2a694816..008f65019 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -23,6 +23,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" @@ -49,6 +50,7 @@ DifferentiationInterfacePolyesterForwardDiffExt = [ "ForwardDiff", "DiffResults", ] +DifferentiationInterfaceReactantExt = "Reactant" DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"] DifferentiationInterfaceSparseArraysExt = "SparseArrays" DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer" @@ -76,6 +78,7 @@ HyperHessians = "0.3" LinearAlgebra = "1" Mooncake = "0.5.25" PolyesterForwardDiff = "0.1.2" +Reactant = "0.2.283" ReverseDiff = "1.15.1" SparseArrays = "1" SparseConnectivityTracer = "0.6.14, 1" diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl index c44529caf..ef4a2e182 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl @@ -216,6 +216,40 @@ end ## Jacobian +struct EnzymeReactantJacobianPrep{SIG, D, P} <: DI.JacobianPrep{SIG} + _sig::Val{SIG} + directions::D + pushforward_prep::P +end + +function prepare_reactant_jacobian( + strict::Val, + f_or_f!y::Tuple, + backend::AutoEnzyme, + x, + contexts::Vararg{DI.Context, C} + ) where {C} + _sig = DI.signature(f_or_f!y..., backend, x, contexts...; strict) + directions = onehot(x) + pushforward_prep = DI.prepare_pushforward_nokwarg( + strict, f_or_f!y..., backend, x, directions, contexts... + ) + return EnzymeReactantJacobianPrep(_sig, directions, pushforward_prep) +end + +function reactant_jacobian( + f_or_f!y::Tuple, + prep::EnzymeReactantJacobianPrep, + backend::AutoEnzyme, + x, + contexts::Vararg{DI.Context, C} + ) where {C} + columns = DI.pushforward( + f_or_f!y..., prep.pushforward_prep, backend, x, prep.directions, contexts... + ) + return DI.stack_vec_col(columns) +end + struct EnzymeForwardOneArgJacobianPrep{SIG, B, DF, DC, O} <: DI.JacobianPrep{SIG} _sig::Val{SIG} _valB::Val{B} @@ -232,6 +266,9 @@ function DI.prepare_jacobian_nokwarg( x, contexts::Vararg{DI.Constant, C} ) where {F, C} + if DI._use_reactant_jacobian(backend) + return prepare_reactant_jacobian(strict, (f,), backend, x, contexts...) + end _sig = DI.signature(f, backend, x, contexts...; strict) y = f(x, map(DI.unwrap, contexts)...) valB = to_val(DI.pick_batchsize(backend, x)) @@ -244,6 +281,29 @@ function DI.prepare_jacobian_nokwarg( ) end +function DI.jacobian( + f::F, + prep::EnzymeReactantJacobianPrep, + backend::AutoEnzyme{<:Union{ForwardMode, Nothing}}, + x, + contexts::Vararg{DI.Context, C}, + ) where {F, C} + DI.check_prep(f, prep, backend, x, contexts...) + return reactant_jacobian((f,), prep, backend, x, contexts...) +end + +function DI.jacobian!( + f::F, + jac, + prep::EnzymeReactantJacobianPrep, + backend::AutoEnzyme{<:Union{ForwardMode, Nothing}}, + x, + contexts::Vararg{DI.Context, C}, + ) where {F, C} + DI.check_prep(f, prep, backend, x, contexts...) + return copyto!(jac, reactant_jacobian((f,), prep, backend, x, contexts...)) +end + function DI.jacobian( f::F, prep::EnzymeForwardOneArgJacobianPrep{SIG, B}, diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl index 629f046a0..47075c2f9 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl @@ -114,3 +114,97 @@ function DI.pushforward!( DI.value_and_pushforward!(f!, y, ty, prep, backend, x, tx, contexts...) return ty end + +## Jacobian + +struct EnzymeReactantTwoArgJacobianPrep{SIG, D, P} <: DI.JacobianPrep{SIG} + _sig::Val{SIG} + directions::D + pushforward_prep::P +end + +function reactant_out_of_place(x, f!::F, y, contexts...) where {F} + new_y = zero(y) + f!(new_y, x, contexts...) + return new_y +end + +function prepare_reactant_twoarg_jacobian( + strict::Val, + f!::F, + y, + backend::AutoEnzyme, + x, + contexts::Vararg{DI.Context, C} + ) where {F, C} + _sig = DI.signature(f!, y, backend, x, contexts...; strict) + directions = onehot(x) + wrapped_contexts = (DI.Constant(f!), DI.Constant(y), contexts...) + pushforward_prep = DI.prepare_pushforward_nokwarg( + strict, reactant_out_of_place, backend, x, directions, wrapped_contexts... + ) + return EnzymeReactantTwoArgJacobianPrep(_sig, directions, pushforward_prep) +end + +function reactant_twoarg_jacobian( + f!::F, + y, + prep::EnzymeReactantTwoArgJacobianPrep, + backend::AutoEnzyme, + x, + contexts::Vararg{DI.Context, C} + ) where {F, C} + wrapped_contexts = (DI.Constant(f!), DI.Constant(y), contexts...) + columns = DI.pushforward( + reactant_out_of_place, + prep.pushforward_prep, + backend, + x, + prep.directions, + wrapped_contexts..., + ) + return DI.stack_vec_col(columns) +end + +function DI.prepare_jacobian_nokwarg( + strict::Val, + f!::F, + y, + backend::AutoEnzyme{<:Union{ForwardMode, Nothing}}, + x, + contexts::Vararg{DI.Context, C} + ) where {F, C} + if DI._use_reactant_jacobian(backend) + return prepare_reactant_twoarg_jacobian(strict, f!, y, backend, x, contexts...) + end + batch_size_settings = DI.pick_batchsize(backend, x) + return DI._prepare_jacobian_aux( + strict, DI.PushforwardFast(), batch_size_settings, y, (f!, y), backend, x, contexts... + ) +end + +function DI.jacobian( + f!::F, + y, + prep::EnzymeReactantTwoArgJacobianPrep, + backend::AutoEnzyme{<:Union{ForwardMode, Nothing}}, + x, + contexts::Vararg{DI.Context, C}, + ) where {F, C} + DI.check_prep(f!, y, prep, backend, x, contexts...) + return reactant_twoarg_jacobian(f!, y, prep, backend, x, contexts...) +end + +function DI.jacobian!( + f!::F, + y, + jac, + prep::EnzymeReactantTwoArgJacobianPrep, + backend::AutoEnzyme{<:Union{ForwardMode, Nothing}}, + x, + contexts::Vararg{DI.Context, C}, + ) where {F, C} + DI.check_prep(f!, y, prep, backend, x, contexts...) + new_jac = reactant_twoarg_jacobian(f!, y, prep, backend, x, contexts...) + return copyto!(jac, new_jac) +end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceReactantExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceReactantExt.jl new file mode 100644 index 000000000..d2c2e219e --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceReactantExt.jl @@ -0,0 +1,9 @@ +module DifferentiationInterfaceReactantExt + +using ADTypes: AutoEnzyme +import DifferentiationInterface as DI +using Reactant: within_compile + +DI._use_reactant_jacobian(::AutoEnzyme) = within_compile() + +end diff --git a/DifferentiationInterface/src/first_order/jacobian.jl b/DifferentiationInterface/src/first_order/jacobian.jl index 4c64b45bc..d60121e48 100644 --- a/DifferentiationInterface/src/first_order/jacobian.jl +++ b/DifferentiationInterface/src/first_order/jacobian.jl @@ -131,6 +131,8 @@ end ## Preparation +_use_reactant_jacobian(::AbstractADType) = false + abstract type StandardJacobianPrep{SIG} <: JacobianPrep{SIG} end struct PushforwardJacobianPrep{ diff --git a/DifferentiationInterface/test/Back/EnzymeReactant/Project.toml b/DifferentiationInterface/test/Back/EnzymeReactant/Project.toml new file mode 100644 index 000000000..e74d651a1 --- /dev/null +++ b/DifferentiationInterface/test/Back/EnzymeReactant/Project.toml @@ -0,0 +1,9 @@ +[deps] +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +DifferentiationInterface = { path = "../../.." } diff --git a/DifferentiationInterface/test/Back/EnzymeReactant/test.jl b/DifferentiationInterface/test/Back/EnzymeReactant/test.jl new file mode 100644 index 000000000..aca25fa71 --- /dev/null +++ b/DifferentiationInterface/test/Back/EnzymeReactant/test.jl @@ -0,0 +1,71 @@ +using DifferentiationInterface: + AutoEnzyme, + Constant, + jacobian, + jacobian!, + prepare_jacobian, + value_and_jacobian +using Enzyme: Enzyme +using Reactant: Reactant, @jit +using Test + +f(x, p) = x .^ 2 .- p +f!(y, x, p) = y .= f(x, p) + +function oop_jacobian(x, p) + backend = AutoEnzyme() + prep = prepare_jacobian(f, backend, x, Constant(p)) + return jacobian(f, prep, backend, x, Constant(p)) +end + +function oop_value_and_jacobian(x, p) + backend = AutoEnzyme() + prep = prepare_jacobian(f, backend, x, Constant(p)) + return value_and_jacobian(f, prep, backend, x, Constant(p)) +end + +function iip_jacobian(x, p) + backend = AutoEnzyme() + y = zero(x) + prep = prepare_jacobian(f!, y, backend, x, Constant(p)) + return jacobian(f!, y, prep, backend, x, Constant(p)) +end + +function iip_value_and_jacobian(x, p) + backend = AutoEnzyme() + y = zero(x) + prep = prepare_jacobian(f!, y, backend, x, Constant(p)) + return value_and_jacobian(f!, y, prep, backend, x, Constant(p)) +end + +function oop_jacobian!(x, p) + backend = AutoEnzyme() + jac = similar(x, length(x), length(x)) + prep = prepare_jacobian(f, backend, x, Constant(p)) + return jacobian!(f, jac, prep, backend, x, Constant(p)) +end + +function iip_jacobian!(x, p) + backend = AutoEnzyme() + y = zero(x) + jac = similar(x, length(x), length(x)) + prep = prepare_jacobian(f!, y, backend, x, Constant(p)) + return jacobian!(f!, y, jac, prep, backend, x, Constant(p)) +end + +@testset "AutoEnzyme Jacobian inside Reactant" begin + x = Reactant.to_rarray(Float32[1, 2]) + p = Reactant.to_rarray(Float32[3, 4]) + expected_jacobian = Float32[2 0; 0 4] + + @test @jit(oop_jacobian(x, p)) ≈ expected_jacobian + value, jac = @jit oop_value_and_jacobian(x, p) + @test value ≈ Float32[-2, 0] + @test jac ≈ expected_jacobian + @test @jit(iip_jacobian(x, p)) ≈ expected_jacobian + value, jac = @jit iip_value_and_jacobian(x, p) + @test value ≈ Float32[-2, 0] + @test jac ≈ expected_jacobian + @test @jit(oop_jacobian!(x, p)) ≈ expected_jacobian + @test @jit(iip_jacobian!(x, p)) ≈ expected_jacobian +end