faiss 0.1.0 → 0.1.1

data/vendor/faiss/IndexPreTransform.h

@@ -0,0 +1,91 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#pragma once
+
+
+
+#include <faiss/Index.h>
+#include <faiss/VectorTransform.h>
+
+namespace faiss {
+
+/** Index that applies a LinearTransform transform on vectors before
+ *  handing them over to a sub-index */
+struct IndexPreTransform: Index {
+
+    std::vector<VectorTransform *> chain;  ///! chain of tranforms
+    Index * index;            ///! the sub-index
+
+    bool own_fields;          ///! whether pointers are deleted in destructor
+
+    explicit IndexPreTransform (Index *index);
+
+    IndexPreTransform ();
+
+    /// ltrans is the last transform before the index
+    IndexPreTransform (VectorTransform * ltrans, Index * index);
+
+    void prepend_transform (VectorTransform * ltrans);
+
+    void train(idx_t n, const float* x) override;
+
+    void add(idx_t n, const float* x) override;
+
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
+
+    void reset() override;
+
+    /** removes IDs from the index. Not supported by all indexes.
+     */
+    size_t remove_ids(const IDSelector& sel) override;
+
+    void search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const override;
+
+
+    /* range search, no attempt is done to change the radius */
+    void range_search (idx_t n, const float* x, float radius,
+                       RangeSearchResult* result) const override;
+
+
+    void reconstruct (idx_t key, float * recons) const override;
+
+    void reconstruct_n (idx_t i0, idx_t ni, float *recons)
+        const override;
+
+    void search_and_reconstruct (idx_t n, const float *x, idx_t k,
+                                 float *distances, idx_t *labels,
+                                 float *recons) const override;
+
+    /// apply the transforms in the chain. The returned float * may be
+    /// equal to x, otherwise it should be deallocated.
+    const float * apply_chain (idx_t n, const float *x) const;
+
+    /// Reverse the transforms in the chain. May not be implemented for
+    /// all transforms in the chain or may return approximate results.
+    void reverse_chain (idx_t n, const float* xt, float* x) const;
+
+
+    /* standalone codec interface */
+    size_t sa_code_size () const override;
+    void sa_encode (idx_t n, const float *x,
+                          uint8_t *bytes) const override;
+    void sa_decode (idx_t n, const uint8_t *bytes,
+                            float *x) const override;
+
+    ~IndexPreTransform() override;
+};
+
+
+} // namespace faiss

data/vendor/faiss/IndexReplicas.cpp

