About JuMP's speed of building an expression

I find that it seems to be not that fast for JuMP to build a “large” expression
I wonder: is this speed normal? Why? Can (or how can) it be improved somehow?

import LinearAlgebra.dot as dot
using JuMP, BenchmarkTools
N = 1000
m = JuMP.Model();
JuMP.@variable(m, x[1:N])
JuMP.@variable(m, y[1:N])
A = rand(N, N)
z, w = rand(N), rand(N)
@btime dot(z, A, w) # 252.100 μs (1 allocation: 16 bytes)
@btime ret = JuMP.@expression(m, dot(x, A, y)) # 45.682 s (22435129 allocations: 105.56 GiB)

Note the performance indices in the comment of the last line.
(Additionally, we can make another test in which A is sparse, here is a readily available code snippet)

We’re missing a method for the three-argument LinearAlgebra.dot, so it falls back to the generic version which allocates a lot.

julia> using JuMP

julia> import LinearAlgebra

julia> N = 300
300

julia> model = Model();

julia> @variable(model, x[1:N]);

julia> @variable(model, y[1:N]);

julia> A = rand(N, N);

julia> @time @expression(model, dot(x, A, y));
  1.174037 seconds (1.28 M allocations: 2.980 GiB, 32.74% gc time)

julia> @time @expression(model, dot(x, A * y));
  0.036104 seconds (4.87 k allocations: 25.038 MiB)

julia> @time @expression(model, x' * A * y);
  0.054529 seconds (4.87 k allocations: 25.038 MiB, 35.70% gc time)

julia> @which dot(x, A, y)
dot(x::AbstractVector, A::AbstractMatrix, y::AbstractVector)
     @ LinearAlgebra ~/.julia/juliaup/julia-1.10.9+0.x64.apple.darwin14/share/julia/stdlib/v1.10/LinearAlgebra/src/generic.jl:923
1 Like

how hard would it be to add 3 arg dot? It seems like it would be useful.

1 Like

In LinearAlgebra, I believe that x' * A * y are different from dot(x, A, y).
The latter has a LAPACK specialized method. x-ref my previous comment.

In LinearAlgebra, I believe that x' * A * y are different from dot(x, A, y)

They’re equivalent for real inputs (which your JuMP code is):

julia> x = rand(10); A = rand(10, 10); y = rand(10);

julia> x' * A * y
7.631533368388598

julia> LinearAlgebra.dot(x, A, y)
7.631533368388598

equivalent in results, but might not in terms of the code it traverses.
Do you have a plan to add a method for 3-arg dot?
I find that the style y' * (A' * x) might be optimal.

julia> @btime x' * A * y;
  52.078 ms (12077 allocations: 37.02 MiB)

julia> @btime y' * (A' * x);
  50.323 ms (12078 allocations: 37.02 MiB)

julia> @btime dot(x, A * y);
  51.340 ms (12076 allocations: 37.02 MiB)

julia> @btime dot(y, A' * x);
  51.997 ms (12077 allocations: 37.02 MiB)

julia> 


julia> @btime x' * A * y;
  52.874 ms (12077 allocations: 37.02 MiB)

julia> @btime y' * (A' * x);
  48.908 ms (12078 allocations: 37.02 MiB)

julia> @btime dot(x, A * y);
  49.874 ms (12076 allocations: 37.02 MiB)

julia> @btime dot(y, A' * x);
  52.799 ms (12077 allocations: 37.02 MiB)

Do you have a plan to add a method for 3-arg dot?

Nope. It’s a bit of a pain to intercept the appropriate methods as the number of arguments grow without introducing method ambiguities.

I don’t think that dot(x, A, y) is a common operation for JuMP models, and I would encourage people to use the more readable x' * A * y (this does appear often, e.g., in x' * Q * x type models). What’s more, x' * A * y is fewer characters to type.

2 Likes

I think it will be used for its being a bilinear term, which could be addressed by Gurobi.

I agree that x' * A * y is a somewhat common operation (but I assume, <5% of JuMP users). My point is that I don’t think anyone (aside from perhaps, you :smile:) knows, cares, or wants to write LinearAlgebra.dot(x, A, y) instead.

1 Like

