Performance of LinearAlgebra.dot

[matbesancon]

TL;DR: using dot or sum to construct a SAF leads to significant performance degradation compared to manually constructing the SAF from a loop comprehension. The "manual" version below does not use linear algebra functions, the non-manual one does, and hence calls MutableArithmetics.

Setting:

  [87dc4568] HiGHS v0.3.1
  [b8f27783] MathOptInterface v0.10.7

MWE:

using BenchmarkTools, SparseArrays
import MathOptInterface
const MOI = MathOptInterface
using HiGHS
using LinearAlgebra
using Random

function build_moi_model(n, optimizer = HiGHS.Optimizer())
    MOI.set(optimizer, MOI.Silent(), true)
    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    (x, _) = MOI.add_constrained_variables(optimizer, fill(MOI.Interval(0.0, 1.0), n * n))
    xmat = reshape(x, n, n)
    for idx in 1:n
       # column constraint
       MOI.add_constraint(
           optimizer,
           sum(1.0 * xmat[:, idx]),
           MOI.EqualTo(1.0),
       )
       # row constraint
       MOI.add_constraint(
           optimizer,
           sum(1.0 * xmat[idx, :]),
           MOI.EqualTo(1.0),
       )
    end
    return optimizer
end

function moi_lmo_manual(direction, optimizer)
    n = size(direction, 1)
    variables = reshape(MOI.get(optimizer, MOI.ListOfVariableIndices()), n, n)
    obj = MOI.ScalarAffineFunction(
        [MOI.ScalarAffineTerm(direction[i], variables[i]) for i in eachindex(direction)],
        0.0,
    )
    MOI.set(optimizer, MOI.ObjectiveFunction{typeof(obj)}(), obj)
    MOI.optimize!(optimizer)
    v = spzeros(n, n)
    @inbounds for i in 1:n
        for j in 1:n
            vij = MOI.get(optimizer, MOI.VariablePrimal(), variables[i,j])
            if vij ≈ 1
                v[i,j] = vij
            end
        end
    end
    return v
end

function moi_lmo(direction, optimizer)
           n = size(direction, 1)
           variables = reshape(MOI.get(optimizer, MOI.ListOfVariableIndices()), n, n)
           obj = dot(direction, variables)
           MOI.set(optimizer, MOI.ObjectiveFunction{typeof(obj)}(), obj)
           MOI.optimize!(optimizer)
           v = spzeros(n, n)
           @inbounds for i in 1:n
               for j in 1:n
                   vij = MOI.get(optimizer, MOI.VariablePrimal(), variables[i,j])
                   if vij ≈ 1
                       v[i,j] = vij
                   end
               end
           end
           return v
end

function f(optimizer, manual)
           direction = zeros(n, n)
           if manual
               moi_lmo_manual(randn!(direction), optimizer)
           else
               moi_lmo(randn!(direction), optimizer)
           end
end

julia> @btime f($optimizer, true)
  15.907 ms (27 allocations: 552.14 KiB)

julia> @btime f($optimizer, false)
  24.955 ms (20034 allocations: 1.33 MiB)

Manual version:

The left 2/3 is on MA:

packages/MutableArithmetics/UtY4H/src/linear_algebra.jl:427, operate [inlined]
packages/MutableArithmetics/UtY4H/src/reduce.jl:34, MethodInstance for MutableArithmetics.fused_map_reduce(::typeof(MutableArithmetics.add_dot), ::Matrix{Float64}, ::Matrix{MathOptInterface.VariableIndex})

The split above is spent roughly 50-50 on these two:

/home/mbesancon/.julia/packages/MutableArithmetics/UtY4H/src/interface.jl:496, buffered_operate_fallback!! [inlined]
/home/mbesancon/.julia/packages/MutableArithmetics/UtY4H/src/interface.jl:188, mutability [inlined]

All the the time in the manual example is spent within the HiGHS lib call

[blegat]

That's unexpected. mutability shouldn't take any time, it should be optimized out at compile time (so it should have time complexity O(0) ^^) because it only depends on types. If the compiler does not do it, maybe we can add some methods of promote_operation in MOI to help.

[odow]

The problem is really this:

julia> using BenchmarkTools

julia> function foo(n)
           d = 1.0:n
           x = MOI.VariableIndex.(1:n)
           return LinearAlgebra.dot(d, x)
       end
foo (generic function with 1 method)

julia> function bar(n)
           d = 1.0:n
           x = MOI.VariableIndex.(1:n)
           return MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(d, x), 0.0)
       end
bar (generic function with 1 method)

julia> @btime foo(100);
  102.966 μs (209 allocations: 8.25 KiB)

julia> @btime bar(100);
  473.226 ns (3 allocations: 2.67 KiB)

[odow]

Likely because of: #87

[odow]

because it only depends on types

Except that MA contains lots of functions with ::Type, which the compiler doesn't run type inference through specialize:

MutableArithmetics.jl/src/interface.jl

Line 47 in c5f4f6d

function promote_operation_fallback(op::Union{AddSubMul,typeof(add_dot)}, T::Type, args::Vararg{Type,N}) where {N}

Moving this to MA, because it seems like a performance issue there.

[matbesancon]

Do you mean it would be fixed by ::Type{T}? Why would the compiler not run type inference here?

[odow]

I meant specialize: https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing

[odow]

I didn't drill down enough to find the root cause, only to find that it's likely a problem that needs fixing in MA.

[odow]

Have a look at this:

using BenchmarkTools
import LinearAlgebra
import MathOptInterface
import MutableArithmetics

const MOI = MathOptInterface
const MA = MutableArithmetics

function direct(n)
    d = 1.0:n
    x = MOI.VariableIndex.(1:n)
    return MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(d, x), 0.0)
end

function no_rewrite(n)
    d = 1.0:n
    x = MOI.VariableIndex.(1:n)
    return sum(d[i] * x[i] for i in 1:n)
end

function rewrite(n)
    d = 1.0:n
    x = MOI.VariableIndex.(1:n)
    return MA.@rewrite(sum(d[i] * x[i] for i in 1:n))
end

function linear_algebra(n)
    d = 1.0:n
    x = MOI.VariableIndex.(1:n)
    return LinearAlgebra.dot(d, x)
end

julia> @btime direct(100);
  470.400 ns (3 allocations: 2.67 KiB)

julia> @btime no_rewrite(100);
  19.837 μs (399 allocations: 262.59 KiB)

julia> @btime rewrite(100);
  34.051 μs (9 allocations: 5.14 KiB)

julia> @btime linear_algebra(100);
  102.763 μs (209 allocations: 8.25 KiB)

That isn't great on any front!