Macro to update expression in-place #3931

joaquimg · 2025-01-30T08:46:56Z

exp += coef * var
exp -= coef * exp2

Are very common operations for JuMP users.

add_to_expression! is there to solve the performance issue here.

However, should we consider a macro to properly rewrite and keep readability?

@rewrite exp += coef * var

It could also be @update, @operate etc

This pattern is heavily used here: https://github.com/GenXProject/GenX.jl/blob/main/src/model/resources/vre_stor/vre_stor.jl

The text was updated successfully, but these errors were encountered:

odow · 2025-01-30T20:25:49Z

There is already MutableArithmetics.@rewrite. We could make it support this case. Alternatively, we could make this work:

julia> x = 1
1

julia> @expression(model, x += 2)
ERROR: UndefVarError: `x` not defined
Stacktrace:
 [1] macro expansion
   @ ~/.julia/dev/MutableArithmetics/src/rewrite.jl:376 [inlined]
 [2] macro expansion
   @ ~/.julia/packages/JuMP/CU7H5/src/macros.jl:264 [inlined]
 [3] macro expansion
   @ ~/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:86 [inlined]
 [4] macro expansion
   @ ~/.julia/packages/JuMP/CU7H5/src/macros.jl:400 [inlined]
 [5] top-level scope

odow · 2025-01-30T20:26:37Z

But also, add_to_expression! exists and is very explicit about what it does. People should use it.

odow · 2025-01-30T20:31:11Z

This case
https://github.com/GenXProject/GenX.jl/blob/dcf3df01b1c23df1bf139f2a10f07468837dd909/src/model/resources/vre_stor/vre_stor.jl#L130-L138

@expression(EP, ePowerBalance_VRE_STOR[t = 1:T, z = 1:Z], JuMP.AffExpr())
for t in 1:T, z in 1:Z
    if !isempty(resources_in_zone_by_rid(gen_VRE_STOR, z))
        ePowerBalance_VRE_STOR[t, z] += sum(EP[:vP][y, t]
        for y in resources_in_zone_by_rid(gen_VRE_STOR,
            z))
    end
end

should be

@expression(EP, ePowerBalance_VRE_STOR[t = 1:T, z = 1:Z], zero(JuMP.AffExpr))
for z in 1:Z
    resources = resources_in_zone_by_rid(gen_VRE_STOR, z)
    if !isempty(resources)
        vP = EP[:vP]
        for t in 1:T
            ePowerBalance_VRE_STOR[t, z] =
                @expression(EP, sum(vP[y, t] for y in resources))
        end            
    end
end

odow · 2025-01-30T20:31:58Z

I don't know that other macro would have helped. It's things like "don't compute resources_in_zone_by_rid(gen_VRE_STOR, z) twice for every t in 1:T.

odow · 2025-01-30T22:27:27Z

We can't make @expression rewrite += because of this case:

julia> a = BigInt(1)
1

julia> b = a
1

julia> b += BigInt(2)
3

julia> a
1

+= is not exactly the same thing as add_to_expression!(x, y)

odow · 2025-01-30T22:28:48Z

Here was my diff to MA

diff --git a/src/rewrite_generic.jl b/src/rewrite_generic.jl
index de2234f..a8c1477 100644
--- a/src/rewrite_generic.jl
+++ b/src/rewrite_generic.jl
@@ -101,6 +101,22 @@ function _rewrite_generic(stack::Expr, expr::Expr)
         root = gensym()
         push!(stack.args, :($root = $if_expr))
         return root, is_mutable
+    elseif Meta.isexpr(expr, :+=, 2)
+        # x += y -> operate!!(add_mul, x, y)
+        y = expr.args[2]
+        rhs = Expr(:call, operate!!, add_mul, esc(expr.args[1]))
+        if Meta.isexpr(y, :call) && y.args[1] == :*
+            for i in 2:length(y.args)
+                arg_root, _ = _rewrite_generic(stack, y.args[i])
+                push!(rhs.args, arg_root)
+            end
+        else
+            y_root, _ = _rewrite_generic(stack, y)
+            push!(rhs.args, y_root)
+        end
+        root = gensym()
+        push!(stack.args, :($root = $rhs))
+        return root, true
     elseif !Meta.isexpr(expr, :call)
         # In situations like `x[i]`, we do not attempt to rewrite. Return `expr`
         # and don't let future callers mutate.

