From 1c384d3a8774ec9906e2e85b5005b22e27489044 Mon Sep 17 00:00:00 2001 From: Linmiao Xu Date: Thu, 16 Jul 2026 08:20:55 +0200 Subject: [PATCH] AVX2 fused SqrClippedReLU + ClippedReLU Speedup measured with profile-build x86-64-avx2: ``` Result of 50 runs ================== base (./sf-master ) = 777661 +/- 8102 test (./sf-fused ) = 796739 +/- 9127 diff = +19077 +/- 3164 speedup = +0.0245 P(speedup > 0) = 1.0000 CPU: 8 x AMD EPYC 7R13 Processor ``` Passed STC: https://tests.stockfishchess.org/tests/view/6a5049335529b8472df8013c LLR: 2.96 (-2.94,2.94) <0.00,2.00> Total: 262624 W: 68444 L: 67815 D: 126365 Ptnml(0-2): 724, 28976, 71325, 29521, 766 closes https://github.com/official-stockfish/Stockfish/pull/6976 No functional change --- src/nnue/layers/affine_transform.h | 19 +++++++++++--- src/nnue/layers/sqr_clipped_relu.h | 41 ++++++++++++++++++++++++++++++ src/nnue/nnue_architecture.h | 29 ++++++++++++++++----- src/nnue/simd.h | 4 +++ 4 files changed, 83 insertions(+), 10 deletions(-) diff --git a/src/nnue/layers/affine_transform.h b/src/nnue/layers/affine_transform.h index 055e7ca5f..8bbde97a9 100644 --- a/src/nnue/layers/affine_transform.h +++ b/src/nnue/layers/affine_transform.h @@ -142,7 +142,7 @@ affine_transform_non_ssse3(i32* output, const i8* weights, const i32* biases, co #endif // !ENABLE_SEQ_OPT -template +template class AffineTransform { public: // Input/output type @@ -170,8 +170,21 @@ class AffineTransform { } static constexpr IndexType get_weight_index_scrambled(IndexType i) { - return (i / 4) % (PaddedInputDimensions / 4) * OutputDimensions * 4 - + i / PaddedInputDimensions * 4 + i % 4; + IndexType inputIndex = i % PaddedInputDimensions; + +#if defined(USE_AVX2_PAIR_ACTIVATIONS) + if constexpr (ScrambledInput) + { + // AVX2 packs operate independently on 128-bit lanes. Keep their interleaved output + // order and rearrange the following layer's weights instead of issuing VPERMD. + const IndexType block = inputIndex / 32; + const IndexType chunk = (inputIndex % 32) / 4; + inputIndex = block * 32 + ((chunk % 2) * 4 + chunk / 2) * 4 + inputIndex % 4; + } +#endif + + return inputIndex / 4 * OutputDimensions * 4 + i / PaddedInputDimensions * 4 + + inputIndex % 4; } static constexpr IndexType get_weight_index(IndexType i) { diff --git a/src/nnue/layers/sqr_clipped_relu.h b/src/nnue/layers/sqr_clipped_relu.h index 4c6a4a6b1..c84f68335 100644 --- a/src/nnue/layers/sqr_clipped_relu.h +++ b/src/nnue/layers/sqr_clipped_relu.h @@ -67,6 +67,47 @@ class SqrClippedReLU { return h; } +#if defined(USE_AVX2_PAIR_ACTIVATIONS) + // Produce the squared and linear clipped activations together, sharing the input loads and + // the initial signed 32-to-16-bit saturating packs. + void propagate_pair(const InputType* input, OutputType* squared, OutputType* clipped) const { + static_assert(WeightScaleBitsLocal >= 5 && WeightScaleBitsLocal <= 8, + "SqrClippedReLU only support WeightScaleBitsLocal between 5 and 8"); + static_assert(InputDimensions % 32 == 0); + + constexpr IndexType NumChunks = InputDimensions / 32; + constexpr int SimdShiftAmount = WeightScaleBitsLocal * 2 + 7 - 16; + + const auto in = reinterpret_cast(input); + auto sqrOut = reinterpret_cast<__m256i*>(squared); + auto clipOut = reinterpret_cast<__m256i*>(clipped); + + const __m256i zero = _mm256_setzero_si256(); + + for (IndexType i = 0; i < NumChunks; ++i) + { + const __m256i words0 = _mm256_packs_epi32(_mm256_load_si256(&in[i * 4 + 0]), + _mm256_load_si256(&in[i * 4 + 1])); + const __m256i words1 = _mm256_packs_epi32(_mm256_load_si256(&in[i * 4 + 2]), + _mm256_load_si256(&in[i * 4 + 3])); + + const __m256i sqr0 = + _mm256_srli_epi16(_mm256_mulhi_epi16(words0, words0), SimdShiftAmount); + const __m256i sqr1 = + _mm256_srli_epi16(_mm256_mulhi_epi16(words1, words1), SimdShiftAmount); + const __m256i sqrPacked = _mm256_packs_epi16(sqr0, sqr1); + _mm256_store_si256(&sqrOut[i], sqrPacked); + + const __m256i clip0 = + _mm256_srli_epi16(_mm256_max_epi16(words0, zero), WeightScaleBitsLocal); + const __m256i clip1 = + _mm256_srli_epi16(_mm256_max_epi16(words1, zero), WeightScaleBitsLocal); + const __m256i clipPacked = _mm256_packs_epi16(clip0, clip1); + _mm256_store_si256(&clipOut[i], clipPacked); + } + } +#endif + // Forward propagation void propagate(const InputType* input, OutputType* output) const { static_assert(WeightScaleBitsLocal >= 5 && WeightScaleBitsLocal <= 8, diff --git a/src/nnue/nnue_architecture.h b/src/nnue/nnue_architecture.h index c06d520c6..f3ccfdabd 100644 --- a/src/nnue/nnue_architecture.h +++ b/src/nnue/nnue_architecture.h @@ -57,14 +57,19 @@ struct NetworkArchitecture { static constexpr IndexType TransformedFeatureDimensions = L1; static constexpr int FC_0_OUTPUTS = L2; static constexpr int FC_1_OUTPUTS = L3; +#if defined(USE_AVX2_PAIR_ACTIVATIONS) + static constexpr bool UsePairedActivations = true; +#else + static constexpr bool UsePairedActivations = false; +#endif - Layers::AffineTransformSparseInput fc_0; - Layers::SqrClippedReLU ac_sqr_0; - Layers::ClippedReLU ac_0; - Layers::AffineTransform fc_1; - Layers::SqrClippedReLU ac_sqr_1; - Layers::ClippedReLU ac_1; - Layers::AffineTransform fc_2; + Layers::AffineTransformSparseInput fc_0; + Layers::SqrClippedReLU ac_sqr_0; + Layers::ClippedReLU ac_0; + Layers::AffineTransform fc_1; + Layers::SqrClippedReLU ac_sqr_1; + Layers::ClippedReLU ac_1; + Layers::AffineTransform fc_2; // Hash value embedded in the evaluation file static constexpr u32 get_hash_value() { @@ -110,12 +115,22 @@ struct NetworkArchitecture { Buffer buffer; fc_0.propagate(transformedFeatures, buffer.fc_0_out, nnzInfo); +#if defined(USE_AVX2_PAIR_ACTIVATIONS) + ac_sqr_0.propagate_pair(buffer.fc_0_out, buffer.concat_buffer, + buffer.concat_buffer + FC_0_OUTPUTS); +#else ac_sqr_0.propagate(buffer.fc_0_out, buffer.concat_buffer); ac_0.propagate(buffer.fc_0_out, buffer.concat_buffer + FC_0_OUTPUTS); +#endif fc_1.propagate(buffer.concat_buffer, buffer.fc_1_out); +#if defined(USE_AVX2_PAIR_ACTIVATIONS) + ac_sqr_1.propagate_pair(buffer.fc_1_out, buffer.concat_buffer + FC_0_OUTPUTS * 2, + buffer.concat_buffer + FC_0_OUTPUTS * 2 + FC_1_OUTPUTS); +#else ac_sqr_1.propagate(buffer.fc_1_out, buffer.concat_buffer + FC_0_OUTPUTS * 2); ac_1.propagate(buffer.fc_1_out, buffer.concat_buffer + FC_0_OUTPUTS * 2 + FC_1_OUTPUTS); +#endif fc_2.propagate(buffer.concat_buffer, buffer.fc_2_out); diff --git a/src/nnue/simd.h b/src/nnue/simd.h index bbaec0800..9a60e33cb 100644 --- a/src/nnue/simd.h +++ b/src/nnue/simd.h @@ -51,6 +51,10 @@ namespace Stockfish::Eval::NNUE::SIMD { +#if defined(USE_AVX2) && !defined(USE_VNNI) && !defined(USE_AVX512) + #define USE_AVX2_PAIR_ACTIVATIONS +#endif + // If vector instructions are enabled, we update and refresh the // accumulator tile by tile such that each tile fits in the CPU's // vector registers.