From 9cd51f9936d113397715df8851fe79f16d27604f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 6 Aug 2025 10:25:22 +0200 Subject: [PATCH 1/7] Convert common sub-functions as common sub-expressions --- src/Nonlinear/parse.jl | 127 ++++++++++++++++++++++++++++++- src/Nonlinear/types.jl | 12 ++- test/Nonlinear/test_Nonlinear.jl | 88 +++++++++++++++++++++ 3 files changed, 224 insertions(+), 3 deletions(-) diff --git a/src/Nonlinear/parse.jl b/src/Nonlinear/parse.jl index 05fd684429..0c0e7f449b 100644 --- a/src/Nonlinear/parse.jl +++ b/src/Nonlinear/parse.jl @@ -36,6 +36,77 @@ function parse_expression(::Model, ::Expression, x::Any, ::Int) ) end +function _extract_subexpression!(expr::Expression, root::Int) + n = length(expr.nodes) + # The whole subexpression is continuous in the tape + first_out = first_value = last_value = nothing + for i in root:n + node = expr.nodes[i] + if i != root && node.parent < root + first_out = i + break + end + index = node.index + if node.type == NODE_VALUE + if isnothing(first_value) + first_value = node.index + last_value = first_value + else + last_value = node.index + end + index -= first_value - 1 + end + expr.nodes[i] = + Node(node.type, index, i == root ? -1 : node.parent - root + 1) + end + if isnothing(first_out) + I = root:n + else + I = root:(first_out-1) + end + if isnothing(first_value) + V = nothing + else + V = first_value:last_value + end + if !isnothing(first_out) + for i in (last(I)+1):n + node = expr.nodes[i] + index = node.index + if node.type == NODE_VALUE && !isnothing(V) + @assert index >= last(V) + index -= length(V) + end + parent = node.parent + if parent > root + @assert parent > last(I) + parent -= length(I) - 1 + end + expr.nodes[i] = Node(node.type, index, parent) + end + end + return I, V +end + +function _extract_subexpression!(data::Model, expr::Expression, root::Int) + parent = expr.nodes[root].parent + I, V = _extract_subexpression!(expr, root) + subexpr = + Expression(expr.nodes[I], isnothing(V) ? Float64[] : expr.values[V]) + push!(data.expressions, subexpr) + index = ExpressionIndex(length(data.expressions)) + expr.nodes[root] = Node(NODE_SUBEXPRESSION, index.value, parent) + if length(I) > 1 + deleteat!(expr.nodes, I[2:end]) + if !isnothing(V) + deleteat!(expr.values, V) + end + else + @assert isnothing(V) + end + return index, I +end + function parse_expression( data::Model, expr::Expression, @@ -46,7 +117,59 @@ function parse_expression( while !isempty(stack) parent_node, arg = pop!(stack) if arg isa MOI.ScalarNonlinearFunction - _parse_without_recursion_inner(stack, data, expr, arg, parent_node) + if haskey(data.cache, arg) + subexpr = data.cache[arg] + if subexpr isa Tuple{Expression,Int} + _expr, _node = subexpr + subexpr, I = _extract_subexpression!(data, _expr, _node) + if expr === _expr + if parent_node > first(I) + @assert parent_node > last(I) + parent_node -= length(I) - 1 + end + for i in eachindex(stack) + _parent_node = stack[i][1] + if _parent_node > first(I) + @assert _parent_node > last(I) + stack[i] = + (_parent_node - length(I) + 1, stack[i][2]) + end + end + end + for (key, val) in data.cache + if val isa Tuple{Expression,Int} + __expr, __node = val + if _expr === __expr && __node > first(I) + if __node <= last(I) + data.cache[key] = ( + data.expressions[subexpr.value], + __node - first(I) + 1, + ) + else + data.cache[key] = + (__expr, __node - length(I) + 1) + end + end + end + end + data.cache[arg] = subexpr + end + parse_expression( + data, + expr, + subexpr::ExpressionIndex, + parent_node, + ) + else + _parse_without_recursion_inner( + stack, + data, + expr, + arg, + parent_node, + ) + data.cache[arg] = (expr, length(expr.nodes)) + end else # We can use recursion here, because ScalarNonlinearFunction only # occur in other ScalarNonlinearFunction. @@ -82,7 +205,7 @@ function _parse_without_recursion_inner(stack, data, expr, x, parent) parent = length(expr.nodes) # Args need to be pushed onto the stack in reverse because the stack is a # first-in last-out datastructure. - for arg in reverse(x.args) + for arg in Iterators.Reverse(x.args) push!(stack, (parent, arg)) end return diff --git a/src/Nonlinear/types.jl b/src/Nonlinear/types.jl index 36ac83a245..3ae8abde91 100644 --- a/src/Nonlinear/types.jl +++ b/src/Nonlinear/types.jl @@ -76,9 +76,10 @@ tree. struct Expression nodes::Vector{Node} values::Vector{Float64} - Expression() = new(Node[], Float64[]) end +Expression() = Expression(Node[], Float64[]) + function Base.:(==)(x::Expression, y::Expression) return x.nodes == y.nodes && x.values == y.values end @@ -165,6 +166,11 @@ mutable struct Model operators::OperatorRegistry # This is a private field, used only to increment the ConstraintIndex. last_constraint_index::Int64 + # This is a private field, used to detect common subexpressions. + cache::Dict{ + MOI.ScalarNonlinearFunction, + Union{ExpressionIndex,Tuple{Expression,Int}}, + } function Model() return new( nothing, @@ -173,6 +179,10 @@ mutable struct Model Float64[], OperatorRegistry(), 0, + Dict{ + MOI.ScalarNonlinearFunction, + Union{ExpressionIndex,Tuple{Expression,Int}}, + }(), ) end end diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index e335c2ced5..0d7d404d7f 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1438,6 +1438,94 @@ function test_intercept_ForwardDiff_MethodError() return end +function test_extract_subexpression() + model = Nonlinear.Model() + x = MOI.VariableIndex(1) + sub = MOI.ScalarNonlinearFunction(:^, Any[x, 3]) + f = MOI.ScalarNonlinearFunction(:+, Any[sub, sub]) + expr = Nonlinear.parse_expression(model, f) + display(expr.nodes) + @test expr == Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + ], + Float64[], + ) + expected_sub = Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 4, -1) + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1) + Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 1) + ], + [3.0], + ) + @test model.expressions == [expected_sub] + @test model.cache[sub] == Nonlinear.ExpressionIndex(1) + + h = MOI.ScalarNonlinearFunction(:*, Any[2, sub, 1]) + g = MOI.ScalarNonlinearFunction(:+, Any[sub, h]) + expr = MOI.Nonlinear.parse_expression(model, g) + expected_g = Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1) + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1) + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, 1) + Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 3) + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 3) + Nonlinear.Node(Nonlinear.NODE_VALUE, 2, 3) + ], + [2.0, 1.0], + ) + @test expr == expected_g + # It should have detected the sub-expressions that was the same as `f` + @test model.expressions == [expected_sub] + # This means that it didn't get to extract from `g`, let's also test + # with extraction by starting with an empty model + + model = Nonlinear.Model() + MOI.Nonlinear.set_objective(model, g) + @test model.objective == expected_g + @test model.expressions == [expected_sub] + # Test that the objective function gets rewritten as we reuse `h` + # Also test that we don't change the parents in the stack of `h` + # by creating a long stack + prod = MOI.ScalarNonlinearFunction(:*, [h, x]) + sum = MOI.ScalarNonlinearFunction(:*, [x, x, x, x, prod]) + expr = Nonlinear.parse_expression(model, sum) + @test isempty(model.objective.values) + @test model.objective.nodes == [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, 1), + ] + @test model.expressions == [ + expected_sub, + Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, -1), + Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_VALUE, 2, 1), + ], + [2.0, 1.0], + ), + ] + @test isempty(expr.values) + @test expr.nodes == [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, -1), + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1), + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1), + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1), + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1), + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, 6), + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 6), + ] + return +end + end # TestNonlinear TestNonlinear.runtests() From 49c78ffd374ad672afe2689ee074f3a1ff943d12 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Mon, 14 Sep 2026 11:41:05 +0200 Subject: [PATCH 2/7] Add test --- test/Nonlinear/test_Nonlinear.jl | 36 ++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index 0d7d404d7f..352fa62753 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1464,6 +1464,42 @@ function test_extract_subexpression() @test model.expressions == [expected_sub] @test model.cache[sub] == Nonlinear.ExpressionIndex(1) + # Test extraction of an expression without values, and ensure that cached + # nested expressions are updated to point into the extracted expression. + model = Nonlinear.Model() + inner = MOI.ScalarNonlinearFunction(:sin, Any[x]) + outer = MOI.ScalarNonlinearFunction(:+, Any[inner, x]) + repeated = MOI.ScalarNonlinearFunction(:*, Any[outer, outer]) + expr = Nonlinear.parse_expression(model, repeated) + @test isempty(model.expressions[1].values) + @test model.cache[inner] == (model.expressions[1], 2) + @test model.cache[outer] == Nonlinear.ExpressionIndex(1) + @test expr.nodes == [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, -1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + ] + + # A zero-argument call consists of one node, so extracting it exercises + # the single-node path, which has no values to delete. + model = Nonlinear.Model() + empty_sum = MOI.ScalarNonlinearFunction(:+, Any[]) + repeated = MOI.ScalarNonlinearFunction(:+, Any[empty_sum, empty_sum]) + expr = Nonlinear.parse_expression(model, repeated) + @test model.expressions == [ + Nonlinear.Expression( + [Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1)], + Float64[], + ), + ] + @test expr.nodes == [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + ] + + model = Nonlinear.Model() + Nonlinear.parse_expression(model, f) h = MOI.ScalarNonlinearFunction(:*, Any[2, sub, 1]) g = MOI.ScalarNonlinearFunction(:+, Any[sub, h]) expr = MOI.Nonlinear.parse_expression(model, g) From 851f93922ae59fea9398c2dbcdfd409d3fce255a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 15 Oct 2026 19:02:18 +0200 Subject: [PATCH 3/7] Use Tuple{Int,Int} --- src/Nonlinear/parse.jl | 43 ++++++++++++++++++-------------- src/Nonlinear/types.jl | 19 +++++++------- test/Nonlinear/test_Nonlinear.jl | 31 ++++++++++++++++++++--- 3 files changed, 62 insertions(+), 31 deletions(-) diff --git a/src/Nonlinear/parse.jl b/src/Nonlinear/parse.jl index 0c0e7f449b..4e73eb33ba 100644 --- a/src/Nonlinear/parse.jl +++ b/src/Nonlinear/parse.jl @@ -114,14 +114,22 @@ function parse_expression( parent_index::Int, ) stack = Tuple{Int,Any}[(parent_index, x)] + expr_index = 0 while !isempty(stack) parent_node, arg = pop!(stack) if arg isa MOI.ScalarNonlinearFunction if haskey(data.cache, arg) - subexpr = data.cache[arg] - if subexpr isa Tuple{Expression,Int} - _expr, _node = subexpr + _expr_index, _node = data.cache[arg] + if _node == 0 + subexpr = ExpressionIndex(_expr_index) + else + _expr = data.cache_expressions[_expr_index] subexpr, I = _extract_subexpression!(data, _expr, _node) + push!( + data.cache_expressions, + data.expressions[subexpr.value], + ) + subexpr_index = length(data.cache_expressions) if expr === _expr if parent_node > first(I) @assert parent_node > last(I) @@ -136,30 +144,23 @@ function parse_expression( end end end - for (key, val) in data.cache - if val isa Tuple{Expression,Int} - __expr, __node = val + for (key, (__expr_index, __node)) in data.cache + if __node != 0 + __expr = data.cache_expressions[__expr_index] if _expr === __expr && __node > first(I) if __node <= last(I) - data.cache[key] = ( - data.expressions[subexpr.value], - __node - first(I) + 1, - ) + data.cache[key] = + (subexpr_index, __node - first(I) + 1) else data.cache[key] = - (__expr, __node - length(I) + 1) + (__expr_index, __node - length(I) + 1) end end end end - data.cache[arg] = subexpr + data.cache[arg] = (subexpr.value, 0) end - parse_expression( - data, - expr, - subexpr::ExpressionIndex, - parent_node, - ) + parse_expression(data, expr, subexpr, parent_node) else _parse_without_recursion_inner( stack, @@ -168,7 +169,11 @@ function parse_expression( arg, parent_node, ) - data.cache[arg] = (expr, length(expr.nodes)) + if expr_index == 0 + push!(data.cache_expressions, expr) + expr_index = length(data.cache_expressions) + end + data.cache[arg] = (expr_index, length(expr.nodes)) end else # We can use recursion here, because ScalarNonlinearFunction only diff --git a/src/Nonlinear/types.jl b/src/Nonlinear/types.jl index 3ae8abde91..161bfac51a 100644 --- a/src/Nonlinear/types.jl +++ b/src/Nonlinear/types.jl @@ -166,11 +166,14 @@ mutable struct Model operators::OperatorRegistry # This is a private field, used only to increment the ConstraintIndex. last_constraint_index::Int64 - # This is a private field, used to detect common subexpressions. - cache::Dict{ - MOI.ScalarNonlinearFunction, - Union{ExpressionIndex,Tuple{Expression,Int}}, - } + # Cache values are (expression index, node index). A zero node index + # identifies an extracted subexpression in `expressions`; otherwise the + # expression index refers to `cache_expressions`. + # We use a separate vector `cache_expressions` instead of storing + # a `Tuple{Expression,Int}` in the `Dict` as creating such tuple + # would trigger a heap allocation for each expression + cache::Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}} + cache_expressions::Vector{Expression} function Model() return new( nothing, @@ -179,10 +182,8 @@ mutable struct Model Float64[], OperatorRegistry(), 0, - Dict{ - MOI.ScalarNonlinearFunction, - Union{ExpressionIndex,Tuple{Expression,Int}}, - }(), + Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}}(), + Expression[], ) end end diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index 352fa62753..46a416fa2a 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1462,7 +1462,7 @@ function test_extract_subexpression() [3.0], ) @test model.expressions == [expected_sub] - @test model.cache[sub] == Nonlinear.ExpressionIndex(1) + @test model.cache[sub] == (1, 0) # Test extraction of an expression without values, and ensure that cached # nested expressions are updated to point into the extracted expression. @@ -1472,14 +1472,39 @@ function test_extract_subexpression() repeated = MOI.ScalarNonlinearFunction(:*, Any[outer, outer]) expr = Nonlinear.parse_expression(model, repeated) @test isempty(model.expressions[1].values) - @test model.cache[inner] == (model.expressions[1], 2) - @test model.cache[outer] == Nonlinear.ExpressionIndex(1) + index, node = model.cache[inner] + @test model.cache_expressions[index] === model.expressions[1] + @test node == 2 + @test model.cache[outer] == (1, 0) @test expr.nodes == [ Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, -1), Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), ] + # Reusing an extracted expression does not register another tape. + num_cached_expressions = length(model.cache_expressions) + reused = Nonlinear.parse_expression(model, outer) + @test length(model.cache_expressions) == num_cached_expressions + @test reused.nodes == [Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, -1)] + + # A nested expression can still be extracted after moving to another tape. + reused = Nonlinear.parse_expression(model, inner) + @test model.cache[inner] == (2, 0) + @test reused.nodes == [Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, -1)] + @test model.expressions[1].nodes[2] == + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, 1) + Nonlinear.set_objective(model, repeated) + evaluator = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(evaluator, [:Grad, :Hess]) + @test MOI.eval_objective(evaluator, [0.5]) ≈ (sin(0.5) + 0.5)^2 + gradient = zeros(1) + MOI.eval_objective_gradient(evaluator, gradient, [0.5]) + @test gradient[1] ≈ 2 * (sin(0.5) + 0.5) * (cos(0.5) + 1) + hessian = zeros(length(MOI.hessian_lagrangian_structure(evaluator))) + MOI.eval_hessian_lagrangian(evaluator, hessian, [0.5], 1.0, Float64[]) + @test hessian ≈ [2 * (cos(0.5) + 1)^2 - 2 * (sin(0.5) + 0.5) * sin(0.5)] + # A zero-argument call consists of one node, so extracting it exercises # the single-node path, which has no values to delete. model = Nonlinear.Model() From 32587e930fc2bbd248f8b8a80f0c62fb28a9e722 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 21 Oct 2026 16:34:01 +0200 Subject: [PATCH 4/7] Remove display line Co-authored-by: Oscar Dowson --- test/Nonlinear/test_Nonlinear.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index 46a416fa2a..214a024213 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1444,7 +1444,6 @@ function test_extract_subexpression() sub = MOI.ScalarNonlinearFunction(:^, Any[x, 3]) f = MOI.ScalarNonlinearFunction(:+, Any[sub, sub]) expr = Nonlinear.parse_expression(model, f) - display(expr.nodes) @test expr == Nonlinear.Expression( [ Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), From c186235aede21f0f152b6c9840082829ec3ad2b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 21 Oct 2026 16:45:09 +0200 Subject: [PATCH 5/7] Split into smaller functions --- test/Nonlinear/test_Nonlinear.jl | 106 +++++++++++++++++++++++-------- 1 file changed, 78 insertions(+), 28 deletions(-) diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index 214a024213..19037a17a8 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1438,6 +1438,17 @@ function test_intercept_ForwardDiff_MethodError() return end +function _extract_subexpression_expected_sub() + return Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 4, -1) + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1) + Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 1) + ], + [3.0], + ) +end + function test_extract_subexpression() model = Nonlinear.Model() x = MOI.VariableIndex(1) @@ -1452,24 +1463,27 @@ function test_extract_subexpression() ], Float64[], ) - expected_sub = Nonlinear.Expression( - [ - Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 4, -1) - Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1) - Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 1) - ], - [3.0], - ) - @test model.expressions == [expected_sub] + @test model.expressions == [_extract_subexpression_expected_sub()] @test model.cache[sub] == (1, 0) + return +end - # Test extraction of an expression without values, and ensure that cached - # nested expressions are updated to point into the extracted expression. +# Parses `(sin(x) + x) * (sin(x) + x)` so that `sin(x) + x` is extracted as +# a subexpression without values, with `sin(x)` cached as a nested node. +function _extract_subexpression_nested_model() model = Nonlinear.Model() + x = MOI.VariableIndex(1) inner = MOI.ScalarNonlinearFunction(:sin, Any[x]) outer = MOI.ScalarNonlinearFunction(:+, Any[inner, x]) repeated = MOI.ScalarNonlinearFunction(:*, Any[outer, outer]) expr = Nonlinear.parse_expression(model, repeated) + return model, expr, x, inner, outer, repeated +end + +# Test extraction of an expression without values, and ensure that cached +# nested expressions are updated to point into the extracted expression. +function test_extract_subexpression_without_values() + model, expr, _, inner, outer, _ = _extract_subexpression_nested_model() @test isempty(model.expressions[1].values) index, node = model.cache[inner] @test model.cache_expressions[index] === model.expressions[1] @@ -1480,14 +1494,23 @@ function test_extract_subexpression() Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), ] + return +end - # Reusing an extracted expression does not register another tape. +# Reusing an extracted expression does not register another tape. +function test_extract_subexpression_reuse() + model, _, _, _, outer, _ = _extract_subexpression_nested_model() num_cached_expressions = length(model.cache_expressions) reused = Nonlinear.parse_expression(model, outer) @test length(model.cache_expressions) == num_cached_expressions @test reused.nodes == [Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, -1)] + return +end - # A nested expression can still be extracted after moving to another tape. +# A nested expression can still be extracted after moving to another tape. +function test_extract_subexpression_nested_after_move() + model, _, x, inner, outer, repeated = _extract_subexpression_nested_model() + Nonlinear.parse_expression(model, outer) reused = Nonlinear.parse_expression(model, inner) @test model.cache[inner] == (2, 0) @test reused.nodes == [Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, -1)] @@ -1503,9 +1526,12 @@ function test_extract_subexpression() hessian = zeros(length(MOI.hessian_lagrangian_structure(evaluator))) MOI.eval_hessian_lagrangian(evaluator, hessian, [0.5], 1.0, Float64[]) @test hessian ≈ [2 * (cos(0.5) + 1)^2 - 2 * (sin(0.5) + 0.5) * sin(0.5)] + return +end - # A zero-argument call consists of one node, so extracting it exercises - # the single-node path, which has no values to delete. +# A zero-argument call consists of one node, so extracting it exercises +# the single-node path, which has no values to delete. +function test_extract_subexpression_single_node() model = Nonlinear.Model() empty_sum = MOI.ScalarNonlinearFunction(:+, Any[]) repeated = MOI.ScalarNonlinearFunction(:+, Any[empty_sum, empty_sum]) @@ -1521,12 +1547,16 @@ function test_extract_subexpression() Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), ] + return +end - model = Nonlinear.Model() - Nonlinear.parse_expression(model, f) +# Returns `x`, `x^3` and `h = 2 * x^3 * 1`, `g = x^3 + h`, as well as the +# expected expression of `g` once `x^3` is extracted. +function _extract_subexpression_shared_functions() + x = MOI.VariableIndex(1) + sub = MOI.ScalarNonlinearFunction(:^, Any[x, 3]) h = MOI.ScalarNonlinearFunction(:*, Any[2, sub, 1]) g = MOI.ScalarNonlinearFunction(:+, Any[sub, h]) - expr = MOI.Nonlinear.parse_expression(model, g) expected_g = Nonlinear.Expression( [ Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1) @@ -1538,19 +1568,39 @@ function test_extract_subexpression() ], [2.0, 1.0], ) + return x, sub, h, g, expected_g +end + +# The subexpression `x^3` of `g` is detected as already extracted from a +# previously parsed function. +function test_extract_subexpression_already_extracted() + x, sub, _, g, expected_g = _extract_subexpression_shared_functions() + model = Nonlinear.Model() + f = MOI.ScalarNonlinearFunction(:+, Any[sub, sub]) + Nonlinear.parse_expression(model, f) + expr = Nonlinear.parse_expression(model, g) @test expr == expected_g - # It should have detected the sub-expressions that was the same as `f` - @test model.expressions == [expected_sub] - # This means that it didn't get to extract from `g`, let's also test - # with extraction by starting with an empty model + @test model.expressions == [_extract_subexpression_expected_sub()] + return +end +# The subexpression `x^3` is extracted from within `g` itself. +function test_extract_subexpression_objective() + _, _, _, g, expected_g = _extract_subexpression_shared_functions() model = Nonlinear.Model() - MOI.Nonlinear.set_objective(model, g) + Nonlinear.set_objective(model, g) @test model.objective == expected_g - @test model.expressions == [expected_sub] - # Test that the objective function gets rewritten as we reuse `h` - # Also test that we don't change the parents in the stack of `h` - # by creating a long stack + @test model.expressions == [_extract_subexpression_expected_sub()] + return +end + +# Test that the objective function gets rewritten as we reuse `h`. Also test +# that we don't change the parents in the stack of `h` by creating a long +# stack. +function test_extract_subexpression_rewrite_objective() + x, _, h, g, _ = _extract_subexpression_shared_functions() + model = Nonlinear.Model() + Nonlinear.set_objective(model, g) prod = MOI.ScalarNonlinearFunction(:*, [h, x]) sum = MOI.ScalarNonlinearFunction(:*, [x, x, x, x, prod]) expr = Nonlinear.parse_expression(model, sum) @@ -1561,7 +1611,7 @@ function test_extract_subexpression() Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, 1), ] @test model.expressions == [ - expected_sub, + _extract_subexpression_expected_sub(), Nonlinear.Expression( [ Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, -1), From d4a380a19dfccb196ac15b6f3283f87ca2b6a221 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 21 Oct 2026 16:55:52 +0200 Subject: [PATCH 6/7] Add docstrings --- src/Nonlinear/parse.jl | 69 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 68 insertions(+), 1 deletion(-) diff --git a/src/Nonlinear/parse.jl b/src/Nonlinear/parse.jl index 4e73eb33ba..9d7f3d4f24 100644 --- a/src/Nonlinear/parse.jl +++ b/src/Nonlinear/parse.jl @@ -36,13 +36,37 @@ function parse_expression(::Model, ::Expression, x::Any, ::Int) ) end +""" + _extract_subexpression!(expr::Expression, root::Int) + +Prepare the extraction of the subtree rooted at `expr.nodes[root]` into a +separate expression. Return `(I, V)` where `I` is the range of `expr.nodes` +and `V` is the range of `expr.values` (or `nothing` if there are no values) +corresponding to the subtree. + +Because the nodes are stored in depth-first order, the subtree is the +contiguous range of nodes `root:(first_out-1)` where `first_out` is the first +node after `root` whose parent is before `root`. Its values are contiguous for +the same reason. + +This function only renumbers the nodes, it does not delete anything: + + * the nodes in `I` are rewritten so that their `parent` (and the `index` of + their `NODE_VALUE`) are relative to the start of `I` (and `V`), so that + `Expression(expr.nodes[I], expr.values[V])` is a valid expression with + `root` as its root; + * the nodes after `I` are rewritten as if the nodes `I[2:end]` and the values + `V` were already deleted, `root` being kept to be replaced by a + `NODE_SUBEXPRESSION`. +""" function _extract_subexpression!(expr::Expression, root::Int) n = length(expr.nodes) - # The whole subexpression is continuous in the tape first_out = first_value = last_value = nothing + # Renumber the nodes of the subtree so that they are relative to `root`. for i in root:n node = expr.nodes[i] if i != root && node.parent < root + # We have left the subtree rooted at `root`. first_out = i break end @@ -69,6 +93,7 @@ function _extract_subexpression!(expr::Expression, root::Int) else V = first_value:last_value end + # Shift the nodes after the subtree as if `I[2:end]` and `V` were deleted. if !isnothing(first_out) for i in (last(I)+1):n node = expr.nodes[i] @@ -78,6 +103,9 @@ function _extract_subexpression!(expr::Expression, root::Int) index -= length(V) end parent = node.parent + # A parent cannot be in `root+1:last(I)` since this node is not in + # the subtree. If `parent == root`, it is now the parent of the + # `NODE_SUBEXPRESSION` so it does not change. if parent > root @assert parent > last(I) parent -= length(I) - 1 @@ -88,6 +116,24 @@ function _extract_subexpression!(expr::Expression, root::Int) return I, V end +""" + _extract_subexpression!(data::Model, expr::Expression, root::Int) + +Move the subtree rooted at `expr.nodes[root]` to a new expression in +`data.expressions` and replace it in `expr` by a single `NODE_SUBEXPRESSION` +node pointing to it. Return `(index, I)` where `index` is the +[`ExpressionIndex`](@ref) of the new subexpression and `I` is the range of +nodes that the subtree used to occupy in `expr`. As `root` is kept, the +nodes after `I` are shifted by `length(I) - 1`, the caller is responsible for +updating any node index into `expr` that it stores. + +This is called when the same `MOI.ScalarNonlinearFunction` is parsed a second +time: its first occurrence was parsed inline into `expr` and is now replaced by +a subexpression that both occurrences can share. + +This costs `O(length(expr.nodes) + length(expr.values))` as we need to shift +the nodes after the subtree and delete the nodes and values of the subtree. +""" function _extract_subexpression!(data::Model, expr::Expression, root::Int) parent = expr.nodes[root].parent I, V = _extract_subexpression!(expr, root) @@ -95,6 +141,9 @@ function _extract_subexpression!(data::Model, expr::Expression, root::Int) Expression(expr.nodes[I], isnothing(V) ? Float64[] : expr.values[V]) push!(data.expressions, subexpr) index = ExpressionIndex(length(data.expressions)) + # We are replacing the subtree by a subexpression so we keep the root node, + # now a `NODE_SUBEXPRESSION`, and remove all the other nodes and values of + # the subtree from the tape. expr.nodes[root] = Node(NODE_SUBEXPRESSION, index.value, parent) if length(I) > 1 deleteat!(expr.nodes, I[2:end]) @@ -119,12 +168,23 @@ function parse_expression( parent_node, arg = pop!(stack) if arg isa MOI.ScalarNonlinearFunction if haskey(data.cache, arg) + # `arg` was already parsed, we reuse it as a subexpression. _expr_index, _node = data.cache[arg] if _node == 0 + # It was already extracted into a subexpression. subexpr = ExpressionIndex(_expr_index) else + # It was parsed inline at node `_node` of `_expr`. This is + # its second occurrence so we extract it into a + # subexpression. This shifts the nodes of `_expr` after + # `I` so we need to update every node index into `_expr` + # that we store: those in `stack` (if `_expr` is the one + # we are parsing into) and those in `data.cache`. _expr = data.cache_expressions[_expr_index] subexpr, I = _extract_subexpression!(data, _expr, _node) + # The nodes of the subtree were moved to the new + # subexpression which may itself contain functions that + # are cached so we register it in `cache_expressions`. push!( data.cache_expressions, data.expressions[subexpr.value], @@ -144,6 +204,11 @@ function parse_expression( end end end + # Cached functions in `I[2:end]` now point into the new + # subexpression and those after `I` are shifted. + # Note that this loops over all the cached functions, so + # extracting `k` subexpressions costs + # `O(k * length(data.cache))`. for (key, (__expr_index, __node)) in data.cache if __node != 0 __expr = data.cache_expressions[__expr_index] @@ -169,6 +234,8 @@ function parse_expression( arg, parent_node, ) + # Cache the node of `arg` in `expr` so that we can extract it + # if we encounter it again. if expr_index == 0 push!(data.cache_expressions, expr) expr_index = length(data.cache_expressions) From ce5f4b5506f3411b9cbf96f53f13de15eb9dbde7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 21 Oct 2026 17:24:14 +0200 Subject: [PATCH 7/7] Deferred removal of first occurence of subexpressions --- src/Nonlinear/model.jl | 4 + src/Nonlinear/parse.jl | 281 ++++++++++++++----------------- src/Nonlinear/types.jl | 11 +- test/Nonlinear/test_Nonlinear.jl | 89 ++++++++++ 4 files changed, 231 insertions(+), 154 deletions(-) diff --git a/src/Nonlinear/model.jl b/src/Nonlinear/model.jl index b570a0dcc3..8db2e34d9c 100644 --- a/src/Nonlinear/model.jl +++ b/src/Nonlinear/model.jl @@ -11,6 +11,9 @@ function MOI.empty!(model::Model) empty!(model.parameters) model.operators = OperatorRegistry() model.last_constraint_index = 0 + empty!(model.cache) + empty!(model.cache_expressions) + empty!(model.cache_pending) return end @@ -322,6 +325,7 @@ function evaluate( expr::Expression; evaluated_expressions = Dict{Int,Float64}(), ) + MOI.Utilities.final_touch(model) storage = zeros(length(expr.nodes)) adj = adjacency_matrix(expr.nodes) children_arr = SparseArrays.rowvals(adj) diff --git a/src/Nonlinear/parse.jl b/src/Nonlinear/parse.jl index 9d7f3d4f24..348056cbe9 100644 --- a/src/Nonlinear/parse.jl +++ b/src/Nonlinear/parse.jl @@ -37,123 +37,135 @@ function parse_expression(::Model, ::Expression, x::Any, ::Int) end """ - _extract_subexpression!(expr::Expression, root::Int) - -Prepare the extraction of the subtree rooted at `expr.nodes[root]` into a -separate expression. Return `(I, V)` where `I` is the range of `expr.nodes` -and `V` is the range of `expr.values` (or `nothing` if there are no values) -corresponding to the subtree. - -Because the nodes are stored in depth-first order, the subtree is the -contiguous range of nodes `root:(first_out-1)` where `first_out` is the first -node after `root` whose parent is before `root`. Its values are contiguous for -the same reason. - -This function only renumbers the nodes, it does not delete anything: - - * the nodes in `I` are rewritten so that their `parent` (and the `index` of - their `NODE_VALUE`) are relative to the start of `I` (and `V`), so that - `Expression(expr.nodes[I], expr.values[V])` is a valid expression with - `root` as its root; - * the nodes after `I` are rewritten as if the nodes `I[2:end]` and the values - `V` were already deleted, `root` being kept to be replaced by a - `NODE_SUBEXPRESSION`. + MOI.Utilities.final_touch(model::Model, index_map = nothing) + +Extract the subexpressions that were detected by [`parse_expression`](@ref) +since the last call to `final_touch`. + +When a `MOI.ScalarNonlinearFunction` is parsed a second time, its first +occurrence, which was parsed inline into some expression, has to be moved into +a new subexpression so that both occurrences can share it. Doing this +immediately would shift the nodes of that expression, so every node index +stored in the cache (and in the stack of the current parse) would need to be +updated for each extraction, which is quadratic. Instead, `parse_expression` +reserves an [`ExpressionIndex`](@ref) for the subexpression, with an empty +placeholder [`Expression`](@ref), and records the node to extract in +`model.cache_pending`. This function then rewrites each affected expression +in a single pass, and remaps the cache in a single pass. + +The subexpressions are invalid until this function is called. It is called by +the [`Evaluator`](@ref) constructors and by [`evaluate`](@ref), and it is a +no-op if there is nothing to extract. """ -function _extract_subexpression!(expr::Expression, root::Int) +function MOI.Utilities.final_touch(model::Model, ::Any = nothing) + pending = model.cache_pending + if isempty(pending) + return + end + # Group the nodes to extract by expression, sorted by node index. + sort!(pending) + maps = Dict{Int,Tuple{Vector{Int},Vector{Int}}}() + first = 1 + while first <= length(pending) + tape = pending[first][1] + last = first + while last < length(pending) && pending[last+1][1] == tape + last += 1 + end + maps[tape] = _extract_subexpressions!(model, tape, pending, first, last) + first = last + 1 + end + empty!(pending) + # Remap the cached nodes of the expressions that were rewritten. + for (key, (tape, node)) in model.cache + if node != 0 && haskey(maps, tape) + new_tape, new_node = maps[tape] + model.cache[key] = (new_tape[node], new_node[node]) + end + end + return +end + +""" + _extract_subexpressions!(model::Model, tape::Int, pending, first, last) + +Extract the nodes `pending[i][2]` of `model.cache_expressions[tape]` into the +subexpressions `model.expressions[pending[i][3]]` for `i in first:last`. The +nodes `pending[first:last]` must be sorted. + +The nodes are stored in depth-first order so a node is visited after its +parent. Each node is therefore appended to the same expression as its parent, +unless it is to be extracted, in which case its parent gets a +`NODE_SUBEXPRESSION` node instead and the node becomes the root of its +subexpression. Nested extractions are handled in the same way. + +The nodes that stay in `model.cache_expressions[tape]` can only move to a +smaller index so we rewrite it in place. + +Return `(new_tape, new_node)` such that node `i` of the original expression is +now node `new_node[i]` of `model.cache_expressions[new_tape[i]]`. +""" +function _extract_subexpressions!( + model::Model, + tape::Int, + pending::Vector{Tuple{Int,Int,Int}}, + first::Int, + last::Int, +) + expr = model.cache_expressions[tape] n = length(expr.nodes) - first_out = first_value = last_value = nothing - # Renumber the nodes of the subtree so that they are relative to `root`. - for i in root:n + new_tape = Vector{Int}(undef, n) + new_node = Vector{Int}(undef, n) + # Number of nodes and values kept in `expr` so far. + num_nodes = num_values = 0 + for i in 1:n node = expr.nodes[i] - if i != root && node.parent < root - # We have left the subtree rooted at `root`. - first_out = i - break + if node.parent == -1 + parent_tape, parent = tape, -1 + else + parent_tape, parent = new_tape[node.parent], new_node[node.parent] end - index = node.index - if node.type == NODE_VALUE - if isnothing(first_value) - first_value = node.index - last_value = first_value + if first <= last && pending[first][2] == i + subexpr = pending[first][3] + first += 1 + # Replace the node by a `NODE_SUBEXPRESSION` in its parent's + # expression and continue with the node as root of `subexpr`. + subexpr_node = Node(NODE_SUBEXPRESSION, subexpr, parent) + if parent_tape == tape + num_nodes += 1 + expr.nodes[num_nodes] = subexpr_node else - last_value = node.index + push!(model.cache_expressions[parent_tape].nodes, subexpr_node) end - index -= first_value - 1 + # `subexpr` may contain cached nodes so it needs to be in + # `cache_expressions`. + push!(model.cache_expressions, model.expressions[subexpr]) + parent_tape, parent = length(model.cache_expressions), -1 end - expr.nodes[i] = - Node(node.type, index, i == root ? -1 : node.parent - root + 1) - end - if isnothing(first_out) - I = root:n - else - I = root:(first_out-1) - end - if isnothing(first_value) - V = nothing - else - V = first_value:last_value - end - # Shift the nodes after the subtree as if `I[2:end]` and `V` were deleted. - if !isnothing(first_out) - for i in (last(I)+1):n - node = expr.nodes[i] - index = node.index - if node.type == NODE_VALUE && !isnothing(V) - @assert index >= last(V) - index -= length(V) + new_tape[i] = parent_tape + index = node.index + if parent_tape == tape + if node.type == NODE_VALUE + num_values += 1 + expr.values[num_values] = expr.values[index] + index = num_values end - parent = node.parent - # A parent cannot be in `root+1:last(I)` since this node is not in - # the subtree. If `parent == root`, it is now the parent of the - # `NODE_SUBEXPRESSION` so it does not change. - if parent > root - @assert parent > last(I) - parent -= length(I) - 1 + num_nodes += 1 + expr.nodes[num_nodes] = Node(node.type, index, parent) + new_node[i] = num_nodes + else + target = model.cache_expressions[parent_tape] + if node.type == NODE_VALUE + push!(target.values, expr.values[index]) + index = length(target.values) end - expr.nodes[i] = Node(node.type, index, parent) + push!(target.nodes, Node(node.type, index, parent)) + new_node[i] = length(target.nodes) end end - return I, V -end - -""" - _extract_subexpression!(data::Model, expr::Expression, root::Int) - -Move the subtree rooted at `expr.nodes[root]` to a new expression in -`data.expressions` and replace it in `expr` by a single `NODE_SUBEXPRESSION` -node pointing to it. Return `(index, I)` where `index` is the -[`ExpressionIndex`](@ref) of the new subexpression and `I` is the range of -nodes that the subtree used to occupy in `expr`. As `root` is kept, the -nodes after `I` are shifted by `length(I) - 1`, the caller is responsible for -updating any node index into `expr` that it stores. - -This is called when the same `MOI.ScalarNonlinearFunction` is parsed a second -time: its first occurrence was parsed inline into `expr` and is now replaced by -a subexpression that both occurrences can share. - -This costs `O(length(expr.nodes) + length(expr.values))` as we need to shift -the nodes after the subtree and delete the nodes and values of the subtree. -""" -function _extract_subexpression!(data::Model, expr::Expression, root::Int) - parent = expr.nodes[root].parent - I, V = _extract_subexpression!(expr, root) - subexpr = - Expression(expr.nodes[I], isnothing(V) ? Float64[] : expr.values[V]) - push!(data.expressions, subexpr) - index = ExpressionIndex(length(data.expressions)) - # We are replacing the subtree by a subexpression so we keep the root node, - # now a `NODE_SUBEXPRESSION`, and remove all the other nodes and values of - # the subtree from the tape. - expr.nodes[root] = Node(NODE_SUBEXPRESSION, index.value, parent) - if length(I) > 1 - deleteat!(expr.nodes, I[2:end]) - if !isnothing(V) - deleteat!(expr.values, V) - end - else - @assert isnothing(V) - end - return index, I + resize!(expr.nodes, num_nodes) + resize!(expr.values, num_values) + return new_tape, new_node end function parse_expression( @@ -174,55 +186,18 @@ function parse_expression( # It was already extracted into a subexpression. subexpr = ExpressionIndex(_expr_index) else - # It was parsed inline at node `_node` of `_expr`. This is - # its second occurrence so we extract it into a - # subexpression. This shifts the nodes of `_expr` after - # `I` so we need to update every node index into `_expr` - # that we store: those in `stack` (if `_expr` is the one - # we are parsing into) and those in `data.cache`. - _expr = data.cache_expressions[_expr_index] - subexpr, I = _extract_subexpression!(data, _expr, _node) - # The nodes of the subtree were moved to the new - # subexpression which may itself contain functions that - # are cached so we register it in `cache_expressions`. + # It was parsed inline at node `_node` of + # `data.cache_expressions[_expr_index]`. This is its second + # occurrence so it needs to be extracted into a + # subexpression. We reserve its index and defer the + # extraction to `MOI.Utilities.final_touch` so that no + # node is moved in the meantime. + push!(data.expressions, Expression()) + subexpr = ExpressionIndex(length(data.expressions)) push!( - data.cache_expressions, - data.expressions[subexpr.value], + data.cache_pending, + (_expr_index, _node, subexpr.value), ) - subexpr_index = length(data.cache_expressions) - if expr === _expr - if parent_node > first(I) - @assert parent_node > last(I) - parent_node -= length(I) - 1 - end - for i in eachindex(stack) - _parent_node = stack[i][1] - if _parent_node > first(I) - @assert _parent_node > last(I) - stack[i] = - (_parent_node - length(I) + 1, stack[i][2]) - end - end - end - # Cached functions in `I[2:end]` now point into the new - # subexpression and those after `I` are shifted. - # Note that this loops over all the cached functions, so - # extracting `k` subexpressions costs - # `O(k * length(data.cache))`. - for (key, (__expr_index, __node)) in data.cache - if __node != 0 - __expr = data.cache_expressions[__expr_index] - if _expr === __expr && __node > first(I) - if __node <= last(I) - data.cache[key] = - (subexpr_index, __node - first(I) + 1) - else - data.cache[key] = - (__expr_index, __node - length(I) + 1) - end - end - end - end data.cache[arg] = (subexpr.value, 0) end parse_expression(data, expr, subexpr, parent_node) diff --git a/src/Nonlinear/types.jl b/src/Nonlinear/types.jl index 161bfac51a..1b2ba5964c 100644 --- a/src/Nonlinear/types.jl +++ b/src/Nonlinear/types.jl @@ -168,12 +168,17 @@ mutable struct Model last_constraint_index::Int64 # Cache values are (expression index, node index). A zero node index # identifies an extracted subexpression in `expressions`; otherwise the - # expression index refers to `cache_expressions`. + # expression index refers to `cache_expressions`. The node indices refer + # to the expressions as they are before the pending extractions are done + # by `MOI.Utilities.final_touch`. # We use a separate vector `cache_expressions` instead of storing # a `Tuple{Expression,Int}` in the `Dict` as creating such tuple # would trigger a heap allocation for each expression cache::Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}} cache_expressions::Vector{Expression} + # Extractions deferred to `MOI.Utilities.final_touch`. Each element is + # (index in `cache_expressions`, node index, index in `expressions`). + cache_pending::Vector{Tuple{Int,Int,Int}} function Model() return new( nothing, @@ -184,6 +189,7 @@ mutable struct Model 0, Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}}(), Expression[], + Tuple{Int,Int,Int}[], ) end end @@ -283,6 +289,7 @@ supported feature is `:ExprGraph`. struct ExprGraphOnly <: AbstractAutomaticDifferentiation end function Evaluator(model::Model, ::ExprGraphOnly, ::Vector{MOI.VariableIndex}) + MOI.Utilities.final_touch(model) return Evaluator(model) end @@ -300,6 +307,7 @@ function Evaluator( ::SparseReverseMode, ordered_variables::Vector{MOI.VariableIndex}, ) + MOI.Utilities.final_touch(model) return Evaluator(model, ReverseAD.NLPEvaluator(model, ordered_variables)) end @@ -316,5 +324,6 @@ function Evaluator( ::SymbolicMode, ordered_variables::Vector{MOI.VariableIndex}, ) + MOI.Utilities.final_touch(model) return Evaluator(model, SymbolicAD.Evaluator(model, ordered_variables)) end diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index 19037a17a8..f3d15280f4 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1455,6 +1455,7 @@ function test_extract_subexpression() sub = MOI.ScalarNonlinearFunction(:^, Any[x, 3]) f = MOI.ScalarNonlinearFunction(:+, Any[sub, sub]) expr = Nonlinear.parse_expression(model, f) + MOI.Utilities.final_touch(model) @test expr == Nonlinear.Expression( [ Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), @@ -1477,6 +1478,7 @@ function _extract_subexpression_nested_model() outer = MOI.ScalarNonlinearFunction(:+, Any[inner, x]) repeated = MOI.ScalarNonlinearFunction(:*, Any[outer, outer]) expr = Nonlinear.parse_expression(model, repeated) + MOI.Utilities.final_touch(model) return model, expr, x, inner, outer, repeated end @@ -1512,6 +1514,7 @@ function test_extract_subexpression_nested_after_move() model, _, x, inner, outer, repeated = _extract_subexpression_nested_model() Nonlinear.parse_expression(model, outer) reused = Nonlinear.parse_expression(model, inner) + MOI.Utilities.final_touch(model) @test model.cache[inner] == (2, 0) @test reused.nodes == [Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, -1)] @test model.expressions[1].nodes[2] == @@ -1536,6 +1539,7 @@ function test_extract_subexpression_single_node() empty_sum = MOI.ScalarNonlinearFunction(:+, Any[]) repeated = MOI.ScalarNonlinearFunction(:+, Any[empty_sum, empty_sum]) expr = Nonlinear.parse_expression(model, repeated) + MOI.Utilities.final_touch(model) @test model.expressions == [ Nonlinear.Expression( [Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1)], @@ -1579,6 +1583,7 @@ function test_extract_subexpression_already_extracted() f = MOI.ScalarNonlinearFunction(:+, Any[sub, sub]) Nonlinear.parse_expression(model, f) expr = Nonlinear.parse_expression(model, g) + MOI.Utilities.final_touch(model) @test expr == expected_g @test model.expressions == [_extract_subexpression_expected_sub()] return @@ -1589,6 +1594,7 @@ function test_extract_subexpression_objective() _, _, _, g, expected_g = _extract_subexpression_shared_functions() model = Nonlinear.Model() Nonlinear.set_objective(model, g) + MOI.Utilities.final_touch(model) @test model.objective == expected_g @test model.expressions == [_extract_subexpression_expected_sub()] return @@ -1604,6 +1610,7 @@ function test_extract_subexpression_rewrite_objective() prod = MOI.ScalarNonlinearFunction(:*, [h, x]) sum = MOI.ScalarNonlinearFunction(:*, [x, x, x, x, prod]) expr = Nonlinear.parse_expression(model, sum) + MOI.Utilities.final_touch(model) @test isempty(model.objective.values) @test model.objective.nodes == [ Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), @@ -1636,6 +1643,88 @@ function test_extract_subexpression_rewrite_objective() return end +# Both `inner` and `outer`, which contains `inner`, are extracted by the same +# call to `final_touch`, so `inner` is moved to the expression of `outer` +# before being extracted. +function test_extract_subexpression_nested_same_final_touch() + model = Nonlinear.Model() + x = MOI.VariableIndex(1) + inner = MOI.ScalarNonlinearFunction(:^, Any[x, 2]) + outer = MOI.ScalarNonlinearFunction(:+, Any[inner, 3]) + repeated = MOI.ScalarNonlinearFunction(:*, Any[outer, outer]) + expr = Nonlinear.parse_expression(model, repeated) + reused = Nonlinear.parse_expression(model, inner) + MOI.Utilities.final_touch(model) + @test expr == Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 3, -1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + ], + Float64[], + ) + @test reused.nodes == [Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, -1)] + @test model.expressions == [ + Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 1, -1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 2, 1), + Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 1), + ], + [3.0], + ), + Nonlinear.Expression( + [ + Nonlinear.Node(Nonlinear.NODE_CALL_MULTIVARIATE, 4, -1), + Nonlinear.Node(Nonlinear.NODE_MOI_VARIABLE, 1, 1), + Nonlinear.Node(Nonlinear.NODE_VALUE, 1, 1), + ], + [2.0], + ), + ] + @test model.cache[outer] == (1, 0) + @test model.cache[inner] == (2, 0) + @test isempty(model.cache_pending) + # Calling `final_touch` again does nothing. + MOI.Utilities.final_touch(model) + @test length(model.expressions) == 2 + return +end + +# The extractions are done by `evaluate` and by the `Evaluator` without +# calling `final_touch` explicitly. +function test_extract_subexpression_implicit_final_touch() + x = MOI.VariableIndex(1) + sub = MOI.ScalarNonlinearFunction(:sin, Any[x]) + f = MOI.ScalarNonlinearFunction(:+, Any[sub, sub]) + model = Nonlinear.Model() + expr = Nonlinear.parse_expression(model, f) + @test model.expressions == [Nonlinear.Expression()] + @test Nonlinear.evaluate(Dict(x => 0.5), model, expr) ≈ 2 * sin(0.5) + @test isempty(model.cache_pending) + model = Nonlinear.Model() + Nonlinear.set_objective(model, f) + evaluator = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(evaluator, [:Grad]) + @test MOI.eval_objective(evaluator, [0.5]) ≈ 2 * sin(0.5) + gradient = zeros(1) + MOI.eval_objective_gradient(evaluator, gradient, [0.5]) + @test gradient[1] ≈ 2 * cos(0.5) + return +end + +function test_extract_subexpression_empty() + x = MOI.VariableIndex(1) + sub = MOI.ScalarNonlinearFunction(:sin, Any[x]) + model = Nonlinear.Model() + Nonlinear.set_objective(model, MOI.ScalarNonlinearFunction(:+, [sub, sub])) + MOI.empty!(model) + @test isempty(model.cache) + @test isempty(model.cache_expressions) + @test isempty(model.cache_pending) + return +end + end # TestNonlinear TestNonlinear.runtests()