From 8be6b142189dfaa4fdee979d863a53522ca46b32 Mon Sep 17 00:00:00 2001 From: anematode Date: Mon, 5 Jan 2026 15:37:49 -0800 Subject: [PATCH] Network loading refactoring closes https://github.com/official-stockfish/Stockfish/pull/6523 No functional change --- src/engine.cpp | 12 ++--- src/nnue/network.h | 7 ++- src/nnue/nnue_common.h | 76 +++++++++++++++-------------- src/nnue/nnue_feature_transformer.h | 24 ++------- 4 files changed, 52 insertions(+), 67 deletions(-) diff --git a/src/engine.cpp b/src/engine.cpp index 355103c7b..01927fc84 100644 --- a/src/engine.cpp +++ b/src/engine.cpp @@ -57,14 +57,10 @@ Engine::Engine(std::optional path) : numaContext(NumaConfig::from_system()), states(new std::deque(1)), threads(), - networks( - numaContext, - // Heap-allocate because sizeof(NN::Networks) is large - std::make_unique( - std::make_unique(NN::EvalFile{EvalFileDefaultNameBig, "None", ""}, - NN::EmbeddedNNUEType::BIG), - std::make_unique(NN::EvalFile{EvalFileDefaultNameSmall, "None", ""}, - NN::EmbeddedNNUEType::SMALL))) { + networks(numaContext, + // Heap-allocate because sizeof(NN::Networks) is large + std::make_unique(NN::EvalFile{EvalFileDefaultNameBig, "None", ""}, + NN::EvalFile{EvalFileDefaultNameSmall, "None", ""})) { pos.set(StartFEN, false, &states->back()); diff --git a/src/nnue/network.h b/src/nnue/network.h index d0e3218ca..cb433718d 100644 --- a/src/nnue/network.h +++ b/src/nnue/network.h @@ -28,7 +28,6 @@ #include #include #include -#include #include "../misc.h" #include "../types.h" @@ -130,9 +129,9 @@ using NetworkSmall = Network; struct Networks { - Networks(std::unique_ptr&& nB, std::unique_ptr&& nS) : - big(std::move(*nB)), - small(std::move(*nS)) {} + Networks(EvalFile bigFile, EvalFile smallFile) : + big(bigFile, EmbeddedNNUEType::BIG), + small(smallFile, EmbeddedNNUEType::SMALL) {} NetworkBig big; NetworkSmall small; diff --git a/src/nnue/nnue_common.h b/src/nnue/nnue_common.h index febe7ca70..27852ac7b 100644 --- a/src/nnue/nnue_common.h +++ b/src/nnue/nnue_common.h @@ -169,52 +169,56 @@ inline void write_little_endian(std::ostream& stream, const IntType* values, std write_little_endian(stream, values[i]); } - // Read N signed integers from the stream s, putting them in the array out. // The stream is assumed to be compressed using the signed LEB128 format. // See https://en.wikipedia.org/wiki/LEB128 for a description of the compression scheme. -template -inline void read_leb_128(std::istream& stream, std::array& out) { +template +inline void read_leb_128_detail(std::istream& stream, + std::array& out, + std::uint32_t& bytes_left, + BufType& buf, + std::uint32_t& buf_pos) { + static_assert(std::is_signed_v, "Not implemented for unsigned types"); + static_assert(sizeof(IntType) <= 4, "Not implemented for types larger than 32 bit"); + + IntType result = 0; + size_t shift = 0, i = 0; + while (i < Count) + { + if (buf_pos == buf.size()) + { + stream.read(reinterpret_cast(buf.data()), + std::min(std::size_t(bytes_left), buf.size())); + buf_pos = 0; + } + + std::uint8_t byte = buf[buf_pos++]; + --bytes_left; + result |= (byte & 0x7f) << (shift % 32); + shift += 7; + + if ((byte & 0x80) == 0) + { + out[i++] = (shift >= 32 || (byte & 0x40) == 0) ? result : result | ~((1 << shift) - 1); + result = 0; + shift = 0; + } + } +} + +template +inline void read_leb_128(std::istream& stream, Arrays&... outs) { // Check the presence of our LEB128 magic string char leb128MagicString[Leb128MagicStringSize]; stream.read(leb128MagicString, Leb128MagicStringSize); assert(strncmp(Leb128MagicString, leb128MagicString, Leb128MagicStringSize) == 0); - static_assert(std::is_signed_v, "Not implemented for unsigned types"); + auto bytes_left = read_little_endian(stream); + std::array buf; + std::uint32_t buf_pos = buf.size(); - const std::uint32_t BUF_SIZE = 4096; - std::uint8_t buf[BUF_SIZE]; - - auto bytes_left = read_little_endian(stream); - - std::uint32_t buf_pos = BUF_SIZE; - for (std::size_t i = 0; i < Count; ++i) - { - IntType result = 0; - size_t shift = 0; - do - { - if (buf_pos == BUF_SIZE) - { - stream.read(reinterpret_cast(buf), std::min(bytes_left, BUF_SIZE)); - buf_pos = 0; - } - - std::uint8_t byte = buf[buf_pos++]; - --bytes_left; - result |= (byte & 0x7f) << shift; - shift += 7; - - if ((byte & 0x80) == 0) - { - out[i] = (sizeof(IntType) * 8 <= shift || (byte & 0x40) == 0) - ? result - : result | ~((1 << shift) - 1); - break; - } - } while (shift < sizeof(IntType) * 8); - } + (read_leb_128_detail(stream, outs, bytes_left, buf, buf_pos), ...); assert(bytes_left == 0); } diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index 99cda2a69..c7b784573 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -152,35 +152,21 @@ class FeatureTransformer { } // Read network parameters - // TODO: This is ugly. Currently LEB128 on the entire L1 necessitates - // reading the weights into a combined array, and then splitting. bool read_parameters(std::istream& stream) { - read_leb_128(stream, biases); + read_leb_128(stream, biases); if (UseThreats) { read_little_endian(stream, threatWeights.data(), ThreatInputDimensions * HalfDimensions); - read_leb_128(stream, weights); + read_leb_128(stream, weights); - auto combinedPsqtWeights = - std::make_unique>(); - - read_leb_128(stream, *combinedPsqtWeights); - - std::copy(combinedPsqtWeights->begin(), - combinedPsqtWeights->begin() + ThreatInputDimensions * PSQTBuckets, - std::begin(threatPsqtWeights)); - - std::copy(combinedPsqtWeights->begin() + ThreatInputDimensions * PSQTBuckets, - combinedPsqtWeights->begin() - + (ThreatInputDimensions + InputDimensions) * PSQTBuckets, - std::begin(psqtWeights)); + read_leb_128(stream, threatPsqtWeights, psqtWeights); } else { - read_leb_128(stream, weights); - read_leb_128(stream, psqtWeights); + read_leb_128(stream, weights); + read_leb_128(stream, psqtWeights); } permute_weights();