From adceb0f8348afa1029faa4a39fc1ae1805785b76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Mon, 21 Apr 2025 20:22:47 +0200 Subject: [PATCH 1/6] =?UTF-8?q?[Nonlinear]=20Merge=20forward=5Fstorage=5F?= =?UTF-8?q?=CF=B5=20with=20reverse=5Fstorage=5F=CF=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ReverseAD/forward_over_reverse.jl | 4 +-- .../ReverseAD/mathoptinterface_api.jl | 28 +++++++++---------- src/Nonlinear/ReverseAD/types.jl | 9 ++---- 3 files changed, 18 insertions(+), 23 deletions(-) diff --git a/src/Nonlinear/ReverseAD/forward_over_reverse.jl b/src/Nonlinear/ReverseAD/forward_over_reverse.jl index bcc1f86a98..e469b6ba0d 100644 --- a/src/Nonlinear/ReverseAD/forward_over_reverse.jl +++ b/src/Nonlinear/ReverseAD/forward_over_reverse.jl @@ -145,7 +145,7 @@ function _hessian_slice_inner(d, ex, input_ϵ, output_ϵ, ::Type{T}) where {T} _forward_eval_ϵ( d, ex, - _reinterpret_unsafe(T, d.forward_storage_ϵ), + _reinterpret_unsafe(T, d.storage_ϵ), _reinterpret_unsafe(T, d.partials_storage_ϵ), input_ϵ, subexpr_forward_values_ϵ, @@ -161,7 +161,7 @@ function _hessian_slice_inner(d, ex, input_ϵ, output_ϵ, ::Type{T}) where {T} _reverse_eval_ϵ( output_ϵ, ex, - _reinterpret_unsafe(T, d.reverse_storage_ϵ), + _reinterpret_unsafe(T, d.storage_ϵ), _reinterpret_unsafe(T, d.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index 4157ac54f5..c76cdcb001 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -138,14 +138,16 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) end # 10 is hardcoded upper bound to avoid excess memory allocation max_chunk = min(max_chunk, 10) + max_expr_with_sub_length = + maximum(d.subexpressions; init = max_expr_length) do subexpr + return length(subexpr.nodes) + end if d.want_hess || want_hess_storage d.input_ϵ = zeros(max_chunk * N) d.output_ϵ = zeros(max_chunk * N) # - len = max_chunk * max_expr_length - d.forward_storage_ϵ = zeros(len) - d.partials_storage_ϵ = zeros(len) - d.reverse_storage_ϵ = zeros(len) + d.partials_storage_ϵ = zeros(max_chunk * max_expr_length) + d.storage_ϵ = zeros(max_chunk * max_expr_with_sub_length) # len = max_chunk * length(d.subexpressions) d.subexpression_forward_values_ϵ = zeros(len) @@ -153,12 +155,8 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) # for k in d.subexpression_order len = max_chunk * length(d.subexpressions[k].nodes) - resize!(d.subexpressions[k].forward_storage_ϵ, len) - fill!(d.subexpressions[k].forward_storage_ϵ, 0.0) resize!(d.subexpressions[k].partials_storage_ϵ, len) fill!(d.subexpressions[k].partials_storage_ϵ, 0.0) - resize!(d.subexpressions[k].reverse_storage_ϵ, len) - fill!(d.subexpressions[k].reverse_storage_ϵ, 0.0) end d.max_chunk = max_chunk if d.want_hess @@ -350,7 +348,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( d, subexpr, - reinterpret(T, subexpr.forward_storage_ϵ), + reinterpret(T, subexpr.storage_ϵ), reinterpret(T, subexpr.partials_storage_ϵ), input_ϵ, subexpr_forward_values_ϵ, @@ -361,13 +359,13 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) subexpr_reverse_values_ϵ = reinterpret(T, d.subexpression_reverse_values_ϵ) fill!(subexpr_reverse_values_ϵ, zero(T)) fill!(d.subexpression_reverse_values, 0.0) - fill!(d.reverse_storage_ϵ, 0.0) + fill!(d.storage_ϵ, 0.0) fill!(output_ϵ, zero(T)) if d.objective !== nothing _forward_eval_ϵ( d, something(d.objective), - reinterpret(T, d.forward_storage_ϵ), + reinterpret(T, d.storage_ϵ), reinterpret(T, d.partials_storage_ϵ), input_ϵ, subexpr_forward_values_ϵ, @@ -376,7 +374,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) _reverse_eval_ϵ( output_ϵ, something(d.objective), - reinterpret(T, d.reverse_storage_ϵ), + reinterpret(T, d.storage_ϵ), reinterpret(T, d.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, @@ -388,7 +386,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) _forward_eval_ϵ( d, con, - reinterpret(T, d.forward_storage_ϵ), + reinterpret(T, d.storage_ϵ), reinterpret(T, d.partials_storage_ϵ), input_ϵ, subexpr_forward_values_ϵ, @@ -397,7 +395,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) _reverse_eval_ϵ( output_ϵ, con, - reinterpret(T, d.reverse_storage_ϵ), + reinterpret(T, d.storage_ϵ), reinterpret(T, d.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, @@ -411,7 +409,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) _reverse_eval_ϵ( output_ϵ, subexpr, - reinterpret(T, subexpr.reverse_storage_ϵ), + reinterpret(T, d.storage_ϵ), reinterpret(T, subexpr.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, diff --git a/src/Nonlinear/ReverseAD/types.jl b/src/Nonlinear/ReverseAD/types.jl index 84688a9e07..d41b10a764 100644 --- a/src/Nonlinear/ReverseAD/types.jl +++ b/src/Nonlinear/ReverseAD/types.jl @@ -11,9 +11,7 @@ struct _SubexpressionStorage forward_storage::Vector{Float64} partials_storage::Vector{Float64} reverse_storage::Vector{Float64} - forward_storage_ϵ::Vector{Float64} partials_storage_ϵ::Vector{Float64} - reverse_storage_ϵ::Vector{Float64} linearity::Linearity function _SubexpressionStorage( @@ -175,11 +173,10 @@ mutable struct NLPEvaluator <: MOI.AbstractNLPEvaluator # so the length should be multiplied by the maximum number of epsilon components disable_2ndorder::Bool # don't offer Hess or HessVec want_hess::Bool - forward_storage_ϵ::Vector{Float64} # (longest expression) - partials_storage_ϵ::Vector{Float64} # (longest expression) - reverse_storage_ϵ::Vector{Float64} # (longest expression) + partials_storage_ϵ::Vector{Float64} # (longest expression excluding subexpressions) + storage_ϵ::Vector{Float64} # (longest expression including subexpressions) input_ϵ::Vector{Float64} # (number of variables) - output_ϵ::Vector{Float64}# (number of variables) + output_ϵ::Vector{Float64} # (number of variables) subexpression_forward_values_ϵ::Vector{Float64} # (number of subexpressions) subexpression_reverse_values_ϵ::Vector{Float64} # (number of subexpressions) hessian_sparsity::Vector{Tuple{Int64,Int64}} From 84f6ded125da1ccd5faa47bfc5f4288eafefcd36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 23 Apr 2025 13:32:11 +0200 Subject: [PATCH 2/6] Fixes --- src/Nonlinear/ReverseAD/forward_over_reverse.jl | 4 ++-- src/Nonlinear/ReverseAD/mathoptinterface_api.jl | 2 +- src/Nonlinear/ReverseAD/types.jl | 2 -- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/Nonlinear/ReverseAD/forward_over_reverse.jl b/src/Nonlinear/ReverseAD/forward_over_reverse.jl index e469b6ba0d..251eb29456 100644 --- a/src/Nonlinear/ReverseAD/forward_over_reverse.jl +++ b/src/Nonlinear/ReverseAD/forward_over_reverse.jl @@ -135,7 +135,7 @@ function _hessian_slice_inner(d, ex, input_ϵ, output_ϵ, ::Type{T}) where {T} subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( d, subexpr, - _reinterpret_unsafe(T, subexpr.forward_storage_ϵ), + _reinterpret_unsafe(T, d.storage_ϵ), _reinterpret_unsafe(T, subexpr.partials_storage_ϵ), input_ϵ, subexpr_forward_values_ϵ, @@ -174,7 +174,7 @@ function _hessian_slice_inner(d, ex, input_ϵ, output_ϵ, ::Type{T}) where {T} _reverse_eval_ϵ( output_ϵ, subexpr, - _reinterpret_unsafe(T, subexpr.reverse_storage_ϵ), + _reinterpret_unsafe(T, d.storage_ϵ), _reinterpret_unsafe(T, subexpr.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index c76cdcb001..39117e271a 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -348,7 +348,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( d, subexpr, - reinterpret(T, subexpr.storage_ϵ), + reinterpret(T, d.storage_ϵ), reinterpret(T, subexpr.partials_storage_ϵ), input_ϵ, subexpr_forward_values_ϵ, diff --git a/src/Nonlinear/ReverseAD/types.jl b/src/Nonlinear/ReverseAD/types.jl index d41b10a764..6526a2d461 100644 --- a/src/Nonlinear/ReverseAD/types.jl +++ b/src/Nonlinear/ReverseAD/types.jl @@ -32,8 +32,6 @@ struct _SubexpressionStorage zeros(N), # partials_storage, zeros(N), # reverse_storage, Float64[], - Float64[], - Float64[], linearity[1], ) end From 707ab59aef5f1775fa725c735e1927a90eec2834 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 23 Apr 2025 14:18:05 +0200 Subject: [PATCH 3/6] Fix --- src/Nonlinear/ReverseAD/mathoptinterface_api.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index 39117e271a..729ea2e730 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -139,8 +139,8 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) # 10 is hardcoded upper bound to avoid excess memory allocation max_chunk = min(max_chunk, 10) max_expr_with_sub_length = - maximum(d.subexpressions; init = max_expr_length) do subexpr - return length(subexpr.nodes) + maximum(d.subexpression_order; init = max_expr_length) do k + return length(d.subexpressions[k].nodes) end if d.want_hess || want_hess_storage d.input_ϵ = zeros(max_chunk * N) From 7e9c6cbeec1b98e2ad734d877690d31a27d7b723 Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Thu, 24 Apr 2025 08:35:19 +1200 Subject: [PATCH 4/6] Update mathoptinterface_api.jl --- src/Nonlinear/ReverseAD/mathoptinterface_api.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index 729ea2e730..22e7827b7c 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -138,10 +138,10 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) end # 10 is hardcoded upper bound to avoid excess memory allocation max_chunk = min(max_chunk, 10) - max_expr_with_sub_length = - maximum(d.subexpression_order; init = max_expr_length) do k - return length(d.subexpressions[k].nodes) - end + max_expr_with_sub_length = max( + max_expr_length, + maximum(x -> length(x.nodes), d.subexpressions; init = 0), + ) if d.want_hess || want_hess_storage d.input_ϵ = zeros(max_chunk * N) d.output_ϵ = zeros(max_chunk * N) From 0ad6caa8a37a7c5a257456080b2c4ff32a413579 Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Thu, 24 Apr 2025 10:04:42 +1200 Subject: [PATCH 5/6] Update mathoptinterface_api.jl --- src/Nonlinear/ReverseAD/mathoptinterface_api.jl | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index 22e7827b7c..42677986fe 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -44,6 +44,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) want_hess_storage = (:HessVec in requested_features) || d.want_hess coloring_storage = Coloring.IndexedSet(N) max_expr_length = 0 + max_expr_with_sub_length = 0 # main_expressions = [c.expression.nodes for (_, c) in d.data.constraints] if d.data.objective !== nothing @@ -71,6 +72,8 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) ) d.subexpressions[k] = subex d.subexpression_linearity[k] = subex.linearity + max_expr_with_sub_length = + max(max_expr_with_sub_length, length(subex.nodes)) if d.want_hess empty!(coloring_storage) _compute_gradient_sparsity!(coloring_storage, subex.nodes) @@ -138,10 +141,8 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) end # 10 is hardcoded upper bound to avoid excess memory allocation max_chunk = min(max_chunk, 10) - max_expr_with_sub_length = max( - max_expr_length, - maximum(x -> length(x.nodes), d.subexpressions; init = 0), - ) + max_expr_with_sub_length = + max(max_expr_with_sub_length, max_expr_length) if d.want_hess || want_hess_storage d.input_ϵ = zeros(max_chunk * N) d.output_ϵ = zeros(max_chunk * N) From 6da0650ea626230de6e5105163cb55a980db8ac5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 24 Apr 2025 08:32:45 +0200 Subject: [PATCH 6/6] Fix format --- src/Nonlinear/ReverseAD/mathoptinterface_api.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index 42677986fe..787eb62673 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -141,8 +141,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) end # 10 is hardcoded upper bound to avoid excess memory allocation max_chunk = min(max_chunk, 10) - max_expr_with_sub_length = - max(max_expr_with_sub_length, max_expr_length) + max_expr_with_sub_length = max(max_expr_with_sub_length, max_expr_length) if d.want_hess || want_hess_storage d.input_ϵ = zeros(max_chunk * N) d.output_ϵ = zeros(max_chunk * N)