Allocations in MA (or possibly JuMP)

[joaquimg]

All using GenX 0.4.2 and their case data: https://github.com/GenXProject/GenX.jl/tree/main/example_systems/8_three_zones_w_colocated_VRE_storage_electrolyzers

we started with:

  [5d317b1e] GenX v0.4.2 `~/.julia/dev/GenX`
  [87dc4568] HiGHS v1.13.0
  [682c06a0] JSON v0.21.4
⌃ [4076af6c] JuMP v1.23.4
⌃ [d8a4904e] MutableArithmetics v1.4.6 ⚲

  5.233504 seconds (19.13 M allocations: 4.140 GiB, 45.43% gc time)

After: jump-dev/MutableArithmetics.jl#303

  [5d317b1e] GenX v0.4.2 `~/.julia/dev/GenX`
  [87dc4568] HiGHS v1.13.0
  [682c06a0] JSON v0.21.4
⌃ [4076af6c] JuMP v1.23.4
⌃ [d8a4904e] MutableArithmetics v1.5.0 ⚲

4.271410 seconds (18.79 M allocations: 3.404 GiB, 37.80% gc time)

(versions 1.5.1 and 1.5.2 lead to the same result above)

But more recently we got some extra allocations back

  [5d317b1e] GenX v0.4.2 `~/.julia/dev/GenX`
  [87dc4568] HiGHS v1.13.0
  [682c06a0] JSON v0.21.4
  [4076af6c] JuMP v1.23.6
  [d8a4904e] MutableArithmetics v1.6.3

4.363094 seconds (20.59 M allocations: 3.540 GiB, 39.73% gc time)

script:

import GenX
import HiGHS
# solver
import JuMP
import HiGHS
function build_genx()
    case = joinpath(@__DIR__, "cases", "8_three_zones_w_colocated_VRE_storage_electrolyzers")
    optimizer = HiGHS.Optimizer

    GenX.print_genx_version() # Log the GenX version
    genx_settings = GenX.get_settings_path(case, "genx_settings.yml") # Settings YAML file path
    writeoutput_settings = GenX.get_settings_path(case, "output_settings.yml") # Write-output settings YAML file path
    mysetup = GenX.configure_settings(genx_settings, writeoutput_settings) # mysetup dictionary stores settings and GenX-specific parameters

    if mysetup["MultiStage"] != 0
        error()
    end

    settings_path = GenX.get_settings_path(case)

    ### Cluster time series inputs if necessary and if specified by the user
    if mysetup["TimeDomainReduction"] == 1
        TDRpath = joinpath(case, mysetup["TimeDomainReductionFolder"])
        system_path = joinpath(case, mysetup["SystemFolder"])
        GenX.prevent_doubled_timedomainreduction(system_path)
        if !GenX.time_domain_reduced_files_exist(TDRpath)
            println("Clustering Time Series Data (Grouped)...")
            GenX.cluster_inputs(case, settings_path, mysetup)
        else
            println("Time Series Data Already Clustered.")
        end
    end

    ### Configure solver
    println("Configuring Solver")
    solver_name = lowercase(get(mysetup, "Solver", ""))
    OPTIMIZER = GenX.configure_solver(settings_path, optimizer; solver_name = solver_name)

    #### Running a case

    ### Load inputs
    println("Loading Inputs")
    myinputs = GenX.load_inputs(mysetup, case)

    println("Generating the Optimization Model")
    time_elapsed = @elapsed EP = GenX.generate_model(mysetup, myinputs, OPTIMIZER)
    println("Time elapsed for model building is")
    println(time_elapsed)
    return
end
build_genx()

[joaquimg]

There are not many details but newer versions seem to have an "extra" copy_if_mutable

(top is newer)

[joaquimg]

I did a binary search and this is not MA fault.
MA 1.6.3 actually allocates even a bit less than 1.5.0.

The issue is something from JuMP 1.23.5 release

[joaquimg]

Probably this: #3883

[odow]

We should investigate ways to opt-in to this only if needed:

JuMP.jl/src/macros.jl

Lines 259 to 261 in 41db3c3

	if !is_mutable
	new_aff = :($_MA.copy_if_mutable($new_aff))
	end

[odow]

For my future reference, see these examples:

julia> using JuMP

julia> model = Model();

julia> @variable(model, x);

julia> @expression(model, expr, 2 * x);

