faiss 0.1.7 → 0.2.3

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

@@ -5,152 +5,149 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
 #include <faiss/gpu/GpuResources.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 struct GpuIndexConfig {
-  inline GpuIndexConfig()
-      : device(0),
-        memorySpace(MemorySpace::Device) {
-  }
-
-  /// GPU device on which the index is resident
-  int device;
-
-  /// What memory space to use for primary storage.
-  /// On Pascal and above (CC 6+) architectures, allows GPUs to use
-  /// more memory than is available on the GPU.
-  MemorySpace memorySpace;
+    inline GpuIndexConfig() : device(0), memorySpace(MemorySpace::Device) {}
+
+    /// GPU device on which the index is resident
+    int device;
+
+    /// What memory space to use for primary storage.
+    /// On Pascal and above (CC 6+) architectures, allows GPUs to use
+    /// more memory than is available on the GPU.
+    MemorySpace memorySpace;
 };
 
 class GpuIndex : public faiss::Index {
- public:
-  GpuIndex(std::shared_ptr<GpuResources> resources,
-           int dims,
-           faiss::MetricType metric,
-           float metricArg,
-           GpuIndexConfig config);
-
-  /// 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
-  void setMinPagingSize(size_t size);
-
-  /// Returns the current minimum data size for paged searches
-  size_t getMinPagingSize() const;
-
-  /// `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(Index::idx_t, const float* x) override;
-
-  /// `x` and `ids` 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_with_ids(Index::idx_t n,
-                    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,
-              const float* x,
-              Index::idx_t k,
-              float* distances,
-              Index::idx_t* labels) const override;
-
-  /// Overridden to force GPU indices to provide their own GPU-friendly
-  /// implementation
-  void compute_residual(const float* x,
-                        float* residual,
-                        Index::idx_t key) const override;
-
-  /// Overridden to force GPU indices to provide their own GPU-friendly
-  /// implementation
-  void compute_residual_n(Index::idx_t n,
-                          const float* xs,
-                          float* residuals,
-                          const Index::idx_t* keys) const override;
-
- protected:
-  /// Copy what we need from the CPU equivalent
-  void copyFrom(const faiss::Index* index);
-
-  /// Copy what we have to the CPU equivalent
-  void copyTo(faiss::Index* index) const;
-
-  /// Does addImpl_ require IDs? If so, and no IDs are provided, we will
-  /// generate them sequentially based on the order in which the IDs are added
-  virtual bool addImplRequiresIDs_() const = 0;
-
-  /// Overridden to actually perform the add
-  /// All data is guaranteed to be resident on our device
-  virtual void addImpl_(int n,
-                        const float* x,
-                        const Index::idx_t* ids) = 0;
-
-  /// Overridden to actually perform the search
-  /// All data is guaranteed to be resident on our device
-  virtual void searchImpl_(int n,
-                           const float* x,
-                           int k,
-                           float* distances,
-                           Index::idx_t* labels) const = 0;
-
-private:
-  /// Handles paged adds if the add set is too large, passes to
-  /// addImpl_ to actually perform the add for the current page
-  void addPaged_(int n,
-                 const float* x,
-                 const Index::idx_t* ids);
-
-  /// Calls addImpl_ for a single page of GPU-resident data
-  void addPage_(int n,
-                const float* x,
-                const Index::idx_t* ids);
-
-  /// Calls searchImpl_ for a single page of GPU-resident data
-  void searchNonPaged_(int n,
-                       const float* x,
-                       int k,
-                       float* outDistancesData,
-                       Index::idx_t* outIndicesData) const;
-
-  /// Calls searchImpl_ for a single page of GPU-resident data,
-  /// handling paging of the data and copies from the CPU
-  void searchFromCpuPaged_(int n,
-                           const float* x,
-                           int k,
-                           float* outDistancesData,
-                           Index::idx_t* outIndicesData) const;
-
- protected:
-  /// Manages streams, cuBLAS handles and scratch memory for devices
-  std::shared_ptr<GpuResources> resources_;
-
-  /// Our configuration options
-  const GpuIndexConfig config_;
-
-  /// Size above which we page copies from the CPU to GPU
-  size_t minPagedSize_;
+   public:
+    GpuIndex(
+            std::shared_ptr<GpuResources> resources,
+            int dims,
+            faiss::MetricType metric,
+            float metricArg,
+            GpuIndexConfig config);
+
+    /// 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
+    void setMinPagingSize(size_t size);
+
+    /// Returns the current minimum data size for paged searches
+    size_t getMinPagingSize() const;
+
+    /// `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(Index::idx_t, const float* x) override;
+
+    /// `x` and `ids` 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_with_ids(Index::idx_t n, 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,
+            const float* x,
+            Index::idx_t k,
+            float* distances,
+            Index::idx_t* labels) const override;
+
+    /// Overridden to force GPU indices to provide their own GPU-friendly
+    /// implementation
+    void compute_residual(const float* x, float* residual, Index::idx_t key)
+            const override;
+
+    /// Overridden to force GPU indices to provide their own GPU-friendly
+    /// implementation
+    void compute_residual_n(
+            Index::idx_t n,
+            const float* xs,
+            float* residuals,
+            const Index::idx_t* keys) const override;
+
+   protected:
+    /// Copy what we need from the CPU equivalent
+    void copyFrom(const faiss::Index* index);
+
+    /// Copy what we have to the CPU equivalent
+    void copyTo(faiss::Index* index) const;
+
+    /// Does addImpl_ require IDs? If so, and no IDs are provided, we will
+    /// generate them sequentially based on the order in which the IDs are added
+    virtual bool addImplRequiresIDs_() const = 0;
+
+    /// Overridden to actually perform the add
+    /// All data is guaranteed to be resident on our device
+    virtual void addImpl_(int n, const float* x, const Index::idx_t* ids) = 0;
+
+    /// Overridden to actually perform the search
+    /// All data is guaranteed to be resident on our device
+    virtual void searchImpl_(
+            int n,
+            const float* x,
+            int k,
+            float* distances,
+            Index::idx_t* labels) const = 0;
+
+   private:
+    /// Handles paged adds if the add set is too large, passes to
+    /// addImpl_ to actually perform the add for the current page
+    void addPaged_(int n, const float* x, const Index::idx_t* ids);
+
+    /// Calls addImpl_ for a single page of GPU-resident data
+    void addPage_(int n, const float* x, const Index::idx_t* ids);
+
+    /// Calls searchImpl_ for a single page of GPU-resident data
+    void searchNonPaged_(
+            int n,
+            const float* x,
+            int k,
+            float* outDistancesData,
+            Index::idx_t* outIndicesData) const;
+
+    /// Calls searchImpl_ for a single page of GPU-resident data,
+    /// handling paging of the data and copies from the CPU
+    void searchFromCpuPaged_(
+            int n,
+            const float* x,
+            int k,
+            float* outDistancesData,
+            Index::idx_t* outIndicesData) const;
+
+   protected:
+    /// Manages streams, cuBLAS handles and scratch memory for devices
+    std::shared_ptr<GpuResources> resources_;
+
+    /// Our configuration options
+    const GpuIndexConfig config_;
+
+    /// Size above which we page copies from the CPU to GPU
+    size_t minPagedSize_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

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

