From a46087ee30ace00e1edd8963ee33b0a1b87031b6 Mon Sep 17 00:00:00 2001 From: maxim Date: Tue, 13 Jun 2023 23:09:05 +0300 Subject: [PATCH] Compressed network parameters Implemented LEB128 (de)compression for the feature transformer. Reduces embedded network size from 70 MiB to 39 Mib. The new nn-78bacfcee510.nnue corresponds to the master net compressed. closes https://github.com/official-stockfish/Stockfish/pull/4617 No functional change --- src/evaluate.h | 2 +- src/nnue/nnue_common.h | 77 +++++++++++++++++++++++++++++ src/nnue/nnue_feature_transformer.h | 12 ++--- 3 files changed, 84 insertions(+), 7 deletions(-) diff --git a/src/evaluate.h b/src/evaluate.h index c35b2f07..33effb1c 100644 --- a/src/evaluate.h +++ b/src/evaluate.h @@ -39,7 +39,7 @@ namespace Eval { // The default net name MUST follow the format nn-[SHA256 first 12 digits].nnue // for the build process (profile-build and fishtest) to work. Do not change the // name of the macro, as it is used in the Makefile. - #define EvalFileDefaultName "nn-cd2ff4716c34.nnue" + #define EvalFileDefaultName "nn-78bacfcee510.nnue" namespace NNUE { diff --git a/src/nnue/nnue_common.h b/src/nnue/nnue_common.h index 12309d26..d338527d 100644 --- a/src/nnue/nnue_common.h +++ b/src/nnue/nnue_common.h @@ -57,6 +57,9 @@ namespace Stockfish::Eval::NNUE { // Size of cache line (in bytes) constexpr std::size_t CacheLineSize = 64; + constexpr const char Leb128MagicString[] = "COMPRESSED_LEB128"; + constexpr const std::size_t Leb128MagicStringSize = sizeof(Leb128MagicString) - 1; + // SIMD width (in bytes) #if defined(USE_AVX2) constexpr std::size_t SimdWidth = 32; @@ -159,6 +162,80 @@ namespace Stockfish::Eval::NNUE { write_little_endian(stream, values[i]); } + template + inline void read_leb_128(std::istream& stream, IntType* out, std::size_t count) { + static_assert(std::is_signed_v, "Not implemented for unsigned types"); + char leb128MagicString[Leb128MagicStringSize]; + stream.read(leb128MagicString, Leb128MagicStringSize); + assert(strncmp(Leb128MagicString, leb128MagicString, Leb128MagicStringSize) == 0); + 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); + } + assert(bytes_left == 0); + } + + template + inline void write_leb_128(std::ostream& stream, const IntType* values, std::size_t count) { + static_assert(std::is_signed_v, "Not implemented for unsigned types"); + stream.write(Leb128MagicString, Leb128MagicStringSize); + std::uint32_t byte_count = 0; + for (std::size_t i = 0; i < count; ++i) { + IntType value = values[i]; + std::uint8_t byte; + do { + byte = value & 0x7f; + value >>= 7; + ++byte_count; + } while ((byte & 0x40) == 0 ? value != 0 : value != -1); + } + write_little_endian(stream, byte_count); + const std::uint32_t BUF_SIZE = 4096; + std::uint8_t buf[BUF_SIZE]; + std::uint32_t buf_pos = 0; + auto flush = [&]() { + if (buf_pos > 0) { + stream.write(reinterpret_cast(buf), buf_pos); + buf_pos = 0; + } + }; + auto write = [&](std::uint8_t byte) { + buf[buf_pos++] = byte; + if (buf_pos == BUF_SIZE) flush(); + }; + for (std::size_t i = 0; i < count; ++i) { + IntType value = values[i]; + while (true) { + std::uint8_t byte = value & 0x7f; + value >>= 7; + if ((byte & 0x40) == 0 ? value == 0 : value == -1) { + write(byte); + break; + } + write(byte | 0x80); + } + } + flush(); + } + } // namespace Stockfish::Eval::NNUE #endif // #ifndef NNUE_COMMON_H_INCLUDED diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index a1888c7a..7571f398 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -253,9 +253,9 @@ namespace Stockfish::Eval::NNUE { // Read network parameters bool read_parameters(std::istream& stream) { - read_little_endian(stream, biases , HalfDimensions ); - read_little_endian(stream, weights , HalfDimensions * InputDimensions); - read_little_endian(stream, psqtWeights, PSQTBuckets * InputDimensions); + read_leb_128(stream, biases , HalfDimensions ); + read_leb_128(stream, weights , HalfDimensions * InputDimensions); + read_leb_128(stream, psqtWeights, PSQTBuckets * InputDimensions); return !stream.fail(); } @@ -263,9 +263,9 @@ namespace Stockfish::Eval::NNUE { // Write network parameters bool write_parameters(std::ostream& stream) const { - write_little_endian(stream, biases , HalfDimensions ); - write_little_endian(stream, weights , HalfDimensions * InputDimensions); - write_little_endian(stream, psqtWeights, PSQTBuckets * InputDimensions); + write_leb_128(stream, biases , HalfDimensions ); + write_leb_128(stream, weights , HalfDimensions * InputDimensions); + write_leb_128(stream, psqtWeights, PSQTBuckets * InputDimensions); return !stream.fail(); } -- 2.39.2