[devmotion]

Fixes #513.

Problem

Multiplying static matrices of wide element types takes very long to compile.
Wide elements include ForwardDiff.Duals, and especially nested Duals.
For example, * of two 8×8 SMatrix{Dual{Dual{Float64,4},4}} takes about 3 s to compile, and 5-arg mul! about 4.6 s.
The reason is the amount of generated code:

• mul_loop unrolls every 8×6 block position.
• muladd_unrolled_all!/muladd_unrolled_chunks! unroll the whole product.

The generated code therefore grows with M·N and with the cost of each element operation.

Change

*, mul! and 5-arg mul! of unwrapped static matrices (SMatrix, MMatrix, SizedMatrix) with isbitstype Number elements now use one kernel:

• Bounded unrolling. Only an mr×nr block of accumulators is unrolled, with mr * nr * sizeof(eltype) ≤ MUL_TILE_BUDGET = 384 bytes. If the whole product fits into the budget, it is a single block; otherwise blocks have at most 8 rows.
  • 384 bytes is the 8×6 Float64 block of mul_loop, so Float64 keeps its block shape.
  • The loops over blocks and over the inner dimension are not unrolled.
  • The generated arithmetic code therefore consists of at most four block variants (edges included), whatever the size of the matrices and of their elements. Copying the inputs and building the result still grow linearly with the number of elements, as before.
  • Products are accumulated with muladd in the same order as before.
• Single block: if the product consists of a single block, * returns the accumulators directly, as mul_loop did. This includes products with more than 8 rows that fit into the budget, e.g. 16×2 Float64.
• mul! into an MArray: the kernel writes through a pointer inside a single GC.@preserve block. The per-element preserves of setindex! are only removed after LLVM's vectorizer runs, so they kept the blocks scalar. This alone makes mul! with Float64 2–5× faster for sizes 4–8.
• Square blocks for wide elements (suggested by @mateuszbaran): elements of at least 32 bytes are not vectorized across rows. For them the block is approximately square (e.g. 3×4 for Dual{3}, 2×2 for Dual{8}), which minimizes the loads per update. Rows are distributed evenly among the blocks, which avoids a thin block at the bottom edge.
• Unchanged: wrapped matrices (Transpose, Adjoint, Symmetric, Hermitian, triangular, Diagonal), transposed destinations, non-isbitstype elements, matrix–vector products, and the BLAS path for large BlasFloat products all keep the existing code paths.

Wrappers could be supported in a follow-up. Copying Symmetric/Hermitian matrices to a plain matrix efficiently was the open question: an unrolled copy has compile time growing with the matrix size, while a copy in a loop is slower at runtime.

Benchmarks

Apple M2 Pro, Julia 1.13.1. Compile time is the first call of a @noinline wrapper, after warming up scalar arithmetic and the other element types. Runtime is from @belapsed, averaged over two runs in fresh processes. Format: master → this PR.

Element type	n	op	compile (ms)	runtime (ns)
Float64	2	*	4.5 → 4.5	3.2 → 3.2
Float64	4	mul!	63.3 → 16.3	12.6 → 5.1
Float64	8	*	34.0 → 40.9	41.2 → 42.0
Float64	8	mul!	80.6 → 21.3	50.4 → 23.0
Float64	8	mul! (5-arg)	65.3 → 61.5	61.3 → 29.4
ComplexF64	8	*	112.4 → 91.0	219.2 → 254.2
Dual{4}	8	*	196.7 → 82.1	379.6 → 321.3
Dual{4}	8	mul!	173.6 → 56.8	428.1 → 241.8
Dual{8}	8	*	451.1 → 29.5	745.2 → 580.3
Dual{Dual{2},2}	8	mul! (5-arg)	1179.4 → 61.2	1731.2 → 1360.4
Dual{Dual{4},4}	4	*	464.7 → 29.0	363.1 → 398.4
Dual{Dual{4},4}	8	*	3310.8 → 29.9	3036.4 → 2622.7
Dual{Dual{4},4}	8	mul!	2220.5 → 31.2	5381.9 → 2189.8
Dual{Dual{4},4}	8	mul! (5-arg)	5146.3 → 47.8	6525.0 → 3645.8

Across all 180 cases of the first script (9 element types, sizes 2–8, 3 operations, plus wrappers):

• Compile time is lower in 157 of 174, and 0.20× in total.
• Runtime is more than 10% faster in 72 cases and more than 10% slower in 22. The geometric mean ratio is 0.86.

Across the 540 cases of @mateuszbaran's script with non-square shapes (9 element types, 20 shapes including 1×8×8, 8×1×8, 16×4×2, 6×10×6, 3 operations):

• Compile time is 0.13× in total.
• Runtime is more than 10% faster in 248 cases and more than 10% slower in 50 (consistently in both runs). The geometric mean ratio is 0.79.

Known regressions (runtime vs master, consistent in both runs):

case	runtime vs master	cause
Dual{Dual{4},4} 2×2 and 3×2×3	* 1.36–1.38×, mul! 1.48–1.53×	1×1 blocks, since each element is 200 bytes
outer products 8×1×8 * with ComplexF64 and Duals	1.20–1.66×	the products are not a single block, so * goes through a temporary
ComplexF64 and Dual{1} * with 7–16 rows (8×8, 7×8×8, 10×3×10, 16×4×2)	1.14–1.28×	same
Dual{3}/Dual{4} * 4×4 and 3×4×12	1.11–1.30×
small mul! (2×2, 3×3) with Float32, ComplexF64, Dual{1}, Dual{3}, Dual{4}, Dual{Dual{2},2}	1.12–1.31×
* 1×8×8 with Float32/Float64	1.19–1.28×
* Float64 8×8 compile time	34 → 41 ms

Bool products are slow to compile and to run, on master as well: 8×8 · 8×48 takes 1.8 s (master) and 2.2 s (this PR) to compile, and about 2 µs to run. true * true is a Bool but muladd of Bools is an Int, so the accumulators have mixed types. This is left for a separate fix.

Alternatives benchmarked and not adopted (same machine; each compared with this PR):

alternative	result
A single block whenever the product is at most 4× the budget (for *)	fixes ComplexF64/Dual{1} * 8×8, 7×8×8 and 8×1×8 (0.65–0.87×), but Float64 8×8 * 1.14×, Dual{3}/Dual{4} 16×4×2 1.37× and Dual{4} 10×3×10 1.33× slower; total compile time +12–23%. Undecided without more x86 numbers
Same, also for mul!	runtime worse overall (geometric mean 1.01–1.02×), e.g. Float64 8×8 mul! 22 → 35 ns; total compile time +33–43%
Compute the block width after distributing rows evenly, so wide elements use the whole budget (Dual{3} 9×9–12×12: 5×2 instead of 5×1)	runtime 1.1–1.6× slower, compile time 1.3–2.3× slower (measured before the switch to square blocks)
Single-column blocks for elements of at least 32 bytes	superseded by square blocks, which fix the same Dual{3} 5×5/6×6 regressions without slowing down 5-arg mul!
Treat elements of less than 4 bytes as 4-byte elements (8×12 blocks instead of up to 8×48)	compile time of Bool/Int8/Int16 8×8 · 8×{24,48} 2–9× lower, but runtime 1.1–1.26× slower, and 8×24 products lose the single-block path. The Bool case is better fixed by a stable accumulator type
Same with 2 bytes	smaller compile time gains, same runtime cost
One helper with do-block closures for the row and column block loops instead of the duplicated code	identical generated code, but the first multi-block product of a session compiles about 13% slower (the closures are specialized on the store closure); @nospecialize does not help