@@ -12,85 +12,88 @@
 #include <faiss/gpu/GpuResources.h>
 #include <memory>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class BinaryFlatIndex;
 
-struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {
-};
+struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {};
 
 /// A GPU version of IndexBinaryFlat for brute-force comparison of bit vectors
 /// via Hamming distance
 class GpuIndexBinaryFlat : public IndexBinary {
- public:
-  /// Construct from a pre-existing faiss::IndexBinaryFlat instance, copying
-  /// data over to the given GPU
-  GpuIndexBinaryFlat(GpuResourcesProvider* resources,
-                     const faiss::IndexBinaryFlat* index,
-                     GpuIndexBinaryFlatConfig config =
-                     GpuIndexBinaryFlatConfig());
-
-  /// Construct an empty instance that can be added to
-  GpuIndexBinaryFlat(GpuResourcesProvider* resources,
-                     int dims,
-                     GpuIndexBinaryFlatConfig config =
-                     GpuIndexBinaryFlatConfig());
-
-  ~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);
-
-  /// Copy ourselves to the given CPU index; will overwrite all data
-  /// in the index instance
-  void copyTo(faiss::IndexBinaryFlat* index) const;
-
-  void add(faiss::IndexBinary::idx_t n,
-           const uint8_t* x) override;
-
-  void reset() override;
-
-  void search(faiss::IndexBinary::idx_t n,
-              const uint8_t* x,
-              faiss::IndexBinary::idx_t k,
-              int32_t* distances,
-              faiss::IndexBinary::idx_t* labels) const override;
-
-  void reconstruct(faiss::IndexBinary::idx_t key,
-                   uint8_t* recons) const override;
-
- protected:
-  /// Called from search when the input data is on the CPU;
-  /// potentially allows for pinned memory usage
-  void searchFromCpuPaged_(int n,
-                           const uint8_t* x,
-                           int k,
-                           int32_t* outDistancesData,
-                           int* outIndicesData) const;
-
-  void searchNonPaged_(int n,
-                       const uint8_t* x,
-                       int k,
-                       int32_t* outDistancesData,
-                       int* outIndicesData) const;
-
- protected:
-  /// Manages streans, cuBLAS handles and scratch memory for devices
-  std::shared_ptr<GpuResources> resources_;
-
-  /// Configuration options
-  const GpuIndexBinaryFlatConfig binaryFlatConfig_;
-
-  /// Holds our GPU data containing the list of vectors
-  std::unique_ptr<BinaryFlatIndex> data_;
+   public:
+    /// Construct from a pre-existing faiss::IndexBinaryFlat instance, copying
+    /// data over to the given GPU
+    GpuIndexBinaryFlat(
+            GpuResourcesProvider* resources,
+            const faiss::IndexBinaryFlat* index,
+            GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig());
+
+    /// Construct an empty instance that can be added to
+    GpuIndexBinaryFlat(
+            GpuResourcesProvider* resources,
+            int dims,
+            GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig());
+
+    ~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);
+
+    /// Copy ourselves to the given CPU index; will overwrite all data
+    /// in the index instance
+    void copyTo(faiss::IndexBinaryFlat* index) const;
+
+    void add(faiss::IndexBinary::idx_t n, const uint8_t* x) override;
+
+    void reset() override;
+
+    void search(
+            faiss::IndexBinary::idx_t n,
+            const uint8_t* x,
+            faiss::IndexBinary::idx_t k,
+            int32_t* distances,
+            faiss::IndexBinary::idx_t* labels) const override;
+
+    void reconstruct(faiss::IndexBinary::idx_t key, uint8_t* recons)
+            const override;
+
+   protected:
+    /// Called from search when the input data is on the CPU;
+    /// potentially allows for pinned memory usage
+    void searchFromCpuPaged_(
+            int n,
+            const uint8_t* x,
+            int k,
+            int32_t* outDistancesData,
+            int* outIndicesData) const;
+
+    void searchNonPaged_(
+            int n,
+            const uint8_t* x,
+            int k,
+            int32_t* outDistancesData,
+            int* outIndicesData) const;
+
+   protected:
+    /// Manages streans, cuBLAS handles and scratch memory for devices
+    std::shared_ptr<GpuResources> resources_;
+
+    /// Configuration options
+    const GpuIndexBinaryFlatConfig binaryFlatConfig_;
+
+    /// Holds our GPU data containing the list of vectors
+    std::unique_ptr<BinaryFlatIndex> data_;
 };
 
-} } // namespace gpu
+} // namespace gpu
+} // namespace faiss