Performance disparity in expression operations

[jd-lara]

I made a test of the performance impact of having Parameters in an affine expression for a different number of expressions in the vector product operation.

The test compares 3 possible cases.

1. Expressions with parameters ParameterJuMP.ParametrizedGenericAffExpr{Float64, JuMP.VariableRef}
2. Expresions without parameters. JuMP.GenericAffExpr{Float64, JuMP.VariableRef}
3. Mix of expressions using JuMP.GenericAffExpr{Float64, JuMP.VariableRef} and JuMP.GenericAffExpr{Float64, JuMP.VariableRef}

Here is some reproducible code:

using JuMP
using ParameterJuMP
using Random

function model_with_params(size_1, size_2)
    m = Model()
    ParameterJuMP.enable_parameters(m)

    set_names_1 = [randstring() for _ in 1:size_1]
    set_names_2 = [randstring() for _ in 1:size_1]
    time = 1:size_2

    z = @variable(m, [v in set_names_1, t in time])
    x = @variable(m, [v in set_names_2, t in time])

    param_array = Matrix(undef, size_1, size_2)
    for i in 1:size_1, t in 1:size_2
        param_array[i, t] = JuMP.@variable(m, variable_type = ParameterJuMP.Param())
        ParameterJuMP.set_value(param_array[i, t], rand())
    end

    expression_container = JuMP.Containers.DenseAxisArray{ParameterJuMP.ParametrizedGenericAffExpr{Float64, JuMP.VariableRef}}(undef, 1:size_1, 1:size_2)

    for i in 1:size_1, t in 1:size_2
        expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + param_array[i, t] + rand()
    end

    col = rand(size_1,)
    _, build_time, build_bytes_alloc, build_sec_in_gc =
            @timed @constraint(m, sum(col[i]*expression_container[i, 1] for i in 1:size_1) == 0.0)
    return build_time, build_bytes_alloc, build_sec_in_gc
end

function model_without_params(size_1, size_2)
    m2 = Model()

    set_names_1 = [randstring() for _ in 1:size_1]
    set_names_2 = [randstring() for _ in 1:size_1]
    time = 1:size_2

    z = @variable(m2, [v in set_names_1, t in time])
    x = @variable(m2, [v in set_names_2, t in time])

    expression_container = JuMP.Containers.DenseAxisArray{JuMP.GenericAffExpr{Float64, JuMP.VariableRef}}(undef, 1:size_1, 1:size_2)

    for i in 1:size_1, t in 1:size_2
        expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + rand()
    end

    col = rand(size_1,)
    _, build_time, build_bytes_alloc, build_sec_in_gc =
            @timed @constraint(m2, sum(col[i]*expression_container[i, 1] for i in 1:size_1) == 0.0)
    return build_time, build_bytes_alloc, build_sec_in_gc
end

function model_with_mixed_params(size_1, size_2)
    m = Model()
    ParameterJuMP.enable_parameters(m)

    set_names_1 = [randstring() for _ in 1:size_1]
    set_names_2 = [randstring() for _ in 1:size_1]
    time = 1:size_2

    z = @variable(m, [v in set_names_1, t in time])
    x = @variable(m, [v in set_names_2, t in time])

    param_array = Matrix(undef, size_1, size_2)
    for i in 1:size_1, t in 1:size_2
        param_array[i, t] = JuMP.@variable(m, variable_type = ParameterJuMP.Param())
        ParameterJuMP.set_value(param_array[i, t], rand())
    end

    expression_container = JuMP.Containers.DenseAxisArray{ParameterJuMP.ParametrizedGenericAffExpr{Float64, JuMP.VariableRef}}(undef, 1:size_1, 1:size_2)

    param_ixs = shuffle(1:size_1)[1:Int(size_1/2)]
    for i in 1:size_1, t in 1:size_2
        if i in param_ixs
            expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + param_array[i, t] + rand()
        else
            expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + rand()
        end
    end

    col = rand(size_1,)
    _, build_time, build_bytes_alloc, build_sec_in_gc =
            @timed @constraint(m, sum(col[i]*expression_container[i, 1] for i in 1:size_1) == 0.0)
    return build_time, build_bytes_alloc, build_sec_in_gc
