mul! performance regression on master

[daviehh]

On the current master branch, mul! of small matrices can be 10x slower than 1.3 (and also allocates)

julia> versioninfo()
Julia Version 1.4.0-DEV.556
Commit 4800158ef5* (2019-12-03 21:18 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin19.0.0)
  CPU: Intel(R) Core(TM) i5-8259U CPU @ 2.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-8.0.1 (ORCJIT, skylake)

test code:

using LinearAlgebra
using BenchmarkTools

ndim = 3

m1 = rand(ComplexF64,ndim,ndim);
m2 = rand(ComplexF64,ndim,ndim);
ou = rand(ComplexF64,ndim,ndim);

@btime mul!($ou, $m1, $m2);

With the release-1.3 version,

33.471 ns (0 allocations: 0 bytes)

on master:

394.428 ns (1 allocation: 16 bytes)

While it's nano-seconds, when one has mul! in the inner-most/hot loop, this can easily translate to big performance degradation when most of the calculation is with such matrix productions. Larger matrices (e.g. ndim = 30) appears unaffected (dispatched to BLAS?) 2.732 μs (0 allocations: 0 bytes) on 1.3 and 2.774 μs (0 allocations: 0 bytes) on master.

[daviehh]

Likely happened some time between mul!/gemm_wrapper! and matmul3x3!:

@btime LinearAlgebra.gemm_wrapper!($ou, $t, $t, $m1, $m2, 1.2, .5im)

388.228 ns (3 allocations: 80 bytes)

ad = LinearAlgebra.MulAddMul(1.2, 0.5im)
t = 'N'
@btime LinearAlgebra.matmul3x3!($ou, $t, $t, $m1, $m2, $ad)

48.857 ns (0 allocations: 0 bytes)

But there's not much between:
https://github.com/JuliaLang/julia/blob/1a6baefb51170027f7ae502ae510dc5023029ed7/stdlib/LinearAlgebra/src/matmul.jl#L552-L578

---

t = 'N'
al = 1.2
bt = .5im
@btime LinearAlgebra.matmul3x3!($ou, $t, $t, $m1, $m2, LinearAlgebra.MulAddMul($al, $bt))
@btime LinearAlgebra.matmul3x3!($ou, $t, $t, $m1, $m2, LinearAlgebra.MulAddMul(1.2, .5im))

returns:

359.817 ns (3 allocations: 80 bytes)
48.784 ns (0 allocations: 0 bytes)

[dkarrasch]

So the issue seems to come from when the constructor MulAddMul is called, see JuliaLang/julia#33743 (comment). In the release version, calling mul!(C, A, B) calls mul!(C, A, B, true, false), which does

alpha, beta = promote(α, β, zero(T))
    if alpha isa T && beta isa T
        return gemm_wrapper!(C, 'N', 'N', A, B, MulAddMul(alpha, beta))
    else
        return generic_matmatmul!(C, 'N', 'N', A, B, MulAddMul(α, β))
    end

In JuliaLang/julia#33743, we construct the MulAddMul one layer below, i.e., inside the gemm_wrapper!, as suggested by @tkf. Unfortunately, the compiler does not propagate the constants, so the construction of the MulAddMul(true, false) does not happen at compile time, and then badly hits performance, as was already noticed by @Jutho in JuliaLang/julia#29634 (comment). Still, JuliaLang/julia#33743 was addressing other performance issues with adjoints/transposes and number types different from Bool and the matrix eltype. Is there any way to help constant propagation here?

[daviehh]

@dkarrasch Thanks for the comment, looks like julia thinks LinearAlgebra.MulAddMul(α, β) might not be type stable? Test with

using LinearAlgebra

ndim = 3

a = rand(ComplexF64,ndim,ndim);
b = rand(ComplexF64,ndim,ndim);
c = rand(ComplexF64,ndim,ndim);

t = 'N'
al = 1.2
bt = .5im

@code_warntype LinearAlgebra.gemm_wrapper!(c, t, t, a, b, al, bt)

[dkarrasch]

Now that you mention it, it's not surprising: the exact type (including its parameters) of MulAddMul is constructed based on the values of α and β, like iszero and isone. The return values of those two functions are AFAIK never determined by type. So maybe the issue is that by shifting the construction by one layer, we are potentially missing a function barrier that we have in the v1.3-release?

[dkarrasch]

No, actually, I think we do have function barriers, matmul2x2! and matmul3x3!, so it really seems to be the constant propagation issue, i.e., whether or not the MulAddMul object can be constructed at compile time.

[dkarrasch]

@daviehh Have you considered using StaticArrays.jl for the small matrix case?

using LinearAlgebra
using BenchmarkTools
using StaticArrays

ndim = 3

m1 = SMatrix{ndim,ndim}(rand(ComplexF64,ndim,ndim))
m2 = SMatrix{ndim,ndim}(rand(ComplexF64,ndim,ndim))
ou = MMatrix{ndim,ndim}(rand(ComplexF64,ndim,ndim))

@btime mul!($ou, $m1, $m2)

[daviehh]

That can be one weird way to get around this issue... performance seems to be the reason why julia has special code for 2x2 and 3x3 matrices multiplication in the first place due to the overhead of sending matrices to BLAS.gemm!(). Maybe rethink when to construct MulAddMul / deal with const. propagation or function barriers for now?

I tried StaticArrays before, for one of my particular project, the code need to tackle both large and small matrices, and it can be messy to have separate code and keep them synced & maintained to support both standard and static arrays.

[KristofferC]

Bump, this needs to be fixed / worked-around with some priority.

[tkf]

It looks like a single @inline fixes it (by helping const prop, I guess). See JuliaLang/julia#34384.

[daviehh]

Thanks for the fix! Since the 2x2 and 3x3 matrix multiplication is dealt with differently, can this be put in the nanosildier's benchmark script? possibly append 2 or 3 to the SIZES at

https://github.com/JuliaCI/BaseBenchmarks.jl/blob/406d5b1d11b498884e58e111d5745b29c9471d40/src/linalg/LinAlgBenchmarks.jl#L13