@@ -0,0 +1,123 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/IndexReplicas.h>
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss {
+
+template <typename IndexT>
+IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(bool threaded)
+    : ThreadedIndex<IndexT>(threaded) {
+}
+
+template <typename IndexT>
+IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(idx_t d, bool threaded)
+    : ThreadedIndex<IndexT>(d, threaded) {
+}
+
+template <typename IndexT>
+IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(int d, bool threaded)
+    : ThreadedIndex<IndexT>(d, threaded) {
+}
+
+template <typename IndexT>
+void
+IndexReplicasTemplate<IndexT>::onAfterAddIndex(IndexT* index) {
+  // Make sure that the parameters are the same for all prior indices, unless
+  // we're the first index to be added
+  if (this->count() > 0 && this->at(0) != index) {
+    auto existing = this->at(0);
+
+    FAISS_THROW_IF_NOT_FMT(index->ntotal == existing->ntotal,
+                           "IndexReplicas: newly added index does "
+                           "not have same number of vectors as prior index; "
+                           "prior index has %ld vectors, new index has %ld",
+                           existing->ntotal, index->ntotal);
+
+    FAISS_THROW_IF_NOT_MSG(index->is_trained == existing->is_trained,
+                           "IndexReplicas: newly added index does "
+                           "not have same train status as prior index");
+  } else {
+    // Set our parameters based on the first index we're adding
+    // (dimension is handled in ThreadedIndex)
+    this->ntotal = index->ntotal;
+    this->verbose = index->verbose;
+    this->is_trained = index->is_trained;
+    this->metric_type = index->metric_type;
+  }
+}
+
+template <typename IndexT>
+void
+IndexReplicasTemplate<IndexT>::train(idx_t n, const component_t* x) {
+  this->runOnIndex([n, x](int, IndexT* index){ index->train(n, x); });
+}
+
+template <typename IndexT>
+void
+IndexReplicasTemplate<IndexT>::add(idx_t n, const component_t* x) {
+  this->runOnIndex([n, x](int, IndexT* index){ index->add(n, x); });
+  this->ntotal += n;
+}
+
+template <typename IndexT>
+void
+IndexReplicasTemplate<IndexT>::reconstruct(idx_t n, component_t* x) const {
+  FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
+
+  // Just pass to the first replica
+  this->at(0)->reconstruct(n, x);
+}
+
+template <typename IndexT>
+void
+IndexReplicasTemplate<IndexT>::search(idx_t n,
+                                      const component_t* x,
+                                      idx_t k,
+                                      distance_t* distances,
+                                      idx_t* labels) const {
+  FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
+
+  if (n == 0) {
+    return;
+  }
+
+  auto dim = this->d;
+  size_t componentsPerVec =
+    sizeof(component_t) == 1 ? (dim + 7) / 8 : dim;
+
+  // Partition the query by the number of indices we have
+  faiss::Index::idx_t queriesPerIndex =
+    (faiss::Index::idx_t) (n + this->count() - 1) /
+    (faiss::Index::idx_t) this->count();
+  FAISS_ASSERT(n / queriesPerIndex <= this->count());
+
+  auto fn =
+    [queriesPerIndex, componentsPerVec,
+     n, x, k, distances, labels](int i, const IndexT* index) {
+      faiss::Index::idx_t base = (faiss::Index::idx_t) i * queriesPerIndex;
+
+      if (base < n) {
+        auto numForIndex = std::min(queriesPerIndex, n - base);
+
+        index->search(numForIndex,
+                      x + base * componentsPerVec,
+                      k,
+                      distances + base * k,
+                      labels + base * k);
+      }
+    };
+
+  this->runOnIndex(fn);
+}
+
+// explicit instantiations
+template struct IndexReplicasTemplate<Index>;
+template struct IndexReplicasTemplate<IndexBinary>;
+
+} // namespace

data/vendor/faiss/IndexReplicas.h

@@ -0,0 +1,76 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/Index.h>
+#include <faiss/IndexBinary.h>
+#include <faiss/impl/ThreadedIndex.h>
+
+namespace faiss {
+
+/// Takes individual faiss::Index instances, and splits queries for
+/// sending to each Index instance, and joins the results together
+/// when done.
+/// Each index is managed by a separate CPU thread.
+template <typename IndexT>
+class IndexReplicasTemplate : public ThreadedIndex<IndexT> {
+ public:
+  using idx_t = typename IndexT::idx_t;
+  using component_t = typename IndexT::component_t;
+  using distance_t = typename IndexT::distance_t;
+
+  /// The dimension that all sub-indices must share will be the dimension of the
+  /// first sub-index added
+  /// @param threaded do we use one thread per sub-index or do queries
+  /// sequentially?
+  explicit IndexReplicasTemplate(bool threaded = true);
+
+  /// @param d the dimension that all sub-indices must share
+  /// @param threaded do we use one thread per sub index or do queries
+  /// sequentially?
+  explicit IndexReplicasTemplate(idx_t d, bool threaded = true);
+
+  /// int version due to the implicit bool conversion ambiguity of int as
+  /// dimension
+  explicit IndexReplicasTemplate(int d, bool threaded = true);
+
+  /// Alias for addIndex()
+  void add_replica(IndexT* index) { this->addIndex(index); }
+
+  /// Alias for removeIndex()
+  void remove_replica(IndexT* index) { this->removeIndex(index); }
+
+  /// faiss::Index API
+  /// All indices receive the same call
+  void train(idx_t n, const component_t* x) override;
+
+  /// faiss::Index API
+  /// All indices receive the same call
+  void add(idx_t n, const component_t* x) override;
+
+  /// faiss::Index API
+  /// Query is partitioned into a slice for each sub-index
+  /// split by ceil(n / #indices) for our sub-indices
+  void search(idx_t n,
+              const component_t* x,
+              idx_t k,
+              distance_t* distances,
+              idx_t* labels) const override;
+
+  /// reconstructs from the first index
+  void reconstruct(idx_t, component_t *v) const override;
+
+ protected:
+  /// Called just after an index is added
+  void onAfterAddIndex(IndexT* index) override;
+};
+
+using IndexReplicas = IndexReplicasTemplate<Index>;
+using IndexBinaryReplicas = IndexReplicasTemplate<IndexBinary>;
+
+} // namespace

