From 2caac75ba06e756509a6bb64096a33ce471f46fb 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/7] 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 98a37e60aa8..aa48f6ebc44 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 45f3ac3496d..5afe688c876 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -570,6 +570,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 @@ -581,6 +585,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 @@ -590,7 +598,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 7b96a8403ca719ba1879035091840aef4bb01eb7 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/7] 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 5afe688c876..0d0e8dfd3fc 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -571,8 +571,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}}[] @@ -585,8 +586,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)) @@ -598,9 +599,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 a80bd272d6b65b9f4fd3e36713cc17d38eab0150 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/7] 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 0d0e8dfd3fc..88efad035f7 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -569,10 +569,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)) @@ -580,13 +579,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 @@ -599,13 +600,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 73b37a7f82332925449bad4acbf2d1a0dc003872 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/7] 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 6e6624eb13b13bb2f2cfe98d2b46e16f3b8edf07 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/7] 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 88efad035f7..e8ad0f3b45b 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -569,7 +569,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] From b7b160dfd4d7b7945d39f7c55295fe3061399761 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 26 Feb 2026 20:48:18 +0100 Subject: [PATCH 6/7] Apply suggestion from @blegat --- src/nlp_expr.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl index e8ad0f3b45b..fcf5da3aeca 100644 --- a/src/nlp_expr.jl +++ b/src/nlp_expr.jl @@ -583,7 +583,7 @@ function moi_function( 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]) + ret.args[i] = moi_function(f.args[i], model) else ret.args[i] = moi_function(f.args[i]) end From c669881b0201eb264b1e2a5826915976c928aedd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 26 Feb 2026 20:54:09 +0100 Subject: [PATCH 7/7] Fix --- src/objective.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/objective.jl b/src/objective.jl index c9dcbbde3ac..6e493045889 100644 --- a/src/objective.jl +++ b/src/objective.jl @@ -277,7 +277,7 @@ end function set_objective_function(model::GenericModel, func::AbstractJuMPScalar) check_belongs_to_model(func, model) - set_objective_function(model, moi_function(func)) + set_objective_function(model, moi_function(func, model)) return end @@ -296,7 +296,7 @@ function set_objective_function( for f in func check_belongs_to_model(f, model) end - set_objective_function(model, moi_function(func)) + set_objective_function(model, moi_function(func, model)) return end