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] 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()