From 5da6eb42bda45dd587cb7ba1efe0717d330c83c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 3 Nov 2026 11:26:49 +0100 Subject: [PATCH] Visit shared subexpressions once in is_canonical MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `is_canonical(::ScalarNonlinearFunction)` is called for every nonlinear constraint added to a `Utilities.Model`. It was exponential in the depth of the expression graph when subexpressions are shared, as in a Runge-Kutta discretization. Benchmark of `is_canonical` on `f_{k+1} = sin(f_k) * f_k + f_k` with `f_0 = sin(x)`, whose tree has O(3^n) nodes but whose DAG has O(n) nodes: | n | master (μs) | master (bytes) | this PR (μs) | this PR (bytes) | |----|-------------|----------------|--------------|-----------------| | 1 | 0.0597 | 80 | 0.271 | 352 | | 4 | 1.59 | 176 | 1.0 | 944 | | 8 | 131 | 464 | 2.09 | 1376 | | 12 | 10800 | 464 | 3.81 | 1840 | --- src/Utilities/functions.jl | 20 ++++++++++++-------- test/Utilities/test_functions.jl | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/src/Utilities/functions.jl b/src/Utilities/functions.jl index 0d78c7089e..3ae1e49d06 100644 --- a/src/Utilities/functions.jl +++ b/src/Utilities/functions.jl @@ -1003,15 +1003,19 @@ is_canonical(::Union{MOI.VariableIndex,MOI.VectorOfVariables}) = true function is_canonical(f::MOI.ScalarNonlinearFunction) # Don't use recursion here. This gets called for all scalar nonlinear # constraints. - stack = Any[arg for arg in f.args] + # A subexpression may appear several times in the expression graph, we + # visit it only once. + visited = Base.IdSet{MOI.ScalarNonlinearFunction}() + stack = MOI.ScalarNonlinearFunction[f] while !isempty(stack) - arg = pop!(stack) - if arg isa MOI.ScalarNonlinearFunction - for a in arg.args - push!(stack, a) - end - else - if !is_canonical(arg) + g = pop!(stack) + for arg in g.args + if arg isa MOI.ScalarNonlinearFunction + if !(arg in visited) + push!(visited, arg) + push!(stack, arg) + end + elseif !is_canonical(arg) return false end end diff --git a/test/Utilities/test_functions.jl b/test/Utilities/test_functions.jl index b49427bf65..a53e5c8b23 100644 --- a/test/Utilities/test_functions.jl +++ b/test/Utilities/test_functions.jl @@ -2101,6 +2101,23 @@ function test_ScalarNonlinearFunction_is_canonical() return end +function test_is_canonical_ScalarNonlinearFunction_common_subexpressions() + x = MOI.VariableIndex(1) + # The tree has 2^N leaves, the expression graph has N nodes + N = 100 + f = MOI.ScalarNonlinearFunction(:log, Any[1.0*x]) + for _ in 1:N + f = MOI.ScalarNonlinearFunction(:+, Any[f, f]) + end + @test MOI.Utilities.is_canonical(f) + g = MOI.ScalarNonlinearFunction(:log, Any[1.0*x+1.0*x]) + for _ in 1:N + g = MOI.ScalarNonlinearFunction(:+, Any[g, g]) + end + @test !MOI.Utilities.is_canonical(g) + return +end + function test_vector_type() for T in (Int, Float64) @test MOI.Utilities.vector_type(T) == Vector{T}