Can't `--trim` with `AutoForwardDiff`, but can with `AutoFiniteDiff`

[RomeoV]

I'm trying to make a simple C-library using SimpleNonlinearSolve.jl and the Julia 1.12 --trim feature.
However, I'm running into issues with the jacobian computation when I use AutoForwardDiff.

Here's a MWE:

# main.jl
module MyModule
using DifferentiationInterface
using FiniteDiff
using ForwardDiff
using LinearAlgebra

function f!(du, u, p)
    du .= p.λ * u
end

function main(ξ=1.0)
    # J = jacobian(u->f(u, p), AutoForwardDiff(), [1.0; 1.0])
    f!_ = (du, u)->f!(du, u, (; λ=ξ))
    # autodiff = AutoFiniteDiff()
    autodiff = AutoForwardDiff()
    Jprep = prepare_jacobian(f!_, zeros(2),        autodiff, zeros(2))
    J =             jacobian(f!_, zeros(2), Jprep, autodiff, zeros(2))
    return sum(J)
end

Base.@ccallable function main(argc::Cint, argv::Ptr{Ptr{Int8}})::Cint
    cstr = unsafe_load(argv, 2)
    jlstr = unsafe_string(cstr)
    λ = parse(Float64, jlstr)
    sol = main(λ)
    println(Core.stdout, sum(sol))
    return 0
end
end

To compile:

juliaup add 1.12
juliaup override set 1.12
julia --project=. -e 'import Pkg; Pkg.add("DifferentiationInterface", "FiniteDiff", "ForwardDiff")'
julia --project=. --depwarn=error ~/.julia/juliaup/julia-1.12.0-beta4+0.x64.linux.gnu/share/julia/juliac.jl --experimental --trim=unsafe-warn --output-exe main_mwe --compile-ccallable --relative-rpath main.jl

This fails when using AutoForwardDiff, but works fine with AutoFiniteDiff.

The error

Verifier error #1: unresolved call from statement MyModule.jacobian(%new()::MyModule.var"#main##0#main##1"{Float64}, %new()::Vector{Float64}, Dif ferentiationInterface.prepare_jacobian_nokwarg($(QuoteNode(Val{true}())), %new()::MyModule.var"#main##0#main##1"{Float64}, %new()::Vector{Float64 }, $(QuoteNode(ADTypes.AutoForwardDiff{nothing, Nothing}(nothing))), %new()::Vector{Float64})::DifferentiationInterfaceForwardDiffExt.ForwardDiff TwoArgJacobianPrep{Tuple{MyModule.var"#main##0#main##1"{Float64}, Vector{Float64}, ADTypes.AutoForwardDiff{nothing, Nothing}, Vector{Float64}, Tu ple{}}, C, Tuple{}} where C<:(ForwardDiff.JacobianConfig{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{ Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{Forwa rdDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A}} where _A), $(QuoteNode(ADTypes.AutoForwardDiff{nothing, Nothing}(nothing))), %new()::Vector{Float64})::Matrix{Float64} Stacktrace: [1] main(ξ::Float64) @ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:17 [2] main(argc::Int32, argv::Ptr{Ptr{Int8}}) @ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:25

Verifier error #2: unresolved call from statement DifferentiationInterfaceForwardDiffExt.JacobianConfig(DifferentiationInterfaceForwardDiffExt.no
thing, y::Vector{Float64}, x::Vector{Float64}, φ ()::ForwardDiff.Chunk, $(QuoteNode(ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Floa
t64}())))::ForwardDiff.JacobianConfig{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.D
ual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule
.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A}} where _A
Stacktrace:
[1] prepare_jacobian_nokwarg(::Val{true}, ::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothi
ng}, ::Vector{Float64})
@ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/alBlj/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl
:386
[2] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
}; strict::Val{true})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:18 [inlined]
[3] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:15 [inlined]
[4] main(ξ::Float64)
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:16
[5] main(argc::Int32, argv::Ptr{Ptr{Int8}})
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:25

