faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/gpu/GpuCloner.h

@@ -9,73 +9,91 @@
 
 #include <vector>
 
+#include <faiss/Clustering.h>
 #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 {
+namespace faiss {
+namespace gpu {
 
 class GpuResourcesProvider;
 
 /// 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;
+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 {
-    GpuResourcesProvider *provider;
+struct ToGpuCloner : faiss::Cloner, GpuClonerOptions {
+    GpuResourcesProvider* provider;
     int device;
 
-    ToGpuCloner(GpuResourcesProvider *prov, int device,
-                const GpuClonerOptions &options);
-
-    Index *clone_Index(const Index *index) override;
+    ToGpuCloner(
+            GpuResourcesProvider* prov,
+            int device,
+            const GpuClonerOptions& options);
 
+    Index* clone_Index(const Index* index) override;
 };
 
 /// Cloner specialized for CPU -> multiple GPUs
-struct ToGpuClonerMultiple: faiss::Cloner, GpuMultipleClonerOptions {
+struct ToGpuClonerMultiple : faiss::Cloner, GpuMultipleClonerOptions {
     std::vector<ToGpuCloner> sub_cloners;
 
-    ToGpuClonerMultiple(std::vector<GpuResourcesProvider *> & provider,
-                        std::vector<int>& devices,
-                        const GpuMultipleClonerOptions &options);
+    ToGpuClonerMultiple(
+            std::vector<GpuResourcesProvider*>& provider,
+            std::vector<int>& devices,
+            const GpuMultipleClonerOptions& options);
 
-    ToGpuClonerMultiple(const std::vector<ToGpuCloner> & sub_cloners,
-                        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);
+    void copy_ivf_shard(
+            const IndexIVF* index_ivf,
+            IndexIVF* idx2,
+            long n,
+            long i);
 
-    Index * clone_Index_to_shards (const Index *index);
+    Index* clone_Index_to_shards(const Index* index);
 
     /// main function
-    Index *clone_Index(const Index *index) override;
+    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(
+        GpuResourcesProvider* provider,
+        int device,
+        const faiss::Index* index,
+        const GpuClonerOptions* options = nullptr);
 
+faiss::Index* index_cpu_to_gpu_multiple(
+        std::vector<GpuResourcesProvider*>& provider,
+        std::vector<int>& devices,
+        const faiss::Index* index,
+        const GpuMultipleClonerOptions* options = nullptr);
 
-/// converts any GPU index inside gpu_index to a CPU index
-faiss::Index * index_gpu_to_cpu(const faiss::Index *gpu_index);
+/// index factory for the ProgressiveDimClustering object
 
-/// converts any CPU index that can be converted to GPU
-faiss::Index * index_cpu_to_gpu(
-       GpuResourcesProvider* provider, int device,
-       const faiss::Index *index,
-       const GpuClonerOptions *options = nullptr);
+struct GpuProgressiveDimIndexFactory : ProgressiveDimIndexFactory {
+    GpuMultipleClonerOptions options;
+    std::vector<GpuResourcesProvider*> vres;
+    std::vector<int> devices;
+    int ncall;
 
-faiss::Index * index_cpu_to_gpu_multiple(
-       std::vector<GpuResourcesProvider*> & provider,
-       std::vector<int> &devices,
-       const faiss::Index *index,
-       const GpuMultipleClonerOptions *options = nullptr);
+    explicit GpuProgressiveDimIndexFactory(int ngpu);
 
+    Index* operator()(int dim) override;
 
+    virtual ~GpuProgressiveDimIndexFactory() override;
+};
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp

@@ -7,22 +7,20 @@
 
 #include <faiss/gpu/GpuClonerOptions.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 GpuClonerOptions::GpuClonerOptions()
-    : indicesOptions(INDICES_64_BIT),
-      useFloat16CoarseQuantizer(false),
-      useFloat16(false),
-      usePrecomputed(false),
-      reserveVecs(0),
-      storeTransposed(false),
-      verbose(false) {
-}
+        : indicesOptions(INDICES_64_BIT),
+          useFloat16CoarseQuantizer(false),
+          useFloat16(false),
+          usePrecomputed(false),
+          reserveVecs(0),
+          storeTransposed(false),
+          verbose(false) {}
 
 GpuMultipleClonerOptions::GpuMultipleClonerOptions()
-    : shard(false),
-      shard_type(1)
-{
-}
+        : shard(false), shard_type(1) {}
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

data/vendor/faiss/faiss/gpu/GpuClonerOptions.h

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

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

@@ -5,111 +5,113 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResourcesProvider;
 
 // Scalar type of the vector data
 enum class DistanceDataType {
-  F32 = 1,
-  F16,
+    F32 = 1,
+    F16,
 };
 
 // Scalar type of the indices data
 enum class IndicesDataType {
-  I64 = 1,
-  I32,
+    I64 = 1,
+    I32,
 };
 
 /// Arguments to brute-force GPU k-nearest neighbor searching
 struct GpuDistanceParams {
-  GpuDistanceParams()
-      : metric(faiss::MetricType::METRIC_L2),
-        metricArg(0),
-        k(0),
-        dims(0),
-        vectors(nullptr),
-        vectorType(DistanceDataType::F32),
-        vectorsRowMajor(true),
-        numVectors(0),
-        vectorNorms(nullptr),
-        queries(nullptr),
-        queryType(DistanceDataType::F32),
-        queriesRowMajor(true),
-        numQueries(0),
-        outDistances(nullptr),
-        ignoreOutDistances(false),
-        outIndicesType(IndicesDataType::I64),
-        outIndices(nullptr) {
-  }
-
-  //
-  // Search parameters
-  //
-
-  // Search parameter: distance metric
-  faiss::MetricType metric;
-
-  // Search parameter: distance metric argument (if applicable)
-  // For metric == METRIC_Lp, this is the p-value
-  float metricArg;
-
-  // Search parameter: return k nearest neighbors
-  int k;
-
-  // Vector dimensionality
-  int dims;
-
-  //
-  // Vectors being queried
-  //
-
-  // If vectorsRowMajor is true, this is
-  // numVectors x dims, with dims innermost; otherwise,
-  // dims x numVectors, with numVectors innermost
-  const void* vectors;
-  DistanceDataType vectorType;
-  bool vectorsRowMajor;
-  int numVectors;
-
-  // Precomputed L2 norms for each vector in `vectors`, which can be optionally
-  // provided in advance to speed computation for METRIC_L2
-  const float* vectorNorms;
-
-  //
-  // The query vectors (i.e., find k-nearest neighbors in `vectors` for each of
-  // the `queries`
-  //
-
-  // If queriesRowMajor is true, this is
-  // numQueries x dims, with dims innermost; otherwise,
-  // dims x numQueries, with numQueries innermost
-  const void* queries;
-  DistanceDataType queryType;
-  bool queriesRowMajor;
-  int numQueries;
-
-  //
-  // Output results
-  //
-
-  // A region of memory size numQueries x k, with k
-  // innermost (row major)
-  float* outDistances;
-
-  // Do we only care abouty the indices reported, rather than the output
-  // distances?
-  bool ignoreOutDistances;
-
-  // A region of memory size numQueries x k, with k
-  // innermost (row major)
-  IndicesDataType outIndicesType;
-  void* outIndices;
+    GpuDistanceParams()
+            : metric(faiss::MetricType::METRIC_L2),
+              metricArg(0),
+              k(0),
+              dims(0),
+              vectors(nullptr),
+              vectorType(DistanceDataType::F32),
+              vectorsRowMajor(true),
+              numVectors(0),
+              vectorNorms(nullptr),
+              queries(nullptr),
+              queryType(DistanceDataType::F32),
+              queriesRowMajor(true),
+              numQueries(0),
+              outDistances(nullptr),
+              ignoreOutDistances(false),
+              outIndicesType(IndicesDataType::I64),
+              outIndices(nullptr) {}
+
+    //
+    // Search parameters
+    //
+
+    /// Search parameter: distance metric
+    faiss::MetricType metric;
+
+    /// Search parameter: distance metric argument (if applicable)
+    /// For metric == METRIC_Lp, this is the p-value
+    float metricArg;
+
+    /// Search parameter: return k nearest neighbors
+    /// If the value provided is -1, then we report all pairwise distances
+    /// without top-k filtering
+    int k;
+
+    /// Vector dimensionality
+    int dims;
+
+    //
+    // Vectors being queried
+    //
+
+    /// If vectorsRowMajor is true, this is
+    /// numVectors x dims, with dims innermost; otherwise,
+    /// dims x numVectors, with numVectors innermost
+    const void* vectors;
+    DistanceDataType vectorType;
+    bool vectorsRowMajor;
+    int numVectors;
+
+    /// Precomputed L2 norms for each vector in `vectors`, which can be
+    /// optionally provided in advance to speed computation for METRIC_L2
+    const float* vectorNorms;
+
+    //
+    // The query vectors (i.e., find k-nearest neighbors in `vectors` for each
+    // of the `queries`
+    //
+
+    /// If queriesRowMajor is true, this is
+    /// numQueries x dims, with dims innermost; otherwise,
+    /// dims x numQueries, with numQueries innermost
+    const void* queries;
+    DistanceDataType queryType;
+    bool queriesRowMajor;
+    int numQueries;
+
+    //
+    // Output results
+    //
+
+    /// A region of memory size numQueries x k, with k
+    /// innermost (row major) if k > 0, or if k == -1, a region of memory of
+    /// size numQueries x numVectors
+    float* outDistances;
+
+    /// Do we only care about the indices reported, rather than the output
+    /// distances? Not used if k == -1 (all pairwise distances)
+    bool ignoreOutDistances;
+
+    /// A region of memory size numQueries x k, with k
+    /// innermost (row major). Not used if k == -1 (all pairwise distances)
+    IndicesDataType outIndicesType;
+    void* outIndices;
 };
 
 /// A wrapper for gpu/impl/Distance.cuh to expose direct brute-force k-nearest
@@ -127,27 +129,29 @@ struct GpuDistanceParams {
 void bfKnn(GpuResourcesProvider* resources, const GpuDistanceParams& args);
 
 /// Deprecated legacy implementation
-void bruteForceKnn(GpuResourcesProvider* 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)
-                   Index::idx_t* outIndices);
-
-} } // namespace
+void bruteForceKnn(
+        GpuResourcesProvider* 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)
+        Index::idx_t* outIndices);
+
+} // namespace gpu
+} // namespace faiss

