X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=src%2Fnnue%2Flayers%2Faffine_transform.h;h=44fa5d00a434f8285dea2357e8e6889cfdd8aed6;hb=HEAD;hp=3fba45ed87dc57e5ace02a366bad1a63e1f1c5d7;hpb=2d0237db3f0e596fb06e3ffbadba84dcc4e018f6;p=stockfish diff --git a/src/nnue/layers/affine_transform.h b/src/nnue/layers/affine_transform.h index 3fba45ed..59a6149f 100644 --- a/src/nnue/layers/affine_transform.h +++ b/src/nnue/layers/affine_transform.h @@ -1,6 +1,6 @@ /* Stockfish, a UCI chess playing engine derived from Glaurung 2.1 - Copyright (C) 2004-2023 The Stockfish developers (see AUTHORS file) + Copyright (C) 2004-2024 The Stockfish developers (see AUTHORS file) Stockfish is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by @@ -39,25 +39,26 @@ namespace Stockfish::Eval::NNUE::Layers { +#if defined(USE_SSSE3) || defined(USE_NEON_DOTPROD) + #define ENABLE_SEQ_OPT +#endif + // Fallback implementation for older/other architectures. // Requires the input to be padded to at least 16 values. -#if !defined(USE_SSSE3) +#ifndef ENABLE_SEQ_OPT + template static void affine_transform_non_ssse3(std::int32_t* output, const std::int8_t* weights, const std::int32_t* biases, const std::uint8_t* input) { - #if defined(USE_SSE2) || defined(USE_NEON_DOTPROD) || defined(USE_NEON) + #if defined(USE_SSE2) || defined(USE_NEON) #if defined(USE_SSE2) // At least a multiple of 16, with SSE2. constexpr IndexType NumChunks = ceil_to_multiple(InputDimensions, 16) / 16; const __m128i Zeros = _mm_setzero_si128(); const auto inputVector = reinterpret_cast(input); - #elif defined(USE_NEON_DOTPROD) - constexpr IndexType NumChunks = ceil_to_multiple(InputDimensions, 16) / 16; - const auto inputVector = reinterpret_cast(input); - #elif defined(USE_NEON) constexpr IndexType NumChunks = ceil_to_multiple(InputDimensions, 16) / 16; const auto inputVector = reinterpret_cast(input); @@ -91,16 +92,8 @@ static void affine_transform_non_ssse3(std::int32_t* output, sum = _mm_add_epi32(sum, sum_second_32); output[i] = _mm_cvtsi128_si32(sum); - #elif defined(USE_NEON_DOTPROD) - int32x4_t sum = {biases[i]}; - const auto row = reinterpret_cast(&weights[offset]); - for (IndexType j = 0; j < NumChunks; ++j) - { - sum = vdotq_s32(sum, inputVector[j], row[j]); - } - output[i] = vaddvq_s32(sum); - #elif defined(USE_NEON) + int32x4_t sum = {biases[i]}; const auto row = reinterpret_cast(&weights[offset]); for (IndexType j = 0; j < NumChunks; ++j) @@ -127,7 +120,8 @@ static void affine_transform_non_ssse3(std::int32_t* output, } #endif } -#endif + +#endif // !ENABLE_SEQ_OPT template class AffineTransform { @@ -162,7 +156,7 @@ class AffineTransform { } static constexpr IndexType get_weight_index(IndexType i) { -#if defined(USE_SSSE3) +#ifdef ENABLE_SEQ_OPT return get_weight_index_scrambled(i); #else return i; @@ -190,32 +184,28 @@ class AffineTransform { // Forward propagation void propagate(const InputType* input, OutputType* output) const { -#if defined(USE_SSSE3) +#ifdef ENABLE_SEQ_OPT if constexpr (OutputDimensions > 1) { - #if defined(USE_AVX512) using vec_t = __m512i; - #define vec_setzero _mm512_setzero_si512 #define vec_set_32 _mm512_set1_epi32 #define vec_add_dpbusd_32 Simd::m512_add_dpbusd_epi32 - #define vec_add_dpbusd_32x2 Simd::m512_add_dpbusd_epi32x2 - #define vec_hadd Simd::m512_hadd #elif defined(USE_AVX2) using vec_t = __m256i; - #define vec_setzero _mm256_setzero_si256 #define vec_set_32 _mm256_set1_epi32 #define vec_add_dpbusd_32 Simd::m256_add_dpbusd_epi32 - #define vec_add_dpbusd_32x2 Simd::m256_add_dpbusd_epi32x2 - #define vec_hadd Simd::m256_hadd #elif defined(USE_SSSE3) using vec_t = __m128i; - #define vec_setzero _mm_setzero_si128 #define vec_set_32 _mm_set1_epi32 #define vec_add_dpbusd_32 Simd::m128_add_dpbusd_epi32 - #define vec_add_dpbusd_32x2 Simd::m128_add_dpbusd_epi32x2 - #define vec_hadd Simd::m128_hadd + #elif defined(USE_NEON_DOTPROD) + using vec_t = int32x4_t; + #define vec_set_32 vdupq_n_s32 + #define vec_add_dpbusd_32(acc, a, b) \ + Simd::dotprod_m128_add_dpbusd_epi32(acc, vreinterpretq_s8_s32(a), \ + vreinterpretq_s8_s32(b)) #endif static constexpr IndexType OutputSimdWidth = sizeof(vec_t) / sizeof(OutputType); @@ -231,46 +221,47 @@ class AffineTransform { for (IndexType k = 0; k < NumRegs; ++k) acc[k] = biasvec[k]; - for (IndexType i = 0; i < NumChunks; i += 2) + for (IndexType i = 0; i < NumChunks; ++i) { - const vec_t in0 = vec_set_32(input32[i + 0]); - const vec_t in1 = vec_set_32(input32[i + 1]); + const vec_t in0 = vec_set_32(input32[i]); const auto col0 = - reinterpret_cast(&weights[(i + 0) * OutputDimensions * 4]); - const auto col1 = - reinterpret_cast(&weights[(i + 1) * OutputDimensions * 4]); + reinterpret_cast(&weights[i * OutputDimensions * 4]); + for (IndexType k = 0; k < NumRegs; ++k) - vec_add_dpbusd_32x2(acc[k], in0, col0[k], in1, col1[k]); + vec_add_dpbusd_32(acc[k], in0, col0[k]); } vec_t* outptr = reinterpret_cast(output); for (IndexType k = 0; k < NumRegs; ++k) outptr[k] = acc[k]; - #undef vec_setzero #undef vec_set_32 #undef vec_add_dpbusd_32 - #undef vec_add_dpbusd_32x2 - #undef vec_hadd } else if constexpr (OutputDimensions == 1) { - - // We cannot use AVX512 for the last layer because there's only 32 inputs and the buffer is not padded to 64 elements. + // We cannot use AVX512 for the last layer because there are only 32 inputs + // and the buffer is not padded to 64 elements. #if defined(USE_AVX2) using vec_t = __m256i; - #define vec_setzero _mm256_setzero_si256 + #define vec_setzero() _mm256_setzero_si256() #define vec_set_32 _mm256_set1_epi32 #define vec_add_dpbusd_32 Simd::m256_add_dpbusd_epi32 - #define vec_add_dpbusd_32x2 Simd::m256_add_dpbusd_epi32x2 #define vec_hadd Simd::m256_hadd #elif defined(USE_SSSE3) using vec_t = __m128i; - #define vec_setzero _mm_setzero_si128 + #define vec_setzero() _mm_setzero_si128() #define vec_set_32 _mm_set1_epi32 #define vec_add_dpbusd_32 Simd::m128_add_dpbusd_epi32 - #define vec_add_dpbusd_32x2 Simd::m128_add_dpbusd_epi32x2 #define vec_hadd Simd::m128_hadd + #elif defined(USE_NEON_DOTPROD) + using vec_t = int32x4_t; + #define vec_setzero() vdupq_n_s32(0) + #define vec_set_32 vdupq_n_s32 + #define vec_add_dpbusd_32(acc, a, b) \ + Simd::dotprod_m128_add_dpbusd_epi32(acc, vreinterpretq_s8_s32(a), \ + vreinterpretq_s8_s32(b)) + #define vec_hadd Simd::neon_m128_hadd #endif const auto inputVector = reinterpret_cast(input); @@ -293,7 +284,6 @@ class AffineTransform { #undef vec_setzero #undef vec_set_32 #undef vec_add_dpbusd_32 - #undef vec_add_dpbusd_32x2 #undef vec_hadd } #else