Verifier error #3: unresolved call from statement DifferentiationInterfaceForwardDiffExt.ForwardDiffTwoArgJacobianPrep($(QuoteNode(Val{Tuple{MyMo
dule.var"#main##0#main##1"{Float64}, Vector{Float64}, ADTypes.AutoForwardDiff{nothing, Nothing}, Vector{Float64}, Tuple{}}}())), DifferentiationI
nterfaceForwardDiffExt.JacobianConfig(DifferentiationInterfaceForwardDiffExt.nothing, y::Vector{Float64}, x::Vector{Float64}, φ ()::ForwardDiff.C
hunk, $(QuoteNode(ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}())))::ForwardDiff.JacobianConfig{ForwardDiff.Tag{MyModule.var
"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float
64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where
_A}} where _A, ()::Tuple{})::DifferentiationInterfaceForwardDiffExt.ForwardDiffTwoArgJacobianPrep{Tuple{MyModule.var"#main##0#main##1"{Float64},
Vector{Float64}, ADTypes.AutoForwardDiff{nothing, Nothing}, Vector{Float64}, Tuple{}}, C, Tuple{}} where C<:(ForwardDiff.JacobianConfig{ForwardDi
ff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main#
#1"{Float64}, Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Floa
t64, _A}, 1} where _A}} where _A)
Stacktrace:
[1] prepare_jacobian_nokwarg(::Val{true}, ::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothi
ng}, ::Vector{Float64})
@ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/alBlj/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl
:388
[2] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
}; strict::Val{true})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:18 [inlined]
[3] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:15 [inlined]
[4] main(ξ::Float64)
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:16
[5] main(argc::Int32, argv::Ptr{Ptr{Int8}})
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:25

Verifier error #4: unresolved call from statement MyModule.jacobian(%new()::MyModule.var"#main##0#main##1"{Float64}, %new()::Vector{Float64}, Dif
ferentiationInterface.prepare_jacobian_nokwarg($(QuoteNode(Val{true}())), %new()::MyModule.var"#main##0#main##1"{Float64}, %new()::Vector{Float64
}, $(QuoteNode(ADTypes.AutoForwardDiff{nothing, Nothing}(nothing))), %new()::Vector{Float64})::DifferentiationInterfaceForwardDiffExt.ForwardDiff
TwoArgJacobianPrep{Tuple{MyModule.var"#main##0#main##1"{Float64}, Vector{Float64}, ADTypes.AutoForwardDiff{nothing, Nothing}, Vector{Float64}, Tu
ple{}}, C, Tuple{}} where C<:(ForwardDiff.JacobianConfig{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{
Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{Forwa
rdDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A}} where _A), $(QuoteNode(ADTypes.AutoForwardDiff{nothing,
Nothing}(nothing))), %new()::Vector{Float64})::Matrix{Float64}
Stacktrace:
[1] main(ξ::Float64)
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:17
[2] main(argc::Int32, argv::Ptr{Ptr{Int8}})
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:25

Verifier error #5: unresolved call from statement DifferentiationInterfaceForwardDiffExt.JacobianConfig(DifferentiationInterfaceForwardDiffExt.no
thing, y::Vector{Float64}, x::Vector{Float64}, φ ()::ForwardDiff.Chunk, $(QuoteNode(ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Floa
t64}())))::ForwardDiff.JacobianConfig{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.D
ual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule
.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where _A}} where _A
Stacktrace:
[1] prepare_jacobian_nokwarg(::Val{true}, ::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothi
ng}, ::Vector{Float64})
@ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/alBlj/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl
:386
[2] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
}; strict::Val{true})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:18 [inlined]
[3] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:15 [inlined]
[4] main(ξ::Float64)
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:16
[5] main(argc::Int32, argv::Ptr{Ptr{Int8}})
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:25

