diff --git a/docs/src/backend.md b/docs/src/backend.md index 7e073675..c2894dcb 100644 --- a/docs/src/backend.md +++ b/docs/src/backend.md @@ -10,11 +10,11 @@ The functions used internally to define the NLPModel API and the possible backen | Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend | | --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- | | `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeADGradient` | -- | -| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | -- | `SparseADJacobian` | -| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | -- | `SparseADHessian`/`SparseReverseADHessian` | -| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | -- | -- | -| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | -- | -- | -| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | -- | -- | +| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `EnzymeADJacobian` | `SparseADJacobian` | +| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | `EnzymeADHessian` | `SparseADHessian`/`SparseReverseADHessian` | +| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | `EnzymeADJprod` | -- | +| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | `EnzymeADJtprod` | -- | +| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | `EnzymeADHvprod` | -- | | `directional_second_derivative` | `ForwardDiffADGHjvprod` | -- | -- | -- | -- | The functions `hess_structure!`, `hess_coord!`, `jac_structure!` and `jac_coord!` defined in `ad.jl` are generic to all the backends for now. diff --git a/docs/src/predefined.md b/docs/src/predefined.md index 15793aae..12e771bf 100644 --- a/docs/src/predefined.md +++ b/docs/src/predefined.md @@ -46,3 +46,12 @@ It is possible to use these pre-defined backends using the keyword argument `bac nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon, backend = :optimized) get_adbackend(nlp) ``` + +The backend `:enzyme` focuses on backend based on [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). + +```@example ex1 +nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon, backend = :enzyme) +get_adbackend(nlp) +``` + +The backend `:zygote` focuses on backend based on [Zygote.jl](https://github.com/FluxML/Zygote.jl). diff --git a/docs/src/sparse.md b/docs/src/sparse.md index 73a60e61..3822bb64 100644 --- a/docs/src/sparse.md +++ b/docs/src/sparse.md @@ -31,7 +31,7 @@ x = rand(T, 2) H = hess(nlp, x) ``` -The backends available for sparse derivatives (`SparseADJacobian`, `SparseADHessian`, and `SparseReverseADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`. +The backends available for sparse derivatives (`SparseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, `SparseReverseADHessian`, and `SparseEnzymeADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`. These arguments specify the sparsity pattern detector and the coloring algorithm, respectively. - A **`detector`** must be of type `ADTypes.AbstractSparsityDetector`. @@ -39,7 +39,7 @@ These arguments specify the sparsity pattern detector and the coloring algorithm Prior to version 0.8.0, the default was `SymbolicSparsityDetector()` from `Symbolics.jl`. - A **`coloring_algorithm`** must be of type `SparseMatrixColorings.GreedyColoringAlgorithm`. - The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian`. + The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian` and `SparseEnzymeADHessian`. These algorithms are provided by the package `SparseMatrixColorings.jl`. The `GreedyColoringAlgorithm{:direct}()` performs column coloring for Jacobians and star coloring for Hessians. diff --git a/src/enzyme.jl b/src/enzyme.jl index db5133fe..3fea7c4f 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -1,6 +1,6 @@ -struct EnzymeADGradient <: ADNLPModels.ADBackend end +struct EnzymeReverseADGradient <: InPlaceADbackend end -function EnzymeADGradient( +function EnzymeReverseADGradient( nvar::Integer, f, ncon::Integer = 0, @@ -8,14 +8,553 @@ function EnzymeADGradient( x0::AbstractVector = rand(nvar), kwargs..., ) - return EnzymeADGradient() + return EnzymeReverseADGradient() +end + +struct EnzymeReverseADJacobian <: ADBackend end + +function EnzymeReverseADJacobian( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + return EnzymeReverseADJacobian() +end + +struct EnzymeReverseADHessian{T} <: ADBackend + seed::Vector{T} + Hv::Vector{T} +end + +function EnzymeReverseADHessian( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + @assert nvar > 0 + nnzh = nvar * (nvar + 1) / 2 + + seed = zeros(T, nvar) + Hv = zeros(T, nvar) + return EnzymeReverseADHessian(seed, Hv) +end + +struct EnzymeReverseADHvprod{T} <: InPlaceADbackend + grad::Vector{T} +end + +function EnzymeReverseADHvprod( + nvar::Integer, + f, + ncon::Integer = 0, + c!::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + grad = zeros(T, nvar) + return EnzymeReverseADHvprod(grad) +end + +struct EnzymeReverseADJprod{T} <: InPlaceADbackend + cx::Vector{T} +end + +function EnzymeReverseADJprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + cx = zeros(T, nvar) + return EnzymeReverseADJprod(cx) +end + +struct EnzymeReverseADJtprod{T} <: InPlaceADbackend + cx::Vector{T} +end + +function EnzymeReverseADJtprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + cx = zeros(T, nvar) + return EnzymeReverseADJtprod(cx) +end + +struct SparseEnzymeADJacobian{R, C, S} <: ADBackend + nvar::Int + ncon::Int + rowval::Vector{Int} + colptr::Vector{Int} + nzval::Vector{R} + result_coloring::C + compressed_jacobian::S + v::Vector{R} + cx::Vector{R} +end + +function SparseEnzymeADJacobian( + nvar, + f, + ncon, + c!; + x0::AbstractVector = rand(nvar), + coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), + detector::AbstractSparsityDetector = TracerSparsityDetector(), + kwargs..., +) + output = similar(x0, ncon) + J = compute_jacobian_sparsity(c!, output, x0, detector = detector) + SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...) +end + +function SparseEnzymeADJacobian( + nvar, + f, + ncon, + c!, + J::SparseMatrixCSC{Bool, Int}; + x0::AbstractVector{T} = rand(nvar), + coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), + kwargs..., +) where {T} + # We should support :row and :bidirectional in the future + problem = ColoringProblem{:nonsymmetric, :column}() + result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T) + + rowval = J.rowval + colptr = J.colptr + nzval = T.(J.nzval) + compressed_jacobian = similar(x0, ncon) + v = similar(x0) + cx = zeros(T, ncon) + + SparseEnzymeADJacobian( + nvar, + ncon, + rowval, + colptr, + nzval, + result_coloring, + compressed_jacobian, + v, + cx, + ) +end + +struct SparseEnzymeADHessian{R, C, S, L} <: ADBackend + nvar::Int + rowval::Vector{Int} + colptr::Vector{Int} + nzval::Vector{R} + result_coloring::C + coloring_mode::Symbol + compressed_hessian_icol::Vector{R} + compressed_hessian::S + v::Vector{R} + y::Vector{R} + grad::Vector{R} + cx::Vector{R} + ℓ::L +end + +function SparseEnzymeADHessian( + nvar, + f, + ncon, + c!; + x0::AbstractVector = rand(nvar), + coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(), + detector::AbstractSparsityDetector = TracerSparsityDetector(), + kwargs..., +) + H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) + SparseEnzymeADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, kwargs...) +end + +function SparseEnzymeADHessian( + nvar, + f, + ncon, + c!, + H::SparseMatrixCSC{Bool, Int}; + x0::AbstractVector{T} = rand(nvar), + coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(), + kwargs..., +) where {T} + problem = ColoringProblem{:symmetric, :column}() + result_coloring = coloring(H, problem, coloring_algorithm, decompression_eltype = T) + + trilH = tril(H) + rowval = trilH.rowval + colptr = trilH.colptr + nzval = T.(trilH.nzval) + if coloring_algorithm isa GreedyColoringAlgorithm{:direct} + coloring_mode = :direct + compressed_hessian_icol = similar(x0) + compressed_hessian = compressed_hessian_icol + else + coloring_mode = :substitution + group = column_groups(result_coloring) + ncolors = length(group) + compressed_hessian_icol = similar(x0) + compressed_hessian = similar(x0, (nvar, ncolors)) + end + v = similar(x0) + y = similar(x0, ncon) + cx = similar(x0, ncon) + grad = similar(x0) + function ℓ(x, y, obj_weight, cx) + res = obj_weight * f(x) + if ncon != 0 + c!(cx, x) + res += sum(cx[i] * y[i] for i = 1:ncon) + end + return res + end + + return SparseEnzymeADHessian( + nvar, + rowval, + colptr, + nzval, + result_coloring, + coloring_mode, + compressed_hessian_icol, + compressed_hessian, + v, + y, + grad, + cx, + ℓ, + ) end @init begin @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin - function ADNLPModels.gradient!(::EnzymeADGradient, g, f, x) - Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + + function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) + g = similar(x) + Enzyme.gradient!(Enzyme.Reverse, g, Enzyme.Const(f), x) + return g + end + + function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) + Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(x, g)) return g end + + jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) + + function hessian(b::EnzymeReverseADHessian, f, x) + T = eltype(x) + n = length(x) + hess = zeros(T, n, n) + fill!(b.seed, zero(T)) + for i in 1:n + b.seed[i] = one(T) + Enzyme.hvp!(b.Hv, Enzyme.Const(f), x, b.seed) + view(hess, :, i) .= b.Hv + b.seed[i] = zero(T) + end + return hess + end + + function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Forward, Enzyme.Const(c!), Enzyme.Duplicated(b.cx, Jv), Enzyme.Duplicated(x, v)) + return Jv + end + + function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(c!), Enzyme.Duplicated(b.cx, Jtv), Enzyme.Duplicated(x, v)) + return Jtv + end + + function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x, + v, + ℓ, + ::Val{:lag}, + y, + obj_weight::Real = one(eltype(x)), + ) + Enzyme.autodiff( + Enzyme.Forward, + Enzyme.Const(Enzyme.gradient!), + Enzyme.Const(Enzyme.Reverse), + Enzyme.DuplicatedNoNeed(b.grad, Hv), + Enzyme.Const(ℓ), + Enzyme.Duplicated(x, v), + Enzyme.Const(y), + ) + return Hv + end + + function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x, + v, + f, + ::Val{:obj}, + obj_weight::Real = one(eltype(x)), + ) + Enzyme.autodiff( + Enzyme.Forward, + Enzyme.Const(Enzyme.gradient!), + Enzyme.Const(Enzyme.Reverse), + Enzyme.DuplicatedNoNeed(b.grad, Hv), + Enzyme.Const(f), + Enzyme.Duplicated(x, v), + ) + return Hv + end + + # Sparse Jacobian + function get_nln_nnzj(b::SparseEnzymeADJacobian, nvar, ncon) + length(b.rowval) + end + + function NLPModels.jac_structure!( + b::SparseEnzymeADJacobian, + nlp::ADModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + rows .= b.rowval + for i = 1:(nlp.meta.nvar) + for j = b.colptr[i]:(b.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols + end + + function sparse_jac_coord!( + c!::Function, + b::SparseEnzymeADJacobian, + x::AbstractVector, + vals::AbstractVector, + ) + # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression + A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval) + + groups = column_groups(b.result_coloring) + for (icol, cols) in enumerate(groups) + # Update the seed + b.v .= 0 + for col in cols + b.v[col] = 1 + end + + # b.compressed_jacobian is just a vector Jv here + # We don't use the vector mode + Enzyme.autodiff( + Enzyme.Forward, + Enzyme.Const(c!), + Enzyme.Duplicated(b.cx, b.compressed_jacobian), + Enzyme.Duplicated(x, b.v) + ) + + # Update the columns of the Jacobian that have the color `icol` + decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring) + end + vals .= b.nzval + return vals + end + + function NLPModels.jac_coord!( + b::SparseEnzymeADJacobian, + nlp::ADModel, + x::AbstractVector, + vals::AbstractVector, + ) + sparse_jac_coord!(nlp.c!, b, x, vals) + return vals + end + + function NLPModels.jac_structure_residual!( + b::SparseEnzymeADJacobian, + nls::AbstractADNLSModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + rows .= b.rowval + for i = 1:(nls.meta.nvar) + for j = b.colptr[i]:(b.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols + end + + function NLPModels.jac_coord_residual!( + b::SparseEnzymeADJacobian, + nls::AbstractADNLSModel, + x::AbstractVector, + vals::AbstractVector, + ) + sparse_jac_coord!(nls.F!, b, x, vals) + return vals + end + + # Sparse Hessian + function get_nln_nnzh(b::SparseEnzymeADHessian, nvar) + return length(b.rowval) + end + + function NLPModels.hess_structure!( + b::SparseEnzymeADHessian, + nlp::ADModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + rows .= b.rowval + for i = 1:(nlp.meta.nvar) + for j = b.colptr[i]:(b.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols + end + + function sparse_hess_coord!( + b::SparseEnzymeADHessian, + x::AbstractVector, + obj_weight, + y::AbstractVector, + vals::AbstractVector, + ) + # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression + A = SparseMatrixCSC(b.nvar, b.nvar, b.colptr, b.rowval, b.nzval) + + groups = column_groups(b.result_coloring) + for (icol, cols) in enumerate(groups) + # Update the seed + b.v .= 0 + for col in cols + b.v[col] = 1 + end + + function _gradient!(dx, ℓ, x, y, obj_weight, cx) + Enzyme.make_zero!(dx) + dcx = Enzyme.make_zero(cx) + res = Enzyme.autodiff( + Enzyme.Reverse, + ℓ, + Enzyme.Active, + Enzyme.Duplicated(x, dx), + Enzyme.Const(y), + Enzyme.Const(obj_weight), + Enzyme.Duplicated(cx, dcx) + ) + return nothing + end + + function _hvp!(res, ℓ, x, v, y, obj_weight, cx) + dcx = Enzyme.make_zero(cx) + Enzyme.autodiff( + Enzyme.Forward, + _gradient!, + res, + Enzyme.Const(ℓ), + Enzyme.Duplicated(x, v), + Enzyme.Const(y), + Enzyme.Const(obj_weight), + Enzyme.Duplicated(cx, dcx), + ) + return nothing + end + + _hvp!( + Enzyme.DuplicatedNoNeed(b.grad, b.compressed_hessian_icol), + b.ℓ, x, b.v, y, obj_weight, b.cx + ) + + if b.coloring_mode == :direct + # Update the coefficients of the lower triangular part of the Hessian that are related to the color `icol` + decompress_single_color!(A, b.compressed_hessian_icol, icol, b.result_coloring, :L) + end + if b.coloring_mode == :substitution + view(b.compressed_hessian, :, icol) .= b.compressed_hessian_icol + end + end + if b.coloring_mode == :substitution + decompress!(A, b.compressed_hessian, b.result_coloring, :L) + end + vals .= b.nzval + return vals + end + + function NLPModels.hess_coord!( + b::SparseEnzymeADHessian, + nlp::ADModel, + x::AbstractVector, + y::AbstractVector, + obj_weight::Real, + vals::AbstractVector, + ) + sparse_hess_coord!(b, x, obj_weight, y, vals) + end + + # Could be optimized! + function NLPModels.hess_coord!( + b::SparseEnzymeADHessian, + nlp::ADModel, + x::AbstractVector, + obj_weight::Real, + vals::AbstractVector, + ) + b.y .= 0 + sparse_hess_coord!(b, x, obj_weight, b.y, vals) + end + + function NLPModels.hess_coord!( + b::SparseEnzymeADHessian, + nlp::ADModel, + x::AbstractVector, + j::Integer, + vals::AbstractVector, + ) + for (w, k) in enumerate(nlp.meta.nln) + b.y[w] = k == j ? 1 : 0 + end + obj_weight = zero(eltype(x)) + sparse_hess_coord!(b, x, obj_weight, b.y, vals) + return vals + end + + function NLPModels.hess_structure_residual!( + b::SparseEnzymeADHessian, + nls::AbstractADNLSModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + return hess_structure!(b, nls, rows, cols) + end + + function NLPModels.hess_coord_residual!( + b::SparseEnzymeADHessian, + nls::AbstractADNLSModel, + x::AbstractVector, + v::AbstractVector, + vals::AbstractVector, + ) + obj_weight = zero(eltype(x)) + sparse_hess_coord!(b, x, obj_weight, v, vals) + end end end diff --git a/src/forward.jl b/src/forward.jl index 652760a6..0c06af96 100644 --- a/src/forward.jl +++ b/src/forward.jl @@ -71,7 +71,7 @@ function Jprod!(::GenericForwardDiffADJprod, Jv, f, x, v, ::Val) return Jv end -struct ForwardDiffADJprod{T, Tag} <: ADNLPModels.InPlaceADbackend +struct ForwardDiffADJprod{T, Tag} <: InPlaceADbackend z::Vector{ForwardDiff.Dual{Tag, T, 1}} cz::Vector{ForwardDiff.Dual{Tag, T, 1}} end @@ -91,7 +91,7 @@ function ForwardDiffADJprod( return ForwardDiffADJprod(z, cz) end -function ADNLPModels.Jprod!(b::ForwardDiffADJprod{T, Tag}, Jv, c!, x, v, ::Val) where {T, Tag} +function Jprod!(b::ForwardDiffADJprod{T, Tag}, Jv, c!, x, v, ::Val) where {T, Tag} map!(ForwardDiff.Dual{Tag}, b.z, x, v) # x + ε * v c!(b.cz, b.z) # c!(cz, x + ε * v) ForwardDiff.extract_derivative!(Tag, Jv, b.cz) # ∇c!(cx, x)ᵀv @@ -113,7 +113,7 @@ function Jtprod!(::GenericForwardDiffADJtprod, Jtv, f, x, v, ::Val) return Jtv end -struct ForwardDiffADJtprod{Tag, GT, S} <: ADNLPModels.InPlaceADbackend +struct ForwardDiffADJtprod{Tag, GT, S} <: InPlaceADbackend cfg::ForwardDiff.GradientConfig{Tag} ψ::GT temp::S @@ -144,7 +144,7 @@ function ForwardDiffADJtprod( return ForwardDiffADJtprod(cfg, ψ, temp, sol) end -function ADNLPModels.Jtprod!( +function Jtprod!( b::ForwardDiffADJtprod{Tag, GT, S}, Jtv, c!, diff --git a/src/predefined_backend.jl b/src/predefined_backend.jl index 740fde7e..c612504b 100644 --- a/src/predefined_backend.jl +++ b/src/predefined_backend.jl @@ -13,8 +13,8 @@ default_backend = Dict( :hessian_residual_backend => SparseADHessian, ) -optimized = Dict( - :gradient_backend => ReverseDiffADGradient, # EnzymeADGradient +optimized_backend = Dict( + :gradient_backend => ReverseDiffADGradient, :hprod_backend => ReverseDiffADHvprod, :jprod_backend => ForwardDiffADJprod, :jtprod_backend => ReverseDiffADJtprod, @@ -28,7 +28,7 @@ optimized = Dict( :hessian_residual_backend => SparseReverseADHessian, ) -generic = Dict( +generic_backend = Dict( :gradient_backend => GenericForwardDiffADGradient, :hprod_backend => GenericForwardDiffADHvprod, :jprod_backend => GenericForwardDiffADJprod, @@ -43,7 +43,41 @@ generic = Dict( :hessian_residual_backend => ForwardDiffADHessian, ) -predefined_backend = Dict(:default => default_backend, :optimized => optimized, :generic => generic) +enzyme_backend = Dict( + :gradient_backend => EnzymeReverseADGradient, + :jprod_backend => EnzymeReverseADJprod, + :jtprod_backend => EnzymeReverseADJtprod, + :hprod_backend => EnzymeReverseADHvprod, + :jacobian_backend => SparseEnzymeADJacobian, + :hessian_backend => SparseEnzymeADHessian, + :ghjvprod_backend => ForwardDiffADGHjvprod, + :jprod_residual_backend => EnzymeReverseADJprod, + :jtprod_residual_backend => EnzymeReverseADJtprod, + :hprod_residual_backend => EnzymeReverseADHvprod, + :jacobian_residual_backend => SparseEnzymeADJacobian, + :hessian_residual_backend => SparseEnzymeADHessian, +) + +zygote_backend = Dict( + :gradient_backend => ZygoteADGradient, + :jprod_backend => ZygoteADJprod, + :jtprod_backend => ZygoteADJtprod, + :hprod_backend => ForwardDiffADHvprod, + :jacobian_backend => ZygoteADJacobian, + :hessian_backend => ZygoteADHessian, + :ghjvprod_backend => ForwardDiffADGHjvprod, + :jprod_residual_backend => ZygoteADJprod, + :jtprod_residual_backend => ZygoteADJtprod, + :hprod_residual_backend => ForwardDiffADHvprod, + :jacobian_residual_backend => ZygoteADJacobian, + :hessian_residual_backend => ZygoteADHessian, +) + +predefined_backend = Dict(:default => default_backend, + :optimized => optimized_backend, + :generic => generic_backend, + :enzyme => enzyme_backend, + :zygote => zygote_backend) """ get_default_backend(meth::Symbol, backend::Symbol; kwargs...) diff --git a/src/reverse.jl b/src/reverse.jl index 992a6bf7..a21e04ca 100644 --- a/src/reverse.jl +++ b/src/reverse.jl @@ -30,7 +30,7 @@ function gradient!(adbackend::ReverseDiffADGradient, g, f, x) return ReverseDiff.gradient!(g, adbackend.cfg, x) end -struct GenericReverseDiffADGradient <: ADNLPModels.ADBackend end +struct GenericReverseDiffADGradient <: ADBackend end function GenericReverseDiffADGradient( nvar::Integer, @@ -43,7 +43,7 @@ function GenericReverseDiffADGradient( return GenericReverseDiffADGradient() end -function ADNLPModels.gradient!(::GenericReverseDiffADGradient, g, f, x) +function gradient!(::GenericReverseDiffADGradient, g, f, x) return ReverseDiff.gradient!(g, f, x) end @@ -139,7 +139,7 @@ function Jtprod!(::GenericReverseDiffADJtprod, Jtv, f, x, v, ::Val) return Jtv end -struct ReverseDiffADJtprod{T, S, GT} <: ADNLPModels.InPlaceADbackend +struct ReverseDiffADJtprod{T, S, GT} <: InPlaceADbackend gtape::GT _tmp_out::Vector{ReverseDiff.TrackedReal{T, T, Nothing}} _rval::S # temporary storage for jtprod @@ -167,7 +167,7 @@ function ReverseDiffADJtprod( return ReverseDiffADJtprod(gtape, _tmp_out, _rval) end -function ADNLPModels.Jtprod!(b::ReverseDiffADJtprod, Jtv, c!, x, v, ::Val) +function Jtprod!(b::ReverseDiffADJtprod, Jtv, c!, x, v, ::Val) ReverseDiff.gradient!((Jtv, b._rval), b.gtape, (x, v)) return Jtv end diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl index 99b9c6ca..cbfb7e78 100644 --- a/src/sparse_hessian.jl +++ b/src/sparse_hessian.jl @@ -1,4 +1,4 @@ -struct SparseADHessian{Tag, R, T, C, H, S, GT} <: ADNLPModels.ADBackend +struct SparseADHessian{Tag, R, T, C, H, S, GT} <: ADBackend nvar::Int rowval::Vector{Int} colptr::Vector{Int} @@ -104,7 +104,7 @@ function SparseADHessian( ) end -struct SparseReverseADHessian{Tagf, Tagψ, R, T, C, H, S, F, P} <: ADNLPModels.ADBackend +struct SparseReverseADHessian{Tagf, Tagψ, R, T, C, H, S, F, P} <: ADBackend nvar::Int rowval::Vector{Int} colptr::Vector{Int} diff --git a/src/sparsity_pattern.jl b/src/sparsity_pattern.jl index daf5d6e9..95147149 100644 --- a/src/sparsity_pattern.jl +++ b/src/sparsity_pattern.jl @@ -79,7 +79,11 @@ end function get_sparsity_pattern(model::ADModel, ::Val{:jacobian}) backend = model.adbackend.jacobian_backend - validate_sparse_backend(backend, SparseADJacobian, "Jacobian") + validate_sparse_backend( + backend, + Union{SparseADJacobian, SparseEnzymeADJacobian}, + "Jacobian", + ) m = model.meta.ncon n = model.meta.nvar colptr = backend.colptr @@ -91,7 +95,11 @@ end function get_sparsity_pattern(model::ADModel, ::Val{:hessian}) backend = model.adbackend.hessian_backend - validate_sparse_backend(backend, Union{SparseADHessian, SparseReverseADHessian}, "Hessian") + validate_sparse_backend( + backend, + Union{SparseADHessian, SparseReverseADHessian, SparseEnzymeADHessian}, + "Hessian", + ) n = model.meta.nvar colptr = backend.colptr rowval = backend.rowval @@ -102,7 +110,11 @@ end function get_sparsity_pattern(model::AbstractADNLSModel, ::Val{:jacobian_residual}) backend = model.adbackend.jacobian_residual_backend - validate_sparse_backend(backend, SparseADJacobian, "Jacobian of the residual") + validate_sparse_backend( + backend, + Union{SparseADJacobian, SparseEnzymeADJacobian}, + "Jacobian of the residual", + ) m = model.nls_meta.nequ n = model.meta.nvar colptr = backend.colptr @@ -116,7 +128,7 @@ function get_sparsity_pattern(model::AbstractADNLSModel, ::Val{:hessian_residual backend = model.adbackend.hessian_residual_backend validate_sparse_backend( backend, - Union{SparseADHessian, SparseReverseADHessian}, + Union{SparseADHessian, SparseReverseADHessian, SparseEnzymeADHessian}, "Hessian of the residual", ) n = model.meta.nvar diff --git a/test/enzyme.jl b/test/enzyme.jl index 504557a0..f9cba09c 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -1,4 +1,5 @@ using LinearAlgebra, SparseArrays, Test +using SparseMatrixColorings using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! @@ -6,78 +7,130 @@ using ADNLPModels: # Automatically loads the code for Enzyme with Requires import Enzyme -#= -ADNLPModels.EmptyADbackend(args...; kwargs...) = ADNLPModels.EmptyADbackend() - -names = OptimizationProblems.meta[!, :name] -list_excluded_enzyme = [ - "brybnd", - "clplatea", - "clplateb", - "clplatec", - "curly", - "curly10", - "curly20", - "curly30", - "elec", - "fminsrf2", - "hs101", - "hs117", - "hs119", - "hs86", - "integreq", - "ncb20", - "ncb20b", - "palmer1c", - "palmer1d", - "palmer2c", - "palmer3c", - "palmer4c", - "palmer5c", - "palmer5d", - "palmer6c", - "palmer7c", - "palmer8c", - "sbrybnd", - "tetra", - "tetra_duct12", - "tetra_duct15", - "tetra_duct20", - "tetra_foam5", - "tetra_gear", - "tetra_hook", - "threepk", - "triangle", - "triangle_deer", - "triangle_pacman", - "triangle_turtle", - "watson", -] -for pb in names - @info pb - (pb in list_excluded_enzyme) && continue - nlp = eval(Meta.parse(pb))( - gradient_backend = ADNLPModels.EnzymeADGradient, - jacobian_backend = ADNLPModels.EmptyADbackend, - hessian_backend = ADNLPModels.EmptyADbackend, - ) - grad(nlp, get_x0(nlp)) +EnzymeReverseAD() = ADNLPModels.ADModelBackend( + ADNLPModels.EnzymeReverseADGradient(), + ADNLPModels.EnzymeReverseADHvprod(zeros(1)), + ADNLPModels.EnzymeReverseADJprod(zeros(1)), + ADNLPModels.EnzymeReverseADJtprod(zeros(1)), + ADNLPModels.EnzymeReverseADJacobian(), + ADNLPModels.EnzymeReverseADHessian(zeros(1), zeros(1)), + ADNLPModels.EnzymeReverseADHvprod(zeros(1)), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), +) + +function mysum!(y, x) + sum!(y, x) + return nothing +end + +function test_autodiff_backend_error() + @testset "Error without loading package - $backend" for backend in [:EnzymeReverseAD] + adbackend = eval(backend)() + # @test_throws ArgumentError gradient(adbackend.gradient_backend, sum, [1.0]) + # @test_throws ArgumentError gradient!(adbackend.gradient_backend, [1.0], sum, [1.0]) + # @test_throws ArgumentError jacobian(adbackend.jacobian_backend, identity, [1.0]) + # @test_throws ArgumentError hessian(adbackend.hessian_backend, sum, [1.0]) + # @test_throws ArgumentError Jprod!( + # adbackend.jprod_backend, + # [1.0], + # [1.0], + # identity, + # [1.0], + # Val(:c), + # ) + # @test_throws ArgumentError Jtprod!( + # adbackend.jtprod_backend, + # [1.0], + # [1.0], + # identity, + # [1.0], + # Val(:c), + # ) + gradient(adbackend.gradient_backend, sum, [1.0]) + gradient!(adbackend.gradient_backend, [1.0], sum, [1.0]) + jacobian(adbackend.jacobian_backend, sum, [1.0]) + hessian(adbackend.hessian_backend, sum, [1.0]) + Jprod!( + adbackend.jprod_backend, + [1.0], + sum!, + [1.0], + [1.0], + Val(:c), + ) + Jtprod!( + adbackend.jtprod_backend, + [1.0], + mysum!, + [1.0], + [1.0], + Val(:c), + ) + end +end + +test_autodiff_backend_error() + +include("sparse_jacobian.jl") +include("sparse_jacobian_nls.jl") +include("sparse_hessian.jl") +include("sparse_hessian_nls.jl") + +list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) + +@testset "Sparse Jacobian" begin + for (backend, kw) in list_sparse_jac_backend + sparse_jacobian(backend, kw) + sparse_jacobian_nls(backend, kw) + end +end + +list_sparse_hess_backend = ( + ( ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), + ), + ( + ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), + ), +) + +@testset "Sparse Hessian" begin + for (backend, kw) in list_sparse_hess_backend + sparse_hessian(backend, kw) + sparse_hessian_nls(backend, kw) + end +end + +for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] + include("nlp/problems/$(lowercase(problem)).jl") +end +for problem in NLPModelsTest.nls_problems + include("nls/problems/$(lowercase(problem)).jl") +end + +include("utils.jl") +include("nlp/basic.jl") +include("nls/basic.jl") +include("nlp/nlpmodelstest.jl") +include("nls/nlpmodelstest.jl") + +@testset "Basic NLP tests using $backend " for backend in (:enzyme,) + test_autodiff_model("$backend", backend = backend) +end + +@testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in (:enzyme,) + nlp_nlpmodelstest(backend) +end + +@testset "Basic NLS tests using $backend " for backend in (:enzyme,) + autodiff_nls_test("$backend", backend = backend) +end + +@testset "Checking NLPModelsTest (NLS) tests with $backend" for backend in (:enzyme,) + nls_nlpmodelstest(backend) end -=# - -#= -ERROR: Duplicated Returns not yet handled -Stacktrace: - [1] autodiff - @.julia\packages\Enzyme\DIkTv\src\Enzyme.jl:209 [inlined] - [2] autodiff(mode::EnzymeCore.ReverseMode, f::OptimizationProblems.ADNLPProblems.var"#f#254"{OptimizationProblems.ADNLPProblems.var"#f#250#255"}, args::Duplicated{Vector{Float64}}) - @ Enzyme.julia\packages\Enzyme\DIkTv\src\Enzyme.jl:248 - [3] gradient!(#unused#::ADNLPModels.EnzymeADGradient, g::Vector{Float64}, f::Function, x::Vector{Float64}) - @ ADNLPModelsDocuments\cvs\ADNLPModels.jl\src\enzyme.jl:17 - [4] grad!(nlp::ADNLPModel{Float64, Vector{Float64}, Vector{Int64}}, x::Vector{Float64}, g::Vector{Float64}) - @ ADNLPModelsDocuments\cvs\ADNLPModels.jl\src\nlp.jl:542 - [5] grad(nlp::ADNLPModel{Float64, Vector{Float64}, Vector{Int64}}, x::Vector{Float64}) - @ NLPModels.julia\packages\NLPModels\XBcWL\src\nlp\api.jl:31 - [6] top-level scope - @ .\REPL[7]:5 -=# diff --git a/test/nlp/basic.jl b/test/nlp/basic.jl index 741c3cb8..07c4d940 100644 --- a/test/nlp/basic.jl +++ b/test/nlp/basic.jl @@ -1,6 +1,6 @@ -mutable struct LinearRegression - x::Vector - y::Vector +mutable struct LinearRegression{T} + x::Vector{T} + y::Vector{T} end function (regr::LinearRegression)(beta) @@ -17,7 +17,7 @@ function test_autodiff_model(name; kwargs...) nlp = ADNLPModel(f, x0, c, [0.0], [0.0]; kwargs...) @test obj(nlp, x0) == f(x0) - x = range(-1, stop = 1, length = 100) + x = range(-1, stop = 1, length = 100) |> collect y = 2x .+ 3 + randn(100) * 0.1 regr = LinearRegression(x, y) nlp = ADNLPModel(regr, ones(2); kwargs...) @@ -343,7 +343,3 @@ function test_autodiff_model(name; kwargs...) @test A == sparse(nlp.clinrows, nlp.clincols, nlp.clinvals) end end - -@testset "Basic tests using $backend " for backend in keys(ADNLPModels.predefined_backend) - test_autodiff_model("$backend", backend = backend) -end diff --git a/test/nlp/nlpmodelstest.jl b/test/nlp/nlpmodelstest.jl index 78bf56ec..5aa96017 100644 --- a/test/nlp/nlpmodelstest.jl +++ b/test/nlp/nlpmodelstest.jl @@ -1,7 +1,5 @@ -@testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in - keys(ADNLPModels.predefined_backend) - @testset "Checking NLPModelsTest tests on problem $problem" for problem in - NLPModelsTest.nlp_problems +function nlp_nlpmodelstest(backend) + @testset "Checking NLPModelsTest tests on problem $problem" for problem in NLPModelsTest.nlp_problems nlp_from_T = eval(Meta.parse(lowercase(problem) * "_autodiff")) nlp_ad = nlp_from_T(; backend = backend) nlp_man = eval(Meta.parse(problem))() @@ -18,8 +16,10 @@ @testset "Check multiple precision" begin multiple_precision_nlp(nlp_from_T, exclude = [], linear_api = true) end - @testset "Check view subarray" begin - view_subarray_nlp(nlp_ad, exclude = []) + if backend != :enzyme + @testset "Check view subarray" begin + view_subarray_nlp(nlp_ad, exclude = []) + end end @testset "Check coordinate memory" begin coord_memory_nlp(nlp_ad, exclude = [], linear_api = true) diff --git a/test/nls/basic.jl b/test/nls/basic.jl index e7ec613b..31ae2539 100644 --- a/test/nls/basic.jl +++ b/test/nls/basic.jl @@ -360,7 +360,3 @@ function autodiff_nls_test(name; kwargs...) @test A == sparse(nls.clinrows, nls.clincols, nls.clinvals) end end - -@testset "Basic NLS tests using $backend " for backend in keys(ADNLPModels.predefined_backend) - autodiff_nls_test("$backend", backend = backend) -end diff --git a/test/nls/nlpmodelstest.jl b/test/nls/nlpmodelstest.jl index e1681f49..d277c0b5 100644 --- a/test/nls/nlpmodelstest.jl +++ b/test/nls/nlpmodelstest.jl @@ -1,7 +1,5 @@ -@testset "Checking NLPModelsTest (NLS) tests with $backend" for backend in - keys(ADNLPModels.predefined_backend) - @testset "Checking NLPModelsTest tests on problem $problem" for problem in - NLPModelsTest.nls_problems +function nls_nlpmodelstest(backend) + @testset "Checking NLPModelsTest tests on problem $problem" for problem in NLPModelsTest.nls_problems nls_from_T = eval(Meta.parse(lowercase(problem) * "_autodiff")) nls_ad = nls_from_T(; backend = backend) nls_man = eval(Meta.parse(problem))() @@ -35,8 +33,10 @@ @testset "Check multiple precision" begin multiple_precision_nls(nls_from_T, exclude = exclude, linear_api = true) end - @testset "Check view subarray" begin - view_subarray_nls.(nlss, exclude = exclude) + if backend != :enzyme + @testset "Check view subarray" begin + view_subarray_nls.(nlss, exclude = exclude) + end end end end diff --git a/test/runtests.jl b/test/runtests.jl index 1d764265..da4bf09a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,14 +23,43 @@ end include("manual.jl") end -@testset "Basic Jacobian derivative test" begin - include("sparse_jacobian.jl") - include("sparse_jacobian_nls.jl") +include("sparse_jacobian.jl") +include("sparse_jacobian_nls.jl") +include("sparse_hessian.jl") +include("sparse_hessian_nls.jl") + +list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), + (ADNLPModels.ForwardDiffADJacobian, Dict())) + +@testset "Sparse Jacobian" begin + for (backend, kw) in list_sparse_jac_backend + sparse_jacobian(backend, kw) + sparse_jacobian_nls(backend, kw) + end end -@testset "Basic Hessian derivative test" begin - include("sparse_hessian.jl") - include("sparse_hessian_nls.jl") +list_sparse_hess_backend = ( + (ADNLPModels.SparseADHessian, Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}())), + ( + ADNLPModels.SparseADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), + ), + ( + ADNLPModels.SparseReverseADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), + ), + ( + ADNLPModels.SparseReverseADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), + ), + (ADNLPModels.ForwardDiffADHessian, Dict()), +) + +@testset "Sparse Hessian" begin + for (backend, kw) in list_sparse_hess_backend + sparse_hessian(backend, kw) + sparse_hessian_nls(backend, kw) + end end for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] @@ -42,6 +71,30 @@ end include("utils.jl") include("nlp/basic.jl") -include("nls/basic.jl") include("nlp/nlpmodelstest.jl") +include("nls/basic.jl") include("nls/nlpmodelstest.jl") + +@testset "Basic NLP tests using $backend " for backend in keys(ADNLPModels.predefined_backend) + (backend == :zygote) && continue + (backend == :enzyme) && continue + test_autodiff_model("$backend", backend = backend) +end + +@testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in keys(ADNLPModels.predefined_backend) + (backend == :zygote) && continue + (backend == :enzyme) && continue + nlp_nlpmodelstest(backend) +end + +@testset "Basic NLS tests using $backend " for backend in keys(ADNLPModels.predefined_backend) + (backend == :zygote) && continue + (backend == :enzyme) && continue + autodiff_nls_test("$backend", backend = backend) +end + +@testset "Checking NLPModelsTest (NLS) tests with $backend" for backend in keys(ADNLPModels.predefined_backend) + (backend == :zygote) && continue + (backend == :enzyme) && continue + nls_nlpmodelstest(backend) +end diff --git a/test/script_OP.jl b/test/script_OP.jl index 37f7c66b..092a3a70 100644 --- a/test/script_OP.jl +++ b/test/script_OP.jl @@ -11,44 +11,48 @@ using JuMP, NLPModelsJuMP names = OptimizationProblems.meta[!, :name] -for pb in names - @info pb - - nlp = try - OptimizationProblems.ADNLPProblems.eval(Meta.parse(pb))(backend = :default, show_time = true) - catch e - println("Error $e with ADNLPModel") - continue - end +function test_OP(backend) + for pb in names + @info pb + + nlp = try + OptimizationProblems.ADNLPProblems.eval(Meta.parse(pb))(backend = backend, show_time = true) + catch e + println("Error $e with ADNLPModel") + continue + end - jum = try - MathOptNLPModel(OptimizationProblems.PureJuMP.eval(Meta.parse(pb))()) - catch e - println("Error $e with JuMP") - continue - end + jum = try + MathOptNLPModel(OptimizationProblems.PureJuMP.eval(Meta.parse(pb))()) + catch e + println("Error $e with JuMP") + continue + end - n, m = nlp.meta.nvar, nlp.meta.ncon - x = 10 * [-(-1.0)^i for i = 1:n] # find a better point in the domain. - v = 10 * [-(-1.0)^i for i = 1:n] - y = 3.14 * ones(m) - - # test the main functions in the API - try - @testset "Test NLPModel API $(nlp.meta.name)" begin - @test grad(nlp, x) ≈ grad(jum, x) - @test hess(nlp, x) ≈ hess(jum, x) - @test hess(nlp, x, y) ≈ hess(jum, x, y) - @test hprod(nlp, x, v) ≈ hprod(jum, x, v) - @test hprod(nlp, x, y, v) ≈ hprod(jum, x, y, v) - if nlp.meta.ncon > 0 - @test jac(nlp, x) ≈ jac(jum, x) - @test jprod(nlp, x, v) ≈ jprod(jum, x, v) - @test jtprod(nlp, x, y) ≈ jtprod(jum, x, y) + n, m = nlp.meta.nvar, nlp.meta.ncon + x = 10 * [-(-1.0)^i for i = 1:n] # find a better point in the domain. + v = 10 * [-(-1.0)^i for i = 1:n] + y = 3.14 * ones(m) + + # test the main functions in the API + try + @testset "Test NLPModel API $(nlp.meta.name)" begin + @test grad(nlp, x) ≈ grad(jum, x) + @test hess(nlp, x) ≈ hess(jum, x) + @test hess(nlp, x, y) ≈ hess(jum, x, y) + @test hprod(nlp, x, v) ≈ hprod(jum, x, v) + @test hprod(nlp, x, y, v) ≈ hprod(jum, x, y, v) + if nlp.meta.ncon > 0 + @test jac(nlp, x) ≈ jac(jum, x) + @test jprod(nlp, x, v) ≈ jprod(jum, x, v) + @test jtprod(nlp, x, y) ≈ jtprod(jum, x, y) + end end + catch e + println("Error $e with API") + continue end - catch e - println("Error $e with API") - continue end end + +test_OP(:default) diff --git a/test/sparse_hessian.jl b/test/sparse_hessian.jl index 250299f0..ca130989 100644 --- a/test/sparse_hessian.jl +++ b/test/sparse_hessian.jl @@ -1,106 +1,87 @@ -list_sparse_hess_backend = ( - (ADNLPModels.SparseADHessian, Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}())), - ( - ADNLPModels.SparseADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - ( - ADNLPModels.SparseReverseADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), - ), - ( - ADNLPModels.SparseReverseADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - (ADNLPModels.ForwardDiffADHessian, Dict()), -) +function sparse_hessian(backend, kw) + @testset "Basic Hessian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64) + c!(cx, x) = begin + cx[1] = x[1] - 1 + cx[2] = 10 * (x[2] - x[1]^2) + cx[3] = x[2] + 1 + cx + end + x0 = T[-1.2; 1.0] + nvar = 2 + ncon = 3 + nlp = ADNLPModel!( + x -> x[1] * x[2]^2 + x[1]^2 * x[2], + x0, + c!, + zeros(T, ncon), + zeros(T, ncon), + hessian_backend = backend; + kw..., + ) -dt = (Float32, Float64) + x = rand(T, 2) + y = rand(T, 3) + rows, cols = zeros(Int, nlp.meta.nnzh), zeros(Int, nlp.meta.nnzh) + vals = zeros(T, nlp.meta.nnzh) + hess_structure!(nlp, rows, cols) + hess_coord!(nlp, x, vals) + @test eltype(vals) == T + H = sparse(rows, cols, vals, nvar, nvar) + @test H == [2*x[2] 0; 2*(x[1] + x[2]) 2*x[1]] -@testset "Basic Hessian derivative with backend=$(backend) and T=$(T)" for T in dt, - (backend, kw) in list_sparse_hess_backend + # Test also the implementation of the backends + b = nlp.adbackend.hessian_backend + obj_weight = 0.5 + @test nlp.meta.nnzh == ADNLPModels.get_nln_nnzh(b, nvar) + ADNLPModels.hess_structure!(b, nlp, rows, cols) + ADNLPModels.hess_coord!(b, nlp, x, obj_weight, vals) + @test eltype(vals) == T + H = sparse(rows, cols, vals, nvar, nvar) + @test H == [x[2] 0; x[1]+x[2] x[1]] + ADNLPModels.hess_coord!(b, nlp, x, y, obj_weight, vals) + @test eltype(vals) == T + H = sparse(rows, cols, vals, nvar, nvar) + @test H == [x[2] 0; x[1]+x[2] x[1]] + y[2] * [-20 0; 0 0] - c!(cx, x) = begin - cx[1] = x[1] - 1 - cx[2] = 10 * (x[2] - x[1]^2) - cx[3] = x[2] + 1 - cx - end - x0 = T[-1.2; 1.0] - nvar = 2 - ncon = 3 - nlp = ADNLPModel!( - x -> x[1] * x[2]^2 + x[1]^2 * x[2], - x0, - c!, - zeros(T, ncon), - zeros(T, ncon), - hessian_backend = backend; - kw..., - ) - - x = rand(T, 2) - y = rand(T, 3) - rows, cols = zeros(Int, nlp.meta.nnzh), zeros(Int, nlp.meta.nnzh) - vals = zeros(T, nlp.meta.nnzh) - hess_structure!(nlp, rows, cols) - hess_coord!(nlp, x, vals) - @test eltype(vals) == T - H = sparse(rows, cols, vals, nvar, nvar) - @test H == [2*x[2] 0; 2*(x[1] + x[2]) 2*x[1]] - - # Test also the implementation of the backends - b = nlp.adbackend.hessian_backend - obj_weight = 0.5 - @test nlp.meta.nnzh == ADNLPModels.get_nln_nnzh(b, nvar) - ADNLPModels.hess_structure!(b, nlp, rows, cols) - ADNLPModels.hess_coord!(b, nlp, x, obj_weight, vals) - @test eltype(vals) == T - H = sparse(rows, cols, vals, nvar, nvar) - @test H == [x[2] 0; x[1]+x[2] x[1]] - ADNLPModels.hess_coord!(b, nlp, x, y, obj_weight, vals) - @test eltype(vals) == T - H = sparse(rows, cols, vals, nvar, nvar) - @test H == [x[2] 0; x[1]+x[2] x[1]] + y[2] * [-20 0; 0 0] + if backend != ADNLPModels.ForwardDiffADHessian + H_sp = get_sparsity_pattern(nlp, :hessian) + @test H_sp == SparseMatrixCSC{Bool, Int}([ + 1 0 + 1 1 + ]) + end - if (backend == ADNLPModels.SparseADHessian) || (backend == ADNLPModels.SparseReverseADHessian) - H_sp = get_sparsity_pattern(nlp, :hessian) - @test H_sp == SparseMatrixCSC{Bool, Int}([ - 1 0 - 1 1 - ]) - end - - nlp = ADNLPModel!( - x -> x[1] * x[2]^2 + x[1]^2 * x[2], - x0, - c!, - zeros(T, ncon), - zeros(T, ncon), - matrix_free = true; - kw..., - ) - @test nlp.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend + nlp = ADNLPModel!( + x -> x[1] * x[2]^2 + x[1]^2 * x[2], + x0, + c!, + zeros(T, ncon), + zeros(T, ncon), + matrix_free = true; + kw..., + ) + @test nlp.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend - n = 4 - x = ones(T, 4) - nlp = ADNLPModel( - x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)), - x, - hessian_backend = backend, - name = "Extended Rosenbrock", - ) - @test hess(nlp, x) == T[802 -400 0 0; -400 1002 -400 0; 0 -400 1002 -400; 0 0 -400 200] + n = 4 + x = ones(T, 4) + nlp = ADNLPModel( + x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)), + x, + hessian_backend = backend, + name = "Extended Rosenbrock", + ) + @test hess(nlp, x) == T[802 -400 0 0; -400 1002 -400 0; 0 -400 1002 -400; 0 0 -400 200] - x = ones(T, 2) - nlp = ADNLPModel(x -> x[1]^2 + x[1] * x[2], x, hessian_backend = backend) - @test hess(nlp, x) == T[2 1; 1 0] + x = ones(T, 2) + nlp = ADNLPModel(x -> x[1]^2 + x[1] * x[2], x, hessian_backend = backend) + @test hess(nlp, x) == T[2 1; 1 0] - nlp = ADNLPModel( - x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)), - x, - name = "Extended Rosenbrock", - matrix_free = true, - ) - @test nlp.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend + nlp = ADNLPModel( + x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)), + x, + name = "Extended Rosenbrock", + matrix_free = true, + ) + @test nlp.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend + end end diff --git a/test/sparse_hessian_nls.jl b/test/sparse_hessian_nls.jl index f0ee5cee..80ac269b 100644 --- a/test/sparse_hessian_nls.jl +++ b/test/sparse_hessian_nls.jl @@ -1,64 +1,45 @@ -list_sparse_hess_backend = ( - (ADNLPModels.SparseADHessian, Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}())), - ( - ADNLPModels.SparseADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - ( - ADNLPModels.SparseReverseADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), - ), - ( - ADNLPModels.SparseReverseADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - (ADNLPModels.ForwardDiffADHessian, Dict()), -) +function sparse_hessian_nls(backend, kw) + @testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64) + F!(Fx, x) = begin + Fx[1] = x[1] - 1 + Fx[2] = 10 * (x[2] - x[1]^2) + Fx[3] = x[2] + 1 + Fx + end + x0 = T[-1.2; 1.0] + nvar = 2 + nequ = 3 + nls = ADNLPModels.ADNLSModel!(F!, x0, 3, hessian_residual_backend = backend; kw...) -dt = (Float32, Float64) + x = rand(T, nvar) + v = rand(T, nequ) + rows, cols = zeros(Int, nls.nls_meta.nnzh), zeros(Int, nls.nls_meta.nnzh) + vals = zeros(T, nls.nls_meta.nnzh) + hess_structure_residual!(nls, rows, cols) + hess_coord_residual!(nls, x, v, vals) + @test eltype(vals) == T + H = Symmetric(sparse(rows, cols, vals, nvar, nvar), :L) + @test H == [-20*v[2] 0; 0 0] -@testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in dt, - (backend, kw) in list_sparse_hess_backend + # Test also the implementation of the backends + b = nls.adbackend.hessian_residual_backend + @test nls.nls_meta.nnzh == ADNLPModels.get_nln_nnzh(b, nvar) + ADNLPModels.hess_structure_residual!(b, nls, rows, cols) + ADNLPModels.hess_coord_residual!(b, nls, x, v, vals) + @test eltype(vals) == T + H = Symmetric(sparse(rows, cols, vals, nvar, nvar), :L) + @test H == [-20*v[2] 0; 0 0] - F!(Fx, x) = begin - Fx[1] = x[1] - 1 - Fx[2] = 10 * (x[2] - x[1]^2) - Fx[3] = x[2] + 1 - Fx - end - x0 = T[-1.2; 1.0] - nvar = 2 - nequ = 3 - nls = ADNLPModels.ADNLSModel!(F!, x0, 3, hessian_residual_backend = backend; kw...) - - x = rand(T, nvar) - v = rand(T, nequ) - rows, cols = zeros(Int, nls.nls_meta.nnzh), zeros(Int, nls.nls_meta.nnzh) - vals = zeros(T, nls.nls_meta.nnzh) - hess_structure_residual!(nls, rows, cols) - hess_coord_residual!(nls, x, v, vals) - @test eltype(vals) == T - H = Symmetric(sparse(rows, cols, vals, nvar, nvar), :L) - @test H == [-20*v[2] 0; 0 0] + if backend != ADNLPModels.ForwardDiffADHessian + H_sp = get_sparsity_pattern(nls, :hessian_residual) + @test H_sp == SparseMatrixCSC{Bool, Int}([ + 1 0 + 0 0 + ]) + end - # Test also the implementation of the backends - b = nls.adbackend.hessian_residual_backend - @test nls.nls_meta.nnzh == ADNLPModels.get_nln_nnzh(b, nvar) - ADNLPModels.hess_structure_residual!(b, nls, rows, cols) - ADNLPModels.hess_coord_residual!(b, nls, x, v, vals) - @test eltype(vals) == T - H = Symmetric(sparse(rows, cols, vals, nvar, nvar), :L) - @test H == [-20*v[2] 0; 0 0] - - if (backend == ADNLPModels.SparseADHessian) || (backend == ADNLPModels.SparseReverseADHessian) - H_sp = get_sparsity_pattern(nls, :hessian_residual) - @test H_sp == SparseMatrixCSC{Bool, Int}([ - 1 0 - 0 0 - ]) + nls = ADNLPModels.ADNLSModel!(F!, x0, 3, matrix_free = true; kw...) + @test nls.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend + @test nls.adbackend.hessian_residual_backend isa ADNLPModels.EmptyADbackend end - - nls = ADNLPModels.ADNLSModel!(F!, x0, 3, matrix_free = true; kw...) - @test nls.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend - @test nls.adbackend.hessian_residual_backend isa ADNLPModels.EmptyADbackend end diff --git a/test/sparse_jacobian.jl b/test/sparse_jacobian.jl index da0a8e5f..4c0801eb 100644 --- a/test/sparse_jacobian.jl +++ b/test/sparse_jacobian.jl @@ -1,65 +1,60 @@ -list_sparse_jac_backend = - ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict())) +function sparse_jacobian(backend, kw) + @testset "Basic Jacobian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64) + c!(cx, x) = begin + cx[1] = x[1] - 1 + cx[2] = 10 * (x[2] - x[1]^2) + cx[3] = x[2] + 1 + cx + end + x0 = T[-1.2; 1.0] + nvar = 2 + ncon = 3 + nlp = ADNLPModel!( + x -> sum(x), + x0, + c!, + zeros(T, ncon), + zeros(T, ncon), + jacobian_backend = backend; + kw..., + ) -dt = (Float32, Float64) - -@testset "Basic Jacobian derivative with backend=$(backend) and T=$(T)" for T in dt, - (backend, kw) in list_sparse_jac_backend - - c!(cx, x) = begin - cx[1] = x[1] - 1 - cx[2] = 10 * (x[2] - x[1]^2) - cx[3] = x[2] + 1 - cx - end - x0 = T[-1.2; 1.0] - nvar = 2 - ncon = 3 - nlp = ADNLPModel!( - x -> sum(x), - x0, - c!, - zeros(T, ncon), - zeros(T, ncon), - jacobian_backend = backend; - kw..., - ) - - x = rand(T, 2) - rows, cols = zeros(Int, nlp.meta.nln_nnzj), zeros(Int, nlp.meta.nln_nnzj) - vals = zeros(T, nlp.meta.nln_nnzj) - jac_nln_structure!(nlp, rows, cols) - jac_nln_coord!(nlp, x, vals) - @test eltype(vals) == T - J = sparse(rows, cols, vals, ncon, nvar) - @test J == [ - 1 0 - -20*x[1] 10 - 0 1 - ] - - # Test also the implementation of the backends - b = nlp.adbackend.jacobian_backend - @test nlp.meta.nnzj == ADNLPModels.get_nln_nnzj(b, nvar, ncon) - ADNLPModels.jac_structure!(b, nlp, rows, cols) - ADNLPModels.jac_coord!(b, nlp, x, vals) - @test eltype(vals) == T - J = sparse(rows, cols, vals, ncon, nvar) - @test J == [ - 1 0 - -20*x[1] 10 - 0 1 - ] + x = rand(T, 2) + rows, cols = zeros(Int, nlp.meta.nln_nnzj), zeros(Int, nlp.meta.nln_nnzj) + vals = zeros(T, nlp.meta.nln_nnzj) + jac_nln_structure!(nlp, rows, cols) + jac_nln_coord!(nlp, x, vals) + @test eltype(vals) == T + J = sparse(rows, cols, vals, ncon, nvar) + @test J == [ + 1 0 + -20*x[1] 10 + 0 1 + ] - if backend == ADNLPModels.SparseADJacobian - J_sp = get_sparsity_pattern(nlp, :jacobian) - @test J_sp == SparseMatrixCSC{Bool, Int}([ + # Test also the implementation of the backends + b = nlp.adbackend.jacobian_backend + @test nlp.meta.nnzj == ADNLPModels.get_nln_nnzj(b, nvar, ncon) + ADNLPModels.jac_structure!(b, nlp, rows, cols) + ADNLPModels.jac_coord!(b, nlp, x, vals) + @test eltype(vals) == T + J = sparse(rows, cols, vals, ncon, nvar) + @test J == [ 1 0 - 1 1 + -20*x[1] 10 0 1 - ]) - end + ] - nlp = ADNLPModel!(x -> sum(x), x0, c!, zeros(T, ncon), zeros(T, ncon), matrix_free = true; kw...) - @test nlp.adbackend.jacobian_backend isa ADNLPModels.EmptyADbackend + if backend != ADNLPModels.ForwardDiffADJacobian + J_sp = get_sparsity_pattern(nlp, :jacobian) + @test J_sp == SparseMatrixCSC{Bool, Int}([ + 1 0 + 1 1 + 0 1 + ]) + end + + nlp = ADNLPModel!(x -> sum(x), x0, c!, zeros(T, ncon), zeros(T, ncon), matrix_free = true; kw...) + @test nlp.adbackend.jacobian_backend isa ADNLPModels.EmptyADbackend + end end diff --git a/test/sparse_jacobian_nls.jl b/test/sparse_jacobian_nls.jl index 4405e7f7..fc6f8e98 100644 --- a/test/sparse_jacobian_nls.jl +++ b/test/sparse_jacobian_nls.jl @@ -1,58 +1,53 @@ -list_sparse_jac_backend = - ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict())) +function sparse_jacobian_nls(backend, kw) + @testset "Basic Jacobian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64) + F!(Fx, x) = begin + Fx[1] = x[1] - 1 + Fx[2] = 10 * (x[2] - x[1]^2) + Fx[3] = x[2] + 1 + Fx + end + x0 = T[-1.2; 1.0] + nvar = 2 + nequ = 3 + nls = ADNLPModels.ADNLSModel!(F!, x0, 3, jacobian_residual_backend = backend; kw...) -dt = (Float32, Float64) - -@testset "Basic Jacobian of residual derivative with backend=$(backend) and T=$(T)" for T in dt, - (backend, kw) in list_sparse_jac_backend - - F!(Fx, x) = begin - Fx[1] = x[1] - 1 - Fx[2] = 10 * (x[2] - x[1]^2) - Fx[3] = x[2] + 1 - Fx - end - x0 = T[-1.2; 1.0] - nvar = 2 - nequ = 3 - nls = ADNLPModels.ADNLSModel!(F!, x0, 3, jacobian_residual_backend = backend; kw...) - - x = rand(T, 2) - rows, cols = zeros(Int, nls.nls_meta.nnzj), zeros(Int, nls.nls_meta.nnzj) - vals = zeros(T, nls.nls_meta.nnzj) - jac_structure_residual!(nls, rows, cols) - jac_coord_residual!(nls, x, vals) - @test eltype(vals) == T - J = sparse(rows, cols, vals, nequ, nvar) - @test J == [ - 1 0 - -20*x[1] 10 - 0 1 - ] - - # Test also the implementation of the backends - b = nls.adbackend.jacobian_residual_backend - @test nls.nls_meta.nnzj == ADNLPModels.get_nln_nnzj(b, nvar, nequ) - ADNLPModels.jac_structure_residual!(b, nls, rows, cols) - ADNLPModels.jac_coord_residual!(b, nls, x, vals) - @test eltype(vals) == T - J = sparse(rows, cols, vals, nequ, nvar) - @test J == [ - 1 0 - -20*x[1] 10 - 0 1 - ] + x = rand(T, 2) + rows, cols = zeros(Int, nls.nls_meta.nnzj), zeros(Int, nls.nls_meta.nnzj) + vals = zeros(T, nls.nls_meta.nnzj) + jac_structure_residual!(nls, rows, cols) + jac_coord_residual!(nls, x, vals) + @test eltype(vals) == T + J = sparse(rows, cols, vals, nequ, nvar) + @test J == [ + 1 0 + -20*x[1] 10 + 0 1 + ] - if backend == ADNLPModels.SparseADJacobian - J_sp = get_sparsity_pattern(nls, :jacobian_residual) - @test J_sp == SparseMatrixCSC{Bool, Int}([ + # Test also the implementation of the backends + b = nls.adbackend.jacobian_residual_backend + @test nls.nls_meta.nnzj == ADNLPModels.get_nln_nnzj(b, nvar, nequ) + ADNLPModels.jac_structure_residual!(b, nls, rows, cols) + ADNLPModels.jac_coord_residual!(b, nls, x, vals) + @test eltype(vals) == T + J = sparse(rows, cols, vals, nequ, nvar) + @test J == [ 1 0 - 1 1 + -20*x[1] 10 0 1 - ]) - end + ] - nls = ADNLPModels.ADNLSModel!(F!, x0, 3, matrix_free = true; kw...) - @test nls.adbackend.jacobian_backend isa ADNLPModels.EmptyADbackend - @test nls.adbackend.jacobian_residual_backend isa ADNLPModels.EmptyADbackend + if backend != ADNLPModels.ForwardDiffADJacobian + J_sp = get_sparsity_pattern(nls, :jacobian_residual) + @test J_sp == SparseMatrixCSC{Bool, Int}([ + 1 0 + 1 1 + 0 1 + ]) + end + + nls = ADNLPModels.ADNLSModel!(F!, x0, 3, matrix_free = true; kw...) + @test nls.adbackend.jacobian_backend isa ADNLPModels.EmptyADbackend + @test nls.adbackend.jacobian_residual_backend isa ADNLPModels.EmptyADbackend + end end diff --git a/test/zygote.jl b/test/zygote.jl index 9c1e6273..023c217d 100644 --- a/test/zygote.jl +++ b/test/zygote.jl @@ -54,25 +54,6 @@ end # Test the argument error without loading the packages test_autodiff_backend_error() -# Additional backends used for tests -push!( - ADNLPModels.predefined_backend, - :zygote_backend => Dict( - :gradient_backend => ADNLPModels.ZygoteADGradient, - :jprod_backend => ADNLPModels.ZygoteADJprod, - :jtprod_backend => ADNLPModels.ZygoteADJtprod, - :hprod_backend => ADNLPModels.ForwardDiffADHvprod, - :jacobian_backend => ADNLPModels.ZygoteADJacobian, - :hessian_backend => ADNLPModels.ZygoteADHessian, - :ghjvprod_backend => ADNLPModels.ForwardDiffADGHjvprod, - :jprod_residual_backend => ADNLPModels.ZygoteADJprod, - :jtprod_residual_backend => ADNLPModels.ZygoteADJtprod, - :hprod_residual_backend => ADNLPModels.ForwardDiffADHvprod, - :jacobian_residual_backend => ADNLPModels.ZygoteADJacobian, - :hessian_residual_backend => ADNLPModels.ZygoteADHessian, - ), -) - # Automatically loads the code for Zygote with Requires import Zygote @@ -81,3 +62,19 @@ include("nlp/basic.jl") include("nls/basic.jl") include("nlp/nlpmodelstest.jl") include("nls/nlpmodelstest.jl") + +@testset "Basic NLP tests using $backend " for backend in (:zygote,) + test_autodiff_model("$backend", backend = backend) +end + +@testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in (:zygote,) + nlp_nlpmodelstest(backend) +end + +@testset "Basic NLS tests using $backend " for backend in (:zygote,) + autodiff_nls_test("$backend", backend = backend) +end + +@testset "Checking NLPModelsTest (NLS) tests with $backend" for backend in (:zygote,) + nls_nlpmodelstest(backend) +end