faiss 0.1.0 → 0.1.1

data/vendor/faiss/gpu/GpuCloner.h

@@ -0,0 +1,82 @@
+/**
+ * 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 <vector>
+
+#include <faiss/Index.h>
+#include <faiss/clone_index.h>
+#include <faiss/gpu/GpuClonerOptions.h>
+#include <faiss/gpu/GpuIndex.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
+
+namespace faiss { namespace gpu {
+
+class GpuResources;
+
+
+/// Cloner specialized for GPU -> CPU
+struct ToCPUCloner: faiss::Cloner {
+    void merge_index(Index *dst, Index *src, bool successive_ids);
+    Index *clone_Index(const Index *index) override;
+};
+
+
+/// Cloner specialized for CPU -> 1 GPU
+struct ToGpuCloner: faiss::Cloner, GpuClonerOptions {
+    GpuResources *resources;
+    int device;
+
+    ToGpuCloner(GpuResources *resources, int device,
+                const GpuClonerOptions &options);
+
+    Index *clone_Index(const Index *index) override;
+
+};
+
+/// Cloner specialized for CPU -> multiple GPUs
+struct ToGpuClonerMultiple: faiss::Cloner, GpuMultipleClonerOptions {
+    std::vector<ToGpuCloner> sub_cloners;
+
+    ToGpuClonerMultiple(std::vector<GpuResources *> & resources,
+                        std::vector<int>& devices,
+                        const GpuMultipleClonerOptions &options);
+
+    ToGpuClonerMultiple(const std::vector<ToGpuCloner> & sub_cloners,
+                        const GpuMultipleClonerOptions &options);
+
+    void copy_ivf_shard (const IndexIVF *index_ivf, IndexIVF *idx2,
+                         long n, long i);
+
+    Index * clone_Index_to_shards (const Index *index);
+
+    /// main function
+    Index *clone_Index(const Index *index) override;
+};
+
+
+
+
+/// converts any GPU index inside gpu_index to a CPU index
+faiss::Index * index_gpu_to_cpu(const faiss::Index *gpu_index);
+
+/// converts any CPU index that can be converted to GPU
+faiss::Index * index_cpu_to_gpu(
+       GpuResources* resources, int device,
+       const faiss::Index *index,
+       const GpuClonerOptions *options = nullptr);
+
+faiss::Index * index_cpu_to_gpu_multiple(
+       std::vector<GpuResources*> & resources,
+       std::vector<int> &devices,
+       const faiss::Index *index,
+       const GpuMultipleClonerOptions *options = nullptr);
+
+
+
+} } // namespace

data/vendor/faiss/gpu/GpuClonerOptions.cpp

@@ -0,0 +1,28 @@
+/**
+ * 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/gpu/GpuClonerOptions.h>
+
+namespace faiss { namespace gpu {
+
+GpuClonerOptions::GpuClonerOptions()
+    : indicesOptions(INDICES_64_BIT),
+      useFloat16CoarseQuantizer(false),
+      useFloat16(false),
+      usePrecomputed(true),
+      reserveVecs(0),
+      storeTransposed(false),
+      verbose(false) {
+}
+
+GpuMultipleClonerOptions::GpuMultipleClonerOptions()
+    : shard(false),
+      shard_type(1)
+{
+}
+
+} } // namespace

data/vendor/faiss/gpu/GpuClonerOptions.h

@@ -0,0 +1,53 @@
+/**
+ * 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/GpuIndicesOptions.h>
+
+namespace faiss { namespace gpu {
+
+/// set some options on how to copy to GPU
+struct GpuClonerOptions {
+  GpuClonerOptions();
+
+  /// how should indices be stored on index types that support indices
+  /// (anything but GpuIndexFlat*)?
+  IndicesOptions indicesOptions;
+
+  /// is the coarse quantizer in float16?
+  bool useFloat16CoarseQuantizer;
+
+  /// for GpuIndexIVFFlat, is storage in float16?
+  /// for GpuIndexIVFPQ, are intermediate calculations in float16?
+  bool useFloat16;
+
+  /// use precomputed tables?
+  bool usePrecomputed;
+
+  /// reserve vectors in the invfiles?
+  long reserveVecs;
+
+  /// For GpuIndexFlat, store data in transposed layout?
+  bool storeTransposed;
+
+  /// Set verbose options on the index
+  bool verbose;
+};
+
+struct GpuMultipleClonerOptions : public GpuClonerOptions {
+  GpuMultipleClonerOptions ();
+
+  /// Whether to shard the index across GPUs, versus replication
+  /// across GPUs
+  bool shard;
+
+  /// IndexIVF::copy_subset_to subset type
+  int shard_type;
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuDistance.h

@@ -0,0 +1,52 @@
+/**
+ * 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>
+
+namespace faiss { namespace gpu {
+
+class GpuResources;
+
+/// A wrapper for gpu/impl/Distance.cuh to expose direct brute-force k-nearest
+/// neighbor searches on an externally-provided region of memory (e.g., from a
+/// pytorch tensor).
+/// The data (vectors, queries, outDistances, outIndices) can be resident on the
+/// GPU or the CPU, but all calculations are performed on the GPU. If the result
+/// buffers are on the CPU, results will be copied back when done.
+///
+/// All GPU computation is performed on the current CUDA device, and ordered
+/// with respect to resources->getDefaultStreamCurrentDevice().
+///
+/// For each vector in `queries`, searches all of `vectors` to find its k
+/// nearest neighbors with respect to the given metric
+void bruteForceKnn(GpuResources* resources,
+                   faiss::MetricType metric,
+                   // If vectorsRowMajor is true, this is
+                   // numVectors x dims, with dims innermost; otherwise,
+                   // dims x numVectors, with numVectors innermost
+                   const float* vectors,
+                   bool vectorsRowMajor,
+                   int numVectors,
+                   // If queriesRowMajor is true, this is
+                   // numQueries x dims, with dims innermost; otherwise,
+                   // dims x numQueries, with numQueries innermost
+                   const float* queries,
+                   bool queriesRowMajor,
+                   int numQueries,
+                   int dims,
+                   int k,
+                   // A region of memory size numQueries x k, with k
+                   // innermost (row major)
+                   float* outDistances,
+                   // A region of memory size numQueries x k, with k
+                   // innermost (row major)
+                   faiss::Index::idx_t* outIndices);
+
+} } // namespace

data/vendor/faiss/gpu/GpuFaissAssert.h

@@ -0,0 +1,29 @@
+/**
+ * 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.
+ */
+
+
+#ifndef GPU_FAISS_ASSERT_INCLUDED
+#define GPU_FAISS_ASSERT_INCLUDED
+
+#include <faiss/impl/FaissAssert.h>
+#include <cuda.h>
+
+///
+/// Assertions
+///
+
+#ifdef __CUDA_ARCH__
+#define GPU_FAISS_ASSERT(X) assert(X)
+#define GPU_FAISS_ASSERT_MSG(X, MSG) assert(X)
+#define GPU_FAISS_ASSERT_FMT(X, FMT, ...) assert(X)
+#else
+#define GPU_FAISS_ASSERT(X) FAISS_ASSERT(X)
+#define GPU_FAISS_ASSERT_MSG(X, MSG) FAISS_ASSERT_MSG(X, MSG)
+#define GPU_FAISS_ASSERT_FMT(X, FMT, ...) FAISS_ASSERT_FMT(X, FMT, __VA_ARGS)
+#endif // __CUDA_ARCH__
+
+#endif

