mirror of
https://github.com/official-stockfish/Stockfish.git
synced 2026-07-22 12:47:08 +00:00
Timothy HerchenandJoost VandeVondele
1554a2ca0b
Passed STC https://tests.stockfishchess.org/tests/view/69f6df65b64b50e29dbed427 LLR: 2.94 (-2.94,2.94) <0.00,2.00> Total: 112224 W: 29195 L: 28782 D: 54247 Ptnml(0-2): 312, 12345, 30381, 12766, 308 Local speedup: ``` Result of 100 runs base (...kfish.master) = 2050387 +/- 5497 test (./stockfish ) = 2069598 +/- 4426 diff = +19211 +/- 7468 speedup = +0.0094 P(speedup > 0) = 1.0000 ``` closes https://github.com/official-stockfish/Stockfish/pull/6797 No functional change
422 lines
16 KiB
C++
422 lines
16 KiB
C++
/*
|
|
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 <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
// Definition of layer AffineTransformSparseInput of NNUE evaluation function
|
|
|
|
#ifndef NNUE_LAYERS_AFFINE_TRANSFORM_SPARSE_INPUT_H_INCLUDED
|
|
#define NNUE_LAYERS_AFFINE_TRANSFORM_SPARSE_INPUT_H_INCLUDED
|
|
|
|
#include <algorithm>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
|
|
#include "../../bitboard.h"
|
|
#include "../../memory.h"
|
|
#include "../simd.h"
|
|
#include "../nnue_common.h"
|
|
|
|
/*
|
|
This file contains the definition for a fully connected layer (aka affine transform) with block sparse input.
|
|
*/
|
|
|
|
namespace Stockfish::Eval::NNUE::Layers {
|
|
|
|
// Sparse input implementation
|
|
template<IndexType InDims, IndexType OutDims>
|
|
class AffineTransformSparseInput {
|
|
public:
|
|
// Input/output type
|
|
using InputType = std::uint8_t;
|
|
using OutputType = std::int32_t;
|
|
|
|
// Number of input/output dimensions
|
|
static constexpr IndexType InputDimensions = InDims;
|
|
static constexpr IndexType OutputDimensions = OutDims;
|
|
|
|
static_assert(OutputDimensions % 16 == 0,
|
|
"Only implemented for OutputDimensions divisible by 16.");
|
|
|
|
static constexpr IndexType PaddedInputDimensions =
|
|
ceil_to_multiple<IndexType>(InputDimensions, MaxSimdWidth);
|
|
static constexpr IndexType PaddedOutputDimensions =
|
|
ceil_to_multiple<IndexType>(OutputDimensions, MaxSimdWidth);
|
|
|
|
#if (USE_SSSE3 | (USE_NEON >= 8))
|
|
static constexpr IndexType ChunkSize = 4;
|
|
#else
|
|
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 (USE_SSSE3 | (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
|
|
static constexpr std::uint32_t get_hash_value(std::uint32_t prevHash) {
|
|
std::uint32_t hashValue = 0xCC03DAE4u;
|
|
hashValue += OutputDimensions;
|
|
hashValue ^= prevHash >> 1;
|
|
hashValue ^= prevHash << 31;
|
|
return hashValue;
|
|
}
|
|
|
|
static constexpr IndexType get_weight_index_scrambled(IndexType i) {
|
|
return (i / ChunkSize) % (PaddedInputDimensions / ChunkSize) * OutputDimensions * ChunkSize
|
|
+ i / PaddedInputDimensions * ChunkSize + i % ChunkSize;
|
|
}
|
|
|
|
static constexpr IndexType get_weight_index(IndexType i) {
|
|
#if (USE_SSSE3 | (USE_NEON >= 8))
|
|
return get_weight_index_scrambled(i);
|
|
#else
|
|
return i;
|
|
#endif
|
|
}
|
|
|
|
// Read network parameters
|
|
bool read_parameters(std::istream& stream) {
|
|
read_little_endian<BiasType>(stream, biases, OutputDimensions);
|
|
for (IndexType i = 0; i < OutputDimensions * PaddedInputDimensions; ++i)
|
|
weights[get_weight_index(i)] = read_little_endian<WeightType>(stream);
|
|
|
|
return !stream.fail();
|
|
}
|
|
|
|
// Write network parameters
|
|
bool write_parameters(std::ostream& stream) const {
|
|
write_little_endian<BiasType>(stream, biases, OutputDimensions);
|
|
|
|
for (IndexType i = 0; i < OutputDimensions * PaddedInputDimensions; ++i)
|
|
write_little_endian<WeightType>(stream, weights[get_weight_index(i)]);
|
|
|
|
return !stream.fail();
|
|
}
|
|
|
|
std::size_t get_content_hash() const {
|
|
std::size_t h = 0;
|
|
hash_combine(h, get_raw_data_hash(biases));
|
|
hash_combine(h, get_raw_data_hash(weights));
|
|
hash_combine(h, get_hash_value(0));
|
|
return h;
|
|
}
|
|
|
|
// Forward propagation
|
|
void propagate(const InputType* input, OutputType* output) const {
|
|
|
|
#if (USE_SSSE3 | (USE_NEON >= 8))
|
|
#if defined(USE_AVX512)
|
|
using invec_t = __m512i;
|
|
using outvec_t = __m512i;
|
|
#define vec_add_32 _mm512_add_epi32
|
|
#define vec_set_32 _mm512_set1_epi32
|
|
#define vec_add_dpbusd_32 SIMD::m512_add_dpbusd_epi32
|
|
#elif defined(USE_AVX2)
|
|
using invec_t = __m256i;
|
|
using outvec_t = __m256i;
|
|
#define vec_add_32 _mm256_add_epi32
|
|
#define vec_set_32 _mm256_set1_epi32
|
|
#define vec_add_dpbusd_32 SIMD::m256_add_dpbusd_epi32
|
|
#elif defined(USE_SSSE3)
|
|
using invec_t = __m128i;
|
|
using outvec_t = __m128i;
|
|
#define vec_set_32 _mm_set1_epi32
|
|
#define vec_add_dpbusd_32 SIMD::m128_add_dpbusd_epi32
|
|
#elif defined(USE_NEON_DOTPROD)
|
|
using invec_t = int8x16_t;
|
|
using outvec_t = int32x4_t;
|
|
#define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a))
|
|
#define vec_add_dpbusd_32 SIMD::dotprod_m128_add_dpbusd_epi32
|
|
#elif defined(USE_NEON)
|
|
using invec_t = int8x16_t;
|
|
using outvec_t = int32x4_t;
|
|
#define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a))
|
|
#define vec_add_dpbusd_32 SIMD::neon_m128_add_dpbusd_epi32
|
|
#endif
|
|
constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType);
|
|
constexpr IndexType NumChunks = ceil_to_multiple<IndexType>(InputDimensions, 8) / ChunkSize;
|
|
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)
|
|
3 * NumAccums;
|
|
#else
|
|
NumAccums;
|
|
#endif
|
|
NNZOutputType nnz[NumChunks];
|
|
IndexType count;
|
|
|
|
// Find indices of nonzero 32-bit blocks
|
|
find_nnz<NumChunks>(input, nnz, count);
|
|
|
|
const outvec_t* biasvec = reinterpret_cast<const outvec_t*>(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
|
|
const std::int8_t* weights_cp = weights;
|
|
#if defined(USE_VNNI) || defined(USE_NEON_DOTPROD)
|
|
#if defined(USE_VNNI)
|
|
for (IndexType k = NumAccums; k < NumRegs; ++k)
|
|
acc[k] = vec_zero();
|
|
#else
|
|
for (IndexType k = NumAccums; k < NumRegs; ++k)
|
|
acc[k] = vdupq_n_s32(0);
|
|
#endif
|
|
|
|
while (start < end - 2)
|
|
{
|
|
const std::ptrdiff_t i0 = *start++;
|
|
const std::ptrdiff_t i1 = *start++;
|
|
const std::ptrdiff_t i2 = *start++;
|
|
const invec_t in0 =
|
|
vec_set_32(load_as<std::int32_t>(input + i0 * sizeof(std::int32_t)));
|
|
const invec_t in1 =
|
|
vec_set_32(load_as<std::int32_t>(input + i1 * sizeof(std::int32_t)));
|
|
const invec_t in2 =
|
|
vec_set_32(load_as<std::int32_t>(input + i2 * sizeof(std::int32_t)));
|
|
const auto col0 =
|
|
reinterpret_cast<const invec_t*>(&weights_cp[i0 * OutputDimensions * ChunkSize]);
|
|
const auto col1 =
|
|
reinterpret_cast<const invec_t*>(&weights_cp[i1 * OutputDimensions * ChunkSize]);
|
|
const auto col2 =
|
|
reinterpret_cast<const invec_t*>(&weights_cp[i2 * OutputDimensions * ChunkSize]);
|
|
for (IndexType k = 0; k < NumAccums; ++k)
|
|
{
|
|
vec_add_dpbusd_32(acc[k], in0, col0[k]);
|
|
vec_add_dpbusd_32(acc[k + NumAccums], in1, col1[k]);
|
|
vec_add_dpbusd_32(acc[k + 2 * NumAccums], in2, col2[k]);
|
|
}
|
|
}
|
|
#if defined(USE_VNNI)
|
|
for (IndexType k = 0; k < NumAccums; ++k)
|
|
acc[k] = vec_add_32(vec_add_32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]);
|
|
#else
|
|
for (IndexType k = 0; k < NumAccums; ++k)
|
|
acc[k] = vaddq_s32(vaddq_s32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]);
|
|
#endif
|
|
#endif
|
|
while (start < end)
|
|
{
|
|
const std::ptrdiff_t i = *start++;
|
|
const invec_t in = vec_set_32(load_as<std::int32_t>(input + i * sizeof(std::int32_t)));
|
|
const auto col =
|
|
reinterpret_cast<const invec_t*>(&weights_cp[i * OutputDimensions * ChunkSize]);
|
|
for (IndexType k = 0; k < NumAccums; ++k)
|
|
vec_add_dpbusd_32(acc[k], in, col[k]);
|
|
}
|
|
|
|
outvec_t* outptr = reinterpret_cast<outvec_t*>(output);
|
|
for (IndexType k = 0; k < NumAccums; ++k)
|
|
outptr[k] = acc[k];
|
|
|
|
#undef vec_set_32
|
|
#undef vec_add_dpbusd_32
|
|
#ifdef vec_add_32
|
|
#undef vec_add_32
|
|
#endif
|
|
#else
|
|
// Use dense implementation for the other architectures.
|
|
affine_transform_non_ssse3<InputDimensions, PaddedInputDimensions, OutputDimensions>(
|
|
output, weights, biases, input);
|
|
#endif
|
|
}
|
|
|
|
private:
|
|
#if (USE_SSSE3 | (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<IndexType NumChunks>
|
|
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<NNZOutputType, std::uint8_t>)
|
|
{
|
|
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<const uint32x4_t*>(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<const vec_uint_t*>(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<const vec128_t*>(&Lookup.offset_indices[nnz]));
|
|
vec128_storeu(reinterpret_cast<vec128_t*>(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;
|
|
|
|
alignas(CacheLineSize) BiasType biases[OutputDimensions];
|
|
alignas(CacheLineSize) WeightType weights[OutputDimensions * PaddedInputDimensions];
|
|
};
|
|
|
|
} // namespace Stockfish::Eval::NNUE::Layers
|
|
|
|
#endif // #ifndef NNUE_LAYERS_AFFINE_TRANSFORM_SPARSE_INPUT_H_INCLUDED
|