Use _SubexpressionStorage inside _FunctionStorage

src/Nonlinear/ReverseAD/forward_over_reverse.jl

@@ -38,7 +38,7 @@ function _eval_hessian(
     scale::Float64,
     nzcount::Int,
 )::Int
-    if ex.linearity == LINEAR
+    if ex.expr.linearity == LINEAR
         @assert length(ex.hess_I) == 0
         return 0
     end
@@ -128,13 +128,9 @@ function _hessian_slice_inner(d, ex, ::Type{T}) where {T}
         _reinterpret_unsafe(T, d.subexpression_forward_values_ϵ)
     for i in ex.dependent_subexpressions
         subexpr = d.subexpressions[i]
-        subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(
-            d,
-            subexpr,
-            _reinterpret_unsafe(T, subexpr.partials_storage_ϵ),
-        )
+        subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(d, subexpr, T)
     end
-    _forward_eval_ϵ(d, ex, _reinterpret_unsafe(T, d.partials_storage_ϵ))
+    _forward_eval_ϵ(d, ex.expr, T)
     # do a reverse pass
     subexpr_reverse_values_ϵ =
         _reinterpret_unsafe(T, d.subexpression_reverse_values_ϵ)
@@ -144,9 +140,8 @@ function _hessian_slice_inner(d, ex, ::Type{T}) where {T}
     end
     _reverse_eval_ϵ(
         output_ϵ,
-        ex,
+        ex.expr,
         _reinterpret_unsafe(T, d.storage_ϵ),
-        _reinterpret_unsafe(T, d.partials_storage_ϵ),
         d.subexpression_reverse_values,
         subexpr_reverse_values_ϵ,
         1.0,
@@ -159,7 +154,6 @@ function _hessian_slice_inner(d, ex, ::Type{T}) where {T}
             output_ϵ,
             subexpr,
             _reinterpret_unsafe(T, d.storage_ϵ),
-            _reinterpret_unsafe(T, subexpr.partials_storage_ϵ),
             d.subexpression_reverse_values,
             subexpr_reverse_values_ϵ,
             d.subexpression_reverse_values[j],
@@ -173,8 +167,8 @@ end
     _forward_eval_ϵ(
         d::NLPEvaluator,
         ex::Union{_FunctionStorage,_SubexpressionStorage},
-        partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}},
-    ) where {N,T}
+        ::Type{P},
+    ) where {N,T,P<:ForwardDiff.Partials{N,T}}
 
 Evaluate the directional derivatives of the expression tree in `ex`.
 
