Skip to content

Commit

Permalink
Compressed network parameters
Browse files Browse the repository at this point in the history
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 #4617

No functional change
  • Loading branch information
maxim authored and vondele committed Jun 19, 2023
1 parent 0187546 commit a46087e
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 7 deletions.
2 changes: 1 addition & 1 deletion src/evaluate.h
Expand Up @@ -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 {

Expand Down
77 changes: 77 additions & 0 deletions src/nnue/nnue_common.h
Expand Up @@ -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;
Expand Down Expand Up @@ -159,6 +162,80 @@ namespace Stockfish::Eval::NNUE {
write_little_endian<IntType>(stream, values[i]);
}

template <typename IntType>
inline void read_leb_128(std::istream& stream, IntType* out, std::size_t count) {
static_assert(std::is_signed_v<IntType>, "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<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 {
if (buf_pos == BUF_SIZE) {
stream.read(reinterpret_cast<char*>(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 <typename IntType>
inline void write_leb_128(std::ostream& stream, const IntType* values, std::size_t count) {
static_assert(std::is_signed_v<IntType>, "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<char*>(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
12 changes: 6 additions & 6 deletions src/nnue/nnue_feature_transformer.h
Expand Up @@ -253,19 +253,19 @@ namespace Stockfish::Eval::NNUE {
// Read network parameters
bool read_parameters(std::istream& stream) {

read_little_endian<BiasType >(stream, biases , HalfDimensions );
read_little_endian<WeightType >(stream, weights , HalfDimensions * InputDimensions);
read_little_endian<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);
read_leb_128<BiasType >(stream, biases , HalfDimensions );
read_leb_128<WeightType >(stream, weights , HalfDimensions * InputDimensions);
read_leb_128<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);

return !stream.fail();
}

// Write network parameters
bool write_parameters(std::ostream& stream) const {

write_little_endian<BiasType >(stream, biases , HalfDimensions );
write_little_endian<WeightType >(stream, weights , HalfDimensions * InputDimensions);
write_little_endian<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);
write_leb_128<BiasType >(stream, biases , HalfDimensions );
write_leb_128<WeightType >(stream, weights , HalfDimensions * InputDimensions);
write_leb_128<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);

return !stream.fail();
}
Expand Down

0 comments on commit a46087e

Please sign in to comment.