Help wanted: all numbers above are from a single Apple Silicon machine. More benchmarks on x86 (AVX2, AVX-512) would be very welcome, in particular to decide on the single block of up to 4× the budget for *. The scripts are below.

Benchmark script (square products and wrappers)

# Compile time (first call) and runtime of `*`, `mul!`, and 5-arg `mul!` for SMatrix inputs.
# Usage: julia --project=<env> bench.jl <outfile.tsv>
using StaticArrays, LinearAlgebra, ForwardDiff, BenchmarkTools
using ForwardDiff: Dual

struct T1 end; struct T2 end
D(N) = Dual{T1,Float64,N}
DD(N) = Dual{T2,Dual{T1,Float64,N},N}
rnd(::Type{T}) where {T<:Union{Float32,Float64,ComplexF64}} = rand(T)
rnd(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(rnd(V), ntuple(_ -> rnd(V), Val(N))...)
rmat(T, M, K) = SMatrix{M,K,T}(ntuple(_ -> rnd(T), Val(M*K)))
label(T) = replace(string(T), "ForwardDiff." => "", "Main." => "", r"Dual\{T\d, " => "Dual{")

@noinline f_mul(A, B) = A * B
@noinline f_mul!(C, A, B) = mul!(C, A, B)
@noinline f_mul5!(C, A, B, α, β) = mul!(C, A, B, α, β)

types = [Float32, Float64, ComplexF64, D(1), D(3), D(4), D(8), DD(2), DD(4)]
for T in types   # warm up scalar arithmetic and the generic machinery
    a = rnd(T); a * a; muladd(a, a, a); a * a + a
    f_mul(rmat(T, 1, 1), rmat(T, 1, 1))
end

open(ARGS[1], "w") do io
    println(io, "op\ttype\tn\tcompile_ms\trun_ns")
    for T in types, n in (2, 3, 4, 5, 6, 8)
        A = rmat(T, n, n); B = rmat(T, n, n); C = MMatrix(rmat(T, n, n))
        α, β = rnd(T), rnd(T)
        for (op, call, bench) in (
                ("*", () -> f_mul(A, B), () -> @belapsed(f_mul($A, $B), seconds = 0.2)),
                ("mul!", () -> f_mul!(C, A, B), () -> @belapsed(f_mul!($C, $A, $B), seconds = 0.2)),
                ("mul!5", () -> f_mul5!(C, A, B, α, β), () -> @belapsed(f_mul5!($C, $A, $B, $α, $β), seconds = 0.2)))
            tc = @elapsed call()
            tr = bench()
            println(io, op, '\t', label(T), '\t', n, '\t', round(1000tc, digits = 1), '\t', round(1e9tr, digits = 1))
            flush(io)
        end
    end
    # wrappers (Float64 and Dual{4})
    for T in (Float64, D(4)), n in (3, 4, 8), (wname, w) in (("transpose", transpose), ("UpperTriangular", UpperTriangular), ("Symmetric", Symmetric))
        A = w(rmat(T, n, n)); B = rmat(T, n, n)
        tc = @elapsed f_mul(A, B)
        tr = @belapsed f_mul($A, $B) seconds = 0.2
        println(io, "$wname(A)*B", '\t', label(T), '\t', n, '\t', round(1000tc, digits = 1), '\t', round(1e9tr, digits = 1))
        flush(io)
    end
end

Benchmark script by @mateuszbaran (non-square products)

using StaticArrays, LinearAlgebra, ForwardDiff, BenchmarkTools
using ForwardDiff: Dual

struct T1 end; struct T2 end
D(N) = Dual{T1,Float64,N}
DD(N) = Dual{T2,Dual{T1,Float64,N},N}
rnd(::Type{T}) where {T<:Union{Float32,Float64,ComplexF64}} = rand(T)
rnd(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(rnd(V), ntuple(_ -> rnd(V), Val(N))...)
rmat(T, M, K) = SMatrix{M,K,T}(ntuple(_ -> rnd(T), Val(M*K)))
label(T) = replace(string(T), "ForwardDiff." => "", "Main." => "", r"Dual\{T\d, " => "Dual{")

@noinline f_mul(A, B) = A * B
@noinline f_mul!(C, A, B) = mul!(C, A, B)
@noinline f_mul5!(C, A, B, α, β) = mul!(C, A, B, α, β)

types = [Float32, Float64, ComplexF64, D(1), D(3), D(4), D(8), DD(2), DD(4)]
for T in types   # warm up scalar arithmetic and the generic machinery
    f_mul(rmat(T, 1, 1), rmat(T, 1, 1))
end

# (M, K, N) for A::M×K, B::K×N, C::M×N
shapes = [
    [(n, n, n) for n in (2, 3, 4, 5, 6, 8)];
    [(2, 3, 4), (3, 2, 3), (4, 8, 2), (8, 2, 4),   # small, various aspect ratios
     (1, 8, 8), (8, 8, 1), (8, 1, 8),              # row × matrix, matrix × column, outer product
     (8, 8, 7), (7, 8, 8), (6, 10, 6),             # partial edge blocks, long inner dimension
     (10, 3, 10), (12, 4, 3), (3, 4, 12), (16, 4, 2)]  # more than 8 rows / wide output
]

open(ARGS[1], "w") do io
    println(io, "op\ttype\tM\tK\tN\tcompile_ms\trun_ns")
    for T in types, (M, K, N) in shapes
        A = rmat(T, M, K); B = rmat(T, K, N); C = MMatrix(rmat(T, M, N))
        α, β = rnd(T), rnd(T)
        for (op, call, bench) in (
                ("*", () -> f_mul(A, B), () -> @belapsed(f_mul($A, $B), seconds = 0.2)),
                ("mul!", () -> f_mul!(C, A, B), () -> @belapsed(f_mul!($C, $A, $B), seconds = 0.2)),
                ("mul!5", () -> f_mul5!(C, A, B, α, β), () -> @belapsed(f_mul5!($C, $A, $B, $α, $β), seconds = 0.2)))
            tc = @elapsed call()
            tr = bench()
            println(io, op, '\t', label(T), '\t', M, '\t', K, '\t', N, '\t', round(1000tc, digits = 1), '\t', round(1e9tr, digits = 1))
            flush(io)
        end
    end
end

Behaviour changes

• mul! results can differ from previous versions in the last bit. Products are accumulated in the same order with muladd, but muladd lets the compiler decide per call site whether to fuse multiply and add.
• mul! into a destination with a different element type converts each entry as setindex! does, e.g. an InexactError for a non-integer result in an Int destination.

Tests

New tests in test/matrix_multiply.jl:

• mul_tile block shapes, including single blocks with more than 8 rows and the square blocks with evenly distributed rows for elements of at least 32 bytes;
• exact results with Int and Complex{Int128} elements for sizes with a single block (also with more than 8 rows) and with edge blocks (including a loop over row blocks followed by a block at the bottom edge), for SMatrix/MMatrix/SizedMatrix inputs and MMatrix/SizedMatrix destinations, 3- and 5-arg;
• zero allocations;
• conversion of elements, and aliasing (mul!(A, A, A));
• an element type larger than the budget;
• Bool and Int8 elements, including accumulators that muladd promotes from Bool to Int;
• empty matrices and a zero inner dimension.

🤖 Generated with Claude Code

[mateuszbaran]

This looks like something generally reasonable but it would be nice if we could get rid of some of the worst runtime regressions. I've used Claude to extended the benchmark with non-square matrices and here are the results:

Total compile ms: master 100896, PR 19418

Consistent run-time regressions (>10%):
  *      Float32                       1x8x8  8.0 -> 11.6 ns (1.63x)
  *      Float32                       7x8x8  40.1 -> 48.3 ns (1.21x)
  mul!   Float32                       7x8x8  37.4 -> 48.3 ns (1.29x)
  *      Float32                      3x4x12  16.1 -> 21.2 ns (1.32x)
  *      Float32                      16x4x2  3.9 -> 8.8 ns (2.26x)
  *      Float64                       4x8x2  3.4 -> 4.3 ns (1.26x)
  *      Float64                       1x8x8  6.5 -> 11.6 ns (1.81x)
  *      Float64                      3x4x12  17.1 -> 20.7 ns (1.21x)
  *      Float64                      16x4x2  7.0 -> 8.6 ns (1.26x)
  mul!5  ComplexF64                    2x3x4  15.2 -> 17.5 ns (1.15x)
  *      ComplexF64                    8x1x8  21.1 -> 29.2 ns (1.39x)
  *      ComplexF64                   16x4x2  49.8 -> 55.5 ns (1.13x)
  mul!   Dual{Float64, 1}              3x3x3  8.4 -> 9.9 ns (1.18x)
  mul!   Dual{Float64, 1}              8x8x8  178.3 -> 218.5 ns (1.23x)
  mul!   Dual{Float64, 1}              2x3x4  7.1 -> 8.9 ns (1.27x)
  mul!5  Dual{Float64, 1}              8x8x1  27.4 -> 33.9 ns (1.24x)
  *      Dual{Float64, 1}              8x1x8  43.9 -> 49.7 ns (1.14x)
  mul!   Dual{Float64, 1}              8x8x7  157.3 -> 181.9 ns (1.17x)
  mul!   Dual{Float64, 1}              7x8x8  153.8 -> 173.5 ns (1.13x)
  mul!   Dual{Float64, 1}             6x10x6  123.3 -> 142.2 ns (1.15x)
  mul!   Dual{Float64, 1}             12x4x3  55.6 -> 64.9 ns (1.17x)
  *      Dual{Float64, 1}             16x4x2  38.5 -> 78.2 ns (2.05x)
  mul!   Dual{Float64, 1}             16x4x2  47.3 -> 62.9 ns (1.34x)
  mul!   Dual{Float64, 3}              2x3x4  13.0 -> 14.6 ns (1.12x)
  mul!   Dual{Float64, 3}              3x2x3  10.1 -> 11.8 ns (1.18x)
  mul!5  Dual{Float64, 3}              8x2x4  82.1 -> 100.9 ns (1.25x)
  mul!5  Dual{Float64, 3}             6x10x6  332.2 -> 369.9 ns (1.15x)
  mul!5  Dual{Float64, 3}             12x4x3  140.4 -> 156.5 ns (1.11x)
  mul!5  Dual{Float64, 3}             16x4x2  122.0 -> 140.0 ns (1.15x)
  mul!   Dual{Float64, 4}              2x2x2  5.5 -> 6.2 ns (1.13x)
  mul!   Dual{Float64, 4}              3x3x3  21.9 -> 25.5 ns (1.19x)
  mul!   Dual{Float64, 4}              4x4x4  56.5 -> 65.1 ns (1.16x)
  mul!   Dual{Float64, 4}              2x3x4  17.3 -> 25.2 ns (1.47x)
  mul!5  Dual{Float64, 4}              8x2x4  94.8 -> 107.0 ns (1.15x)
  mul!   Dual{Float64, 4}            10x3x10  281.9 -> 352.3 ns (1.28x)
  mul!   Dual{Float64, 4}             12x4x3  129.6 -> 153.6 ns (1.23x)
  *      Dual{Float64, 8}              3x3x3  40.3 -> 52.0 ns (1.31x)
  mul!   Dual{Float64, 8}              3x3x3  32.4 -> 37.1 ns (1.19x)
  *      Dual{Float64, 8}              4x4x4  93.2 -> 125.5 ns (1.35x)
  mul!   Dual{Float64, 8}              4x4x4  85.6 -> 101.5 ns (1.21x)
  mul!   Dual{Float64, 8}              6x6x6  294.2 -> 323.8 ns (1.12x)
  mul!   Dual{Float64, 8}              8x8x8  680.8 -> 814.9 ns (1.20x)
  *      Dual{Float64, 8}              2x3x4  38.0 -> 53.3 ns (1.42x)
  mul!   Dual{Float64, 8}              2x3x4  26.0 -> 36.1 ns (1.40x)
  *      Dual{Float64, 8}              3x2x3  25.2 -> 31.5 ns (1.25x)
  *      Dual{Float64, 8}              4x8x2  103.2 -> 117.8 ns (1.14x)
  mul!   Dual{Float64, 8}              4x8x2  88.1 -> 102.6 ns (1.17x)
  *      Dual{Float64, 8}              8x2x4  72.1 -> 95.6 ns (1.49x)
  mul!5  Dual{Float64, 8}              8x2x4  122.2 -> 182.6 ns (1.50x)
  *      Dual{Float64, 8}              8x8x1  88.8 -> 118.6 ns (1.34x)
  mul!   Dual{Float64, 8}              8x8x1  80.1 -> 101.8 ns (1.29x)
  mul!5  Dual{Float64, 8}              8x1x8  265.3 -> 324.6 ns (1.27x)
  mul!   Dual{Float64, 8}              8x8x7  618.7 -> 710.6 ns (1.18x)
  mul!   Dual{Float64, 8}             12x4x3  203.9 -> 233.0 ns (1.25x)
  *      Dual{Float64, 8}             16x4x2  220.1 -> 253.1 ns (1.20x)
  mul!   Dual{Float64, 8}             16x4x2  173.5 -> 211.4 ns (1.28x)
  mul!   Dual{Dual{Float64, 2}, 2}     2x2x2  18.7 -> 25.4 ns (1.37x)
  *      Dual{Dual{Float64, 2}, 2}     2x3x4  62.0 -> 70.0 ns (1.14x)
  *      Dual{Dual{Float64, 2}, 2}     3x2x3  43.6 -> 60.4 ns (1.44x)
  mul!   Dual{Dual{Float64, 2}, 2}     3x2x3  34.3 -> 42.9 ns (1.26x)
  mul!   Dual{Dual{Float64, 4}, 4}     2x2x2  64.0 -> 78.7 ns (1.25x)
  *      Dual{Dual{Float64, 4}, 4}     3x2x3  160.6 -> 195.4 ns (1.22x)
  mul!   Dual{Dual{Float64, 4}, 4}     3x2x3  146.2 -> 166.6 ns (1.24x)
Consistent compile-time regressions (>10%):
  *      Float32                       8x8x1  10.6 -> 12.3 ms (1.17x)
  mul!5  Float32                       8x8x1  13.8 -> 15.6 ms (1.21x)
  *      Float32                      12x4x3  22.1 -> 27.4 ms (1.32x)
  *      Float32                      16x4x2  26.5 -> 56.4 ms (2.88x)
  *      Float64                       8x8x8  32.5 -> 40.0 ms (1.28x)
  *      Float64                       2x3x4  10.1 -> 11.2 ms (1.12x)
  mul!5  Float64                       8x8x1  13.1 -> 15.5 ms (1.21x)
  *      Float64                       8x1x8  19.7 -> 29.4 ms (1.51x)
  *      Float64                       8x8x7  28.0 -> 36.2 ms (1.31x)
  *      Float64                       7x8x8  32.2 -> 38.8 ms (1.26x)
  *      Float64                     10x3x10  49.3 -> 61.7 ms (1.28x)
  *      Float64                      12x4x3  22.5 -> 28.0 ms (1.29x)
  *      Float64                      16x4x2  20.9 -> 24.4 ms (1.21x)
  *      ComplexF64                    5x5x5  35.4 -> 45.2 ms (1.31x)
  *      ComplexF64                    6x6x6  41.3 -> 51.6 ms (1.26x)
  *      ComplexF64                    8x8x8  80.5 -> 97.7 ms (1.31x)
  mul!   ComplexF64                    8x8x1  17.2 -> 21.4 ms (1.28x)
  mul!5  ComplexF64                    8x8x1  18.6 -> 22.6 ms (1.25x)
  *      ComplexF64                    8x1x8  41.9 -> 51.9 ms (1.26x)
  *      ComplexF64                   6x10x6  49.8 -> 57.7 ms (1.37x)
  mul!5  ComplexF64                   12x4x3  54.7 -> 68.1 ms (1.36x)
  *      Dual{Float64, 1}              5x5x5  37.1 -> 52.1 ms (1.48x)
  mul!5  Dual{Float64, 1}              3x2x3  27.2 -> 44.3 ms (1.85x)
  *      Dual{Float64, 1}              8x2x4  37.6 -> 47.0 ms (1.26x)
  mul!   Dual{Float64, 1}              8x8x1  17.6 -> 23.2 ms (1.36x)
  mul!5  Dual{Float64, 1}              8x8x1  18.5 -> 23.5 ms (1.31x)
  *      Dual{Float64, 1}             6x10x6  48.8 -> 57.6 ms (1.18x)
  *      Dual{Float64, 1}             12x4x3  48.2 -> 64.3 ms (1.33x)
  mul!5  Dual{Float64, 1}             12x4x3  47.5 -> 65.9 ms (1.41x)
  *      Dual{Float64, 3}              2x2x2  13.9 -> 15.4 ms (1.13x)
  *      Dual{Float64, 3}              4x4x4  36.4 -> 45.3 ms (1.26x)
  *      Dual{Float64, 3}              1x8x8  24.1 -> 26.8 ms (1.12x)
  mul!   Dual{Float64, 3}              8x8x1  21.1 -> 29.5 ms (1.73x)
  mul!5  Dual{Float64, 3}              8x8x1  30.9 -> 39.5 ms (1.51x)
  mul!   Dual{Float64, 3}             3x4x12  45.3 -> 57.6 ms (1.28x)
  *      Dual{Float64, 4}              3x3x3  28.8 -> 38.2 ms (1.33x)
  mul!5  Dual{Float64, 4}              4x8x2  31.4 -> 39.1 ms (1.27x)
  mul!   Dual{Float64, 4}              8x8x1  18.2 -> 35.4 ms (1.97x)
  mul!5  Dual{Float64, 4}              8x8x1  26.7 -> 45.8 ms (1.75x)
  *      Dual{Float64, 8}              1x8x8  41.5 -> 56.0 ms (1.35x)
  mul!5  Dual{Float64, 8}              1x8x8  53.3 -> 65.7 ms (1.26x)
  mul!   Dual{Float64, 8}              8x8x1  19.6 -> 27.7 ms (1.48x)
  *      Dual{Dual{Float64, 2}, 2}     1x8x8  69.6 -> 77.9 ms (1.21x)
  mul!5  Dual{Dual{Float64, 2}, 2}     1x8x8  90.4 -> 100.1 ms (1.19x)
  mul!   Dual{Dual{Float64, 2}, 2}     8x8x1  20.9 -> 35.2 ms (1.82x)
  mul!   Dual{Dual{Float64, 4}, 4}     8x8x1  28.0 -> 32.5 ms (1.31x)

New benchmark script:

using StaticArrays, LinearAlgebra, ForwardDiff, BenchmarkTools
using ForwardDiff: Dual

struct T1 end; struct T2 end
D(N) = Dual{T1,Float64,N}
DD(N) = Dual{T2,Dual{T1,Float64,N},N}
rnd(::Type{T}) where {T<:Union{Float32,Float64,ComplexF64}} = rand(T)
rnd(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(rnd(V), ntuple(_ -> rnd(V), Val(N))...)
rmat(T, M, K) = SMatrix{M,K,T}(ntuple(_ -> rnd(T), Val(M*K)))
label(T) = replace(string(T), "ForwardDiff." => "", "Main." => "", r"Dual\{T\d, " => "Dual{")

@noinline f_mul(A, B) = A * B
@noinline f_mul!(C, A, B) = mul!(C, A, B)
@noinline f_mul5!(C, A, B, α, β) = mul!(C, A, B, α, β)

types = [Float32, Float64, ComplexF64, D(1), D(3), D(4), D(8), DD(2), DD(4)]
for T in types   # warm up scalar arithmetic and the generic machinery
    f_mul(rmat(T, 1, 1), rmat(T, 1, 1))
end

# (M, K, N) for A::M×K, B::K×N, C::M×N
shapes = [
    [(n, n, n) for n in (2, 3, 4, 5, 6, 8)];
    [(2, 3, 4), (3, 2, 3), (4, 8, 2), (8, 2, 4),   # small, various aspect ratios
     (1, 8, 8), (8, 8, 1), (8, 1, 8),              # row × matrix, matrix × column, outer product
     (8, 8, 7), (7, 8, 8), (6, 10, 6),             # partial edge blocks, long inner dimension
     (10, 3, 10), (12, 4, 3), (3, 4, 12), (16, 4, 2)]  # more than 8 rows / wide output
]

open(ARGS[1], "w") do io
    println(io, "op\ttype\tM\tK\tN\tcompile_ms\trun_ns")
    for T in types, (M, K, N) in shapes
        A = rmat(T, M, K); B = rmat(T, K, N); C = MMatrix(rmat(T, M, N))
        α, β = rnd(T), rnd(T)
        for (op, call, bench) in (
                ("*", () -> f_mul(A, B), () -> @belapsed(f_mul($A, $B), seconds = 0.2)),
                ("mul!", () -> f_mul!(C, A, B), () -> @belapsed(f_mul!($C, $A, $B), seconds = 0.2)),
                ("mul!5", () -> f_mul5!(C, A, B, α, β), () -> @belapsed(f_mul5!($C, $A, $B, $α, $β), seconds = 0.2)))
            tc = @elapsed call()
            tr = bench()
            println(io, op, '\t', label(T), '\t', M, '\t', K, '\t', N, '\t', round(1000tc, digits = 1), '\t', round(1e9tr, digits = 1))
            flush(io)
        end
    end
end

versioninfo:

julia> versioninfo()
Julia Version 1.13.1
Commit 96ca370cf0e (2026-09-25 19:34 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 16 × 12th Gen Intel(R) Core(TM) i5-12600H
  WORD_SIZE: 64
  LLVM: libLLVM-20.1.8 (ORCJIT, alderlake)
  GC: Built with stock GC
Threads: 1 default, 1 interactive, 1 GC (on 16 virtual cores)

I can try getting Claude to improve that on my machine.

[devmotion]

Thanks for the benchmarks! Most of the worst cases don't show up on my machine: eg Float32 16x4x2 is 7.0 ns on master and 7.3 ns with this PR (Apple M2 Pro), so it seems x86-specific.

I tried to fully unroll small products, i.e., all loops if M * N * K * sizeof(eltype) <= 2048 (keeping the block shape), on the branch dmw/tiled-mul-unroll. On the M2 it fixes some of the cases with small K (eg * with Float64 16x4x2: 14.5 -> 11.7 ns, ComplexF64 8x1x8: 63.3 -> 44.2 ns) but overall it's a wash: in the unrolled cases the runtime is unchanged on average (41 cases >10% faster, 50 >10% slower, mostly Float32 and Dual{Float64,3}/Dual{Float64,4}), and compile times increase by ~35%. Total compile time with your script is 153 s (master), 24 s (this PR), and 30 s (with unrolling).

Could you run your script with that branch on your machine? If it helps on x86, we could include it, possibly with a different threshold (MUL_UNROLL_BUDGET in src/matrix_multiply.jl). Of course, if you want to experiment with MUL_TILE_BUDGET, the row cap in mul_tile, or the 32 byte cutoff on x86, that would be very welcome as well.

[mateuszbaran]

I've let Claude try to figure it out on another machine (desktop x86) and after a lot of tests it looks like a small improvement can be achieved with a limited tweak:
tiled-mul-fix.patch

Here is the summary:

The machine this time is

Julia Version 1.13.1
Commit 96ca370cf0e (2026-09-25 19:34 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 32 × AMD Ryzen 9 9950X3D 16-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-20.1.8 (ORCJIT, znver5)
  GC: Built with stock GC

Can you test what effect that patch has on your machine?

[devmotion]

Thanks for the patch! I benchmarked it on the M2 Pro with both scripts (the one from the PR description, 180 cases, and your extended one, 540 cases). I split it into its two parts and also tried two variations. Every variant ran twice in fresh processes, and a case counts as changed only if it changed by more than 10% in both runs. (An A/A run of this PR against itself showed no such runtime changes. Total compile time, however, varied by up to 16% between runs.)

compared with this PR	runtime (geometric mean), 180 / 540 cases	faster / slower by >10%	total compile time
square blocks only	0.985 / 0.995	17/5, 46/38	+5–11%
your patch	0.990 / 0.990	19/9, 63/42	+12–23%
your patch, single block also in mul!	1.017 / 1.012	13/21, 61/67	+33–43%
square blocks, single block only within the budget	0.963 / 0.977	15/5, 49/30	−4 to −7%

What I saw:

• Square blocks help 5-arg mul! with wide elements a lot (e.g. Dual{3} 6×6: 221 → 148 ns, Dual{4} 8×8: 420 → 302 ns), and they fix the Dual{3} 5×5/6×6 regressions. They make * and 3-arg mul! with 7–8 rows 10–25% slower (e.g. Dual{3} 8×8 mul!: 210 → 255 ns), but overall they're a clear improvement.
• A single block of up to 4× the budget fixes several * cases (ComplexF64 and Dual{1} 8×8, 7×8×8, 8×1×8: 0.65–0.87×). On the M2 it also makes others slower: Float64 8×8 * 40 → 46 ns, Dual{3}/Dual{4} 16×4×2 1.37×, Dual{4} 10×3×10 1.33×. Extending it to mul! is clearly worse (e.g. Float64 8×8 mul!: 22 → 35 ns), so the large blocks seem to spill.
• The 16×4×2 regressions came from the 8-row limit, not from the budget: 16×2 Float64 accumulators are only 256 bytes. So the variant I pushed applies the single-block rule only when the whole product fits within the budget. It does that for both * and mul!: Float64 16×4×2 * 14.4 → 10.2 ns, mul! 11.4 → 7.1 ns.

So I pushed your square blocks, plus the single-block rule within the budget, in mul_tile (no separate function needed). I included your tests, adjusted to the 1× rule, and the 17×3×13 case. The code is identical to the variant in the last row of the table.

Could you run your script on this version on the Ryzen? It would be good to know whether the 4× rule helps on x86 enough to be worth it for *. On the M2 it's a mix of wins and losses.

[mateuszbaran]

I've run a long job with benchmarks and below is the summary. I think this PR is already good enough for smaller element sizes, and here Claude proposed a slightly different rule to tackle larger sizes. If that also works for you then I think we could go with that or some slight variation that works better for M2. We definitely can't have a 100% win without any regressions here without much more complex kernel selection/generation logic, and I'm not sure if that would be worth the necessary effort.

--- Claude output here

Block sizes of mul_tiled! for wide elements (≥ 32 bytes)

Benchmarks of mul_tile (src/matrix_multiply.jl) at commit 77040a8 ("Use square blocks for wide elements and single blocks within the budget").

Machine: AMD Ryzen 9 9950X3D, Julia 1.13.1, one thread, pinned to one core (taskset -c 4 julia -t1).

Timings: each value is the minimum of two runs. Runs of the variants being compared alternated (A, B, A, B), and the median run-to-run variation was 0.5%.

Summary

• Without the single-block rule (M * N * s <= budget && return M, N), the cases it affects get 1.25× slower (geometric mean). For elements under 32 bytes the rule is clearly worth keeping (e.g. 16×4×2 Float64: 2–3×). For elements of 32 bytes or more the effect is mixed.
• For elements of 32 bytes or more, single-row blocks (1, nr) were faster than both the approximately square blocks and the single block. The exception is N == 1, where single-row blocks degenerate to 1×1.
• Proposed rule (patch below): for s >= 32 && N > 1, use (1, min(N, 4 * budget ÷ (3 * s))). Otherwise use the existing rule without the square/even-rows special case.
• In an out-of-sample run on bench_dmw.jl (wide types only, plus 7 extra shapes), the proposed rule is 1.105× faster than 77040a8 (geometric mean over 486 cases). 235 cases are more than 10% faster and 65 are more than 10% slower. First-call compile time is about 13% lower.
• Elements under 32 bytes are not affected.

1. Is the single-block rule beneficial?

I ran bench_dmw.jl with and without the line M * N * s <= budget && return M, N. The line changes the generated code in only 24 of the 540 cases. In the other 516 the geometric-mean ratio was 0.998, which is a useful control.

Run times in ns, written as "with rule → without rule". The ratio is (without rule) / (with rule), so > 1 means the rule helps.

case	tile with / without	*	mul!	mul!5
Float32 16×4×2	(16,2) / (8,2)	2.7 → 5.2 (1.93)	2.4 → 4.9 (2.04)	3.5 → 8.6 (2.46)
Float64 16×4×2	(16,2) / (8,2)	5.0 → 9.8 (1.96)	3.0 → 9.3 (3.10)	3.8 → 11.4 (3.00)
Float32 12×4×3	(12,3) / (8,3)	5.9 → 5.8 (0.98)	4.3 → 3.6 (0.84)	5.2 → 4.6 (0.88)
Float64 12×4×3	(12,3) / (8,3)	7.1 → 6.6 (0.93)	4.7 → 4.0 (0.85)	5.5 → 6.0 (1.09)
Dual{Float64,3} 4×8×2	(4,2) / (2,2)	54.1 → 47.2 (0.87)	54.2 → 46.8 (0.86)	65.5 → 50.0 (0.76)
Dual{Float64,3} 8×8×1	(8,1) / (3,1)	59.0 → 56.5 (0.96)	58.7 → 56.2 (0.96)	33.4 → 39.8 (1.19)
Dual{Float64,4} 4×8×2	(4,2) / (2,2)	53.3 → 64.9 (1.22)	53.0 → 62.6 (1.18)	47.1 → 45.7 (0.97)
Dual{Float64,4} 8×8×1	(8,1) / (3,1)	46.6 → 83.0 (1.78)	45.7 → 78.0 (1.71)	64.0 → 46.9 (0.73)

Over these 24 cases, removing the rule makes them 1.25× slower (geometric mean).

Caveat: the small Float32/Float64 cases are not reliable at the level of single cases. Some cases with identical generated code differed by up to 2× between the two package versions (e.g. * Float32 8×1×8: 7.8 ns vs 4.0 ns), and the difference repeated in every run. That points to code placement, not tiling. The Dual timings were stable to within 1%.

2. Exhaustive tile sweep for wide elements (in-sample)

I patched a copy of the package so that mul_tile returns a forced tile (MUL_TILE_OVERRIDE). Then I timed every tile with mr ≤ min(M, 8), nr ≤ N, mr * nr ≤ 12 and mr * nr * s ≤ 2048. Each tile used a fresh Dual tag type, so it got freshly generated code, and every result was checked against Matrix(A) * Matrix(B).

• Types: Dual{Float64,N} for N = 3, 4, 6, 8 (32, 40, 56 and 72 bytes), Dual{Dual{Float64,2},2} (72 bytes) and Dual{Dual{Float64,4},4} (200 bytes).
• Shapes (M, K, N): 2×2×2, 3×3×3, 4×4×4, 5×5×5, 2×3×4, 3×2×3, 4×8×2, 2×8×4, 8×2×4, 8×8×1, 1×8×8, 8×1×8, 6×4×2, 3×8×3, 4×4×3, 8×8×8.
• Ops: *, mul!, mul!5.
• In total, 288 (type, shape, op) cases.

Each candidate rule is scored as the geometric mean, over the 288 cases, of the time of its tile divided by the time of the fastest tile for that case.

rule for s >= 32	slowdown vs. best tile
current (single block if within budget, else near-square)	1.297
current without the single block	1.300
near-square, budget 64 … 768	1.27 … 1.35
column blocks (min(M, 8, b ÷ s), 1), b = 32 … 1024	≥ 1.30
fill a row first, then extend over rows (nr = min(N, b ÷ s), mr = (b ÷ s) ÷ nr)	≥ 1.235
always (1, 1)	1.35
row blocks (1, min(N, b ÷ s)), b = 384	1.216
row blocks (1, min(N, b ÷ s)), b = 512	1.194
row blocks, b = 512, plus single block if M * N * s ≤ 128 … 384	1.195 … 1.219
best fixed tile for all cases, (1, 8) clamped to the shape	1.201

Single-row blocks win across all element types except Dual{Dual{Float64,2},2}, where they were slightly worse (1.272 vs 1.245).

The tile landscape is very irregular. The best tile often differs between *, mul! and mul!5 for the same shape, and there are large outliers that probably come from LLVM's loop vectorizer. For example, * with Dual{Float64,3} 8×8×1 and a 1×1 tile takes 15.8 ns, while every other tile takes 50 ns or more. Optimizing for single cases is therefore pointless, so I only compared simple rules with few parameters.

3. Out-of-sample validation

Benchmark: bench_dmw.jl with types = [D(3), D(4), D(6), D(8), DD(2), DD(4)] and 7 shapes added that the sweep did not tune on: (4, 2, 8), (2, 4, 8), (7, 3, 5), (5, 7, 3), (1, 4, 6), (6, 4, 1), (4, 6, 4). That gives 27 shapes × 6 types × 3 ops = 486 cases.

Variants (77040a8 as the baseline):

• row384: s >= 32 && return 1, min(N, max(1, budget ÷ s))
• row512: same with 4 * budget ÷ 3
• hybrid (proposed): row512, but only for N > 1

Speedup = t(77040a8) / t(variant):

	row384	row512	hybrid
geometric mean, 486 cases	1.064	1.090	1.105
> 10% faster / > 10% slower	175 / 74	238 / 78	235 / 65
compile time of first call (geometric-mean ratio)	1.37	1.22	1.15

Why N == 1 is excluded. I split the row512 results by whether 77040a8 uses a single block and whether N == 1. N == 1 was the only consistent loss: geometric means of 0.87, 0.82 and 0.77 for *, mul! and mul!5 when 77040a8 used a single block. Single-row blocks degenerate to 1×1 blocks there, so the hybrid falls back to the existing rule for N == 1.

Hybrid speedup by type and op:

group	geometric mean	min	max
Dual{Float64,3}	1.122	0.55	2.74
Dual{Float64,4}	1.164	0.56	2.64
Dual{Float64,6}	1.123	0.71	2.55
Dual{Float64,8}	1.098	0.55	1.60
Dual{Dual{Float64,2},2}	1.036	0.68	1.71
Dual{Dual{Float64,4},4}	1.093	0.85	1.34
*	1.120	0.55	2.74
mul!	1.138	0.67	2.64
mul!5	1.060	0.58	1.79

Largest regressions of the hybrid (ns):

case	77040a8	hybrid
* Dual{Float64,8} 3×2×3	32.5	59.4
* Dual{Float64,3} 3×2×3	14.4	26.0
* Dual{Float64,4} 3×2×3	17.4	31.1
mul!5 Dual{Float64,4} 3×3×3	25.9	44.5
* Dual{Float64,8} 8×1×8	128.9	218.8
mul!5 Dual{Float64,4} 16×4×2	98.5	157.6
* Dual{Float64,8} 4×2×8	107.5	166.2
* Dual{Float64,4} 3×3×3	23.3	36.0

Most of these are small * products with a small inner dimension, where the single block of 77040a8 (or its square block) happens to vectorize well. Restoring the single block for such products did not help in the sweep, so the patch has no special case for them.

Caveats

• Effective budget. For elements of 32 bytes or more, the blocks may hold up to 4/3 * MUL_TILE_BUDGET bytes of accumulators, so the invariant mr * nr * s <= budget no longer holds for them. With the plain budget (row384), the speedup drops from 1.090 to 1.064, mostly because Dual{Dual{Float64,4},4} (200 bytes) falls back to 1×1 blocks.
• Mechanism unclear. I don't have a confirmed explanation of why single rows are faster. The argument for square blocks (fewer loads per update) does not seem to matter for Duals, whose arithmetic dominates the loads.
• One machine. The results come from a single Zen 5 machine (AVX-512) and one Julia version. They should be checked on other CPUs, e.g. AVX2-only ones.
• Test coverage of edges. With the patch, the Complex{Int128} (32 bytes) exact-arithmetic tests no longer produce blocks at the bottom edge (mr == 1), and produce right-edge blocks only for N > 16. The edge code is still covered by the Int cases.

Patch

The patch is against 77040a8 and is also in benchmark/mul_tile_rows.patch. All 1988 tests in test/matrix_multiply.jl pass with it.

mul_tile_rows.patch

diff --git a/src/matrix_multiply.jl b/src/matrix_multiply.jl
index a35dcce..618df56 100644
--- a/src/matrix_multiply.jl
+++ b/src/matrix_multiply.jl
@@ -260,23 +260,21 @@ const MUL_TILE_BUDGET = 384
 Size `(mr, nr)` of the block of accumulators of `mul_tiled!` for a product of size `M`×`N`
 with elements of `s >= 1` bytes.
 
-The block satisfies `mr * nr * s <= budget`, unless a single element exceeds the budget
-(then `mr = nr = 1`). If the whole product fits into the budget, it is a single block.
-Otherwise the block has at most 8 rows.
-
 Elements of at least 32 bytes (the size of AVX2 registers) are not vectorized across rows.
-For them the block is approximately square, which minimizes the number of loads of elements
-of the factors per update of an accumulator (`mr + nr` loads for `mr * nr` updates), and
-the rows are distributed evenly among the blocks: Otherwise the block at the bottom edge
-might consist of only a few rows, i.e., few independent accumulators.
+For them, unless `N == 1`, the block is a single row with at most `4 * budget ÷ (3 * s)`
+columns. In benchmarks with `ForwardDiff.Dual` numbers, such blocks were faster than
+approximately square blocks (which minimize the number of loads of elements of the factors
+per update of an accumulator) and than a single block for products within the budget.
+
+Otherwise the block satisfies `mr * nr * s <= budget`, unless a single element exceeds the
+budget (then `mr = nr = 1`). If the whole product fits into the budget, it is a single
+block. Otherwise the block has at most 8 rows.
 """
 function mul_tile(M::Int, N::Int, s::Int, budget::Int)
+    s >= 32 && N > 1 && return 1, min(N, max(1, 4 * budget ÷ (3 * s)))
     M * N * s <= budget && return M, N
-    mr = min(M, 8, max(1, s >= 32 ? isqrt(budget ÷ s) : budget ÷ s))
+    mr = min(M, 8, max(1, budget ÷ s))
     nr = min(N, max(1, budget ÷ (s * mr)))
-    if s >= 32
-        mr = cld(M, cld(M, mr))
-    end
     return mr, nr
 end
 
diff --git a/test/matrix_multiply.jl b/test/matrix_multiply.jl
index 010c623..f1194c6 100644
--- a/test/matrix_multiply.jl
+++ b/test/matrix_multiply.jl
@@ -455,22 +455,24 @@ Base.:*(a::WideInt, b::WideInt) = WideInt(a.x * b.x)
             @test StaticArrays.mul_tile(8, 8, 8, budget) == (8, 6)
             @test StaticArrays.mul_tile(8, 8, 1000, budget) == (1, 1)
             @test StaticArrays.mul_tile(9, 9, 8, budget) == (8, 6)
-            @test StaticArrays.mul_tile(6, 6, 72, budget) == (2, 2)
-            @test StaticArrays.mul_tile(7, 8, 32, budget) == (3, 4)
+            @test StaticArrays.mul_tile(6, 6, 72, budget) == (1, 6)
+            @test StaticArrays.mul_tile(7, 8, 32, budget) == (1, 8)
+            @test StaticArrays.mul_tile(8, 8, 200, budget) == (1, 2)
+            @test StaticArrays.mul_tile(8, 1, 32, budget) == (8, 1)
+            @test StaticArrays.mul_tile(8, 1, 72, budget) == (5, 1)
             @test StaticArrays.mul_tile(16, 3, 8, budget) == (16, 3)
             @test StaticArrays.mul_tile(16, 4, 8, budget) == (8, 4)
             for M in (1:10..., 16, 17), N in 1:10, s in (1, 4, 8, 16, 32, 72, 200, 1000)
                 mr, nr = StaticArrays.mul_tile(M, N, s, budget)
-                if M * N * s <= budget
+                if s >= 32 && N > 1
+                    # Single row
+                    @test mr == 1 && nr == min(N, max(1, 4 * budget ÷ (3 * s)))
+                elseif M * N * s <= budget
                     # Single block
                     @test (mr, nr) == (M, N)
                 else
                     @test 1 <= mr <= min(M, 8) && 1 <= nr <= N
                     @test mr * nr * s <= budget || mr == nr == 1
-                    if s >= 32
-                        # Approximately square blocks, whose rows are distributed evenly
-                        @test cld(M, mr) == cld(M, min(M, max(1, isqrt(budget ÷ s))))
-                    end
                 end
             end
         end

Tile sweep script (needs a MUL_TILE_OVERRIDE hook in mul_tile)

Hook added to src/matrix_multiply.jl for the sweep only:

const MUL_TILE_OVERRIDE = Ref{Any}(nothing)
function mul_tile(M::Int, N::Int, s::Int, budget::Int)
    o = MUL_TILE_OVERRIDE[]; o !== nothing && return o::Tuple{Int,Int}
    # ... unchanged ...

Sweep script (MAXACC=12 julia -t1 sweep.jl out.tsv):

using StaticArrays, LinearAlgebra, ForwardDiff, BenchmarkTools
using ForwardDiff: Dual
const OV = StaticArrays.MUL_TILE_OVERRIDE
struct Tg{m,n} end; struct Tg2{m,n} end
# A fresh tag per tile forces freshly generated code for every tile
kinds = [("D3", (m,n)->Dual{Tg{m,n},Float64,3}), ("D4", (m,n)->Dual{Tg{m,n},Float64,4}),
         ("D6", (m,n)->Dual{Tg{m,n},Float64,6}), ("D8", (m,n)->Dual{Tg{m,n},Float64,8}),
         ("DD2", (m,n)->Dual{Tg2{m,n},Dual{Tg{m,n},Float64,2},2}), ("DD4", (m,n)->Dual{Tg2{m,n},Dual{Tg{m,n},Float64,4},4})]
rnd(::Type{Float64}) = rand()
rnd(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(rnd(V), ntuple(_ -> rnd(V), Val(N))...)
rmat(T, M, K) = SMatrix{M,K,T}(ntuple(_ -> rnd(T), Val(M*K)))
@noinline f_mul(A, B) = A * B
@noinline f_mul!(C, A, B) = mul!(C, A, B)
@noinline f_mul5!(C, A, B, α, β) = mul!(C, A, B, α, β)
shapes = [(2,2,2),(3,3,3),(4,4,4),(5,5,5),(2,3,4),(3,2,3),(4,8,2),(2,8,4),(8,2,4),(8,8,1),(1,8,8),(8,1,8),(6,4,2),(3,8,3),(4,4,3),(8,8,8)]
maxacc = parse(Int, get(ENV, "MAXACC", "12"))
open(ARGS[1], "w") do io
    println(io, "op\ttype\tM\tK\tN\tmr\tnr\trun_ns")
    for (kl, kT) in kinds, (M, K, N) in shapes, mr in 1:min(M, 8), nr in 1:N
        mr * nr <= maxacc || continue
        T = kT(mr, nr); sizeof(T) * mr * nr <= 2048 || continue
        A = rmat(T, M, K); B = rmat(T, K, N); C = MMatrix(rmat(T, M, N)); α, β = rnd(T), rnd(T)
        OV[] = (mr, nr)
        R = f_mul(A, B); f_mul!(C, A, B); C2 = copy(C); f_mul5!(C, A, B, α, β)
        OV[] = nothing
        ref = Matrix(A) * Matrix(B)
        (R ≈ ref && C2 ≈ ref && C ≈ C2 * β + α * ref) || error("wrong result $kl $M $K $N $mr $nr")
        for (op, t) in (("*", @belapsed(f_mul($A, $B), seconds = 0.05)),
                        ("mul!", @belapsed(f_mul!($C, $A, $B), seconds = 0.05)),
                        ("mul!5", @belapsed(f_mul5!($C, $A, $B, $α, $β), seconds = 0.05)))
            println(io, op, '\t', kl, '\t', M, '\t', K, '\t', N, '\t', mr, '\t', nr, '\t', round(1e9t, digits = 1))
        end
        flush(io)
    end
end

[devmotion]

Thanks for the thorough analysis! I ran your patch (applied to 77040a8) on the M2 Pro with the same protocol as last time: base and patch alternating, two rounds in fresh processes, and a case counts as changed only if it changed by more than 10% in both rounds. I used three scripts: the one from the PR description (180 cases), your extended one (540 cases), and your out-of-sample setup (wide types D(3), D(4), D(6), D(8), DD(2), DD(4), plus your 7 extra shapes: 486 cases).

On the M2, the single-row rule is a net loss for wide elements. The control types (Float32, Float64, ComplexF64, Dual{1}) stay at 0.97–1.01, as expected.

patch / 77040a8 (< 1 is faster)	geometric mean	faster / slower by >10% (both rounds)
PR description script, 180 cases	1.033	8 / 28
your extended script, 540 cases	1.018	40 / 70
your out-of-sample wide setup, 486 cases	1.036	64 / 122

To see where the differences come from, I split the wide setup by how the patch changes the block:

change of the block	*	mul!	mul!5
unchanged (control, 12 cases per op)	0.99	1.00	1.00
single block within the budget → single rows	1.33	1.04	1.15
approximately square → single rows	1.02	0.99	1.15
DD(4): 1×1 → 1×2	1.00	1.01	0.94
N == 1: 2×1 → up to budget ÷ s rows	0.79	0.94	0.75

So on the M2, single rows are roughly neutral for * and mul! with larger products. They lose clearly in 5-arg mul! (e.g. Dual{4} 8×8×8: 289 → 480 ns) and when they replace the single block (e.g. Dual{4} 3×3×3 *: 20 → 33 ns), consistent with your largest regressions.

The N == 1 row is interesting, though. Those cases don't use single rows at all: your patch drops the square cap, so they get blocks of min(M, 8, budget ÷ s) rows instead of about isqrt(budget ÷ s). For example, Dual{6} 8×8×1 mul!5 goes 83 → 53 ns, Dual{8} 8×8×1 mul!5 97 → 66 ns, and DD(2) 8×8×1 mul!5 158 → 113 ns. The only loss is DD(2) 6×4×1 mul! (34 → 39 ns). Since the square cap is only motivated when nr can grow as well, this makes sense to me.

I also tried to generalize this ("fill the budget"), but neither generalization holds up on the M2:

• Using the remaining budget for more rows whenever the block spans all columns also affects N = 2, 3, where it's mostly slower (e.g. Dual{3} 16×4×2 mul!: 59 → 75 ns).
• Distributing the rows evenly loses most of the N == 1 gain (4 + 4 instead of 6 + 2 rows: Dual{6} 8×8×1 mul!5 0.84 instead of 0.64), so more accumulators in the main block matter more than a balanced edge block.
• At least two accumulators for DD(4) gives only 0.95 in mul!5, but almost doubles the compile time of these cases, which runs against the goal of this PR.

So the same single-row rule is about 10% faster on Zen 5 and about 3.6% slower on the M2. Combined with the irregular tile landscape you found, I don't think we can get a clear win for wide elements on both machines with a simple rule, apart from the N == 1 case. I'd suggest keeping the current rule except for this change:

function mul_tile(M::Int, N::Int, s::Int, budget::Int)
    M * N * s <= budget && return M, N
    s >= 32 && N == 1 && return min(M, 8, max(1, budget ÷ s)), 1
    ...

Your hybrid already uses exactly these blocks for N == 1. Could you check the N == 1 cases of your hybrid run against 77040a8 on the Ryzen? If they hold up there too, I'll add the change (with tests) to this PR. Further tuning for wide elements could be left for a follow-up, ideally with data from more machines (e.g. an AVX2-only x86).

[mateuszbaran]

Thanks for testing. I've checked your suggestion on my machine but it doesn't really help. The current state (commit 77040a8) is still on average an improvement over the master branch so I think it's worth merging but first I'd like to give other people a chance to provide their input, so I'd suggest leaving this open as ready for review for a couple of days before merging.

Below is the Claude's summary.

---

Thanks for running this on the M2! I checked the N == 1 rule on the Ryzen. Unfortunately it doesn't hold up as well here.

Setup: Ryzen 9 9950X3D, Julia 1.13.1, one thread. I compared three variants, all based on 77040a8:

• base: 77040a8 as it is.
• rows: the single-row patch.
• n1: 77040a8 plus s >= 32 && N == 1 && return min(M, 8, max(1, budget ÷ s)), 1 after the single-block line.

The variants ran simultaneously, each pinned to its own core. There were 3 rounds in fresh processes, with the core assignment rotated between rounds. Each value is the minimum over rounds, and a case counts as changed only if it changed by more than 10% in all 3 rounds. The median run-to-run variation was 1.3–1.8%. I ran both bench_dmw.jl (540 cases) and the wide out-of-sample setup (486 cases).

Overall

geometric mean (< 1 is faster)	bench_dmw.jl, 540 cases	wide setup, 486 cases
rows / 77040a8	0.948 (137 faster / 36 slower)	0.911 (212 / 58)
n1 / 77040a8	0.997 (5 / 0)	0.999 (6 / 4)

The control types (Float32, Float64, ComplexF64, Dual{1}) are at 0.996–1.001 for all variants. In the cases where n1 changes the block, n1 and rows are within 1% of each other, so they really generate the same blocks there.

N == 1 cases (15 cases of the wide setup)

op	n1 / 77040a8	faster / slower by >10%	your M2 numbers
*	0.91	3 / 1	0.79
mul!	0.91	2 / 1	0.94
mul!5	1.05	1 / 2	0.75

The gains in * and mul! mostly come from DD(2). Several of your largest mul!5 wins go the other way on Zen 5, and D(6) 8×8×1 flips between ops:

case	77040a8 (ns)	n1 (ns)	ratio
DD(2) 8×8×1, mul!	225	144	0.64
DD(2) 8×8×1, *	232	153	0.66
D(6) 8×8×1, mul!5	67	50	0.75
DD(2) 8×8×1, mul!5	134	148	1.10
D(8) 8×8×1, mul!5	73	82	1.11
D(6) 8×8×1, mul!	72	87	1.20
D(8) 6×4×1, mul!5	34	41	1.22
D(6) 8×8×1, *	75	98	1.31

Compile time: in these cases the first call takes about 2.1–2.4× longer with the N == 1 rule, for example DD(2) 8×8×1 mul!5 takes 29 → 80 ms. With 2-row blocks, 8 rows are a loop over a single block body. With 5 + 3 or 6 + 2 rows, there are two separately unrolled block bodies. Do you see the same on the M2?

Single rows, split as in your table

For comparison, here are the rows / 77040a8 ratios on Zen 5, with your M2 numbers in the last column:

change of the block	*	mul!	mul!5	M2: * / mul! / mul!5
unchanged	1.00	1.00	1.00	0.99 / 1.00 / 1.00
single block → single rows	1.01	0.85	0.99	1.33 / 1.04 / 1.15
approximately square → single rows	0.89	0.88	0.92	1.02 / 0.99 / 1.15
N == 1	0.91	0.91	1.05	0.79 / 0.94 / 0.75

So the two microarchitectures disagree exactly where the M2 loses most: 5-arg mul! and products that fit into the budget.

Suggestion

The N == 1 rule is roughly neutral overall on Zen 5. It has ±30% outliers in both directions and about twice the compile time for the cases it changes. So I'd keep 77040a8 as it is, since it doesn't regress on either machine. The tuning for wide elements can be left for a follow-up with data from more machines, as you suggested (e.g. an AVX2-only x86). If you'd still prefer to include the N == 1 change, I'm fine with that too: it's small and doesn't affect anything else.