Fix mutable allocations in n-ary multiplication

[odow]

I thought @joaquimg's suggestion (3) in NREL-Sienna/PowerSimulations.jl#1218 was weird. It's because we weren't optimizing this case.

Before

julia> model = Model();

julia> @variable(model, x)
x

julia> a, b = 2.0, 3.0
(2.0, 3.0)

julia> @allocated @expression(model, a * x * b)
1728

After

julia> model = Model();

julia> @variable(model, x)
x

julia> a, b = 2.0, 3.0
(2.0, 3.0)

julia> @allocated @expression(model, a * x * b)
1152

Before code:

julia> @macroexpand @expression(model, a * x * b)
quote
    #= REPL[66]: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"#313###301" = 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"#314###302" = (MutableArithmetics.operate)(*, a, x)
                                    var"#315###303" = (MutableArithmetics.operate)(*, var"#314###302", b)
                                end
                                #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:376 =#
                                var"#315###303"
                            end
                        end
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:265 =#
                    var"#316###304" = (JuMP.flatten!)((MutableArithmetics).copy_if_mutable(var"#313###301"))
                end
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:89 =#
                JuMP._replace_zero(model, var"#316###304")
            end
        end
    end
end

Issue is this line, which allocates y = a * x and then y * b

                                    var"#314###302" = (MutableArithmetics.operate)(*, a, x)
                                    var"#315###303" = (MutableArithmetics.operate)(*, var"#314###302", b)

[odow]

I guess this changed in the nonlinear rewrite:

julia> @macroexpand MA.@rewrite(x[1] * x[2] * x[3], move_factors_into_sums = false)
quote
    #= /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"#328###307" = (MutableArithmetics.operate)(*, x[1], x[2])
            var"#329###308" = (MutableArithmetics.operate)(*, var"#328###307", x[3])
        end
        #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:376 =#
        var"#329###308"
    end
end

julia> @macroexpand MA.@rewrite(x[1] * x[2] * x[3], move_factors_into_sums = true)
quote
    #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:365 =#
    let
        #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:366 =#
        begin
            begin
                var"#332###310" = MutableArithmetics.Zero()
                var"#333###309" = (MutableArithmetics.operate!!)(MutableArithmetics.add_mul, var"#332###310", x[1], x[2], x[3])
            end
        end
        #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:367 =#
        var"#333###309"
    end
end

I'm surprised we didn't notice earlier.

[odow]

I guess we didn't notice because we don't hit this code path if it is inside a summation, so this is less of a problem than I first thought:

julia> model = Model();

julia> @variable(model, x[1:2])
2-element Vector{VariableRef}:
 x[1]
 x[2]

julia> a, b = 2.0, 3.0
(2.0, 3.0)

julia> @macroexpand @expression(model, sum(a * x[i] * b for i in 1:2))
quote
    #= REPL[83]: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"#378###326" = 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"#379###327" = MutableArithmetics.Zero()
                                    for i = 1:2
                                        #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite_generic.jl:318 =#
                                        var"#379###327" = (MutableArithmetics.operate!!)(MutableArithmetics.add_mul, var"#379###327", a, x[i], b)
                                    end
                                end
                                #= /Users/oscar/.julia/dev/MutableArithmetics/src/rewrite.jl:376 =#
                                var"#379###327"
                            end
                        end
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:265 =#
                    var"#380###328" = (JuMP.flatten!)(var"#378###326")
                end
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@expression.jl:89 =#
                JuMP._replace_zero(model, var"#380###328")
            end
        end
    end
end