Skip to content

arm64: NEON small-matrix GEMM kernels for VORTEXM4 - #6155

Open
tesch1 wants to merge 4 commits into
OpenMathLib:developfrom
tesch1:sme2-small-kernels
Open

tesch1 wants to merge 4 commits into
OpenMathLib:developfrom
tesch1:sme2-small-kernels

Conversation

@tesch1

@tesch1 tesch1 commented Oct 22, 2026

Copy link
Copy Markdown
Contributor

Addresses the small-problem part of #6149.

Before: on VORTEXM4, SGEMM/DGEMM below 3000 multiply-adds ran plain C loops inside the SME2 kernel (s2_small() in sme2_gemm_impl.h), which read and write C once per depth step: 5-45 GFLOPS. Row-major SGEMM below 64^3 went to the SME1 direct kernel, which takes about 150 ns per call at 8^3 (1 GFLOPS at 4^3, 7 at 8^3). The target's small-matrix kernels (SMALL_MATRIX_OPT) could not help, because interface/gemm.c sends VORTEXM4 to the SME kernels before it checks the small-matrix permit.

After: new NEON small-matrix kernels keep blocks of C in registers over the whole depth, without packing, and reach about 100 GFLOPS (fp32) and 50 GFLOPS (fp64) from 16^3. VORTEXM4 uses them below 16000 (fp32) and 12000 (fp64) multiply-adds, where the SME2 kernels become faster.

How:

  • kernel/arm64/gemm_small_kernel_neon.h with gemm_small_kernel_{nn,nt,tn,tt}_neon.c (also the B0 variants): blocks of two vectors by four columns with lane FMAs for A op(B), alpha applied to op(B) before the products as in the reference GEMM; dot products along the depth for A^T B; plain loops for A^T B^T. Row tails of one vector or more are an overlapped block whose rows already written are not stored again.
  • gemm_small_kernel_permit_vortexm4.c, selected in KERNEL.VORTEXM4. Vector shapes (M or N = 1) stay on the SME2 kernels: forwarded to GEMV instead, LAPACK test SHS 11 (xeigtsts < nep.in, order 16, type 20) fails on M4, although the same test passes on a VORTEX build that uses the same gemv_t.S.
  • interface/gemm.c (shared): the small-matrix permit is checked before the SME kernels. Targets without their own permit (the generic one returns 0) are unchanged.
  • sgemm_direct_performant.c: row-major SGEMM below 10000 multiply-adds leaves the SME1 direct kernel when SME2 is usable.
  • sme_dgemm_kernel.c: the SME1 kernel special case below 5000 multiply-adds is removed; those problems now take the small-matrix kernels before they reach it.

The kernels are plain NEON and could serve other arm64 targets (VORTEX, NEOVERSEN1, ARMV9SME) through their KERNEL files, but their thresholds would need measurements on that hardware, so this PR selects them for VORTEXM4 only.

Results

M4 Pro, one thread, square, GFLOPS, develop 791b02a -> this PR (column-major with beta = 1, row-major with beta = 0):

n SGEMM col SGEMM row DGEMM col DGEMM row
4 4.9 -> 12 1.0 -> 10 5.0 -> 12 4.5 -> 12
8 9.9 -> 60 6.8 -> 56 9.8 -> 40 9.0 -> 44
12 26 -> 87 25 -> 87 18 -> 48 17 -> 53
16 36 -> 100 59 -> 108 42 -> 49 44 -> 53
20 44 -> 96 79 -> 100 40 -> 50 56 -> 53
24 and up unchanged unchanged unchanged unchanged

For comparison, Accelerate on the same machine, SGEMM column-major: 11 / 28 / 30 / 56 / 71 GFLOPS at n = 4 / 8 / 12 / 16 / 20.