julia> @macroexpand @expression(model, y += x)
quote
    #= REPL[20]:1 =#
    JuMP._valid_model(model, :model)
    begin
        #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:399 =#
        let model = model
            #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:400 =#
            begin
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:86 =#
                begin
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:264 =#
                    var"#112###246" = begin
                            #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:374 =#
                            let
                                #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:375 =#
                                begin
                                    #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:371 =#
                                    var"#113###247" = (MutableArithmetics.operate!!)(MutableArithmetics.add_mul, y, x)
                                end
                                #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:376 =#
                                var"#113###247"
                            end
                        end
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:265 =#
                    var"#114###248" = (JuMP.flatten!)(var"#112###246")
                end
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:89 =#
                JuMP._replace_zero(model, var"#114###248")
            end
        end
    end
end

julia> @macroexpand @expression(model, y += 2 * x)
quote
    #= REPL[21]:1 =#
    JuMP._valid_model(model, :model)
    begin
        #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:399 =#
        let model = model
            #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:400 =#
            begin
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:86 =#
                begin
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:264 =#
                    var"#120###249" = begin
                            #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:374 =#
                            let
                                #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:375 =#
                                begin
                                    #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:371 =#
                                    var"#121###250" = (MutableArithmetics.operate!!)(MutableArithmetics.add_mul, y, 2, x)
                                end
                                #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:376 =#
                                var"#121###250"
                            end
                        end
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:265 =#
                    var"#122###251" = (JuMP.flatten!)(var"#120###249")
                end
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:89 =#
                JuMP._replace_zero(model, var"#122###251")
            end
        end
    end
end

odow · 2025-01-30T22:41:21Z

I think this is another issue that is really asking for a way to accurately profile where time is spent in JuMP. The issue is that the Julia profiler doesn't really help if you have lots of dispersed += or a * x because none of these individual lines shows up on the profiler. But collectively, the style of writing JuMP code is bad. It leads to more GC pressure, which slows things down.

See GenXProject/GenX.jl#815 (comment), where half the remaining time is GC.

@joaquimg have you tried profiling allocations instead of time?

joaquimg · 2025-01-31T01:27:19Z

More details about GenX is what I am currently looking at.

The issue is that the Julia profiler doesn't really help if you have lots of dispersed += or a * x because none of these individual lines shows up on the profiler. But collectively, the style of writing JuMP code is bad. It leads to more GC pressure, which slows things down.

yup

I have been checking for improvements by looking at allocations reduction, so we can reduce GC pressure. Still thinking about the best way to report.

joaquimg · 2025-01-31T01:45:12Z

We can't make @expression rewrite += because of this case

I dont fully get it.

The example is showing that b is assigned to a new expression, which is the same as in JuMP: exp += 2.

Using @expression would do it in-place, so yes b and a would end up with the same value. This is what I expected...

in JuMP:

Your example reads the same:

julia> a = 0 * x + 1
1

julia> b = a
1

julia> b += 2 # new expression created
3

julia> a
1

and add_to_expression! has the desired effect

julia> a = 0 * x + 1
1

julia> b = a
1

julia> add_to_expression!(b, 2)
3

julia> a
3

so I would expect:

julia> a = 0 * x + 1
1

julia> b = a
1

julia> @expression(model, b += 2) # in-place
3

julia> a
3

On the other hand, if the argument is coming from: "macro never modify their inputs (except model)", then I understand.

In which case another macro would make sense: @inplace, @expression!, @update

odow · 2025-01-31T01:46:26Z

I have been checking for improvements by looking at allocations reduction

See https://www.youtube.com/watch?v=BFvpwC8hEWQ

joaquimg · 2025-01-31T01:49:06Z

