From 9d7940d3719e3db34a490f55029db4486bfedb89 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 7 May 2026 14:01:58 +0200 Subject: [PATCH 1/2] add Enzyme w fwdsplit support --- Project.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Project.toml b/Project.toml index 3184c1e7..02230569 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,9 @@ Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +[sources] +Enzyme = {url = "https://github.com/EnzymeAD/Enzyme.jl", rev = "vc/fwdsplit"} + [compat] Enzyme = "0.13.50" Krylov = "0.10.1" From f732abcdf4f604ae7485aa410761e5b86c1e6e2d Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 7 May 2026 14:13:41 +0200 Subject: [PATCH 2/2] First attempt at SplitOperator --- src/Ariadne.jl | 116 ++++++++++++++++++++++++++++++++ test/split_jacobian.jl | 149 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 265 insertions(+) create mode 100644 test/split_jacobian.jl diff --git a/src/Ariadne.jl b/src/Ariadne.jl index 8e7ffe59..bd2481b1 100644 --- a/src/Ariadne.jl +++ b/src/Ariadne.jl @@ -104,6 +104,111 @@ function mul!(out, J′::Union{Adjoint{<:Any, <:JacobianOperator}, Transpose{<:A return nothing end +## +# SplitJacobianOperator - Uses Enzyme's ForwardModeSplit for efficient Newton-Krylov +## + +""" + SplitJacobianOperator + +Uses Enzyme's ForwardModeSplit to cache the primal computation and avoid +re-running it on every Krylov iteration. The primal `f!(res, u, p)` is +computed once during `prepare!` and the tape is reused for all JVP queries. +""" +struct SplitJacobianOperator{F, F′, A, P, P′, FwdThunk, DerivThunk} <: AbstractJacobianOperator + f::F # f!(res, u, p) + f′::F′ # shadow for f (nothing if guaranteed_const) + res::A # residual buffer (shared with workspace) + res′::A # shadow for res during prepare! (zero tangent; allocated once) + u::A # state + u′::A # shadow for u during prepare! (zero tangent; allocated once) + p::P + p′::P′ # shadow for p (nothing if p is const) + fwd_thunk::FwdThunk + deriv_thunk::DerivThunk + tape::Base.RefValue{Any} # tape stored after prepare! +end + +""" + SplitJacobianOperator(f::F, res::A, u::A, p::P; assume_p_const=false) + +Creates a SplitJacobianOperator using Enzyme's ForwardModeSplit mode. +Requires Enzyme with ForwardSplitNoPrimal support (PR #3024). +""" +function SplitJacobianOperator(f::F, res::A, u::A, p::P; assume_p_const = false) where {F, A, P} + f_const = Enzyme.Compiler.guaranteed_const(F) + p_const = assume_p_const || Enzyme.Compiler.guaranteed_const(P) + + f′ = f_const ? nothing : Enzyme.make_zero(f) + p′ = p_const ? nothing : Enzyme.make_zero(p) + res′ = Enzyme.make_zero(res) + u′ = Enzyme.make_zero(u) + + FA = f_const ? Const{F} : Duplicated{F} + PA = p_const ? Const{P} : Duplicated{P} + + # Check if ForwardSplitNoPrimal is available + if !isdefined(Enzyme, :ForwardSplitNoPrimal) + error("SplitJacobianOperator requires Enzyme with ForwardSplitNoPrimal support (PR #3024)") + end + + fwd_thunk, deriv_thunk = Enzyme.autodiff_thunk( + Enzyme.ForwardSplitNoPrimal, FA, Const, + Duplicated{A}, Duplicated{A}, PA + ) + + return SplitJacobianOperator( + f, f′, res, res′, u, u′, p, p′, + fwd_thunk, deriv_thunk, Ref{Any}(nothing) + ) +end + +""" + prepare!(J::SplitJacobianOperator) + +Runs the augmented forward pass to populate `J.res` with `f!(J.res, J.u, J.p)` +and stores the tape for subsequent derivative passes. +""" +function prepare!(J::SplitJacobianOperator) + # res′ and u′ stay zero → tangent for the augmented pass is zero + # (we only want the primal result in res and the tape; not a JVP) + result = J.fwd_thunk( + maybe_duplicated(J.f, J.f′), + Duplicated(J.res, J.res′), + Duplicated(J.u, J.u′), + maybe_duplicated(J.p, J.p′), + ) + # Handle both (tape,) and tape return formats defensively + J.tape[] = result isa Tuple ? first(result) : result + return nothing +end + +batch_size(::SplitJacobianOperator) = 1 + +Base.size(J::SplitJacobianOperator) = (length(J.res), length(J.u)) +Base.eltype(J::SplitJacobianOperator) = eltype(J.u) +Base.length(J::SplitJacobianOperator) = prod(size(J)) + +function mul!(out, J::SplitJacobianOperator, v) + J.deriv_thunk( + maybe_duplicated(J.f, J.f′), + Duplicated(J.res, reshape(out, size(J.res))), + Duplicated(J.u, reshape(v, size(J.u))), + maybe_duplicated(J.p, J.p′), + J.tape[], + ) + return nothing +end + +LinearAlgebra.adjoint(J::SplitJacobianOperator) = Adjoint(J) +LinearAlgebra.transpose(J::SplitJacobianOperator) = Transpose(J) + +# Note: Adjoint/transpose operations for SplitJacobianOperator would need reverse mode +# which is not part of ForwardModeSplit. For now, these are not implemented. +function mul!(out, J′::Union{Adjoint{<:Any, <:SplitJacobianOperator}, Transpose{<:Any, <:SplitJacobianOperator}}, v) + error("Adjoint/transpose operations not yet implemented for SplitJacobianOperator") +end + function init_cache(x, ::Val{N}) where {N} if !Enzyme.Compiler.guaranteed_const(typeof(x)) @@ -383,6 +488,17 @@ function evaluate!(ws::NewtonKrylovWorkspace) return norm(ws.res) end +""" + Ariadne.evaluate!(ws::NewtonKrylovWorkspace{<:Any,<:Any,<:Any,<:SplitJacobianOperator}) -> norm_res + +Specialized evaluate! for SplitJacobianOperator that calls `prepare!` instead of +evaluating the function directly. This both populates `ws.res` and caches the tape. +""" +function evaluate!(ws::NewtonKrylovWorkspace{<:Any, <:Any, <:Any, <:SplitJacobianOperator}) + prepare!(ws.J) + return norm(ws.res) +end + ## # LineSearches ## diff --git a/test/split_jacobian.jl b/test/split_jacobian.jl new file mode 100644 index 00000000..49bb93b2 --- /dev/null +++ b/test/split_jacobian.jl @@ -0,0 +1,149 @@ +using Test +using Ariadne +import Ariadne: SplitJacobianOperator, prepare! +using Enzyme, LinearAlgebra + +# Skip tests if ForwardSplitNoPrimal is not available +const FORWARD_SPLIT_AVAILABLE = isdefined(Enzyme, :ForwardSplitNoPrimal) + +function F!(res, x, _) + res[1] = x[1]^2 + x[2]^2 - 2 + res[2] = exp(x[1] - 1) + x[2]^2 - 2 + return nothing +end + +function F(x, p) + res = similar(x) + F!(res, x, p) + return res +end + +@testset "SplitJacobianOperator" begin + if !FORWARD_SPLIT_AVAILABLE + @test_skip "ForwardSplitNoPrimal not available, skipping SplitJacobianOperator tests" + return + end + + @testset "constructor" begin + x = [3.0, 5.0] + res = zeros(2) + p = nothing + + J_split = SplitJacobianOperator(F!, res, x, p) + + @test size(J_split) == (2, 2) + @test length(J_split) == 4 + @test eltype(J_split) == Float64 + @test J_split.tape[] === nothing # not prepared yet + end + + @testset "prepare! and mul!" begin + x = [3.0, 5.0] + res = zeros(2) + p = nothing + + J_split = SplitJacobianOperator(F!, res, x, p) + + # Before prepare!, res should be zeros + @test res == [0.0, 0.0] + + # After prepare!, res should contain F!(res, x, p) and tape should be stored + prepare!(J_split) + expected_res = F(x, p) + @test res ≈ expected_res + @test J_split.tape[] !== nothing + + # Test JVP computation + v = [1.0, 0.0] + out = zeros(2) + mul!(out, J_split, v) + + # Compare with regular JacobianOperator result + J_regular = Ariadne.JacobianOperator(F!, copy(res), x, p) + out_regular = zeros(2) + mul!(out_regular, J_regular, v) + + @test out ≈ out_regular + end + + @testset "consistency with JacobianOperator" begin + x = [3.0, 5.0] + res = zeros(2) + p = nothing + + # Create both operators + J_split = SplitJacobianOperator(F!, res, x, p) + prepare!(J_split) + + res_regular = copy(res) + J_regular = Ariadne.JacobianOperator(F!, res_regular, x, p) + + # Test on multiple directions + for v in ([1.0, 0.0], [0.0, 1.0], [1.0, 1.0], rand(2)) + out_split = zeros(2) + out_regular = zeros(2) + + mul!(out_split, J_split, v) + mul!(out_regular, J_regular, v) + + @test out_split ≈ out_regular rtol = 1.0e-12 + end + end + + @testset "collect" begin + x = [3.0, 5.0] + res = zeros(2) + p = nothing + + J_split = SplitJacobianOperator(F!, res, x, p) + prepare!(J_split) + + # Compare collected matrix with Enzyme jacobian + J_matrix_split = collect(J_split) + J_matrix_enz = jacobian(Forward, x -> F(x, nothing), x) |> only + + @test J_matrix_split ≈ J_matrix_enz rtol = 1.0e-12 + end + + @testset "adjoint/transpose error" begin + x = [3.0, 5.0] + res = zeros(2) + p = nothing + + J_split = SplitJacobianOperator(F!, res, x, p) + prepare!(J_split) + + v = [1.0, 0.0] + out = zeros(2) + + # Adjoint and transpose should throw errors + @test_throws ErrorException mul!(out, adjoint(J_split), v) + @test_throws ErrorException mul!(out, transpose(J_split), v) + end + + @testset "parameter handling" begin + # Test with actual parameters + function F_param!(res, x, p) + res[1] = x[1]^2 + x[2]^2 - p[1] + res[2] = exp(x[1] - 1) + x[2]^2 - p[2] + return nothing + end + + x = [3.0, 5.0] + res = zeros(2) + p = [2.0, 2.0] + + # Test with assume_p_const=true + J_split = SplitJacobianOperator(F_param!, res, x, p; assume_p_const = true) + @test J_split.p′ === nothing + + prepare!(J_split) + @test res ≈ [3.0^2 + 5.0^2 - 2.0, exp(3.0 - 1) + 5.0^2 - 2.0] + + # Test with assume_p_const=false (if p is not guaranteed const by Enzyme) + if !Enzyme.Compiler.guaranteed_const(typeof(p)) + J_split2 = SplitJacobianOperator(F_param!, copy(res), x, p; assume_p_const = false) + @test J_split2.p′ !== nothing + end + end +end