faiss 0.1.3 → 0.2.0

data/vendor/faiss/faiss/IndexScalarQuantizer.cpp

@@ -192,7 +192,7 @@ void IndexIVFScalarQuantizer::encode_vectors(idx_t n, const float* x,
     size_t coarse_size = include_listnos ? coarse_code_size () : 0;
     memset(codes, 0, (code_size + coarse_size) * n);
 
-#pragma omp parallel if(n > 1)
+#pragma omp parallel if(n > 1000)
     {
         std::vector<float> residual (d);
 
@@ -222,7 +222,7 @@ void IndexIVFScalarQuantizer::sa_decode (idx_t n, const uint8_t *codes,
     std::unique_ptr<ScalarQuantizer::Quantizer> squant (sq.select_quantizer ());
     size_t coarse_size = coarse_code_size ();
 
-#pragma omp parallel if(n > 1)
+#pragma omp parallel if(n > 1000)
     {
         std::vector<float> residual (d);
 

data/vendor/faiss/faiss/IndexScalarQuantizer.h

@@ -82,7 +82,7 @@ struct IndexScalarQuantizer: Index {
 
 
  /** An IVF implementation where the components of the residuals are
- * encoded with a scalar uniform quantizer. All distance computations
+ * encoded with a scalar quantizer. All distance computations
  * are asymmetric, so the encoded vectors are decoded and approximate
  * distances are computed.
  */

data/vendor/faiss/faiss/gpu/GpuDistance.h

@@ -148,6 +148,6 @@ void bruteForceKnn(GpuResourcesProvider* resources,
                    float* outDistances,
                    // A region of memory size numQueries x k, with k
                    // innermost (row major)
-                   faiss::Index::idx_t* outIndices);
+                   Index::idx_t* outIndices);
 
 } } // namespace

data/vendor/faiss/faiss/gpu/GpuIndex.h

@@ -36,9 +36,12 @@ class GpuIndex : public faiss::Index {
            float metricArg,
            GpuIndexConfig config);
 
-  inline int getDevice() const {
-    return device_;
-  }
+  /// Returns the device that this index is resident on
+  int getDevice() const;
+
+  /// Returns a reference to our GpuResources object that manages memory, stream
+  /// and handle resources on the GPU
+  std::shared_ptr<GpuResources> getResources();
 
   /// Set the minimum data size for searches (in MiB) for which we use
   /// CPU -> GPU paging
@@ -50,7 +53,7 @@ class GpuIndex : public faiss::Index {
   /// `x` can be resident on the CPU or any GPU; copies are performed
   /// as needed
   /// Handles paged adds if the add set is too large; calls addInternal_
-  void add(faiss::Index::idx_t, const float* x) override;
+  void add(Index::idx_t, const float* x) override;
 
   /// `x` and `ids` can be resident on the CPU or any GPU; copies are
   /// performed as needed
@@ -59,6 +62,13 @@ class GpuIndex : public faiss::Index {
                     const float* x,
                     const Index::idx_t* ids) override;
 
+  /// `x` and `labels` can be resident on the CPU or any GPU; copies are
+  /// performed as needed
+  void assign(Index::idx_t n,
+              const float* x,
+              Index::idx_t* labels,
+              Index::idx_t k = 1) const override;
+
   /// `x`, `distances` and `labels` can be resident on the CPU or any
   /// GPU; copies are performed as needed
   void search(Index::idx_t n,
@@ -136,11 +146,8 @@ private:
   /// Manages streams, cuBLAS handles and scratch memory for devices
   std::shared_ptr<GpuResources> resources_;
 
-  /// The GPU device we are resident on
-  const int device_;
-
-  /// The memory space of our primary storage on the GPU
-  const MemorySpace memorySpace_;
+  /// Our configuration options
+  const GpuIndexConfig config_;
 
   /// Size above which we page copies from the CPU to GPU
   size_t minPagedSize_;

data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h

@@ -38,6 +38,13 @@ class GpuIndexBinaryFlat : public IndexBinary {
 
   ~GpuIndexBinaryFlat() override;
 
+  /// Returns the device that this index is resident on
+  int getDevice() const;
+
+  /// Returns a reference to our GpuResources object that manages memory, stream
+  /// and handle resources on the GPU
+  std::shared_ptr<GpuResources> getResources();
+
   /// Initialize ourselves from the given CPU index; will overwrite
   /// all data in ourselves
   void copyFrom(const faiss::IndexBinaryFlat* index);
@@ -80,7 +87,7 @@ class GpuIndexBinaryFlat : public IndexBinary {
   std::shared_ptr<GpuResources> resources_;
 
   /// Configuration options
-  GpuIndexBinaryFlatConfig config_;
+  const GpuIndexBinaryFlatConfig binaryFlatConfig_;
 
   /// Holds our GPU data containing the list of vectors
   std::unique_ptr<BinaryFlatIndex> data_;

data/vendor/faiss/faiss/gpu/GpuIndexFlat.h

@@ -21,7 +21,7 @@ struct IndexFlatIP;
 
 namespace faiss { namespace gpu {
 
-struct FlatIndex;
+class FlatIndex;
 
 struct GpuIndexFlatConfig : public GpuIndexConfig {
   inline GpuIndexFlatConfig()
@@ -87,27 +87,27 @@ class GpuIndexFlat : public GpuIndex {
   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;
+  void add(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;
+  void reconstruct(Index::idx_t key, float* out) const override;
 
   /// Batch reconstruction method
-  void reconstruct_n(faiss::Index::idx_t i0,
-                     faiss::Index::idx_t num,
+  void reconstruct_n(Index::idx_t i0,
+                     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;
+                        Index::idx_t key) const override;
 
   /// Compute residual (batch mode)
-  void compute_residual_n(faiss::Index::idx_t n,
+  void compute_residual_n(Index::idx_t n,
                           const float* xs,
                           float* residuals,
-                          const faiss::Index::idx_t* keys) const override;
+                          const Index::idx_t* keys) const override;
 
   /// For internal access
   inline FlatIndex* getGpuData() { return data_.get(); }
@@ -126,11 +126,11 @@ class GpuIndexFlat : public GpuIndex {
                    const float* x,
                    int k,
                    float* distances,
-                   faiss::Index::idx_t* labels) const override;
+                   Index::idx_t* labels) const override;
 
  protected:
-  /// Our config object
-  const GpuIndexFlatConfig config_;
+  /// Our configuration options
+  const GpuIndexFlatConfig flatConfig_;
 
   /// Holds our GPU data containing the list of vectors
   std::unique_ptr<FlatIndex> data_;

data/vendor/faiss/faiss/gpu/GpuIndexIVF.h

@@ -56,6 +56,22 @@ class GpuIndexIVF : public GpuIndex {
   /// Returns the number of inverted lists we're managing
   int getNumLists() const;
 
+  /// Returns the number of vectors present in a particular inverted list
+  virtual int getListLength(int listId) const = 0;
+
+  /// Return the encoded vector data contained in a particular inverted list,
+  /// for debugging purposes.
+  /// If gpuFormat is true, the data is returned as it is encoded in the
+  /// GPU-side representation.
+  /// Otherwise, it is converted to the CPU format.
+  /// compliant format, while the native GPU format may differ.
+  virtual std::vector<uint8_t>
+  getListVectorData(int listId, bool gpuFormat = false) const = 0;
+
+  /// Return the vector indices contained in a particular inverted list, for
+  /// debugging purposes.
+  virtual std::vector<Index::idx_t> getListIndices(int listId) const = 0;
+
   /// Return the quantizer we're using
   GpuIndexFlat* getQuantizer();
 
@@ -67,7 +83,7 @@ class GpuIndexIVF : public GpuIndex {
 
  protected:
   bool addImplRequiresIDs_() const override;
-  void trainQuantizer_(faiss::Index::idx_t n, const float* x);
+  void trainQuantizer_(Index::idx_t n, const float* x);
 
  public:
   /// Exposing this like the CPU version for manipulation
@@ -83,7 +99,8 @@ class GpuIndexIVF : public GpuIndex {
   GpuIndexFlat* quantizer;
 
  protected:
-  GpuIndexIVFConfig ivfConfig_;
+  /// Our configuration options
+  const GpuIndexIVFConfig ivfConfig_;
 };
 
 } } // namespace

data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h

@@ -19,6 +19,13 @@ class IVFFlat;
 class GpuIndexFlat;
 
 struct GpuIndexIVFFlatConfig : public GpuIndexIVFConfig {
+  inline GpuIndexIVFFlatConfig()
+      : interleavedLayout(true) {
+  }
+
+  /// Use the alternative memory layout for the IVF lists
+  /// (currently the default)
+  bool interleavedLayout;
 };
 
 /// Wrapper around the GPU implementation that looks like
@@ -56,10 +63,28 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
   /// to exactly the amount needed. Returns space reclaimed in bytes
   size_t reclaimMemory();
 
+  /// Clears out all inverted lists, but retains the coarse centroid information
   void reset() override;
 
+  /// Trains the coarse quantizer based on the given vector data
   void train(Index::idx_t n, const float* x) override;
 
+  /// Returns the number of vectors present in a particular inverted list
+  int getListLength(int listId) const override;
+
+  /// Return the encoded vector data contained in a particular inverted list,
+  /// for debugging purposes.
+  /// If gpuFormat is true, the data is returned as it is encoded in the
+  /// GPU-side representation.
+  /// Otherwise, it is converted to the CPU format.
+  /// compliant format, while the native GPU format may differ.
+  std::vector<uint8_t>
+  getListVectorData(int listId, bool gpuFormat = false) const override;
+
+  /// Return the vector indices contained in a particular inverted list, for
+  /// debugging purposes.
+  std::vector<Index::idx_t> getListIndices(int listId) const override;
+
  protected:
   /// Called from GpuIndex for add/add_with_ids
   void addImpl_(int n,
@@ -73,8 +98,9 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
                    float* distances,
                    Index::idx_t* labels) const override;
 
- private:
-  GpuIndexIVFFlatConfig ivfFlatConfig_;
+ protected:
+  /// Our configuration options
+  const GpuIndexIVFFlatConfig ivfFlatConfig_;
 
   /// Desired inverted list memory reservation
   size_t reserveMemoryVecs_;

data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h

@@ -23,7 +23,7 @@ struct GpuIndexIVFPQConfig : public GpuIndexIVFConfig {
   inline GpuIndexIVFPQConfig()
       : useFloat16LookupTables(false),
         usePrecomputedTables(false),
-        alternativeLayout(false),
+        interleavedLayout(false),
         useMMCodeDistance(false) {
   }
 
@@ -38,7 +38,7 @@ struct GpuIndexIVFPQConfig : public GpuIndexIVFConfig {
 
   /// Use the alternative memory layout for the IVF lists
   /// WARNING: this is a feature under development, do not use!
-  bool alternativeLayout;
+  bool interleavedLayout;
 
   /// Use GEMM-backed computation of PQ code distances for the no precomputed
   /// table version of IVFPQ.
@@ -108,19 +108,24 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
   /// product centroid information
   void reset() override;
 
+  /// Trains the coarse and product quantizer based on the given vector data
   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;
+  /// Returns the number of vectors present in a particular inverted list
+  int getListLength(int listId) const override;
 
-  /// For debugging purposes, return the list codes of a particular
-  /// list
-  std::vector<unsigned char> getListCodes(int listId) const;
+  /// Return the encoded vector data contained in a particular inverted list,
+  /// for debugging purposes.
+  /// If gpuFormat is true, the data is returned as it is encoded in the
+  /// GPU-side representation.
+  /// Otherwise, it is converted to the CPU format.
+  /// compliant format, while the native GPU format may differ.
+  std::vector<uint8_t>
+  getListVectorData(int listId, bool gpuFormat = false) const override;
 
-  /// For debugging purposes, return the list indices of a particular
-  /// list
-  std::vector<long> getListIndices(int listId) const;
+  /// Return the vector indices contained in a particular inverted list, for
+  /// debugging purposes.
+  std::vector<Index::idx_t> getListIndices(int listId) const override;
 
  protected:
   /// Called from GpuIndex for add/add_with_ids
@@ -135,13 +140,18 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
                    float* distances,
                    Index::idx_t* labels) const override;
 
- private:
+  /// Throws errors if configuration settings are improper
   void verifySettings_() const;
 
+  /// Trains the PQ quantizer based on the given vector data
   void trainResidualQuantizer_(Index::idx_t n, const float* x);
 
- private:
-  GpuIndexIVFPQConfig ivfpqConfig_;
+ protected:
+  /// Our configuration options that we were initialized with
+  const GpuIndexIVFPQConfig ivfpqConfig_;
+
+  /// Runtime override: whether or not we use precomputed tables
+  bool usePrecomputedTables_;
 
   /// Number of sub-quantizers per encoded vector
   int subQuantizers_;

data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h

@@ -18,6 +18,13 @@ class IVFFlat;
 class GpuIndexFlat;
 
 struct GpuIndexIVFScalarQuantizerConfig : public GpuIndexIVFConfig {
+  inline GpuIndexIVFScalarQuantizerConfig()
+      : interleavedLayout(true) {
+  }
+
+  /// Use the alternative memory layout for the IVF lists
+  /// (currently the default)
+  bool interleavedLayout;
 };
 
 /// Wrapper around the GPU implementation that looks like
@@ -61,10 +68,29 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
   /// to exactly the amount needed. Returns space reclaimed in bytes
   size_t reclaimMemory();
 
+  /// Clears out all inverted lists, but retains the coarse and scalar quantizer
+  /// information
   void reset() override;
 
+  /// Trains the coarse and scalar quantizer based on the given vector data
   void train(Index::idx_t n, const float* x) override;
 
+  /// Returns the number of vectors present in a particular inverted list
+  int getListLength(int listId) const override;
+
+  /// Return the encoded vector data contained in a particular inverted list,
+  /// for debugging purposes.
+  /// If gpuFormat is true, the data is returned as it is encoded in the
+  /// GPU-side representation.
+  /// Otherwise, it is converted to the CPU format.
+  /// compliant format, while the native GPU format may differ.
+  std::vector<uint8_t>
+  getListVectorData(int listId, bool gpuFormat = false) const override;
+
+  /// Return the vector indices contained in a particular inverted list, for
+  /// debugging purposes.
+  std::vector<Index::idx_t> getListIndices(int listId) const override;
+
  protected:
   /// Called from GpuIndex for add/add_with_ids
   void addImpl_(int n,
@@ -88,8 +114,9 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
   /// Exposed like the CPU version
   bool by_residual;
 
- private:
-  GpuIndexIVFScalarQuantizerConfig ivfSQConfig_;
+protected:
+  /// Our configuration options
+  const GpuIndexIVFScalarQuantizerConfig ivfSQConfig_;
 
   /// Desired inverted list memory reservation
   size_t reserveMemoryVecs_;

data/vendor/faiss/faiss/gpu/GpuResources.h

@@ -198,6 +198,10 @@ class GpuResources {
   /// given device
   virtual cudaStream_t getDefaultStream(int device) = 0;
 
+  /// Overrides the default stream for a device to the user-supplied stream. The
+  /// resources object does not own this stream (i.e., it will not destroy it).
+  virtual void setDefaultStream(int device, cudaStream_t stream) = 0;
+
   /// Returns the set of alternative streams that we use for the given device
   virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0;
 

data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp

@@ -101,12 +101,8 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
   for (auto& entry : defaultStreams_) {
     DeviceScope scope(entry.first);
 
-    auto it = userDefaultStreams_.find(entry.first);
-    if (it == userDefaultStreams_.end()) {
-      // The user did not specify this stream, thus we are the ones
-      // who have created it
-      CUDA_VERIFY(cudaStreamDestroy(entry.second));
-    }
+    // We created these streams, so are responsible for destroying them
+    CUDA_VERIFY(cudaStreamDestroy(entry.second));
   }
 
   for (auto& entry : alternateStreams_) {
@@ -210,16 +206,47 @@ StandardGpuResourcesImpl::setPinnedMemory(size_t size) {
 
 void
 StandardGpuResourcesImpl::setDefaultStream(int device, cudaStream_t stream) {
-  auto it = defaultStreams_.find(device);
-  if (it != defaultStreams_.end()) {
-    // Replace this stream with the user stream
-    CUDA_VERIFY(cudaStreamDestroy(it->second));
-    it->second = stream;
+  if (isInitialized(device)) {
+     // A new series of calls may not be ordered with what was the previous
+     // stream, so if the stream being specified is different, then we need to
+     // ensure ordering between the two (new stream waits on old).
+    auto it = userDefaultStreams_.find(device);
+    cudaStream_t prevStream = nullptr;
+
+    if (it != userDefaultStreams_.end()) {
+      prevStream = it->second;
+    } else {
+      FAISS_ASSERT(defaultStreams_.count(device));
+      prevStream = defaultStreams_[device];
+    }
+
+    if (prevStream != stream) {
+      streamWait({stream}, {prevStream});
+    }
   }
 
   userDefaultStreams_[device] = stream;
 }
 
+void
+StandardGpuResourcesImpl::revertDefaultStream(int device) {
+  if (isInitialized(device)) {
+    auto it = userDefaultStreams_.find(device);
+
+    if (it != userDefaultStreams_.end()) {
+      // There was a user stream set that we need to synchronize against
+      cudaStream_t prevStream = userDefaultStreams_[device];
+
+      FAISS_ASSERT(defaultStreams_.count(device));
+      cudaStream_t newStream = defaultStreams_[device];
+
+      streamWait({newStream}, {prevStream});
+    }
+  }
+
+  userDefaultStreams_.erase(device);
+}
+
 void
 StandardGpuResourcesImpl::setDefaultNullStreamAllDevices() {
   for (int dev = 0; dev < getNumDevices(); ++dev) {
@@ -274,14 +301,8 @@ StandardGpuResourcesImpl::initializeForDevice(int device) {
 
   // Create streams
   cudaStream_t defaultStream = 0;
-  auto it = userDefaultStreams_.find(device);
-  if (it != userDefaultStreams_.end()) {
-    // We already have a stream provided by the user
-    defaultStream = it->second;
-  } else {
-    CUDA_VERIFY(cudaStreamCreateWithFlags(&defaultStream,
-                                          cudaStreamNonBlocking));
-  }
+  CUDA_VERIFY(cudaStreamCreateWithFlags(&defaultStream,
+                                        cudaStreamNonBlocking));
 
   defaultStreams_[device] = defaultStream;
 
@@ -308,15 +329,14 @@ StandardGpuResourcesImpl::initializeForDevice(int device) {
   FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
   blasHandles_[device] = blasHandle;
 
-  // Enable tensor core support if available
-#if CUDA_VERSION >= 9000 && CUDA_VERSION < 11000
-  // This flag was deprecated in CUDA 11
-  if (getTensorCoreSupport(device)) {
-    cublasSetMathMode(blasHandle, CUBLAS_TENSOR_OP_MATH);
-  }
-#endif
+  // For CUDA 10 on V100, enabling tensor core usage would enable automatic
+  // rounding down of inputs to f16 (though accumulate in f32) which results in
+  // unacceptable loss of precision in general.
+  // For CUDA 11 / A100, only enable tensor core support if it doesn't result in
+  // a loss of precision.
 #if CUDA_VERSION >= 11000
-  cublasSetMathMode(blasHandle, CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION);
+  cublasSetMathMode(blasHandle,
+                    CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION);
 #endif
 
   FAISS_ASSERT(allocs_.count(device) == 0);
@@ -341,6 +361,14 @@ StandardGpuResourcesImpl::getBlasHandle(int device) {
 cudaStream_t
 StandardGpuResourcesImpl::getDefaultStream(int device) {
   initializeForDevice(device);
+
+  auto it = userDefaultStreams_.find(device);
+  if (it != userDefaultStreams_.end()) {
+    // There is a user override stream set
+    return it->second;
+  }
+
+  // Otherwise, our base default stream
   return defaultStreams_[device];
 }
 
@@ -539,6 +567,11 @@ StandardGpuResources::setDefaultStream(int device, cudaStream_t stream) {
   res_->setDefaultStream(device, stream);
 }
 
+void
+StandardGpuResources::revertDefaultStream(int device) {
+  res_->revertDefaultStream(device);
+}
+
 void
 StandardGpuResources::setDefaultNullStreamAllDevices() {
   res_->setDefaultNullStreamAllDevices();