X-Git-Url: https://git.sesse.net/?p=stockfish;a=blobdiff_plain;f=src%2Fnnue%2Flayers%2Faffine_transform.h;fp=src%2Fnnue%2Flayers%2Faffine_transform.h;h=9a3b778e6bbbedec7cb8b6d409c5d226f7569206;hp=fc1926912df1f9a15a731db50f47dd86e4f8d33e;hb=e8d64af1230fdac65bb0da246df3e7abe82e0838;hpb=f90274d8ce1aad4ad0595aacbceb74b6cbe306a8 diff --git a/src/nnue/layers/affine_transform.h b/src/nnue/layers/affine_transform.h index fc192691..9a3b778e 100644 --- a/src/nnue/layers/affine_transform.h +++ b/src/nnue/layers/affine_transform.h @@ -69,62 +69,15 @@ namespace Stockfish::Eval::NNUE::Layers { if (!previousLayer.read_parameters(stream)) return false; for (std::size_t i = 0; i < OutputDimensions; ++i) biases[i] = read_little_endian(stream); -#if !defined (USE_SSSE3) for (std::size_t i = 0; i < OutputDimensions * PaddedInputDimensions; ++i) +#if !defined (USE_SSSE3) weights[i] = read_little_endian(stream); #else - std::unique_ptr indexMap = std::make_unique(OutputDimensions * PaddedInputDimensions); - for (std::size_t i = 0; i < OutputDimensions * PaddedInputDimensions; ++i) { - const uint32_t scrambledIdx = + weights[ (i / 4) % (PaddedInputDimensions / 4) * OutputDimensions * 4 + i / PaddedInputDimensions * 4 + - i % 4; - weights[scrambledIdx] = read_little_endian(stream); - indexMap[scrambledIdx] = i; - } - - // Determine if eights of weight and input products can be summed using 16bits - // without saturation. We assume worst case combinations of 0 and 127 for all inputs. - if (OutputDimensions > 1 && !stream.fail()) - { - canSaturate16.count = 0; -#if !defined(USE_VNNI) - for (IndexType i = 0; i < PaddedInputDimensions; i += 16) - for (IndexType j = 0; j < OutputDimensions; ++j) - for (int x = 0; x < 2; ++x) - { - WeightType* w = &weights[i * OutputDimensions + j * 4 + x * 2]; - int sum[2] = {0, 0}; - for (int k = 0; k < 8; ++k) - { - IndexType idx = k / 2 * OutputDimensions * 4 + k % 2; - sum[w[idx] < 0] += w[idx]; - } - for (int sign : { -1, 1 }) - while (sign * sum[sign == -1] > 258) - { - int maxK = 0, maxW = 0; - for (int k = 0; k < 8; ++k) - { - IndexType idx = k / 2 * OutputDimensions * 4 + k % 2; - if (maxW < sign * w[idx]) - maxK = k, maxW = sign * w[idx]; - } - - IndexType idx = maxK / 2 * OutputDimensions * 4 + maxK % 2; - sum[sign == -1] -= w[idx]; - const uint32_t scrambledIdx = idx + i * OutputDimensions + j * 4 + x * 2; - canSaturate16.add(j, i + maxK / 2 * 4 + maxK % 2 + x * 2, w[idx], indexMap[scrambledIdx]); - w[idx] = 0; - } - } - - // Non functional optimization for faster more linear access - std::sort(canSaturate16.ids, canSaturate16.ids + canSaturate16.count, - [](const typename CanSaturate::Entry& e1, const typename CanSaturate::Entry& e2) - { return e1.in == e2.in ? e1.out < e2.out : e1.in < e2.in; }); -#endif - } + i % 4 + ] = read_little_endian(stream); #endif return !stream.fail(); @@ -148,8 +101,6 @@ namespace Stockfish::Eval::NNUE::Layers { i % 4 ]; } - for (int i = 0; i < canSaturate16.count; ++i) - unscrambledWeights[canSaturate16.ids[i].wIdx] = canSaturate16.ids[i].w; for (std::size_t i = 0; i < OutputDimensions * PaddedInputDimensions; ++i) write_little_endian(stream, unscrambledWeights[i]); @@ -194,11 +145,11 @@ namespace Stockfish::Eval::NNUE::Layers { __m512i product1 = _mm512_maddubs_epi16(a1, b1); __m512i product2 = _mm512_maddubs_epi16(a2, b2); __m512i product3 = _mm512_maddubs_epi16(a3, b3); - product0 = _mm512_add_epi16(product0, product1); - product2 = _mm512_add_epi16(product2, product3); - product0 = _mm512_add_epi16(product0, product2); + product0 = _mm512_adds_epi16(product0, product1); product0 = _mm512_madd_epi16(product0, Ones512); - acc = _mm512_add_epi32(acc, product0); + product2 = _mm512_adds_epi16(product2, product3); + product2 = _mm512_madd_epi16(product2, Ones512); + acc = _mm512_add_epi32(acc, _mm512_add_epi32(product0, product2)); #endif }; @@ -236,11 +187,11 @@ namespace Stockfish::Eval::NNUE::Layers { __m256i product1 = _mm256_maddubs_epi16(a1, b1); __m256i product2 = _mm256_maddubs_epi16(a2, b2); __m256i product3 = _mm256_maddubs_epi16(a3, b3); - product0 = _mm256_add_epi16(product0, product1); - product2 = _mm256_add_epi16(product2, product3); - product0 = _mm256_add_epi16(product0, product2); + product0 = _mm256_adds_epi16(product0, product1); product0 = _mm256_madd_epi16(product0, Ones256); - acc = _mm256_add_epi32(acc, product0); + product2 = _mm256_adds_epi16(product2, product3); + product2 = _mm256_madd_epi16(product2, Ones256); + acc = _mm256_add_epi32(acc, _mm256_add_epi32(product0, product2)); #endif }; @@ -267,11 +218,11 @@ namespace Stockfish::Eval::NNUE::Layers { __m128i product1 = _mm_maddubs_epi16(a1, b1); __m128i product2 = _mm_maddubs_epi16(a2, b2); __m128i product3 = _mm_maddubs_epi16(a3, b3); - product0 = _mm_add_epi16(product0, product1); - product2 = _mm_add_epi16(product2, product3); - product0 = _mm_add_epi16(product0, product2); + product0 = _mm_adds_epi16(product0, product1); product0 = _mm_madd_epi16(product0, Ones128); - acc = _mm_add_epi32(acc, product0); + product2 = _mm_adds_epi16(product2, product3); + product2 = _mm_madd_epi16(product2, Ones128); + acc = _mm_add_epi32(acc, _mm_add_epi32(product0, product2)); }; #endif @@ -300,6 +251,8 @@ namespace Stockfish::Eval::NNUE::Layers { #endif #if defined (USE_SSSE3) + // Different layout, we process 4 inputs at a time, always. + static_assert(InputDimensions % 4 == 0); const auto output = reinterpret_cast(buffer); const auto inputVector = reinterpret_cast(input); @@ -310,7 +263,7 @@ namespace Stockfish::Eval::NNUE::Layers { // because then it is also an input dimension. if constexpr (OutputDimensions % OutputSimdWidth == 0) { - constexpr IndexType NumChunks = PaddedInputDimensions / 4; + constexpr IndexType NumChunks = InputDimensions / 4; const auto input32 = reinterpret_cast(input); vec_t* outptr = reinterpret_cast(output); @@ -329,8 +282,6 @@ namespace Stockfish::Eval::NNUE::Layers { for (int j = 0; j * OutputSimdWidth < OutputDimensions; ++j) vec_add_dpbusd_32x4(outptr[j], in0, col0[j], in1, col1[j], in2, col2[j], in3, col3[j]); } - for (int i = 0; i < canSaturate16.count; ++i) - output[canSaturate16.ids[i].out] += input[canSaturate16.ids[i].in] * canSaturate16.ids[i].w; } else if constexpr (OutputDimensions == 1) { @@ -377,17 +328,21 @@ namespace Stockfish::Eval::NNUE::Layers { auto output = reinterpret_cast(buffer); #if defined(USE_SSE2) - constexpr IndexType NumChunks = PaddedInputDimensions / SimdWidth; + // At least a multiple of 16, with SSE2. + static_assert(InputDimensions % SimdWidth == 0); + constexpr IndexType NumChunks = InputDimensions / SimdWidth; const __m128i Zeros = _mm_setzero_si128(); const auto inputVector = reinterpret_cast(input); #elif defined(USE_MMX) - constexpr IndexType NumChunks = PaddedInputDimensions / SimdWidth; + static_assert(InputDimensions % SimdWidth == 0); + constexpr IndexType NumChunks = InputDimensions / SimdWidth; const __m64 Zeros = _mm_setzero_si64(); const auto inputVector = reinterpret_cast(input); #elif defined(USE_NEON) - constexpr IndexType NumChunks = PaddedInputDimensions / SimdWidth; + static_assert(InputDimensions % SimdWidth == 0); + constexpr IndexType NumChunks = InputDimensions / SimdWidth; const auto inputVector = reinterpret_cast(input); #endif @@ -473,25 +428,6 @@ namespace Stockfish::Eval::NNUE::Layers { alignas(CacheLineSize) BiasType biases[OutputDimensions]; alignas(CacheLineSize) WeightType weights[OutputDimensions * PaddedInputDimensions]; -#if defined (USE_SSSE3) - struct CanSaturate { - int count; - struct Entry { - uint32_t wIdx; - uint16_t out; - uint16_t in; - int8_t w; - } ids[PaddedInputDimensions * OutputDimensions * 3 / 4]; - - void add(int i, int j, int8_t w, uint32_t wIdx) { - ids[count].wIdx = wIdx; - ids[count].out = i; - ids[count].in = j; - ids[count].w = w; - ++count; - } - } canSaturate16; -#endif }; } // namespace Stockfish::Eval::NNUE::Layers