From 8132dbcb5a998b4d151a33db29135cb0fb829dbf Mon Sep 17 00:00:00 2001 From: anematode Date: Thu, 2 Apr 2026 20:23:33 +0200 Subject: [PATCH] split affine_transform on VNNI as well ``` LLR: 2.93 (-2.94,2.94) <0.00,2.00> Total: 184544 W: 47990 L: 47474 D: 89080 Ptnml(0-2): 551, 20238, 50206, 20698, 579 ``` Torom measured ``` sf_base = 2538696 +/- 2915 (95%) sf_test = 2546510 +/- 3011 (95%) diff = 7814 +/- 4146 (95%) speedup = 0.30782% +/- 0.163% (95%) ``` I get something similar. The benefit would be larger if we ever decide to further increase the L2 size I think closes https://github.com/official-stockfish/Stockfish/pull/6683 No functional change --- src/nnue/layers/affine_transform.h | 39 +++++++++++++++++++++++++++--- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/src/nnue/layers/affine_transform.h b/src/nnue/layers/affine_transform.h index 6cd44e19c..9c7699ac3 100644 --- a/src/nnue/layers/affine_transform.h +++ b/src/nnue/layers/affine_transform.h @@ -201,10 +201,12 @@ class AffineTransform { #if defined(USE_AVX512) using vec_t = __m512i; #define vec_set_32 _mm512_set1_epi32 + #define vec_add_32 _mm512_add_epi32 #define vec_add_dpbusd_32 SIMD::m512_add_dpbusd_epi32 #elif defined(USE_AVX2) using vec_t = __m256i; #define vec_set_32 _mm256_set1_epi32 + #define vec_add_32 _mm256_add_epi32 #define vec_add_dpbusd_32 SIMD::m256_add_dpbusd_epi32 #elif defined(USE_SSSE3) using vec_t = __m128i; @@ -223,14 +225,45 @@ class AffineTransform { static_assert(OutputDimensions % OutputSimdWidth == 0); constexpr IndexType NumChunks = ceil_to_multiple(InputDimensions, 8) / 4; - constexpr IndexType NumRegs = OutputDimensions / OutputSimdWidth; + constexpr IndexType NumAccums = OutputDimensions / OutputSimdWidth; + + #if defined(USE_VNNI) + constexpr IndexType NumRegs = 2 * NumAccums; + #else + constexpr IndexType NumRegs = NumAccums; + #endif const vec_t* biasvec = reinterpret_cast(biases); vec_t acc[NumRegs]; - for (IndexType k = 0; k < NumRegs; ++k) + for (IndexType k = 0; k < NumAccums; ++k) acc[k] = biasvec[k]; + for (IndexType k = NumAccums; k < NumRegs; ++k) + acc[k] = vec_set_32(0); - for (IndexType i = 0; i < NumChunks; ++i) + IndexType i = 0; + #if defined(USE_VNNI) + for (; i < NumChunks; i += 2) + { + const vec_t in0 = + vec_set_32(load_as(input + i * sizeof(std::int32_t))); + const vec_t in1 = + vec_set_32(load_as(input + (i + 1) * sizeof(std::int32_t))); + const auto col0 = + reinterpret_cast(&weights[i * OutputDimensions * 4]); + const auto col1 = + reinterpret_cast(&weights[(i + 1) * OutputDimensions * 4]); + + for (IndexType k = 0; k < NumAccums; ++k) + { + vec_add_dpbusd_32(acc[k], in0, col0[k]); + vec_add_dpbusd_32(acc[k + NumAccums], in1, col1[k]); + } + } + + for (IndexType k = 0; k < NumAccums; ++k) + acc[k] = vec_add_32(acc[k], acc[k + NumAccums]); + #endif + for (; i < NumChunks; ++i) { const vec_t in0 = vec_set_32(load_as(input + i * sizeof(std::int32_t)));