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 05fd684429..348056cbe9 100644 --- a/src/Nonlinear/parse.jl +++ b/src/Nonlinear/parse.jl @@ -36,6 +36,138 @@ function parse_expression(::Model, ::Expression, x::Any, ::Int) ) end +""" + 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 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) + 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 node.parent == -1 + parent_tape, parent = tape, -1 + else + parent_tape, parent = new_tape[node.parent], new_node[node.parent] + end + 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 + push!(model.cache_expressions[parent_tape].nodes, subexpr_node) + end + # `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 + 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 + 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 + push!(target.nodes, Node(node.type, index, parent)) + new_node[i] = length(target.nodes) + end + end + resize!(expr.nodes, num_nodes) + resize!(expr.values, num_values) + return new_tape, new_node +end + function parse_expression( data::Model, expr::Expression, @@ -43,10 +175,48 @@ 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 - _parse_without_recursion_inner(stack, data, expr, arg, parent_node) + 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 + # `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_pending, + (_expr_index, _node, subexpr.value), + ) + data.cache[arg] = (subexpr.value, 0) + end + parse_expression(data, expr, subexpr, parent_node) + else + _parse_without_recursion_inner( + stack, + data, + expr, + 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) + end + data.cache[arg] = (expr_index, length(expr.nodes)) + end else # We can use recursion here, because ScalarNonlinearFunction only # occur in other ScalarNonlinearFunction. @@ -82,7 +252,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..1b2ba5964c 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,19 @@ mutable struct Model operators::OperatorRegistry # This is a private field, used only to increment the ConstraintIndex. 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`. 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, @@ -173,6 +187,9 @@ mutable struct Model Float64[], OperatorRegistry(), 0, + Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}}(), + Expression[], + Tuple{Int,Int,Int}[], ) end end @@ -272,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 @@ -289,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 @@ -305,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 e335c2ced5..f3d15280f4 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1438,6 +1438,293 @@ 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) + 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), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + Nonlinear.Node(Nonlinear.NODE_SUBEXPRESSION, 1, 1), + ], + Float64[], + ) + @test model.expressions == [_extract_subexpression_expected_sub()] + @test model.cache[sub] == (1, 0) + return +end + +# 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) + MOI.Utilities.final_touch(model) + 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] + @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), + ] + return +end + +# 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. +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] == + 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)] + return +end + +# 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]) + expr = Nonlinear.parse_expression(model, repeated) + MOI.Utilities.final_touch(model) + @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), + ] + return +end + +# 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]) + 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], + ) + 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) + MOI.Utilities.final_touch(model) + @test expr == expected_g + @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() + Nonlinear.set_objective(model, g) + MOI.Utilities.final_touch(model) + @test model.objective == expected_g + @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) + MOI.Utilities.final_touch(model) + @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 == [ + _extract_subexpression_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 + +# 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()