faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h

@@ -0,0 +1,300 @@
+// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+#pragma once
+
+// This file contains a custom fast implementation of faiss::Index::sa_decode()
+//   function for the following index families:
+//   * IVF256,PQ[1]x8np
+//   * Residual[1]x8,PQ[2]x8
+//   * IVF[2^9-2^16 bit],PQ[1]x8 (such as IVF1024,PQ16np)
+//   * Residual1x[9-16 bit],PQ[1]x8 (such as Residual1x9,PQ8)
+//   * PQ[1]x8
+// Additionally, AVX2 and ARM versions support
+//   * Residual[1]x8,PQ[2]x10
+//   * Residual[1]x8,PQ[2]x16
+//   * Residual[1]x10,PQ[2]x10
+//   * Residual[1]x10,PQ[2]x16
+//   * Residual[1]x16,PQ[2]x10
+//   * Residual[1]x16,PQ[2]x16
+//   * Residual1x[9-16 bit],PQ[1]x10 (such as Residual1x9,PQ16x10)
+//   * * (use with COARSE_BITS=16)
+//   * Residual1x[9-16 bit],PQ[1]x16 (such as Residual1x9,PQ16x16)
+//   * * (use with COARSE_BITS=16)
+//   * PQ[1]x10
+//   * PQ[1]x16
+// Unfortunately, currently Faiss does not support something like
+//   IVF256,PQ16x10np
+//
+// The goal was to achieve the maximum performance, so the template version it
+// is. The provided index families share the same code for sa_decode.
+//
+// The front-end code provides two high-level structures.
+//
+// First one:
+//   {
+//     template <
+//        intptr_t DIM,
+//        intptr_t COARSE_SIZE,
+//        intptr_t FINE_SIZE,
+//        intptr_t COARSE_BITS = 8
+//        intptr_t FINE_BITS = 8>
+//     struct Index2LevelDecoder { /*...*/ };
+//   }
+// * DIM is the dimensionality of data
+// * COARSE_SIZE is the dimensionality of the coarse quantizer (IVF, Residual)
+// * FINE_SIZE is the dimensionality of the ProductQuantizer dsq
+// * COARSE_BITS is the number of bits that are needed to represent a coarse
+//   quantizer code.
+// * FINE_BITS is the number of bits that are needed to represent a fine
+//   quantizer code.
+// For example, "IVF256,PQ8np" for 160-dim data translates into
+//   Index2LevelDecoder<160,160,20,8>
+// For example, "Residual4x8,PQ16" for 256-dim data translates into
+//   Index2LevelDecoder<256,64,1,8>
+// For example, "IVF1024,PQ16np" for 256-dim data translates into
+//   Index2LevelDecoder<256,256,16,10>. But as there are only 1 coarse code
+//   element, Index2LevelDecoder<256,256,16,16> can be used as a faster
+//   decoder.
+// For example, "Residual4x10,PQ16x10np" for 256-dim data translates into
+//   Index2LevelDecoder<256,64,16,10,10>
+//
+// Additional supported values for COARSE_BITS and FINE_BITS may be added later.
+//
+// Second one:
+//   {
+//     template <
+//        intptr_t DIM,
+//        intptr_t FINE_SIZE,
+//        intptr_t FINE_BITS = 8>
+//     struct IndexPQDecoder { /*...*/ };
+//   }
+// * DIM is the dimensionality of data
+// * FINE_SIZE is the dimensionality of the ProductQuantizer dsq
+// * FINE_BITS is the number of bits that are needed to represent a fine
+//   quantizer code.
+// For example, "PQ8np" for 160-dim data translates into
+//   IndexPQDecoder<160,20>
+//
+// Unlike the general purpose version in faiss::Index::sa_decode(),
+//   this version provides the following functions (please note that
+//   pqCoarseCentroids params are not available for IndexPQDecoder,
+//   but the functionality is the same as for Index2LevelDecoder):
+//
+// * ::store(), which is similar to sa_decode(1, input, output),
+//   The method signature is the following:
+//   {
+//     void store(
+//       const float* const __restrict pqCoarseCentroids,
+//       const float* const __restrict pqFineCentroids,
+//       const uint8_t* const __restrict code,
+//       float* const __restrict outputStore);
+//   }
+//
+// * ::accum(), which is used to create a linear combination
+//   of decoded vectors:
+//   {
+//     const faiss::Index* const index;
+//     const uint8_t* const input;
+//     float weight;
+//
+//     std::vector<float> buffer(d, 0);
+//
+//     index->sa_decode(1, input, buffer.data());
+//     for (size_t iDim = 0; iDim < d; iDim++)
+//       output[iDim] += weight * buffer[iDim];
+//   }
+//   The method signature is the following:
+//   {
+//    static void accum(
+//      const float* const __restrict pqCoarseCentroids,
+//      const float* const __restrict pqFineCentroids,
+//      const uint8_t* const __restrict code,
+//      const float weight,
+//      float* const __restrict outputAccum);
+//   }
+//
+// * There is an additional overload for ::accum() that decodes two vectors
+//   per call. This provides an additional speedup because of a CPU
+//   superscalar architecture:
+//   {
+//     const faiss::Index* const index;
+//     const uint8_t* const input0;
+//     float weight0;
+//     const uint8_t* const input1;
+//     float weight1;
+//
+//     std::vector<float> buffer(d, 0);
+//
+//     index->sa_decode(1, input0, buffer.data());
+//     for (size_t iDim = 0; iDim < d; iDim++)
+//       output[iDim] += weight0 * buffer[iDim];
+//
+//     index->sa_decode(1, input1, buffer.data());
+//     for (size_t iDim = 0; iDim < d; iDim++)
+//       output[iDim] += weight1 * buffer[iDim];
+//   }
+//   If each code uses its own coarse quantizer centroids table and its own fine
+//   quantizer centroids table, then the following overload can be used:
+//   {
+//    static void accum(
+//      const float* const __restrict pqCoarseCentroids0,
+//      const float* const __restrict pqFineCentroids0,
+//      const uint8_t* const __restrict code0,
+//      const float weight0,
+//      const float* const __restrict pqCoarseCentroids1,
+//      const float* const __restrict pqFineCentroids1,
+//      const uint8_t* const __restrict code1,
+//      const float weight1,
+//      float* const __restrict outputAccum);
+//   }
+//   If codes share the coarse quantizer centroids table and also share
+//   the fine quantizer centroids table, then the following overload can be
+//   used:
+//   {
+//    static void accum(
+//      const float* const __restrict pqCoarseCentroids,
+//      const float* const __restrict pqFineCentroids,
+//      const uint8_t* const __restrict code0,
+//      const float weight0,
+//      const uint8_t* const __restrict code1,
+//      const float weight1,
+//      float* const __restrict outputAccum);
+//   }
+//
+// * And one more overload for ::accum() that decodes and accumulates
+//   three vectors per call.
+//   {
+//     const faiss::Index* const index;
+//     const uint8_t* const input0;
+//     float weight0;
+//     const uint8_t* const input1;
+//     float weight1;
+//     const uint8_t* const input2;
+//     float weight2;
+//
+//     std::vector<float> buffer(d, 0);
+//
+//     index->sa_decode(1, input0, buffer.data());
+//     for (size_t iDim = 0; iDim < d; iDim++)
+//       output[iDim] += weight0 * buffer[iDim];
+//
+//     index->sa_decode(1, input1, buffer.data());
+//     for (size_t iDim = 0; iDim < d; iDim++)
+//       output[iDim] += weight1 * buffer[iDim];
+//
+//     index->sa_decode(1, input2, buffer.data());
+//     for (size_t iDim = 0; iDim < d; iDim++)
+//       output[iDim] += weight2 * buffer[iDim];
+//   }
+//
+//   If each code uses its own coarse quantizer centroids table and its own fine
+//   quantizer centroids table, then the following overload can be used:
+//   {
+//    static void accum(
+//      const float* const __restrict pqCoarseCentroids0,
+//      const float* const __restrict pqFineCentroids0,
+//      const uint8_t* const __restrict code0,
+//      const float weight0,
+//      const float* const __restrict pqCoarseCentroids1,
+//      const float* const __restrict pqFineCentroids1,
+//      const uint8_t* const __restrict code1,
+//      const float weight1,
+//      const float* const __restrict pqCoarseCentroids2,
+//      const float* const __restrict pqFineCentroids2,
+//      const uint8_t* const __restrict code2,
+//      const float weight2,
+//      float* const __restrict outputAccum);
+//   }
+//   If codes share the coarse quantizer centroids table and also share
+//   the fine quantizer centroids table, then the following overload can be
+//   used:
+//   {
+//    static void accum(
+//      const float* const __restrict pqCoarseCentroids,
+//      const float* const __restrict pqFineCentroids,
+//      const uint8_t* const __restrict code0,
+//      const float weight0,
+//      const uint8_t* const __restrict code1,
+//      const float weight1,
+//      const uint8_t* const __restrict code2,
+//      const float weight2,
+//      float* const __restrict outputAccum);
+//   }
+//
+// The provided version is not multithreaded.
+//
+// Currently, an AVX2+FMA implementation is available. AVX512 version is also
+//   doable, but it was found to be slower than AVX2 for real world applications
+//   that I needed.
+//
+////////////////////////////////////////////////////////////////////////////////////
+//
+// It is possible to use an additional index wrapper on top of IVFPQ /
+// Residual+PQ, known as IndexRowwiseMinMax / IndexRowwiseMinMaxFP16. Index
+// wrapper that performs rowwise normalization to [0,1], preserving the
+// coefficients. This is a vector codec index only.
+// For more details please refer to the description in
+// faiss/IndexRowwiseMinMax.h file.
+//
+// If such a wrapper is used, then the quantizer will look like, say,
+//    MinMaxFP16,IVF256,PQ32np
+//  or
+//    MinMax,PQ16np
+// In this case, please use the following contruction for the decoding,
+// basically, wrapping a kernel in a kernel:
+//   {
+//      using SubT = faiss::cppcontrib::Index2LevelDecoder<128, 128, 2>;
+//      using T = faiss::cppcontrib::IndexMinMaxFP16Decoder<SubT>;
+//      // do T::store(...) or T::accum(...)
+//   }
+//
+// T::accum(...) contains an additional function variable which is
+// used for accumulating scaling. Thus, the code pattern is the following:
+//   {
+//     const float* const __restrict pqCoarseCentroidsQ;
+//     const float* const __restrict pqFineCentroidsQ;
+//     const uint8_t* const __restrict input;
+//     const float* const __restrict weights;
+//     float* const __restrict output;
+//     float outputAccumMin = 0;
+//
+//     for (size_t i = 0; i < n; i++) {
+//         T::accum(
+//                 pqCoarseCentroidsQ,
+//                 pqFineCentroidsQ,
+//                 input + i * code_size,
+//                 weights[i],
+//                 output,
+//                 outputAccumMin);
+//     }
+//     for (size_t j = 0; j < d; j++)
+//         output[j] += outputAccumMin;
+//   }
+// This is similar to the following regular pseudo-code:
+//   {
+//     const faiss::Index* const index;
+//     const uint8_t* const __restrict input;
+//     const float* const __restrict weights;
+//     float* const __restrict output;
+//
+//     for (size_t i = 0; i < n; i++) {
+//       std::vector<float> buffer(d, 0);
+//
+//       index->sa_decode(1, input + i * code_size, buffer.data());
+//       for (size_t j = 0; j < d; j++)
+//         output[j] += weights[i] * buffer[j];
+//     }
+
+#include <faiss/cppcontrib/sa_decode/MinMax-inl.h>
+#include <faiss/cppcontrib/sa_decode/MinMaxFP16-inl.h>
+
+#ifdef __AVX2__
+#include <faiss/cppcontrib/sa_decode/Level2-avx2-inl.h>
+#include <faiss/cppcontrib/sa_decode/PQ-avx2-inl.h>
+#elif defined(__ARM_NEON)
+#include <faiss/cppcontrib/sa_decode/Level2-neon-inl.h>
+#include <faiss/cppcontrib/sa_decode/PQ-neon-inl.h>
+#else
+#include <faiss/cppcontrib/sa_decode/Level2-inl.h>
+#include <faiss/cppcontrib/sa_decode/PQ-inl.h>
+#endif

