Skip to content

arm64: SME GEMM kernels with packing buffers per thread (thread safety) - #6157

Open
tesch1 wants to merge 2 commits into
OpenMathLib:developfrom
tesch1:sme1-buffers
Open

tesch1 wants to merge 2 commits into
OpenMathLib:developfrom
tesch1:sme1-buffers

Conversation

@tesch1

@tesch1 tesch1 commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Before: the SME GEMM kernels kernel/arm64/sme_[sdcz]gemm_kernel.c (used by ARMV9SME for all four precisions, and on VORTEXM4 for CGEMM/ZGEMM and some small DGEMM shapes) pack their panels into static arrays (the thread_local is commented out, probably because the arrays are 5-20 MB). Two user threads that call GEMM at the same time, or the kernels' own #pragma omp parallel for tasks in an OpenMP build, then overwrite each other's panels. On an Apple M4 Pro (static VORTEXM4 build, develop 791b02a), with two pthreads each calling the same GEMM on its own data and comparing with the single-threaded result:

test develop this PR
CGEMM 300^3, 50 calls per thread 49/50 and 50/50 wrong 0/50, 0/50
ZGEMM 300^3, 50 calls per thread 50/50 and 50/50 wrong 0/50, 0/50
DGEMM 16^3 (SME kernel path), 200000 calls per thread 15053 and 17465 wrong 0, 0

After: each call (and each OpenMP task) packs into a buffer of the calling thread. A new header kernel/arm64/sme_gemm_buffer.h keeps one buffer per thread (pthread key, freed when the thread exits) for problems that need up to 4 MB, sized from M, N and K, and takes larger buffers from blas_memory_alloc, which is per call and already used by the level-3 driver. Without pthreads (Windows) every call uses blas_memory_alloc.

The per-call cost does not change measurably (one thread, DGEMM 16^3: 186 -> 189 ns; CGEMM 16^3: 641 -> 667 ns, 64^3: 8213 -> 8284 ns, within noise).

The existing thread-safety tests in cpp_thread_test cover DGEMM and DGEMV at sizes that do not reach these kernels on VORTEXM4, and CI has no SME hardware, which is probably why this was not caught. A second commit adds cpp_thread_test/zgemm_thread_safety.cpp (the DGEMM test with complex numbers, 384 x 384, make and CMake): on the M4 it fails on develop and passes with this PR.

Tests

  • The two-thread tests above (C, Z, D), and the new zgemm_tester (default arguments, and 12 threads x 4 rounds).
  • make: BLAS/CBLAS tests, utest 144/144 and 1564/1564; make lapack-test: 0 numerical errors.
  • Randomized complex level-3 check against a double reference (CGEMM/ZGEMM included): about 20 million checks per precision, 0 failures.

Only tested on M4 (VORTEXM4); the ARMV9SME path uses the same code but no ARMV9SME hardware was available.

AI assistance

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

sme_[sdcz]gemm_kernel.c packed into static arrays, so two user
threads (or the kernels' own OpenMP tasks) calling them at the same
time overwrote each other's panels: on Apple M4, ZGEMM 300^3 from
two threads gave 100/100 wrong results and DGEMM 16^3 about 8%.
The panels now go to a buffer of the calling thread, cached per
thread up to 4 MB and taken from blas_memory_alloc above that.
The same structure as the DGEMM test, with complex alpha and beta and
384 x 384 matrices. It fails on Apple M4 without the per-thread
packing buffers of the SME GEMM kernels.

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