faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/impl/kmeans1d.cpp

@@ -278,13 +278,15 @@ double kmeans1d(const float* x, size_t n, size_t nclusters, float* centroids) {
     ****************************************************/
 
     // for imbalance factor
-    double tot = 0.0, uf = 0.0;
+    double tot = 0.0;
+    double uf = 0.0;
 
     idx_t end = n;
     for (idx_t k = nclusters - 1; k >= 0; k--) {
-        idx_t start = T.at(k, end - 1);
-        float sum = std::accumulate(&arr[start], &arr[end], 0.0f);
-        idx_t size = end - start;
+        const idx_t start = T.at(k, end - 1);
+        const float sum =
+                std::accumulate(arr.data() + start, arr.data() + end, 0.0f);
+        const idx_t size = end - start;
         FAISS_THROW_IF_NOT_FMT(
                 size > 0, "Cluster %d: size %d", int(k), int(size));
         centroids[k] = sum / size;

data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp

@@ -122,30 +122,70 @@ void pq4_pack_codes_range(
     }
 }
 
+namespace {
+
+// get the specific address of the vector inside a block
+// shift is used for determine the if the saved in bits 0..3 (false) or
+// bits 4..7 (true)
+uint8_t get_vector_specific_address(
+        size_t bbs,
+        size_t vector_id,
+        size_t sq,
+        bool& shift) {
+    // get the vector_id inside the block
+    vector_id = vector_id % bbs;
+    shift = vector_id > 15;
+    vector_id = vector_id & 15;
+
+    // get the address of the vector in sq
+    size_t address;
+    if (vector_id < 8) {
+        address = vector_id << 1;
+    } else {
+        address = ((vector_id - 8) << 1) + 1;
+    }
+    if (sq & 1) {
+        address += 16;
+    }
+    return (sq >> 1) * bbs + address;
+}
+
+} // anonymous namespace
+
 uint8_t pq4_get_packed_element(
         const uint8_t* data,
         size_t bbs,
         size_t nsq,
-        size_t i,
+        size_t vector_id,
         size_t sq) {
     // move to correct bbs-sized block
-    data += (i / bbs * (nsq / 2) + sq / 2) * bbs;
-    sq = sq & 1;
-    i = i % bbs;
-
-    // another step
-    data += (i / 32) * 32;
-    i = i % 32;
-
-    if (sq == 1) {
-        data += 16;
+    // number of blocks * block size
+    data += (vector_id / bbs) * (((nsq + 1) / 2) * bbs);
+    bool shift;
+    size_t address = get_vector_specific_address(bbs, vector_id, sq, shift);
+    if (shift) {
+        return data[address] >> 4;
+    } else {
+        return data[address] & 15;
     }
-    const uint8_t iperm0[16] = {
-            0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15};
-    if (i < 16) {
-        return data[iperm0[i]] & 15;
+}
+
+void pq4_set_packed_element(
+        uint8_t* data,
+        uint8_t code,
+        size_t bbs,
+        size_t nsq,
+        size_t vector_id,
+        size_t sq) {
+    // move to correct bbs-sized block
+    // number of blocks * block size
+    data += (vector_id / bbs) * (((nsq + 1) / 2) * bbs);
+    bool shift;
+    size_t address = get_vector_specific_address(bbs, vector_id, sq, shift);
+    if (shift) {
+        data[address] = (code << 4) | (data[address] & 15);
     } else {
-        return data[iperm0[i - 16]] >> 4;
+        data[address] = code | (data[address] & ~15);
     }
 }
 

data/vendor/faiss/faiss/impl/pq4_fast_scan.h

@@ -26,7 +26,7 @@ namespace faiss {
  *  The unused bytes are set to 0.
  *
  * @param codes   input codes, size (ntotal, ceil(M / 2))
- * @param nototal number of input codes
+ * @param ntotal  number of input codes
  * @param nb      output number of codes (ntotal rounded up to a multiple of
  *                bbs)
  * @param M2      number of sub-quantizers (=M rounded up to a muliple of 2)
@@ -61,14 +61,27 @@ void pq4_pack_codes_range(
 
 /** get a single element from a packed codes table
  *
- * @param i        vector id
+ * @param vector_id        vector id
  * @param sq       subquantizer (< nsq)
  */
 uint8_t pq4_get_packed_element(
         const uint8_t* data,
         size_t bbs,
         size_t nsq,
-        size_t i,
+        size_t vector_id,
+        size_t sq);
+
+/** set a single element "code" into a packed codes table
+ *
+ * @param vector_id       vector id
+ * @param sq       subquantizer (< nsq)
+ */
+void pq4_set_packed_element(
+        uint8_t* data,
+        uint8_t code,
+        size_t bbs,
+        size_t nsq,
+        size_t vector_id,
         size_t sq);
 
 /** Pack Look-up table for consumption by the kernel.
@@ -88,8 +101,9 @@ void pq4_pack_LUT(int nq, int nsq, const uint8_t* src, uint8_t* dest);
  * @param nsq     number of sub-quantizers (muliple of 2)
  * @param codes   packed codes array
  * @param LUT     packed look-up table
+ * @param scaler  scaler to scale the encoded norm
  */
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop(
         int nq,
         size_t nb,
@@ -97,7 +111,8 @@ void pq4_accumulate_loop(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res);
+        ResultHandler& res,
+        const Scaler& scaler);
 
 /* qbs versions, supported only for bbs=32.
  *
@@ -141,20 +156,22 @@ int pq4_pack_LUT_qbs_q_map(
 
 /** Run accumulation loop.
  *
- * @param qbs     4-bit encded number of queries
+ * @param qbs     4-bit encoded number of queries
  * @param nb      number of database codes (mutliple of bbs)
  * @param nsq     number of sub-quantizers
  * @param codes   encoded database vectors (packed)
  * @param LUT     look-up table (packed)
  * @param res     call-back for the resutls
+ * @param scaler  scaler to scale the encoded norm
  */
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop_qbs(
         int qbs,
         size_t nb,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res);
+        ResultHandler& res,
+        const Scaler& scaler);
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/pq4_fast_scan.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/simd_result_handlers.h>
 
 namespace faiss {
@@ -26,12 +27,13 @@ namespace {
  * writes results in a ResultHandler
  */
 
-template <int NQ, int BB, class ResultHandler>
+template <int NQ, int BB, class ResultHandler, class Scaler>
 void kernel_accumulate_block(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     // distance accumulators
     simd16uint16 accu[NQ][BB][4];
 
@@ -44,7 +46,7 @@ void kernel_accumulate_block(
         }
     }
 
-    for (int sq = 0; sq < nsq; sq += 2) {
+    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
         simd32uint8 lut_cache[NQ];
         for (int q = 0; q < NQ; q++) {
             lut_cache[q] = simd32uint8(LUT);
@@ -72,6 +74,35 @@ void kernel_accumulate_block(
         }
     }
 
+    for (int sq = 0; sq < scaler.nscale; sq += 2) {
+        simd32uint8 lut_cache[NQ];
+        for (int q = 0; q < NQ; q++) {
+            lut_cache[q] = simd32uint8(LUT);
+            LUT += 32;
+        }
+
+        for (int b = 0; b < BB; b++) {
+            simd32uint8 c = simd32uint8(codes);
+            codes += 32;
+            simd32uint8 mask(15);
+            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
+            simd32uint8 clo = c & mask;
+
+            for (int q = 0; q < NQ; q++) {
+                simd32uint8 lut = lut_cache[q];
+
+                simd32uint8 res0 = scaler.lookup(lut, clo);
+                accu[q][b][0] += scaler.scale_lo(res0); // handle vectors 0..7
+                accu[q][b][1] += scaler.scale_hi(res0); // handle vectors 8..15
+
+                simd32uint8 res1 = scaler.lookup(lut, chi);
+                accu[q][b][2] += scaler.scale_lo(res1); // handle vectors 16..23
+                accu[q][b][3] +=
+                        scaler.scale_hi(res1); //  handle vectors 24..31
+            }
+        }
+    }
+
     for (int q = 0; q < NQ; q++) {
         for (int b = 0; b < BB; b++) {
             accu[q][b][0] -= accu[q][b][1] << 8;
@@ -85,17 +116,18 @@ void kernel_accumulate_block(
     }
 }
 
-template <int NQ, int BB, class ResultHandler>
+template <int NQ, int BB, class ResultHandler, class Scaler>
 void accumulate_fixed_blocks(
         size_t nb,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     constexpr int bbs = 32 * BB;
     for (int64_t j0 = 0; j0 < nb; j0 += bbs) {
         FixedStorageHandler<NQ, 2 * BB> res2;
-        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2);
+        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
         res.set_block_origin(0, j0);
         res2.to_other_handler(res);
         codes += bbs * nsq / 2;
@@ -104,7 +136,7 @@ void accumulate_fixed_blocks(
 
 } // anonymous namespace
 
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop(
         int nq,
         size_t nb,
@@ -112,15 +144,16 @@ void pq4_accumulate_loop(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
     FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
     FAISS_THROW_IF_NOT(nb % bbs == 0);
 
-#define DISPATCH(NQ, BB)                                           \
-    case NQ * 1000 + BB:                                           \
-        accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res); \
+#define DISPATCH(NQ, BB)                                                   \
+    case NQ * 1000 + BB:                                                   \
+        accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res, scaler); \
         break
 
     switch (nq * 1000 + bbs / 32) {
@@ -141,20 +174,28 @@ void pq4_accumulate_loop(
 
 // explicit template instantiations
 
-#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map)         \
-    template void pq4_accumulate_loop<TH<C, with_id_map>>( \
-            int,                                           \
-            size_t,                                        \
-            int,                                           \
-            int,                                           \
-            const uint8_t*,                                \
-            const uint8_t*,                                \
-            TH<C, with_id_map>&);
-
-#define INSTANTIATE_3(C, with_id_map)                           \
-    INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map) \
-    INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map)         \
-    INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map)
+#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map, S)         \
+    template void pq4_accumulate_loop<TH<C, with_id_map>, S>( \
+            int,                                              \
+            size_t,                                           \
+            int,                                              \
+            int,                                              \
+            const uint8_t*,                                   \
+            const uint8_t*,                                   \
+            TH<C, with_id_map>&,                              \
+            const S&);
+
+using DS = DummyScaler;
+using NS = NormTableScaler;
+
+#define INSTANTIATE_3(C, with_id_map)                               \
+    INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map, DS) \
+    INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map, DS)         \
+    INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map, DS)    \
+                                                                    \
+    INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map, NS) \
+    INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map, NS)         \
+    INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map, NS)
 
 using Csi = CMax<uint16_t, int>;
 INSTANTIATE_3(Csi, false);

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/pq4_fast_scan.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/simd_result_handlers.h>
 #include <faiss/utils/simdlib.h>
 
@@ -27,15 +28,17 @@ namespace {
  * writes results in a ResultHandler
  */
 
-template <int NQ, class ResultHandler>
+template <int NQ, class ResultHandler, class Scaler>
 void kernel_accumulate_block(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     // dummy alloc to keep the windows compiler happy
     constexpr int NQA = NQ > 0 ? NQ : 1;
     // distance accumulators
+    // layout: accu[q][b]: distance accumulator for vectors 8*b..8*b+7
     simd16uint16 accu[NQA][4];
 
     for (int q = 0; q < NQ; q++) {
@@ -45,7 +48,7 @@ void kernel_accumulate_block(
     }
 
     // _mm_prefetch(codes + 768, 0);
-    for (int sq = 0; sq < nsq; sq += 2) {
+    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
         // prefetch
         simd32uint8 c(codes);
         codes += 32;
@@ -71,6 +74,31 @@ void kernel_accumulate_block(
         }
     }
 
+    for (int sq = 0; sq < scaler.nscale; sq += 2) {
+        // prefetch
+        simd32uint8 c(codes);
+        codes += 32;
+
+        simd32uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
+        simd32uint8 clo = c & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 2 quantizers
+            simd32uint8 lut(LUT);
+            LUT += 32;
+
+            simd32uint8 res0 = scaler.lookup(lut, clo);
+            accu[q][0] += scaler.scale_lo(res0); // handle vectors 0..7
+            accu[q][1] += scaler.scale_hi(res0); // handle vectors 8..15
+
+            simd32uint8 res1 = scaler.lookup(lut, chi);
+            accu[q][2] += scaler.scale_lo(res1); // handle vectors 16..23
+            accu[q][3] += scaler.scale_hi(res1); //  handle vectors 24..31
+        }
+    }
+
     for (int q = 0; q < NQ; q++) {
         accu[q][0] -= accu[q][1] << 8;
         simd16uint16 dis0 = combine2x2(accu[q][0], accu[q][1]);
@@ -81,13 +109,14 @@ void kernel_accumulate_block(
 }
 
 // handle at most 4 blocks of queries
-template <int QBS, class ResultHandler>
+template <int QBS, class ResultHandler, class Scaler>
 void accumulate_q_4step(
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT0,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     constexpr int Q1 = QBS & 15;
     constexpr int Q2 = (QBS >> 4) & 15;
     constexpr int Q3 = (QBS >> 8) & 15;
@@ -97,21 +126,21 @@ void accumulate_q_4step(
     for (int64_t j0 = 0; j0 < ntotal2; j0 += 32) {
         FixedStorageHandler<SQ, 2> res2;
         const uint8_t* LUT = LUT0;
-        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2);
+        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2, scaler);
         LUT += Q1 * nsq * 16;
         if (Q2 > 0) {
             res2.set_block_origin(Q1, 0);
-            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2, scaler);
             LUT += Q2 * nsq * 16;
         }
         if (Q3 > 0) {
             res2.set_block_origin(Q1 + Q2, 0);
-            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2, scaler);
             LUT += Q3 * nsq * 16;
         }
         if (Q4 > 0) {
             res2.set_block_origin(Q1 + Q2 + Q3, 0);
-            kernel_accumulate_block<Q4>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q4>(nsq, codes, LUT, res2, scaler);
         }
         res.set_block_origin(0, j0);
         res2.to_other_handler(res);
@@ -119,29 +148,31 @@ void accumulate_q_4step(
     }
 }
 
-template <int NQ, class ResultHandler>
+template <int NQ, class ResultHandler, class Scaler>
 void kernel_accumulate_block_loop(
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     for (int64_t j0 = 0; j0 < ntotal2; j0 += 32) {
         res.set_block_origin(0, j0);
         kernel_accumulate_block<NQ, ResultHandler>(
-                nsq, codes + j0 * nsq / 2, LUT, res);
+                nsq, codes + j0 * nsq / 2, LUT, res, scaler);
     }
 }
 
 // non-template version of accumulate kernel -- dispatches dynamically
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void accumulate(
         int nq,
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     assert(nsq % 2 == 0);
     assert(is_aligned_pointer(codes));
     assert(is_aligned_pointer(LUT));
@@ -149,7 +180,7 @@ void accumulate(
 #define DISPATCH(NQ)                                     \
     case NQ:                                             \
         kernel_accumulate_block_loop<NQ, ResultHandler>( \
-                ntotal2, nsq, codes, LUT, res);          \
+                ntotal2, nsq, codes, LUT, res, scaler);  \
         return
 
     switch (nq) {
@@ -165,23 +196,24 @@ void accumulate(
 
 } // namespace
 
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop_qbs(
         int qbs,
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT0,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     assert(nsq % 2 == 0);
     assert(is_aligned_pointer(codes));
     assert(is_aligned_pointer(LUT0));
 
     // try out optimized versions
     switch (qbs) {
-#define DISPATCH(QBS)                                            \
-    case QBS:                                                    \
-        accumulate_q_4step<QBS>(ntotal2, nsq, codes, LUT0, res); \
+#define DISPATCH(QBS)                                                    \
+    case QBS:                                                            \
+        accumulate_q_4step<QBS>(ntotal2, nsq, codes, LUT0, res, scaler); \
         return;
         DISPATCH(0x3333); // 12
         DISPATCH(0x2333); // 11
@@ -219,9 +251,10 @@ void pq4_accumulate_loop_qbs(
             int nq = qi & 15;
             qi >>= 4;
             res.set_block_origin(i0, j0);
-#define DISPATCH(NQ)                                                      \
-    case NQ:                                                              \
-        kernel_accumulate_block<NQ, ResultHandler>(nsq, codes, LUT, res); \
+#define DISPATCH(NQ)                                \
+    case NQ:                                        \
+        kernel_accumulate_block<NQ, ResultHandler>( \
+                nsq, codes, LUT, res, scaler);      \
         break
             switch (nq) {
                 DISPATCH(1);
@@ -241,9 +274,23 @@ void pq4_accumulate_loop_qbs(
 
 // explicit template instantiations
 
-#define INSTANTIATE_ACCUMULATE_Q(RH)           \
-    template void pq4_accumulate_loop_qbs<RH>( \
-            int, size_t, int, const uint8_t*, const uint8_t*, RH&);
+#define INSTANTIATE_ACCUMULATE_Q(RH)                            \
+    template void pq4_accumulate_loop_qbs<RH, DummyScaler>(     \
+            int,                                                \
+            size_t,                                             \
+            int,                                                \
+            const uint8_t*,                                     \
+            const uint8_t*,                                     \
+            RH&,                                                \
+            const DummyScaler&);                                \
+    template void pq4_accumulate_loop_qbs<RH, NormTableScaler>( \
+            int,                                                \
+            size_t,                                             \
+            int,                                                \
+            const uint8_t*,                                     \
+            const uint8_t*,                                     \
+            RH&,                                                \
+            const NormTableScaler&);
 
 using Csi = CMax<uint16_t, int>;
 INSTANTIATE_ACCUMULATE_Q(SingleResultHandler<Csi>)
@@ -293,7 +340,8 @@ void accumulate_to_mem(
         uint16_t* accu) {
     FAISS_THROW_IF_NOT(ntotal2 % 32 == 0);
     StoreResultHandler handler(accu, ntotal2);
-    accumulate(nq, ntotal2, nsq, codes, LUT, handler);
+    DummyScaler scaler;
+    accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler);
 }
 
 int pq4_preferred_qbs(int n) {