From 4c78ba892049d5cb4faafd41de8443d37474ed63 Mon Sep 17 00:00:00 2001 From: Linmiao Xu Date: Thu, 16 Jul 2026 08:20:08 +0200 Subject: [PATCH] AVX-512 squared clipped ReLU Passed STC: https://tests.stockfishchess.org/tests/view/6a4dd1075529b8472df7fd5e LLR: 3.44 (-2.94,2.94) <0.00,2.00> Total: 327136 W: 84861 L: 84124 D: 158151 Ptnml(0-2): 692, 34323, 92827, 35008, 718 closes https://github.com/official-stockfish/Stockfish/pull/6974 No functional change --- src/nnue/layers/sqr_clipped_relu.h | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/nnue/layers/sqr_clipped_relu.h b/src/nnue/layers/sqr_clipped_relu.h index 5da8e183e..4c6a4a6b1 100644 --- a/src/nnue/layers/sqr_clipped_relu.h +++ b/src/nnue/layers/sqr_clipped_relu.h @@ -75,7 +75,24 @@ class SqrClippedReLU { // MulHi strips the lower 16 bits (i.e. shift by 16) so we need to shift out the remaining. [[maybe_unused]] constexpr int SimdShiftAmount = WeightScaleBitsLocal * 2 + 7 - 16; -#if defined(USE_SSE2) +#if defined(USE_AVX512) + static_assert(InputDimensions % 32 == 0); + constexpr IndexType NumChunks = InputDimensions / 32; + const auto in = reinterpret_cast(input); + const auto out = reinterpret_cast<__m256i*>(output); + for (IndexType i = 0; i < NumChunks; ++i) + { + const __m256i words0 = _mm512_cvtsepi32_epi16(_mm512_load_si512(&in[i * 2 + 0])); + const __m256i words1 = _mm512_cvtsepi32_epi16(_mm512_load_si512(&in[i * 2 + 1])); + __m512i words = _mm512_inserti64x4(_mm512_castsi256_si512(words0), words1, 1); + + words = _mm512_srli_epi16(_mm512_mulhi_epi16(words, words), SimdShiftAmount); + + _mm256_store_si256(&out[i], _mm512_cvtsepi16_epi8(words)); + } + constexpr IndexType Start = NumChunks * 32; + +#elif defined(USE_SSE2) constexpr IndexType NumChunks = InputDimensions / 16; const auto in = reinterpret_cast(input); const auto out = reinterpret_cast<__m128i*>(output);