From d8e9515f8836d1102b74f64a5bc419ddcba5250f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 22 Jul 2025 08:57:53 +0200 Subject: [PATCH 1/5] Cache subexpressions when building MOI.ScalarNonlinearExpression --- src/JuMP.jl | 3 +++ src/nlp_expr.jl | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/src/JuMP.jl b/src/JuMP.jl index 823896ff4f3..27e06d6313e 100644 --- a/src/JuMP.jl +++ b/src/JuMP.jl @@ -138,6 +138,8 @@ mutable struct GenericModel{T<:Real} <: AbstractModel # A dictionary to store timing information from the JuMP macros. enable_macro_timing::Bool macro_times::Dict{Tuple{LineNumberNode,String},Float64} + # We use `Any` as key because we haven't defined `GenericNonlinearExpr` yet + subexpressions::Dict{Any,MOI.ScalarNonlinearFunction} end value_type(::Type{GenericModel{T}}) where {T} = T @@ -251,6 +253,7 @@ function direct_generic_model( Dict{Any,MOI.ConstraintIndex}(), false, Dict{Tuple{LineNumberNode,String},Float64}(), + Dict{Any,MOI.ScalarNonlinearFunction}(), ) end diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl index a3d777b8923..53080f83fa8 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -553,6 +553,10 @@ end moi_function(x::Number) = x function moi_function(f::GenericNonlinearExpr{V}) where {V} + model = owner_model(f) + if haskey(model.subexpressions, f) + return model.subexpressions[f] + end ret = MOI.ScalarNonlinearFunction(f.head, similar(f.args)) stack = Tuple{MOI.ScalarNonlinearFunction,Int,GenericNonlinearExpr{V}}[] for i in length(f.args):-1:1 @@ -564,6 +568,10 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V} end while !isempty(stack) parent, i, arg = pop!(stack) + if haskey(model.subexpressions, arg) + parent.args[i] = model.subexpressions[arg] + continue + end child = MOI.ScalarNonlinearFunction(arg.head, similar(arg.args)) parent.args[i] = child for j in length(arg.args):-1:1 @@ -573,7 +581,9 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V} child.args[j] = moi_function(arg.args[j]) end end + model.subexpressions[arg] = child end + model.subexpressions[f] = ret return ret end From ee3be7dfae3589dbeb81be172f4f0249a09219f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 22 Jul 2025 10:22:06 +0200 Subject: [PATCH 2/5] Fix --- src/nlp_expr.jl | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl index 53080f83fa8..fe2270858d3 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -554,8 +554,9 @@ moi_function(x::Number) = x function moi_function(f::GenericNonlinearExpr{V}) where {V} model = owner_model(f) - if haskey(model.subexpressions, f) - return model.subexpressions[f] + cache = isnothing(model) ? nothing : model.subexpressions + if !isnothing(cache) && haskey(cache, f) + return cache[f] end ret = MOI.ScalarNonlinearFunction(f.head, similar(f.args)) stack = Tuple{MOI.ScalarNonlinearFunction,Int,GenericNonlinearExpr{V}}[] @@ -568,8 +569,8 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V} end while !isempty(stack) parent, i, arg = pop!(stack) - if haskey(model.subexpressions, arg) - parent.args[i] = model.subexpressions[arg] + if !isnothing(cache) && haskey(cache, arg) + parent.args[i] = cache[arg] continue end child = MOI.ScalarNonlinearFunction(arg.head, similar(arg.args)) @@ -581,9 +582,13 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V} child.args[j] = moi_function(arg.args[j]) end end - model.subexpressions[arg] = child + if !isnothing(cache) + cache[arg] = child + end + end + if !isnothing(cache) + cache[f] = ret end - model.subexpressions[f] = ret return ret end From 58fc02505d20356768efc236c9f8f9ef6e3af30b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 22 Jul 2025 15:28:13 +0200 Subject: [PATCH 3/5] Merge check_belongs_to_model with moi_function --- src/constraints.jl | 12 ++++++++++-- src/nlp_expr.jl | 19 ++++++++----------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/src/constraints.jl b/src/constraints.jl index 157907ccd5e..c75a4683c48 100644 --- a/src/constraints.jl +++ b/src/constraints.jl @@ -760,6 +760,10 @@ function moi_function(constraint::AbstractConstraint) return moi_function(jump_function(constraint)) end +function moi_function(constraint::AbstractConstraint, model) + return moi_function(jump_function(constraint), model) +end + """ moi_set(constraint::AbstractConstraint) @@ -1016,6 +1020,11 @@ function _moi_add_constraint( return MOI.add_constraint(model, f, s) end +function moi_function(f, model) + check_belongs_to_model(f, model) + return moi_function(f) +end + """ add_constraint( model::GenericModel, @@ -1032,10 +1041,9 @@ function add_constraint( name::String = "", ) con = model_convert(model, con) + func, set = moi_function(con, model), moi_set(con) # The type of backend(model) is unknown so we directly redirect to another # function. - check_belongs_to_model(con, model) - func, set = moi_function(con), moi_set(con) cindex = _moi_add_constraint( backend(model), func, diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl index fe2270858d3..a7c1ec21a62 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -552,10 +552,9 @@ end moi_function(x::Number) = x -function moi_function(f::GenericNonlinearExpr{V}) where {V} - model = owner_model(f) - cache = isnothing(model) ? nothing : model.subexpressions - if !isnothing(cache) && haskey(cache, f) +function moi_function(f::GenericNonlinearExpr{V}, model::JuMP.GenericModel) where {V} + cache = model.subexpressions + if haskey(cache, f) return cache[f] end ret = MOI.ScalarNonlinearFunction(f.head, similar(f.args)) @@ -563,13 +562,15 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V} for i in length(f.args):-1:1 if f.args[i] isa GenericNonlinearExpr{V} push!(stack, (ret, i, f.args[i])) + elseif f.args[i] isa AbstractJuMPScalar + ret.args[i] = moi_function(model, f.args[i]) else ret.args[i] = moi_function(f.args[i]) end end while !isempty(stack) parent, i, arg = pop!(stack) - if !isnothing(cache) && haskey(cache, arg) + if haskey(cache, arg) parent.args[i] = cache[arg] continue end @@ -582,13 +583,9 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V} child.args[j] = moi_function(arg.args[j]) end end - if !isnothing(cache) - cache[arg] = child - end - end - if !isnothing(cache) - cache[f] = ret + cache[arg] = child end + cache[f] = ret return ret end From 4de1d1c66958e57ae15d3625a115b2003bf56b7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 22 Jul 2025 16:03:20 +0200 Subject: [PATCH 4/5] Fix --- src/constraints.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/constraints.jl b/src/constraints.jl index c75a4683c48..1f9561612f5 100644 --- a/src/constraints.jl +++ b/src/constraints.jl @@ -1020,6 +1020,12 @@ function _moi_add_constraint( return MOI.add_constraint(model, f, s) end +function check_belongs_to_model(f::Vector, model) + for func in f + check_belongs_to_model(func, model) + end +end + function moi_function(f, model) check_belongs_to_model(f, model) return moi_function(f) From 0bd8a48751c06f5329bc59e1a409b236724f45b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 22 Jul 2025 16:06:35 +0200 Subject: [PATCH 5/5] Fix format --- src/nlp_expr.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl index a7c1ec21a62..f398bdff7be 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -552,7 +552,10 @@ end moi_function(x::Number) = x -function moi_function(f::GenericNonlinearExpr{V}, model::JuMP.GenericModel) where {V} +function moi_function( + f::GenericNonlinearExpr{V}, + model::JuMP.GenericModel, +) where {V} cache = model.subexpressions if haskey(cache, f) return cache[f]