If I were doing quadratic programming and were an experienced Julia user, I’d probably write the objective like that right away: dot(x, A, x) + dot(b, x) + c

1 Like
  1. B * x is not a scalar
  2. constant c don’t need to appear on the objective

My bad, fixed the typo to dot(b, x). And I am aware that a constant does not change a quadratic objective :wink:

@gdalle I agree it would be could be a common way users would write it. The issue is that we have a lot of methods to implement

julia> methods(LinearAlgebra.dot, (AbstractVector, AbstractMatrix, AbstractVector))
# 14 methods for generic function "dot" from LinearAlgebra:
  [1] dot(x::AbstractVector, transA::Transpose{<:Real}, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/generic.jl:950
  [2] dot(x::AbstractVector, D::Diagonal, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/diagonal.jl:960
  [3] dot(x::AbstractVector, adjA::Adjoint, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/generic.jl:949
  [4] dot(x::AbstractVector, B::Bidiagonal, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/bidiag.jl:800
  [5] dot(x::AbstractVector, A::UpperTriangular, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:706
  [6] dot(x::AbstractVector, A::Tridiagonal, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/tridiag.jl:887
  [7] dot(x::AbstractVector, H::UpperHessenberg, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/hessenberg.jl:340
  [8] dot(x::AbstractVector, A::UnitLowerTriangular, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:769
  [9] dot(x::AbstractVector, A::UnitUpperTriangular, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:727
 [10] dot(x::AbstractVector, A::LowerTriangular, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:749
 [11] dot(x::AbstractVector, S::SymTridiagonal, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/tridiag.jl:227
 [12] dot(x::AbstractVector, A::Union{Hermitian{var"#s5029", var"#s5028"}, Symmetric{var"#s5029", var"#s5028"}, Symmetric{Complex{var"#s5029"}, var"#s5028"}} where {var"#s5029"<:Real, var"#s5028"<:Diagonal}, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/special.jl:326
 [13] dot(x::AbstractVector, A::Union{Hermitian{T, S}, Hermitian{Complex{T}, S}, Symmetric{T, S}} where {T<:Real, S}, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/symmetric.jl:550
 [14] dot(x::AbstractVector, A::AbstractMatrix, y::AbstractVector)
     @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/generic.jl:931

Ideally, we should rewrite almost all of LinearAlgebra and SparseArrays in MutableArithmetics, make sure the relevant calls go towards our implementation via multiple dispatch whenever some of the element types are subtypes of AbstractMutable (this magic is done in the dispatch.jl file), make sure it works on all Julia version we support without creating any method ambiguities. That’s even more work than just reimplementing LinearAlgebra and SparseArrays ^^
In the long term, I think it would be best for these packages to have default implementations that work well with mutable types by already using add_mul!! so it would mean moving MutableArithmetics or part of it as dependency of these packages. Maybe now that these have been moved out of Julia, it can start being a possibility.

3 Likes

I didn’t get your full idea.
The dispatch of LinearAlgebra.jl had better depends on MutableArithmetics, or the MutableArithmetics follows the dispatch pattern of LinearAlgebra.jl?

PS At least in terms of documentation/tutorials, JuMP is more agreeable.

I meant that if LinearAlgebra was using MutableArithmetics internally it its default implementations then we wouldn’t need to touch at the dispatch part of LinearAlgebra and we could get rid of the dispatch.jl file of MutableArithmetics.

2 Likes

Great. I support.
In a larger sense, Matrices in LinearAlgebra.jl and SparseArrays.jl should hold more objects than only <:Real. It may also hold JuMP.VariableRef, or Symobolics.Num, etc.
The special structure is widely applicable, e.g. TriDiagonal, Symmetric, etc.
Things can be generalized. Only for regular numerical arithmetics, it then dispatch to LAPACK, or SUITESPARSE, etc.

Yes, they can hold anything. The only issue is that their default implementations use += instead or a mutable muladd. See for instance:

This would also significantly speed up linear algebra with BigFloat, no?

Yes, that would be a good argument in favor

Perhaps you should suggest it at LinearAlgebra then? I’d support it. Also, I’ve contributed a few functions to LinearAlgebra, I’d be happy to rewrite them in MutableArithmeticsform.

1 Like