Macro to update expression in-place

[joaquimg]

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

[odow]

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]

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

[odow]

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]

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]

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]

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]

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]

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]

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]

I have been checking for improvements by looking at allocations reduction

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

[joaquimg]

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]

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]

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]

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]

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]

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]

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]

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]

@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]

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]

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]

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.