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