diff --git a/src/Nonlinear/parse.jl b/src/Nonlinear/parse.jl index 05fd684429..9d7f3d4f24 100644 --- a/src/Nonlinear/parse.jl +++ b/src/Nonlinear/parse.jl @@ -36,6 +36,126 @@ 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) + 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 + 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 + # 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) + 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 + end + expr.nodes[i] = Node(node.type, index, parent) + 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 +end + function parse_expression( data::Model, expr::Expression, @@ -43,10 +163,85 @@ 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 `_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], + ) + 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) + 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 +277,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..161bfac51a 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,14 @@ 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`. + # 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, @@ -173,6 +182,8 @@ mutable struct Model Float64[], OperatorRegistry(), 0, + Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}}(), + Expression[], ) end end diff --git a/test/Nonlinear/test_Nonlinear.jl b/test/Nonlinear/test_Nonlinear.jl index e335c2ced5..19037a17a8 100644 --- a/test/Nonlinear/test_Nonlinear.jl +++ b/test/Nonlinear/test_Nonlinear.jl @@ -1438,6 +1438,204 @@ 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) + @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) + 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) + @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) + @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) + @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) + @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) + @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 + end # TestNonlinear TestNonlinear.runtests()