data/vendor/faiss/IndexScalarQuantizer.cpp

@@ -0,0 +1,317 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IndexScalarQuantizer.h>
+
+#include <cstdio>
+#include <algorithm>
+
+#include <omp.h>
+
+#include <faiss/utils/utils.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/ScalarQuantizer.h>
+
+namespace faiss {
+
+
+
+/*******************************************************************
+ * IndexScalarQuantizer implementation
+ ********************************************************************/
+
+IndexScalarQuantizer::IndexScalarQuantizer
+                      (int d, ScalarQuantizer::QuantizerType qtype,
+                       MetricType metric):
+          Index(d, metric),
+          sq (d, qtype)
+{
+    is_trained =
+        qtype == ScalarQuantizer::QT_fp16 ||
+        qtype == ScalarQuantizer::QT_8bit_direct;
+    code_size = sq.code_size;
+}
+
+
+IndexScalarQuantizer::IndexScalarQuantizer ():
+    IndexScalarQuantizer(0, ScalarQuantizer::QT_8bit)
+{}
+
+void IndexScalarQuantizer::train(idx_t n, const float* x)
+{
+    sq.train(n, x);
+    is_trained = true;
+}
+
+void IndexScalarQuantizer::add(idx_t n, const float* x)
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    codes.resize ((n + ntotal) * code_size);
+    sq.compute_codes (x, &codes[ntotal * code_size], n);
+    ntotal += n;
+}
+
+
+void IndexScalarQuantizer::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    FAISS_THROW_IF_NOT (metric_type == METRIC_L2 ||
+                        metric_type == METRIC_INNER_PRODUCT);
+
+#pragma omp parallel
+    {
+        InvertedListScanner* scanner = sq.select_InvertedListScanner
+            (metric_type, nullptr, true);
+        ScopeDeleter1<InvertedListScanner> del(scanner);
+
+#pragma omp for
+        for (size_t i = 0; i < n; i++) {
+            float * D = distances + k * i;
+            idx_t * I = labels + k * i;
+            // re-order heap
+            if (metric_type == METRIC_L2) {
+                maxheap_heapify (k, D, I);
+            } else {
+                minheap_heapify (k, D, I);
+            }
+            scanner->set_query (x + i * d);
+            scanner->scan_codes (ntotal, codes.data(),
+                                 nullptr, D, I, k);
+
+            // re-order heap
+            if (metric_type == METRIC_L2) {
+                maxheap_reorder (k, D, I);
+            } else {
+                minheap_reorder (k, D, I);
+            }
+        }
+    }
+
+}
+
+
+DistanceComputer *IndexScalarQuantizer::get_distance_computer () const
+{
+    ScalarQuantizer::SQDistanceComputer *dc =
+        sq.get_distance_computer (metric_type);
+    dc->code_size = sq.code_size;
+    dc->codes = codes.data();
+    return dc;
+}
+
+
+void IndexScalarQuantizer::reset()
+{
+    codes.clear();
+    ntotal = 0;
+}
+
+void IndexScalarQuantizer::reconstruct_n(
+             idx_t i0, idx_t ni, float* recons) const
+{
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer ());
+    for (size_t i = 0; i < ni; i++) {
+        squant->decode_vector(&codes[(i + i0) * code_size], recons + i * d);
+    }
+}
+
+void IndexScalarQuantizer::reconstruct(idx_t key, float* recons) const
+{
+    reconstruct_n(key, 1, recons);
+}
+
+/* Codec interface */
+size_t IndexScalarQuantizer::sa_code_size () const
+{
+    return sq.code_size;
+}
+
+void IndexScalarQuantizer::sa_encode (idx_t n, const float *x,
+                      uint8_t *bytes) const
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    sq.compute_codes (x, bytes, n);
+}
+
+void IndexScalarQuantizer::sa_decode (idx_t n, const uint8_t *bytes,
+                                              float *x) const
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    sq.decode(bytes, x, n);
+}
+
+
+
+/*******************************************************************
+ * IndexIVFScalarQuantizer implementation
+ ********************************************************************/
+
+IndexIVFScalarQuantizer::IndexIVFScalarQuantizer (
+            Index *quantizer, size_t d, size_t nlist,
+            ScalarQuantizer::QuantizerType qtype,
+            MetricType metric, bool encode_residual)
+    : IndexIVF(quantizer, d, nlist, 0, metric),
+      sq(d, qtype),
+      by_residual(encode_residual)
+{
+    code_size = sq.code_size;
+    // was not known at construction time
+    invlists->code_size = code_size;
+    is_trained = false;
+}
+
+IndexIVFScalarQuantizer::IndexIVFScalarQuantizer ():
+    IndexIVF(),
+    by_residual(true)
+{
+}
+
+void IndexIVFScalarQuantizer::train_residual (idx_t n, const float *x)
+{
+    sq.train_residual(n, x, quantizer, by_residual, verbose);
+}
+
+void IndexIVFScalarQuantizer::encode_vectors(idx_t n, const float* x,
+                                             const idx_t *list_nos,
+                                             uint8_t * codes,
+                                             bool include_listnos) const
+{
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant (sq.select_quantizer ());
+    size_t coarse_size = include_listnos ? coarse_code_size () : 0;
+    memset(codes, 0, (code_size + coarse_size) * n);
+
+#pragma omp parallel if(n > 1)
+    {
+        std::vector<float> residual (d);
+
+#pragma omp for
+        for (size_t i = 0; i < n; i++) {
+            int64_t list_no = list_nos [i];
+            if (list_no >= 0) {
+                const float *xi = x + i * d;
+                uint8_t *code = codes + i * (code_size + coarse_size);
+                if (by_residual) {
+                    quantizer->compute_residual (
+                          xi, residual.data(), list_no);
+                    xi = residual.data ();
+                }
+                if (coarse_size) {
+                    encode_listno (list_no, code);
+                }
+                squant->encode_vector (xi, code + coarse_size);
+            }
+        }
+    }
+}
+
+void IndexIVFScalarQuantizer::sa_decode (idx_t n, const uint8_t *codes,
+                                                 float *x) const
+{
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant (sq.select_quantizer ());
+    size_t coarse_size = coarse_code_size ();
+
+#pragma omp parallel if(n > 1)
+    {
+        std::vector<float> residual (d);
+
+#pragma omp for
+        for (size_t i = 0; i < n; i++) {
+            const uint8_t *code = codes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno (code);
+            float *xi = x + i * d;
+            squant->decode_vector (code + coarse_size, xi);
+            if (by_residual) {
+                quantizer->reconstruct (list_no, residual.data());
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] += residual[j];
+                }
+            }
+        }
+    }
+}
+
+
+
+void IndexIVFScalarQuantizer::add_with_ids
+       (idx_t n, const float * x, const idx_t *xids)
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    std::unique_ptr<int64_t []> idx (new int64_t [n]);
+    quantizer->assign (n, x, idx.get());
+    size_t nadd = 0;
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer ());
+
+#pragma omp parallel reduction(+: nadd)
+    {
+        std::vector<float> residual (d);
+        std::vector<uint8_t> one_code (code_size);
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
+
+        // each thread takes care of a subset of lists
+        for (size_t i = 0; i < n; i++) {
+            int64_t list_no = idx [i];
+            if (list_no >= 0 && list_no % nt == rank) {
+                int64_t id = xids ? xids[i] : ntotal + i;
+
+                const float * xi = x + i * d;
+                if (by_residual) {
+                    quantizer->compute_residual (xi, residual.data(), list_no);
+                    xi = residual.data();
+                }
+
+                memset (one_code.data(), 0, code_size);
+                squant->encode_vector (xi, one_code.data());
+
+                invlists->add_entry (list_no, id, one_code.data());
+
+                nadd++;
+
+            }
+        }
+    }
+    ntotal += n;
+}
+
+
+
+
+
+InvertedListScanner* IndexIVFScalarQuantizer::get_InvertedListScanner
+    (bool store_pairs) const
+{
+    return sq.select_InvertedListScanner (metric_type, quantizer, store_pairs,
+                                          by_residual);
+}
+
+
+void IndexIVFScalarQuantizer::reconstruct_from_offset (int64_t list_no,
+                                                       int64_t offset,
+                                                       float* recons) const
+{
+    std::vector<float> centroid(d);
+    quantizer->reconstruct (list_no, centroid.data());
+
+    const uint8_t* code = invlists->get_single_code (list_no, offset);
+    sq.decode (code, recons, 1);
+    for (int i = 0; i < d; ++i) {
+        recons[i] += centroid[i];
+    }
+}
+
+
+
+
+} // namespace faiss