end


model_with_params(10, 48)
model_without_params(10, 48)
model_with_mixed_params(10, 48)

results = Dict()
results_params = Dict()
results_mixed = Dict()
for i in [100, 1000, 5000, 10_000, 50_000, 100_000, 500_000]
    println(i)
    @show results[i] = model_without_params(i, 48)
    @show results_params[i] = model_with_params(i, 48)
    @show results_mixed = model_with_mixed_params(i, 48)
end

Results:

100
results[i] = model_without_params(i, 48) = (0.00015558, 29488, 0.0)
results_params[i] = model_with_params(i, 48) = (0.000182823, 76048, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.000172758, 67904, 0.0)
1000
results[i] = model_without_params(i, 48) = (0.000626027, 197072, 0.0)
results_params[i] = model_with_params(i, 48) = (0.001597523, 560976, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.001377474, 462640, 0.0)
5000
results[i] = model_without_params(i, 48) = (0.00356614, 1241216, 0.0)
results_params[i] = model_with_params(i, 48) = (0.008178023, 3541104, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.006022408, 2940048, 0.0)
10000
results[i] = model_without_params(i, 48) = (0.006917794, 2477760, 0.0)
results_params[i] = model_with_params(i, 48) = (0.019982686, 7071424, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.014490844, 6131968, 0.0)
50000
results[i] = model_without_params(i, 48) = (0.048876917, 10544256, 0.0)
results_params[i] = model_with_params(i, 48) = (0.234299942, 31453584, 0.11880938)
results_mixed = model_with_mixed_params(i, 48) = (0.322152996, 25348384, 0.203175162)
100000
results[i] = model_without_params(i, 48) = (0.106661881, 18983120, 0.0)
results_params[i] = model_with_params(i, 48) = (7.070552948, 54502512, 6.839998649)
results_mixed = model_with_mixed_params(i, 48) = (0.384903389, 48584096, 0.173239921)

[jd-lara]

The 500 000 case finished in 4 seconds for GenericAffExpr and crashed my julia session for the cases with parameters.

[odow]

I think your results are getting mixed with when the GC runs.

using JuMP
using ParameterJuMP
using Random

function model_with_params(m)
    model = Model()
    @variable(model, x[1:m])
    @variable(model, y[1:m])
    @variable(model, p[1:m] == 1.0, ParameterJuMP.Param())
    @expression(model, expr[i=1:m], x[i] + y[i] + p[i] + rand())
    col = rand(m)
    @time @constraint(model, sum(col[i] * expr[i] for i in 1:m) == 0.0)
    return
end

function model_without_params(m)
    model = Model()
    @variable(model, x[1:m])
    @variable(model, y[1:m])
    @variable(model, p[1:m] == 1.0)
    @expression(model, expr[i=1:m], x[i] + y[i] + p[i] + rand())
    col = rand(m)
    @time @constraint(model, sum(col[i] * expr[i] for i in 1:m) == 0.0)
    return
end

model_with_params(10)
model_without_params(10)
for i in [100, 1000, 10_000, 50_000, 100_000, 500_000]
    println(i)
    GC.gc()
    model_without_params(i)
    GC.gc()
    model_with_params(i)
end

Yields

100
  0.000066 seconds (33 allocations: 35.188 KiB)
  0.000130 seconds (102 allocations: 69.453 KiB)
1000
  0.000607 seconds (45 allocations: 299.047 KiB)
  0.001099 seconds (137 allocations: 543.016 KiB)
10000
  0.006241 seconds (57 allocations: 3.534 MiB)
  0.027691 seconds (184 allocations: 6.740 MiB, 53.47% gc time)
50000
  0.038874 seconds (65 allocations: 13.946 MiB)
  0.072116 seconds (214 allocations: 29.975 MiB)
100000
  0.081927 seconds (70 allocations: 25.032 MiB)
  0.178597 seconds (228 allocations: 51.979 MiB)
500000
  0.631749 seconds (79 allocations: 116.645 MiB)
  1.029031 seconds (259 allocations: 202.852 MiB)

So there is a ~2x penalty to using parameters. Which makes sense: every expression has two dictionaries.

[jd-lara]