data/vendor/faiss/gpu/GpuIndex.h

@@ -0,0 +1,148 @@
+/**
+ * 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/gpu/utils/MemorySpace.h>
+
+namespace faiss { namespace gpu {
+
+class GpuResources;
+
+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;
+};
+
+class GpuIndex : public faiss::Index {
+ public:
+  GpuIndex(GpuResources* resources,
+           int dims,
+           faiss::MetricType metric,
+           GpuIndexConfig config);
+
+  inline int getDevice() const {
+    return device_;
+  }
+
+  inline GpuResources* getResources() {
+    return resources_;
+  }
+
+  /// 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(faiss::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`, `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:
+  /// 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
+  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_;
+
+  /// Size above which we page copies from the CPU to GPU
+  size_t minPagedSize_;
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuIndexBinaryFlat.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/IndexBinaryFlat.h>
+#include <faiss/gpu/GpuIndex.h>
+
+namespace faiss { namespace gpu {
+
+class BinaryFlatIndex;
+class GpuResources;
+
+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(GpuResources* resources,
+                     const faiss::IndexBinaryFlat* index,
+                     GpuIndexBinaryFlatConfig config =
+                     GpuIndexBinaryFlatConfig());
+
+  /// Construct an empty instance that can be added to
+  GpuIndexBinaryFlat(GpuResources* resources,
+                     int dims,
+                     GpuIndexBinaryFlatConfig config =
+                     GpuIndexBinaryFlatConfig());
+
+  ~GpuIndexBinaryFlat() override;
+
+  /// 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
+  GpuResources* resources_;
+
+  /// Configuration options
+  GpuIndexBinaryFlatConfig 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
+  BinaryFlatIndex* data_;
+};
+
+} } // namespace gpu