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))