Allocations when using `mul!` + `transpose`

[rvignolo]

Hi!

I have been looking at this and cannot figure out if this is an issue or is expected behaviour:

julia> using LinearAlgebra
julia> using BenchmarkTools

julia> @btime mul!(cache, mat, transpose(mat)) setup=(N=3; mat=rand(N,N); cache=zeros(N,N))
  59.471 ns (1 allocation: 16 bytes)
3×3 Array{Float64,2}:
 2.2594   1.68041  1.98825
 1.68041  1.56487  1.46579
 1.98825  1.46579  1.75028

julia> @btime mul!(cache, mat, transpose(mat)) setup=(N=4; mat=rand(N,N); cache=zeros(N,N))
  282.589 ns (0 allocations: 0 bytes)
4×4 Array{Float64,2}:
 1.21238   0.646434  0.315025  0.831406
 0.646434  0.829583  0.355537  0.450968
 0.315025  0.355537  0.241286  0.398706
 0.831406  0.450968  0.398706  1.09966

For N = 3 there are allocations and for N = 4 there are not. For N = 2 there are also the same amount of allocations as for N = 3. For all the other cases I tested, there are not allocations.

julia> versioninfo()
Julia Version 1.5.0
Commit 96786e22cc (2020-08-01 23:44 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin18.7.0)
  CPU: Intel(R) Core(TM) i7-4558U CPU @ 2.80GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, haswell)
Environment:
  JULIA_EDITOR = "/Applications/Visual Studio Code.app/Contents/Resources/app/bin/code"
  JULIA_NUM_THREADS =

[dkarrasch]

For two-by-two and three-by-three matrices a plain Julia multiplication method is called, starting from four-by-four BLAS methods get called. I think the allocation is due to the creation of the MulAddMul object in the syrk_wrapper!, which is not inlined and hence alpha=true and beta=false don't get constant-propagated. This behavior is a much-discussed compromise to avoid creation of the MulAddMul object in the code path to BLAS.syrk!. I'm not sure we can do much about it.

However, if you happen to know that your use case is really about 3x3 or 2x2 matrices, you may wish to consider StaticArrays.jl, or call lower level code manually:

julia> @btime LinearAlgebra.matmul3x3!(cache, 'N', 'T', mat, mat, LinearAlgebra.MulAddMul(true, false)) setup=(N=3; mat=rand(N,N); cache=zeros(N,N))
  15.792 ns (0 allocations: 0 bytes)
3×3 Array{Float64,2}:
 0.95553   0.361062  0.687642
 0.361062  0.248967  0.490268
 0.687642  0.490268  1.11368

[rvignolo]

Hi, thank you very much! I am using static arrays for these cases, as you suggested. However, I wanted to make sure nothing "wrong" was happening under the hood. Again, thank you!

[rvignolo]

@dkarrasch I was thinking about this and I could not figure out why, when using two different matrices of the same size, I don't get allocations:

julia> @btime mul!(cache, mat1, mat2) setup=(N=3; mat1=rand(N,N); mat2 = rand(N,N); cache=zeros(N,N))
  30.077 ns (0 allocations: 0 bytes)
3×3 Array{Float64,2}:
 0.225271  0.719219  0.996567
 0.173524  0.713072  0.537267
 0.141159  0.5759    0.376142

[dkarrasch]

It's because in this case, mul!(cache, mat1, mat2) calls mul!(cache, mat1, mat2, true, false), which constant-propagates true and false into the MulAddMul object passed on to gemm_wrapper!. The relevant discussion is in JuliaLang/julia#33743. I believe we missed to apply the optimization from JuliaLang/julia#34601 to the syrk and herk cases. I'm currently testing whether this helps, and will have a PR shortly.