@@ -186,10 +180,11 @@ This assumes that `_reverse_model(d, x)` has already been called.
 """
 function _forward_eval_ϵ(
     d::NLPEvaluator,
-    ex::Union{_FunctionStorage,_SubexpressionStorage},
-    partials_storage_ϵ::AbstractVector{P},
+    ex::_SubexpressionStorage,
+    ::Type{P},
 ) where {N,T,P<:ForwardDiff.Partials{N,T}}
     storage_ϵ = _reinterpret_unsafe(P, d.storage_ϵ)
+    partials_storage_ϵ = _reinterpret_unsafe(P, ex.partials_storage_ϵ)
     x_values_ϵ = _reinterpret_unsafe(P, d.input_ϵ)
     subexpression_values_ϵ =
         _reinterpret_unsafe(P, d.subexpression_forward_values_ϵ)
@@ -370,14 +365,15 @@ end
 # to compute hessian-vector products.
 function _reverse_eval_ϵ(
     output_ϵ::AbstractVector{ForwardDiff.Partials{N,T}},
-    ex::Union{_FunctionStorage,_SubexpressionStorage},
+    ex::_SubexpressionStorage,
     reverse_storage_ϵ,
-    partials_storage_ϵ,
     subexpression_output,
     subexpression_output_ϵ,
     scale::T,
     scale_ϵ::ForwardDiff.Partials{N,T},
 ) where {N,T}
+    partials_storage_ϵ =
+        _reinterpret_unsafe(ForwardDiff.Partials{N,T}, ex.partials_storage_ϵ)
     @assert length(reverse_storage_ϵ) >= length(ex.nodes)
     @assert length(partials_storage_ϵ) >= length(ex.nodes)
     if ex.nodes[1].type == Nonlinear.NODE_VARIABLE

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -65,11 +65,12 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     for k in d.subexpression_order
         # Only load expressions which actually are used
         d.subexpression_forward_values[k] = NaN
-        subex = _SubexpressionStorage(
-            d.data.expressions[k],
-            d.subexpression_linearity,
+        expr = d.data.expressions[k]
+        subex, _ = _subexpression_and_linearity(
+            expr,
             moi_index_to_consecutive_index,
-            d.want_hess,
+            Float64[],
+            d,
         )
         d.subexpressions[k] = subex
         d.subexpression_linearity[k] = subex.linearity
@@ -101,43 +102,54 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
         end
     end
     max_chunk = 1
+    shared_partials_storage_ϵ = Float64[]
     if d.data.objective !== nothing
+        expr = something(d.data.objective)
+        subexpr, linearity = _subexpression_and_linearity(
+            expr,
+            moi_index_to_consecutive_index,
+            shared_partials_storage_ϵ,
+            d,
+        )
         objective = _FunctionStorage(
-            main_expressions[1],
-            something(d.data.objective).values,
+            subexpr,
             N,
             coloring_storage,
             d.want_hess,
             d.subexpressions,
             individual_order[1],
-            d.subexpression_linearity,
             subexpression_edgelist,
             subexpression_variables,
-            moi_index_to_consecutive_index,
+            linearity,
         )
-        max_expr_length = max(max_expr_length, length(objective.nodes))
+        max_expr_length = max(max_expr_length, length(expr.nodes))
         max_chunk = max(max_chunk, size(objective.seed_matrix, 2))
         d.objective = objective
     end
     for (k, (_, constraint)) in enumerate(d.data.constraints)
         idx = d.data.objective !== nothing ? k + 1 : k
+        expr = constraint.expression
+        subexpr, linearity = _subexpression_and_linearity(
+            expr,
+            moi_index_to_consecutive_index,
+            shared_partials_storage_ϵ,
+            d,
+        )
         push!(
             d.constraints,
             _FunctionStorage(
-                main_expressions[idx],
-                constraint.expression.values,
+                subexpr,
                 N,
                 coloring_storage,
                 d.want_hess,
                 d.subexpressions,
                 individual_order[idx],
-                d.subexpression_linearity,
                 subexpression_edgelist,
                 subexpression_variables,
-                moi_index_to_consecutive_index,
+                linearity,
             ),
         )
-        max_expr_length = max(max_expr_length, length(d.constraints[end].nodes))
+        max_expr_length = max(max_expr_length, length(expr.nodes))
         max_chunk = max(max_chunk, size(d.constraints[end].seed_matrix, 2))
     end
     max_chunk = min(max_chunk, MAX_CHUNK)
@@ -146,7 +158,8 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
         d.input_ϵ = zeros(max_chunk * N)
         d.output_ϵ = zeros(max_chunk * N)
         #
-        d.partials_storage_ϵ = zeros(max_chunk * max_expr_length)
+        resize!(shared_partials_storage_ϵ, max_chunk * max_expr_length)
+        fill!(shared_partials_storage_ϵ, 0.0)
         d.storage_ϵ = zeros(max_chunk * max_expr_with_sub_length)
         #
         len = max_chunk * length(d.subexpressions)
@@ -178,7 +191,7 @@ function MOI.eval_objective(d::NLPEvaluator, x)
         error("No nonlinear objective.")
     end
     _reverse_mode(d, x)
-    return something(d.objective).forward_storage[1]
+    return something(d.objective).expr.forward_storage[1]
 end
 
 function MOI.eval_objective_gradient(d::NLPEvaluator, g, x)
@@ -194,7 +207,7 @@ end
 function MOI.eval_constraint(d::NLPEvaluator, g, x)
     _reverse_mode(d, x)
     for i in 1:length(d.constraints)
-        g[i] = d.constraints[i].forward_storage[1]
+        g[i] = d.constraints[i].expr.forward_storage[1]
     end
     return
 end
@@ -345,11 +358,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
     subexpr_forward_values_ϵ = reinterpret(T, d.subexpression_forward_values_ϵ)
     for i in d.subexpression_order
         subexpr = d.subexpressions[i]
-        subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(
-            d,
-            subexpr,
-            reinterpret(T, subexpr.partials_storage_ϵ),
-        )
+        subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(d, subexpr, T)
     end
     # we only need to do one reverse pass through the subexpressions as well
     subexpr_reverse_values_ϵ = reinterpret(T, d.subexpression_reverse_values_ϵ)
@@ -358,29 +367,23 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
     fill!(d.storage_ϵ, 0.0)
     fill!(output_ϵ, zero(T))
     if d.objective !== nothing
-        _forward_eval_ϵ(
-            d,
-            something(d.objective),
-            reinterpret(T, d.partials_storage_ϵ),
-        )
+        _forward_eval_ϵ(d, something(d.objective).expr, T)
         _reverse_eval_ϵ(
             output_ϵ,
-            something(d.objective),
-            reinterpret(T, d.storage_ϵ),
-            reinterpret(T, d.partials_storage_ϵ),
+            something(d.objective).expr,
+            _reinterpret_unsafe(T, d.storage_ϵ),
             d.subexpression_reverse_values,
             subexpr_reverse_values_ϵ,
             σ,
             zero(T),
         )
     end
     for (i, con) in enumerate(d.constraints)
-        _forward_eval_ϵ(d, con, reinterpret(T, d.partials_storage_ϵ))
+        _forward_eval_ϵ(d, con.expr, T)
         _reverse_eval_ϵ(
             output_ϵ,
-            con,
+            con.expr,
             reinterpret(T, d.storage_ϵ),
-            reinterpret(T, d.partials_storage_ϵ),
             d.subexpression_reverse_values,
             subexpr_reverse_values_ϵ,
             μ[i],
@@ -394,7 +397,6 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
             output_ϵ,
             subexpr,
             reinterpret(T, d.storage_ϵ),
-            reinterpret(T, subexpr.partials_storage_ϵ),
             d.subexpression_reverse_values,
             subexpr_reverse_values_ϵ,
             d.subexpression_reverse_values[j],

src/Nonlinear/ReverseAD/reverse_mode.jl

@@ -39,20 +39,20 @@ function _reverse_mode(d::NLPEvaluator, x)
             _forward_eval(d.subexpressions[k], d, x)
     end
     if d.objective !== nothing
-        _forward_eval(d.objective::_FunctionStorage, d, x)
+        _forward_eval(something(d.objective).expr, d, x)
     end
     for con in d.constraints
-        _forward_eval(con, d, x)
+        _forward_eval(con.expr, d, x)
     end
     # Phase II
     for k in d.subexpression_order
         _reverse_eval(d.subexpressions[k])
     end
     if d.objective !== nothing
-        _reverse_eval(d.objective::_FunctionStorage)
+        _reverse_eval(something(d.objective).expr)
     end
     for con in d.constraints
-        _reverse_eval(con)
+        _reverse_eval(con.expr)
     end
     # If a JuMP model uses the legacy nonlinear interface, then JuMP constructs
     # a NLPEvaluator at the start of a call to `JuMP.optimize!` and it passes in
@@ -81,7 +81,7 @@ end
 
 """
     _forward_eval(
-        f::Union{_FunctionStorage,_SubexpressionStorage},
+        f::_SubexpressionStorage,
         d::NLPEvaluator,
         x::AbstractVector{T},
     ) where {T}
