[Nonlinear] Merge forward_storage_ϵ with reverse_storage_ϵ

Author: blegat

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_ϵ,

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -138,27 +138,25 @@ 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)
         d.subexpression_reverse_values_ϵ = zeros(len)
         #
         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_ϵ,

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}}