Also, analogous to @expression updating things in-place is

julia> a = [1]
1-element Vector{Int64}:
 1

julia> b = a
1-element Vector{Int64}:
 1

julia> @. b += [2]
1-element Vector{Int64}:
 3

julia> a
1-element Vector{Int64}:
 3

joaquimg · 2025-01-31T01:50:05Z

See https://www.youtube.com/watch?v=BFvpwC8hEWQ

Thanks! I have been using PProf allocation tracker. I will check is there are other cool stuff there.

odow · 2025-01-31T01:54:23Z

The rewrite macros shouldn't change the semantics of the code that runs. They should strictly be a performance improvement only. That's why we're so careful to track which things are mutable (meaning we created the object, and can thus change it without worrying that the user holds a reference to it).

@. b += [2]

Ah now this is very subtle. Because you are updating the elements of b, as well as the binding b itself.

joaquimg · 2025-01-31T02:01:17Z

The rewrite macros shouldn't change the semantics of the code that runs.

This is a fair point, it is the JuMP rule. So I agree we should not do += inside @expression.

Still, a new macro might make sense.

odow · 2025-01-31T03:10:49Z

You mean this macro:

julia> using JuMP

julia> macro add_to(expr)
           @assert Meta.isexpr(expr, :+=)
           x, y = expr.args
           ret = Expr(:call, JuMP.add_to_expression!, x)
           if Meta.isexpr(y, :call) && y.args[1] == :*
               append!(ret.args, y.args[2:end])
           else
               push!(ret.args, y)
           end
           return esc(ret)
       end
@add_to (macro with 1 method)

julia> model = Model()
A JuMP Model
├ solver: none
├ objective_sense: FEASIBILITY_SENSE
├ num_variables: 0
├ num_constraints: 0
└ Names registered in the model: none

julia> @variable(model, x)
x

julia> y = zero(AffExpr)
0

julia> @macroexpand @add_to y += 2 * x
:((JuMP.add_to_expression!)(y, 2, x))

julia> @macroexpand @add_to y += x
:((JuMP.add_to_expression!)(y, x))

julia> @macroexpand @add_to y += x / 2
:((JuMP.add_to_expression!)(y, x / 2))

This doesn't seem like it offers much benefit over the functional form.

I think the bar for adding a new macro should really be: can this be achieved any other way?

So far, our only utility macro is @force_nonlinear which is a very specialised form that needs to intercept how we rewrite the constraints. It's not to make things more readable.

joaquimg · 2025-01-31T03:30:49Z

Besides:

@add_to y += 2 * x

A more sophisticated @add_to would be nice for:

@add_to y += sum(x[i] for i in 1:N)

sure you could:

add_to_expression!(y, @expression(model, sum(x[i] for i in 1:N)))

but you would be potentially allocating a massive array and then throwing it away.
Moreover, @add_to is way more ergonomic.

odow · 2025-01-31T03:40:23Z

We don't have access to model in the @add_to macro's scope.

I think just no. I'm strongly against this. add_to_expression! is enough. If people don't find it readable they can do:

# Equivalent to `y += 2 * x`, but it modifies `y` in-place without allocating.
add_to_expression!(y, 2, x)

joaquimg · 2025-01-31T03:52:00Z

We don't have access to model in the @add_to macro's scope.

y is an expression so there is owner_model.

Recording my latest idea here:

@add_to_expression!(y, 2 * x)
@add_to_expression!(y, sum(x[i] for i in 1:N))

odow · 2025-01-31T22:03:12Z

@add_to_expression!(y, 2 * x)

There doesn't seem much benefit compared to

add_to_expression!(y, 2, x)

@add_to_expression!(y, sum(x[i] for i in 1:N))

Okay, I understand the motivation for this one, but still:

expr = @expression(model, sum(x[i] for i in 1:N))
add_to_expression!(y, expr)

I tend to think that people writing y += sum(x[i] for i in 1:N) are still going to write it, even if we add a fancy new macro. The issue isn't that there is a work-around, it's that they don't even know it is a problem in the first place.

