From b615dd8ac39297a9e02624e4ce4965c133b3bed9 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 26 Nov 2024 14:26:57 -0600 Subject: [PATCH 01/44] Enzyme WIP --- src/enzyme.jl | 167 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 158 insertions(+), 9 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index db5133fe..69b983ba 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -1,6 +1,9 @@ -struct EnzymeADGradient <: ADNLPModels.ADBackend end +struct EnzymeReverseADJacobian <: ADBackend end +struct EnzymeReverseADHessian <: ADBackend end -function EnzymeADGradient( +struct EnzymeReverseADGradient <: ADNLPModels.ADBackend end + +function EnzymeReverseADGradient( nvar::Integer, f, ncon::Integer = 0, @@ -8,14 +11,160 @@ function EnzymeADGradient( x0::AbstractVector = rand(nvar), kwargs..., ) - return EnzymeADGradient() + return EnzymeReverseADGradient() +end + +function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) + Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + return g +end + +function EnzymeReverseADJacobian( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + return EnzymeReverseADJacobian() +end + +jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) + +function EnzymeReverseADHessian( + nvar::Integer, + + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + @assert nvar > 0 + nnzh = nvar * (nvar + 1) / 2 + return EnzymeReverseADHessian() 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) - return g - end +function hessian(::EnzymeReverseADHessian, f, x) + seed = similar(x) + hess = zeros(eltype(x), length(x), length(x)) + fill!(seed, zero(x)) + for i in 1:length(x) + seed[i] = one(x) + Enzyme.hvp!(view(hess, i, :), f, x, seed) + seed[i] = zero(x) end + return hess +end + +struct EnzymeReverseADJprod <: InPlaceADBackend + x::Vector{Float64} +end + +function EnzymeReverseADJprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + x = zeros(nvar) + return EnzymeReverseADJprod(x) +end + +function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Forward, c!, Duplicated(b.x, Jv), Enzyme.Duplicated(x, v)) + return Jv +end + +struct EnzymeReverseADJtprod <: InPlaceADBackend + x::Vector{Float64} +end + +function EnzymeReverseADJtprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + x = zeros(nvar) + return EnzymeReverseADJtprod(x) +end + +function Jtvprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Reverse, c!, Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) + return Jtv +end + +struct EnzymeReverseADHprod <: InPlaceADBackend + grad::Vector{Float64} +end + +function EnzymeReverseADHvprod( + nvar::Integer, + f, + ncon::Integer = 0, + c!::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + grad = zeros(nvar) + return EnzymeReverseADHprod(grad) +end + +function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) + # What to do with args? + Enzyme.autodiff( + Forward, + gradient!, + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(f), + Duplicated(x, v), + ) + return Hv +end + +function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x::AbstractVector{T}, + v, + ℓ, + ::Val{:lag}, + y, + obj_weight::Real = one(T), +) + Enzyme.autodiff( + Forward, + gradient!, + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(ℓ), + Duplicated(x, v), + Const(y), + ) + + return Hv +end + +function Hvprod!( + b::EnzymeReverseADHvprod{T, S, Tagf}, + Hv, + x, + v, + f, + ::Val{:obj}, + obj_weight::Real = one(T), +) + Enzyme.autodiff( + Forward, + gradient!, + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(f), + Duplicated(x, v), + Const(y), + ) + return Hv end From 4a186be0bf68646c81b8f90a761e35adb2cda1bb Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 26 Nov 2024 14:42:10 -0600 Subject: [PATCH 02/44] cont. --- Project.toml | 1 + src/enzyme.jl | 16 ++++++++-------- test/enzyme.jl | 8 ++++++-- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/Project.toml b/Project.toml index afa56b1f..860f3615 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ version = "0.8.9" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" diff --git a/src/enzyme.jl b/src/enzyme.jl index 69b983ba..1750e689 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -56,7 +56,7 @@ function hessian(::EnzymeReverseADHessian, f, x) return hess end -struct EnzymeReverseADJprod <: InPlaceADBackend +struct EnzymeReverseADJprod <: InPlaceADbackend x::Vector{Float64} end @@ -76,7 +76,7 @@ function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) return Jv end -struct EnzymeReverseADJtprod <: InPlaceADBackend +struct EnzymeReverseADJtprod <: InPlaceADbackend x::Vector{Float64} end @@ -96,7 +96,7 @@ function Jtvprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) return Jtv end -struct EnzymeReverseADHprod <: InPlaceADBackend +struct EnzymeReverseADHvprod <: InPlaceADbackend grad::Vector{Float64} end @@ -109,7 +109,7 @@ function EnzymeReverseADHvprod( kwargs..., ) where {T} grad = zeros(nvar) - return EnzymeReverseADHprod(grad) + return EnzymeReverseADHvprod(grad) end function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) @@ -128,12 +128,12 @@ end function Hvprod!( b::EnzymeReverseADHvprod, Hv, - x::AbstractVector{T}, + x, v, ℓ, ::Val{:lag}, y, - obj_weight::Real = one(T), + obj_weight::Real = one(eltype(x)), ) Enzyme.autodiff( Forward, @@ -149,13 +149,13 @@ function Hvprod!( end function Hvprod!( - b::EnzymeReverseADHvprod{T, S, Tagf}, + b::EnzymeReverseADHvprod, Hv, x, v, f, ::Val{:obj}, - obj_weight::Real = one(T), + obj_weight::Real = one(eltype(x)), ) Enzyme.autodiff( Forward, diff --git a/test/enzyme.jl b/test/enzyme.jl index 504557a0..5654cec7 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -3,8 +3,12 @@ using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! -# Automatically loads the code for Enzyme with Requires -import Enzyme +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 #= ADNLPModels.EmptyADbackend(args...; kwargs...) = ADNLPModels.EmptyADbackend() From 52a79bd806e9f44484d57f37adaad93250133e11 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 26 Nov 2024 15:56:30 -0600 Subject: [PATCH 03/44] cont. --- Project.toml | 3 +++ src/ADNLPModels.jl | 2 +- src/enzyme.jl | 22 +++++++++++----- test/enzyme.jl | 64 +++++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 83 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index 860f3615..a397f6a4 100644 --- a/Project.toml +++ b/Project.toml @@ -7,7 +7,10 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" Requires = "ae029012-a4dd-5104-9daa-d747884805df" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" diff --git a/src/ADNLPModels.jl b/src/ADNLPModels.jl index a50d1005..8a474401 100644 --- a/src/ADNLPModels.jl +++ b/src/ADNLPModels.jl @@ -7,7 +7,7 @@ using LinearAlgebra, SparseArrays using ADTypes: ADTypes, AbstractColoringAlgorithm, AbstractSparsityDetector using SparseConnectivityTracer: TracerSparsityDetector using SparseMatrixColorings -using ForwardDiff, ReverseDiff +using ForwardDiff, ReverseDiff, Enzyme # JSO using NLPModels diff --git a/src/enzyme.jl b/src/enzyme.jl index 1750e689..6227648f 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -14,6 +14,12 @@ function EnzymeReverseADGradient( return EnzymeReverseADGradient() end +function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) + g = similar(x) + Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + return g +end + function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) return g @@ -47,11 +53,13 @@ end function hessian(::EnzymeReverseADHessian, f, x) seed = similar(x) hess = zeros(eltype(x), length(x), length(x)) - fill!(seed, zero(x)) + fill!(seed, zero(eltype(x))) + tmp = similar(x) for i in 1:length(x) - seed[i] = one(x) - Enzyme.hvp!(view(hess, i, :), f, x, seed) - seed[i] = zero(x) + seed[i] = one(eltype(seed)) + Enzyme.hvp!(tmp, f, x, seed) + hess[:, i] .= tmp + seed[i] = zero(eltype(seed)) end return hess end @@ -72,7 +80,9 @@ function EnzymeReverseADJprod( end function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Forward, c!, Duplicated(b.x, Jv), Enzyme.Duplicated(x, v)) + @show c!(x) + @show Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(x, v)) + error("This is BAD") return Jv end @@ -91,7 +101,7 @@ function EnzymeReverseADJtprod( return EnzymeReverseADJtprod(x) end -function Jtvprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) +function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) Enzyme.autodiff(Enzyme.Reverse, c!, Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) return Jtv end diff --git a/test/enzyme.jl b/test/enzyme.jl index 5654cec7..60932fa5 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -10,8 +10,70 @@ for problem in NLPModelsTest.nls_problems include("nls/problems/$(lowercase(problem)).jl") end +EnzymeReverseAD() = ADNLPModels.ADModelBackend( + ADNLPModels.EnzymeReverseADGradient(), + ADNLPModels.EnzymeReverseADHvprod(zeros(1)), + ADNLPModels.EnzymeReverseADJprod(zeros(1)), + ADNLPModels.EnzymeReverseADJtprod(zeros(1)), + ADNLPModels.EnzymeReverseADJacobian(), + ADNLPModels.EnzymeReverseADHessian(), + ADNLPModels.EnzymeReverseADHvprod(zeros(1)), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), + ADNLPModels.EmptyADbackend(), +) + +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, identity, [1.0]) + hessian(adbackend.hessian_backend, sum, [1.0]) + Jprod!( + adbackend.jprod_backend, + [1.0], + identity, + [1.0], + [1.0], + Val(:c), + ) + # Jtprod!( + # adbackend.jtprod_backend, + # [1.0], + # identity, + # [1.0], + # [1.0], + # Val(:c), + # ) + end +end + +test_autodiff_backend_error() #= -ADNLPModels.EmptyADbackend(args...; kwargs...) = ADNLPModels.EmptyADbackend() +# ADNLPModels.EmptyADbackend(args...; kwargs...) = ADNLPModels.EmptyADbackend() names = OptimizationProblems.meta[!, :name] list_excluded_enzyme = [ From 5de25771662b2f88afa51796b42fdb9fbc78c8bc Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Wed, 27 Nov 2024 12:03:11 -0600 Subject: [PATCH 04/44] Enzyme WIP --- src/enzyme.jl | 21 ++++--- test/enzyme.jl | 121 ++++++++++++-------------------------- test/nlp/basic.jl | 14 ++--- test/nlp/nlpmodelstest.jl | 29 ++++----- 4 files changed, 69 insertions(+), 116 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 6227648f..c2371d70 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -1,7 +1,7 @@ struct EnzymeReverseADJacobian <: ADBackend end struct EnzymeReverseADHessian <: ADBackend end -struct EnzymeReverseADGradient <: ADNLPModels.ADBackend end +struct EnzymeReverseADGradient <: InPlaceADbackend end function EnzymeReverseADGradient( nvar::Integer, @@ -16,12 +16,13 @@ end function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) g = similar(x) - Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + Enzyme.gradient!(Reverse, g, Const(f), x) return g end function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) - Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) return g end @@ -57,7 +58,7 @@ function hessian(::EnzymeReverseADHessian, f, x) tmp = similar(x) for i in 1:length(x) seed[i] = one(eltype(seed)) - Enzyme.hvp!(tmp, f, x, seed) + Enzyme.hvp!(tmp, Const(f), x, seed) hess[:, i] .= tmp seed[i] = zero(eltype(seed)) end @@ -80,9 +81,7 @@ function EnzymeReverseADJprod( end function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) - @show c!(x) - @show Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(x, v)) - error("This is BAD") + Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x,Jv), Duplicated(x, v)) return Jv end @@ -102,7 +101,7 @@ function EnzymeReverseADJtprod( end function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Reverse, c!, Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) + Enzyme.autodiff(Enzyme.Reverse, Const(c!), Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) return Jtv end @@ -126,7 +125,7 @@ function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) # What to do with args? Enzyme.autodiff( Forward, - gradient!, + Const(Enzyme.gradient!), Const(Reverse), DuplicatedNoNeed(b.grad, Hv), Const(f), @@ -147,7 +146,7 @@ function Hvprod!( ) Enzyme.autodiff( Forward, - gradient!, + Const(Enzyme.gradient!), Const(Reverse), DuplicatedNoNeed(b.grad, Hv), Const(ℓ), @@ -169,7 +168,7 @@ function Hvprod!( ) Enzyme.autodiff( Forward, - gradient!, + Const(Enzyme.gradient!), Const(Reverse), DuplicatedNoNeed(b.grad, Hv), Const(f), diff --git a/test/enzyme.jl b/test/enzyme.jl index 60932fa5..51c87ec8 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -24,7 +24,10 @@ EnzymeReverseAD() = ADNLPModels.ADModelBackend( 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)() @@ -50,100 +53,50 @@ function test_autodiff_backend_error() # ) gradient(adbackend.gradient_backend, sum, [1.0]) gradient!(adbackend.gradient_backend, [1.0], sum, [1.0]) - jacobian(adbackend.jacobian_backend, identity, [1.0]) + jacobian(adbackend.jacobian_backend, sum, [1.0]) hessian(adbackend.hessian_backend, sum, [1.0]) Jprod!( adbackend.jprod_backend, [1.0], - identity, + sum!, + [1.0], + [1.0], + Val(:c), + ) + Jtprod!( + adbackend.jtprod_backend, + [1.0], + mysum!, [1.0], [1.0], Val(:c), ) - # Jtprod!( - # adbackend.jtprod_backend, - # [1.0], - # identity, - # [1.0], - # [1.0], - # Val(:c), - # ) end end test_autodiff_backend_error() -#= -# 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)) -end -=# +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.EnzymeReverseADJacobian, + :hessian_backend => ADNLPModels.EnzymeReverseADHessian, + :ghjvprod_backend => ADNLPModels.ForwardDiffADGHjvprod, + :jprod_residual_backend => ADNLPModels.EnzymeReverseADJprod, + :jtprod_residual_backend => ADNLPModels.EnzymeReverseADJtprod, + :hprod_residual_backend => ADNLPModels.EnzymeReverseADHvprod, + :jacobian_residual_backend => ADNLPModels.EnzymeReverseADJacobian, + :hessian_residual_backend => ADNLPModels.EnzymeReverseADHessian, + ), +) + +include("utils.jl") +include("nlp/basic.jl") +include("nls/basic.jl") +include("nlp/nlpmodelstest.jl") +include("nls/nlpmodelstest.jl") -#= -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..d09d2912 100644 --- a/test/nlp/basic.jl +++ b/test/nlp/basic.jl @@ -17,13 +17,13 @@ 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) - y = 2x .+ 3 + randn(100) * 0.1 - regr = LinearRegression(x, y) - nlp = ADNLPModel(regr, ones(2); kwargs...) - β = [ones(100) x] \ y - @test abs(obj(nlp, β) - norm(y .- β[1] - β[2] * x)^2 / 2) < 1e-12 - @test norm(grad(nlp, β)) < 1e-12 + # x = range(-1, stop = 1, length = 100) + # y = 2x .+ 3 + randn(100) * 0.1 + # regr = LinearRegression(x, y) + # nlp = ADNLPModel(regr, ones(2); kwargs...) + # β = [ones(100) x] \ y + # @test abs(obj(nlp, β) - norm(y .- β[1] - β[2] * x)^2 / 2) < 1e-12 + # @test norm(grad(nlp, β)) < 1e-12 test_getter_setter(nlp) diff --git a/test/nlp/nlpmodelstest.jl b/test/nlp/nlpmodelstest.jl index 78bf56ec..cd3dc31d 100644 --- a/test/nlp/nlpmodelstest.jl +++ b/test/nlp/nlpmodelstest.jl @@ -1,5 +1,6 @@ -@testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in - keys(ADNLPModels.predefined_backend) +# @testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in +# keys(ADNLPModels.predefined_backend) +backend = :enzyme_backend @testset "Checking NLPModelsTest tests on problem $problem" for problem in NLPModelsTest.nlp_problems nlp_from_T = eval(Meta.parse(lowercase(problem) * "_autodiff")) @@ -12,17 +13,17 @@ @testset "Check Consistency" begin consistent_nlps(nlps, exclude = [], linear_api = true, reimplemented = ["jtprod"]) end - @testset "Check dimensions" begin - check_nlp_dimensions(nlp_ad, exclude = [], linear_api = true) - end - @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 = []) - end - @testset "Check coordinate memory" begin - coord_memory_nlp(nlp_ad, exclude = [], linear_api = true) - end + # @testset "Check dimensions" begin + # check_nlp_dimensions(nlp_ad, exclude = [], linear_api = true) + # end + # @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 = []) + # end + # @testset "Check coordinate memory" begin + # coord_memory_nlp(nlp_ad, exclude = [], linear_api = true) + # end end end From aee2282f89eca2b7424d5284473d17ee513fa376 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 13:04:11 -0600 Subject: [PATCH 05/44] Add a sparse Jacobian with Enzyme.jl --- docs/src/sparse.md | 4 +- src/enzyme.jl | 2 +- src/sparse_jacobian.jl | 133 ++++++++++++++++++++++++++++++------ src/sparsity_pattern.jl | 12 +++- test/sparse_jacobian.jl | 6 +- test/sparse_jacobian_nls.jl | 6 +- 6 files changed, 132 insertions(+), 31 deletions(-) diff --git a/docs/src/sparse.md b/docs/src/sparse.md index 73a60e61..b8535979 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`, and `SparseReverseADHessian`) 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`. 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 c2371d70..7577ba92 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -81,7 +81,7 @@ function EnzymeReverseADJprod( end function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x,Jv), Duplicated(x, v)) + Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v)) return Jv end diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl index db027de0..f84b42eb 100644 --- a/src/sparse_jacobian.jl +++ b/src/sparse_jacobian.jl @@ -64,25 +64,6 @@ function SparseADJacobian( ) end -function get_nln_nnzj(b::SparseADJacobian, nvar, ncon) - length(b.rowval) -end - -function NLPModels.jac_structure!( - b::SparseADJacobian, - 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!( ℓ!::Function, b::SparseADJacobian{Tag}, @@ -111,8 +92,116 @@ function sparse_jac_coord!( return vals 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} + buffer::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) + buffer = zeros(T, ncon) + + SparseEnzymeADJacobian( + nvar, + ncon, + rowval, + colptr, + nzval, + result_coloring, + compressed_jacobian, + v, + buffer, + ) +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, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), 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 get_nln_nnzj(b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nvar, ncon) + length(b.rowval) +end + +function NLPModels.jac_structure!( + b::Union{SparseADJacobian, 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 NLPModels.jac_coord!( - b::SparseADJacobian, + b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nlp::ADModel, x::AbstractVector, vals::AbstractVector, @@ -122,7 +211,7 @@ function NLPModels.jac_coord!( end function NLPModels.jac_structure_residual!( - b::SparseADJacobian, + b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nls::AbstractADNLSModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}, @@ -137,7 +226,7 @@ function NLPModels.jac_structure_residual!( end function NLPModels.jac_coord_residual!( - b::SparseADJacobian, + b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nls::AbstractADNLSModel, x::AbstractVector, vals::AbstractVector, diff --git a/src/sparsity_pattern.jl b/src/sparsity_pattern.jl index daf5d6e9..cc86f541 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 @@ -102,7 +106,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 diff --git a/test/sparse_jacobian.jl b/test/sparse_jacobian.jl index da0a8e5f..349ea980 100644 --- a/test/sparse_jacobian.jl +++ b/test/sparse_jacobian.jl @@ -1,5 +1,7 @@ list_sparse_jac_backend = - ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict())) + ((ADNLPModels.SparseADJacobian, Dict()), + (ADNLPModels.SparseEnzymeADJacobian, Dict()), + (ADNLPModels.ForwardDiffADJacobian, Dict())) dt = (Float32, Float64) @@ -51,7 +53,7 @@ dt = (Float32, Float64) 0 1 ] - if backend == ADNLPModels.SparseADJacobian + if backend != ADNLPModels.ForwardDiffADJacobian J_sp = get_sparsity_pattern(nlp, :jacobian) @test J_sp == SparseMatrixCSC{Bool, Int}([ 1 0 diff --git a/test/sparse_jacobian_nls.jl b/test/sparse_jacobian_nls.jl index 4405e7f7..55d4af10 100644 --- a/test/sparse_jacobian_nls.jl +++ b/test/sparse_jacobian_nls.jl @@ -1,5 +1,7 @@ list_sparse_jac_backend = - ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict())) + ((ADNLPModels.SparseADJacobian, Dict()), + (ADNLPModels.SparseEnzymeADJacobian, Dict()), + (ADNLPModels.ForwardDiffADJacobian, Dict())) dt = (Float32, Float64) @@ -43,7 +45,7 @@ dt = (Float32, Float64) 0 1 ] - if backend == ADNLPModels.SparseADJacobian + if backend != ADNLPModels.ForwardDiffADJacobian J_sp = get_sparsity_pattern(nls, :jacobian_residual) @test J_sp == SparseMatrixCSC{Bool, Int}([ 1 0 From 2e08e9aec71a54a7e1b9e7b5f98a5feb562fc6f8 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 13:46:35 -0600 Subject: [PATCH 06/44] Update tests --- src/enzyme.jl | 144 +++++++++++++++++++++++++++++++++++++++++ src/sparse_jacobian.jl | 99 ++-------------------------- test/enzyme.jl | 1 - 3 files changed, 149 insertions(+), 95 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 7577ba92..52e16f6a 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -177,3 +177,147 @@ function Hvprod!( ) return Hv end + +# Sparse Jacobian +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} + buffer::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) + buffer = zeros(T, ncon) + + SparseEnzymeADJacobian( + nvar, + ncon, + rowval, + colptr, + nzval, + result_coloring, + compressed_jacobian, + v, + buffer, + ) +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, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), 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 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 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 diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl index f84b42eb..360c519b 100644 --- a/src/sparse_jacobian.jl +++ b/src/sparse_jacobian.jl @@ -92,101 +92,12 @@ function sparse_jac_coord!( return vals 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} - buffer::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) - buffer = zeros(T, ncon) - - SparseEnzymeADJacobian( - nvar, - ncon, - rowval, - colptr, - nzval, - result_coloring, - compressed_jacobian, - v, - buffer, - ) -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, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), 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 get_nln_nnzj(b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nvar, ncon) +function get_nln_nnzj(b::SparseADJacobian, nvar, ncon) length(b.rowval) end function NLPModels.jac_structure!( - b::Union{SparseADJacobian, SparseEnzymeADJacobian}, + b::SparseADJacobian, nlp::ADModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}, @@ -201,7 +112,7 @@ function NLPModels.jac_structure!( end function NLPModels.jac_coord!( - b::Union{SparseADJacobian, SparseEnzymeADJacobian}, + b::SparseADJacobian, nlp::ADModel, x::AbstractVector, vals::AbstractVector, @@ -211,7 +122,7 @@ function NLPModels.jac_coord!( end function NLPModels.jac_structure_residual!( - b::Union{SparseADJacobian, SparseEnzymeADJacobian}, + b::SparseADJacobian, nls::AbstractADNLSModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}, @@ -226,7 +137,7 @@ function NLPModels.jac_structure_residual!( end function NLPModels.jac_coord_residual!( - b::Union{SparseADJacobian, SparseEnzymeADJacobian}, + b::SparseADJacobian, nls::AbstractADNLSModel, x::AbstractVector, vals::AbstractVector, diff --git a/test/enzyme.jl b/test/enzyme.jl index 51c87ec8..08552352 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -99,4 +99,3 @@ include("nlp/basic.jl") include("nls/basic.jl") include("nlp/nlpmodelstest.jl") include("nls/nlpmodelstest.jl") - From 298e218350c3145d74e4148b6f726a6dcfb85c06 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 13:55:54 -0600 Subject: [PATCH 07/44] Update tests --- test/enzyme.jl | 7 +++++++ test/runtests.jl | 2 ++ test/sparse_jacobian.jl | 10 ++++++---- test/sparse_jacobian_nls.jl | 10 ++++++---- 4 files changed, 21 insertions(+), 8 deletions(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index 08552352..b000c937 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -99,3 +99,10 @@ include("nlp/basic.jl") include("nls/basic.jl") include("nlp/nlpmodelstest.jl") include("nls/nlpmodelstest.jl") + +const test_enzyme = true + +include("sparse_jacobian.jl") +include("sparse_jacobian_nls.jl") +# include("sparse_hessian.jl") +# include("sparse_hessian_nls.jl") diff --git a/test/runtests.jl b/test/runtests.jl index 1d764265..94d14ffa 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,6 +4,8 @@ using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! +const test_enzyme = false + @testset "Test sparsity pattern of Jacobian and Hessian" begin f(x) = sum(x .^ 2) c(x) = x diff --git a/test/sparse_jacobian.jl b/test/sparse_jacobian.jl index 349ea980..b8a22259 100644 --- a/test/sparse_jacobian.jl +++ b/test/sparse_jacobian.jl @@ -1,7 +1,9 @@ -list_sparse_jac_backend = - ((ADNLPModels.SparseADJacobian, Dict()), - (ADNLPModels.SparseEnzymeADJacobian, Dict()), - (ADNLPModels.ForwardDiffADJacobian, Dict())) +if test_enzyme + list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) +else + list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), + (ADNLPModels.ForwardDiffADJacobian, Dict())) +end dt = (Float32, Float64) diff --git a/test/sparse_jacobian_nls.jl b/test/sparse_jacobian_nls.jl index 55d4af10..cecb05f9 100644 --- a/test/sparse_jacobian_nls.jl +++ b/test/sparse_jacobian_nls.jl @@ -1,7 +1,9 @@ -list_sparse_jac_backend = - ((ADNLPModels.SparseADJacobian, Dict()), - (ADNLPModels.SparseEnzymeADJacobian, Dict()), - (ADNLPModels.ForwardDiffADJacobian, Dict())) +if test_enzyme + list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) +else + list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), + (ADNLPModels.ForwardDiffADJacobian, Dict())) +end dt = (Float32, Float64) From 2cc2372039557e2d004661b43b1819c1f5d1b53b Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 14:02:21 -0600 Subject: [PATCH 08/44] Update nlpmodelstest.jl --- test/nlp/nlpmodelstest.jl | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/test/nlp/nlpmodelstest.jl b/test/nlp/nlpmodelstest.jl index cd3dc31d..78bf56ec 100644 --- a/test/nlp/nlpmodelstest.jl +++ b/test/nlp/nlpmodelstest.jl @@ -1,6 +1,5 @@ -# @testset "Checking NLPModelsTest (NLP) tests with $backend" for backend in -# keys(ADNLPModels.predefined_backend) -backend = :enzyme_backend +@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 nlp_from_T = eval(Meta.parse(lowercase(problem) * "_autodiff")) @@ -13,17 +12,17 @@ backend = :enzyme_backend @testset "Check Consistency" begin consistent_nlps(nlps, exclude = [], linear_api = true, reimplemented = ["jtprod"]) end - # @testset "Check dimensions" begin - # check_nlp_dimensions(nlp_ad, exclude = [], linear_api = true) - # end - # @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 = []) - # end - # @testset "Check coordinate memory" begin - # coord_memory_nlp(nlp_ad, exclude = [], linear_api = true) - # end + @testset "Check dimensions" begin + check_nlp_dimensions(nlp_ad, exclude = [], linear_api = true) + end + @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 = []) + end + @testset "Check coordinate memory" begin + coord_memory_nlp(nlp_ad, exclude = [], linear_api = true) + end end end From f768b4e8ba3b2555a232d2034de57071ff26987f Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 14:04:00 -0600 Subject: [PATCH 09/44] Fix Project.toml --- Project.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/Project.toml b/Project.toml index a397f6a4..afa56b1f 100644 --- a/Project.toml +++ b/Project.toml @@ -4,13 +4,9 @@ version = "0.8.9" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" -NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" -NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" Requires = "ae029012-a4dd-5104-9daa-d747884805df" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" From 668ce5382cc369c500a02eb4490c28bfba14b6d6 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 14:05:59 -0600 Subject: [PATCH 10/44] Fix the diff --- src/ADNLPModels.jl | 2 +- src/sparse_jacobian.jl | 38 +++++++++++++++++++------------------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/src/ADNLPModels.jl b/src/ADNLPModels.jl index 8a474401..a50d1005 100644 --- a/src/ADNLPModels.jl +++ b/src/ADNLPModels.jl @@ -7,7 +7,7 @@ using LinearAlgebra, SparseArrays using ADTypes: ADTypes, AbstractColoringAlgorithm, AbstractSparsityDetector using SparseConnectivityTracer: TracerSparsityDetector using SparseMatrixColorings -using ForwardDiff, ReverseDiff, Enzyme +using ForwardDiff, ReverseDiff # JSO using NLPModels diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl index 360c519b..db027de0 100644 --- a/src/sparse_jacobian.jl +++ b/src/sparse_jacobian.jl @@ -64,6 +64,25 @@ function SparseADJacobian( ) end +function get_nln_nnzj(b::SparseADJacobian, nvar, ncon) + length(b.rowval) +end + +function NLPModels.jac_structure!( + b::SparseADJacobian, + 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!( ℓ!::Function, b::SparseADJacobian{Tag}, @@ -92,25 +111,6 @@ function sparse_jac_coord!( return vals end -function get_nln_nnzj(b::SparseADJacobian, nvar, ncon) - length(b.rowval) -end - -function NLPModels.jac_structure!( - b::SparseADJacobian, - 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 NLPModels.jac_coord!( b::SparseADJacobian, nlp::ADModel, From 293a879723476d46944ce720d6180416677826f5 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 14:17:59 -0600 Subject: [PATCH 11/44] Fix the tests --- src/enzyme.jl | 362 +++++++++++++++++++++++++------------------------ test/enzyme.jl | 3 + 2 files changed, 187 insertions(+), 178 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 52e16f6a..51ab0aa7 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -1,6 +1,3 @@ -struct EnzymeReverseADJacobian <: ADBackend end -struct EnzymeReverseADHessian <: ADBackend end - struct EnzymeReverseADGradient <: InPlaceADbackend end function EnzymeReverseADGradient( @@ -14,17 +11,7 @@ function EnzymeReverseADGradient( return EnzymeReverseADGradient() end -function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) - g = similar(x) - # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - Enzyme.gradient!(Reverse, g, Const(f), x) - return g -end - -function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) - Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - return g -end +struct EnzymeReverseADJacobian <: ADBackend end function EnzymeReverseADJacobian( nvar::Integer, @@ -36,7 +23,7 @@ function EnzymeReverseADJacobian( return EnzymeReverseADJacobian() end -jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) +struct EnzymeReverseADHessian <: ADBackend end function EnzymeReverseADHessian( nvar::Integer, @@ -51,18 +38,20 @@ function EnzymeReverseADHessian( return EnzymeReverseADHessian() end -function hessian(::EnzymeReverseADHessian, f, x) - seed = similar(x) - hess = zeros(eltype(x), length(x), length(x)) - fill!(seed, zero(eltype(x))) - tmp = similar(x) - for i in 1:length(x) - seed[i] = one(eltype(seed)) - Enzyme.hvp!(tmp, Const(f), x, seed) - hess[:, i] .= tmp - seed[i] = zero(eltype(seed)) - end - return hess +struct EnzymeReverseADHvprod <: InPlaceADbackend + grad::Vector{Float64} +end + +function EnzymeReverseADHvprod( + nvar::Integer, + f, + ncon::Integer = 0, + c!::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + grad = zeros(nvar) + return EnzymeReverseADHvprod(grad) end struct EnzymeReverseADJprod <: InPlaceADbackend @@ -80,11 +69,6 @@ function EnzymeReverseADJprod( return EnzymeReverseADJprod(x) end -function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v)) - return Jv -end - struct EnzymeReverseADJtprod <: InPlaceADbackend x::Vector{Float64} end @@ -100,85 +84,6 @@ function EnzymeReverseADJtprod( return EnzymeReverseADJtprod(x) end -function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Reverse, Const(c!), Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) - return Jtv -end - -struct EnzymeReverseADHvprod <: InPlaceADbackend - grad::Vector{Float64} -end - -function EnzymeReverseADHvprod( - nvar::Integer, - f, - ncon::Integer = 0, - c!::Function = (args...) -> []; - x0::AbstractVector{T} = rand(nvar), - kwargs..., -) where {T} - grad = zeros(nvar) - return EnzymeReverseADHvprod(grad) -end - -function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) - # What to do with args? - Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(f), - Duplicated(x, v), - ) - return Hv -end - -function Hvprod!( - b::EnzymeReverseADHvprod, - Hv, - x, - v, - ℓ, - ::Val{:lag}, - y, - obj_weight::Real = one(eltype(x)), -) - Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(ℓ), - Duplicated(x, v), - Const(y), - ) - - return Hv -end - -function Hvprod!( - b::EnzymeReverseADHvprod, - Hv, - x, - v, - f, - ::Val{:obj}, - obj_weight::Real = one(eltype(x)), -) - Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(f), - Duplicated(x, v), - Const(y), - ) - return Hv -end - -# Sparse Jacobian struct SparseEnzymeADJacobian{R, C, S} <: ADBackend nvar::Int ncon::Int @@ -240,84 +145,185 @@ function SparseEnzymeADJacobian( ) 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 +@init begin + @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin + + function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) + g = similar(x) + # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + Enzyme.gradient!(Reverse, g, Const(f), x) + return g + end + + function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) + Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + return g + end + + jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) + + function hessian(::EnzymeReverseADHessian, f, x) + seed = similar(x) + hess = zeros(eltype(x), length(x), length(x)) + fill!(seed, zero(eltype(x))) + tmp = similar(x) + for i in 1:length(x) + seed[i] = one(eltype(seed)) + Enzyme.hvp!(tmp, Const(f), x, seed) + hess[:, i] .= tmp + seed[i] = zero(eltype(seed)) end + return hess + end - # b.compressed_jacobian is just a vector Jv here - # We don't use the vector mode - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), Duplicated(x, b.v)) + function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v)) + return Jv + end - # Update the columns of the Jacobian that have the color `icol` - decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring) + function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Reverse, Const(c!), Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) + return Jtv end - vals .= b.nzval - return vals -end -function get_nln_nnzj(b::SparseEnzymeADJacobian, nvar, ncon) - length(b.rowval) -end + function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) + # What to do with args? + Enzyme.autodiff( + Forward, + Const(Enzyme.gradient!), + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(f), + Duplicated(x, v), + ) + return Hv + 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 + function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x, + v, + ℓ, + ::Val{:lag}, + y, + obj_weight::Real = one(eltype(x)), + ) + Enzyme.autodiff( + Forward, + Const(Enzyme.gradient!), + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(ℓ), + Duplicated(x, v), + Const(y), + ) + + return Hv + end + + function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x, + v, + f, + ::Val{:obj}, + obj_weight::Real = one(eltype(x)), + ) + Enzyme.autodiff( + Forward, + Const(Enzyme.gradient!), + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(f), + Duplicated(x, v), + Const(y), + ) + return Hv + end + + # Sparse Jacobian + 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, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), 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 - return rows, cols -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 get_nln_nnzj(b::SparseEnzymeADJacobian, nvar, ncon) + length(b.rowval) + 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 + 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 - 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 + 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 end + diff --git a/test/enzyme.jl b/test/enzyme.jl index b000c937..e0636e3e 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -3,6 +3,9 @@ using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! +# Automatically loads the code for Enzyme with Requires +import Enzyme + for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] include("nlp/problems/$(lowercase(problem)).jl") end From 875f250ca57f91558ddce34c2f91aa1f0837a880 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 18:24:25 -0600 Subject: [PATCH 12/44] Fix the tests --- src/enzyme.jl | 309 +++++++++++++++++++++++++------------------------- 1 file changed, 155 insertions(+), 154 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 51ab0aa7..c402cb06 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -148,182 +148,183 @@ end @init begin @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin - function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) - g = similar(x) - # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - Enzyme.gradient!(Reverse, g, Const(f), x) - return g - end + function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) + g = similar(x) + # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + Enzyme.gradient!(Reverse, g, Const(f), x) + return g + end - function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) - Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - return g - end + function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) + Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + return g + end - jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) - - function hessian(::EnzymeReverseADHessian, f, x) - seed = similar(x) - hess = zeros(eltype(x), length(x), length(x)) - fill!(seed, zero(eltype(x))) - tmp = similar(x) - for i in 1:length(x) - seed[i] = one(eltype(seed)) - Enzyme.hvp!(tmp, Const(f), x, seed) - hess[:, i] .= tmp - seed[i] = zero(eltype(seed)) + jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) + + function hessian(::EnzymeReverseADHessian, f, x) + seed = similar(x) + hess = zeros(eltype(x), length(x), length(x)) + fill!(seed, zero(eltype(x))) + tmp = similar(x) + for i in 1:length(x) + seed[i] = one(eltype(seed)) + Enzyme.hvp!(tmp, Const(f), x, seed) + hess[:, i] .= tmp + seed[i] = zero(eltype(seed)) + end + return hess end - return hess - end - function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v)) - return Jv - end + function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v)) + return Jv + end - function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Reverse, Const(c!), Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) - return Jtv - end + function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) + Enzyme.autodiff(Enzyme.Reverse, Const(c!), Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) + return Jtv + end - function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) - # What to do with args? - Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(f), - Duplicated(x, v), - ) - return Hv - end + function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) + # What to do with args? + Enzyme.autodiff( + Forward, + Const(Enzyme.gradient!), + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(f), + Duplicated(x, v), + ) + return Hv + end - function Hvprod!( - b::EnzymeReverseADHvprod, - Hv, - x, - v, - ℓ, - ::Val{:lag}, - y, - obj_weight::Real = one(eltype(x)), - ) - Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(ℓ), - Duplicated(x, v), - Const(y), + function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x, + v, + ℓ, + ::Val{:lag}, + y, + obj_weight::Real = one(eltype(x)), ) + Enzyme.autodiff( + Forward, + Const(Enzyme.gradient!), + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(ℓ), + Duplicated(x, v), + Const(y), + ) + + return Hv + end - return Hv - end - - function Hvprod!( - b::EnzymeReverseADHvprod, - Hv, - x, - v, - f, - ::Val{:obj}, - obj_weight::Real = one(eltype(x)), - ) - Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(f), - Duplicated(x, v), - Const(y), + function Hvprod!( + b::EnzymeReverseADHvprod, + Hv, + x, + v, + f, + ::Val{:obj}, + obj_weight::Real = one(eltype(x)), ) - return Hv - end + Enzyme.autodiff( + Forward, + Const(Enzyme.gradient!), + Const(Reverse), + DuplicatedNoNeed(b.grad, Hv), + Const(f), + Duplicated(x, v), + Const(y), + ) + return Hv + end - # Sparse Jacobian - 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 + # Sparse Jacobian + 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, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), 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 - - # b.compressed_jacobian is just a vector Jv here - # We don't use the vector mode - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), 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) + vals .= b.nzval + return vals end - vals .= b.nzval - return vals - end - function get_nln_nnzj(b::SparseEnzymeADJacobian, nvar, ncon) - length(b.rowval) - end + 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 + 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 - return rows, cols - 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_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 + 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 - 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 + function NLPModels.jac_coord_residual!( + b::SparseEnzymeADJacobian, + nls::AbstractADNLSModel, + x::AbstractVector, + vals::AbstractVector, + ) + sparse_jac_coord!(nls.F!, b, x, vals) + return vals + end end end From 9865f1a1e2b19fc0e03b24ecd90c9b80dde513fe Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 27 Nov 2024 18:34:46 -0600 Subject: [PATCH 13/44] Fix the tests --- src/enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index c402cb06..4b02b27f 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -151,7 +151,7 @@ end function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) g = similar(x) # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - Enzyme.gradient!(Reverse, g, Const(f), x) + Enzyme.gradient!(Enzyme.Reverse, g, Const(f), x) return g end From 22e5dd7a924f9d1bf8197bca4bc2c0055dd6da2d Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 29 Nov 2024 23:21:53 -0600 Subject: [PATCH 14/44] Fix tests --- src/enzyme.jl | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 4b02b27f..82591806 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -148,19 +148,20 @@ end @init begin @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin + import Enzyme: Const, Reverse, Forward, Duplicated, DuplicatedNoNeed + function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) g = similar(x) - # Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - Enzyme.gradient!(Enzyme.Reverse, g, Const(f), x) + Enzyme.gradient!(Reverse, g, Const(f), x) return g end function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) - Enzyme.autodiff(Enzyme.Reverse, Const(f), Active, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + Enzyme.autodiff(Reverse, Const(f), Active, Duplicated(x, g)) return g end - jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) + jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Reverse, f, x) function hessian(::EnzymeReverseADHessian, f, x) seed = similar(x) @@ -182,7 +183,7 @@ end end function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Reverse, Const(c!), Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) + Enzyme.autodiff(Reverse, Const(c!), Duplicated(b.x, Jtv), Duplicated(x, v)) return Jtv end From 9b63fb6a14ecf428339162f1ca0363925069c815 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 29 Nov 2024 23:45:05 -0600 Subject: [PATCH 15/44] fix tests again --- src/enzyme.jl | 56 +++++++++++++++++++++++++-------------------------- 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 82591806..b03df03c 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -148,20 +148,18 @@ end @init begin @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin - import Enzyme: Const, Reverse, Forward, Duplicated, DuplicatedNoNeed - function ADNLPModels.gradient(::EnzymeReverseADGradient, f, x) g = similar(x) - Enzyme.gradient!(Reverse, g, Const(f), x) + Enzyme.gradient!(Enzyme.Reverse, g, Enzyme.Const(f), x) return g end function ADNLPModels.gradient!(::EnzymeReverseADGradient, g, f, x) - Enzyme.autodiff(Reverse, Const(f), Active, Duplicated(x, g)) + Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(x, g)) return g end - jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Reverse, f, x) + jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) function hessian(::EnzymeReverseADHessian, f, x) seed = similar(x) @@ -170,7 +168,7 @@ end tmp = similar(x) for i in 1:length(x) seed[i] = one(eltype(seed)) - Enzyme.hvp!(tmp, Const(f), x, seed) + Enzyme.hvp!(tmp, Enzyme.Const(f), x, seed) hess[:, i] .= tmp seed[i] = zero(eltype(seed)) end @@ -178,24 +176,24 @@ end end function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val) - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v)) + Enzyme.autodiff(Enzyme.Forward, Enzyme.Const(c!), Enzyme.Duplicated(b.x, Jv), Enzyme.Duplicated(x, v)) return Jv end function Jtprod!(b::EnzymeReverseADJtprod, Jtv, c!, x, v, ::Val) - Enzyme.autodiff(Reverse, Const(c!), Duplicated(b.x, Jtv), Duplicated(x, v)) + Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(c!), Enzyme.Duplicated(b.x, Jtv), Enzyme.Duplicated(x, v)) return Jtv end function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) # What to do with args? Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(f), - Duplicated(x, v), + 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 @@ -211,13 +209,13 @@ end obj_weight::Real = one(eltype(x)), ) Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(ℓ), - Duplicated(x, v), - Const(y), + 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 @@ -233,13 +231,13 @@ end obj_weight::Real = one(eltype(x)), ) Enzyme.autodiff( - Forward, - Const(Enzyme.gradient!), - Const(Reverse), - DuplicatedNoNeed(b.grad, Hv), - Const(f), - Duplicated(x, v), - Const(y), + Enzyme.Forward, + Enzyme.Const(Enzyme.gradient!), + Enzyme.Const(Enzyme.Reverse), + Enzyme.DuplicatedNoNeed(b.grad, Hv), + Enzyme.Const(f), + Enzyme.Duplicated(x, v), + Enzyme.Const(y), ) return Hv end @@ -264,7 +262,7 @@ end # b.compressed_jacobian is just a vector Jv here # We don't use the vector mode - Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), Duplicated(x, b.v)) + Enzyme.autodiff(Enzyme.Forward, Enzyme.Const(c!), Enzyme.Duplicated(b.buffer, 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) From c86d8489c48988012795c24c77a4141a8d459fd9 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 29 Nov 2024 23:56:11 -0600 Subject: [PATCH 16/44] fix tests again --- test/enzyme.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index e0636e3e..f58e2ad3 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -97,15 +97,15 @@ push!( ), ) -include("utils.jl") -include("nlp/basic.jl") -include("nls/basic.jl") -include("nlp/nlpmodelstest.jl") -include("nls/nlpmodelstest.jl") - const test_enzyme = true include("sparse_jacobian.jl") include("sparse_jacobian_nls.jl") # include("sparse_hessian.jl") # include("sparse_hessian_nls.jl") + +include("utils.jl") +include("nlp/basic.jl") +include("nls/basic.jl") +include("nlp/nlpmodelstest.jl") +include("nls/nlpmodelstest.jl") From 659c5b3ccd638601d5aec232d8ef1bfe4cbc72c5 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 1 Dec 2024 18:46:34 -0600 Subject: [PATCH 17/44] Add SparseEnzymeADHessian --- docs/src/sparse.md | 4 +- src/enzyme.jl | 244 ++++++++++++++++++++++++++++++++----- src/sparsity_pattern.jl | 8 +- test/enzyme.jl | 4 +- test/sparse_hessian.jl | 46 ++++--- test/sparse_hessian_nls.jl | 46 ++++--- 6 files changed, 280 insertions(+), 72 deletions(-) diff --git a/docs/src/sparse.md b/docs/src/sparse.md index b8535979..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`, `SparseEnzymeADJacobian`, `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`, `SparseEnzymeADJacobian` 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 b03df03c..65038786 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -145,6 +145,76 @@ function SparseEnzymeADJacobian( ) end +struct SparseEnzymeADHessian{R, C, S} <: ADNLPModels.ADBackend + nvar::Int + rowval::Vector{Int} + colptr::Vector{Int} + nzval::Vector{R} + result_coloring::C + coloring_mode::Symbol + compressed_hessian::S + v::Vector{R} + y::Vector{R} + grad::Vector{R} +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 = similar(x0) + else + coloring_mode = :substitution + group = column_groups(result_coloring) + ncolors = length(group) + compressed_hessian = similar(x0, (nvar, ncolors)) + end + v = similar(x0) + y = similar(x0, ncon) + grad = similar(x0) + + return SparseEnzymeADHessian( + nvar, + rowval, + colptr, + nzval, + result_coloring, + coloring_mode, + compressed_hessian, + v, + y, + ) +end + @init begin @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin @@ -185,19 +255,6 @@ end return Jtv end - function Hvprod!(b::EnzymeReverseADHvprod, Hv, x, v, f, args...) - # What to do with args? - 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 - function Hvprod!( b::EnzymeReverseADHvprod, Hv, @@ -217,7 +274,6 @@ end Enzyme.Duplicated(x, v), Enzyme.Const(y), ) - return Hv end @@ -237,12 +293,30 @@ end Enzyme.DuplicatedNoNeed(b.grad, Hv), Enzyme.Const(f), Enzyme.Duplicated(x, v), - Enzyme.Const(y), ) 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, @@ -271,18 +345,24 @@ end return vals end - function get_nln_nnzj(b::SparseEnzymeADJacobian, nvar, ncon) - length(b.rowval) + 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!( + function NLPModels.jac_structure_residual!( b::SparseEnzymeADJacobian, - nlp::ADModel, + nls::AbstractADNLSModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}, ) rows .= b.rowval - for i = 1:(nlp.meta.nvar) + for i = 1:(nls.meta.nvar) for j = b.colptr[i]:(b.colptr[i + 1] - 1) cols[j] = i end @@ -290,24 +370,29 @@ end return rows, cols end - function NLPModels.jac_coord!( + function NLPModels.jac_coord_residual!( b::SparseEnzymeADJacobian, - nlp::ADModel, + nls::AbstractADNLSModel, x::AbstractVector, vals::AbstractVector, ) - sparse_jac_coord!(nlp.c!, b, x, vals) + sparse_jac_coord!(nls.F!, b, x, vals) return vals end - function NLPModels.jac_structure_residual!( - b::SparseEnzymeADJacobian, - nls::AbstractADNLSModel, + # 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:(nls.meta.nvar) + for i = 1:(nlp.meta.nvar) for j = b.colptr[i]:(b.colptr[i + 1] - 1) cols[j] = i end @@ -315,15 +400,110 @@ end return rows, cols end - function NLPModels.jac_coord_residual!( - b::SparseEnzymeADJacobian, - nls::AbstractADNLSModel, + function sparse_hess_coord!( + b::SparseEnzymeADHessian, x::AbstractVector, + obj_weight, + y::AbstractVector, vals::AbstractVector, ) - sparse_jac_coord!(nls.F!, b, x, vals) + # 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 + + # column icol of the compressed hessian + 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.Duplicated(x, b.v), + Enzyme.Const(y), + ) + + 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, compressed_hessian_icol, icol, b.result_coloring, :L) + 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/sparsity_pattern.jl b/src/sparsity_pattern.jl index cc86f541..95147149 100644 --- a/src/sparsity_pattern.jl +++ b/src/sparsity_pattern.jl @@ -95,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 @@ -124,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 f58e2ad3..df2694bf 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -101,8 +101,8 @@ const test_enzyme = true include("sparse_jacobian.jl") include("sparse_jacobian_nls.jl") -# include("sparse_hessian.jl") -# include("sparse_hessian_nls.jl") +include("sparse_hessian.jl") +include("sparse_hessian_nls.jl") include("utils.jl") include("nlp/basic.jl") diff --git a/test/sparse_hessian.jl b/test/sparse_hessian.jl index 250299f0..3a8204e6 100644 --- a/test/sparse_hessian.jl +++ b/test/sparse_hessian.jl @@ -1,19 +1,31 @@ -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()), -) +if test_enzyme + list_sparse_hess_backend = ( + ( ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), + ), + ( + ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), + ), + ) +else + 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()), + ) +end dt = (Float32, Float64) @@ -63,7 +75,7 @@ dt = (Float32, Float64) 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.SparseADHessian) || (backend == ADNLPModels.SparseReverseADHessian) + if backend != ADNLPModels.ForwardDiffADHessian H_sp = get_sparsity_pattern(nlp, :hessian) @test H_sp == SparseMatrixCSC{Bool, Int}([ 1 0 diff --git a/test/sparse_hessian_nls.jl b/test/sparse_hessian_nls.jl index f0ee5cee..5a81b238 100644 --- a/test/sparse_hessian_nls.jl +++ b/test/sparse_hessian_nls.jl @@ -1,19 +1,31 @@ -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()), -) +if test_enzyme + list_sparse_hess_backend = ( + ( ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), + ), + ( + ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), + ), + ) +else + 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()), + ) +end dt = (Float32, Float64) @@ -50,7 +62,7 @@ dt = (Float32, Float64) H = Symmetric(sparse(rows, cols, vals, nvar, nvar), :L) @test H == [-20*v[2] 0; 0 0] - if (backend == ADNLPModels.SparseADHessian) || (backend == ADNLPModels.SparseReverseADHessian) + if backend != ADNLPModels.ForwardDiffADHessian H_sp = get_sparsity_pattern(nls, :hessian_residual) @test H_sp == SparseMatrixCSC{Bool, Int}([ 1 0 From 2b3e80344d129769ad7271b34ea5ab796198e719 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 1 Dec 2024 18:55:53 -0600 Subject: [PATCH 18/44] Add SparseEnzymeADHessian --- test/enzyme.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/enzyme.jl b/test/enzyme.jl index df2694bf..5a471036 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! From 2bfb90cbe460759e080e312fe849dc8085b96dff Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 1 Dec 2024 19:16:10 -0600 Subject: [PATCH 19/44] Add SparseEnzymeADHessian --- src/enzyme.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/enzyme.jl b/src/enzyme.jl index 65038786..0ecbbe34 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -212,6 +212,7 @@ function SparseEnzymeADHessian( compressed_hessian, v, y, + grad, ) end From 279579531f4ebe105815429ce9e7c8179ac6feff Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 1 Dec 2024 19:54:29 -0600 Subject: [PATCH 20/44] Add SparseEnzymeADHessian --- src/enzyme.jl | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 0ecbbe34..4880ac8a 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -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} @@ -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( @@ -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, @@ -213,6 +217,8 @@ function SparseEnzymeADHessian( v, y, grad, + buffer, + ℓ, ) end @@ -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 From 1900a0d528a1d42d27216450d30fff3011ea41bf Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 1 Dec 2024 19:59:03 -0600 Subject: [PATCH 21/44] Add SparseEnzymeADHessian --- src/enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 4880ac8a..988d57b5 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -439,7 +439,7 @@ end Enzyme.Duplicated(x, b.v), Enzyme.Const(y), Enzyme.Const(obj_weight), - Enzyme.Const(buffer), + Enzyme.Const(b.buffer), ) if b.coloring_mode == :direct From a95f5159ef5b8646333296bc88ef5c24b5777661 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 1 Dec 2024 20:46:39 -0600 Subject: [PATCH 22/44] nlp.f(x) -> f(x) --- src/enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 988d57b5..9d4a45f7 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -204,7 +204,7 @@ function SparseEnzymeADHessian( 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) + ℓ(x, y, obj_weight, buffer) = obj_weight * f(x) + dot(c!(buffer, x), y) return SparseEnzymeADHessian( nvar, From 9ccd2f4a513d38e4a3ccb3ad22c7695923aa1bdf Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Mon, 2 Dec 2024 12:11:00 -0600 Subject: [PATCH 23/44] Fix Enzyme in sparse Hessian --- src/enzyme.jl | 41 +++++++++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 9d4a45f7..79b1581e 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -429,17 +429,38 @@ end compressed_hessian_icol = (b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol) - # AD with Enzyme.jl - Enzyme.autodiff( - Enzyme.Forward, - Enzyme.Const(Enzyme.gradient!), - Enzyme.Const(Enzyme.Reverse), + function _gradient!(dx, f, x, y, obj_weight, buffer) + Enzyme.make_zero!(dx) + res = Enzyme.autodiff( + Enzyme.Reverse, + f, + Enzyme.Active, + Enzyme.Duplicated(x, dx), + Enzyme.Const(y), + Enzyme.Const(obj_weight), + Enzyme.Const(buffer) + ) + return nothing + end + + function _hvp!(res, f, x, v, y, obj_weight, buffer) + # grad = Enzyme.make_zero(x) + Enzyme.autodiff( + Enzyme.Forward, + _gradient!, + res, + Enzyme.Const(f), + Enzyme.Duplicated(x, v), + Enzyme.Const(y), + Enzyme.Const(obj_weight), + Enzyme.Const(buffer), + ) + return nothing + end + + _hvp!( Enzyme.DuplicatedNoNeed(b.grad, compressed_hessian_icol), - Enzyme.Const(b.ℓ), - Enzyme.Duplicated(x, b.v), - Enzyme.Const(y), - Enzyme.Const(obj_weight), - Enzyme.Const(b.buffer), + b.ℓ, x, b.v, y, obj_weight, b.buffer ) if b.coloring_mode == :direct From e1c774703ea50ff46ef0120d20e3aefd9e18bdae Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 2 Dec 2024 13:15:22 -0600 Subject: [PATCH 24/44] Update the sparse Hessian --- src/enzyme.jl | 101 +++++++++++++++++++++++++++----------------------- 1 file changed, 55 insertions(+), 46 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 79b1581e..a3110ad8 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -23,23 +23,29 @@ function EnzymeReverseADJacobian( return EnzymeReverseADJacobian() end -struct EnzymeReverseADHessian <: ADBackend 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..., ) @assert nvar > 0 nnzh = nvar * (nvar + 1) / 2 - return EnzymeReverseADHessian() + + seed = zeros(T, nvar) + Hv = zeros(T, nvar) + return EnzymeReverseADHessian(seed, Hv) end -struct EnzymeReverseADHvprod <: InPlaceADbackend - grad::Vector{Float64} +struct EnzymeReverseADHvprod{T} <: InPlaceADbackend + grad::Vector{T} end function EnzymeReverseADHvprod( @@ -50,12 +56,12 @@ function EnzymeReverseADHvprod( x0::AbstractVector{T} = rand(nvar), kwargs..., ) where {T} - grad = zeros(nvar) + grad = zeros(T, nvar) return EnzymeReverseADHvprod(grad) end -struct EnzymeReverseADJprod <: InPlaceADbackend - x::Vector{Float64} +struct EnzymeReverseADJprod{T} <: InPlaceADbackend + cx::Vector{T} end function EnzymeReverseADJprod( @@ -65,12 +71,12 @@ function EnzymeReverseADJprod( c::Function = (args...) -> []; kwargs..., ) - x = zeros(nvar) - return EnzymeReverseADJprod(x) + cx = zeros(T, nvar) + return EnzymeReverseADJprod(cx) end -struct EnzymeReverseADJtprod <: InPlaceADbackend - x::Vector{Float64} +struct EnzymeReverseADJtprod{T} <: InPlaceADbackend + cx::Vector{T} end function EnzymeReverseADJtprod( @@ -80,8 +86,8 @@ function EnzymeReverseADJtprod( c::Function = (args...) -> []; kwargs..., ) - x = zeros(nvar) - return EnzymeReverseADJtprod(x) + cx = zeros(T, nvar) + return EnzymeReverseADJtprod(cx) end struct SparseEnzymeADJacobian{R, C, S} <: ADBackend @@ -93,7 +99,7 @@ struct SparseEnzymeADJacobian{R, C, S} <: ADBackend result_coloring::C compressed_jacobian::S v::Vector{R} - buffer::Vector{R} + cx::Vector{R} end function SparseEnzymeADJacobian( @@ -130,7 +136,7 @@ function SparseEnzymeADJacobian( nzval = T.(J.nzval) compressed_jacobian = similar(x0, ncon) v = similar(x0) - buffer = zeros(T, ncon) + cx = zeros(T, ncon) SparseEnzymeADJacobian( nvar, @@ -141,7 +147,7 @@ function SparseEnzymeADJacobian( result_coloring, compressed_jacobian, v, - buffer, + cx, ) end @@ -152,11 +158,12 @@ struct SparseEnzymeADHessian{R, C, S, L} <: ADNLPModels.ADBackend 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} - buffer::Vector{R} + cx::Vector{R} ℓ::L end @@ -193,18 +200,20 @@ function SparseEnzymeADHessian( nzval = T.(trilH.nzval) if coloring_algorithm isa GreedyColoringAlgorithm{:direct} coloring_mode = :direct - compressed_hessian = similar(x0) + 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) - buffer = similar(x0, ncon) + cx = similar(x0, ncon) grad = similar(x0) - ℓ(x, y, obj_weight, buffer) = obj_weight * f(x) + dot(c!(buffer, x), y) + ℓ(x, y, obj_weight, cx) = obj_weight * f(x) + dot(c!(cx, x), y) return SparseEnzymeADHessian( nvar, @@ -213,11 +222,12 @@ function SparseEnzymeADHessian( nzval, result_coloring, coloring_mode, + compressed_hessian_icol, compressed_hessian, v, y, grad, - buffer, + cx, ℓ, ) end @@ -238,27 +248,27 @@ end jacobian(::EnzymeReverseADJacobian, f, x) = Enzyme.jacobian(Enzyme.Reverse, f, x) - function hessian(::EnzymeReverseADHessian, f, x) - seed = similar(x) - hess = zeros(eltype(x), length(x), length(x)) - fill!(seed, zero(eltype(x))) - tmp = similar(x) - for i in 1:length(x) - seed[i] = one(eltype(seed)) - Enzyme.hvp!(tmp, Enzyme.Const(f), x, seed) - hess[:, i] .= tmp - seed[i] = zero(eltype(seed)) + 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 + seed[i] = one(T) + Enzyme.hvp!(b.Hv, Enzyme.Const(f), x, b.seed) + view(hess, :, i) .= b.Hv + 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.x, Jv), Enzyme.Duplicated(x, v)) + 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.x, Jtv), Enzyme.Duplicated(x, v)) + Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(c!), Enzyme.Duplicated(b.cx, Jtv), Enzyme.Duplicated(x, v)) return Jtv end @@ -343,7 +353,7 @@ 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.buffer, b.compressed_jacobian), Enzyme.Duplicated(x, b.v)) + 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) @@ -425,11 +435,7 @@ end b.v[col] = 1 end - # column icol of the compressed hessian - compressed_hessian_icol = - (b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol) - - function _gradient!(dx, f, x, y, obj_weight, buffer) + function _gradient!(dx, f, x, y, obj_weight, cx) Enzyme.make_zero!(dx) res = Enzyme.autodiff( Enzyme.Reverse, @@ -438,12 +444,12 @@ end Enzyme.Duplicated(x, dx), Enzyme.Const(y), Enzyme.Const(obj_weight), - Enzyme.Const(buffer) + Enzyme.Const(cx) ) return nothing end - function _hvp!(res, f, x, v, y, obj_weight, buffer) + function _hvp!(res, f, x, v, y, obj_weight, cx) # grad = Enzyme.make_zero(x) Enzyme.autodiff( Enzyme.Forward, @@ -453,19 +459,22 @@ end Enzyme.Duplicated(x, v), Enzyme.Const(y), Enzyme.Const(obj_weight), - Enzyme.Const(buffer), + Enzyme.Const(cx), ) return nothing end _hvp!( - Enzyme.DuplicatedNoNeed(b.grad, compressed_hessian_icol), - b.ℓ, x, b.v, y, obj_weight, b.buffer + 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, compressed_hessian_icol, icol, b.result_coloring, :L) + 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 From b98429b907338b06a8f7d8d180ece75ffa0b16f4 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 2 Dec 2024 13:36:30 -0600 Subject: [PATCH 25/44] Fix the backends for Enzyme.jl --- src/enzyme.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index a3110ad8..a67899de 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -35,7 +35,7 @@ function EnzymeReverseADHessian( c::Function = (args...) -> []; x0::AbstractVector{T} = rand(nvar), kwargs..., -) +) where {T} @assert nvar > 0 nnzh = nvar * (nvar + 1) / 2 @@ -69,8 +69,9 @@ function EnzymeReverseADJprod( f, ncon::Integer = 0, c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), kwargs..., -) +) where {T} cx = zeros(T, nvar) return EnzymeReverseADJprod(cx) end @@ -84,8 +85,9 @@ function EnzymeReverseADJtprod( f, ncon::Integer = 0, c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), kwargs..., -) +) where {T} cx = zeros(T, nvar) return EnzymeReverseADJtprod(cx) end From a81336098eb8fcb66423fe8321085204fddeac35 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 2 Dec 2024 13:40:05 -0600 Subject: [PATCH 26/44] Update test/enzyme.jl --- test/enzyme.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index 5a471036..49de78ff 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -7,13 +7,6 @@ using ADNLPModels: # Automatically loads the code for Enzyme with Requires import Enzyme -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 - EnzymeReverseAD() = ADNLPModels.ADModelBackend( ADNLPModels.EnzymeReverseADGradient(), ADNLPModels.EnzymeReverseADHvprod(zeros(1)), @@ -105,6 +98,13 @@ include("sparse_jacobian_nls.jl") include("sparse_hessian.jl") include("sparse_hessian_nls.jl") +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") From e8a2d0c074706d82acf5727e6af24a808b037edc Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 2 Dec 2024 14:02:19 -0600 Subject: [PATCH 27/44] Update test/enzyme.jl --- test/enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index 49de78ff..33197cb2 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -13,7 +13,7 @@ EnzymeReverseAD() = ADNLPModels.ADModelBackend( ADNLPModels.EnzymeReverseADJprod(zeros(1)), ADNLPModels.EnzymeReverseADJtprod(zeros(1)), ADNLPModels.EnzymeReverseADJacobian(), - ADNLPModels.EnzymeReverseADHessian(), + ADNLPModels.EnzymeReverseADHessian(zeros(1), zeros(1)), ADNLPModels.EnzymeReverseADHvprod(zeros(1)), ADNLPModels.EmptyADbackend(), ADNLPModels.EmptyADbackend(), From d11ce2f26afb29b0cdcda415a29bb5a357465649 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 2 Dec 2024 14:10:21 -0600 Subject: [PATCH 28/44] Update test/enzyme.jl --- src/enzyme.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index a67899de..a5090385 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -256,10 +256,10 @@ end hess = zeros(T, n, n) fill!(b.seed, zero(T)) for i in 1:n - seed[i] = one(T) + b.seed[i] = one(T) Enzyme.hvp!(b.Hv, Enzyme.Const(f), x, b.seed) view(hess, :, i) .= b.Hv - seed[i] = zero(T) + b.seed[i] = zero(T) end return hess end From e75011858697cea7d5654a4167f21f11185955c2 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Tue, 3 Dec 2024 00:00:14 -0600 Subject: [PATCH 29/44] Update the function for the lagrangian --- src/enzyme.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index a5090385..72e90ec6 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -215,7 +215,10 @@ function SparseEnzymeADHessian( y = similar(x0, ncon) cx = similar(x0, ncon) grad = similar(x0) - ℓ(x, y, obj_weight, cx) = obj_weight * f(x) + dot(c!(cx, x), y) + function ℓ(x, y, obj_weight, cx) + c!(cx, x) + obj_weight * f(x) + sum(cx[i] * y[i] for i = 1:ncon) + end return SparseEnzymeADHessian( nvar, From 09b4793d54c22273bb678f5b215d148f3abb95ce Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Tue, 3 Dec 2024 00:05:21 -0600 Subject: [PATCH 30/44] Update the function for the lagrangian --- src/enzyme.jl | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 72e90ec6..5b032c81 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -216,8 +216,12 @@ function SparseEnzymeADHessian( cx = similar(x0, ncon) grad = similar(x0) function ℓ(x, y, obj_weight, cx) - c!(cx, x) - obj_weight * f(x) + sum(cx[i] * y[i] for i = 1:ncon) + 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( From 8ca8044483e8b88952d424631b59d0b7e5ad3eec Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Tue, 3 Dec 2024 00:15:48 -0600 Subject: [PATCH 31/44] Update the function for the lagrangian --- src/enzyme.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 5b032c81..8439834b 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -217,7 +217,7 @@ function SparseEnzymeADHessian( grad = similar(x0) function ℓ(x, y, obj_weight, cx) res = obj_weight * f(x) - if ncon == 0 + if ncon != 0 c!(cx, x) res += sum(cx[i] * y[i] for i = 1:ncon) end From 9f1c2cc32b73f62ec7ac3378cfbd1a17172a7ab2 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Tue, 3 Dec 2024 00:34:03 -0600 Subject: [PATCH 32/44] comment some tests --- src/enzyme.jl | 7 ++++++- test/enzyme.jl | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 8439834b..96f516aa 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -362,7 +362,12 @@ 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)) + 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) diff --git a/test/enzyme.jl b/test/enzyme.jl index 33197cb2..73608bcf 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -95,8 +95,8 @@ const test_enzyme = true include("sparse_jacobian.jl") include("sparse_jacobian_nls.jl") -include("sparse_hessian.jl") -include("sparse_hessian_nls.jl") +# include("sparse_hessian.jl") +# include("sparse_hessian_nls.jl") for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] include("nlp/problems/$(lowercase(problem)).jl") From fed8fa3e1d30527a57578a56d6562a21e1567443 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Tue, 3 Dec 2024 00:58:14 -0600 Subject: [PATCH 33/44] Uncomment some tests --- test/enzyme.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index 73608bcf..33197cb2 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -95,8 +95,8 @@ const test_enzyme = true include("sparse_jacobian.jl") include("sparse_jacobian_nls.jl") -# include("sparse_hessian.jl") -# include("sparse_hessian_nls.jl") +include("sparse_hessian.jl") +include("sparse_hessian_nls.jl") for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] include("nlp/problems/$(lowercase(problem)).jl") From 11ea3eb8ff925b6437fb287d81d6e5dee037c0ee Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 3 Dec 2024 12:43:47 -0600 Subject: [PATCH 34/44] Fix cx in Hessian --- src/enzyme.jl | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 96f516aa..4e2741f4 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -449,31 +449,32 @@ end b.v[col] = 1 end - function _gradient!(dx, f, x, y, obj_weight, cx) + function _gradient!(dx, ℓ, x, y, obj_weight, cx) Enzyme.make_zero!(dx) + dcx = make_zero(cx) res = Enzyme.autodiff( Enzyme.Reverse, - f, + ℓ, Enzyme.Active, Enzyme.Duplicated(x, dx), Enzyme.Const(y), Enzyme.Const(obj_weight), - Enzyme.Const(cx) + Enzyme.Duplicated(cx, dcx) ) return nothing end - function _hvp!(res, f, x, v, y, obj_weight, cx) - # grad = Enzyme.make_zero(x) + function _hvp!(res, ℓ, x, v, y, obj_weight, cx) + dcx = make_zero(cx) Enzyme.autodiff( Enzyme.Forward, _gradient!, res, - Enzyme.Const(f), + Enzyme.Const(ℓ), Enzyme.Duplicated(x, v), Enzyme.Const(y), Enzyme.Const(obj_weight), - Enzyme.Const(cx), + Enzyme.Duplicated(cx, dcx), ) return nothing end From 0fd47a22eed32adf0bbe871b4a602e0b6f026982 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 3 Dec 2024 13:06:25 -0600 Subject: [PATCH 35/44] Fix --- src/enzyme.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 4e2741f4..202e9172 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -451,7 +451,7 @@ end function _gradient!(dx, ℓ, x, y, obj_weight, cx) Enzyme.make_zero!(dx) - dcx = make_zero(cx) + dcx = Enzyme.make_zero(cx) res = Enzyme.autodiff( Enzyme.Reverse, ℓ, @@ -465,7 +465,7 @@ end end function _hvp!(res, ℓ, x, v, y, obj_weight, cx) - dcx = make_zero(cx) + dcx = Enzyme.make_zero(cx) Enzyme.autodiff( Enzyme.Forward, _gradient!, From 5915c81f88417f853ec878c529d8f09dfa24b9bb Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 4 Dec 2024 10:08:57 -0600 Subject: [PATCH 36/44] Fix the functor --- test/enzyme.jl | 2 +- test/nlp/basic.jl | 6 +++--- test/nlp/nlpmodelstest.jl | 6 ++++-- test/nls/nlpmodelstest.jl | 6 ++++-- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index 33197cb2..f8be171d 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -75,7 +75,7 @@ test_autodiff_backend_error() push!( ADNLPModels.predefined_backend, - :enzyme_backend => Dict( + :enzyme => Dict( :gradient_backend => ADNLPModels.EnzymeReverseADGradient, :jprod_backend => ADNLPModels.EnzymeReverseADJprod, :jtprod_backend => ADNLPModels.EnzymeReverseADJtprod, diff --git a/test/nlp/basic.jl b/test/nlp/basic.jl index d09d2912..dd54decd 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) diff --git a/test/nlp/nlpmodelstest.jl b/test/nlp/nlpmodelstest.jl index 78bf56ec..786f52e3 100644 --- a/test/nlp/nlpmodelstest.jl +++ b/test/nlp/nlpmodelstest.jl @@ -18,8 +18,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/nlpmodelstest.jl b/test/nls/nlpmodelstest.jl index e1681f49..92afbcc0 100644 --- a/test/nls/nlpmodelstest.jl +++ b/test/nls/nlpmodelstest.jl @@ -35,8 +35,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 From 0355f5fc9fcac57b2881528b5ed51fb396ba336d Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 20 Dec 2024 17:14:31 +0100 Subject: [PATCH 37/44] Uncomment some code --- test/Project.toml | 2 ++ test/nlp/basic.jl | 14 +++++++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index e6ae782e..502b8f56 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,6 @@ [deps] +ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a" +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/nlp/basic.jl b/test/nlp/basic.jl index dd54decd..2b087ad2 100644 --- a/test/nlp/basic.jl +++ b/test/nlp/basic.jl @@ -17,13 +17,13 @@ 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) - # y = 2x .+ 3 + randn(100) * 0.1 - # regr = LinearRegression(x, y) - # nlp = ADNLPModel(regr, ones(2); kwargs...) - # β = [ones(100) x] \ y - # @test abs(obj(nlp, β) - norm(y .- β[1] - β[2] * x)^2 / 2) < 1e-12 - # @test norm(grad(nlp, β)) < 1e-12 + x = range(-1, stop = 1, length = 100) + y = 2x .+ 3 + randn(100) * 0.1 + regr = LinearRegression(x, y) + nlp = ADNLPModel(regr, ones(2); kwargs...) + β = [ones(100) x] \ y + @test abs(obj(nlp, β) - norm(y .- β[1] - β[2] * x)^2 / 2) < 1e-12 + @test norm(grad(nlp, β)) < 1e-12 test_getter_setter(nlp) From f870f70d8e6787b14a1e2796ced042fd4240e43e Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 20 Dec 2024 17:17:24 +0100 Subject: [PATCH 38/44] fix test/Project.toml --- test/Project.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 502b8f56..e6ae782e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,4 @@ [deps] -ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617" From 921ad6c51a0551c1a10325506b6c7cccde1fa76e Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 21 Dec 2024 20:52:56 +0100 Subject: [PATCH 39/44] Update the tests for Enzyme and Zygote --- src/enzyme.jl | 2 +- src/forward.jl | 8 ++-- src/predefined_backend.jl | 42 +++++++++++++++++++-- src/reverse.jl | 8 ++-- src/sparse_hessian.jl | 4 +- test/enzyme.jl | 59 +++++++++++++++++------------ test/nlp/basic.jl | 4 -- test/nlp/nlpmodelstest.jl | 6 +-- test/nls/basic.jl | 4 -- test/nls/nlpmodelstest.jl | 6 +-- test/runtests.jl | 50 ++++++++++++++++++++++--- test/script_OP.jl | 74 +++++++++++++++++++------------------ test/sparse_hessian.jl | 33 +---------------- test/sparse_hessian_nls.jl | 30 +-------------- test/sparse_jacobian.jl | 11 +----- test/sparse_jacobian_nls.jl | 11 +----- test/zygote.jl | 35 ++++++++---------- 17 files changed, 193 insertions(+), 194 deletions(-) diff --git a/src/enzyme.jl b/src/enzyme.jl index 202e9172..3fea7c4f 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -153,7 +153,7 @@ function SparseEnzymeADJacobian( ) end -struct SparseEnzymeADHessian{R, C, S, L} <: ADNLPModels.ADBackend +struct SparseEnzymeADHessian{R, C, S, L} <: ADBackend nvar::Int rowval::Vector{Int} colptr::Vector{Int} 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..1595d054 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 => EnzymeReverseADJacobian, + :hessian_backend => EnzymeReverseADHessian, + :ghjvprod_backend => ForwardDiffADGHjvprod, + :jprod_residual_backend => EnzymeReverseADJprod, + :jtprod_residual_backend => EnzymeReverseADJtprod, + :hprod_residual_backend => EnzymeReverseADHvprod, + :jacobian_residual_backend => EnzymeReverseADJacobian, + :hessian_residual_backend => EnzymeReverseADHessian, +) + +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/test/enzyme.jl b/test/enzyme.jl index f8be171d..1ed36669 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -21,10 +21,12 @@ EnzymeReverseAD() = ADNLPModels.ADModelBackend( 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)() @@ -73,30 +75,25 @@ end test_autodiff_backend_error() -push!( - ADNLPModels.predefined_backend, - :enzyme => Dict( - :gradient_backend => ADNLPModels.EnzymeReverseADGradient, - :jprod_backend => ADNLPModels.EnzymeReverseADJprod, - :jtprod_backend => ADNLPModels.EnzymeReverseADJtprod, - :hprod_backend => ADNLPModels.EnzymeReverseADHvprod, - :jacobian_backend => ADNLPModels.EnzymeReverseADJacobian, - :hessian_backend => ADNLPModels.EnzymeReverseADHessian, - :ghjvprod_backend => ADNLPModels.ForwardDiffADGHjvprod, - :jprod_residual_backend => ADNLPModels.EnzymeReverseADJprod, - :jtprod_residual_backend => ADNLPModels.EnzymeReverseADJtprod, - :hprod_residual_backend => ADNLPModels.EnzymeReverseADHvprod, - :jacobian_residual_backend => ADNLPModels.EnzymeReverseADJacobian, - :hessian_residual_backend => ADNLPModels.EnzymeReverseADHessian, - ), -) - -const test_enzyme = true +@testset "Sparse Jacobian" begin + list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) + include("sparse_jacobian.jl") + include("sparse_jacobian_nls.jl") +end -include("sparse_jacobian.jl") -include("sparse_jacobian_nls.jl") -include("sparse_hessian.jl") -include("sparse_hessian_nls.jl") +@testset "Sparse Hessian" begin + list_sparse_hess_backend = ( + ( ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), + ), + ( + ADNLPModels.SparseEnzymeADHessian, + Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), + ), + ) + include("sparse_hessian.jl") + include("sparse_hessian_nls.jl") +end for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] include("nlp/problems/$(lowercase(problem)).jl") @@ -110,3 +107,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 (: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 diff --git a/test/nlp/basic.jl b/test/nlp/basic.jl index 2b087ad2..119d1734 100644 --- a/test/nlp/basic.jl +++ b/test/nlp/basic.jl @@ -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 786f52e3..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))() 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 92afbcc0..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))() diff --git a/test/runtests.jl b/test/runtests.jl index 94d14ffa..9704c9e9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,8 +4,6 @@ using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! -const test_enzyme = false - @testset "Test sparsity pattern of Jacobian and Hessian" begin f(x) = sum(x .^ 2) c(x) = x @@ -25,12 +23,30 @@ end include("manual.jl") end -@testset "Basic Jacobian derivative test" begin +@testset "Sparse Jacobian" begin + list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), + (ADNLPModels.ForwardDiffADJacobian, Dict())) include("sparse_jacobian.jl") include("sparse_jacobian_nls.jl") end -@testset "Basic Hessian derivative test" begin +@testset "Sparse Hessian" begin + 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()), + ) include("sparse_hessian.jl") include("sparse_hessian_nls.jl") end @@ -44,6 +60,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 3a8204e6..802553e1 100644 --- a/test/sparse_hessian.jl +++ b/test/sparse_hessian.jl @@ -1,35 +1,4 @@ -if test_enzyme - list_sparse_hess_backend = ( - ( ADNLPModels.SparseEnzymeADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), - ), - ( - ADNLPModels.SparseEnzymeADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - ) -else - 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()), - ) -end - -dt = (Float32, Float64) - -@testset "Basic Hessian derivative with backend=$(backend) and T=$(T)" for T in dt, +@testset "Basic Hessian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), (backend, kw) in list_sparse_hess_backend c!(cx, x) = begin diff --git a/test/sparse_hessian_nls.jl b/test/sparse_hessian_nls.jl index 5a81b238..066cf949 100644 --- a/test/sparse_hessian_nls.jl +++ b/test/sparse_hessian_nls.jl @@ -1,35 +1,7 @@ if test_enzyme list_sparse_hess_backend = ( ( ADNLPModels.SparseEnzymeADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), - ), - ( - ADNLPModels.SparseEnzymeADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - ) -else - 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()), - ) -end - -dt = (Float32, Float64) - -@testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in dt, +@testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), (backend, kw) in list_sparse_hess_backend F!(Fx, x) = begin diff --git a/test/sparse_jacobian.jl b/test/sparse_jacobian.jl index b8a22259..dd24e5c8 100644 --- a/test/sparse_jacobian.jl +++ b/test/sparse_jacobian.jl @@ -1,13 +1,4 @@ -if test_enzyme - list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) -else - list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), - (ADNLPModels.ForwardDiffADJacobian, Dict())) -end - -dt = (Float32, Float64) - -@testset "Basic Jacobian derivative with backend=$(backend) and T=$(T)" for T in dt, +@testset "Basic Jacobian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), (backend, kw) in list_sparse_jac_backend c!(cx, x) = begin diff --git a/test/sparse_jacobian_nls.jl b/test/sparse_jacobian_nls.jl index cecb05f9..f8249b86 100644 --- a/test/sparse_jacobian_nls.jl +++ b/test/sparse_jacobian_nls.jl @@ -1,13 +1,4 @@ -if test_enzyme - list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) -else - list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), - (ADNLPModels.ForwardDiffADJacobian, Dict())) -end - -dt = (Float32, Float64) - -@testset "Basic Jacobian of residual derivative with backend=$(backend) and T=$(T)" for T in dt, +@testset "Basic Jacobian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), (backend, kw) in list_sparse_jac_backend F!(Fx, x) = begin 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 From d62b30888349dc04a97e5684fb3a75750bad0b9e Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 21 Dec 2024 21:15:25 +0100 Subject: [PATCH 40/44] Update the documentation --- docs/src/backend.md | 10 +++++----- docs/src/predefined.md | 14 ++++++++++++++ src/predefined_backend.jl | 8 ++++---- 3 files changed, 23 insertions(+), 9 deletions(-) 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..6068d518 100644 --- a/docs/src/predefined.md +++ b/docs/src/predefined.md @@ -46,3 +46,17 @@ 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). + +```@example ex1 +nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon, backend = :zygote) +get_adbackend(nlp) +``` diff --git a/src/predefined_backend.jl b/src/predefined_backend.jl index 1595d054..c612504b 100644 --- a/src/predefined_backend.jl +++ b/src/predefined_backend.jl @@ -48,14 +48,14 @@ enzyme_backend = Dict( :jprod_backend => EnzymeReverseADJprod, :jtprod_backend => EnzymeReverseADJtprod, :hprod_backend => EnzymeReverseADHvprod, - :jacobian_backend => EnzymeReverseADJacobian, - :hessian_backend => EnzymeReverseADHessian, + :jacobian_backend => SparseEnzymeADJacobian, + :hessian_backend => SparseEnzymeADHessian, :ghjvprod_backend => ForwardDiffADGHjvprod, :jprod_residual_backend => EnzymeReverseADJprod, :jtprod_residual_backend => EnzymeReverseADJtprod, :hprod_residual_backend => EnzymeReverseADHvprod, - :jacobian_residual_backend => EnzymeReverseADJacobian, - :hessian_residual_backend => EnzymeReverseADHessian, + :jacobian_residual_backend => SparseEnzymeADJacobian, + :hessian_residual_backend => SparseEnzymeADHessian, ) zygote_backend = Dict( From b4e88ab333927571cbfcc124a7ab2668e43526da Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 21 Dec 2024 21:40:28 +0100 Subject: [PATCH 41/44] Fix tests --- test/enzyme.jl | 39 +++++---- test/runtests.jl | 55 ++++++++----- test/sparse_hessian.jl | 158 ++++++++++++++++++------------------ test/sparse_hessian_nls.jl | 81 +++++++++--------- test/sparse_jacobian.jl | 108 ++++++++++++------------ test/sparse_jacobian_nls.jl | 94 ++++++++++----------- 6 files changed, 277 insertions(+), 258 deletions(-) diff --git a/test/enzyme.jl b/test/enzyme.jl index 1ed36669..f9cba09c 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -75,24 +75,35 @@ 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 - list_sparse_jac_backend = ((ADNLPModels.SparseEnzymeADJacobian, Dict()),) - include("sparse_jacobian.jl") - include("sparse_jacobian_nls.jl") + 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 - list_sparse_hess_backend = ( - ( ADNLPModels.SparseEnzymeADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:direct}()), - ), - ( - ADNLPModels.SparseEnzymeADHessian, - Dict(:coloring_algorithm => GreedyColoringAlgorithm{:substitution}()), - ), - ) - include("sparse_hessian.jl") - include("sparse_hessian_nls.jl") + 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"] diff --git a/test/runtests.jl b/test/runtests.jl index 9704c9e9..da4bf09a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,32 +23,43 @@ end include("manual.jl") end +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 - list_sparse_jac_backend = ((ADNLPModels.SparseADJacobian, Dict()), - (ADNLPModels.ForwardDiffADJacobian, Dict())) - include("sparse_jacobian.jl") - include("sparse_jacobian_nls.jl") + for (backend, kw) in list_sparse_jac_backend + sparse_jacobian(backend, kw) + sparse_jacobian_nls(backend, kw) + end end +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 - 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()), - ) - include("sparse_hessian.jl") - include("sparse_hessian_nls.jl") + 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"] diff --git a/test/sparse_hessian.jl b/test/sparse_hessian.jl index 802553e1..77c0bf29 100644 --- a/test/sparse_hessian.jl +++ b/test/sparse_hessian.jl @@ -1,87 +1,87 @@ -@testset "Basic Hessian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), - (backend, kw) in list_sparse_hess_backend +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..., + ) - 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]] + 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] + # 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.ForwardDiffADHessian + 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 066cf949..6a02b6ae 100644 --- a/test/sparse_hessian_nls.jl +++ b/test/sparse_hessian_nls.jl @@ -1,48 +1,45 @@ -if test_enzyme - list_sparse_hess_backend = ( - ( ADNLPModels.SparseEnzymeADHessian, -@testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), - (backend, kw) in list_sparse_hess_backend +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...) - 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] - 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] + # 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] - # 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.ForwardDiffADHessian + H_sp = get_sparsity_pattern(nls, :hessian_residual) + @test H_sp == SparseMatrixCSC{Bool, Int}([ + 1 0 + 0 0 + ]) + end - if backend != ADNLPModels.ForwardDiffADHessian - 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 dd24e5c8..4c0801eb 100644 --- a/test/sparse_jacobian.jl +++ b/test/sparse_jacobian.jl @@ -1,60 +1,60 @@ -@testset "Basic Jacobian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), - (backend, kw) in list_sparse_jac_backend +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..., + ) - 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.ForwardDiffADJacobian - 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 f8249b86..fc6f8e98 100644 --- a/test/sparse_jacobian_nls.jl +++ b/test/sparse_jacobian_nls.jl @@ -1,53 +1,53 @@ -@testset "Basic Jacobian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), - (backend, kw) in list_sparse_jac_backend +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...) - 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.ForwardDiffADJacobian - 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 From 7eb8c160aae4dcdfef3f19a4532fd7f45bd35043 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 21 Dec 2024 21:48:41 +0100 Subject: [PATCH 42/44] Update nlp/basic.jl --- test/nlp/basic.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/nlp/basic.jl b/test/nlp/basic.jl index 119d1734..07c4d940 100644 --- a/test/nlp/basic.jl +++ b/test/nlp/basic.jl @@ -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...) From 3f1ff12b00d726b2353d068dce261583542fea5a Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 21 Dec 2024 22:16:59 +0100 Subject: [PATCH 43/44] FIx a typo --- test/sparse_hessian.jl | 2 +- test/sparse_hessian_nls.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/sparse_hessian.jl b/test/sparse_hessian.jl index 77c0bf29..ca130989 100644 --- a/test/sparse_hessian.jl +++ b/test/sparse_hessian.jl @@ -1,5 +1,5 @@ function sparse_hessian(backend, kw) - @testset "Basic Hessian derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), + @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) diff --git a/test/sparse_hessian_nls.jl b/test/sparse_hessian_nls.jl index 6a02b6ae..80ac269b 100644 --- a/test/sparse_hessian_nls.jl +++ b/test/sparse_hessian_nls.jl @@ -1,5 +1,5 @@ function sparse_hessian_nls(backend, kw) - @testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in (Float32, Float64), + @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) From ed926edf38324fa64990952cf74fe7072156efe6 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 21 Dec 2024 22:47:53 +0100 Subject: [PATCH 44/44] Fix documentation --- docs/src/predefined.md | 5 ----- 1 file changed, 5 deletions(-) diff --git a/docs/src/predefined.md b/docs/src/predefined.md index 6068d518..12e771bf 100644 --- a/docs/src/predefined.md +++ b/docs/src/predefined.md @@ -55,8 +55,3 @@ get_adbackend(nlp) ``` The backend `:zygote` focuses on backend based on [Zygote.jl](https://github.com/FluxML/Zygote.jl). - -```@example ex1 -nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon, backend = :zygote) -get_adbackend(nlp) -```