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?

  1. 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_gemm gives 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 a Barrier, which it also passes before an exception from its own part leaves; OpenMP builds run the parts in a parallel for. Fixed-size blockings, which own their buffers, keep the shared session, and so does every build or scalar pair without the SME kernel.
  2. 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.

Edited by Michael Tesch

Merge request reports

Loading
Loading