data/vendor/faiss/faiss/gpu/GpuFaissAssert.h

@@ -5,12 +5,11 @@
  * 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>
+#include <faiss/impl/FaissAssert.h>
 
 ///
 /// Assertions

data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h

@@ -0,0 +1,60 @@
+/**
+ * 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/impl/LocalSearchQuantizer.h>
+
+#include <memory>
+
+namespace faiss {
+namespace gpu {
+
+class GpuResourcesProvider;
+struct IcmEncoderShards;
+
+/** Perform LSQ encoding on GPU.
+ *
+ * Split input vectors to different devices and call IcmEncoderImpl::encode
+ * to encode them
+ */
+class GpuIcmEncoder : public lsq::IcmEncoder {
+   public:
+    GpuIcmEncoder(
+            const LocalSearchQuantizer* lsq,
+            const std::vector<GpuResourcesProvider*>& provs,
+            const std::vector<int>& devices);
+
+    ~GpuIcmEncoder();
+
+    GpuIcmEncoder(const GpuIcmEncoder&) = delete;
+    GpuIcmEncoder& operator=(const GpuIcmEncoder&) = delete;
+
+    void set_binary_term() override;
+
+    void encode(
+            int32_t* codes,
+            const float* x,
+            std::mt19937& gen,
+            size_t n,
+            size_t ils_iters) const override;
+
+   private:
+    std::unique_ptr<IcmEncoderShards> shards;
+};
+
+struct GpuIcmEncoderFactory : public lsq::IcmEncoderFactory {
+    explicit GpuIcmEncoderFactory(int ngpus = 1);
+
+    lsq::IcmEncoder* get(const LocalSearchQuantizer* lsq) override;
+
+    std::vector<GpuResourcesProvider*> provs;
+    std::vector<int> devices;
+};
+
+} // namespace gpu
+} // namespace faiss