joaquimg · 2025-02-01T00:16:23Z

I agree, @expression + add_to_expression! aint terrible, but allocates the double of @add_to_expression!.

I tend to think that people writing y += sum(x[i] for i in 1:N) are still going to write it, even if we add a fancy new macro. The issue isn't that there is a work-around, it's that they don't even know it is a problem in the first place.

I agree in part.

Many just don't know. The solution for this group would be an operation counter warning (for all operations, the previous warning was just for two expressions). Maybe we could even compute the aggregate time and allocations to tell the user: "you lost x time and y memory because you were too lazy".
The syntax: y += sum(x[i] for i in 1:N) is very nice, readable an mathy. @add_to y += sum(x[i] for i in 1:N) would be nice because the rewrite recommended by the above-suggested warning would be extremely easy to implement. This would minimize: "They are really telling me to rewrite it all?". @add_to_expression!(y, sum(x[i] for i in 1:N)) is a compromise to keep the new thing close to the "lower-level" suggestion.

odow · 2025-02-02T20:36:00Z

but allocates the double of

Then do:

for i in 1:N
    add_to_expression!(y, x[i])
end

The solution for this group would be an operation counter warning (for all operations, the previous warning was just for two expressions). Maybe we could even compute the aggregate time and allocations to tell the user: "you lost x time and y memory because you were too lazy".

Yip. See p131 of https://www.osti.gov/servlets/purl/1771935

We could also just document TimerOutputs and tell people to use that where possible.

odow · 2025-02-04T23:45:08Z

The concrete proposal for discussion is a suggestion to add some variation of this comment:

#3931 (comment)

julia> using JuMP

julia> macro add_to(expr)
           @assert Meta.isexpr(expr, :+=)
           x, y = expr.args
           ret = Expr(:call, JuMP.add_to_expression!, x)
           if Meta.isexpr(y, :call) && y.args[1] == :*
               append!(ret.args, y.args[2:end])
           else
               push!(ret.args, y)
           end
           return esc(ret)
       end
@add_to (macro with 1 method)

julia> model = Model()
A JuMP Model
├ solver: none
├ objective_sense: FEASIBILITY_SENSE
├ num_variables: 0
├ num_constraints: 0
└ Names registered in the model: none

julia> @variable(model, x)
x

julia> y = zero(AffExpr)
0

julia> @macroexpand @add_to y += 2 * x
:((JuMP.add_to_expression!)(y, 2, x))

julia> @macroexpand @add_to y += x
:((JuMP.add_to_expression!)(y, x))

julia> @macroexpand @add_to y += x / 2
:((JuMP.add_to_expression!)(y, x / 2))

Pros:

It would let people write += code more efficiently than they otherwise would.
Looks nice

Cons:

A new macro
Users still need to know about it
They could use add_to_expression! instead.

blegat · 2025-02-08T19:13:18Z

I would just do this with MA.@rewrite, no need for a different macro name

odow · 2025-02-08T21:50:23Z

We can't do it in MA.@rewrite because this isn't a universally valid rewrite. See my comment #3931 (comment)

blegat · 2025-02-10T08:53:12Z

I see. So it's kind of like operate!!. The operation y += x means the Julia variable should always alias to the result of y + x and the object that y used to alias before the operation may be modified. We could call it @rewrite!! where !! means that this macro may mutate the lhs.

odow · 2025-02-10T21:40:50Z

If we add something like this, I don't really think that MutableArithmetics is a good home for it. I'd like it to be user-friendly at the JuMP level. People using MA for performance can opt-in to operate!! instead of +=. It's energy modeler people like GenX that find add_to_expression! annoying instead of +=.

Even if we add this, we still have to teach them to use it. And "add this extra package" is a large step (even if it is already a transitive dependency).

blegat · 2025-02-11T08:39:38Z

