From 2cbf39f845a036456402bfbe00a51126872825f3 Mon Sep 17 00:00:00 2001 From: Linmiao Xu Date: Mon, 6 Jul 2026 11:33:51 +0200 Subject: [PATCH] Split NEON sparse affine accumulation three ways Speedup measured locally with apple-silicon M3 Pro: ``` Result of 50 runs ================== base (./sf_master ) = 1817788 +/- 4388 test (...on_sparse_3c) = 1839826 +/- 5482 diff = +22039 +/- 3215 speedup = +0.0121 P(speedup > 0) = 1.0000 CPU: 11 x arm Hyperthreading: off ``` Passed STC: https://tests.stockfishchess.org/tests/view/6a496a49f97ff95f78795b3d LLR: 2.92 (-2.94,2.94) <0.00,2.00> Total: 44384 W: 11659 L: 11351 D: 21374 Ptnml(0-2): 86, 4673, 12406, 4901, 126 closes https://github.com/official-stockfish/Stockfish/pull/6951 No functional change --- .../layers/affine_transform_sparse_input.h | 68 +++++++++++++++++-- 1 file changed, 64 insertions(+), 4 deletions(-) diff --git a/src/nnue/layers/affine_transform_sparse_input.h b/src/nnue/layers/affine_transform_sparse_input.h index bdce63918..168bf2f32 100644 --- a/src/nnue/layers/affine_transform_sparse_input.h +++ b/src/nnue/layers/affine_transform_sparse_input.h @@ -169,7 +169,7 @@ class AffineTransformSparseInput { // If we're using high-latency dot product instructions, split the accumulators // to create 3 separate dependency chains and merge at the end constexpr IndexType NumRegs = - #if defined(USE_VNNI) && defined(USE_AVX512) + #if (defined(USE_VNNI) && defined(USE_AVX512)) || defined(USE_NEON_DOTPROD) 3 * NumAccums; #else NumAccums; @@ -180,6 +180,11 @@ class AffineTransformSparseInput { for (IndexType k = 0; k < NumAccums; ++k) acc[k] = biasvec[k]; + #if defined(USE_NEON_DOTPROD) + for (IndexType k = NumAccums; k < NumRegs; ++k) + acc[k] = vdupq_n_s32(0); + #endif + // convince GCC to not do weird pointer arithmetic in the following loops const i8* weights_cp = weights; @@ -250,6 +255,55 @@ class AffineTransformSparseInput { asm("" : "+r"(base_addr), "+r"(weights_base)); // opt barrier #endif + #if defined(USE_NEON_DOTPROD) + while (bits) + { + const isize i0 = pop_lsb(bits); + if (!bits) + { + const invec_t in0 = vec_set_32(load_as(base_addr + i0 * sizeof(i32))); + const auto col0 = reinterpret_cast( + &weights_base[i0 * OutputDimensions * ChunkSize]); + for (IndexType l = 0; l < NumAccums; ++l) + vec_add_dpbusd_32(acc[l], in0, col0[l]); + break; + } + + const isize i1 = pop_lsb(bits); + if (!bits) + { + const invec_t in0 = vec_set_32(load_as(base_addr + i0 * sizeof(i32))); + const invec_t in1 = vec_set_32(load_as(base_addr + i1 * sizeof(i32))); + const auto col0 = reinterpret_cast( + &weights_base[i0 * OutputDimensions * ChunkSize]); + const auto col1 = reinterpret_cast( + &weights_base[i1 * OutputDimensions * ChunkSize]); + for (IndexType l = 0; l < NumAccums; ++l) + { + vec_add_dpbusd_32(acc[l], in0, col0[l]); + vec_add_dpbusd_32(acc[l + NumAccums], in1, col1[l]); + } + break; + } + + const isize i2 = pop_lsb(bits); + const invec_t in0 = vec_set_32(load_as(base_addr + i0 * sizeof(i32))); + const invec_t in1 = vec_set_32(load_as(base_addr + i1 * sizeof(i32))); + const invec_t in2 = vec_set_32(load_as(base_addr + i2 * sizeof(i32))); + const auto col0 = reinterpret_cast( + &weights_base[i0 * OutputDimensions * ChunkSize]); + const auto col1 = reinterpret_cast( + &weights_base[i1 * OutputDimensions * ChunkSize]); + const auto col2 = reinterpret_cast( + &weights_base[i2 * OutputDimensions * ChunkSize]); + for (IndexType l = 0; l < NumAccums; ++l) + { + vec_add_dpbusd_32(acc[l], in0, col0[l]); + vec_add_dpbusd_32(acc[l + NumAccums], in1, col1[l]); + vec_add_dpbusd_32(acc[l + 2 * NumAccums], in2, col2[l]); + } + } + #else while (bits) { isize i = pop_lsb(bits); @@ -257,16 +311,22 @@ class AffineTransformSparseInput { auto col = reinterpret_cast(&weights_base[i * OutputDimensions * ChunkSize]); - #ifdef FIX_GCC15_MISOPTIMIZATION + #ifdef FIX_GCC15_MISOPTIMIZATION asm("" : "+r"(col), "+r"(input_addr)); - #undef FIX_GCC15_MISOPTIMIZATION - #endif + #undef FIX_GCC15_MISOPTIMIZATION + #endif const invec_t in = vec_set_32(load_as(input_addr)); for (IndexType l = 0; l < NumAccums; ++l) vec_add_dpbusd_32(acc[l], in, col[l]); } + #endif } + + #if defined(USE_NEON_DOTPROD) + for (IndexType l = 0; l < NumAccums; ++l) + acc[l] = vaddq_s32(vaddq_s32(acc[l], acc[l + NumAccums]), acc[l + 2 * NumAccums]); + #endif #endif outvec_t* outptr = reinterpret_cast(output); for (IndexType k = 0; k < NumAccums; ++k)