git.sesse.net Git - plocate/blob - turbopfor.h

   1 #ifndef _TURBOPFOR_H
   2 #define _TURBOPFOR_H 1
   3
   4 // A reimplementation of parts of the TurboPFor codecs, using the same
   5 // storage format. These are not as fast as the reference implementation
   6 // (about 1/3 of the performance), and do not support the same breadth of
   7 // codecs (in particular, only delta-plus-1 is implemented, and only 32-bit
   8 // docids are tested), but aim to be more portable and easier-to-understand.
   9 // In particular, they will compile on x86 without SSE4.1 or AVX support.
  10 //
  11 // The main reference is https://michael.stapelberg.ch/posts/2019-02-05-turbopfor-analysis/,
  12 // although some implementation details have been worked out by studying the
  13 // TurboPFor code.
  14
  15 #include <assert.h>
  16 #include <endian.h>
  17 #include <stdint.h>
  18 #include <string.h>
  19 #include <limits.h>
  20
  21 #include <algorithm>
  22
  23 template<class Docid>
  24 Docid read_le(const void *in)
  25 {
  26         Docid val;
  27         memcpy(&val, in, sizeof(val));
  28         if constexpr (sizeof(Docid) == 8) {
  29                 return le64toh(val);
  30         } else if constexpr (sizeof(Docid) == 4) {
  31                 return le32toh(val);
  32         } else if constexpr (sizeof(Docid) == 2) {
  33                 return le16toh(val);
  34         } else if constexpr (sizeof(Docid) == 1) {
  35                 return val;
  36         } else {
  37                 assert(false);
  38         }
  39 }
  40
  41 // Reads a single value with an encoding that looks a bit like PrefixVarint.
  42 // It's unclear why this doesn't use the varbyte encoding.
  43 template<class Docid>
  44 const unsigned char *read_baseval(const unsigned char *in, Docid *out)
  45 {
  46         //fprintf(stderr, "baseval: 0x%02x 0x%02x 0x%02x 0x%02x\n", in[0], in[1], in[2], in[3]);
  47         if (*in < 128) {
  48                 *out = *in;
  49                 return in + 1;
  50         } else if (*in < 192) {
  51                 *out = ((uint32_t(in[0]) << 8) | uint32_t(in[1])) & 0x3fff;
  52                 return in + 2;
  53         } else if (*in < 224) {
  54                 *out = ((uint32_t(in[0]) << 16) |
  55                         (uint32_t(in[2]) << 8) |
  56                         (uint32_t(in[1]))) & 0x1fffff;
  57                 return in + 3;
  58         } else {
  59                 assert(false);  // Not implemented.
  60         }
  61 }
  62
  63 template<class Docid>
  64 const unsigned char *read_vb(const unsigned char *in, Docid *out)
  65 {
  66         if (*in <= 176) {
  67                 *out = *in;
  68                 return in + 1;
  69         } else if (*in <= 240) {
  70                 *out = ((uint32_t(in[0] - 177) << 8) | uint32_t(in[1])) + 177;
  71                 return in + 2;
  72         } else if (*in <= 248) {
  73                 *out = ((uint32_t(in[0] - 241) << 16) | read_le<uint16_t>(in + 1)) + 16561;
  74                 return in + 3;
  75         } else if (*in == 249) {
  76                 *out = (uint32_t(in[1])) |
  77                         (uint32_t(in[2]) << 8) |
  78                         (uint32_t(in[3]) << 16);
  79                 return in + 4;
  80         } else if (*in == 250) {
  81                 *out = read_le<uint32_t>(in + 1);
  82                 return in + 5;
  83         } else {
  84                 assert(false);
  85         }
  86 }
  87
  88 struct BitReader {
  89 public:
  90         BitReader(const unsigned char *in, unsigned bits)
  91                 : in(in), bits(bits), mask((1U << bits) - 1) {}
  92         uint32_t read()
  93         {
  94                 uint32_t val = (read_le<uint32_t>(in) >> bits_used) & mask;
  95
  96                 bits_used += bits;
  97                 in += bits_used / 8;
  98                 bits_used %= 8;
  99
 100                 return val;
 101         }
 102
 103 private:
 104         const unsigned char *in;
 105         const unsigned bits;
 106         const unsigned mask;
 107         unsigned bits_used = 0;
 108 };
 109
 110 template<unsigned NumStreams>
 111 struct InterleavedBitReader {
 112 public:
 113         InterleavedBitReader(const unsigned char *in, unsigned bits)
 114                 : in(in), bits(bits), mask((1U << bits) - 1) {}
 115         uint32_t read()
 116         {
 117                 uint32_t val;
 118                 if (bits_used + bits > 32) {
 119                         val = (read_le<uint32_t>(in) >> bits_used) | (read_le<uint32_t>(in + Stride) << (32 - bits_used));
 120                 } else {
 121                         val = (read_le<uint32_t>(in) >> bits_used);
 122                 }
 123
 124                 bits_used += bits;
 125                 in += Stride * (bits_used / 32);
 126                 bits_used %= 32;
 127
 128                 return val & mask;
 129         }
 130
 131 private:
 132         static constexpr unsigned Stride = NumStreams * sizeof(uint32_t);
 133         const unsigned char *in;
 134         const unsigned bits;
 135         const unsigned mask;
 136         unsigned bits_used = 0;
 137 };
 138
 139 // Does not properly account for overflow.
 140 inline unsigned div_round_up(unsigned val, unsigned div)
 141 {
 142         return (val + div - 1) / div;
 143 }
 144
 145 inline unsigned bytes_for_packed_bits(unsigned num, unsigned bit_width)
 146 {
 147         return div_round_up(num * bit_width, CHAR_BIT);
 148 }
 149
 150 // Constant block. Layout:
 151 //
 152 //  - Bit width (6 bits) | type << 6
 153 //  - Base values (<bits> bits, rounded up to nearest byte)
 154 template<class Docid>
 155 const unsigned char *decode_constant(const unsigned char *in, unsigned num, Docid *out)
 156 {
 157         const unsigned bit_width = *in++ & 0x3f;
 158         Docid val = read_le<Docid>(in);
 159         if (bit_width < sizeof(Docid) * 8) {
 160                 val &= ((1U << bit_width) - 1);
 161         }
 162
 163         Docid *prev_out = out - 1;
 164         for (unsigned i = 0; i < num; ++i) {
 165                 out[i] = val + prev_out[i] + 1;
 166         }
 167         return in + div_round_up(bit_width, 8);
 168 }
 169
 170 // FOR block (ie., PFor without exceptions). Layout:
 171 //
 172 //  - Bit width (6 bits) | type << 6
 173 //  - Base values (<num> values of <bits> bits, rounded up to a multiple of 32 values)
 174 template<class Docid>
 175 const unsigned char *decode_for(const unsigned char *in, unsigned num, Docid *out)
 176 {
 177         const unsigned bit_width = *in++ & 0x3f;
 178
 179         Docid *prev_out = out - 1;
 180         BitReader bs(in, bit_width);
 181         for (unsigned i = 0; i < num; ++i) {
 182                 out[i] = bs.read() + prev_out[i] + 1;
 183         }
 184         return in + bytes_for_packed_bits(num, bit_width);
 185 }
 186
 187 // Like decode_for(), but the values are organized in four independent streams,
 188 // for SIMD (presumably SSE2). Supports a whole block only.
 189 template<unsigned BlockSize, class Docid>
 190 const unsigned char *decode_for_interleaved(const unsigned char *in, Docid *out)
 191 {
 192         const unsigned bit_width = *in++ & 0x3f;
 193
 194         InterleavedBitReader<4> bs0(in + 0 * sizeof(uint32_t), bit_width);
 195         InterleavedBitReader<4> bs1(in + 1 * sizeof(uint32_t), bit_width);
 196         InterleavedBitReader<4> bs2(in + 2 * sizeof(uint32_t), bit_width);
 197         InterleavedBitReader<4> bs3(in + 3 * sizeof(uint32_t), bit_width);
 198         for (unsigned i = 0; i < BlockSize / 4; ++i) {
 199                 out[i * 4 + 0] = bs0.read();
 200                 out[i * 4 + 1] = bs1.read();
 201                 out[i * 4 + 2] = bs2.read();
 202                 out[i * 4 + 3] = bs3.read();
 203         }
 204         Docid *prev_out = out - 1;
 205         for (unsigned i = 0; i < BlockSize; ++i) {
 206                 out[i] += prev_out[i] + 1;
 207         }
 208         return in + bytes_for_packed_bits(BlockSize, bit_width);
 209 }
 210
 211 // PFor block with bitmap exceptions. Layout:
 212 //
 213 //  - Bit width (6 bits) | type << 6
 214 //  - Exception bit width (8 bits)
 215 //  - Bitmap of which values have exceptions (<num> bits, rounded up to a byte)
 216 //  - Exceptions (<num_exc> values of <bits_exc> bits, rounded up to a byte)
 217 //  - Base values (<num> values of <bits> bits, rounded up to a byte)
 218 template<class Docid>
 219 const unsigned char *decode_pfor_bitmap(const unsigned char *in, unsigned num, Docid *out)
 220 {
 221         memset(out, 0, num * sizeof(Docid));
 222
 223         const unsigned bit_width = *in++ & 0x3f;
 224         const unsigned exception_bit_width = *in++;
 225
 226         // Decode exceptions.
 227         {
 228                 const uint64_t *exception_bitmap_ptr = reinterpret_cast<const uint64_t *>(in);
 229                 in += div_round_up(num, 8);
 230
 231                 int num_exceptions = 0;
 232
 233                 BitReader bs(in, exception_bit_width);
 234                 for (unsigned i = 0; i < num; i += 64, ++exception_bitmap_ptr) {
 235                         uint64_t exceptions = read_le<uint64_t>(exception_bitmap_ptr);
 236                         if (num - i < 64) {
 237                                 // We've read some bytes past the end, so clear out the junk bits.
 238                                 exceptions &= (1ULL << (num - i)) - 1;
 239                         }
 240                         for (; exceptions != 0; exceptions &= exceptions - 1, ++num_exceptions) {
 241                                 unsigned idx = (ffsll(exceptions) - 1) + i;
 242                                 out[idx] = bs.read() << bit_width;
 243                         }
 244                 }
 245                 in += bytes_for_packed_bits(num_exceptions, exception_bit_width);
 246         }
 247
 248         // Decode the base values, and delta-decode.
 249         Docid *prev_out = out - 1;
 250         BitReader bs(in, bit_width);
 251         for (unsigned i = 0; i < num; ++i) {
 252                 out[i] = (out[i] | bs.read()) + prev_out[i] + 1;
 253         }
 254         return in + bytes_for_packed_bits(num, bit_width);
 255 }
 256
 257 // Like decode_pfor_bitmap(), but the base values are organized in four
 258 // independent streams, for SIMD (presumably SSE2). Supports a whole block only.
 259 template<unsigned BlockSize, class Docid>
 260 const unsigned char *decode_pfor_bitmap_interleaved(const unsigned char *in, Docid *out)
 261 {
 262         memset(out, 0, BlockSize * sizeof(Docid));
 263
 264         const unsigned bit_width = *in++ & 0x3f;
 265         const unsigned exception_bit_width = *in++;
 266
 267         // Decode exceptions.
 268         {
 269                 const uint64_t *exception_bitmap_ptr = reinterpret_cast<const uint64_t *>(in);
 270                 in += div_round_up(BlockSize, 8);
 271
 272                 int num_exceptions = 0;
 273
 274                 BitReader bs(in, exception_bit_width);
 275                 for (unsigned i = 0; i < BlockSize; i += 64, ++exception_bitmap_ptr) {
 276                         uint64_t exceptions = read_le<uint64_t>(exception_bitmap_ptr);
 277                         for (; exceptions != 0; exceptions &= exceptions - 1, ++num_exceptions) {
 278                                 unsigned idx = (ffsll(exceptions) - 1) + i;
 279                                 out[idx] = bs.read() << bit_width;
 280                         }
 281                 }
 282                 in += bytes_for_packed_bits(num_exceptions, exception_bit_width);
 283         }
 284
 285         // Decode the base values, and delta-decode.
 286         InterleavedBitReader<4> bs0(in + 0 * sizeof(uint32_t), bit_width);
 287         InterleavedBitReader<4> bs1(in + 1 * sizeof(uint32_t), bit_width);
 288         InterleavedBitReader<4> bs2(in + 2 * sizeof(uint32_t), bit_width);
 289         InterleavedBitReader<4> bs3(in + 3 * sizeof(uint32_t), bit_width);
 290         for (unsigned i = 0; i < BlockSize / 4; ++i) {
 291                 out[i * 4 + 0] |= bs0.read();
 292                 out[i * 4 + 1] |= bs1.read();
 293                 out[i * 4 + 2] |= bs2.read();
 294                 out[i * 4 + 3] |= bs3.read();
 295         }
 296         Docid *prev_out = out - 1;
 297         for (unsigned i = 0; i < BlockSize; ++i) {
 298                 out[i] += prev_out[i] + 1;
 299         }
 300         return in + bytes_for_packed_bits(BlockSize, bit_width);
 301 }
 302
 303 // PFor block with variable-byte exceptions. Layout:
 304 //
 305 //  - Bit width (6 bits) | type << 6
 306 //  - Number of exceptions (8 bits)
 307 //  - Base values (<num> values of <bits> bits, rounded up to a byte)
 308 //  - Exceptions:
 309 //    - If first byte is 255, <num_exc> 32-bit values (does not include the 255 byte)
 310 //    - Else, <num_exc> varbyte-encoded values (includes the non-255 byte)
 311 //  - Indexes of exceptions (<num_exc> bytes).
 312 template<unsigned BlockSize, class Docid>
 313 const unsigned char *decode_pfor_vb(const unsigned char *in, unsigned num, Docid *out)
 314 {
 315         //fprintf(stderr, "in=%p out=%p num=%u\n", in, out, num);
 316
 317         const unsigned bit_width = *in++ & 0x3f;
 318         unsigned num_exceptions = *in++;
 319
 320         // Decode the base values.
 321         BitReader bs(in, bit_width);
 322         for (unsigned i = 0; i < num; ++i) {
 323                 out[i] = bs.read();
 324         }
 325         in += bytes_for_packed_bits(num, bit_width);
 326
 327         // Decode exceptions.
 328         Docid exceptions[BlockSize];
 329         if (*in == 255) {
 330                 ++in;
 331                 for (unsigned i = 0; i < num_exceptions; ++i) {
 332                         exceptions[i] = read_le<Docid>(in);
 333                         in += sizeof(Docid);
 334                 }
 335         } else {
 336                 for (unsigned i = 0; i < num_exceptions; ++i) {
 337                         in = read_vb(in, &exceptions[i]);
 338                 }
 339         }
 340         // Apply exceptions.
 341         for (unsigned i = 0; i < num_exceptions; ++i) {
 342                 unsigned idx = *in++;
 343                 out[idx] |= exceptions[i] << bit_width;
 344         }
 345
 346         // Delta-decode.
 347         Docid *prev_out = out - 1;
 348         for (unsigned i = 0; i < num; ++i) {
 349                 out[i] = out[i] + prev_out[i] + 1;
 350         }
 351
 352         return in;
 353 }
 354
 355 // Like decode_pfor_vb(), but the base values are organized in four
 356 // independent streams, for SIMD (presumably SSE2). Supports a whole block only.
 357 template<unsigned BlockSize, class Docid>
 358 const unsigned char *decode_pfor_vb_interleaved(const unsigned char *in, Docid *out)
 359 {
 360         const unsigned bit_width = *in++ & 0x3f;
 361         unsigned num_exceptions = *in++;
 362
 363         // Decode the base values.
 364         InterleavedBitReader<4> bs0(in + 0 * sizeof(uint32_t), bit_width);
 365         InterleavedBitReader<4> bs1(in + 1 * sizeof(uint32_t), bit_width);
 366         InterleavedBitReader<4> bs2(in + 2 * sizeof(uint32_t), bit_width);
 367         InterleavedBitReader<4> bs3(in + 3 * sizeof(uint32_t), bit_width);
 368         for (unsigned i = 0; i < BlockSize / 4; ++i) {
 369                 out[i * 4 + 0] = bs0.read();
 370                 out[i * 4 + 1] = bs1.read();
 371                 out[i * 4 + 2] = bs2.read();
 372                 out[i * 4 + 3] = bs3.read();
 373         }
 374         in += bytes_for_packed_bits(BlockSize, bit_width);
 375
 376         // Decode exceptions.
 377         Docid exceptions[BlockSize];
 378         if (*in == 255) {
 379                 ++in;
 380                 for (unsigned i = 0; i < num_exceptions; ++i) {
 381                         exceptions[i] = read_le<Docid>(in);
 382                         in += sizeof(Docid);
 383                 }
 384         } else {
 385                 for (unsigned i = 0; i < num_exceptions; ++i) {
 386                         in = read_vb(in, &exceptions[i]);
 387                 }
 388         }
 389
 390         // Apply exceptions.
 391         for (unsigned i = 0; i < num_exceptions; ++i) {
 392                 unsigned idx = *in++;
 393                 out[idx] |= exceptions[i] << bit_width;
 394         }
 395
 396         // Delta-decode.
 397         Docid *prev_out = out - 1;
 398         for (unsigned i = 0; i < BlockSize; ++i) {
 399                 out[i] = out[i] + prev_out[i] + 1;
 400         }
 401
 402         return in;
 403 }
 404
 405 enum BlockType {
 406         FOR = 0,
 407         PFOR_VB = 1,
 408         PFOR_BITMAP = 2,
 409         CONSTANT = 3
 410 };
 411
 412 template<unsigned BlockSize, class Docid>
 413 const unsigned char *decode_pfor_delta1(const unsigned char *in, unsigned num, bool interleaved, Docid *out)
 414 {
 415         if (num == 0) {
 416                 return in;
 417         }
 418         in = read_baseval(in, out++);
 419
 420         for (unsigned i = 1; i < num; i += BlockSize, out += BlockSize) {
 421                 const unsigned num_this_block = std::min<unsigned>(num - i, BlockSize);
 422                 switch (in[0] >> 6) {
 423                 case BlockType::FOR:
 424                         if (interleaved && num_this_block == BlockSize) {
 425                                 dprintf("%d+%d: blocktype=%d (for, interleaved), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 426                                 in = decode_for_interleaved<BlockSize>(in, out);
 427                         } else {
 428                                 dprintf("%d+%d: blocktype=%d (for), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 429                                 in = decode_for(in, num_this_block, out);
 430                         }
 431                         break;
 432                 case BlockType::PFOR_VB:
 433                         if (interleaved && num_this_block == BlockSize) {
 434                                 dprintf("%d+%d: blocktype=%d (pfor + vb, interleaved), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 435                                 in = decode_pfor_vb_interleaved<BlockSize>(in, out);
 436                         } else {
 437                                 dprintf("%d+%d: blocktype=%d (pfor + vb), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 438                                 in = decode_pfor_vb<BlockSize>(in, num_this_block, out);
 439                         }
 440                         break;
 441                 case BlockType::PFOR_BITMAP:
 442                         if (interleaved && num_this_block == BlockSize) {
 443                                 dprintf("%d+%d: blocktype=%d (pfor + bitmap, interleaved), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 444                                 in = decode_pfor_bitmap_interleaved<BlockSize>(in, out);
 445                         } else {
 446                                 dprintf("%d+%d: blocktype=%d (pfor + bitmap), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 447                                 in = decode_pfor_bitmap(in, num_this_block, out);
 448                         }
 449                         break;
 450                 case BlockType::CONSTANT:
 451                         dprintf("%d+%d: blocktype=%d (constant), bitwidth=%d\n", i, num_this_block, in[0] >> 6, in[0] & 0x3f);
 452                         in = decode_constant(in, num_this_block, out);
 453                         break;
 454                 }
 455         }
 456
 457         return in;
 458 }
 459
 460 #endif  // !defined(_TURBOPFOR_H)