Convert common sub-functions as common sub-expressions

src/Nonlinear/parse.jl

@@ -36,17 +36,145 @@ 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,
     x::MOI.ScalarNonlinearFunction,
     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)
+                _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)
+                            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, (__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,
+                )
+                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 +210,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

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

test/Nonlinear/test_Nonlinear.jl

@@ -1438,6 +1438,155 @@ 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] == (1, 0)
+
+    # 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)
+    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()
+    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)
+    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()