add DifferentionInterface based JacobianOperator

Project.toml

@@ -4,12 +4,14 @@ authors = ["Valentin Churavy <v.churavy@gmail.com>"]
 version = "0.1.3-DEV"
 
 [deps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
+DifferentiationInterface = "0.7.14"
 Enzyme = "0.13.50"
 Krylov = "0.10.1"
 LinearAlgebra = "1.10"

src/Ariadne.jl

@@ -4,204 +4,28 @@ export newton_krylov, newton_krylov!, NewtonKrylovWorkspace
 
 using Krylov
 using LinearAlgebra, SparseArrays
-using Enzyme
 
 ##
 # JacobianOperator
 ##
 import LinearAlgebra: mul!
 
-function init_cache(x)
-    if !Enzyme.Compiler.guaranteed_const(typeof(x))
-        Enzyme.make_zero(x)
-    else
-        return nothing
-    end
-end
-
-function maybe_duplicated(x::T, x′::Union{Nothing, T}) where {T}
-    if x′ === nothing
-        return Const(x)
-    else
-        Enzyme.remake_zero!(x′)
-        return Duplicated(x, x′)
-    end
-end
-
 abstract type AbstractJacobianOperator end
 
+# Interface:
+# Base.size(J::AbstractJacobianOperator)
+# Base.eltype(J::AbstractJacobianOperator)
+# Base.length(J::AbstractJacobianOperator)
+# mul!(out, J::AbstractJacobianOperator, v)
+# LinearAlgebra.adjoint(J::AbstractJacobianOperator)
+# LinearAlgebra.transpose(J::AbstractJacobianOperator)
+# mul!(out, J′::Union{Adjoint{<:Any, <:AbstractJacobianOperator}, Transpose{<:Any, <:AbstractJacobianOperator}}, v)
 
-"""
-    JacobianOperator
-
-Efficient implementation of `J(f,x,p) * v` and `v * J(f, x,p)'`
-"""
-struct JacobianOperator{F, F′, A, P, P′} <: AbstractJacobianOperator
-    f::F # F!(res, u, p)
-    f′::F′ # cache
-    res::A
-    u::A
-    p::P
-    p′::P′ # cache
-end
-
-"""
-    JacobianOperator(f::F, res, u, p; assume_p_const::Bool = false)
-
-Creates a Jacobian operator for `f!(res, u, p)` where `res` is the residual,
-`u` is the state variable, and `p` are the parameters.
-
-If `assume_p_const` is `true`, the parameters `p` are assumed to be constant
-during the Jacobian computation, which can improve performance by not requiring the
-shadow for `p`.
-"""
-function JacobianOperator(f::F, res, u, p; assume_p_const::Bool = false) where {F}
-    f′ = init_cache(f)
-    if assume_p_const
-        p′ = nothing
-    else
-        p′ = init_cache(p)
-    end
-    return JacobianOperator(f, f′, res, u, p, p′)
-end
-
-batch_size(::JacobianOperator) = 1
-
-Base.size(J::JacobianOperator) = (length(J.res), length(J.u))
-Base.eltype(J::JacobianOperator) = eltype(J.u)
-Base.length(J::JacobianOperator) = prod(size(J))
-
-function mul!(out, J::JacobianOperator, v)
-    autodiff(
-        Forward,
-        maybe_duplicated(J.f, J.f′), Const,
-        Duplicated(J.res, reshape(out, size(J.res))),
-        Duplicated(J.u, reshape(v, size(J.u))),
-        maybe_duplicated(J.p, J.p′)
-    )
-    return nothing
-end
-
-LinearAlgebra.adjoint(J::JacobianOperator) = Adjoint(J)
-LinearAlgebra.transpose(J::JacobianOperator) = Transpose(J)
-
-# Jᵀ(y, u) = ForwardDiff.gradient!(y, x -> dot(F(x), u), xk)
-# or just reverse mode
-
-function mul!(out, J′::Union{Adjoint{<:Any, <:JacobianOperator}, Transpose{<:Any, <:JacobianOperator}}, v)
-    J = parent(J′)
-    # TODO: provide cache for `copy(v)`
-    # Enzyme zeros input derivatives and that confuses the solvers.
-    # If `out` is non-zero we might get spurious gradients
-    fill!(out, 0)
-    autodiff(
-        Reverse,
-        maybe_duplicated(J.f, J.f′), Const,
-        Duplicated(J.res, reshape(copy(v), size(J.res))),
-        Duplicated(J.u, reshape(out, size(J.u))),
-        maybe_duplicated(J.p, J.p′)
-    )
-    return nothing
-end
-
-
-function init_cache(x, ::Val{N}) where {N}
-    if !Enzyme.Compiler.guaranteed_const(typeof(x))
-        return ntuple(_ -> Enzyme.make_zero(x), Val(N))
-    else
-        return nothing
-    end
-end
-
-function maybe_duplicated(x::T, x′::Union{Nothing, NTuple{N, T}}, ::Val{N}) where {T, N}
-    if x′ === nothing
-        return Const(x)
-    else
-        Enzyme.remake_zero!(x′)
-        return BatchDuplicated(x, x′)
-    end
-end
-
-"""
-    BatchedJacobianOperator{N}
-
-
-"""
-struct BatchedJacobianOperator{N, F, A, P} <: AbstractJacobianOperator
-    f::F # F!(res, u, p)
-    f′::Union{Nothing, NTuple{N, F}} # cache
-    res::A
-    u::A
-    p::P
-    p′::Union{Nothing, NTuple{N, P}} # cache
-    function BatchedJacobianOperator{N}(f::F, res, u, p) where {F, N}
-        f′ = init_cache(f, Val(N))
-        p′ = init_cache(p, Val(N))
-        return new{N, F, typeof(u), typeof(p)}(f, f′, res, u, p, p′)
-    end
-end
-
-batch_size(::BatchedJacobianOperator{N}) where {N} = N
-
-Base.size(J::BatchedJacobianOperator) = (length(J.res), length(J.u))
-Base.eltype(J::BatchedJacobianOperator) = eltype(J.u)
-Base.length(J::BatchedJacobianOperator) = prod(size(J))
-
-LinearAlgebra.adjoint(J::BatchedJacobianOperator) = Adjoint(J)
-LinearAlgebra.transpose(J::BatchedJacobianOperator) = Transpose(J)
-
-if VERSION >= v"1.11.0"
 
-    function tuple_of_vectors(M::Matrix{T}, shape) where {T}
-        n, m = size(M)
-        return ntuple(m) do i
-            vec = Base.wrap(Array, memoryref(M.ref, (i - 1) * n + 1), (n,))
-            reshape(vec, shape)
-        end
-    end
-
-    function mul!(Out, J::BatchedJacobianOperator{N}, V) where {N}
-        @assert size(Out, 2) == size(V, 2)
-        out = tuple_of_vectors(Out, size(J.res))
-        v = tuple_of_vectors(V, size(J.u))
-
-        @assert N == length(out)
-        autodiff(
-            Forward,
-            maybe_duplicated(J.f, J.f′, Val(N)), Const,
-            BatchDuplicated(J.res, out),
-            BatchDuplicated(J.u, v),
-            maybe_duplicated(J.p, J.p′, Val(N))
-        )
-        return nothing
-    end
-
-    function mul!(Out, J′::Union{Adjoint{<:Any, <:BatchedJacobianOperator{N}}, Transpose{<:Any, <:BatchedJacobianOperator{N}}}, V) where {N}
-        J = parent(J′)
-        @assert size(Out, 2) == size(V, 2)
-
-        # If `out` is non-zero we might get spurious gradients
-        fill!(Out, 0)
-
-        # TODO: provide cache for `copy(v)`
-        # Enzyme zeros input derivatives and that confuses the solvers.
-        V = copy(V)
-
-        out = tuple_of_vectors(Out, size(J.u))
-        v = tuple_of_vectors(V, size(J.res))
-
-        @assert N == length(out)
+include("operators/enzyme.jl")
+include("operators/di.jl")
 
-        autodiff(
-            Reverse,
-            maybe_duplicated(J.f, J.f′, Val(N)), Const,
-            BatchDuplicated(J.res, v),
-            BatchDuplicated(J.u, out),
-            maybe_duplicated(J.p, J.p′, Val(N))
-        )
-        return nothing
-    end
-end # VERSION >= v"1.11.0"
+const JacobianOperator = EnzymeJacobianOperator
 
 function Base.collect(JOp::Union{Adjoint{<:Any, <:AbstractJacobianOperator}, Transpose{<:Any, <:AbstractJacobianOperator}, AbstractJacobianOperator})
     N, M = size(JOp)
@@ -359,15 +183,25 @@ struct NewtonKrylovWorkspace{F, A, P, JOp <: AbstractJacobianOperator, KW}
 end
 
 function NewtonKrylovWorkspace(
-        F!, u::AbstractArray, p, res::AbstractArray, ::Val{Algo} = Val(:gmres);
-        assume_p_const::Bool = false
+        F!, u::AbstractArray, p, res::AbstractArray, ::Val{Algo} = Val(:gmres), operator::Type{Op} = JacobianOperator;
+        assume_p_const::Bool = false,
+        backend = nothing,
     ) where {Algo}
     # res .= 0 might ignore ghost cells
     # memory allocated with similar might contain NaN/Inf
     Enzyme.make_zero!(res)
     neg_res = similar(res)
     Enzyme.make_zero!(neg_res)
-    J = JacobianOperator(F!, res, u, p; assume_p_const)
+    if operator <: DIJacobianOperator
+        if backend === nothing
+            error("DIJacobianOperator requires a differentiation backend (e.g. backend = ADTypes.AutoEnzyme())")
+        end
+        J = operator(backend, F!, res, u, p)
+    elseif operator <: EnzymeJacobianOperator
+        J = operator(F!, res, u, p; assume_p_const)
+    else
+        error("Unknown Jacobian operator type: $operator")
+    end
     kc = KrylovConstructor(res)
     krylov = krylov_workspace(Val(Algo), kc)
     return NewtonKrylovWorkspace(F!, u, res, neg_res, p, J, krylov)
@@ -438,9 +272,16 @@ Takes an in-place residual function `F!(res, u, p)`.
 
 $(KWARGS_DOCS)
 """
-function newton_krylov!(F!, u₀::AbstractArray, p = nothing, M::Int = length(u₀); algo::Symbol = :gmres, assume_p_const::Bool = false, kwargs...)
+function newton_krylov!(
+        F!, u₀::AbstractArray, p = nothing, M::Int = length(u₀);
+        algo::Symbol = :gmres,
+        assume_p_const::Bool = false,
+        backend = nothing,
+        operator::Type{<:AbstractJacobianOperator} = JacobianOperator,
+        kwargs...
+    )
     res = similar(u₀, M)
-    ws = NewtonKrylovWorkspace(F!, u₀, p, res, Val(algo); assume_p_const)
+    ws = NewtonKrylovWorkspace(F!, u₀, p, res, Val(algo), operator; assume_p_const, backend)
     return newton_krylov!(ws; kwargs...)
 end
 
@@ -461,9 +302,11 @@ function newton_krylov!(
         F!, u::AbstractArray, p, res::AbstractArray;
         algo::Symbol = :gmres,
         assume_p_const::Bool = false,
+        backend = nothing,
+        operator::Type{<:AbstractJacobianOperator} = JacobianOperator,
         kwargs...,
     )
-    ws = NewtonKrylovWorkspace(F!, u, p, res, Val(algo); assume_p_const)
+    ws = NewtonKrylovWorkspace(F!, u, p, res, Val(algo), operator; assume_p_const, backend)
     return newton_krylov!(ws; kwargs...)
 end
 

src/operators/di.jl

@@ -0,0 +1,44 @@
+import DifferentiationInterface as DI
+
+"""
+    DIJacobianOperator
+"""
+struct DIJacobianOperator{F, A, P} <: AbstractJacobianOperator
+    f::F # F!(res, u, p)
+    res::A
+    u::A
+    p::P
+    prep
+    backend
+end
+
+"""
+    DIJacobianOperator(f::F, res, u, p)
+
+Creates a Jacobian operator for `f!(res, u, p)` where `res` is the residual,
+`u` is the state variable, and `p` are the parameters.
+"""
+function DIJacobianOperator(backend, f::F, res, u, p) where {F}
+    tu = zero(u) # dummy tangent
+    prep = DI.prepare_pushforward(f, res, backend, u, (tu,), DI.ConstantOrCache(p))
+
+    return DIJacobianOperator(f, res, u, p, prep, backend)
+end
+
+Base.size(J::DIJacobianOperator) = (length(J.res), length(J.u))
+Base.eltype(J::DIJacobianOperator) = eltype(J.u)
+Base.length(J::DIJacobianOperator) = prod(size(J))
+
+function mul!(out, J::DIJacobianOperator, v)
+    DI.pushforward!(
+        J.f,
+        J.res,
+        (out,),
+        J.prep,
+        J.backend,
+        J.u,
+        (v,), # TODO: Must we zero this?
+        DI.ConstantOrCache(J.p)
+    )
+    return nothing
+end