diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 219a812f..33d28bb9 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -11,18 +11,18 @@ steps: julia --color=yes --project=test -e 'include("test/gpu.jl")' timeout_in_minutes: 30 - # - label: "CPUs -- Enzyme.jl" - # plugins: - # - JuliaCI/julia#v1: - # version: "1.10" - # agents: - # queue: "juliaecosystem" - # os: "linux" - # arch: "x86_64" - # command: | - # julia --color=yes --project=test -e 'using Pkg; Pkg.add("Enzyme"); Pkg.develop(path="."); Pkg.instantiate()' - # julia --color=yes --project=test -e 'include("test/enzyme.jl")' - # timeout_in_minutes: 30 + - label: "CPUs -- Enzyme.jl" + plugins: + - JuliaCI/julia#v1: + version: "1.10" + agents: + queue: "juliaecosystem" + os: "linux" + arch: "x86_64" + command: | + julia --color=yes --project=test -e 'using Pkg; Pkg.add("Enzyme"); Pkg.develop(path="."); Pkg.instantiate()' + julia --color=yes --project=test -e 'include("test/enzyme.jl")' + timeout_in_minutes: 30 - label: "CPUs -- Zygote.jl" plugins: diff --git a/src/ADNLPModels.jl b/src/ADNLPModels.jl index a50d1005..9874adf8 100644 --- a/src/ADNLPModels.jl +++ b/src/ADNLPModels.jl @@ -207,7 +207,6 @@ Return the lagrangian function `ℓ(x) = obj_weight * f(x) + c(x)ᵀy`. function get_lag(nlp::AbstractADNLPModel, b::ADBackend, obj_weight::Real) return ℓ(x; obj_weight = obj_weight) = obj_weight * nlp.f(x) end - function get_lag(nlp::AbstractADNLPModel, b::ADBackend, obj_weight::Real, y::AbstractVector) if nlp.meta.nnln == 0 return get_lag(nlp, b, obj_weight) diff --git a/src/enzyme.jl b/src/enzyme.jl index 2469fb1a..fa6d9ab9 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -23,9 +23,10 @@ function EnzymeReverseADJacobian( return EnzymeReverseADJacobian() end -struct EnzymeReverseADHessian{T} <: ADBackend +struct EnzymeReverseADHessian{T,F} <: ADBackend seed::Vector{T} Hv::Vector{T} + f::F end function EnzymeReverseADHessian( @@ -41,11 +42,20 @@ function EnzymeReverseADHessian( seed = zeros(T, nvar) Hv = zeros(T, nvar) - return EnzymeReverseADHessian(seed, Hv) + return EnzymeReverseADHessian(seed, Hv, f) end -struct EnzymeReverseADHvprod{T} <: InPlaceADbackend - grad::Vector{T} +struct EnzymeReverseADHvprod{V, F, C, L} <: InPlaceADbackend + grad::V # length nvar, gradient buffer (primal in DuplicatedNoNeed) + hvbuf::V # length nvar, Hv output buffer (shadow in DuplicatedNoNeed) + xbuf::V # length nvar, input x buffer + vbuf::V # length nvar, input v buffer (tangent direction) + cx::V # length ncon, constraint output buffer + ybuf::V # length ncon, multiplier buffer for jth_hprod + f::F + c!::C + ℓ::L + ncon::Int end function EnzymeReverseADHvprod( @@ -56,12 +66,34 @@ function EnzymeReverseADHvprod( x0::AbstractVector{T} = rand(nvar), kwargs..., ) where {T} - grad = zeros(T, nvar) - return EnzymeReverseADHvprod(grad) + grad = fill!(similar(x0), zero(T)) + hvbuf = similar(x0) + xbuf = similar(x0) + vbuf = similar(x0) + cx = fill!(similar(x0, ncon), zero(T)) + ybuf = fill!(similar(x0, ncon), zero(T)) + + function ℓ(x, y, obj_weight, cx) + if ncon != 0 + c!(cx, x) + end + res = obj_weight * f(x) + if ncon != 0 + for i = 1:ncon + res += cx[i] * y[i] + end + end + return res + end + + return EnzymeReverseADHvprod(grad, hvbuf, xbuf, vbuf, cx, ybuf, f, c!, ℓ, ncon) end -struct EnzymeReverseADJprod{T} <: InPlaceADbackend - cx::Vector{T} +struct EnzymeReverseADJprod{V} <: InPlaceADbackend + cx::V # length ncon, primal output buffer + xbuf::V # length nvar, input x buffer + vbuf::V # length nvar, input v buffer (tangent direction) + jvbuf::V # length ncon, output Jv buffer end function EnzymeReverseADJprod( @@ -72,12 +104,18 @@ function EnzymeReverseADJprod( x0::AbstractVector{T} = rand(nvar), kwargs..., ) where {T} - cx = zeros(T, nvar) - return EnzymeReverseADJprod(cx) + cx = fill!(similar(x0, ncon), zero(T)) + xbuf = similar(x0) + vbuf = similar(x0) + jvbuf = fill!(similar(x0, ncon), zero(T)) + return EnzymeReverseADJprod(cx, xbuf, vbuf, jvbuf) end -struct EnzymeReverseADJtprod{T} <: InPlaceADbackend - cx::Vector{T} +struct EnzymeReverseADJtprod{V} <: InPlaceADbackend + cx::V # length ncon, primal output buffer + xbuf::V # length nvar, input x buffer + vbuf::V # length ncon, cotangent seed buffer + jtvbuf::V # length nvar, output Jtv buffer end function EnzymeReverseADJtprod( @@ -88,11 +126,14 @@ function EnzymeReverseADJtprod( x0::AbstractVector{T} = rand(nvar), kwargs..., ) where {T} - cx = zeros(T, nvar) - return EnzymeReverseADJtprod(cx) + cx = fill!(similar(x0, ncon), zero(T)) + xbuf = similar(x0) + vbuf = fill!(similar(x0, ncon), zero(T)) + jtvbuf = similar(x0) + return EnzymeReverseADJtprod(cx, xbuf, vbuf, jtvbuf) end -struct SparseEnzymeADJacobian{R, C, S} <: ADBackend +struct SparseEnzymeADJacobian{R, C, S, V} <: ADBackend nvar::Int ncon::Int rowval::Vector{Int} @@ -100,8 +141,9 @@ struct SparseEnzymeADJacobian{R, C, S} <: ADBackend nzval::Vector{R} result_coloring::C compressed_jacobian::S - v::Vector{R} - cx::Vector{R} + v::V + cx::V + xbuf::V end function SparseEnzymeADJacobian( @@ -152,7 +194,8 @@ function SparseEnzymeADJacobian( timer = @elapsed begin v = similar(x0) - cx = zeros(T, ncon) + cx = fill!(similar(x0, ncon), zero(T)) + xbuf = similar(x0) end show_time && println(" • Allocation of the AD buffers for the sparse Jacobian: $timer seconds.") @@ -166,23 +209,26 @@ function SparseEnzymeADJacobian( compressed_jacobian, v, cx, + xbuf, ) end -struct SparseEnzymeADHessian{R, C, S, L} <: ADBackend +struct SparseEnzymeADHessian{R, C, S, L, F, V} <: 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_icol::V compressed_hessian::S - v::Vector{R} - y::Vector{R} - grad::Vector{R} - cx::Vector{R} + v::V + y::V + grad::V + cx::V + f::F ℓ::L + xbuf::V end function SparseEnzymeADHessian( @@ -245,12 +291,17 @@ function SparseEnzymeADHessian( y = similar(x0, ncon) cx = similar(x0, ncon) grad = similar(x0) + xbuf = 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 + res = obj_weight * f(x) + if ncon != 0 + for i = 1:ncon + res += cx[i] * y[i] + end end return res end @@ -270,19 +321,74 @@ function SparseEnzymeADHessian( y, grad, cx, + f, ℓ, + xbuf, ) end @init begin @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin + function _gradient!(dx, f, x) + Enzyme.make_zero!(dx) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse), + f, + Enzyme.Active, + Enzyme.Duplicated(x, dx), + ) + return nothing + end + + function _hvp!(res, f, x, v) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Forward), + _gradient!, + res, + Enzyme.Const(f), + Enzyme.Duplicated(x, v), + ) + return nothing + end + + function _gradient!(dx, ℓ, x, y, obj_weight, cx) + Enzyme.make_zero!(dx) + dcx = Enzyme.make_zero(cx) + Enzyme.autodiff( + Enzyme.set_runtime_activity(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.set_runtime_activity(Enzyme.Forward), + _gradient!, + res, + Enzyme.Const(ℓ), + Enzyme.Duplicated(x, v), + Enzyme.Const(y), + Enzyme.Const(obj_weight), + Enzyme.Duplicated(cx, dcx), + ) + return nothing + end + function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) g = similar(x) - Enzyme.gradient!(Enzyme.Reverse, g, Enzyme.Const(f), x) + Enzyme.autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(x, g)) return g end function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) + Enzyme.make_zero!(g) Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(x, g)) return g end @@ -296,7 +402,8 @@ end fill!(b.seed, zero(T)) for i = 1:n b.seed[i] = one(T) - Enzyme.hvp!(b.Hv, Enzyme.Const(f), x, b.seed) + grad = Enzyme.make_zero(x) + _hvp!(Enzyme.DuplicatedNoNeed(grad, b.Hv), f, x, b.seed) view(hess, :, i) .= b.Hv b.seed[i] = zero(T) end @@ -304,22 +411,38 @@ end end function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) + copyto!(b.xbuf, x) + copyto!(b.vbuf, v) Enzyme.autodiff( Enzyme.Forward, Enzyme.Const(c!), - Enzyme.Duplicated(b.cx, Jv), - Enzyme.Duplicated(x, v), + Enzyme.Duplicated(b.cx, b.jvbuf), + Enzyme.Duplicated(b.xbuf, b.vbuf), ) + copyto!(Jv, b.jvbuf) return Jv end + # Wrapper that calls c!(y, x) but returns nothing. + # Enzyme reverse mode requires functions to return nothing (not their output array), + # otherwise it errors with "Duplicated Returns not yet handled". + function _void_c!(c!, y, x) + c!(y, x) + return nothing + end + function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) + copyto!(b.xbuf, x) + copyto!(b.vbuf, v) + Enzyme.make_zero!(b.jtvbuf) Enzyme.autodiff( Enzyme.Reverse, + Enzyme.Const(_void_c!), Enzyme.Const(c!), - Enzyme.Duplicated(b.cx, Jtv), - Enzyme.Duplicated(x, v), + Enzyme.Duplicated(b.cx, b.vbuf), + Enzyme.Duplicated(b.xbuf, b.jtvbuf), ) + copyto!(Jtv, b.jtvbuf) return Jtv end @@ -328,20 +451,15 @@ end Hv, x, v, - ℓ, + ℓ_unused, ::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), - ) + copyto!(b.xbuf, x) + copyto!(b.vbuf, v) + _hvp!(Enzyme.DuplicatedNoNeed(b.grad, b.hvbuf), b.ℓ, b.xbuf, b.vbuf, y, obj_weight, b.cx) + copyto!(Hv, b.hvbuf) return Hv end @@ -350,18 +468,65 @@ end Hv, x, v, - f, + f_unused, ::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), + copyto!(b.xbuf, x) + copyto!(b.vbuf, v) + _hvp!(Enzyme.DuplicatedNoNeed(b.grad, b.hvbuf), b.f, b.xbuf, b.vbuf) + @. Hv = obj_weight * b.hvbuf + return Hv + end + + # jth_hprod: Hessian-vector product for the j-th constraint. + # Uses the Lagrangian with y = e_j (unit vector) and obj_weight = 0, + # avoiding the closure x -> c(x)[j] that Enzyme can't handle. + function NLPModels.hprod!( + b::EnzymeReverseADHvprod, + nlp::ADModel, + x::AbstractVector, + v::AbstractVector, + j::Integer, + Hv::AbstractVector, + ) + copyto!(b.xbuf, x) + copyto!(b.vbuf, v) + b.cx .= 0 + # Build y = e_{j-nlin} (unit vector for nonlinear constraint index) + b.ybuf .= 0 + k = 0 + for i in nlp.meta.nln + k += 1 + if i == j + b.ybuf[k] = one(eltype(x)) + break + end + end + _hvp!( + Enzyme.DuplicatedNoNeed(b.grad, b.hvbuf), + b.ℓ, b.xbuf, b.vbuf, b.ybuf, zero(eltype(x)), b.cx, ) + copyto!(Hv, b.hvbuf) + return Hv + end + + # hprod_residual: Hessian-vector product for the i-th residual. + # Uses forward-over-reverse on F_i(x) = F(x)[i]. + function NLPModels.hprod_residual!( + b::EnzymeReverseADHvprod, + nls::AbstractADNLSModel, + x::AbstractVector, + v::AbstractVector, + i::Integer, + Hv::AbstractVector, + ) + F = get_F(nls) # out-of-place version + Fi(x) = F(x)[i] + copyto!(b.xbuf, x) + copyto!(b.vbuf, v) + _hvp!(Enzyme.DuplicatedNoNeed(b.grad, b.hvbuf), Fi, b.xbuf, b.vbuf) + copyto!(Hv, b.hvbuf) return Hv end @@ -394,6 +559,10 @@ end # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval) + # Enzyme.Duplicated requires primal and shadow to have the same type. + # Copy x into a pre-allocated buffer to ensure type match with b.v. + copyto!(b.xbuf, x) + groups = column_groups(b.result_coloring) for (icol, cols) in enumerate(groups) # Update the seed @@ -408,7 +577,7 @@ end Enzyme.Forward, Enzyme.Const(c!), Enzyme.Duplicated(b.cx, b.compressed_jacobian), - Enzyme.Duplicated(x, b.v), + Enzyme.Duplicated(b.xbuf, b.v), ) # Update the columns of the Jacobian that have the color `icol` @@ -483,6 +652,10 @@ end # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression A = SparseMatrixCSC(b.nvar, b.nvar, b.colptr, b.rowval, b.nzval) + # Enzyme.Duplicated requires primal and shadow to have the same type. + # Copy x into a pre-allocated buffer to ensure type match with b.v. + copyto!(b.xbuf, x) + groups = column_groups(b.result_coloring) for (icol, cols) in enumerate(groups) # Update the seed @@ -491,40 +664,10 @@ end 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.xbuf, b.v, y, obj_weight, diff --git a/test/Project.toml b/test/Project.toml index e6ae782e..da81f119 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617" diff --git a/test/enzyme.jl b/test/enzyme.jl index a844166e..0d1a9df6 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -7,14 +7,30 @@ using ADNLPModels: # Automatically loads the code for Enzyme with Requires import Enzyme +# Dummy constraint and Lagrangian functions for constructing an EnzymeReverseADHvprod +# with no constraints. These are only used by EnzymeReverseAD() below, which builds +# a minimal ADModelBackend for the low-level smoke tests (gradient, jacobian, …). +# Real models get their own ℓ constructed from f and c! at model-creation time. +_noop_c!(y, x) = nothing +_noop_ℓ(x, y, obj_weight, cx) = zero(eltype(x)) + +# Construct an EnzymeReverseADHvprod with pre-allocated buffers of size `n` and +# no constraints (ncon = 0). Fields: grad, hvbuf, xbuf, vbuf, cx, ybuf, f, c!, ℓ, ncon. +function _make_enzyme_hvprod(n) + ADNLPModels.EnzymeReverseADHvprod( + zeros(n), zeros(n), zeros(n), zeros(n), zeros(0), zeros(0), + identity, _noop_c!, _noop_ℓ, 0, + ) +end + EnzymeReverseAD() = ADNLPModels.ADModelBackend( ADNLPModels.EnzymeReverseADGradient(), - ADNLPModels.EnzymeReverseADHvprod(zeros(1)), - ADNLPModels.EnzymeReverseADJprod(zeros(1)), - ADNLPModels.EnzymeReverseADJtprod(zeros(1)), + _make_enzyme_hvprod(1), + ADNLPModels.EnzymeReverseADJprod(zeros(1), zeros(1), zeros(1), zeros(1)), + ADNLPModels.EnzymeReverseADJtprod(zeros(1), zeros(1), zeros(1), zeros(1)), ADNLPModels.EnzymeReverseADJacobian(), - ADNLPModels.EnzymeReverseADHessian(zeros(1), zeros(1)), - ADNLPModels.EnzymeReverseADHvprod(zeros(1)), + ADNLPModels.EnzymeReverseADHessian(zeros(1), zeros(1), identity), + _make_enzyme_hvprod(1), ADNLPModels.EmptyADbackend(), ADNLPModels.EmptyADbackend(), ADNLPModels.EmptyADbackend(), @@ -28,28 +44,8 @@ function mysum!(y, x) end function test_autodiff_backend_error() - @testset "Error without loading package - $backend" for backend in [:EnzymeReverseAD] + @testset "Enzyme basic operations - $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]) @@ -61,6 +57,26 @@ end test_autodiff_backend_error() +push!( + ADNLPModels.predefined_backend, + :enzyme_backend => Dict( + :gradient_backend => ADNLPModels.EnzymeReverseADGradient, + :jprod_backend => ADNLPModels.EnzymeReverseADJprod, + :jtprod_backend => ADNLPModels.EnzymeReverseADJtprod, + :hprod_backend => ADNLPModels.EnzymeReverseADHvprod, + :jacobian_backend => ADNLPModels.SparseEnzymeADJacobian, + :hessian_backend => ADNLPModels.SparseEnzymeADHessian, + :ghjvprod_backend => ADNLPModels.ForwardDiffADGHjvprod, + :jprod_residual_backend => ADNLPModels.EnzymeReverseADJprod, + :jtprod_residual_backend => ADNLPModels.EnzymeReverseADJtprod, + :hprod_residual_backend => ADNLPModels.EnzymeReverseADHvprod, + :jacobian_residual_backend => ADNLPModels.SparseEnzymeADJacobian, + :hessian_residual_backend => ADNLPModels.SparseEnzymeADHessian, + ), +) + +const test_enzyme = true + include("sparse_jacobian.jl") include("sparse_jacobian_nls.jl") include("sparse_hessian.jl") @@ -78,18 +94,20 @@ end list_sparse_hess_backend = ( ( ADNLPModels.SparseEnzymeADHessian, + "star coloring", Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), ), ( ADNLPModels.SparseEnzymeADHessian, + "acyclic coloring", 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) + for (backend, info, kw) in list_sparse_hess_backend + sparse_hessian(backend, info, kw) + sparse_hessian_nls(backend, info, kw) end end @@ -106,18 +124,18 @@ 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) +@testset "Basic NLP tests using enzyme_backend" begin + test_autodiff_model("enzyme_backend", backend = :enzyme_backend) end -@testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in (:enzyme,) - nlp_nlpmodelstest(backend) +@testset "Basic NLS tests using enzyme_backend" begin + autodiff_nls_test("enzyme_backend", backend = :enzyme_backend) end -@testset "Basic NLS tests using $backend " for backend in (:enzyme,) - autodiff_nls_test("$backend", backend = backend) +@testset "Checking NLPModelsTest (NLP) tests with enzyme_backend" begin + nlp_nlpmodelstest(:enzyme_backend) end -@testset "Checking NLPModelsTest (NLS) tests with $backend" for backend in (:enzyme,) - nls_nlpmodelstest(backend) +@testset "Checking NLPModelsTest (NLS) tests with enzyme_backend" begin + nls_nlpmodelstest(:enzyme_backend) end diff --git a/test/nlp/nlpmodelstest.jl b/test/nlp/nlpmodelstest.jl index 6be6611a..9f260bc1 100644 --- a/test/nlp/nlpmodelstest.jl +++ b/test/nlp/nlpmodelstest.jl @@ -18,7 +18,7 @@ function nlp_nlpmodelstest(backend) @testset "Check multiple precision" begin multiple_precision_nlp(nlp_from_T, exclude = [], linear_api = true) end - if backend != :enzyme + if backend != :enzyme && backend != :enzyme_backend @testset "Check view subarray" begin view_subarray_nlp(nlp_ad, exclude = []) end diff --git a/test/nls/nlpmodelstest.jl b/test/nls/nlpmodelstest.jl index f6b29882..07217013 100644 --- a/test/nls/nlpmodelstest.jl +++ b/test/nls/nlpmodelstest.jl @@ -42,7 +42,7 @@ function nls_nlpmodelstest(backend) @testset "Check multiple precision" begin multiple_precision_nls(nls_from_T, exclude = exclude, linear_api = true) end - if backend != :enzyme + if backend != :enzyme && backend != :enzyme_backend @testset "Check view subarray" begin view_subarray_nls.(nlss, exclude = exclude) end diff --git a/test/sparse_hessian.jl b/test/sparse_hessian.jl index 98c0cf72..4ab16d82 100644 --- a/test/sparse_hessian.jl +++ b/test/sparse_hessian.jl @@ -1,6 +1,5 @@ function sparse_hessian(backend, info, kw) @testset "Basic Hessian derivative with backend=$(backend) -- $info -- T=$(T)" for T in ( - Float32, Float64, ) c!(cx, x) = begin @@ -34,6 +33,7 @@ function sparse_hessian(backend, info, kw) # Test also the implementation of the backends b = nlp.adbackend.hessian_backend + @show b obj_weight = 0.5 @test nlp.meta.nnzh == ADNLPModels.get_nln_nnzh(b, nvar) ADNLPModels.hess_structure!(b, nlp, rows, cols) @@ -65,16 +65,29 @@ function sparse_hessian(backend, info, kw) ) @test nlp.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend - n = 4 - x = ones(T, 4) + # n = 4 + x0 = ones(T, 4) + function f(x) + n = length(x) + sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)) + # res = 0 + # n = length(x) + # for i in 1:(n-1) + # res += 100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 + # end + # res + end nlp = ADNLPModel( - x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)), - x, + # x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:(n - 1)), + # x -> sum(100 * (x[i + 1] - x[i]^2)^2 + (x[i] - 1)^2 for i = 1:3), + # x -> 100 * (x[2] - x[1]^2)^2 + (x[1] - 1)^2, + f, + x0, hessian_backend = backend, name = "Extended Rosenbrock"; kw..., ) - @test hess(nlp, x) == T[802 -400 0 0; -400 1002 -400 0; 0 -400 1002 -400; 0 0 -400 200] + @test hess(nlp, x0) == 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; kw...)