Skip to content

arm64: faster TRMM on VORTEXM4 (four-slice ZA stores, diagonal tiles, streaming copies) - #6159

Open
tesch1 wants to merge 5 commits into
OpenMathLib:developfrom
tesch1:sme2-trmm
Open

tesch1 wants to merge 5 commits into
OpenMathLib:developfrom
tesch1:sme2-trmm

Conversation

@tesch1

@tesch1 tesch1 commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Addresses #6147. Stacked on #6154 (its two commits are the first two here); only the last three commits are new.

Before: on VORTEXM4 (with #6154), one thread, STRMM at n = 1024 ran at about 830-1000 GFLOPS against about 1550 for Accelerate, and at n = 256 at 340-370 against about 1350. Timing the parts separately showed the TRMM kernel as the main cost: for fp32 at n = 256 it needed about 4x the time per flop of the GEMM kernel, although the code is the same apart from the store. The difference is s2_c_store_alpha() (C = alpha AB, used by the TRMM kernel), which read ZA one slice at a time.

After: all eight STRMM/DTRMM variants are faster at every size measured, by 1.3-3.8x; the left side reaches 0.95x of Accelerate at n = 1024 (fp32).

How:

  • sme2_gemm_tile.h: s2_c_store_alpha() reads four ZA slices per MOVA, as s2_c_store() does. This is the main fix; it also helps the GEMM kernel when alpha is not +-1.
  • gemm_kernel_sme2.c: for left TRMM, the depth steps at the diagonal block skip the outer products of tiles whose rows of A are zero (s2p_mac_diag()). The block functions get the diagonal position as one more argument, unused in the GEMM kernel.
  • trmm_[nt]copy_sme2.c, sme2_gemm_copy.h: a TRMM panel copy (kept part, zero part and diagonal fix) is one streaming-mode call above 512 elements, with predicated stores for the zero side and the unit diagonal, instead of core stores between streaming-mode calls.

Only VORTEXM4 files change. Not included: the outer TRMM copies (right side, B panels) are fixed to the generic files in kernel/Makefile.L3. A streaming-mode version of them gives another 1.15-1.7x on the right side (fp32, n = 64-1024), but needs optional [SD]TRMM??COPY_N variables in kernel/Makefile.L3 and kernel/CMakeLists.txt, like the existing _M ones; I can propose that separately if it is wanted.

Results

M4 Pro, one thread, non-unit, GFLOPS, best of two runs (noise about 3%). Base = develop + #6154. Columns: side (L/R), uplo (U/L), trans (N/T).

n LUN LUT LLN LLT RUN RUT RLN RLT
s base 64 81 80 82 81 88 85 88 88
s new 64 292 300 314 280 207 199 215 198
s base 256 371 370 370 369 364 339 362 343
s new 256 995 1011 1027 976 740 666 747 665
s base 1024 995 995 986 1001 858 835 868 829
s new 1024 1438 1468 1466 1453 1266 1173 1281 1184
d base 64 42 42 42 42 61 48 48 61
d new 64 137 146 151 130 136 105 103 138
d base 256 155 155 155 155 151 150 147 151
d new 256 330 340 340 326 238 239 237 238
d base 1024 309 306 310 312 300 302 303 300
d new 1024 392 397 400 398 380 375 375 376

Accelerate on the same machine, STRMM: about 430-560 (n = 64), 1290-1400 (n = 256), 1520-1560 (n = 1024); DTRMM about 215-240, 390-405, 390-440. GEMM, SYMM, SYRK and TRSM stay within 2% on one thread and within noise with threads.

Tests

  • make: BLAS/CBLAS tests, utest 144/144 and 1564/1564; make lapack-test: 0 numerical errors.
  • Randomized GEMM/SYMM/SYRK/SYR2K/TRMM/TRSM check against naive references (all side/uplo/trans/diag variants): s and d, 1000 trials each, 2000 trials up to dimension 40, 300 trials up to 1300 (multi-block and threaded paths): 0 failures.
  • CMake (-DTARGET=VORTEXM4 -DNOFORTRAN=1 -DBUILD_TESTING=ON, ctest 14/14), DYNAMIC_ARCH=1 NO_FORTRAN=1 (randomized check, 0 failures) and GCC 16 (TARGET=VORTEXM4; BLAS tests and utest pass).

AI assistance

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

The SME unit moves data 2-4x faster than NEON or the generic copies.
In sme2_gemm_copy.h, s2c_rows() uses predicated streaming vectors
and s2c_cols() transposes through ZA, for pieces of 512 elements or
more; smaller pieces and CPUs without SME2 stay on NEON. The new
B-side copies (S/DGEMMONCOPY, S/DGEMMOTCOPY) use the same helpers
with the layout of the generic copies (UNROLL_N, then 8/4/2/1).

Part of OpenMathLib#6146
…ement tests

The diagonal block goes through the panel copy; then each row gets
its zero side and its (unit or inverted) diagonal element.

Part of OpenMathLib#6146
s2_c_store_alpha() (C = alpha ZA in the TRMM kernel, C += alpha ZA in
the GEMM kernel for alpha other than 1 and -1) read ZA one slice at a
time. It now reads four slices per MOVA, as s2_c_store(). This makes
the TRMM kernel about 1.5x (n = 1024) to 2.4x (n = 256) faster.

Part of OpenMathLib#6147
In the 64-row block on the diagonal of a left TRMM, the A panel rows
of tile Q are zero for VL * Q (lower: the last) steps. The kernel now
leaves out the outer products of those tiles; about 4% (fp32) and 12%
(fp64) faster at n = 256.

Part of OpenMathLib#6147
The triangle copies zeroed with memset and fixed the diagonal block
with scalar stores on the core, between two streaming copies. The
kernel then had to read these lines back from the core's L1. Now a
copy of more than 512 elements is one streaming-mode call: the kept
part with s2c_rows_s/s2c_cols_s, the zero part and the diagonal block
with predicated streaming stores.

Part of OpenMathLib#6147

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