julia> @macroexpand @constraint(model, 0 <= expr <= 1)
quote
    #= REPL[28]: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/@constraint.jl:171 =#
                begin
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:482 =#
                    var"#135###261" = (MutableArithmetics).copy_if_mutable(expr)
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:483 =#
                    var"#136###262" = (MutableArithmetics).copy_if_mutable(0)
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:484 =#
                    var"#137###263" = (MutableArithmetics).copy_if_mutable(1)
                end
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:172 =#
                var"#138#build" = JuMP.model_convert(model, JuMP.build_constraint(JuMP.Containers.var"#error_fn#98"{String}("At REPL[28]:1: `@constraint(model, 0 <= expr <= 1)`: "), var"#135###261", var"#136###262", var"#137###263"))
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:173 =#
                JuMP.add_constraint(model, var"#138#build", "")
            end
        end
    end
end

There's also this example that we could be cleverer about:

julia> using JuMP

julia> model = Model();

julia> @variable(model, x);

julia> @expression(model, expr, 2 * x);

julia> @macroexpand @constraint(model, expr <= 0)
quote
    #= REPL[33]: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/@constraint.jl:171 =#
                begin
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:264 =#
                    var"#155###267" = begin
                            #= /Users/oscar/.julia/packages/MutableArithmetics/bmgkg/src/rewrite.jl:374 =#
                            let
                                #= /Users/oscar/.julia/packages/MutableArithmetics/bmgkg/src/rewrite.jl:375 =#
                                begin
                                    #= /Users/oscar/.julia/packages/MutableArithmetics/bmgkg/src/rewrite.jl:371 =#
                                    var"#156###268" = (MutableArithmetics.copy_if_mutable)(expr)
                                    var"#157###269" = (MutableArithmetics.operate!!)(MutableArithmetics.sub_mul, var"#156###268", 0)
                                end
                                #= /Users/oscar/.julia/packages/MutableArithmetics/bmgkg/src/rewrite.jl:376 =#
                                var"#157###269"
                            end
                        end
                    #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros.jl:265 =#
                    var"#158###270" = (JuMP.flatten!)(var"#155###267")
                end
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:172 =#
                var"#159#build" = JuMP.model_convert(model, JuMP.build_constraint(JuMP.Containers.var"#error_fn#98"{String}("At REPL[33]:1: `@constraint(model, expr <= 0)`: "), JuMP._functionize(var"#158###270"), LessThanZero()))
                #= /Users/oscar/.julia/packages/JuMP/CU7H5/src/macros/@constraint.jl:173 =#
                JuMP.add_constraint(model, var"#159#build", "")
            end
        end
    end
end

[odow]

I don't know that there's much we can do here:

• In the common case, f(x) <= g(x), we create a new mutable expression by rewriting f(x) - g(x) <= 0.
• Then, we move the constant from the function into the set
  

  JuMP.jl/src/macros/@constraint.jl

  Lines 1023 to 1038 in 5d28e63

  	# !!! warning
  	# This method assumes that we can mutate `expr`. Ensure that this is the
  	# case upstream of this call site.
  	function build_constraint(
  	::Function,
  	expr::Union{Number,GenericAffExpr,GenericQuadExpr},
  	set::MOI.AbstractScalarSet,
  	)
  	if MOI.Utilities.supports_shift_constant(typeof(set))
  	expr, offset = _clear_constant!(expr)
  	new_set = MOI.Utilities.shift_constant(set, -offset)
  	return ScalarConstraint(expr, new_set)
  	else
  	return ScalarConstraint(expr, set)
  	end
  	end

• The extra copy_if_mutable ensures that if the user passes in an expression like 0 <= expr[1] <= 1, we create a copy so that it can be mutated. This is a necessary bug fix.
• There are some cases, like [expr] in Zeros() where we hit the extra copy_if_mutable, and it isn't strictly necessary. So currently, these copies are redundant. However, it means that if we ever modify the build_constraint methods to modify expr in-place, then we don't accidentally introduce a new bug.

[odow]

This actually has some interesting consequences:

julia> model = Model();

julia> @variable(model, x);

julia> @expression(model, a[1:1], x+1)
1-element Vector{AffExpr}:
 x + 1

julia> @expression(model, b[i in 1:1], a[i])
1-element Vector{AffExpr}:
 x + 1

julia> add_to_expression!(a[1], 1)
x + 2

julia> a
1-element Vector{AffExpr}:
 x + 2

julia> b
1-element Vector{AffExpr}:
 x + 1

I think this is arguably the correct behavior, but JuMP 1.23.4 had:

julia> using JuMP

julia> model = Model();

julia> @variable(model, x);

julia> @expression(model, a[1:1], x+1)
1-element Vector{AffExpr}:
 x + 1

julia> @expression(model, b[i in 1:1], a[i])
1-element Vector{AffExpr}:
 x + 1

julia> add_to_expression!(a[1], 1)
x + 2

julia> a
1-element Vector{AffExpr}:
 x + 2

julia> b
1-element Vector{AffExpr}:
 x + 2

[odow]

Closing because I don't think there is much to do here. The regression is a bug fix.