SME: one disjoint part of the result per SME unit, minimum task size per scalar type
Multi-threading for the SME GEMM backend, split out of !3164 (merged) as discussed there. Builds on !3164 (merged): until that merges, this branch contains its commits, and the diff against master shows both; the change of this MR is the last commit (Parallelizer.h, one method in GeneralMatrixMatrix.h, PreprocessorDirectives.dox, test/product_threaded.cpp).
On Apple M4 Pro, at default threading (a thread pool of one thread per core), this MR takes the MpGEMM paper's 24 workloads from 2081 to 2450 GFLOPS column-major (1.18x) and from 1733 to 2329 row-major (1.34x), square sizes 64^3 to 384^3 by 1.4x to 6.4x, and float 1024^3 and 2048^3 on the pool from about 2600-2750 to 2970-3140 GFLOPS. The design comes from MTGEMM-A, the reproduction of the MpGEMM paper that !3164 (merged) references, whose two-thread mode runs one independent product per SME unit.
What does this implement/fix?
- One disjoint part of the result per SME unit. Where the SME unit count is known (
nbSmeUnits() > 0, from !3164 (merged)) and the blocking holds no preallocated buffers,parallelize_gemmgives each thread a disjoint part of the result and runs the single-thread algorithm on it, instead of the shared parallel session. In the shared session each thread packs part of the LHS and reads the parts the others packed, and waits for them at every depth step; with one SME unit per core cluster, each with its own L2, the other cluster's packed LHS is a remote read. The split takes the side that duplicates less packing, on the kernel's view of the product (a row-major result is its transpose): the kernel's columns (each thread packs the whole LHS, a copy or nothing when read in place) unless its rows are at least twice as many (each thread packs the whole RHS, a transposing pack for a column-major one). The rows are split only when they also give at least 8 LHS panels per part, since narrower results already run close to the kernel's peak on one unit; otherwise the columns are split, and no part is empty. On the thread pool the caller runs the last part, spins for the others for up to about 65536 checks (the parts take about as long as each other, and a blocking wake-up costs a short product more) and then blocks on aBarrier, which it also passes before an exception from its own part leaves; OpenMP builds run the parts in aparallel for. Fixed-size blockings, which own their buffers, keep the shared session, and so does every build or scalar pair without the SME kernel. - A minimum task size per scalar type. Where the SME unit count is known, a product on the SME kernel is split only past
EIGEN_SME_MIN_TASK_SIZE(2^20) FP32 multiply-adds per thread, instead of the generic 50000: the SME kernel does far more work per unit of time, so a thread needs more of it to repay the handoff. Double and complex count four FP32 multiply-adds per multiply-add, complex sixteen (an FP64 outer product covers a quarter of the elements of an FP32 one at the same vector length, and a complex multiply-add is four real ones).
Results
Default threading, float, GFLOPS, geometric means, from the MTGEMM-A harness (minimum of 5 trials, C.noalias() += A * B column-major, C = A * B row-major); Accelerate with its own thread count, Eigen on a pool of one thread per core:
| workloads | order | Accelerate | MTGEMM-A | !3164 (merged) | this MR | this MR / !3164 (merged) | this MR / Accelerate |
|---|---|---|---|---|---|---|---|
| M = 64 (1-6) | col | 1836 | 2153 | 1499 | 1744 | 1.16x | 0.95x |
| M = 128 (7-12) | col | 2368 | 2715 | 2048 | 2288 | 1.12x | 0.97x |
| M = 4096 (13-18) | col | 2903 | 3285 | 2773 | 3287 | 1.19x | 1.13x |
| N = 256 (19-24) | col | 2472 | 3097 | 2202 | 2747 | 1.25x | 1.11x |
| all 24 | col | 2363 | 2777 | 2081 | 2450 | 1.18x | 1.04x |
| squares 512-4096 | col | 2878 | 3406 | 2631 | 3225 | 1.23x | 1.12x |
| M = 64 (1-6) | row | 1467 | 2076 | 1086 | 1605 | 1.48x | 1.09x |
| M = 128 (7-12) | row | 2127 | 2628 | 1507 | 2150 | 1.43x | 1.01x |
| M = 4096 (13-18) | row | 3255 | 3367 | 2176 | 3168 | 1.46x | 0.97x |
| N = 256 (19-24) | row | 3085 | 3161 | 2534 | 2691 | 1.06x | 0.87x |
| all 24 | row | 2366 | 2760 | 1733 | 2329 | 1.34x | 0.98x |
| squares 512-4096 | row | 3193 | 3473 | 2563 | 3125 | 1.22x | 0.98x |
Square sizes, column-major:
| n | Accelerate | !3164 (merged) | this MR | this MR / !3164 (merged) |
|---|---|---|---|---|
| 64 | 924 | 109 | 700 | 6.42x |
| 96 | 1258 | 278 | 1029 | 3.70x |
| 128 | 1431 | 537 | 1245 | 2.32x |
| 192 | 1593 | 918 | 1635 | 1.78x |
| 256 | 1684 | 1180 | 2219 | 1.88x |
| 384 | 1787 | 1907 | 2734 | 1.43x |
Over the MTGEMM-A shape sets (squares 4-384, thin shapes, and M x N grids from 4 to 4096 at K = 512 and 4096; 540 cells), this MR is 1.34x of !3164 (merged) alone by geometric mean; the cells at or below 0.9x measured level on re-runs. The grid cells below 0.9x of Accelerate go from 230 to 73 of 484.
GEMM on the thread pool, median of 3 repetitions, bench_gemm_threads (float, capped at the unit count):
| float n | pool threads | !3164 (merged) (GFLOPS) | this MR (GFLOPS) | ratio |
|---|---|---|---|---|
| 1024 | 1 | 1569 | 1596 | 1.02x |
| 1024 | 2 | 2622 | 2966 | 1.13x |
| 1024 | 4 | 2627 | 3143 | 1.20x |
| 1024 | 8 | 2608 | 2980 | 1.14x |
| 1024 | 12 | 2616 | 3015 | 1.15x |
| 2048 | 1 | 1618 | 1620 | 1.00x |
| 2048 | 2 | 2751 | 3117 | 1.13x |
| 2048 | 4 | 2690 | 3111 | 1.16x |
| 2048 | 8 | 2709 | 3108 | 1.15x |
| 2048 | 12 | 2724 | 3102 | 1.14x |
A second unit pays off from about 2^21 FP32 multiply-adds; a higher threshold (2^21 per thread) was measured and loses 16-33% on shapes such as 128^3 row-major and 4 x 1024 x 512 while winning only on narrow row splits, which the 8-panel rule covers.
Additional information
Machine: Apple M4 Pro, 8 performance cores in two clusters of 4 (one SME unit each), macOS 27.0, SVL 512. Compiler: Apple clang 21.0.0, -O3 -DNDEBUG -march=armv8.6-a+sme2+sme-f64f64, EIGEN_GEMM_THREADPOOL. The host was not idle during the default-threading runs (1-minute load 3.7-4.0, including a simulator); the two builds ran alternately on each shape set.
Tests: product_threaded runs its threaded tests twice, with setNbSmeUnits(2) (disjoint parts) and setNbSmeUnits(0) (the shared session), so both paths run on every host; it gains both split directions and storage orders compared with a serial product, every scalar the SME kernel serves (float, double, complex, complex), a split side with fewer chunks than threads, a narrow row split that stays on one thread, a block destination accumulated with a scale factor, and a fixed-size result that keeps the shared session. product_sme gains a disjoint-parts test that reaches the OpenMP branch in an EIGEN_TEST_OPENMP build; built with libomp on M4, product_sme parts 1-4 pass. The product and decomposition tests pass under the SME flags. Changed lines are clang-format 17 clean and scripts/check_style.py --diff reports nothing.