Conversation
On SME2 cores with a 512-bit streaming vector length (checked at run time), sme_sgemm_kernel and sme_dgemm_kernel run a blocked SME2 GEMM after Deng et al., "Demystifying ARM SME to Optimize General Matrix Multiplications" (arXiv:2512.21473), as measured in MTGEMM-A: blocking from the paper's analytical model, the ZA transposition of A in packing, online packing of B, a 1 x NT-tile micro-kernel with multi-vector loads plus half-width and edge kernels, four-row ZA moves for C, core-side L2 prefetch, all four transpose cases, and one thread per SME unit through exec_blas from 2^22 multiply-adds. The existing kernels stay as the fallback. - #pragma clang attribute enables SME2 for these functions only (Apple clang >= 17, LLVM >= 18); no build flags change. - Small problems: a per-thread packing buffer; interface/gemm.c keeps problems below 20^3 (fp64 16^3) multiply-adds on NEON; the SME1 direct sgemm keeps row-major problems below 64^3. - M, N or K of 2^30 or more stay on the existing kernels; a failed blas_memory_alloc falls back to one thread and malloc. - SME units: one per performance cluster on Apple; two per Super cluster on M5 Pro/Max. - cpuid: detect the M4 Pro (hw.cpufamily 0x17d5b93a) and later Apple cores with FEAT_SME as VORTEXM4.
For real single and double precision on SME targets, kernel/arm64/sme_level3.h splits the symmetric or triangular dimension recursively: off-diagonal blocks are GEMM calls on SME_[SD]GEMM_KERNEL, diagonal blocks of at most 64 run as one GEMM on a full copy (SYMM, SYRK, SYR2K, TRMM), and TRSM solves blocks of at most 32 by substitution in a local tile (NEON transposes for the left side). The interface files only call the hooks, after the argument checks and develop's SME1 direct kernels, from 2e5 multiply-adds; the level-3 driver and its kernels are unchanged. TRMM falls back to the driver when no buffer is available.
Author
|
Superseded by #6074, which does the same work without changes to the interface files. This PR ran SYMM, SYRK, SYR2K, TRMM and TRSM on the SME GEMM kernel through hooks in interface/{symm,syrk,syr2k,trsm}.c, which included a header from kernel/arm64/, and it added an M4-tuned size threshold to interface/gemm.c. #6074 adds SME2 kernels for the level-3 driver instead (KERNEL.VORTEXM4 and param.h), so the LAPACK routines on the driver also run on SME2. It keeps the small-size limit inside the kernel, and it limits level-3 threads to one per SME unit in num_cpu_avail(). #6074 also compares the speed of both PRs; this PR stays faster for some fp64 routines and for very small SGEMMs. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #6073.
On SME2 hardware with a 512-bit streaming vector length (checked at run time),
sme_sgemm_kernelandsme_dgemm_kernelrun a blocked SME2 GEMM after C. Deng, W. Yang, J. Fang, D. Dong, Demystifying ARM SME toOptimize General Matrix Multiplications (arXiv:2512.21473), as reproduced and measured in MTGEMM-A
(https://github.com/tesch1/mtgemm-a). The kernels from #5971 stay as the fallback for other SME hardware.
SYMM, SYRK, SYR2K, TRMM and TRSM, which ran on the NEON kernels of the level-3 driver, now use the SME GEMM kernel.
Apple M4 Pro (GFLOPS; GEMM rows are geometric means over the shape set; "default" is each library's own thread
count; all measured on the same machine with the same harness; row-major / column-major where two values are
given):
Not slower than develop anywhere measured: GEMM on 348 shape/configuration pairs and, per call, 54 square
sizes from 4^3 to 256^3 in both orders and precisions; SYMM, SYRK, SYR2K, TRMM and TRSM at n = 64-4096 on 1 thread
and default threads, and every variant at n = 1024. The only ratios below 1.0 (0.95x-0.99x per call) are sizes
where both libraries run the same code: row-major SGEMM below 64^3 stays on develop's SME1 direct kernel, fp64 16^3
on its SME kernel. That direct kernel mallocs a scratch copy on every call and its speed depends on where the
allocation lands; develop itself gives 82-137 GFLOPS at 24^3 and 230-400 at 32^3 from process to process.
Where this PR is still slower:
fp64 4^3-16^3 at 0.50x-0.93x, fp64 24^3-96^3 at 0.60x-0.93x), shapes with M or N = 1 (0.36x-0.58x), fp64
16 x 16 x 4096, 16 x 4096 x 4096 and 4096 x 16 x 4096 (0.74x-0.75x), and fp32 4096 x 24576 x 1536 at default
threads (0.89x). Every other shape is at 0.95x or
better, and the geometric means of the paper's workloads and large squares are 1.0x-1.3x Accelerate.
at 64, which stays on develop's NEON path below the threshold, 0.04x-0.07x; TRMM, TRSM 0.16x-0.76x). At n = 1024,
SYMM, SYRK, SYR2K and TRMM are 0.72x-0.97x on one thread and 0.62x-1.20x at default threads; from n = 2048
0.74x-1.13x (1.61x for fp32 SYRK at 2048 at default threads, where Accelerate drops). TRSM from n = 1024 is
0.74x-1.14x on one thread and 0.47x-1.10x at default threads, where its base solves still run on one thread;
over the variants at n = 1024, fp32 577-591 GFLOPS (triangle on the left) and 673-696 (right) against
Accelerate's 791-879 and 840-949; fp64 253-258 and 273-275 against 263-290 and 207-252.
squares; per call 0.93x-0.95x at its worst small sizes and faster at 185 of 216 (NEON and the older kernels
take the smallest problems); 0.97x-0.98x on fp32 thin shapes (memory-bound; within 2% when timed in one
process).
Full tables and raw data: https://github.com/tesch1/mtgemm-a (README, "What this means for OpenBLAS").
Commits
kernel/arm64/sme2_gemm_impl.h(new), included bysme_sgemm_kernel.candsme_dgemm_kernel.c: blockingfrom the paper's analytical model, the ZA transposition of A in packing, online packing of B during the first
micro-kernel pass, a 1 x NT-tile micro-kernel with multi-vector loads (plus half-width and edge kernels with
predicates for every column tail wider than one vector), four-row ZA moves for C (also partial-width),
core-side L2 prefetch, all four transpose cases (vectorised copy for column-contiguous A, alpha applied after
the ZA transposition), and one thread per SME unit through
exec_blasfrom 2^22 multiply-adds.#pragma clang attributeenables SME2 for these functions only, so no build flags change; the path iscompiled with Apple clang >= 17 or LLVM >= 18 (
sme2_gemm_detect.halso has the run-time check).the ZA transposition; a per-thread packing buffer (freed at thread exit; the stack and
blas_memory_allocwere slower);
interface/gemm.ckeeps problems below 20^3 multiply-adds (fp64 16^3) on NEON; the SME1 directSGEMM keeps row-major problems below 64^3 (checked first, so small calls cost what they did); fp64 below 5000
with M and N multiples of 16 keeps the existing SME kernel.
(
hw.perflevel1.nameis "Performance") and have been reported to have two SME units for one Super cluster,so the count doubles there (not measured here).
int;if
blas_memory_allocfails, the two-thread split runs on one thread and the workspace comes frommalloc.cpuid_arm64.c: the M4 Pro (hw.cpufamily0x17d5b93a) was not recognised, so builds withoutTARGETfellback to ARMV8 without SME; later Apple cores reporting FEAT_SME now map to VORTEXM4.
kernel/arm64/sme_level3.h(new), for real types: the symmetric/triangular dimension is split recursively,off-diagonal blocks are GEMM calls on
SME_[SD]GEMM_KERNEL, diagonal blocks (<= 64) run as one GEMM on afull copy (SYMM, SYRK, SYR2K, TRMM), and TRSM solves blocks <= 32 by substitution in a local tile of 16 rows
(right side) or 16 columns (left side, moved with NEON 4 x 4 / 2 x 2 transposes); the generic TRSM kernel
runs at 10-16 GFLOPS on such blocks. Every side/uplo/trans/diag, from 2e5 multiply-adds.
interface/{symm,syrk,syr2k,trsm}.ceach get only anARCH_ARM64-guarded include and one guarded call tos3_symm_hook,s3_syrk_hookors3_trxm_hook(8 lines per file). TRMM falls back to the level-3 driverif no buffer is available.
Why the level-3 routines are not in the level-3 driver
The driver runs SYMM/SYRK/SYR2K/TRMM/TRSM through their own pack and micro-kernel routines, which share
GEMM_UNROLL_M/Nwith the GEMM kernel of the target (the constraint that held up #5011). The SME2 kernel packsinside its own blocking, so sharing unroll sizes with it would change every other kernel of VORTEXM4 and
ARMV9SME. Recursive blocking over the complete GEMM kernel needs no new pack routines, no
gotoblas_tmembers andno build-system changes; it is the usual way to put these routines on a fast GEMM. The code is in
kernel/arm64/sme_level3.h; the interface files only call it, after the argument checks and quick returns andafter develop's SME1 direct kernels. Those still run first where they apply: row-major fp32 SSYMM (left side),
SSYRK and STRMM (left side, non-unit) with packed leading dimensions, and row-major SGEMM below 64^3.
Relation to #5011
#5011 (WIP SME2 SGEMM through the level-3 driver) stopped at the point this PR works around: SYMM and TRMM
share the GEMM unroll sizes, so a new GEMM kernel shape breaks them. On cores with a 512-bit streaming vector
length (all Apple SME cores so far) this PR covers SGEMM and DGEMM and the level-3 routines with SYMM/TRMM
passing, so #5011 could be closed in its favour for those cores. #5011 also targets other vector lengths,
which this PR does not handle (they keep the #5971 kernels); that part would remain open.
Testing
Machine: Apple M4 Pro (8 performance + 4 efficiency cores; two SME units, one per performance cluster),
macOS 27 (Darwin 27.0.0), Apple clang 21.0.0.
Untested on M5, M5 Pro/Max, M6 and all other SME hardware. On those, the new paths run only if the core
reports SME2 and a 512-bit streaming vector length (checked at run time); otherwise develop's kernels run
unchanged.
DYNAMIC_ARCHalready selectsvortexm4for every Apple core with FEAT_SME, andgetarchnow mapsunknown Apple cores with FEAT_SME to VORTEXM4. The thresholds and the SME unit count were tuned or derived on the
M4 Pro only; the M5 Pro/Max unit count comes from a third-party report and is not measured here.
Results from other machines would be welcome.
Build configurations
TARGET=VORTEXM4 USE_THREAD=1 USE_OPENMP=0 NUM_THREADS=12 NO_LAPACK=1DYNAMIC_ARCH=1 USE_OPENMP=1 NUM_THREADS=56 TARGET=VORTEX NO_SVE=1, Apple clang with-Xpreprocessor -fopenmpand libomp, as the Homebrew formula buildsvortexm4selected at run time, new paths takenTARGETgetarchnow reports VORTEXM4 on the M4 Pro (was ARMV8)OpenBLAS's own tests
Both macOS configurations above, built with
makeand gfortran-16 (NO_LAPACK=1), passtest(s/d/c/z blat1-3),ctest(cblas s/d/c/z, both orders) andutest(125 + 1500 tests). Their default sizes (up to 31) stay belowthe new level-3 paths, so
sblat3anddblat3were also run with N = 0 1 2 3 7 31 33 63 65 (65 is the testers'NMAX), which reaches them: all six routines pass in both precisions and both configurations.
GEMM checks (
test_openblas_gemm.c)12,448 calls of
cblas_sgemmandcblas_dgemm, each compared element by element with a long-double reference;tolerance 4 (k + 4) eps (eps = 1.2e-7 fp32, 2.3e-16 fp64), inputs uniform in [-1, 1]:
panel and tail boundary of the kernels (16/32/64 fp32, 8/16/64 fp64), plus (256, 300, 260), (513, 129, 700),
(70, 1100, 300), (1030, 40, 1500), (600, 700, 1300), which exercise blocking and the two-thread split;
In addition every call of the GEMM benchmark (348 shape and configuration pairs) is checked against MTGEMM-A:
340 are bit for bit identical (as is develop on all 348); the other eight, fp32 column-major 4^3-16^3 on one
thread and at default threads, go to the NEON kernels and differ by at most 8.4e-8 sqrt(K).
Small sizes, per call
bench_smalltimes 54 square sizes from 4^3 to 256^3 per call (minimum over many short batches, three passes overall sizes, since a pass can start on an efficiency core whose SME unit is several times slower), for develop, this
PR and MTGEMM-A, one thread, both orders and precisions. The tables are under "Performance details".
Level-3 checks (
test_openblas_l3.c)2,816 calls, each against a long-double reference, in fp32 and fp64:
(128, 128)}: below, at and above the recursion blocks (32 for TRSM, 64 otherwise) and odd splits;
fails the check; every element of B or C outside the result (other triangle, padding) must come back bit for
bit; TRSM uses a diagonally dominant A.
Develop passes the same test; the tests are in https://github.com/tesch1/mtgemm-a (
bench/).Performance details
Harness: minimum of five trials of at least 50 ms each, after a 12 s wait (idle OpenBLAS pool threads spin for
2^28 timer ticks, about 11 s at Apple's 24 MHz, after start-up and slow the first calls by up to 25%); threads at
QoS user-interactive; runs start only below a one-minute load average of 4. Timings of the same build vary by
up to about 3% between sessions.
GEMM, per shape
Accelerate / develop / this PR / MTGEMM-A, GFLOPS; "row" is row-major
C = A*B(beta 0), "col" column-majorC += A*B(beta 1).The paper's 24 workloads and squares 512-4096, fp32, 1 thread
The paper's 24 workloads and squares 512-4096, fp32, default threads
The paper's 24 workloads and squares 512-4096, fp64, 1 thread
Small squares, fp32, 1 thread
Small squares, fp32, default threads
Small squares, fp64, 1 thread
Thin shapes, fp32, 1 thread
Thin shapes, fp32, default threads
Thin shapes, fp64, 1 thread
SYMM, SYRK, SYR2K, TRMM, TRSM, per size
Column-major, lower, left side, no transpose, non-unit diagonal, beta 1; flops 2n^3 for GEMM, SYMM, SYR2K and n^3 for SYRK, TRMM, TRSM. Accelerate / develop / this PR, GFLOPS; the last column is this PR over Accelerate.
fp32, 1 thread
fp64, 1 thread
fp32, default threads
fp64, default threads
GEMM per call, small squares
Nanoseconds per call, one thread, minimum over many short batches in three passes (
bench_small); row-major beta 0, column-major beta 1. A ratio above 1 means this PR is faster.row-major, fp32
row-major, fp64
col-major, fp32
col-major, fp64
Credit
Design: Deng et al. (arXiv:2512.21473). MTGEMM-A is the measured reproduction this port comes from. The
existing kernels of #5971 (ported from vlovero's ARMv9.2-GEMM) and the SME1 direct kernels (#5084, #5380) remain
in place as fallbacks and for small problems.
Prepared with AI coding assistance.