Reduce matmul latency by splitting small matmul

[jishnub]

This splits the matmul2x2 and matmul3x3 into components that depend on MulAddMul and those that don't depend on it. This improves compilation time, as the MulAddMul-independent methods won't need to be recompiled in the @stable_muladdmul branches.

TTFX (each call timed in a separate session):

julia> using LinearAlgebra

julia> A = rand(2,2); B = Symmetric(rand(2,2)); C = zeros(2,2);

julia> @time mul!(C, A, B);
  1.927468 seconds (5.67 M allocations: 282.523 MiB, 12.09% gc time, 100.00% compilation time) # nightly v"1.12.0-DEV.492"
  1.282717 seconds (4.46 M allocations: 228.816 MiB, 4.58% gc time, 100.00% compilation time) # This PR

julia> A = rand(2,2); B = rand(2,2); C = zeros(2,2);

julia> @time mul!(C, A, B);
  1.653368 seconds (5.75 M allocations: 291.586 MiB, 13.94% gc time, 100.00% compilation time) # nightly
  1.148330 seconds (4.46 M allocations: 230.714 MiB, 4.47% gc time, 100.00% compilation time) # This PR

Edit: Not inlining the function seems to incur a runtime perfomance cost.

julia> using LinearAlgebra

julia> A = rand(3,3); B = rand(size(A)...); C = zeros(size(A));

julia> @btime mul!($C, $A, $B);
  23.923 ns (0 allocations: 0 bytes) # nightly
  31.732 ns (0 allocations: 0 bytes) # This PR

Adding @inline annotations resolves this difference, but this reintroduces the compilation latency. The tradeoff is perhaps ok, as users may use StaticArrays for performance-critical matrix multiplications.

[dkarrasch]

What if you added aggressive constant propagation here? tA and tB should be known as constants at the call site, because they are obtained from unpeeling outer wrappers. Does that also increase latency again?

[jishnub]

I had tried this, but this doesn't seem to change the compilation time at all. Some of the branches are certainly eliminated, but the bulk of the latency probably arises elsewhere.

This doesn't make things worse either, though, so perhaps we may include this, just in case the other issues are resolved in the future. This may also let us use lazy versions of the wrappers instead of copies, as this would be type-stable if the other branches are eliminated.

[dkarrasch]

Awesome! The compilation latency improvement should hold for any initial calls to mul!, right? Not only for small matrices, because those code paths dependent on values of objects unknown at compile time?

@nanosoldier runbenchmarks(["linalg"], vs=":master")

[dkarrasch]

@nanosoldier runbenchmarks("linalg", vs = ":master")

[nanosoldier]

Your benchmark job has completed - no performance regressions were detected. A full report can be found here.

[jishnub]

Yes, this improves the compilation time for matrices of any size

julia> A = rand(30,30); B = rand(size(A)...); C = zeros(size(A));

julia> @time mul!(C, A, B); # similar to above
  1.272428 seconds (4.25 M allocations: 221.403 MiB, 19.92% gc time, 99.93% compilation time)