faiss 0.1.0 → 0.1.1

data/vendor/faiss/gpu/GpuIndexFlat.h

@@ -0,0 +1,190 @@
+/**
+ * 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/gpu/GpuIndex.h>
+
+namespace faiss {
+
+struct IndexFlat;
+struct IndexFlatL2;
+struct IndexFlatIP;
+
+}
+
+namespace faiss { namespace gpu {
+
+struct FlatIndex;
+
+struct GpuIndexFlatConfig : public GpuIndexConfig {
+  inline GpuIndexFlatConfig()
+      : useFloat16(false),
+        useFloat16Accumulator(false),
+        storeTransposed(false) {
+  }
+
+  /// Whether or not data is stored as float16
+  bool useFloat16;
+
+  /// Whether or not all math is performed in float16, if useFloat16 is
+  /// specified. If true, we use cublasHgemm, supported only on CC
+  /// 5.3+. Otherwise, we use cublasSgemmEx.
+  bool useFloat16Accumulator;
+
+  /// Whether or not data is stored (transparently) in a transposed
+  /// layout, enabling use of the NN GEMM call, which is ~10% faster.
+  /// This will improve the speed of the flat index, but will
+  /// substantially slow down any add() calls made, as all data must
+  /// be transposed, and will increase storage requirements (we store
+  /// data in both transposed and non-transposed layouts).
+  bool storeTransposed;
+};
+
+/// Wrapper around the GPU implementation that looks like
+/// faiss::IndexFlat; copies over centroid data from a given
+/// faiss::IndexFlat
+class GpuIndexFlat : public GpuIndex {
+ public:
+  /// Construct from a pre-existing faiss::IndexFlat instance, copying
+  /// data over to the given GPU
+  GpuIndexFlat(GpuResources* resources,
+               const faiss::IndexFlat* index,
+               GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  /// Construct an empty instance that can be added to
+  GpuIndexFlat(GpuResources* resources,
+               int dims,
+               faiss::MetricType metric,
+               GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  ~GpuIndexFlat() override;
+
+  /// Initialize ourselves from the given CPU index; will overwrite
+  /// all data in ourselves
+  void copyFrom(const faiss::IndexFlat* index);
+
+  /// Copy ourselves to the given CPU index; will overwrite all data
+  /// in the index instance
+  void copyTo(faiss::IndexFlat* index) const;
+
+  /// Returns the number of vectors we contain
+  size_t getNumVecs() const;
+
+  /// Clears all vectors from this index
+  void reset() override;
+
+  /// This index is not trained, so this does nothing
+  void train(Index::idx_t n, const float* x) override;
+
+  /// Overrides to avoid excessive copies
+  void add(faiss::Index::idx_t, const float* x) override;
+
+  /// Reconstruction methods; prefer the batch reconstruct as it will
+  /// be more efficient
+  void reconstruct(faiss::Index::idx_t key, float* out) const override;
+
+  /// Batch reconstruction method
+  void reconstruct_n(faiss::Index::idx_t i0,
+                     faiss::Index::idx_t num,
+                     float* out) const override;
+
+  /// Compute residual
+  void compute_residual(const float* x,
+                        float* residual,
+                        faiss::Index::idx_t key) const override;
+
+  /// Compute residual (batch mode)
+  void compute_residual_n(faiss::Index::idx_t n,
+                          const float* xs,
+                          float* residuals,
+                          const faiss::Index::idx_t* keys) const override;
+
+  /// For internal access
+  inline FlatIndex* getGpuData() { return data_; }
+
+ protected:
+  /// Flat index does not require IDs as there is no storage available for them
+  bool addImplRequiresIDs_() const override;
+
+  /// Called from GpuIndex for add
+  void addImpl_(int n,
+                const float* x,
+                const Index::idx_t* ids) override;
+
+  /// Called from GpuIndex for search
+  void searchImpl_(int n,
+                   const float* x,
+                   int k,
+                   float* distances,
+                   faiss::Index::idx_t* labels) const override;
+
+ private:
+  /// Checks user settings for consistency
+  void verifySettings_() const;
+
+ protected:
+  /// Our config object
+  const GpuIndexFlatConfig config_;
+
+  /// Holds our GPU data containing the list of vectors; is managed via raw
+  /// pointer so as to allow non-CUDA compilers to see this header
+  FlatIndex* data_;
+};
+
+/// Wrapper around the GPU implementation that looks like
+/// faiss::IndexFlatL2; copies over centroid data from a given
+/// faiss::IndexFlat
+class GpuIndexFlatL2 : public GpuIndexFlat {
+ public:
+  /// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
+  /// data over to the given GPU
+  GpuIndexFlatL2(GpuResources* resources,
+                 faiss::IndexFlatL2* index,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  /// Construct an empty instance that can be added to
+  GpuIndexFlatL2(GpuResources* resources,
+                 int dims,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  /// Initialize ourselves from the given CPU index; will overwrite
+  /// all data in ourselves
+  void copyFrom(faiss::IndexFlat* index);
+
+  /// Copy ourselves to the given CPU index; will overwrite all data
+  /// in the index instance
+  void copyTo(faiss::IndexFlat* index);
+};
+
+/// Wrapper around the GPU implementation that looks like
+/// faiss::IndexFlatIP; copies over centroid data from a given
+/// faiss::IndexFlat
+class GpuIndexFlatIP : public GpuIndexFlat {
+ public:
+  /// Construct from a pre-existing faiss::IndexFlatIP instance, copying
+  /// data over to the given GPU
+  GpuIndexFlatIP(GpuResources* resources,
+                 faiss::IndexFlatIP* index,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  /// Construct an empty instance that can be added to
+  GpuIndexFlatIP(GpuResources* resources,
+                 int dims,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  /// Initialize ourselves from the given CPU index; will overwrite
+  /// all data in ourselves
+  void copyFrom(faiss::IndexFlat* index);
+
+  /// Copy ourselves to the given CPU index; will overwrite all data
+  /// in the index instance
+  void copyTo(faiss::IndexFlat* index);
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuIndexIVF.h

@@ -0,0 +1,89 @@
+/**
+ * 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/gpu/GpuIndex.h>
+#include <faiss/gpu/GpuIndexFlat.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
+#include <faiss/Clustering.h>
+
+namespace faiss { struct IndexIVF; }
+
+namespace faiss { namespace gpu {
+
+class GpuIndexFlat;
+class GpuResources;
+
+struct GpuIndexIVFConfig : public GpuIndexConfig {
+  inline GpuIndexIVFConfig()
+      : indicesOptions(INDICES_64_BIT) {
+  }
+
+  /// Index storage options for the GPU
+  IndicesOptions indicesOptions;
+
+  /// Configuration for the coarse quantizer object
+  GpuIndexFlatConfig flatConfig;
+};
+
+class GpuIndexIVF : public GpuIndex {
+ public:
+  GpuIndexIVF(GpuResources* resources,
+              int dims,
+              faiss::MetricType metric,
+              int nlist,
+              GpuIndexIVFConfig config = GpuIndexIVFConfig());
+
+  ~GpuIndexIVF() override;
+
+ private:
+  /// Shared initialization functions
+  void init_();
+
+ public:
+  /// Copy what we need from the CPU equivalent
+  void copyFrom(const faiss::IndexIVF* index);
+
+  /// Copy what we have to the CPU equivalent
+  void copyTo(faiss::IndexIVF* index) const;
+
+  /// Returns the number of inverted lists we're managing
+  int getNumLists() const;
+
+  /// Return the quantizer we're using
+  GpuIndexFlat* getQuantizer();
+
+  /// Sets the number of list probes per query
+  void setNumProbes(int nprobe);
+
+  /// Returns our current number of list probes per query
+  int getNumProbes() const;
+
+ protected:
+  bool addImplRequiresIDs_() const override;
+  void trainQuantizer_(faiss::Index::idx_t n, const float* x);
+
+ public:
+  /// Exposing this like the CPU version for manipulation
+  ClusteringParameters cp;
+
+  /// Exposing this like the CPU version for query
+  int nlist;
+
+  /// Exposing this like the CPU version for manipulation
+  int nprobe;
+
+  /// Exposeing this like the CPU version for query
+  GpuIndexFlat* quantizer;
+
+ protected:
+  GpuIndexIVFConfig ivfConfig_;
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuIndexIVFFlat.h

@@ -0,0 +1,85 @@
+/**
+ * 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/gpu/GpuIndexIVF.h>
+
+namespace faiss { struct IndexIVFFlat; }
+
+namespace faiss { namespace gpu {
+
+class IVFFlat;
+class GpuIndexFlat;
+
+struct GpuIndexIVFFlatConfig : public GpuIndexIVFConfig {
+};
+
+/// Wrapper around the GPU implementation that looks like
+/// faiss::IndexIVFFlat
+class GpuIndexIVFFlat : public GpuIndexIVF {
+ public:
+  /// Construct from a pre-existing faiss::IndexIVFFlat instance, copying
+  /// data over to the given GPU, if the input index is trained.
+  GpuIndexIVFFlat(GpuResources* resources,
+                  const faiss::IndexIVFFlat* index,
+                  GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
+
+  /// Constructs a new instance with an empty flat quantizer; the user
+  /// provides the number of lists desired.
+  GpuIndexIVFFlat(GpuResources* resources,
+                  int dims,
+                  int nlist,
+                  faiss::MetricType metric,
+                  GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
+
+  ~GpuIndexIVFFlat() override;
+
+  /// Reserve GPU memory in our inverted lists for this number of vectors
+  void reserveMemory(size_t numVecs);
+
+  /// Initialize ourselves from the given CPU index; will overwrite
+  /// all data in ourselves
+  void copyFrom(const faiss::IndexIVFFlat* index);
+
+  /// Copy ourselves to the given CPU index; will overwrite all data
+  /// in the index instance
+  void copyTo(faiss::IndexIVFFlat* index) const;
+
+  /// After adding vectors, one can call this to reclaim device memory
+  /// to exactly the amount needed. Returns space reclaimed in bytes
+  size_t reclaimMemory();
+
+  void reset() override;
+
+  void train(Index::idx_t n, const float* x) override;
+
+ protected:
+  /// Called from GpuIndex for add/add_with_ids
+  void addImpl_(int n,
+                const float* x,
+                const Index::idx_t* ids) override;
+
+  /// Called from GpuIndex for search
+  void searchImpl_(int n,
+                   const float* x,
+                   int k,
+                   float* distances,
+                   Index::idx_t* labels) const override;
+
+ private:
+  GpuIndexIVFFlatConfig ivfFlatConfig_;
+
+  /// Desired inverted list memory reservation
+  size_t reserveMemoryVecs_;
+
+  /// Instance that we own; contains the inverted list
+  IVFFlat* index_;
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuIndexIVFPQ.h

@@ -0,0 +1,143 @@
+/**
+ * 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/gpu/GpuIndexIVF.h>
+#include <vector>
+
+namespace faiss { struct IndexIVFPQ; }
+
+namespace faiss { namespace gpu {
+
+class GpuIndexFlat;
+class IVFPQ;
+
+struct GpuIndexIVFPQConfig : public GpuIndexIVFConfig {
+  inline GpuIndexIVFPQConfig()
+      : useFloat16LookupTables(false),
+        usePrecomputedTables(false) {
+  }
+
+  /// Whether or not float16 residual distance tables are used in the
+  /// list scanning kernels. When subQuantizers * 2^bitsPerCode >
+  /// 16384, this is required.
+  bool useFloat16LookupTables;
+
+  /// Whether or not we enable the precomputed table option for
+  /// search, which can substantially increase the memory requirement.
+  bool usePrecomputedTables;
+};
+
+/// IVFPQ index for the GPU
+class GpuIndexIVFPQ : public GpuIndexIVF {
+ public:
+  /// Construct from a pre-existing faiss::IndexIVFPQ instance, copying
+  /// data over to the given GPU, if the input index is trained.
+  GpuIndexIVFPQ(GpuResources* resources,
+                const faiss::IndexIVFPQ* index,
+                GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
+
+  /// Construct an empty index
+  GpuIndexIVFPQ(GpuResources* resources,
+                int dims,
+                int nlist,
+                int subQuantizers,
+                int bitsPerCode,
+                faiss::MetricType metric,
+                GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
+
+  ~GpuIndexIVFPQ() override;
+
+  /// Reserve space on the GPU for the inverted lists for `num`
+  /// vectors, assumed equally distributed among
+
+  /// Initialize ourselves from the given CPU index; will overwrite
+  /// all data in ourselves
+  void copyFrom(const faiss::IndexIVFPQ* index);
+
+  /// Copy ourselves to the given CPU index; will overwrite all data
+  /// in the index instance
+  void copyTo(faiss::IndexIVFPQ* index) const;
+
+  /// Reserve GPU memory in our inverted lists for this number of vectors
+  void reserveMemory(size_t numVecs);
+
+  /// Enable or disable pre-computed codes
+  void setPrecomputedCodes(bool enable);
+
+  /// Are pre-computed codes enabled?
+  bool getPrecomputedCodes() const;
+
+  /// Return the number of sub-quantizers we are using
+  int getNumSubQuantizers() const;
+
+  /// Return the number of bits per PQ code
+  int getBitsPerCode() const;
+
+  /// Return the number of centroids per PQ code (2^bits per code)
+  int getCentroidsPerSubQuantizer() const;
+
+  /// After adding vectors, one can call this to reclaim device memory
+  /// to exactly the amount needed. Returns space reclaimed in bytes
+  size_t reclaimMemory();
+
+  /// Clears out all inverted lists, but retains the coarse and
+  /// product centroid information
+  void reset() override;
+
+  void train(Index::idx_t n, const float* x) override;
+
+  /// For debugging purposes, return the list length of a particular
+  /// list
+  int getListLength(int listId) const;
+
+  /// For debugging purposes, return the list codes of a particular
+  /// list
+  std::vector<unsigned char> getListCodes(int listId) const;
+
+  /// For debugging purposes, return the list indices of a particular
+  /// list
+  std::vector<long> getListIndices(int listId) const;
+
+ protected:
+  /// Called from GpuIndex for add/add_with_ids
+  void addImpl_(int n,
+                const float* x,
+                const Index::idx_t* ids) override;
+
+  /// Called from GpuIndex for search
+  void searchImpl_(int n,
+                   const float* x,
+                   int k,
+                   float* distances,
+                   Index::idx_t* labels) const override;
+
+ private:
+  void verifySettings_() const;
+
+  void trainResidualQuantizer_(Index::idx_t n, const float* x);
+
+ private:
+  GpuIndexIVFPQConfig ivfpqConfig_;
+
+  /// Number of sub-quantizers per encoded vector
+  int subQuantizers_;
+
+  /// Bits per sub-quantizer code
+  int bitsPerCode_;
+
+  /// Desired inverted list memory reservation
+  size_t reserveMemoryVecs_;
+
+  /// The product quantizer instance that we own; contains the
+  /// inverted lists
+  IVFPQ* index_;
+};
+
+} } // namespace