faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/IndexFlat.h

@@ -18,7 +18,9 @@ namespace faiss {
 
 /** Index that stores the full vectors and performs exhaustive search */
 struct IndexFlat : IndexFlatCodes {
-    explicit IndexFlat(idx_t d, MetricType metric = METRIC_L2);
+    explicit IndexFlat(
+            idx_t d, ///< dimensionality of the input vectors
+            MetricType metric = METRIC_L2);
 
     void search(
             idx_t n,
@@ -76,8 +78,25 @@ struct IndexFlatIP : IndexFlat {
 };
 
 struct IndexFlatL2 : IndexFlat {
+    // Special cache for L2 norms.
+    // If this cache is set, then get_distance_computer() returns
+    // a special version that computes the distance using dot products
+    // and l2 norms.
+    std::vector<float> cached_l2norms;
+
+    /**
+     * @param d dimensionality of the input vectors
+     */
     explicit IndexFlatL2(idx_t d) : IndexFlat(d, METRIC_L2) {}
     IndexFlatL2() {}
+
+    // override for l2 norms cache.
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
+
+    // compute L2 norms
+    void sync_l2norms();
+    // clear L2 norms
+    void clear_l2norms();
 };
 
 /// optimized version for 1D "vectors".

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -12,6 +12,8 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/extra_distances.h>
 
 namespace faiss {
 
@@ -70,11 +72,6 @@ void IndexFlatCodes::reconstruct(idx_t key, float* recons) const {
     reconstruct_n(key, 1, recons);
 }
 
-FlatCodesDistanceComputer* IndexFlatCodes::get_FlatCodesDistanceComputer()
-        const {
-    FAISS_THROW_MSG("not implemented");
-}
-
 void IndexFlatCodes::check_compatible_for_merge(const Index& otherIndex) const {
     // minimal sanity checks
     const IndexFlatCodes* other =
@@ -103,4 +100,172 @@ CodePacker* IndexFlatCodes::get_CodePacker() const {
     return new CodePackerFlat(code_size);
 }
 
+void IndexFlatCodes::permute_entries(const idx_t* perm) {
+    std::vector<uint8_t> new_codes(codes.size());
+
+    for (idx_t i = 0; i < ntotal; i++) {
+        memcpy(new_codes.data() + i * code_size,
+               codes.data() + perm[i] * code_size,
+               code_size);
+    }
+    std::swap(codes, new_codes);
+}
+
+namespace {
+
+template <class VD>
+struct GenericFlatCodesDistanceComputer : FlatCodesDistanceComputer {
+    const IndexFlatCodes& codec;
+    const VD vd;
+    // temp buffers
+    std::vector<uint8_t> code_buffer;
+    std::vector<float> vec_buffer;
+    const float* query = nullptr;
+
+    GenericFlatCodesDistanceComputer(const IndexFlatCodes* codec, const VD& vd)
+            : FlatCodesDistanceComputer(codec->codes.data(), codec->code_size),
+              codec(*codec),
+              vd(vd),
+              code_buffer(codec->code_size * 4),
+              vec_buffer(codec->d * 4) {}
+
+    void set_query(const float* x) override {
+        query = x;
+    }
+
+    float operator()(idx_t i) override {
+        codec.sa_decode(1, codes + i * code_size, vec_buffer.data());
+        return vd(query, vec_buffer.data());
+    }
+
+    float distance_to_code(const uint8_t* code) override {
+        codec.sa_decode(1, code, vec_buffer.data());
+        return vd(query, vec_buffer.data());
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        codec.sa_decode(1, codes + i * code_size, vec_buffer.data());
+        codec.sa_decode(1, codes + j * code_size, vec_buffer.data() + vd.d);
+        return vd(vec_buffer.data(), vec_buffer.data() + vd.d);
+    }
+
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        uint8_t* cp = code_buffer.data();
+        for (idx_t i : {idx0, idx1, idx2, idx3}) {
+            memcpy(cp, codes + i * code_size, code_size);
+            cp += code_size;
+        }
+        // potential benefit is if batch decoding is more efficient than 1 by 1
+        // decoding
+        codec.sa_decode(4, code_buffer.data(), vec_buffer.data());
+        dis0 = vd(query, vec_buffer.data());
+        dis1 = vd(query, vec_buffer.data() + vd.d);
+        dis2 = vd(query, vec_buffer.data() + 2 * vd.d);
+        dis3 = vd(query, vec_buffer.data() + 3 * vd.d);
+    }
+};
+
+struct Run_get_distance_computer {
+    using T = FlatCodesDistanceComputer*;
+
+    template <class VD>
+    FlatCodesDistanceComputer* f(const VD& vd, const IndexFlatCodes* codec) {
+        return new GenericFlatCodesDistanceComputer<VD>(codec, vd);
+    }
+};
+
+template <class BlockResultHandler>
+struct Run_search_with_decompress {
+    using T = void;
+
+    template <class VectorDistance>
+    void f(VectorDistance& vd,
+           const IndexFlatCodes* index_ptr,
+           const float* xq,
+           BlockResultHandler& res) {
+        // Note that there seems to be a clang (?) bug that "sometimes" passes
+        // the const Index & parameters by value, so to be on the safe side,
+        // it's better to use pointers.
+        const IndexFlatCodes& index = *index_ptr;
+        size_t ntotal = index.ntotal;
+        using SingleResultHandler =
+                typename BlockResultHandler::SingleResultHandler;
+        using DC = GenericFlatCodesDistanceComputer<VectorDistance>;
+#pragma omp parallel // if (res.nq > 100)
+        {
+            std::unique_ptr<DC> dc(new DC(&index, vd));
+            SingleResultHandler resi(res);
+#pragma omp for
+            for (int64_t q = 0; q < res.nq; q++) {
+                resi.begin(q);
+                dc->set_query(xq + vd.d * q);
+                for (size_t i = 0; i < ntotal; i++) {
+                    if (res.is_in_selection(i)) {
+                        float dis = (*dc)(i);
+                        resi.add_result(dis, i);
+                    }
+                }
+                resi.end();
+            }
+        }
+    }
+};
+
+struct Run_search_with_decompress_res {
+    using T = void;
+
+    template <class ResultHandler>
+    void f(ResultHandler& res, const IndexFlatCodes* index, const float* xq) {
+        Run_search_with_decompress<ResultHandler> r;
+        dispatch_VectorDistance(
+                index->d,
+                index->metric_type,
+                index->metric_arg,
+                r,
+                index,
+                xq,
+                res);
+    }
+};
+
+} // anonymous namespace
+
+FlatCodesDistanceComputer* IndexFlatCodes::get_FlatCodesDistanceComputer()
+        const {
+    Run_get_distance_computer r;
+    return dispatch_VectorDistance(d, metric_type, metric_arg, r, this);
+}
+
+void IndexFlatCodes::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    Run_search_with_decompress_res r;
+    const IDSelector* sel = params ? params->sel : nullptr;
+    dispatch_knn_ResultHandler(
+            n, distances, labels, k, metric_type, sel, r, this, x);
+}
+
+void IndexFlatCodes::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    const IDSelector* sel = params ? params->sel : nullptr;
+    Run_search_with_decompress_res r;
+    dispatch_range_ResultHandler(result, radius, metric_type, sel, r, this, x);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.h

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <faiss/Index.h>
@@ -34,7 +32,6 @@ struct IndexFlatCodes : Index {
 
     void reset() override;
 
-    /// reconstruction using the codec interface
     void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
@@ -46,19 +43,41 @@ struct IndexFlatCodes : Index {
      * different from the usual ones: the new ids are shifted */
     size_t remove_ids(const IDSelector& sel) override;
 
-    /** a FlatCodesDistanceComputer offers a distance_to_code method */
+    /** a FlatCodesDistanceComputer offers a distance_to_code method
+     *
+     * The default implementation explicitly decodes the vector with sa_decode.
+     */
     virtual FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const;
 
     DistanceComputer* get_distance_computer() const override {
         return get_FlatCodesDistanceComputer();
     }
 
+    /** Search implemented by decoding */
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
     // returns a new instance of a CodePacker
     CodePacker* get_CodePacker() const;
 
     void check_compatible_for_merge(const Index& otherIndex) const override;
 
     virtual void merge_from(Index& otherIndex, idx_t add_id = 0) override;
+
+    // permute_entries. perm of size ntotal maps new to old positions
+    void permute_entries(const idx_t* perm);
 };
 
 } // namespace faiss