Skip to content

Robust options for Enzyme #378

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
amontoison merged 4 commits into from
Feb 25, 2026
Merged

Robust options for Enzyme #378

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
303d293
Robust options for Enzyme
michel2323 Feb 17, 2026
c0a15ac
Fixes
michel2323 Feb 18, 2026
d2b8dd0
Use ForwardDiff over Enzyme Reverse for second-order
michel2323 Feb 18, 2026
5836266
Update sparse_hessian.jl
amontoison Feb 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 18 additions & 5 deletions ext/ADNLPModelsEnzymeExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,10 @@ using ADNLPModels, NLPModels
using SparseMatrixColorings
using Enzyme

# Configure Enzyme for robustness
Enzyme.API.strictAliasing!(false) # handle Union types in Printf/@sprintf
Enzyme.API.looseTypeAnalysis!(true) # handle unresolved types in complex structs
Copy link

Copilot AI Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These Enzyme configuration calls are executed at module top-level. In Julia, top-level side effects may run during precompilation and won’t reliably re-run when loading the precompiled extension, so the settings may not actually be applied in user sessions. Consider moving this configuration into __init__() (and optionally guarding with isdefined/try to avoid hard load failures) so it is applied at runtime when the extension is loaded. Also note this mutates global Enzyme state for the whole Julia process; if possible, scope it (or make it opt-in) to avoid impacting other packages using Enzyme.

Suggested change
# Configure Enzyme for robustness
Enzyme.API.strictAliasing!(false) # handle Union types in Printf/@sprintf
Enzyme.API.looseTypeAnalysis!(true) # handle unresolved types in complex structs
function __init__()
# Configure Enzyme for robustness at runtime (not during precompilation)
try
if isdefined(Enzyme, :API)
Enzyme.API.strictAliasing!(false) # handle Union types in Printf/@sprintf
Enzyme.API.looseTypeAnalysis!(true) # handle unresolved types in complex structs
end
catch err
# Avoid hard failures if Enzyme's API or availability changes
@debug "Failed to configure Enzyme in ADNLPModelsEnzymeExt.__init__" exception = (err, catch_backtrace())
end
end

Copilot uses AI. Check for mistakes.

function _gradient!(dx, f, x)
Enzyme.make_zero!(dx)
Enzyme.autodiff(
Expand Down Expand Up @@ -70,12 +74,21 @@ end

function ADNLPModels.gradient!(::ADNLPModels.EnzymeReverseADGradient, g, f, x)
Enzyme.make_zero!(g)
Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(x, g))
Enzyme.autodiff(
Enzyme.set_runtime_activity(Enzyme.Reverse),
Enzyme.Const(f),
Enzyme.Active,
Enzyme.Duplicated(x, g),
)
return g
end

ADNLPModels.jacobian(::ADNLPModels.EnzymeReverseADJacobian, f, x) =
Enzyme.jacobian(Enzyme.Reverse, f, x)
Enzyme.jacobian(
Enzyme.set_runtime_activity(Enzyme.Reverse),
f,
x
)

function ADNLPModels.hessian(b::ADNLPModels.EnzymeReverseADHessian, f, x)
T = eltype(x)
Expand All @@ -96,7 +109,7 @@ function ADNLPModels.Jprod!(b::ADNLPModels.EnzymeReverseADJprod, Jv, c!, x, v, :
copyto!(b.xbuf, x)
copyto!(b.vbuf, v)
Enzyme.autodiff(
Enzyme.Forward,
Enzyme.set_runtime_activity(Enzyme.Forward),
Enzyme.Const(c!),
Enzyme.Duplicated(b.cx, b.jvbuf),
Enzyme.Duplicated(b.xbuf, b.vbuf),
Expand All @@ -118,7 +131,7 @@ function ADNLPModels.Jtprod!(b::ADNLPModels.EnzymeReverseADJtprod, Jtv, c!, x, v
copyto!(b.vbuf, v)
Enzyme.make_zero!(b.jtvbuf)
Enzyme.autodiff(
Enzyme.Reverse,
Enzyme.set_runtime_activity(Enzyme.Reverse),
Enzyme.Const(_void_c!),
Enzyme.Const(c!),
Enzyme.Duplicated(b.cx, b.vbuf),
Expand Down Expand Up @@ -261,7 +274,7 @@ function sparse_jac_coord!(
# b.compressed_jacobian is just a vector Jv here
# We don't use the vector mode
Enzyme.autodiff(
Enzyme.Forward,
Enzyme.set_runtime_activity(Enzyme.Forward),
Enzyme.Const(c!),
Enzyme.Duplicated(b.cx, b.compressed_jacobian),
Enzyme.Duplicated(b.xbuf, b.v),
Expand Down
Loading
⚡