Thanks, that makes sense.

[odow]

@jd-lara can you try the following instead?

Replace

@constraint(m, sum(col[i]*expression_container[i, 1] for i in 1:size_1) == 0.0)

with

ex = @expression(m, sum(col[i]*expression_container[i, 1] for i in 1:size_1))
c = ScalarConstraint(ex, MOI.EqualTo(0.0))
add_constraint(m, c, "")

[jd-lara]

@odow I have updated the test as follows:

using JuMP
using ParameterJuMP
using Random

function _make_constraint(model::JuMP.Model, col::Vector{Float64}, expression_container::JuMP.Containers.DenseAxisArray{ParameterJuMP.ParametrizedGenericAffExpr{Float64, JuMP.VariableRef}}, size_1::Int)
    ex = @expression(model, sum(col[i]*expression_container[i, 1] for i in 1:size_1))
    c = ScalarConstraint(ex, MOI.EqualTo(0.0))
    add_constraint(model, c, "")
end

function _make_constraint(model::JuMP.Model, col::Vector{Float64}, expression_container::JuMP.Containers.DenseAxisArray{JuMP.GenericAffExpr{Float64, JuMP.VariableRef}}, size_1::Int)
    ex = @expression(model, sum(col[i]*expression_container[i, 1] for i in 1:size_1))
    c = ScalarConstraint(ex, MOI.EqualTo(0.0))
    add_constraint(model, c, "")
end

function model_with_params(size_1, size_2)
    m = Model()
    ParameterJuMP.enable_parameters(m)

    set_names_1 = [randstring() for _ in 1:size_1]
    set_names_2 = [randstring() for _ in 1:size_1]
    time = 1:size_2

    z = @variable(m, [v in set_names_1, t in time])
    x = @variable(m, [v in set_names_2, t in time])

    param_array = Matrix(undef, size_1, size_2)
    for i in 1:size_1, t in 1:size_2
        param_array[i, t] = JuMP.@variable(m, variable_type = ParameterJuMP.Param())
        ParameterJuMP.set_value(param_array[i, t], rand())
    end

    expression_container = JuMP.Containers.DenseAxisArray{ParameterJuMP.ParametrizedGenericAffExpr{Float64, JuMP.VariableRef}}(undef, 1:size_1, 1:size_2)

    for i in 1:size_1, t in 1:size_2
        expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + param_array[i, t] + rand()
    end

    _, build_time, build_bytes_alloc, build_sec_in_gc =
            @timed _make_constraint(m, rand(size_1,), expression_container, size_1)
    return build_time, build_bytes_alloc, build_sec_in_gc
end

function model_without_params(size_1, size_2)
    m2 = Model()

    set_names_1 = [randstring() for _ in 1:size_1]
    set_names_2 = [randstring() for _ in 1:size_1]
    time = 1:size_2

    z = @variable(m2, [v in set_names_1, t in time])
    x = @variable(m2, [v in set_names_2, t in time])

    expression_container = JuMP.Containers.DenseAxisArray{JuMP.GenericAffExpr{Float64, JuMP.VariableRef}}(undef, 1:size_1, 1:size_2)

    for i in 1:size_1, t in 1:size_2
        expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + rand()
    end

    _, build_time, build_bytes_alloc, build_sec_in_gc =
            @timed _make_constraint(m2, rand(size_1,), expression_container, size_1)
    return build_time, build_bytes_alloc, build_sec_in_gc
end

function model_with_mixed_params(size_1, size_2)
    m = Model()
    ParameterJuMP.enable_parameters(m)

    set_names_1 = [randstring() for _ in 1:size_1]
    set_names_2 = [randstring() for _ in 1:size_1]
    time = 1:size_2

    z = @variable(m, [v in set_names_1, t in time])
    x = @variable(m, [v in set_names_2, t in time])

    param_array = Matrix(undef, size_1, size_2)
    for i in 1:size_1, t in 1:size_2
        param_array[i, t] = JuMP.@variable(m, variable_type = ParameterJuMP.Param())
        ParameterJuMP.set_value(param_array[i, t], rand())
    end

    expression_container = JuMP.Containers.DenseAxisArray{ParameterJuMP.ParametrizedGenericAffExpr{Float64, JuMP.VariableRef}}(undef, 1:size_1, 1:size_2)

    param_ixs = shuffle(1:size_1)[1:Int(size_1/2)]
    for i in 1:size_1, t in 1:size_2
        if i in param_ixs
            expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + param_array[i, t] + rand()
        else
            expression_container[i, t] = z[set_names_1[i], t] + 5.0*x[set_names_2[i], t] + rand()
        end
    end


    _, build_time, build_bytes_alloc, build_sec_in_gc =
            @timed _make_constraint(m, rand(size_1,), expression_container, size_1)
    return build_time, build_bytes_alloc, build_sec_in_gc