I think the concern is that operate!!(+, a, b) is much less readable than a += b. The issue with JuMP macros like @expression is that it's not so transparent what they do. Do they create sub-expressions exploited by JuMP's AD ? Is it just a synonym to Containers.@container ? Does it rewrite the expression to exploit mutable arithmetics ?
If the users does not write a += b and wants something more efficient, it's not a lot to ask for them to understand they need to exploit the mutability of JuMP expressions. So if the solution is to do import MutableArithmetics as MA and then use MA.@rewrite!!(a += b), I feel it's the most readable thing you can think of. MA.@rewrite!! clearly says that it rewrites the expression to exploit the mutability. It also says that it does not do anything that is JuMP-specific since it is in MA which also helps readability.

odow · 2025-02-11T19:04:53Z

@lbonaldo any thoughts on this issue?

lbonaldo · 2025-02-11T20:29:00Z

My intuition is that performance and memory allocation are not something that users will immediately think about when getting started with JuMP or prototyping new models for a research project. They are likely to treat JuMP objects as regular variables without thinking much about the cost of operations between them. However, when dealing with very large models, these aspects become crucial.

In our group, we're gradually transitioning to a more idiomatic JuMP style, thanks largely to @RuaridhMacd and your suggestions. We've been replacing all instances of the += operator with add_to_expression! or wrappers around it. There’s still work to do to update the entire codebase, but reviewers and maintainers are now on board with this approach. In general, I also don’t think it’s particularly difficult for users to adopt add_to_expression! instead of +=.

That’s just my intuition— I'm tagging @RuaridhMacd who has more experience and may have additional insights.

odow · 2025-02-11T21:45:50Z

If you're happy to use add_to_expression! then I'm very strongly leaning towards closing this issue.

We can address the underlying issue with better profiling support. Even a documentation checklist of things to investigate could be useful.

RuaridhMacd · 2025-02-11T21:51:18Z

Thanks @lbonaldo

Yes, I agree that 'add_to_expression!' and its derivatives are working well for GenX. It's no more difficult than += once people are aware of the issue

lbonaldo · 2025-02-11T23:53:43Z

yes, tutorials on how to effectively profile JuMP code and follow JuMP best practices would be very helpful.

odow added the Status: Needs developer call This should be discussed on a monthly developer call label Feb 4, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Macro to update expression in-place #3931

Macro to update expression in-place #3931

joaquimg commented Jan 30, 2025 •

edited

Loading

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

joaquimg commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025 •

edited

Loading

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Feb 1, 2025

odow commented Feb 2, 2025

odow commented Feb 4, 2025

blegat commented Feb 8, 2025 •

edited

Loading

odow commented Feb 8, 2025

blegat commented Feb 10, 2025 •

edited

Loading

odow commented Feb 10, 2025

blegat commented Feb 11, 2025

odow commented Feb 11, 2025

lbonaldo commented Feb 11, 2025 •

edited

Loading

odow commented Feb 11, 2025

RuaridhMacd commented Feb 11, 2025

lbonaldo commented Feb 11, 2025

Macro to update expression in-place #3931

Macro to update expression in-place #3931

Comments

joaquimg commented Jan 30, 2025 • edited Loading

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

odow commented Jan 30, 2025

joaquimg commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025 • edited Loading

odow commented Jan 31, 2025

joaquimg commented Jan 31, 2025

odow commented Jan 31, 2025

joaquimg commented Feb 1, 2025

odow commented Feb 2, 2025

odow commented Feb 4, 2025

blegat commented Feb 8, 2025 • edited Loading

odow commented Feb 8, 2025

blegat commented Feb 10, 2025 • edited Loading

odow commented Feb 10, 2025

blegat commented Feb 11, 2025

odow commented Feb 11, 2025

lbonaldo commented Feb 11, 2025 • edited Loading

odow commented Feb 11, 2025

RuaridhMacd commented Feb 11, 2025

lbonaldo commented Feb 11, 2025

joaquimg commented Jan 30, 2025 •

edited

Loading

joaquimg commented Jan 31, 2025 •

edited

Loading

blegat commented Feb 8, 2025 •

edited

Loading

blegat commented Feb 10, 2025 •

edited

Loading

lbonaldo commented Feb 11, 2025 •

edited

Loading