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