From fff35786bf4a5941a017066842f6977398ddac7e Mon Sep 17 00:00:00 2001 From: anematode Date: Tue, 19 May 2026 18:36:11 +0200 Subject: [PATCH] Use bitset representation for nnz and move computation into feature transformer Passed avx2/bmi2 STC: https://tests.stockfishchess.org/tests/view/69ff99eb9392f0c317213eda LLR: 2.93 (-2.94,2.94) <0.00,2.00> Total: 20992 W: 5469 L: 5192 D: 10331 Ptnml(0-2): 38, 2197, 5751, 2470, 40 Passed avxvnni STC: https://tests.stockfishchess.org/tests/view/6a00022b9392f0c317213f7c LLR: 2.93 (-2.94,2.94) <0.00,2.00> Total: 47328 W: 12328 L: 12009 D: 22991 Ptnml(0-2): 112, 5148, 12847, 5423, 134 Passed NEON STC: https://tests.stockfishchess.org/tests/view/69ff99d69392f0c317213ed8 LLR: 2.96 (-2.94,2.94) <0.00,2.00> Total: 29600 W: 7664 L: 7376 D: 14560 Ptnml(0-2): 48, 3074, 8277, 3344, 57 Passed avx512icl non-regression: https://tests.stockfishchess.org/tests/view/6a002b699392f0c317213f91 LLR: 2.94 (-2.94,2.94) <-1.75,0.25> Total: 90112 W: 23130 L: 22975 D: 44007 Ptnml(0-2): 192, 9939, 24633, 10106, 186 Measurements from vondele (Neoverse V2, armv8-dotprod): ==== master ==== ==== Bench: 2344696 ==== 1 Nodes/second : 280724660 2 Nodes/second : 280647282 3 Nodes/second : 282055192 Average (over 3): 281142378 ==== 50a44640a3 ==== ==== Bench: 2344696 ==== 1 Nodes/second : 284271937 2 Nodes/second : 285638071 3 Nodes/second : 284349426 Average (over 3): 284753144 The patch's benefit is non-uniform and is ~0 for avx512icl unfortunately -- although I think we should be able to find something there.... ## Background/explanation The idea here is to move the non-zero block computation back to the previous layer, and overlap the work better. Then, to avoid trying to emulate compress instructions on targets not supporting them (i.e., everything except for AVX512), we use a `pop_lsb` loop on a bitset enumerating the non-zero blocks, rather than first writing them out as indices. An early AVX2 implementation worked on fishtest (https://tests.stockfishchess.org/tests/view/69c3410534b6988b1e472db4) and linrock demonstrated that it also worked for NEON (https://tests.stockfishchess.org/tests/view/69ce0dc73ddc0eccd617188f). Thanks to him for encouraging me to push this idea over the finish line. To abstract over the particular NNZ representation (bitset or index list), we use `NNZInfo` and `NNZCursor`. The cursor is used to write to one or the other perspective of the NNZ list. I'd appreciate ideas on making the code cleaner/more readable as imo it's still a bit ugly. ## Fixing GCC 15 regression Separately, vondele noticed a regression in ARM performance from GCC 15 caused by suboptimal codegen. See https://discord.com/channels/435943710472011776/813919248455827515/1502837886381461605 for more info, but the easiest fix that I could come up with was inserting a couple `asm` optimization barriers. Single threaded data: ``` average: stockfish.master.13.3 1041619 average: stockfish.master.15.2 999214 average: stockfish.patched.15.2 1031369 ``` There's definitely regressions elsewhere as well, which we shld chase down, but at least this should unblock the arm64 universal binary work. closes https://github.com/official-stockfish/Stockfish/pull/6814 No functional change --- src/Makefile | 2 +- .../layers/affine_transform_sparse_input.h | 249 ++++-------------- src/nnue/network.cpp | 18 +- src/nnue/nnue_architecture.h | 6 +- src/nnue/nnue_feature_transformer.h | 45 ++-- src/nnue/nnz_helper.h | 163 ++++++++++++ src/nnue/simd.h | 3 +- 7 files changed, 266 insertions(+), 220 deletions(-) create mode 100644 src/nnue/nnz_helper.h diff --git a/src/Makefile b/src/Makefile index 2d3835311..2809b248a 100644 --- a/src/Makefile +++ b/src/Makefile @@ -81,7 +81,7 @@ HEADERS = benchmark.h bitboard.h evaluate.h misc.h movegen.h movepick.h history. nnue/layers/affine_transform.h nnue/layers/affine_transform_sparse_input.h \ nnue/layers/clipped_relu.h nnue/layers/sqr_clipped_relu.h nnue/nnue_accumulator.h \ nnue/nnue_architecture.h nnue/nnue_common.h nnue/nnue_feature_transformer.h nnue/simd.h \ - position.h search.h syzygy/tbprobe.h thread.h thread_win32_osx.h timeman.h \ + nnue/nnz_helper.h position.h search.h syzygy/tbprobe.h thread.h thread_win32_osx.h timeman.h \ tt.h tune.h types.h uci.h ucioption.h perft.h nnue/network.h engine.h score.h numa.h memory.h shm.h shm_linux.h OBJS = $(notdir $(SRCS:.cpp=.o)) diff --git a/src/nnue/layers/affine_transform_sparse_input.h b/src/nnue/layers/affine_transform_sparse_input.h index b0651c99f..5fb3bf2c9 100644 --- a/src/nnue/layers/affine_transform_sparse_input.h +++ b/src/nnue/layers/affine_transform_sparse_input.h @@ -31,6 +31,7 @@ #include "../../memory.h" #include "../simd.h" #include "../nnue_common.h" +#include "../nnz_helper.h" /* This file contains the definition for a fully connected layer (aka affine transform) with block sparse input. @@ -64,35 +65,6 @@ class AffineTransformSparseInput { static constexpr IndexType ChunkSize = 1; #endif -#if defined(USE_NEON) - using NNZOutputType = std::conditional_t<(InDims <= 1024), std::uint8_t, std::uint16_t>; -#else - using NNZOutputType = std::uint16_t; -#endif - -#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) - alignas(CacheLineSize) static constexpr struct OffsetIndices { - - NNZOutputType offset_indices[256][8]; - - constexpr OffsetIndices() : - offset_indices() { - for (int i = 0; i < 256; ++i) - { - std::uint64_t j = i, k = 0; - while (j) - { - offset_indices[i][k++] = constexpr_lsb(j); - j &= j - 1; - } - while (k < 8) - offset_indices[i][k++] = 0; - } - } - - } Lookup = {}; -#endif - using OutputBuffer = OutputType[PaddedOutputDimensions]; // Hash value embedded in the evaluation file @@ -145,7 +117,9 @@ class AffineTransformSparseInput { } // Forward propagation - void propagate(const InputType* input, OutputType* output) const { + void propagate(const InputType* input, + OutputType* output, + [[maybe_unused]] const NNZInfo& nnzInfo) const { #if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) #if defined(USE_AVX512) @@ -189,41 +163,31 @@ class AffineTransformSparseInput { #define vec_add_dpbusd_32 SIMD::lsx_m128_add_dpbusd_epi32 #endif constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType); - constexpr IndexType NumChunks = ceil_to_multiple(InputDimensions, 8) / ChunkSize; - constexpr IndexType NumAccums = OutputDimensions / OutputSimdWidth; + constexpr IndexType NumAccums = OutputDimensions / OutputSimdWidth; // If we're using high-latency dot product instructions, split the accumulators // to create 3 separate dependency chains and merge at the end constexpr IndexType NumRegs = - #if defined(USE_VNNI) || defined(USE_NEON_DOTPROD) || defined(USE_LASX) + #if defined(USE_VNNI) && defined(USE_AVX512) 3 * NumAccums; #else NumAccums; #endif - NNZOutputType nnz[NumChunks]; - IndexType count; - - // Find indices of nonzero 32-bit blocks - find_nnz(input, nnz, count); const outvec_t* biasvec = reinterpret_cast(biases); outvec_t acc[NumRegs]; for (IndexType k = 0; k < NumAccums; ++k) acc[k] = biasvec[k]; - const auto* start = nnz; - const auto* end = nnz + count; - - // convince GCC to not do weird pointer arithmetic in the following loop + // convince GCC to not do weird pointer arithmetic in the following loops const std::int8_t* weights_cp = weights; - #if defined(USE_VNNI) || defined(USE_NEON_DOTPROD) || defined(USE_LASX) - #if defined(USE_NEON_DOTPROD) - for (IndexType k = NumAccums; k < NumRegs; ++k) - acc[k] = vdupq_n_s32(0); - #else + + #if defined(USE_AVX512) + const auto* start = nnzInfo.nnz; + const auto* end = nnzInfo.nnz + nnzInfo.count; + for (IndexType k = NumAccums; k < NumRegs; ++k) acc[k] = vec_zero(); - #endif - + #if defined(USE_VNNI) while (start < end - 2) { const std::ptrdiff_t i0 = *start++; @@ -248,14 +212,11 @@ class AffineTransformSparseInput { vec_add_dpbusd_32(acc[k + 2 * NumAccums], in2, col2[k]); } } - #if defined(USE_NEON_DOTPROD) - for (IndexType k = 0; k < NumAccums; ++k) - acc[k] = vaddq_s32(vaddq_s32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]); - #else + for (IndexType k = 0; k < NumAccums; ++k) acc[k] = vec_add_32(vec_add_32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]); #endif - #endif + while (start < end) { const std::ptrdiff_t i = *start++; @@ -265,7 +226,49 @@ class AffineTransformSparseInput { for (IndexType k = 0; k < NumAccums; ++k) vec_add_dpbusd_32(acc[k], in, col[k]); } + #else + static_assert(InputDimensions % 256 == 0); + for (IndexType k = 0; k < InputDimensions / 256; ++k) + { + uint64_t bits = load_as(nnzInfo.bitset + k * 8); + ptrdiff_t base = k * 64; + + auto* base_addr = input + base * sizeof(std::int32_t); + auto* weights_base = &weights_cp[base * OutputDimensions * ChunkSize]; + + #if defined(USE_NEON_DOTPROD) && defined(__GNUC__) && !defined(__clang__) + // GCC 15 pessimizes the following code on ARM64 by eliding the intermediate + // computation of key pointers (base_addr, weights_base, col, input_addr), leading + // to a lot of redundant indexing arithmetic in the while (bits) loop. The + // optimization barriers force these pointers to be calculated and used. + #if __GNUC__ >= 15 + #define FIX_GCC15_MISOPTIMIZATION + #endif + #endif + + #ifdef FIX_GCC15_MISOPTIMIZATION + asm("" : "+r"(base_addr), "+r"(weights_base)); // opt barrier + #endif + + while (bits) + { + ptrdiff_t i = pop_lsb(bits); + const auto* input_addr = base_addr + i * sizeof(std::int32_t); + auto col = + reinterpret_cast(&weights_base[i * OutputDimensions * ChunkSize]); + + #ifdef FIX_GCC15_MISOPTIMIZATION + asm("" : "+r"(col), "+r"(input_addr)); + #undef FIX_GCC15_MISOPTIMIZATION + #endif + + const invec_t in = vec_set_32(load_as(input_addr)); + for (IndexType l = 0; l < NumAccums; ++l) + vec_add_dpbusd_32(acc[l], in, col[l]); + } + } + #endif outvec_t* outptr = reinterpret_cast(output); for (IndexType k = 0; k < NumAccums; ++k) outptr[k] = acc[k]; @@ -283,144 +286,6 @@ class AffineTransformSparseInput { } private: -#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) - - #if defined(__GNUC__) || defined(__clang__) - #define RESTRICT __restrict__ - #elif defined(_MSC_VER) - #define RESTRICT __restrict - #else - #define RESTRICT - #endif - - // Find indices of nonzero 32-bit blocks in a packed - // std::uint8_t buffer. NumChunks is the number of blocks. - template - static void find_nnz(const std::uint8_t* RESTRICT input, - NNZOutputType* RESTRICT out, - IndexType& count_out) { - - #if defined(USE_AVX512ICL) - - constexpr IndexType SimdWidthIn = 64; // 512 bits - constexpr IndexType SimdWidthOut = 32; // 512 bits / 16 bits - constexpr IndexType SimdChunks = NumChunks / SimdWidthOut; - const __m512i increment = _mm512_set1_epi16(SimdWidthOut); - __m512i base = _mm512_set_epi16( // Same permute order as _mm512_packus_epi32() - 31, 30, 29, 28, 15, 14, 13, 12, 27, 26, 25, 24, 11, 10, 9, 8, 23, 22, 21, 20, 7, 6, 5, 4, - 19, 18, 17, 16, 3, 2, 1, 0); - - IndexType count = 0; - for (IndexType i = 0; i < SimdChunks; ++i) - { - const __m512i inputV0 = _mm512_load_si512(input + i * 2 * SimdWidthIn); - const __m512i inputV1 = _mm512_load_si512(input + i * 2 * SimdWidthIn + SimdWidthIn); - - // Get a bitmask and gather non zero indices - const __m512i inputV01 = _mm512_packs_epi32(inputV0, inputV1); - const __mmask32 nnzMask = _mm512_test_epi16_mask(inputV01, inputV01); - - // Avoid _mm512_mask_compressstoreu_epi16() as it's 256 uOps on Zen4 - __m512i nnz = _mm512_maskz_compress_epi16(nnzMask, base); - _mm512_storeu_si512(out + count, nnz); - - count += popcount(nnzMask); - base = _mm512_add_epi16(base, increment); - } - count_out = count; - - #elif defined(USE_AVX512) - - constexpr IndexType SimdWidth = 16; // 512 bits / 32 bits - constexpr IndexType SimdChunks = NumChunks / SimdWidth; - const __m512i increment = _mm512_set1_epi32(SimdWidth); - __m512i base = _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); - - IndexType count = 0; - for (IndexType i = 0; i < SimdChunks; ++i) - { - const __m512i inputV = _mm512_load_si512(input + i * SimdWidth * sizeof(std::uint32_t)); - - // Get a bitmask and gather non zero indices - const __mmask16 nnzMask = _mm512_test_epi32_mask(inputV, inputV); - const __m512i nnzV = _mm512_maskz_compress_epi32(nnzMask, base); - _mm512_mask_cvtepi32_storeu_epi16(out + count, 0xFFFF, nnzV); - count += popcount(nnzMask); - base = _mm512_add_epi32(base, increment); - } - count_out = count; - - #else - #if defined(USE_NEON) - // NEON path using uint8_t NNZOutputType - if constexpr (std::is_same_v) - { - static_assert(NumChunks <= 256, "NumChunks must be <= 256"); - - static constexpr std::uint16_t nnzMask[8] = {1, 2, 4, 8, 16, 32, 64, 128}; - constexpr IndexType SimdChunks = NumChunks / 8; - const auto inputVector = reinterpret_cast(input); - - IndexType count = 0; - uint64_t base = 0ULL; - const uint64_t increment = 0x0808080808080808ULL; - for (IndexType i = 0; i < SimdChunks; ++i) - { - const uint32x4_t v0 = inputVector[i * 2]; - const uint32x4_t v1 = inputVector[i * 2 + 1]; - const uint16x8_t nnz = - vcombine_u16(vqmovn_u32(vtstq_u32(v0, v0)), vqmovn_u32(vtstq_u32(v1, v1))); - const uint16_t lookup = vaddvq_u16(vandq_u16(nnz, vld1q_u16(nnzMask))); - - uint64_t offsets; - std::memcpy(&offsets, Lookup.offset_indices[lookup], sizeof(offsets)); - const uint64_t indices = offsets + base; - std::memcpy(out + count, &indices, sizeof(indices)); - - count += popcount(lookup); - base += increment; - } - count_out = count; - } - else - #endif - { - using namespace SIMD; - - constexpr IndexType InputSimdWidth = sizeof(vec_uint_t) / sizeof(std::int32_t); - // Outputs are processed 8 elements at a time, even if the SIMD width is narrower - constexpr IndexType SimdChunkSize = 8; - constexpr IndexType SimdChunks = NumChunks / SimdChunkSize; - constexpr IndexType InputsPerChunk = SimdChunkSize / InputSimdWidth; - - static_assert(InputsPerChunk > 0, "SIMD width too wide"); - - const auto inputVector = reinterpret_cast(input); - IndexType count = 0; - vec128_t base = vec128_zero; - const vec128_t increment = vec128_set_16(8); - for (IndexType i = 0; i < SimdChunks; ++i) - { - // bitmask of nonzero values in this chunk - unsigned nnz = 0; - for (IndexType j = 0; j < InputsPerChunk; ++j) - { - const vec_uint_t inputChunk = inputVector[i * InputsPerChunk + j]; - nnz |= unsigned(vec_nnz(inputChunk)) << (j * InputSimdWidth); - } - const vec128_t offsets = - vec128_load(reinterpret_cast(&Lookup.offset_indices[nnz])); - vec128_storeu(reinterpret_cast(out + count), vec128_add(base, offsets)); - count += popcount(nnz); - base = vec128_add(base, increment); - } - count_out = count; - } - #endif - } - #undef RESTRICT -#endif - using BiasType = OutputType; using WeightType = std::int8_t; diff --git a/src/nnue/network.cpp b/src/nnue/network.cpp index e2c77da48..ab8084924 100644 --- a/src/nnue/network.cpp +++ b/src/nnue/network.cpp @@ -35,6 +35,7 @@ #include "nnue_architecture.h" #include "nnue_common.h" #include "nnue_misc.h" +#include "nnz_helper.h" // Macro to embed the default efficiently updatable neural network (NNUE) file // data in the engine binary (using incbin.h, by Dale Weiler). @@ -149,10 +150,12 @@ NetworkOutput Network::evaluate(const Position& pos, ASSERT_ALIGNED(transformedFeatures, alignment); - const int bucket = (pos.count() - 1) / 4; - const auto psqt = - featureTransformer.transform(pos, accumulatorStack, cache, transformedFeatures, bucket); - const auto positional = network[bucket].propagate(transformedFeatures); + NNZInfo nnzInfo; + + const int bucket = (pos.count() - 1) / 4; + const auto psqt = featureTransformer.transform(pos, accumulatorStack, cache, + transformedFeatures, bucket, nnzInfo); + const auto positional = network[bucket].propagate(transformedFeatures, nnzInfo); return {static_cast(psqt / OutputScale), static_cast(positional / OutputScale)}; } @@ -211,9 +214,10 @@ NnueEvalTrace Network::trace_evaluate(const Position& pos, t.correctBucket = (pos.count() - 1) / 4; for (IndexType bucket = 0; bucket < LayerStacks; ++bucket) { - const auto materialist = - featureTransformer.transform(pos, accumulatorStack, cache, transformedFeatures, bucket); - const auto positional = network[bucket].propagate(transformedFeatures); + NNZInfo nnzInfo; + const auto materialist = featureTransformer.transform(pos, accumulatorStack, cache, + transformedFeatures, bucket, nnzInfo); + const auto positional = network[bucket].propagate(transformedFeatures, nnzInfo); t.psqt[bucket] = static_cast(materialist / OutputScale); t.positional[bucket] = static_cast(positional / OutputScale); diff --git a/src/nnue/nnue_architecture.h b/src/nnue/nnue_architecture.h index 02c38df6d..f90b50757 100644 --- a/src/nnue/nnue_architecture.h +++ b/src/nnue/nnue_architecture.h @@ -32,6 +32,7 @@ #include "layers/clipped_relu.h" #include "layers/sqr_clipped_relu.h" #include "nnue_common.h" +#include "nnz_helper.h" namespace Stockfish::Eval::NNUE { @@ -94,7 +95,8 @@ struct NetworkArchitecture { && fc_2.write_parameters(stream); } - std::int32_t propagate(const TransformedFeatureType* transformedFeatures) const { + std::int32_t propagate(const TransformedFeatureType* transformedFeatures, + const NNZInfo& nnzInfo) const { struct alignas(CacheLineSize) Buffer { alignas(CacheLineSize) typename decltype(fc_0)::OutputBuffer fc_0_out; alignas(CacheLineSize) typename decltype(ac_sqr_0)::OutputType @@ -109,7 +111,7 @@ struct NetworkArchitecture { Buffer buffer; - fc_0.propagate(transformedFeatures, buffer.fc_0_out); + fc_0.propagate(transformedFeatures, buffer.fc_0_out, nnzInfo); ac_sqr_0.propagate(buffer.fc_0_out, buffer.ac_sqr_0_out); ac_0.propagate(buffer.fc_0_out, buffer.ac_0_out); std::memcpy(buffer.ac_sqr_0_out + FC_0_OUTPUTS, buffer.ac_0_out, diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index f74ecae7f..a9c294447 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -204,11 +204,12 @@ class FeatureTransformer { } // Convert input features - std::int32_t transform(const Position& pos, - AccumulatorStack& accumulatorStack, - AccumulatorCaches& cache, - OutputType* output, - int bucket) const { + std::int32_t transform(const Position& pos, + AccumulatorStack& accumulatorStack, + AccumulatorCaches& cache, + OutputType* output, + int bucket, + NNZInfo& nnzInfo) const { using namespace SIMD; accumulatorStack.evaluate(pos, *this, cache); @@ -232,6 +233,8 @@ class FeatureTransformer { { const IndexType offset = (HalfDimensions / 2) * p; + [[maybe_unused]] auto cursor = nnzInfo.make_cursor(p); + #if defined(VECTOR) constexpr IndexType OutputChunkSize = MaxChunkSize; @@ -310,22 +313,30 @@ class FeatureTransformer { reinterpret_cast(&(threatAccumulation[perspectives[p]][0])); const vec_t* tin1 = reinterpret_cast( &(threatAccumulation[perspectives[p]][HalfDimensions / 2])); - for (IndexType j = 0; j < NumOutputChunks; ++j) + for (IndexType j = 0; j < NumOutputChunks; j += 2) { - const vec_t acc0a = vec_add_16(in0[j * 2 + 0], tin0[j * 2 + 0]); - const vec_t acc0b = vec_add_16(in0[j * 2 + 1], tin0[j * 2 + 1]); - const vec_t acc1a = vec_add_16(in1[j * 2 + 0], tin1[j * 2 + 0]); - const vec_t acc1b = vec_add_16(in1[j * 2 + 1], tin1[j * 2 + 1]); + vec_t packed[2]; + for (IndexType k = 0; k < 2; ++k) + { + const IndexType i = (j + k) * 2; - const vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, One), Zero), shift); - const vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, One), Zero), shift); - const vec_t sum1a = vec_min_16(acc1a, One); - const vec_t sum1b = vec_min_16(acc1b, One); + vec_t acc0a = vec_add_16(in0[i + 0], tin0[i + 0]); + vec_t acc0b = vec_add_16(in0[i + 1], tin0[i + 1]); + vec_t acc1a = vec_add_16(in1[i + 0], tin1[i + 0]); + vec_t acc1b = vec_add_16(in1[i + 1], tin1[i + 1]); - const vec_t pa = vec_mulhi_16(sum0a, sum1a); - const vec_t pb = vec_mulhi_16(sum0b, sum1b); + vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, One), Zero), shift); + vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, One), Zero), shift); + vec_t sum1a = vec_min_16(acc1a, One); + vec_t sum1b = vec_min_16(acc1b, One); - out[j] = vec_packus_16(pa, pb); + vec_t pa = vec_mulhi_16(sum0a, sum1a); + vec_t pb = vec_mulhi_16(sum0b, sum1b); + + packed[k] = out[j + k] = vec_packus_16(pa, pb); + } + + cursor.record2(packed[0], packed[1]); } #else diff --git a/src/nnue/nnz_helper.h b/src/nnue/nnz_helper.h new file mode 100644 index 000000000..3db005191 --- /dev/null +++ b/src/nnue/nnz_helper.h @@ -0,0 +1,163 @@ +/* + Stockfish, a UCI chess playing engine derived from Glaurung 2.1 + Copyright (C) 2004-2026 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 + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + Stockfish is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . +*/ + +#ifndef NNZ_HELPER_H_INCLUDED +#define NNZ_HELPER_H_INCLUDED + +#include + +#include "nnue_common.h" +#include "simd.h" + +namespace Stockfish::Eval::NNUE { + +template +struct NNZInfo { + +#if defined(USE_AVX512) + unsigned count = 0; + // indices of non-zero chunks + uint16_t nnz[Dimensions / 4]; + + #ifdef USE_AVX512ICL + alignas(64) static constexpr auto Indices = []() { + std::array, 2> indices{}; + for (int i = 0; i < 2; ++i) + { + indices[i] = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23, + 8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; + for (uint16_t& m : indices[i]) + m += i * Dimensions / 8; + } + return indices; + }(); + #else + alignas(64) static constexpr auto Indices = []() { + std::array, 2> indices{}; + for (int i = 0; i < 2; ++i) + { + indices[i] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; + for (uint32_t& m : indices[i]) + m += i * Dimensions / 8; + } + return indices; + }(); + #endif + + struct NNZCursor { + NNZInfo& info; + __m512i indices; + unsigned count; + + NNZCursor(NNZInfo& info_, bool perspective, unsigned count_) : + info(info_), + count(count_) { + indices = _mm512_load_si512(&Indices[perspective]); + } + + void record2(SIMD::vec_t neurons1, SIMD::vec_t neurons2) { + #if defined(USE_AVX512ICL) + const __m512i increment = _mm512_set1_epi16(32); + + // Get a bitmask and gather non zero indices + const __m512i inputV01 = _mm512_packs_epi32(neurons1, neurons2); + const __mmask32 nnzMask = _mm512_test_epi16_mask(inputV01, inputV01); + + // Avoid _mm512_mask_compressstoreu_epi16() as it's 256 uOps on Zen4 + __m512i nnz = _mm512_maskz_compress_epi16(nnzMask, indices); + _mm512_storeu_si512(info.nnz + count, nnz); + + count += popcount(nnzMask); + indices = _mm512_add_epi16(indices, increment); + #else + for (auto neurons : {neurons1, neurons2}) + { + const __m512i increment = _mm512_set1_epi32(16); + // Get a bitmask and gather non zero indices + const __mmask16 nnzMask = _mm512_test_epi32_mask(neurons, neurons); + const __m512i nnzV = _mm512_maskz_compress_epi32(nnzMask, indices); + _mm512_mask_cvtepi32_storeu_epi16(info.nnz + count, 0xFFFF, nnzV); + + count += popcount(nnzMask); + indices = _mm512_add_epi32(indices, increment); + } + #endif + } + + ~NNZCursor() { info.count = count; } + }; + + NNZCursor make_cursor(bool perspective) { return {*this, perspective, count}; } +#else + // Each 8-bit chunk + uint8_t bitset[(Dimensions + 31) / 32]; + + struct NNZCursor { + uint8_t* out; + + NNZCursor(NNZInfo& info, bool perspective) { + out = info.bitset + perspective * Dimensions / 64; + } + + #if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) + void record2(SIMD::vec_t neurons1, SIMD::vec_t neurons2) { + using namespace SIMD; + + #ifdef USE_NEON + alignas(16) static constexpr uint16_t Mask8[8] = {1, 16, 2, 32, 4, 64, 8, 128}; + + uint32x4_t n1 = vreinterpretq_u32_s16(neurons1); + uint32x4_t n2 = vreinterpretq_u32_s16(neurons2); + + const uint32x4_t t1 = vtstq_u32(n1, n1); + const uint32x4_t t2 = vtstq_u32(n2, n2); + + const uint16x8_t packed = + vtrn1q_u16(vreinterpretq_u16_u32(t1), vreinterpretq_u16_u32(t2)); + const uint16x8_t bits = vandq_u16(packed, vld1q_u16(Mask8)); + + *out++ = vaddvq_u16(bits); + #else + auto m1 = vec_nnz(neurons1); + auto m2 = vec_nnz(neurons2); + + if (sizeof(neurons1) == 16) + { + *out++ = m1 + (m2 << 4); + } + else + { + size_t bytes = sizeof(neurons1) / 32; + memcpy(out, &m1, bytes); + out += bytes; + memcpy(out, &m2, bytes); + out += bytes; + } + #endif + } + #elif defined(VECTOR) + void record2(SIMD::vec_t, SIMD::vec_t) {} + #endif + }; + + NNZCursor make_cursor(bool perspective) { return {*this, perspective}; } +#endif +}; +} // namespace Stockfish + +#endif diff --git a/src/nnue/simd.h b/src/nnue/simd.h index 44cafe569..d1592e88b 100644 --- a/src/nnue/simd.h +++ b/src/nnue/simd.h @@ -217,7 +217,8 @@ using vec_uint_t __attribute__((may_alias)) = uint32x4_t; #define vec_zero_psqt() psqt_vec_t{0} static constexpr std::uint32_t Mask[4] = {1, 2, 4, 8}; - #define vec_nnz(a) vaddvq_u32(vandq_u32(vtstq_u32(a, a), vld1q_u32(Mask))) + #define vec_nnz(a) \ + vaddvq_u32(vandq_u32(vtstq_u32((uint32x4_t) a, (uint32x4_t) a), vld1q_u32(Mask))) #define vec128_zero vdupq_n_u16(0) #define vec128_set_16(a) vdupq_n_u16(a) #define vec128_load(a) vld1q_u16(reinterpret_cast(a))