About JuMP's speed of building an expression

WalterMadelim · May 15, 2025, 4:49am

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)

odow · May 15, 2025, 5:21am

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

Oscar_Smith · May 15, 2025, 5:22am

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

WalterMadelim · May 15, 2025, 8:27am

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.

odow · May 15, 2025, 9:17pm

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

WalterMadelim · May 15, 2025, 11:50pm

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)

odow · May 15, 2025, 11:57pm

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.

WalterMadelim · May 16, 2025, 12:42am

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

odow · May 16, 2025, 12:59am

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 ) knows, cares, or wants to write LinearAlgebra.dot(x, A, y) instead.

gdalle · May 16, 2025, 5:24am

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

WalterMadelim · May 16, 2025, 5:56am

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

gdalle · May 16, 2025, 5:57am

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

blegat · May 16, 2025, 6:15am

@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.

WalterMadelim · May 16, 2025, 7:12am

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.

blegat · May 16, 2025, 7:15am

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.

WalterMadelim · May 16, 2025, 7:59am

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.

blegat · May 16, 2025, 8:15am

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

github.com/JuliaLang/LinearAlgebra.jl

src/matmul.jl

8a1ecfb22


      
          if nC != mA
              throw(DimensionMismatch(lazy"output matrix has size: $(size(C)), but should have size $((mA, mA))"))
          end
          
          _rmul_or_fill!(C, β)
          @inbounds if !conjugate
              if aat
                  for k ∈ 1:n, j ∈ 1:m
                      αA_jk = A[j, k] * α
                      for i ∈ 1:j
                          C[i, j] += A[i, k] * αA_jk
                      end
                  end
              else
                  for j ∈ 1:n, i ∈ 1:j
                      temp = A[1, i] * A[1, j]
                      for k ∈ 2:m
                          temp += A[k, i] * A[k, j]
                      end
                      C[i, j] += temp * α
                  end

araujoms · May 19, 2025, 12:50pm

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

blegat · May 20, 2025, 4:44pm

Yes, that would be a good argument in favor

araujoms · May 21, 2025, 10:13am

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.

Topic		Replies	Views
Computing variance of variables is very slow in JuMP Optimization (Mathematical) question , jump , optimization	3	159	May 25, 2024
Efficient linear algebra for large-scale problems in JuMP Optimization (Mathematical) jump , linearalgebra , optimization	6	1536	May 8, 2018
Slow model construction in JuMP Optimization (Mathematical) jump	14	806	March 14, 2023
JuMP "getvalue" suddenly very slow Performance question	2	514	August 29, 2018
Matrix multiplication time with JuMP variables Optimization (Mathematical) jump	5	128	August 29, 2024

About JuMP's speed of building an expression

Related topics