faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h

@@ -149,11 +149,11 @@ struct QuantizerTemplate<
 /*******************************************************************
  * Lloyd-Max scalar quantizer
  *******************************************************************/
+// Each QuantizerXBase<SL> holds the scalar code of a quantizer. It takes the
+// SIMD level of the file that uses it, so each per-level file compiles its own
+// copy under its own name.
 template <int NBits, SIMDLevel SL>
-struct QuantizerLloydMax;
-
-template <int NBits>
-struct QuantizerLloydMax<NBits, SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+struct QuantizerLloydMaxBase : ScalarQuantizer::SQuantizer {
     static_assert(NBits >= 1 && NBits <= 8);
 
     static constexpr size_t kCentroidsCount = size_t(1) << NBits;
@@ -164,7 +164,7 @@ struct QuantizerLloydMax<NBits, SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
     const float* centroids;
     const float* boundaries;
 
-    QuantizerLloydMax(size_t d_in, const std::vector<float>& trained)
+    QuantizerLloydMaxBase(size_t d_in, const std::vector<float>& trained)
             : d(d_in), centroids(nullptr), boundaries(nullptr) {
         FAISS_THROW_IF_NOT(trained.size() == 2 * kCentroidsCount - 1);
         centroids = trained.data();
@@ -220,8 +220,8 @@ struct QuantizerLloydMax<NBits, SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
 };
 
 template <int NBits, SIMDLevel SL>
-struct QuantizerLloydMax : QuantizerLloydMax<NBits, SIMDLevel::NONE> {
-    using QuantizerLloydMax<NBits, SIMDLevel::NONE>::QuantizerLloydMax;
+struct QuantizerLloydMax : QuantizerLloydMaxBase<NBits, SL> {
+    using QuantizerLloydMaxBase<NBits, SL>::QuantizerLloydMaxBase;
 };
 
 template <int NBits, SIMDLevel SL>
@@ -232,13 +232,10 @@ using QuantizerTurboQuantMSE = QuantizerLloydMax<NBits, SL>;
  *******************************************************************/
 
 template <SIMDLevel SL>
-struct QuantizerFP16;
-
-template <>
-struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+struct QuantizerFP16Base : ScalarQuantizer::SQuantizer {
     const size_t d;
 
-    QuantizerFP16(size_t d_in, const std::vector<float>& /* unused */)
+    QuantizerFP16Base(size_t d_in, const std::vector<float>& /* unused */)
             : d(d_in) {}
 
     void encode_vector(const float* x, uint8_t* code) const final {
@@ -247,7 +244,7 @@ struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
         }
     }
 
-    void decode_vector(const uint8_t* code, float* x) const final {
+    void decode_vector(const uint8_t* code, float* x) const override {
         for (size_t i = 0; i < d; i++) {
             x[i] = decode_fp16(((uint16_t*)code)[i]);
         }
@@ -261,8 +258,8 @@ struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
 };
 
 template <SIMDLevel SL>
-struct QuantizerFP16 : QuantizerFP16<SIMDLevel::NONE> {
-    using QuantizerFP16<SIMDLevel::NONE>::QuantizerFP16;
+struct QuantizerFP16 : QuantizerFP16Base<SL> {
+    using QuantizerFP16Base<SL>::QuantizerFP16Base;
 };
 
 /*******************************************************************
@@ -270,13 +267,10 @@ struct QuantizerFP16 : QuantizerFP16<SIMDLevel::NONE> {
  *******************************************************************/
 
 template <SIMDLevel SL>
-struct QuantizerBF16;
-
-template <>
-struct QuantizerBF16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+struct QuantizerBF16Base : ScalarQuantizer::SQuantizer {
     const size_t d;
 
-    QuantizerBF16(size_t d_in, const std::vector<float>& /* unused */)
+    QuantizerBF16Base(size_t d_in, const std::vector<float>& /* unused */)
             : d(d_in) {}
 
     void encode_vector(const float* x, uint8_t* code) const override {
@@ -295,8 +289,8 @@ struct QuantizerBF16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
 };
 
 template <SIMDLevel SL>
-struct QuantizerBF16 : QuantizerBF16<SIMDLevel::NONE> {
-    using QuantizerBF16<SIMDLevel::NONE>::QuantizerBF16;
+struct QuantizerBF16 : QuantizerBF16Base<SL> {
+    using QuantizerBF16Base<SL>::QuantizerBF16Base;
 };
 
 template <>
@@ -309,22 +303,19 @@ struct QuantizerBF16<SIMDLevel::AVX512_SPR>;
  *******************************************************************/
 
 template <SIMDLevel SL>
-struct Quantizer8bitDirect;
-
-template <>
-struct Quantizer8bitDirect<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+struct Quantizer8bitDirectBase : ScalarQuantizer::SQuantizer {
     const size_t d;
 
-    Quantizer8bitDirect(size_t d_in, const std::vector<float>& /* unused */)
+    Quantizer8bitDirectBase(size_t d_in, const std::vector<float>& /* unused */)
             : d(d_in) {}
 
-    void encode_vector(const float* x, uint8_t* code) const final {
+    void encode_vector(const float* x, uint8_t* code) const override {
         for (size_t i = 0; i < d; i++) {
             code[i] = (uint8_t)x[i];
         }
     }
 
-    void decode_vector(const uint8_t* code, float* x) const final {
+    void decode_vector(const uint8_t* code, float* x) const override {
         for (size_t i = 0; i < d; i++) {
             x[i] = code[i];
         }
@@ -338,8 +329,8 @@ struct Quantizer8bitDirect<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
 };
 
 template <SIMDLevel SL>
-struct Quantizer8bitDirect : Quantizer8bitDirect<SIMDLevel::NONE> {
-    using Quantizer8bitDirect<SIMDLevel::NONE>::Quantizer8bitDirect;
+struct Quantizer8bitDirect : Quantizer8bitDirectBase<SL> {
+    using Quantizer8bitDirectBase<SL>::Quantizer8bitDirectBase;
 };
 
 /*******************************************************************
@@ -347,14 +338,10 @@ struct Quantizer8bitDirect : Quantizer8bitDirect<SIMDLevel::NONE> {
  *******************************************************************/
 
 template <SIMDLevel SL>
-struct Quantizer8bitDirectSigned;
-
-template <>
-struct Quantizer8bitDirectSigned<SIMDLevel::NONE>
-        : ScalarQuantizer::SQuantizer {
+struct Quantizer8bitDirectSignedBase : ScalarQuantizer::SQuantizer {
     const size_t d;
 
-    Quantizer8bitDirectSigned(
+    Quantizer8bitDirectSignedBase(
             size_t d_in,
             const std::vector<float>& /* unused */)
             : d(d_in) {}
@@ -365,7 +352,7 @@ struct Quantizer8bitDirectSigned<SIMDLevel::NONE>
         }
     }
 
-    void decode_vector(const uint8_t* code, float* x) const final {
+    void decode_vector(const uint8_t* code, float* x) const override {
         for (size_t i = 0; i < d; i++) {
             x[i] = code[i] - 128;
         }
@@ -379,8 +366,8 @@ struct Quantizer8bitDirectSigned<SIMDLevel::NONE>
 };
 
 template <SIMDLevel SL>
-struct Quantizer8bitDirectSigned : Quantizer8bitDirectSigned<SIMDLevel::NONE> {
-    using Quantizer8bitDirectSigned<SIMDLevel::NONE>::Quantizer8bitDirectSigned;
+struct Quantizer8bitDirectSigned : Quantizer8bitDirectSignedBase<SL> {
+    using Quantizer8bitDirectSignedBase<SL>::Quantizer8bitDirectSignedBase;
 };
 
 /*******************************************************************
@@ -403,11 +390,7 @@ struct SQTurboQFactors {
 FAISS_PACK_STRUCTS_END
 
 template <int NBits, SIMDLevel SL>
-struct QuantizerTurboQuantFull;
-
-template <int NBits>
-struct QuantizerTurboQuantFull<NBits, SIMDLevel::NONE>
-        : ScalarQuantizer::SQuantizer {
+struct QuantizerTurboQuantFull : ScalarQuantizer::SQuantizer {
     static_assert(NBits >= 2 && NBits <= 5);
 
     static constexpr int kMSEBits = NBits - 1;
@@ -658,13 +641,6 @@ struct QuantizerTurboQuantFull<NBits, SIMDLevel::NONE>
     }
 };
 
-template <int NBits, SIMDLevel SL>
-struct QuantizerTurboQuantFull
-        : QuantizerTurboQuantFull<NBits, SIMDLevel::NONE> {
-    using QuantizerTurboQuantFull<NBits, SIMDLevel::NONE>::
-            QuantizerTurboQuantFull;
-};
-
 /*******************************************************************
  * Selection function
  *******************************************************************/

data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h

@@ -35,6 +35,19 @@ namespace scalar_quantizer {
 using QuantizerType = ScalarQuantizer::QuantizerType;
 using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer;
 
+// True when query_to_codes_batch_4 matches four query_to_code calls. The SIMD
+// DCTemplate batch differs from the single path only for the uniform
+// quantizers, whose single path predecodes the query, so those stay on the
+// scalar scan and batched results never differ from scalar.
+template <class DC>
+constexpr bool sq_batch_4_is_exact() {
+    if constexpr (requires { DC::has_decode_raw(); }) {
+        return !DC::has_decode_raw();
+    } else {
+        return true;
+    }
+}
+
 /*******************************************************************
  * IVFSQScannerIP / IVFSQScannerL2 — moved from anonymous namespace
  * in ScalarQuantizer.cpp
@@ -74,13 +87,35 @@ struct IVFSQScannerIP : InvertedListScanner {
         return accu0 + dc.query_to_code(code);
     }
 
+    void distance_to_codes_batch_4(
+            const uint8_t* code_0,
+            const uint8_t* code_1,
+            const uint8_t* code_2,
+            const uint8_t* code_3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) const {
+        dc.query_to_codes_batch_4(
+                code_0, code_1, code_2, code_3, dis0, dis1, dis2, dis3);
+        dis0 += accu0;
+        dis1 += accu0;
+        dis2 += accu0;
+        dis3 += accu0;
+    }
+
     size_t scan_codes(
             size_t list_size,
             const uint8_t* codes,
             const idx_t* ids,
             ResultHandler& handler) const override {
-        return run_scan_codes_fix_C<CMin<float, idx_t>>(
-                *this, list_size, codes, ids, handler);
+        if constexpr (sq_batch_4_is_exact<DCClass>()) {
+            return run_scan_codes4_fix_C<CMin<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_C<CMin<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        }
     }
 };
 
@@ -133,13 +168,31 @@ struct IVFSQScannerL2 : InvertedListScanner {
         return dc.query_to_code(code);
     }
 
+    void distance_to_codes_batch_4(
+            const uint8_t* code_0,
+            const uint8_t* code_1,
+            const uint8_t* code_2,
+            const uint8_t* code_3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) const {
+        dc.query_to_codes_batch_4(
+                code_0, code_1, code_2, code_3, dis0, dis1, dis2, dis3);
+    }
+
     size_t scan_codes(
             size_t list_size,
             const uint8_t* codes,
             const idx_t* ids,
             ResultHandler& handler) const override {
-        return run_scan_codes_fix_C<CMax<float, idx_t>>(
-                *this, list_size, codes, ids, handler);
+        if constexpr (sq_batch_4_is_exact<DCClass>()) {
+            return run_scan_codes4_fix_C<CMax<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_C<CMax<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        }
     }
 };
 
@@ -175,14 +228,9 @@ struct IVFCoarseDistanceScanner : InvertedListScanner {
 
     void set_query(const float* /*query_vector*/) override {}
 
-    void set_list(idx_t list_no_in, float coarse_dis_in) override {
-        this->list_no = list_no_in;
-        this->coarse_dis = coarse_dis_in;
-    }
+    void set_list(idx_t list_no_in, float coarse_dis_in) override;
 
-    float distance_to_code(const uint8_t* /*code*/) const override {
-        return coarse_dis;
-    }
+    float distance_to_code(const uint8_t* /*code*/) const override;
 };
 
 } // namespace scalar_quantizer