Fix mul! performance regression for 2x2 and 3x3 matrices

[tkf]

This is a quick patch to fix JuliaLang/LinearAlgebra.jl#684.

On my laptop, the result of the benchmark @btime mul!($ou, $m1, $m2) from JuliaLang/LinearAlgebra.jl#684 is:

• Before (07a16d6): 335.059 ns (1 allocation: 16 bytes)
• After (9274757): 37.268 ns (0 allocations: 0 bytes)

[goggle]

I can confirm these benchmarks:

using LinearAlgebra
using BenchmarkTools
for ndim = 2:3
    m1 = rand(ComplexF64, ndim, ndim)
    m2 = rand(ComplexF64, ndim, ndim)
    ou = similar(m1)
    @btime mul!($ou, $m1, $m2)
end

on master:

364.115 ns (1 allocation: 16 bytes)
387.926 ns (1 allocation: 16 bytes)

on PR branch:

18.957 ns (0 allocations: 0 bytes)
35.137 ns (0 allocations: 0 bytes)

Version info:

Julia Version 1.3.1
Commit 2d5741174c (2019-12-30 21:36 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-3770 CPU @ 3.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.1 (ORCJIT, ivybridge)

[daviehh]

One additional thing: it appears while this @inline fixed the 3-argument mul!, the issue with constant propagation of 5-argument mul! may still remain:

using LinearAlgebra
using BenchmarkTools


ndim = 3

al = .5;
bt = .7im;

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

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

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

@btime mul!($ou, $m1, $m2, .5, .7im);

41.007 ns (0 allocations: 0 bytes)
334.309 ns (3 allocations: 80 bytes)
54.090 ns (0 allocations: 0 bytes)

a bit better if one also incoperates #29634 (comment)

40.915 ns (0 allocations: 0 bytes)
77.679 ns (1 allocation: 32 bytes)
54.872 ns (0 allocations: 0 bytes)

[KristofferC]

How about just copy pasting two versions of 2x2 and 3x3 mul? One for 5 arg and one for 3 arg? Relying on constant propagation seems brittle.

[daviehh]

There are many places where the 2x2 and 3x3 mul is called inside the various wrappers, so copy-pasting to all may introduce bloat/make it less maintainable for future changes... one can do the short-circuit MulAddMul alternatively #34394, is it easier to improve upon that?

[tkf]

MulAddMul thing was useful to improve performance for structured matrices; e.g., #29634 (comment)

I think the branching trick #29634 (comment) is an easier solution.

[StefanKarpinski]

Inlining all of gemm_wrapper! seems very excessive: all we really want to inline is the part that checks if the 2x2 and 3x3 cases should be called. If you want to inline this, it should be refactored so that the only part that's actually inlined is this:

if mA == 2 && nA == 2 && nB == 2
    return matmul2x2!(C, tA, tB, A, B, MulAddMul(α, β))
end
if mA == 3 && nA == 3 && nB == 3
    return matmul3x3!(C, tA, tB, A, B, MulAddMul(α, β))
end

It should be further streamlined to something like this:

if 2 == size(A, 1) && 2 == size(A, 2) && 2 == size(B, 1) && 2 == size(B, 2)
    return matmul2x2!(C, tA, tB, A, B, MulAddMul(α, β))
elseif 3 == size(A, 1) && 3 == size(A, 2) && 3 == size(B, 1) && 3 == size(B, 2)
    return matmul3x3!(C, tA, tB, A, B, MulAddMul(α, β))
end

After that there can be a non-inlined call to a function that checks dimension compatibility, looks at the stride arguments, etc. and does the relatively slow, expensive GEMM call.

[tkf]

I think 1c614e3 implements what you suggested.

I also created another PR #34601 that does not add new @inline.

[StefanKarpinski]

Great! Does this version still fix the performance issue?

[KristofferC]

Seems to have some CI problems:

Test Failed at /buildworker/worker/tester_linux64/build/share/julia/stdlib/v1.5/LinearAlgebra/test/matmul.jl:586
  Expression: mul!(copy(C), A, B, ��, 1.0) == C
   Evaluated: [NaN NaN NaN; NaN NaN NaN; NaN NaN NaN] == [1.0 1.0 1.0; 1.0 1.0 1.0; 1.0 1.0 1.0]

[StefanKarpinski]