Tests

  • make: BLAS/CBLAS tests, utest 144/144 and 1564/1564.
  • make lapack-test: 0 numerical errors.
  • Randomized SGEMM/DGEMM against a double reference: all transposes, both orders, alpha in {1, -1, 0.5, 2.5, 0}, beta in {0, 1, -0.5}, padded leading dimensions, m/n/k 1-30 (20000 trials) and 1-80 (3000 trials), some k 200-600: 0 failures.
  • Randomized GEMM/SYMM/SYRK/SYR2K/TRMM/TRSM check (dimensions to 40 and to 600): 0 failures.
  • CMake (-DTARGET=VORTEXM4 -DNOFORTRAN=1 -DBUILD_TESTING=ON, ctest 14/14) and DYNAMIC_ARCH=1 NO_FORTRAN=1 (randomized GEMM check, 10000 trials, 0 failures).

AI assistance

AI coding tools (Claude Code) helped to prepare this PR. The changes were reviewed, tested and benchmarked as described above.

SGEMM/DGEMM below 16000 (fp32) and 12000 (fp64) multiply-adds,
apart from vector shapes, go to new SMALL_MATRIX_OPT kernels that
keep blocks of C in NEON registers without packing: lane FMAs for
A op(B), dot products for A^T B, plain loops for A^T B^T. In
interface/gemm.c the target's small-matrix permit is now checked
before the SME kernels, which on VORTEXM4 came first and left the
small-matrix kernels unused; targets without a permit are unchanged.

Part of OpenMathLib#6149
…kernel on VORTEXM4

At 8^3 the direct kernel takes about 150 ns per call on Apple M4,
the small-matrix kernels about 20 ns.

Part of OpenMathLib#6149
Small DGEMM now runs on the small-matrix kernels before it reaches
sme_dgemm_kernel, so the SME1 kernel for problems below 5000
multiply-adds (M and N multiples of 16) is no longer needed.
@tesch1

tesch1 commented Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

@0tomhoyt, since you know these kernels and have an M3 Pro: a review would be welcome, if you have time.

A speed test on the M3 would also help, but only with one extra step, because this PR selects the kernels for VORTEXM4 alone. Adding these lines to kernel/arm64/KERNEL.VORTEX gives a VORTEX build that uses them below the same thresholds (the permit file only sets thresholds; nothing in it is M4-specific):

SGEMM_SMALL_M_PERMIT = gemm_small_kernel_permit_vortexm4.c
SGEMM_SMALL_K_NN     = gemm_small_kernel_nn_neon.c
SGEMM_SMALL_K_NT     = gemm_small_kernel_nt_neon.c
SGEMM_SMALL_K_TN     = gemm_small_kernel_tn_neon.c
SGEMM_SMALL_K_TT     = gemm_small_kernel_tt_neon.c
SGEMM_SMALL_K_B0_NN  = gemm_small_kernel_nn_neon.c
SGEMM_SMALL_K_B0_NT  = gemm_small_kernel_nt_neon.c
SGEMM_SMALL_K_B0_TN  = gemm_small_kernel_tn_neon.c
SGEMM_SMALL_K_B0_TT  = gemm_small_kernel_tt_neon.c
DGEMM_SMALL_M_PERMIT = gemm_small_kernel_permit_vortexm4.c
DGEMM_SMALL_K_NN     = gemm_small_kernel_nn_neon.c
DGEMM_SMALL_K_NT     = gemm_small_kernel_nt_neon.c
DGEMM_SMALL_K_TN     = gemm_small_kernel_tn_neon.c
DGEMM_SMALL_K_TT     = gemm_small_kernel_tt_neon.c
DGEMM_SMALL_K_B0_NN  = gemm_small_kernel_nn_neon.c
DGEMM_SMALL_K_B0_NT  = gemm_small_kernel_nt_neon.c
DGEMM_SMALL_K_B0_TN  = gemm_small_kernel_tn_neon.c
DGEMM_SMALL_K_B0_TT  = gemm_small_kernel_tt_neon.c

Square SGEMM/DGEMM at n = 4 to 48, one thread, against the same build without the lines, would show whether VORTEX should use them too, and up to which size. No obligation, of course.

…ernels

Each gemm_small_kernel_*_neon.c declares its B0 and beta variants
with #ifdef B0, as kernel/generic/gemm_small_matrix_kernel_*.c do.
The header functions are static inline, so the files that use only
some of them compile without unused-function warnings.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant