Skip to content

Commit a46087e

Browse files
maximvondele
maxim
authored andcommitted
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 #4617 No functional change
1 parent 0187546 commit a46087e

File tree

3 files changed

+84
-7
lines changed

3 files changed

+84
-7
lines changed

src/evaluate.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ namespace Eval {
3939
// The default net name MUST follow the format nn-[SHA256 first 12 digits].nnue
4040
// for the build process (profile-build and fishtest) to work. Do not change the
4141
// name of the macro, as it is used in the Makefile.
42-
#define EvalFileDefaultName "nn-cd2ff4716c34.nnue"
42+
#define EvalFileDefaultName "nn-78bacfcee510.nnue"
4343

4444
namespace NNUE {
4545

src/nnue/nnue_common.h

+77
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ namespace Stockfish::Eval::NNUE {
5757
// Size of cache line (in bytes)
5858
constexpr std::size_t CacheLineSize = 64;
5959

60+
constexpr const char Leb128MagicString[] = "COMPRESSED_LEB128";
61+
constexpr const std::size_t Leb128MagicStringSize = sizeof(Leb128MagicString) - 1;
62+
6063
// SIMD width (in bytes)
6164
#if defined(USE_AVX2)
6265
constexpr std::size_t SimdWidth = 32;
@@ -159,6 +162,80 @@ namespace Stockfish::Eval::NNUE {
159162
write_little_endian<IntType>(stream, values[i]);
160163
}
161164

165+
template <typename IntType>
166+
inline void read_leb_128(std::istream& stream, IntType* out, std::size_t count) {
167+
static_assert(std::is_signed_v<IntType>, "Not implemented for unsigned types");
168+
char leb128MagicString[Leb128MagicStringSize];
169+
stream.read(leb128MagicString, Leb128MagicStringSize);
170+
assert(strncmp(Leb128MagicString, leb128MagicString, Leb128MagicStringSize) == 0);
171+
const std::uint32_t BUF_SIZE = 4096;
172+
std::uint8_t buf[BUF_SIZE];
173+
auto bytes_left = read_little_endian<std::uint32_t>(stream);
174+
std::uint32_t buf_pos = BUF_SIZE;
175+
for (std::size_t i = 0; i < count; ++i) {
176+
IntType result = 0;
177+
size_t shift = 0;
178+
do {
179+
if (buf_pos == BUF_SIZE) {
180+
stream.read(reinterpret_cast<char*>(buf), std::min(bytes_left, BUF_SIZE));
181+
buf_pos = 0;
182+
}
183+
std::uint8_t byte = buf[buf_pos++];
184+
--bytes_left;
185+
result |= (byte & 0x7f) << shift;
186+
shift += 7;
187+
if ((byte & 0x80) == 0) {
188+
out[i] = sizeof(IntType) * 8 <= shift || (byte & 0x40) == 0 ? result : result | ~((1 << shift) - 1);
189+
break;
190+
}
191+
} while (shift < sizeof(IntType) * 8);
192+
}
193+
assert(bytes_left == 0);
194+
}
195+
196+
template <typename IntType>
197+
inline void write_leb_128(std::ostream& stream, const IntType* values, std::size_t count) {
198+
static_assert(std::is_signed_v<IntType>, "Not implemented for unsigned types");
199+
stream.write(Leb128MagicString, Leb128MagicStringSize);
200+
std::uint32_t byte_count = 0;
201+
for (std::size_t i = 0; i < count; ++i) {
202+
IntType value = values[i];
203+
std::uint8_t byte;
204+
do {
205+
byte = value & 0x7f;
206+
value >>= 7;
207+
++byte_count;
208+
} while ((byte & 0x40) == 0 ? value != 0 : value != -1);
209+
}
210+
write_little_endian(stream, byte_count);
211+
const std::uint32_t BUF_SIZE = 4096;
212+
std::uint8_t buf[BUF_SIZE];
213+
std::uint32_t buf_pos = 0;
214+
auto flush = [&]() {
215+
if (buf_pos > 0) {
216+
stream.write(reinterpret_cast<char*>(buf), buf_pos);
217+
buf_pos = 0;
218+
}
219+
};
220+
auto write = [&](std::uint8_t byte) {
221+
buf[buf_pos++] = byte;
222+
if (buf_pos == BUF_SIZE) flush();
223+
};
224+
for (std::size_t i = 0; i < count; ++i) {
225+
IntType value = values[i];
226+
while (true) {
227+
std::uint8_t byte = value & 0x7f;
228+
value >>= 7;
229+
if ((byte & 0x40) == 0 ? value == 0 : value == -1) {
230+
write(byte);
231+
break;
232+
}
233+
write(byte | 0x80);
234+
}
235+
}
236+
flush();
237+
}
238+
162239
} // namespace Stockfish::Eval::NNUE
163240

164241
#endif // #ifndef NNUE_COMMON_H_INCLUDED

src/nnue/nnue_feature_transformer.h

+6-6
Original file line numberDiff line numberDiff line change
@@ -253,19 +253,19 @@ namespace Stockfish::Eval::NNUE {
253253
// Read network parameters
254254
bool read_parameters(std::istream& stream) {
255255

256-
read_little_endian<BiasType >(stream, biases , HalfDimensions );
257-
read_little_endian<WeightType >(stream, weights , HalfDimensions * InputDimensions);
258-
read_little_endian<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);
256+
read_leb_128<BiasType >(stream, biases , HalfDimensions );
257+
read_leb_128<WeightType >(stream, weights , HalfDimensions * InputDimensions);
258+
read_leb_128<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);
259259

260260
return !stream.fail();
261261
}
262262

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

266-
write_little_endian<BiasType >(stream, biases , HalfDimensions );
267-
write_little_endian<WeightType >(stream, weights , HalfDimensions * InputDimensions);
268-
write_little_endian<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);
266+
write_leb_128<BiasType >(stream, biases , HalfDimensions );
267+
write_leb_128<WeightType >(stream, weights , HalfDimensions * InputDimensions);
268+
write_leb_128<PSQTWeightType>(stream, psqtWeights, PSQTBuckets * InputDimensions);
269269

270270
return !stream.fail();
271271
}

0 commit comments

Comments
 (0)