From 24a551899ed5be1533e7a01f65d649469f81bb51 Mon Sep 17 00:00:00 2001 From: Ryan Hirsch Date: Fri, 29 May 2026 19:05:08 +0200 Subject: [PATCH] Speed up feature transform on neon MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit STC: LLR: 2.95 (-2.94,2.94) <0.00,2.00> Total: 109344 W: 27919 L: 27520 D: 53905 Ptnml(0-2): 200, 11718, 30434, 12123, 197 https://tests.stockfishchess.org/tests/view/6a123ed0818cacc1db0ac39e Updated to anematode’s suggestion - apparently clang already optimized it to this. Apple M1 Result of 100 runs base (...kfish.master) = 1307276 +/- 10930 test (./stockfish ) = 1320940 +/- 11279 diff = +13664 +/- 2223 speedup = +0.0105 P(speedup > 0) = 1.0000 closes https://github.com/official-stockfish/Stockfish/pull/6855 No functional change --- src/nnue/nnue_feature_transformer.h | 37 +++++++++++++++++------------ 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index 356beb8cf..2df0454b1 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -234,8 +234,11 @@ class FeatureTransformer { static_assert((HalfDimensions / 2) % OutputChunkSize == 0); constexpr IndexType NumOutputChunks = HalfDimensions / 2 / OutputChunkSize; - const vec_t Zero = vec_zero(); - const vec_t FtMax = vec_set_16(FtMaxVal); + #if !defined(USE_NEON) + const vec_t Zero = vec_zero(); + const vec_t FtMax = vec_set_16(FtMaxVal); + constexpr int shift = 7; + #endif const vec_t* in0 = reinterpret_cast(&(accumulation[perspectives[p]][0])); const vec_t* in1 = @@ -289,19 +292,6 @@ class FeatureTransformer { // 8 bits. Shifting it by 7 bits left will no longer occupy the // signed bit, so we are safe. - // Note that on NEON processors, we shift left by 6 instead - // because the instruction "vqdmulhq_s16" also doubles the - // return value after the multiplication, adding an extra shift - // to the left by 1, so we compensate by shifting less before - // the multiplication. - - constexpr int shift = - #if defined(USE_SSE2) || defined(USE_LASX) || defined(USE_LSX) - 7; - #else - 6; - #endif - const vec_t* tin0 = reinterpret_cast(&(threatAccumulation[perspectives[p]][0])); const vec_t* tin1 = reinterpret_cast( @@ -318,6 +308,21 @@ class FeatureTransformer { vec_t acc1a = vec_add_16(in1[i + 0], tin1[i + 0]); vec_t acc1b = vec_add_16(in1[i + 1], tin1[i + 1]); + #if defined(USE_NEON) + + // The NEON path relies on unsigned saturation for crelu + static_assert(FtMaxVal == 255); + + uint16x8_t mul0 = vmull_u8(vqmovun_s16(acc0a), vqmovun_s16(acc1a)); + uint16x8_t mul1 = vmull_u8(vqmovun_s16(acc0b), vqmovun_s16(acc1b)); + + uint8x16x2_t uzp = + vuzpq_u8(vreinterpretq_u8_u16(mul0), vreinterpretq_u8_u16(mul1)); + uint8x16_t pab = vshrq_n_u8(uzp.val[1], 1); + packed[k] = out[j + k] = reinterpret_cast(pab); + + #else + vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, FtMax), Zero), shift); vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, FtMax), Zero), shift); vec_t sum1a = vec_min_16(acc1a, FtMax); @@ -327,6 +332,8 @@ class FeatureTransformer { vec_t pb = vec_mulhi_16(sum0b, sum1b); packed[k] = out[j + k] = vec_packus_16(pa, pb); + + #endif } cursor.record2(packed[0], packed[1]);