Network loading refactoring

closes https://github.com/official-stockfish/Stockfish/pull/6523

No functional change
This commit is contained in:
anematode
2026-01-06 12:02:37 +01:00
committed by Joost VandeVondele
parent d678f839d8
commit 8be6b14218
4 changed files with 52 additions and 67 deletions
+3 -7
View File
@@ -57,14 +57,10 @@ Engine::Engine(std::optional<std::string> path) :
numaContext(NumaConfig::from_system()),
states(new std::deque<StateInfo>(1)),
threads(),
networks(
numaContext,
networks(numaContext,
// Heap-allocate because sizeof(NN::Networks) is large
std::make_unique<NN::Networks>(
std::make_unique<NN::NetworkBig>(NN::EvalFile{EvalFileDefaultNameBig, "None", ""},
NN::EmbeddedNNUEType::BIG),
std::make_unique<NN::NetworkSmall>(NN::EvalFile{EvalFileDefaultNameSmall, "None", ""},
NN::EmbeddedNNUEType::SMALL))) {
std::make_unique<NN::Networks>(NN::EvalFile{EvalFileDefaultNameBig, "None", ""},
NN::EvalFile{EvalFileDefaultNameSmall, "None", ""})) {
pos.set(StartFEN, false, &states->back());
+3 -4
View File
@@ -28,7 +28,6 @@
#include <string>
#include <string_view>
#include <tuple>
#include <utility>
#include "../misc.h"
#include "../types.h"
@@ -130,9 +129,9 @@ using NetworkSmall = Network<SmallNetworkArchitecture, SmallFeatureTransformer>;
struct Networks {
Networks(std::unique_ptr<NetworkBig>&& nB, std::unique_ptr<NetworkSmall>&& 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;
+30 -26
View File
@@ -169,52 +169,56 @@ inline void write_little_endian(std::ostream& stream, const IntType* values, std
write_little_endian<IntType>(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<typename IntType, std::size_t Count>
inline void read_leb_128(std::istream& stream, std::array<IntType, Count>& out) {
// Check the presence of our LEB128 magic string
char leb128MagicString[Leb128MagicStringSize];
stream.read(leb128MagicString, Leb128MagicStringSize);
assert(strncmp(Leb128MagicString, leb128MagicString, Leb128MagicStringSize) == 0);
template<typename BufType, typename IntType, std::size_t Count>
inline void read_leb_128_detail(std::istream& stream,
std::array<IntType, Count>& out,
std::uint32_t& bytes_left,
BufType& buf,
std::uint32_t& buf_pos) {
static_assert(std::is_signed_v<IntType>, "Not implemented for unsigned types");
static_assert(sizeof(IntType) <= 4, "Not implemented for types larger than 32 bit");
const std::uint32_t BUF_SIZE = 4096;
std::uint8_t buf[BUF_SIZE];
auto bytes_left = read_little_endian<std::uint32_t>(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
size_t shift = 0, i = 0;
while (i < Count)
{
if (buf_pos == BUF_SIZE)
if (buf_pos == buf.size())
{
stream.read(reinterpret_cast<char*>(buf), std::min(bytes_left, BUF_SIZE));
stream.read(reinterpret_cast<char*>(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;
result |= (byte & 0x7f) << (shift % 32);
shift += 7;
if ((byte & 0x80) == 0)
{
out[i] = (sizeof(IntType) * 8 <= shift || (byte & 0x40) == 0)
? result
: result | ~((1 << shift) - 1);
break;
out[i++] = (shift >= 32 || (byte & 0x40) == 0) ? result : result | ~((1 << shift) - 1);
result = 0;
shift = 0;
}
} while (shift < sizeof(IntType) * 8);
}
}
template<typename... Arrays>
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);
auto bytes_left = read_little_endian<std::uint32_t>(stream);
std::array<std::uint8_t, 8192> buf;
std::uint32_t buf_pos = buf.size();
(read_leb_128_detail(stream, outs, bytes_left, buf, buf_pos), ...);
assert(bytes_left == 0);
}
+5 -19
View File
@@ -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<BiasType>(stream, biases);
read_leb_128(stream, biases);
if (UseThreats)
{
read_little_endian<ThreatWeightType>(stream, threatWeights.data(),
ThreatInputDimensions * HalfDimensions);
read_leb_128<WeightType>(stream, weights);
read_leb_128(stream, weights);
auto combinedPsqtWeights =
std::make_unique<std::array<PSQTWeightType, TotalInputDimensions * PSQTBuckets>>();
read_leb_128<PSQTWeightType>(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<WeightType>(stream, weights);
read_leb_128<PSQTWeightType>(stream, psqtWeights);
read_leb_128(stream, weights);
read_leb_128(stream, psqtWeights);
}
permute_weights();