Skip to content

Commit

Permalink
Add SparseEnzymeADHessian
Browse files Browse the repository at this point in the history
  • Loading branch information
amontoison committed Dec 2, 2024
1 parent d5fd76d commit 4376882
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions src/enzyme.jl
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ function SparseEnzymeADJacobian(
)
end

struct SparseEnzymeADHessian{R, C, S} <: ADNLPModels.ADBackend
struct SparseEnzymeADHessian{R, C, S, L} <: ADNLPModels.ADBackend
nvar::Int
rowval::Vector{Int}
colptr::Vector{Int}
Expand All @@ -156,6 +156,8 @@ struct SparseEnzymeADHessian{R, C, S} <: ADNLPModels.ADBackend
v::Vector{R}
y::Vector{R}
grad::Vector{R}
buffer::Vector{R}
::L
end

function SparseEnzymeADHessian(
Expand Down Expand Up @@ -200,7 +202,9 @@ function SparseEnzymeADHessian(
end
v = similar(x0)
y = similar(x0, ncon)
buffer = similar(x0, ncon)
grad = similar(x0)
(x, y, obj_weight, buffer) = obj_weight * nlp.f(x) + dot(c!(buffer, x), y)

return SparseEnzymeADHessian(
nvar,
Expand All @@ -213,6 +217,8 @@ function SparseEnzymeADHessian(
v,
y,
grad,
buffer,
ℓ,
)
end

Expand Down Expand Up @@ -423,18 +429,17 @@ end
compressed_hessian_icol =
(b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol)

# Lagrangian
= get_lag(nlp, b, obj_weight, y)

# AD with Enzyme.jl
Enzyme.autodiff(
Enzyme.Forward,
Enzyme.Const(Enzyme.gradient!),
Enzyme.Const(Enzyme.Reverse),
Enzyme.DuplicatedNoNeed(b.grad, compressed_hessian_icol),
Enzyme.Const(ℓ),
Enzyme.Const(b.ℓ),
Enzyme.Duplicated(x, b.v),
Enzyme.Const(y),
Enzyme.Const(obj_weight),
Enzyme.Const(buffer),
)

if b.coloring_mode == :direct
Expand Down

0 comments on commit 4376882

Please sign in to comment.