Verifier error #6: unresolved call from statement DifferentiationInterfaceForwardDiffExt.ForwardDiffTwoArgJacobianPrep($(QuoteNode(Val{Tuple{MyMo
dule.var"#main##0#main##1"{Float64}, Vector{Float64}, ADTypes.AutoForwardDiff{nothing, Nothing}, Vector{Float64}, Tuple{}}}())), DifferentiationI
nterfaceForwardDiffExt.JacobianConfig(DifferentiationInterfaceForwardDiffExt.nothing, y::Vector{Float64}, x::Vector{Float64}, φ ()::ForwardDiff.C
hunk, $(QuoteNode(ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}())))::ForwardDiff.JacobianConfig{ForwardDiff.Tag{MyModule.var
"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float
64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A}, 1} where
_A}} where _A, ()::Tuple{})::DifferentiationInterfaceForwardDiffExt.ForwardDiffTwoArgJacobianPrep{Tuple{MyModule.var"#main##0#main##1"{Float64},
Vector{Float64}, ADTypes.AutoForwardDiff{nothing, Nothing}, Vector{Float64}, Tuple{}}, C, Tuple{}} where C<:(ForwardDiff.JacobianConfig{ForwardDi
ff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main#
#1"{Float64}, Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{MyModule.var"#main##0#main##1"{Float64}, Float64}, Floa
t64, _A}, 1} where _A}} where _A)
Stacktrace:
[1] prepare_jacobian_nokwarg(::Val{true}, ::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothi
ng}, ::Vector{Float64})
@ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/alBlj/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl
:388
[2] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
}; strict::Val{true})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:18 [inlined]
[3] prepare_jacobian(::MyModule.var"#main##0#main##1"{Float64}, ::Vector{Float64}, ::ADTypes.AutoForwardDiff{nothing, Nothing}, ::Vector{Float64
})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/alBlj/src/first_order/jacobian.jl:15 [inlined]
[4] main(ξ::Float64)
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:16
[5] main(argc::Int32, argv::Ptr{Ptr{Int8}})
@ MyModule ~/Documents/julia_playground/trim_simplenonlinearsolve/main_DI.jl:25

Trim verify finished with 6 errors, 0 warnings.


PS: When compiled with AutoForwardDiff, this can be run with LD_LIBRARY_PATH=~/.julia/juliaup/julia-1.12.0-beta4+0.x64.linux.gnu/lib/julia/ ./main_mwe 5.0

[MasonProtter]

Crossposting from Zulip: it seems that the issue is that prepare_jacobian is type unstable with AutoForwardDiff. Not sure how avoidable that is or isn't, but one way to work around that type instability (and generally improve performance) would be

# main.jl
module MyModule
using DifferentiationInterface
using FiniteDiff
using ForwardDiff
using LinearAlgebra

function f!(du, u, p)
    du .= p.λ * u
end

make_f!(ξ) = (du, u) -> f!(du, u, (; λ=ξ))

const backend = AutoForwardDiff()
const Jprep = prepare_jacobian(make_f!(1.0), zeros(2), backend, zeros(2))

function main(ξ=1.0)
    f!_ = make_f!(ξ)
    J   = jacobian(f!_, zeros(2), Jprep, backend, zeros(2))
    return sum(J)
end

Base.@ccallable function main(argc::Cint, argv::Ptr{Ptr{Int8}})::Cint
    cstr = unsafe_load(argv, 2)
    jlstr = unsafe_string(cstr)
    λ = parse(Float64, jlstr)
    sol = main(λ)
    println(Core.stdout, sum(sol))
    return 0
end

end

[RomeoV]

@MasonProtter helped me on Zulip with this version, which solves the issues

# main.jl
module MyModule
using DifferentiationInterface
using FiniteDiff
using ForwardDiff
using LinearAlgebra

function f!(du, u, p)
    du .= p.λ * u
end

make_f!(ξ) = (du, u) -> f!(du, u, (; λ=ξ))

const backend = AutoForwardDiff()
const Jprep = prepare_jacobian(make_f!(1.0), zeros(2), backend, zeros(2))

function main(ξ=1.0)
    f!_ = make_f!(ξ)
    J   = jacobian(f!_, zeros(2), Jprep, backend, zeros(2))
    return sum(J)
end

Base.@ccallable function main(argc::Cint, argv::Ptr{Ptr{Int8}})::Cint
    cstr = unsafe_load(argv, 2)
    jlstr = unsafe_string(cstr)
    λ = parse(Float64, jlstr)
    sol = main(λ)
    println(Core.stdout, sum(sol))
    return 0
end

end

[gdalle]

The type instability is pretty much unavoidable there, unless you set the chunk size explicitly with AutoForwardDiff(chunk size=...). You lose flexibility but you gain type stability.

[RomeoV]

I have also tried this with AutoForwardDiff(chunksize=1), but it didn't seem to help much.

[gdalle]

If this is due to a type instability, can you isolate it without involving trimmed binaries?

• That being said, in general type stability of preparation with DI is not guaranteed, which is why preparation exists in the first place: to isolate expensive statements in a function barrier outside of hot loops