@@ -98,10 +98,7 @@ Forward-mode evaluation of an expression tree given in `f`.
    associate storage with each edge of the DAG.
 """
 function _forward_eval(
-    # !!! warning
-    #     This Union depends upon _FunctionStorage and _SubexpressionStorage
-    #     having similarly named fields.
-    f::Union{_FunctionStorage,_SubexpressionStorage},
+    f::_SubexpressionStorage,
     d::NLPEvaluator,
     x::AbstractVector{T},
 )::T where {T}
@@ -290,19 +287,14 @@ function _forward_eval(
 end
 
 """
-    _reverse_eval(f::Union{_FunctionStorage,_SubexpressionStorage})
+    _reverse_eval(f::_SubexpressionStorage)
 
 Reverse-mode evaluation of an expression tree given in `f`.
 
  * This function assumes `f.partials_storage` is already updated.
  * This function assumes that `f.reverse_storage` has been initialized with 0.0.
 """
-function _reverse_eval(
-    # !!! warning
-    #     This Union depends upon _FunctionStorage and _SubexpressionStorage
-    #     having similarly named fields.
-    f::Union{_FunctionStorage,_SubexpressionStorage},
-)
+function _reverse_eval(f::_SubexpressionStorage)
     @assert length(f.reverse_storage) >= length(f.nodes)
     @assert length(f.partials_storage) >= length(f.nodes)
     # f.nodes is already in order such that parents always appear before
@@ -361,9 +353,15 @@ end
 
 function _extract_reverse_pass_inner(
     output::AbstractVector{T},
-    # !!! warning
-    #     This Union depends upon _FunctionStorage and _SubexpressionStorage
-    #     having similarly named fields.
+    f::_FunctionStorage,
+    subexpressions::AbstractVector{T},
+    scale::T,
+) where {T}
+    return _extract_reverse_pass_inner(output, f.expr, subexpressions, scale)
+end
+
+function _extract_reverse_pass_inner(
+    output::AbstractVector{T},
     f::Union{_FunctionStorage,_SubexpressionStorage},
     subexpressions::AbstractVector{T},
     scale::T,