end

# FIrst compile run
model_with_params(10, 48)
model_without_params(10, 48)
model_with_mixed_params(10, 48)

results = Dict()
results_params = Dict()
results_mixed = Dict()
for i in [100, 1000, 5000, 10_000, 50_000, 100_000]
    println(i)
    @show results[i] = model_without_params(i, 48)
    GC.gc()
    @show results_params[i] = model_with_params(i, 48)
    GC.gc()
    @show results_mixed = model_with_mixed_params(i, 48)
end

Looks like up to 10000 expressions, the penalty is about 2x as expected and over 50000 the penalty is 4x or more

100
results[i] = model_without_params(i, 48) = (9.7006e-5, 30384, 0.0)
results_params[i] = model_with_params(i, 48) = (0.000171875, 76912, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.000131348, 68880, 0.0)
1000
results[i] = model_without_params(i, 48) = (0.000831869, 205200, 0.0)
results_params[i] = model_with_params(i, 48) = (0.001245529, 569072, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.001065928, 470624, 0.0)
5000
results[i] = model_without_params(i, 48) = (0.003808976, 1281296, 0.0)
results_params[i] = model_with_params(i, 48) = (0.008973684, 3581152, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.005882702, 2980096, 0.0)
10000
results[i] = model_without_params(i, 48) = (0.008080068, 2557840, 0.0)
results_params[i] = model_with_params(i, 48) = (0.01753953, 7151472, 0.0)
results_mixed = model_with_mixed_params(i, 48) = (0.01423789, 6212016, 0.0)
50000
results[i] = model_without_params(i, 48) = (0.050672775, 10944336, 0.0)
results_params[i] = model_with_params(i, 48) = (0.228111582, 31853632, 0.116043022)
results_mixed = model_with_mixed_params(i, 48) = (0.085427278, 25748320, 0.0)
100000
results[i] = model_without_params(i, 48) = (0.117079187, 19783200, 0.0)
results_params[i] = model_with_params(i, 48) = (0.428364249, 55302560, 0.13931007)
results_mixed = model_with_mixed_params(i, 48) = (0.668327524, 49384144, 0.377572102)

[odow]

That's all just GC differences though. If you subtract the GC time, you get 2.5x again:

julia> (0.428364249 - 0.13931007) / 0.117079187
2.4688775725782923

julia> (0.228111582 - 0.116043022) / 0.050672775
2.211612843385822

julia> (0.668327524 - 0.377572102) / 0.117079187
2.4834082764855543

[jd-lara]

So the question is how to reduce the GC times.

[odow]

It'd be good to confirm why the 2x exists first. Is it creating the extra dictionary? Extra key lookups? Resizing the two dictionaries?

One option that might be faster is to refactor everything into a single dictionary with Tuple{Bool,VariableRef} keys.

That might also help GC by creating less dictionaries.

[joaquimg]

I dont understand why having two dicts would be much worse than having a single larger one.
Dis you run any profiling?

[jd-lara]

I have been doing @code_warntype to check type stability but I haven't explored more since this issue.

[joaquimg]

Have you tried: https://github.com/timholy/ProfileView.jl ?

[jd-lara]

@joaquimg I made some observations with @kdheepak with pprof() but haven't found anything definitive

[jd-lara]

Profile using Parameters

Profile without parameters

[jd-lara]

@joaquimg and @odow I think the culprit of this issue is in this line

https://github.com/jump-dev/JuMP.jl/blob/2ff880d46161a8502d9ca378d098d60d1c9864c3/src/aff_expr.jl#L26