faiss 0.3.1 → 0.3.3

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -31,6 +31,8 @@ namespace {
  * writes results in a ResultHandler
  */
 
+#ifndef __AVX512F__
+
 template <int NQ, class ResultHandler, class Scaler>
 void kernel_accumulate_block(
         int nsq,
@@ -111,6 +113,451 @@ void kernel_accumulate_block(
     }
 }
 
+#else
+
+// a special version for NQ=1.
+// Despite the function being large in the text form, it compiles to a very
+//    compact assembler code.
+FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
+template <class ResultHandler, class Scaler>
+void kernel_accumulate_block_avx512_nq1(
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        ResultHandler& res,
+        const Scaler& scaler) {
+    // NQ is kept in order to match the similarity to baseline function
+    constexpr int NQ = 1;
+    // distance accumulators. We can accept more for NQ=1
+    // layout: accu[q][b]: distance accumulator for vectors 32*b..32*b+15
+    simd32uint16 accu[NQ][4];
+    // layout: accu[q][b]: distance accumulator for vectors 32*b+16..32*b+31
+    simd32uint16 accu1[NQ][4];
+
+    for (int q = 0; q < NQ; q++) {
+        for (int b = 0; b < 4; b++) {
+            accu[q][b].clear();
+            accu1[q][b].clear();
+        }
+    }
+
+    // process "nsq - scaler.nscale" part
+    const int nsq_minus_nscale = nsq - scaler.nscale;
+    const int nsq_minus_nscale_8 = (nsq_minus_nscale / 8) * 8;
+    const int nsq_minus_nscale_4 = (nsq_minus_nscale / 4) * 4;
+
+    // process in chunks of 8
+    for (int sq = 0; sq < nsq_minus_nscale_8; sq += 8) {
+        // prefetch
+        simd64uint8 c(codes);
+        codes += 64;
+
+        simd64uint8 c1(codes);
+        codes += 64;
+
+        simd64uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd64uint8 chi = simd64uint8(simd32uint16(c) >> 4) & mask;
+        simd64uint8 clo = c & mask;
+
+        simd64uint8 c1hi = simd64uint8(simd32uint16(c1) >> 4) & mask;
+        simd64uint8 c1lo = c1 & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd64uint8 lut(LUT);
+            LUT += 64;
+
+            {
+                simd64uint8 res0 = lut.lookup_4_lanes(clo);
+                simd64uint8 res1 = lut.lookup_4_lanes(chi);
+
+                accu[q][0] += simd32uint16(res0);
+                accu[q][1] += simd32uint16(res0) >> 8;
+
+                accu[q][2] += simd32uint16(res1);
+                accu[q][3] += simd32uint16(res1) >> 8;
+            }
+        }
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd64uint8 lut(LUT);
+            LUT += 64;
+
+            {
+                simd64uint8 res0 = lut.lookup_4_lanes(c1lo);
+                simd64uint8 res1 = lut.lookup_4_lanes(c1hi);
+
+                accu1[q][0] += simd32uint16(res0);
+                accu1[q][1] += simd32uint16(res0) >> 8;
+
+                accu1[q][2] += simd32uint16(res1);
+                accu1[q][3] += simd32uint16(res1) >> 8;
+            }
+        }
+    }
+
+    // process leftovers: a single chunk of size 4
+    if (nsq_minus_nscale_8 != nsq_minus_nscale_4) {
+        // prefetch
+        simd64uint8 c(codes);
+        codes += 64;
+
+        simd64uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd64uint8 chi = simd64uint8(simd32uint16(c) >> 4) & mask;
+        simd64uint8 clo = c & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd64uint8 lut(LUT);
+            LUT += 64;
+
+            simd64uint8 res0 = lut.lookup_4_lanes(clo);
+            simd64uint8 res1 = lut.lookup_4_lanes(chi);
+
+            accu[q][0] += simd32uint16(res0);
+            accu[q][1] += simd32uint16(res0) >> 8;
+
+            accu[q][2] += simd32uint16(res1);
+            accu[q][3] += simd32uint16(res1) >> 8;
+        }
+    }
+
+    // process leftovers: a single chunk of size 2
+    if (nsq_minus_nscale_4 != nsq_minus_nscale) {
+        // 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 = lut.lookup_2_lanes(clo);
+            simd32uint8 res1 = lut.lookup_2_lanes(chi);
+
+            accu[q][0] += simd32uint16(simd16uint16(res0));
+            accu[q][1] += simd32uint16(simd16uint16(res0) >> 8);
+
+            accu[q][2] += simd32uint16(simd16uint16(res1));
+            accu[q][3] += simd32uint16(simd16uint16(res1) >> 8);
+        }
+    }
+
+    // process "sq" part
+    const int nscale = scaler.nscale;
+    const int nscale_8 = (nscale / 8) * 8;
+    const int nscale_4 = (nscale / 4) * 4;
+
+    // process in chunks of 8
+    for (int sq = 0; sq < nscale_8; sq += 8) {
+        // prefetch
+        simd64uint8 c(codes);
+        codes += 64;
+
+        simd64uint8 c1(codes);
+        codes += 64;
+
+        simd64uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd64uint8 chi = simd64uint8(simd32uint16(c) >> 4) & mask;
+        simd64uint8 clo = c & mask;
+
+        simd64uint8 c1hi = simd64uint8(simd32uint16(c1) >> 4) & mask;
+        simd64uint8 c1lo = c1 & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd64uint8 lut(LUT);
+            LUT += 64;
+
+            {
+                simd64uint8 res0 = scaler.lookup(lut, clo);
+                accu[q][0] += scaler.scale_lo(res0); // handle vectors 0..15
+                accu[q][1] += scaler.scale_hi(res0); // handle vectors 16..31
+
+                simd64uint8 res1 = scaler.lookup(lut, chi);
+                accu[q][2] += scaler.scale_lo(res1); // handle vectors 32..47
+                accu[q][3] += scaler.scale_hi(res1); //  handle vectors 48..63
+            }
+        }
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd64uint8 lut(LUT);
+            LUT += 64;
+
+            {
+                simd64uint8 res0 = scaler.lookup(lut, c1lo);
+                accu1[q][0] += scaler.scale_lo(res0); // handle vectors 0..7
+                accu1[q][1] += scaler.scale_hi(res0); // handle vectors 8..15
+
+                simd64uint8 res1 = scaler.lookup(lut, c1hi);
+                accu1[q][2] += scaler.scale_lo(res1); // handle vectors 16..23
+                accu1[q][3] += scaler.scale_hi(res1); //  handle vectors 24..31
+            }
+        }
+    }
+
+    // process leftovers: a single chunk of size 4
+    if (nscale_8 != nscale_4) {
+        // prefetch
+        simd64uint8 c(codes);
+        codes += 64;
+
+        simd64uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd64uint8 chi = simd64uint8(simd32uint16(c) >> 4) & mask;
+        simd64uint8 clo = c & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd64uint8 lut(LUT);
+            LUT += 64;
+
+            simd64uint8 res0 = scaler.lookup(lut, clo);
+            accu[q][0] += scaler.scale_lo(res0); // handle vectors 0..15
+            accu[q][1] += scaler.scale_hi(res0); // handle vectors 16..31
+
+            simd64uint8 res1 = scaler.lookup(lut, chi);
+            accu[q][2] += scaler.scale_lo(res1); // handle vectors 32..47
+            accu[q][3] += scaler.scale_hi(res1); //  handle vectors 48..63
+        }
+    }
+
+    // process leftovers: a single chunk of size 2
+    if (nscale_4 != nscale) {
+        // 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] +=
+                    simd32uint16(scaler.scale_lo(res0)); // handle vectors 0..7
+            accu[q][1] +=
+                    simd32uint16(scaler.scale_hi(res0)); // handle vectors 8..15
+
+            simd32uint8 res1 = scaler.lookup(lut, chi);
+            accu[q][2] += simd32uint16(
+                    scaler.scale_lo(res1)); // handle vectors 16..23
+            accu[q][3] += simd32uint16(
+                    scaler.scale_hi(res1)); //  handle vectors 24..31
+        }
+    }
+
+    for (int q = 0; q < NQ; q++) {
+        for (int b = 0; b < 4; b++) {
+            accu[q][b] += accu1[q][b];
+        }
+    }
+
+    for (int q = 0; q < NQ; q++) {
+        accu[q][0] -= accu[q][1] << 8;
+        simd16uint16 dis0 = combine4x2(accu[q][0], accu[q][1]);
+        accu[q][2] -= accu[q][3] << 8;
+        simd16uint16 dis1 = combine4x2(accu[q][2], accu[q][3]);
+        res.handle(q, 0, dis0, dis1);
+    }
+}
+
+// general-purpose case
+FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
+template <int NQ, class ResultHandler, class Scaler>
+void kernel_accumulate_block_avx512_nqx(
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        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
+    simd32uint16 accu[NQA][4];
+
+    for (int q = 0; q < NQ; q++) {
+        for (int b = 0; b < 4; b++) {
+            accu[q][b].clear();
+        }
+    }
+
+    // process "nsq - scaler.nscale" part
+    const int nsq_minus_nscale = nsq - scaler.nscale;
+    const int nsq_minus_nscale_4 = (nsq_minus_nscale / 4) * 4;
+
+    // process in chunks of 8
+    for (int sq = 0; sq < nsq_minus_nscale_4; sq += 4) {
+        // prefetch
+        simd64uint8 c(codes);
+        codes += 64;
+
+        simd64uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd64uint8 chi = simd64uint8(simd32uint16(c) >> 4) & mask;
+        simd64uint8 clo = c & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd32uint8 lut_a(LUT);
+            simd32uint8 lut_b(LUT + NQ * 32);
+
+            simd64uint8 lut(lut_a, lut_b);
+            LUT += 32;
+
+            {
+                simd64uint8 res0 = lut.lookup_4_lanes(clo);
+                simd64uint8 res1 = lut.lookup_4_lanes(chi);
+
+                accu[q][0] += simd32uint16(res0);
+                accu[q][1] += simd32uint16(res0) >> 8;
+
+                accu[q][2] += simd32uint16(res1);
+                accu[q][3] += simd32uint16(res1) >> 8;
+            }
+        }
+
+        LUT += NQ * 32;
+    }
+
+    // process leftovers: a single chunk of size 2
+    if (nsq_minus_nscale_4 != nsq_minus_nscale) {
+        // 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 = lut.lookup_2_lanes(clo);
+            simd32uint8 res1 = lut.lookup_2_lanes(chi);
+
+            accu[q][0] += simd32uint16(simd16uint16(res0));
+            accu[q][1] += simd32uint16(simd16uint16(res0) >> 8);
+
+            accu[q][2] += simd32uint16(simd16uint16(res1));
+            accu[q][3] += simd32uint16(simd16uint16(res1) >> 8);
+        }
+    }
+
+    // process "sq" part
+    const int nscale = scaler.nscale;
+    const int nscale_4 = (nscale / 4) * 4;
+
+    // process in chunks of 4
+    for (int sq = 0; sq < nscale_4; sq += 4) {
+        // prefetch
+        simd64uint8 c(codes);
+        codes += 64;
+
+        simd64uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd64uint8 chi = simd64uint8(simd32uint16(c) >> 4) & mask;
+        simd64uint8 clo = c & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 4 quantizers
+            simd32uint8 lut_a(LUT);
+            simd32uint8 lut_b(LUT + NQ * 32);
+
+            simd64uint8 lut(lut_a, lut_b);
+            LUT += 32;
+
+            {
+                simd64uint8 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
+
+                simd64uint8 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
+            }
+        }
+
+        LUT += NQ * 32;
+    }
+
+    // process leftovers: a single chunk of size 2
+    if (nscale_4 != nscale) {
+        // 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] +=
+                    simd32uint16(scaler.scale_lo(res0)); // handle vectors 0..7
+            accu[q][1] +=
+                    simd32uint16(scaler.scale_hi(res0)); // handle vectors 8..15
+
+            simd32uint8 res1 = scaler.lookup(lut, chi);
+            accu[q][2] += simd32uint16(
+                    scaler.scale_lo(res1)); // handle vectors 16..23
+            accu[q][3] += simd32uint16(
+                    scaler.scale_hi(res1)); //  handle vectors 24..31
+        }
+    }
+
+    for (int q = 0; q < NQ; q++) {
+        accu[q][0] -= accu[q][1] << 8;
+        simd16uint16 dis0 = combine4x2(accu[q][0], accu[q][1]);
+        accu[q][2] -= accu[q][3] << 8;
+        simd16uint16 dis1 = combine4x2(accu[q][2], accu[q][3]);
+        res.handle(q, 0, dis0, dis1);
+    }
+}
+
+template <int NQ, class ResultHandler, class Scaler>
+void kernel_accumulate_block(
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        ResultHandler& res,
+        const Scaler& scaler) {
+    if constexpr (NQ == 1) {
+        kernel_accumulate_block_avx512_nq1<ResultHandler, Scaler>(
+                nsq, codes, LUT, res, scaler);
+    } else {
+        kernel_accumulate_block_avx512_nqx<NQ, ResultHandler, Scaler>(
+                nsq, codes, LUT, res, scaler);
+    }
+}
+
+#endif
+
 // handle at most 4 blocks of queries
 template <int QBS, class ResultHandler, class Scaler>
 void accumulate_q_4step(
@@ -304,7 +751,7 @@ void pq4_accumulate_loop_qbs(
         SIMDResultHandler& res,
         const NormTableScaler* scaler) {
     Run_pq4_accumulate_loop_qbs consumer;
-    dispatch_SIMDResultHanlder(res, consumer, qbs, nb, nsq, codes, LUT, scaler);
+    dispatch_SIMDResultHandler(res, consumer, qbs, nb, nsq, codes, LUT, scaler);
 }
 
 /***************************************************************

data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -95,7 +95,7 @@ void accum_and_store_tab(
         for (size_t ij = 1; ij < M; ij++) {
             reg += cbs[ij][kk];
         }
-        output[b * K + kk] = reg;
+        output[kk] = reg;
     }
 }
 
@@ -152,7 +152,7 @@ void accum_and_add_tab(
         for (size_t ij = 1; ij < M; ij++) {
             reg += cbs[ij][kk];
         }
-        output[b * K + kk] += reg;
+        output[kk] += reg;
     }
 }
 
@@ -664,8 +664,6 @@ void refine_beam_mp(
     std::unique_ptr<Index> assign_index;
     if (rq.assign_index_factory) {
         assign_index.reset((*rq.assign_index_factory)(rq.d));
-    } else {
-        assign_index.reset(new IndexFlatL2(rq.d));
     }
 
     // main loop
@@ -701,7 +699,9 @@ void refine_beam_mp(
                 assign_index.get(),
                 rq.approx_topk_mode);
 
-        assign_index->reset();
+        if (assign_index != nullptr) {
+            assign_index->reset();
+        }
 
         std::swap(codes_ptr, new_codes_ptr);
         std::swap(residuals_ptr, new_residuals_ptr);
@@ -809,7 +809,7 @@ void refine_beam_LUT_mp(
                 rq.codebook_offsets.data(),
                 query_cp + rq.codebook_offsets[m],
                 rq.total_codebook_size,
-                rq.cent_norms.data() + rq.codebook_offsets[m],
+                rq.centroid_norms.data() + rq.codebook_offsets[m],
                 m,
                 codes_ptr,
                 distances_ptr,

data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -25,7 +25,7 @@ namespace faiss {
  * It allows low-level access to the encoding function, exposed mainly for unit
  * tests.
  *
- * @param n              number of vectors to hanlde
+ * @param n              number of vectors to handle
  * @param residuals      vectors to encode, size (n, beam_size, d)
  * @param cent           centroids, size (K, d)
  * @param beam_size      input beam size