From 7c2c67655d42ae5917e705fd41e82f3c20e75ca3 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 31 Jan 2025 20:15:43 -0600 Subject: [PATCH] Use show_time for sparse AD backends --- src/ad.jl | 20 ++--- src/enzyme.jl | 108 ++++++++++++++--------- src/sparse_hessian.jl | 191 +++++++++++++++++++++++------------------ src/sparse_jacobian.jl | 39 ++++++--- 4 files changed, 209 insertions(+), 149 deletions(-) diff --git a/src/ad.jl b/src/ad.jl index 306e430e..b950898f 100644 --- a/src/ad.jl +++ b/src/ad.jl @@ -87,7 +87,7 @@ function ADModelBackend( hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend} hessian_backend else - HB(nvar, f, ncon, c!; kwargs...) + HB(nvar, f, ncon, c!; show_time, kwargs...) end end show_time && println("hessian backend $HB: $b seconds;") @@ -170,7 +170,7 @@ function ADModelBackend( jacobian_backend = if jacobian_backend isa Union{AbstractNLPModel, ADBackend} jacobian_backend else - JB(nvar, f, ncon, c!; kwargs...) + JB(nvar, f, ncon, c!; show_time, kwargs...) end end show_time && println("jacobian backend $JB: $b seconds;") @@ -180,7 +180,7 @@ function ADModelBackend( hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend} hessian_backend else - HB(nvar, f, ncon, c!; kwargs...) + HB(nvar, f, ncon, c!; show_time, kwargs...) end end show_time && println("hessian backend $HB: $b seconds;") @@ -263,7 +263,7 @@ function ADModelNLSBackend( hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend} hessian_backend else - HB(nvar, f, ncon, c!; kwargs...) + HB(nvar, f, ncon, c!; show_time, kwargs...) end end show_time && println("hessian backend $HB: $b seconds;") @@ -304,7 +304,7 @@ function ADModelNLSBackend( if jacobian_residual_backend isa Union{AbstractNLPModel, ADBackend} jacobian_residual_backend else - JBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...) + JBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...) end end show_time && println("jacobian_residual backend $JBLS: $b seconds;") @@ -314,7 +314,7 @@ function ADModelNLSBackend( hessian_residual_backend = if hessian_residual_backend isa Union{AbstractNLPModel, ADBackend} hessian_residual_backend else - HBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...) + HBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...) end end show_time && println("hessian_residual backend $HBLS: $b seconds. \n") @@ -410,7 +410,7 @@ function ADModelNLSBackend( jacobian_backend = if jacobian_backend isa Union{AbstractNLPModel, ADBackend} jacobian_backend else - JB(nvar, f, ncon, c!; kwargs...) + JB(nvar, f, ncon, c!; show_time, kwargs...) end end show_time && println("jacobian backend $JB: $b seconds;") @@ -420,7 +420,7 @@ function ADModelNLSBackend( hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend} hessian_backend else - HB(nvar, f, ncon, c!; kwargs...) + HB(nvar, f, ncon, c!; show_time, kwargs...) end end show_time && println("hessian backend $HB: $b seconds;") @@ -471,7 +471,7 @@ function ADModelNLSBackend( if jacobian_residual_backend isa Union{AbstractNLPModel, ADBackend} jacobian_residual_backend else - JBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...) + JBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...) end end show_time && println("jacobian_residual backend $JBLS: $b seconds;") @@ -481,7 +481,7 @@ function ADModelNLSBackend( hessian_residual_backend = if hessian_residual_backend isa Union{AbstractNLPModel, ADBackend} hessian_residual_backend else - HBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...) + HBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...) end end show_time && println("hessian_residual backend $HBLS: $b seconds. \n") diff --git a/src/enzyme.jl b/src/enzyme.jl index d3182705..dea0f8de 100644 --- a/src/enzyme.jl +++ b/src/enzyme.jl @@ -112,11 +112,15 @@ function SparseEnzymeADJacobian( x0::AbstractVector = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), detector::AbstractSparsityDetector = TracerSparsityDetector(), + show_time::Bool = false, kwargs..., ) - output = similar(x0, ncon) - J = compute_jacobian_sparsity(c!, output, x0, detector = detector) - SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...) + timer = @elapsed begin + output = similar(x0, ncon) + J = compute_jacobian_sparsity(c!, output, x0, detector = detector) + end + show_time && println(" • Sparsity pattern detection of the Jacobian: $timer seconds.") + SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, show_time, kwargs...) end function SparseEnzymeADJacobian( @@ -127,18 +131,26 @@ function SparseEnzymeADJacobian( J::SparseMatrixCSC{Bool, Int}; x0::AbstractVector{T} = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), + show_time::Bool = false, 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) + timer = @elapsed begin + # 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) + end + show_time && println(" • Coloring of the sparse Jacobian: $timer seconds.") - rowval = J.rowval - colptr = J.colptr - nzval = T.(J.nzval) - compressed_jacobian = similar(x0, ncon) - v = similar(x0) - cx = zeros(T, ncon) + timer = @elapsed begin + v = similar(x0) + cx = zeros(T, ncon) + end + show_time && println(" • Allocation of the AD buffers for the sparse Jacobian: $timer seconds.") SparseEnzymeADJacobian( nvar, @@ -177,10 +189,14 @@ function SparseEnzymeADHessian( x0::AbstractVector = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(), detector::AbstractSparsityDetector = TracerSparsityDetector(), + show_time::Bool = false, kwargs..., ) - H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) - SparseEnzymeADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, kwargs...) + timer = @elapsed begin + H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) + end + show_time && println(" • Sparsity pattern detection of the Hessian: $timer seconds.") + SparseEnzymeADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, show_time, kwargs...) end function SparseEnzymeADHessian( @@ -191,38 +207,48 @@ function SparseEnzymeADHessian( H::SparseMatrixCSC{Bool, Int}; x0::AbstractVector{T} = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(), + show_time::Bool = false, kwargs..., ) where {T} - problem = ColoringProblem{:symmetric, :column}() - result_coloring = coloring(H, problem, coloring_algorithm, decompression_eltype = T) - - trilH = tril(H) - rowval = trilH.rowval - colptr = trilH.colptr - nzval = T.(trilH.nzval) - if coloring_algorithm isa GreedyColoringAlgorithm{:direct} - coloring_mode = :direct - compressed_hessian_icol = similar(x0) - compressed_hessian = compressed_hessian_icol - else - coloring_mode = :substitution - group = column_groups(result_coloring) - ncolors = length(group) - compressed_hessian_icol = similar(x0) - compressed_hessian = similar(x0, (nvar, ncolors)) + timer = @elapsed begin + problem = ColoringProblem{:symmetric, :column}() + result_coloring = coloring(H, problem, coloring_algorithm, decompression_eltype = T) + + trilH = tril(H) + rowval = trilH.rowval + colptr = trilH.colptr + nzval = T.(trilH.nzval) + if coloring_algorithm isa GreedyColoringAlgorithm{:direct} + coloring_mode = :direct + compressed_hessian_icol = similar(x0) + compressed_hessian = compressed_hessian_icol + else + coloring_mode = :substitution + group = column_groups(result_coloring) + ncolors = length(group) + compressed_hessian_icol = similar(x0) + compressed_hessian = similar(x0, (nvar, ncolors)) + end end - v = similar(x0) - y = similar(x0, ncon) - cx = similar(x0, ncon) - grad = similar(x0) - function ℓ(x, y, obj_weight, cx) - res = obj_weight * f(x) - if ncon != 0 - c!(cx, x) - res += sum(cx[i] * y[i] for i = 1:ncon) + show_time && println(" • Coloring of the sparse Hessian: $timer seconds.") + + timer = @elapsed begin + v = similar(x0) + y = similar(x0, ncon) + cx = similar(x0, ncon) + grad = similar(x0) + + function ℓ(x, y, obj_weight, cx) + res = obj_weight * f(x) + if ncon != 0 + c!(cx, x) + res += sum(cx[i] * y[i] for i = 1:ncon) + end + return res end - return res end + show_time && println(" • Allocation of the AD buffers for the sparse Hessian: $timer seconds.") + return SparseEnzymeADHessian( nvar, diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl index cbfb7e78..f226ae0d 100644 --- a/src/sparse_hessian.jl +++ b/src/sparse_hessian.jl @@ -24,10 +24,14 @@ function SparseADHessian( x0::AbstractVector = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), detector::AbstractSparsityDetector = TracerSparsityDetector(), + show_time::Bool = false, kwargs..., ) - H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) - SparseADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, kwargs...) + timer = @elapsed begin + H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) + end + show_time && println(" • Sparsity pattern detection of the Hessian: $timer seconds.") + SparseADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, show_time, kwargs...) end function SparseADHessian( @@ -38,52 +42,60 @@ function SparseADHessian( H::SparseMatrixCSC{Bool, Int64}; x0::S = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), + show_time::Bool = false, kwargs..., ) where {S} T = eltype(S) - 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 - seed = BitVector(undef, nvar) - - function lag(z; nvar = nvar, ncon = ncon, f = f, c! = c!) - cx, x, y, ob = view(z, 1:ncon), - view(z, (ncon + 1):(nvar + ncon)), - view(z, (nvar + ncon + 1):(nvar + ncon + ncon)), - z[end] - if ncon > 0 - c!(cx, x) - return ob * f(x) + dot(cx, y) + timer = @elapsed begin + 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 - return ob * f(x) + coloring_mode = :substitution + group = column_groups(result_coloring) + ncolors = length(group) + compressed_hessian = similar(x0, (nvar, ncolors)) end + seed = BitVector(undef, nvar) end - ntotal = nvar + 2 * ncon + 1 - sol = similar(x0, ntotal) - lz = Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(lag), T}, T, 1}}(undef, ntotal) - glz = similar(lz) - cfg = ForwardDiff.GradientConfig(lag, lz) - function ∇φ!(gz, z; lag = lag, cfg = cfg) - ForwardDiff.gradient!(gz, lag, z, cfg) - return gz + show_time && println(" • Coloring of the sparse Hessian: $timer seconds.") + + timer = @elapsed begin + function lag(z; nvar = nvar, ncon = ncon, f = f, c! = c!) + cx, x, y, ob = view(z, 1:ncon), + view(z, (ncon + 1):(nvar + ncon)), + view(z, (nvar + ncon + 1):(nvar + ncon + ncon)), + z[end] + if ncon > 0 + c!(cx, x) + return ob * f(x) + dot(cx, y) + else + return ob * f(x) + end + end + + ntotal = nvar + 2 * ncon + 1 + sol = similar(x0, ntotal) + lz = Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(lag), T}, T, 1}}(undef, ntotal) + glz = similar(lz) + cfg = ForwardDiff.GradientConfig(lag, lz) + function ∇φ!(gz, z; lag = lag, cfg = cfg) + ForwardDiff.gradient!(gz, lag, z, cfg) + return gz + end + longv = fill!(S(undef, ntotal), 0) + Hvp = fill!(S(undef, ntotal), 0) + y = fill!(S(undef, ncon), 0) end - longv = fill!(S(undef, ntotal), 0) - Hvp = fill!(S(undef, ntotal), 0) - y = fill!(S(undef, ncon), 0) + show_time && println(" • Allocation of the AD buffers for the sparse Hessian: $timer seconds.") return SparseADHessian( nvar, @@ -133,10 +145,14 @@ function SparseReverseADHessian( x0::AbstractVector = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(), detector::AbstractSparsityDetector = TracerSparsityDetector(), + show_time::Bool = false, kwargs..., ) - H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) - SparseReverseADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, kwargs...) + timer = @elapsed begin + H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector) + end + show_time && println(" • Sparsity pattern detection of the Hessian: $timer seconds.") + SparseReverseADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, show_time, kwargs...) end function SparseReverseADHessian( @@ -147,55 +163,62 @@ function SparseReverseADHessian( H::SparseMatrixCSC{Bool, Int}; x0::AbstractVector{T} = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(), + show_time::Bool = false, 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)) + timer = @elapsed begin + 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 + seed = BitVector(undef, nvar) end - seed = BitVector(undef, nvar) + show_time && println(" • Coloring of the sparse Hessian: $timer seconds.") # unconstrained Hessian - tagf = ForwardDiff.Tag{typeof(f), T} - z = Vector{ForwardDiff.Dual{tagf, T, 1}}(undef, nvar) - gz = similar(z) - f_tape = ReverseDiff.GradientTape(f, z) - cfgf = ReverseDiff.compile(f_tape) - ∇f!(gz, z; cfg = cfgf) = ReverseDiff.gradient!(gz, cfg, z) - - # constraints - ψ(x, u) = begin # ; tmp_out = _tmp_out - ncon = length(u) - tmp_out = similar(x, ncon) - c!(tmp_out, x) - dot(tmp_out, u) - end - tagψ = ForwardDiff.Tag{typeof(ψ), T} - zψ = Vector{ForwardDiff.Dual{tagψ, T, 1}}(undef, nvar) - yψ = fill!(similar(zψ, ncon), zero(T)) - ψ_tape = ReverseDiff.GradientConfig((zψ, yψ)) - cfgψ = ReverseDiff.compile(ReverseDiff.GradientTape(ψ, (zψ, yψ), ψ_tape)) - - gzψ = similar(zψ) - gyψ = similar(yψ) - function ∇l!(gz, gy, z, y; cfg = cfgψ) - ReverseDiff.gradient!((gz, gy), cfg, (z, y)) + timer = @elapsed begin + tagf = ForwardDiff.Tag{typeof(f), T} + z = Vector{ForwardDiff.Dual{tagf, T, 1}}(undef, nvar) + gz = similar(z) + f_tape = ReverseDiff.GradientTape(f, z) + cfgf = ReverseDiff.compile(f_tape) + ∇f!(gz, z; cfg = cfgf) = ReverseDiff.gradient!(gz, cfg, z) + + # constraints + ψ(x, u) = begin # ; tmp_out = _tmp_out + ncon = length(u) + tmp_out = similar(x, ncon) + c!(tmp_out, x) + dot(tmp_out, u) + end + tagψ = ForwardDiff.Tag{typeof(ψ), T} + zψ = Vector{ForwardDiff.Dual{tagψ, T, 1}}(undef, nvar) + yψ = fill!(similar(zψ, ncon), zero(T)) + ψ_tape = ReverseDiff.GradientConfig((zψ, yψ)) + cfgψ = ReverseDiff.compile(ReverseDiff.GradientTape(ψ, (zψ, yψ), ψ_tape)) + + gzψ = similar(zψ) + gyψ = similar(yψ) + function ∇l!(gz, gy, z, y; cfg = cfgψ) + ReverseDiff.gradient!((gz, gy), cfg, (z, y)) + end + Hv_temp = similar(x0) + y = similar(x0, ncon) end - Hv_temp = similar(x0) + show_time && println(" • Allocation of the AD buffers for the sparse Hessian: $timer seconds.") - y = similar(x0, ncon) return SparseReverseADHessian( nvar, rowval, diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl index db027de0..51c2e14c 100644 --- a/src/sparse_jacobian.jl +++ b/src/sparse_jacobian.jl @@ -19,11 +19,15 @@ function SparseADJacobian( x0::AbstractVector = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), detector::AbstractSparsityDetector = TracerSparsityDetector(), + show_time::Bool = false, kwargs..., ) - output = similar(x0, ncon) - J = compute_jacobian_sparsity(c!, output, x0, detector = detector) - SparseADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...) + timer = @elapsed begin + output = similar(x0, ncon) + J = compute_jacobian_sparsity(c!, output, x0, detector = detector) + end + show_time && println(" • Sparsity pattern detection of the Jacobian: $timer seconds.") + SparseADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, show_time, kwargs...) end function SparseADJacobian( @@ -34,21 +38,28 @@ function SparseADJacobian( J::SparseMatrixCSC{Bool, Int}; x0::AbstractVector{T} = rand(nvar), coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(), + show_time::Bool = false, 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) + timer = @elapsed begin + # 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) - seed = BitVector(undef, nvar) + rowval = J.rowval + colptr = J.colptr + nzval = T.(J.nzval) + compressed_jacobian = similar(x0, ncon) + seed = BitVector(undef, nvar) + end + show_time && println(" • Coloring of the sparse Jacobian: $timer seconds.") - tag = ForwardDiff.Tag{typeof(c!), T} - z = Vector{ForwardDiff.Dual{tag, T, 1}}(undef, nvar) - cz = similar(z, ncon) + timer = @elapsed begin + tag = ForwardDiff.Tag{typeof(c!), T} + z = Vector{ForwardDiff.Dual{tag, T, 1}}(undef, nvar) + cz = similar(z, ncon) + end + show_time && println(" • Allocation of the AD buffers for the sparse Jacobian: $timer seconds.") SparseADJacobian( nvar,