]> git.sesse.net Git - plocate/blob - turbopfor-encode.h
clang-format.
[plocate] / turbopfor-encode.h
1 #ifndef _TURBOPFOR_ENCODE_H
2 #define _TURBOPFOR_ENCODE_H
3
4 // Much like turbopfor.h (and shares all of the same caveats), except this is
5 // for encoding. It is _much_ slower than the reference implementation, but we
6 // encode only during build, and most time in build is spent in other things
7 // than encoding posting lists, so it only costs ~5-10% overall. Does not use
8 // any special character sets, and generally isn't optimized at all.
9 //
10 // It encodes about 0.01% denser than the reference encoder (averaged over
11 // a real plocate corpus), probably since it has a slower but more precise
12 // method for estimating the cost of a PFOR + varbyte block.
13
14 #include "turbopfor-common.h"
15
16 #include <algorithm>
17 #include <assert.h>
18 #include <limits.h>
19 #include <stdint.h>
20 #include <string.h>
21
22 template<class Docid>
23 void write_le(Docid val, void *out)
24 {
25         if constexpr (sizeof(Docid) == 8) {
26                 val = htole64(val);
27         } else if constexpr (sizeof(Docid) == 4) {
28                 val = htole32(val);
29         } else if constexpr (sizeof(Docid) == 2) {
30                 val = htole16(val);
31         } else if constexpr (sizeof(Docid) == 1) {
32                 // No change.
33         } else {
34                 assert(false);
35         }
36         memcpy(out, &val, sizeof(val));
37 }
38
39 // Corresponds to read_baseval.
40 template<class Docid>
41 unsigned char *write_baseval(Docid in, unsigned char *out)
42 {
43         if (in < 128) {
44                 *out = in;
45                 return out + 1;
46         } else if (in < 0x4000) {
47                 out[0] = (in >> 8) | 0x80;
48                 out[1] = in & 0xff;
49                 return out + 2;
50         } else if (in < 0x200000) {
51                 out[0] = (in >> 16) | 0xc0;
52                 out[1] = in & 0xff;
53                 out[2] = (in >> 8) & 0xff;
54                 return out + 3;
55         } else {
56                 assert(false);  // Not implemented.
57         }
58 }
59
60 // Writes a varbyte-encoded exception.
61 template<class Docid>
62 unsigned char *write_vb(Docid val, unsigned char *out)
63 {
64         if (val <= 176) {
65                 *out++ = val;
66                 return out;
67         } else if (val <= 16560) {
68                 val -= 177;
69                 *out++ = (val >> 8) + 177;
70                 *out++ = val & 0xff;
71                 return out;
72         } else if (val <= 540848) {
73                 val -= 16561;
74                 *out = (val >> 16) + 241;
75                 write_le<uint16_t>(val & 0xffff, out + 1);
76                 return out + 3;
77         } else if (val <= 16777215) {
78                 *out = 249;
79                 write_le<uint32_t>(val, out + 1);
80                 return out + 4;
81         } else {
82                 *out = 250;
83                 write_le<uint32_t>(val, out + 1);
84                 return out + 5;
85         }
86 }
87
88 template<class Docid>
89 inline unsigned num_bits(Docid x)
90 {
91         if (x == 0) {
92                 return 0;
93         } else {
94                 return sizeof(Docid) * CHAR_BIT - __builtin_clz(x);
95         }
96 }
97
98 struct BitWriter {
99 public:
100         BitWriter(unsigned char *out, unsigned bits)
101                 : out(out), bits(bits) {}
102         void write(uint32_t val)
103         {
104                 cur_val |= val << bits_used;
105                 write_le<uint32_t>(cur_val, out);
106
107                 bits_used += bits;
108                 cur_val >>= (bits_used / 8) * 8;
109                 out += bits_used / 8;
110                 bits_used %= 8;
111         }
112
113 private:
114         unsigned char *out;
115         const unsigned bits;
116         unsigned bits_used = 0;
117         unsigned cur_val = 0;
118 };
119
120 template<unsigned NumStreams>
121 struct InterleavedBitWriter {
122 public:
123         InterleavedBitWriter(unsigned char *out, unsigned bits)
124                 : out(out), bits(bits) {}
125         void write(uint32_t val)
126         {
127                 cur_val |= uint64_t(val) << bits_used;
128                 if (bits_used + bits >= 32) {
129                         write_le<uint32_t>(cur_val & 0xffffffff, out);
130                         out += Stride;
131                         cur_val >>= 32;
132                         bits_used -= 32;  // Underflow, but will be fixed below.
133                 }
134                 write_le<uint32_t>(cur_val, out);
135                 bits_used += bits;
136         }
137
138 private:
139         static constexpr unsigned Stride = NumStreams * sizeof(uint32_t);
140         unsigned char *out;
141         const unsigned bits;
142         unsigned bits_used = 0;
143         uint64_t cur_val = 0;
144 };
145
146 // Bitpacks a set of values (making sure the top bits are lopped off).
147 // If interleaved is set, makes SSE2-compatible interleaving (this is
148 // only allowed for full blocks).
149 template<class Docid>
150 unsigned char *encode_bitmap(const Docid *in, unsigned num, unsigned bit_width, bool interleaved, unsigned char *out)
151 {
152         unsigned mask = mask_for_bits(bit_width);
153         if (interleaved) {
154                 InterleavedBitWriter<4> bs0(out + 0 * sizeof(uint32_t), bit_width);
155                 InterleavedBitWriter<4> bs1(out + 1 * sizeof(uint32_t), bit_width);
156                 InterleavedBitWriter<4> bs2(out + 2 * sizeof(uint32_t), bit_width);
157                 InterleavedBitWriter<4> bs3(out + 3 * sizeof(uint32_t), bit_width);
158                 assert(num % 4 == 0);
159                 for (unsigned i = 0; i < num / 4; ++i) {
160                         bs0.write(in[i * 4 + 0] & mask);
161                         bs1.write(in[i * 4 + 1] & mask);
162                         bs2.write(in[i * 4 + 2] & mask);
163                         bs3.write(in[i * 4 + 3] & mask);
164                 }
165         } else {
166                 BitWriter bs(out, bit_width);
167                 for (unsigned i = 0; i < num; ++i) {
168                         bs.write(in[i] & mask);
169                 }
170         }
171         return out + bytes_for_packed_bits(num, bit_width);
172 }
173
174 // See decode_for() for the format.
175 template<class Docid>
176 unsigned char *encode_for(const Docid *in, unsigned num, unsigned bit_width, bool interleaved, unsigned char *out)
177 {
178         return encode_bitmap(in, num, bit_width, interleaved, out);
179 }
180
181 // See decode_pfor_bitmap() for the format.
182 template<class Docid>
183 unsigned char *encode_pfor_bitmap(const Docid *in, unsigned num, unsigned bit_width, unsigned exception_bit_width, bool interleaved, unsigned char *out)
184 {
185         *out++ = exception_bit_width;
186
187         // Bitmap of exceptions.
188         {
189                 BitWriter bs(out, 1);
190                 for (unsigned i = 0; i < num; ++i) {
191                         bs.write((in[i] >> bit_width) != 0);
192                 }
193                 out += bytes_for_packed_bits(num, 1);
194         }
195
196         // Exceptions.
197         {
198                 BitWriter bs(out, exception_bit_width);
199                 unsigned num_exceptions = 0;
200                 for (unsigned i = 0; i < num; ++i) {
201                         if ((in[i] >> bit_width) != 0) {
202                                 bs.write(in[i] >> bit_width);
203                                 ++num_exceptions;
204                         }
205                 }
206                 out += bytes_for_packed_bits(num_exceptions, exception_bit_width);
207         }
208
209         // Base values.
210         out = encode_bitmap(in, num, bit_width, interleaved, out);
211
212         return out;
213 }
214
215 // See decode_pfor_vb() for the format.
216 template<class Docid>
217 unsigned char *encode_pfor_vb(const Docid *in, unsigned num, unsigned bit_width, bool interleaved, unsigned char *out)
218 {
219         unsigned num_exceptions = 0;
220         for (unsigned i = 0; i < num; ++i) {
221                 if ((in[i] >> bit_width) != 0) {
222                         ++num_exceptions;
223                 }
224         }
225         *out++ = num_exceptions;
226
227         // Base values.
228         out = encode_bitmap(in, num, bit_width, interleaved, out);
229
230         // Exceptions.
231         for (unsigned i = 0; i < num; ++i) {
232                 unsigned val = in[i] >> bit_width;
233                 if (val != 0) {
234                         out = write_vb(val, out);
235                 }
236         }
237
238         // Exception indexes.
239         for (unsigned i = 0; i < num; ++i) {
240                 unsigned val = in[i] >> bit_width;
241                 if (val != 0) {
242                         *out++ = i;
243                 }
244         }
245
246         return out;
247 }
248
249 // Find out which block type would be the smallest for the given data.
250 template<class Docid>
251 BlockType decide_block_type(const Docid *in, unsigned num, unsigned *bit_width, unsigned *exception_bit_width)
252 {
253         // Check if the block is constant.
254         bool constant = true;
255         for (unsigned i = 1; i < num; ++i) {
256                 if (in[i] != in[0]) {
257                         constant = false;
258                         break;
259                 }
260         }
261         if (constant) {
262                 *bit_width = num_bits(in[0]);
263                 return BlockType::CONSTANT;
264         }
265
266         // Build up a histogram of bit sizes.
267         unsigned histogram[sizeof(Docid) * CHAR_BIT + 1] = { 0 };
268         unsigned max_bits = 0;
269         for (unsigned i = 0; i < num; ++i) {
270                 unsigned bits = num_bits(in[i]);
271                 ++histogram[bits];
272                 max_bits = std::max(max_bits, bits);
273         }
274
275         // Straight-up FOR.
276         unsigned best_cost = bytes_for_packed_bits(num, max_bits);
277         unsigned best_bit_width = max_bits;
278
279         // Try PFOR with bitmap exceptions.
280         const unsigned bitmap_cost = bytes_for_packed_bits(num, 1);
281         unsigned num_exceptions = 0;
282         for (unsigned exception_bit_width = 1; exception_bit_width <= max_bits; ++exception_bit_width) {
283                 unsigned test_bit_width = max_bits - exception_bit_width;
284                 num_exceptions += histogram[test_bit_width + 1];
285
286                 // 1 byte for signaling exception bit width, then the bitmap,
287                 // then the base values, then the exceptions.
288                 unsigned cost = 1 + bitmap_cost + bytes_for_packed_bits(num, test_bit_width) +
289                         bytes_for_packed_bits(num_exceptions, exception_bit_width);
290                 if (cost < best_cost) {
291                         best_cost = cost;
292                         best_bit_width = test_bit_width;
293                 }
294         }
295
296         // Try PFOR with varbyte exceptions.
297         bool best_is_varbyte = false;
298         for (unsigned test_bit_width = 0; test_bit_width < max_bits; ++test_bit_width) {
299                 // 1 byte for signaling number of exceptions, plus the base values,
300                 // and then we count up the varbytes and indexes. (This is precise
301                 // but very slow.)
302                 unsigned cost = 1 + bytes_for_packed_bits(num, test_bit_width);
303                 for (unsigned i = 0; i < num && cost < best_cost; ++i) {
304                         unsigned val = in[i] >> test_bit_width;
305                         if (val == 0) {
306                                 // Not stored, and then also no index.
307                         } else if (val <= 176) {
308                                 cost += 2;
309                         } else if (val <= 16560) {
310                                 cost += 3;
311                         } else if (val <= 540848) {
312                                 cost += 4;
313                         } else if (val <= 16777215) {
314                                 cost += 5;
315                         } else {
316                                 cost += 6;
317                         }
318                 }
319                 if (cost < best_cost) {
320                         best_cost = cost;
321                         best_bit_width = test_bit_width;
322                         best_is_varbyte = true;
323                 }
324         }
325
326         // TODO: Consider the last-resort option of just raw storage (255).
327
328         if (best_is_varbyte) {
329                 *bit_width = best_bit_width;
330                 return BlockType::PFOR_VB;
331         } else if (best_bit_width == max_bits) {
332                 *bit_width = max_bits;
333                 return BlockType::FOR;
334         } else {
335                 *bit_width = best_bit_width;
336                 *exception_bit_width = max_bits - best_bit_width;
337                 return BlockType::PFOR_BITMAP;
338         }
339 }
340
341 // The basic entry point. Takes one block of integers (which already must
342 // be delta-minus-1-encoded) and packs it into TurboPFor format.
343 // interleaved corresponds to the interleaved parameter in decode_pfor_delta1()
344 // or the ā€œ128vā€ infix in the reference code's function names; such formats
345 // are much faster to decode, so for full blocks, you probably want it.
346 // The interleaved flag isn't stored anywhere; it's implicit whether you
347 // want to use it for full blocks or not.
348 //
349 // The first value must already be written using write_baseval() (so the delta
350 // coding starts from the second value). Returns the end of the string.
351 // May write 4 bytes past the end.
352 template<unsigned BlockSize, class Docid>
353 unsigned char *encode_pfor_single_block(const Docid *in, unsigned num, bool interleaved, unsigned char *out)
354 {
355         assert(num > 0);
356         if (interleaved) {
357                 assert(num == BlockSize);
358         }
359
360         unsigned bit_width, exception_bit_width;
361         BlockType block_type = decide_block_type(in, num, &bit_width, &exception_bit_width);
362         *out++ = (block_type << 6) | bit_width;
363
364         switch (block_type) {
365         case BlockType::CONSTANT: {
366                 unsigned bit_width = num_bits(in[0]);
367                 write_le<Docid>(in[0], out);
368                 return out + div_round_up(bit_width, 8);
369         }
370         case BlockType::FOR:
371                 return encode_for(in, num, bit_width, interleaved, out);
372         case BlockType::PFOR_BITMAP:
373                 return encode_pfor_bitmap(in, num, bit_width, exception_bit_width, interleaved, out);
374         case BlockType::PFOR_VB:
375                 return encode_pfor_vb(in, num, bit_width, interleaved, out);
376         default:
377                 assert(false);
378         }
379 }
380
381 #endif  // !defined(_TURBOPFOR_ENCODE_H)