data/vendor/faiss/faiss/cppcontrib/detail/CoarseBitType.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <cstdint>
+
+namespace faiss {
+namespace cppcontrib {
+namespace detail {
+
+template <int COARSE_BITS>
+struct CoarseBitType {};
+
+template <>
+struct CoarseBitType<8> {
+    using bit_type = uint8_t;
+};
+
+template <>
+struct CoarseBitType<16> {
+    using bit_type = uint16_t;
+};
+
+} // namespace detail
+} // namespace cppcontrib
+} // namespace faiss

data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h

@@ -0,0 +1,195 @@
+#pragma once
+
+#include <cstdint>
+
+namespace faiss {
+namespace cppcontrib {
+namespace detail {
+
+namespace {
+
+template <intptr_t N_ELEMENTS, intptr_t CPOS>
+struct Uint8Reader {
+    static_assert(CPOS < N_ELEMENTS, "CPOS should be less than N_ELEMENTS");
+
+    static intptr_t get(const uint8_t* const __restrict codes) {
+        // Read using 4-bytes, if possible.
+        // Reading using 8-byte takes too many registers somewhy.
+
+        constexpr intptr_t ELEMENT_TO_READ = CPOS / 4;
+        constexpr intptr_t SUB_ELEMENT = CPOS % 4;
+
+        switch (SUB_ELEMENT) {
+            case 0: {
+                if (N_ELEMENTS > CPOS + 3) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 4);
+                    return (code32 & 0x000000FF);
+                } else {
+                    return codes[CPOS];
+                }
+            }
+            case 1: {
+                if (N_ELEMENTS > CPOS + 2) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 4);
+                    return (code32 & 0x0000FF00) >> 8;
+                } else {
+                    return codes[CPOS];
+                }
+            }
+            case 2: {
+                if (N_ELEMENTS > CPOS + 1) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 4);
+                    return (code32 & 0x00FF0000) >> 16;
+                } else {
+                    return codes[CPOS];
+                }
+            }
+            case 3: {
+                if (N_ELEMENTS > CPOS) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 4);
+                    return (code32) >> 24;
+                } else {
+                    return codes[CPOS];
+                }
+            }
+        }
+    }
+};
+
+// reduces the number of read operations from RAM
+///////////////////////////////////////////////
+// 76543210 76543210 76543210 76543210 76543210
+// 00000000 00
+//            111111 1111
+//                       2222 222222
+//                                  33 33333333
+template <intptr_t N_ELEMENTS, intptr_t CPOS>
+struct Uint10Reader {
+    static_assert(CPOS < N_ELEMENTS, "CPOS should be less than N_ELEMENTS");
+
+    static intptr_t get(const uint8_t* const __restrict codes) {
+        // Read using 4-bytes or 2-bytes.
+
+        constexpr intptr_t ELEMENT_TO_READ = CPOS / 4;
+        constexpr intptr_t SUB_ELEMENT = CPOS % 4;
+
+        switch (SUB_ELEMENT) {
+            case 0: {
+                if (N_ELEMENTS > CPOS + 2) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 5);
+                    return (code32 & 0b0000001111111111);
+                } else {
+                    const uint16_t code16 = *reinterpret_cast<const uint16_t*>(
+                            codes + ELEMENT_TO_READ * 5 + 0);
+                    return (code16 & 0b0000001111111111);
+                }
+            }
+            case 1: {
+                if (N_ELEMENTS > CPOS + 1) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 5);
+                    return (code32 & 0b000011111111110000000000) >> 10;
+                } else {
+                    const uint16_t code16 = *reinterpret_cast<const uint16_t*>(
+                            codes + ELEMENT_TO_READ * 5 + 1);
+                    return (code16 & 0b0000111111111100) >> 2;
+                }
+            }
+            case 2: {
+                if (N_ELEMENTS > CPOS) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 5);
+                    return (code32 & 0b00111111111100000000000000000000) >> 20;
+                } else {
+                    const uint16_t code16 = *reinterpret_cast<const uint16_t*>(
+                            codes + ELEMENT_TO_READ * 5 + 2);
+                    return (code16 & 0b0011111111110000) >> 4;
+                }
+            }
+            case 3: {
+                const uint16_t code16 = *reinterpret_cast<const uint16_t*>(
+                        codes + ELEMENT_TO_READ * 5 + 3);
+                return (code16 & 0b1111111111000000) >> 6;
+            }
+        }
+    }
+};
+
+// reduces the number of read operations from RAM
+template <intptr_t N_ELEMENTS, intptr_t CPOS>
+struct Uint16Reader {
+    static_assert(CPOS < N_ELEMENTS, "CPOS should be less than N_ELEMENTS");
+
+    static intptr_t get(const uint8_t* const __restrict codes) {
+        // Read using 4-bytes or 2-bytes.
+        // Reading using 8-byte takes too many registers somewhy.
+
+        constexpr intptr_t ELEMENT_TO_READ = CPOS / 2;
+        constexpr intptr_t SUB_ELEMENT = CPOS % 2;
+
+        switch (SUB_ELEMENT) {
+            case 0: {
+                if (N_ELEMENTS > CPOS + 1) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 4);
+                    return (code32 & 0x0000FFFF);
+                } else {
+                    const uint16_t* const __restrict codesFp16 =
+                            reinterpret_cast<const uint16_t*>(codes);
+                    return codesFp16[CPOS];
+                }
+            }
+            case 1: {
+                if (N_ELEMENTS > CPOS) {
+                    const uint32_t code32 = *reinterpret_cast<const uint32_t*>(
+                            codes + ELEMENT_TO_READ * 4);
+                    return code32 >> 16;
+                } else {
+                    const uint16_t* const __restrict codesFp16 =
+                            reinterpret_cast<const uint16_t*>(codes);
+                    return codesFp16[CPOS];
+                }
+            }
+        }
+    }
+};
+
+//
+template <intptr_t N_ELEMENTS, intptr_t CODE_BITS, intptr_t CPOS>
+struct UintReaderImplType {};
+
+template <intptr_t N_ELEMENTS, intptr_t CPOS>
+struct UintReaderImplType<N_ELEMENTS, 8, CPOS> {
+    using reader_type = Uint8Reader<N_ELEMENTS, CPOS>;
+};
+
+template <intptr_t N_ELEMENTS, intptr_t CPOS>
+struct UintReaderImplType<N_ELEMENTS, 10, CPOS> {
+    using reader_type = Uint10Reader<N_ELEMENTS, CPOS>;
+};
+
+template <intptr_t N_ELEMENTS, intptr_t CPOS>
+struct UintReaderImplType<N_ELEMENTS, 16, CPOS> {
+    using reader_type = Uint16Reader<N_ELEMENTS, CPOS>;
+};
+
+} // namespace
+
+// reduces the number of read operations from RAM
+template <intptr_t DIM, intptr_t CODE_SIZE, intptr_t CODE_BITS, intptr_t CPOS>
+using UintReader =
+        typename UintReaderImplType<DIM / CODE_SIZE, CODE_BITS, CPOS>::
+                reader_type;
+
+template <intptr_t N_ELEMENTS, intptr_t CODE_BITS, intptr_t CPOS>
+using UintReaderRaw =
+        typename UintReaderImplType<N_ELEMENTS, CODE_BITS, CPOS>::reader_type;
+
+} // namespace detail
+} // namespace cppcontrib
+} // namespace faiss