faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/gpu/GpuCloner.cpp

@@ -7,11 +7,16 @@
 
 #include <faiss/gpu/GpuCloner.h>
 #include <faiss/impl/FaissAssert.h>
+#include <memory>
 #include <typeinfo>
 
 #include <faiss/gpu/StandardGpuResources.h>
 
+#include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexFlat.h>
+#if defined USE_NVIDIA_RAFT
+#include <faiss/IndexHNSW.h>
+#endif
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
@@ -21,6 +26,10 @@
 #include <faiss/IndexShardsIVF.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/gpu/GpuIndex.h>
+#include <faiss/gpu/GpuIndexBinaryFlat.h>
+#if defined USE_NVIDIA_RAFT
+#include <faiss/gpu/GpuIndexCagra.h>
+#endif
 #include <faiss/gpu/GpuIndexFlat.h>
 #include <faiss/gpu/GpuIndexIVFFlat.h>
 #include <faiss/gpu/GpuIndexIVFPQ.h>
@@ -82,7 +91,15 @@ Index* ToCPUCloner::clone_Index(const Index* index) {
         // objective is to make a single component out of them
         // (inverse op of ToGpuClonerMultiple)
 
-    } else if (auto ish = dynamic_cast<const IndexShards*>(index)) {
+    }
+#if defined USE_NVIDIA_RAFT
+    else if (auto icg = dynamic_cast<const GpuIndexCagra*>(index)) {
+        IndexHNSWCagra* res = new IndexHNSWCagra();
+        icg->copyTo(res);
+        return res;
+    }
+#endif
+    else if (auto ish = dynamic_cast<const IndexShards*>(index)) {
         int nshard = ish->count();
         FAISS_ASSERT(nshard > 0);
         Index* res = clone_Index(ish->at(0));
@@ -121,6 +138,7 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         GpuIndexFlatConfig config;
         config.device = device;
         config.useFloat16 = useFloat16;
+        config.use_raft = use_raft;
         return new GpuIndexFlat(provider, ifl, config);
     } else if (
             dynamic_cast<const IndexScalarQuantizer*>(index) &&
@@ -129,6 +147,8 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         GpuIndexFlatConfig config;
         config.device = device;
         config.useFloat16 = true;
+        FAISS_THROW_IF_NOT_MSG(
+                !use_raft, "this type of index is not implemented for RAFT");
         GpuIndexFlat* gif = new GpuIndexFlat(
                 provider, index->d, index->metric_type, config);
         // transfer data by blocks
@@ -146,6 +166,8 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.device = device;
         config.indicesOptions = indicesOptions;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
+        config.use_raft = use_raft;
+        config.allowCpuCoarseQuantizer = allowCpuCoarseQuantizer;
 
         GpuIndexIVFFlat* res = new GpuIndexIVFFlat(
                 provider, ifl->d, ifl->nlist, ifl->metric_type, config);
@@ -162,6 +184,8 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.device = device;
         config.indicesOptions = indicesOptions;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
+        FAISS_THROW_IF_NOT_MSG(
+                !use_raft, "this type of index is not implemented for RAFT");
 
         GpuIndexIVFScalarQuantizer* res = new GpuIndexIVFScalarQuantizer(
                 provider,
@@ -194,6 +218,9 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
         config.useFloat16LookupTables = useFloat16;
         config.usePrecomputedTables = usePrecomputed;
+        config.use_raft = use_raft;
+        config.interleavedLayout = use_raft;
+        config.allowCpuCoarseQuantizer = allowCpuCoarseQuantizer;
 
         GpuIndexIVFPQ* res = new GpuIndexIVFPQ(provider, ipq, config);
 
@@ -202,9 +229,25 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         }
 
         return res;
-    } else {
-        // default: use CPU cloner
-        return Cloner::clone_Index(index);
+    }
+#if defined USE_NVIDIA_RAFT
+    else if (auto icg = dynamic_cast<const faiss::IndexHNSWCagra*>(index)) {
+        GpuIndexCagraConfig config;
+        config.device = device;
+        GpuIndexCagra* res =
+                new GpuIndexCagra(provider, icg->d, icg->metric_type, config);
+        res->copyFrom(icg);
+        return res;
+    }
+#endif
+    else {
+        // use CPU cloner for IDMap and PreTransform
+        auto index_idmap = dynamic_cast<const IndexIDMap*>(index);
+        auto index_pt = dynamic_cast<const IndexPreTransform*>(index);
+        if (index_idmap || index_pt) {
+            return Cloner::clone_Index(index);
+        }
+        FAISS_THROW_MSG("This index type is not implemented on GPU.");
     }
 }
 
@@ -229,7 +272,7 @@ ToGpuClonerMultiple::ToGpuClonerMultiple(
         : GpuMultipleClonerOptions(options) {
     FAISS_THROW_IF_NOT(provider.size() == devices.size());
     for (size_t i = 0; i < provider.size(); i++) {
-        sub_cloners.push_back(ToGpuCloner(provider[i], devices[i], options));
+        sub_cloners.emplace_back(provider[i], devices[i], options);
     }
 }
 
@@ -298,8 +341,8 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
             !dynamic_cast<const IndexFlat*>(quantizer)) {
             // then we flatten the coarse quantizer so that everything remains
             // on GPU
-            new_quantizer.reset(
-                    new IndexFlat(quantizer->d, quantizer->metric_type));
+            new_quantizer = std::make_unique<IndexFlat>(
+                    quantizer->d, quantizer->metric_type);
             std::vector<float> centroids(quantizer->d * quantizer->ntotal);
             quantizer->reconstruct_n(0, quantizer->ntotal, centroids.data());
             new_quantizer->add(quantizer->ntotal, centroids.data());
@@ -309,6 +352,7 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
 
     std::vector<faiss::Index*> shards(n);
 
+#pragma omp parallel for
     for (idx_t i = 0; i < n; i++) {
         // make a shallow copy
         if (reserveVecs) {
@@ -321,7 +365,7 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
                     const_cast<Index*>(quantizer),
                     index_ivfpq->d,
                     index_ivfpq->nlist,
-                    index_ivfpq->code_size,
+                    index_ivfpq->pq.M,
                     index_ivfpq->pq.nbits);
             idx2.metric_type = index_ivfpq->metric_type;
             idx2.pq = index_ivfpq->pq;
@@ -473,5 +517,76 @@ Index* GpuProgressiveDimIndexFactory::operator()(int dim) {
     return index_cpu_to_gpu_multiple(vres, devices, &index, &options);
 }
 
+/*********************************************
+ * Cloning binary indexes
+ *********************************************/
+
+faiss::IndexBinary* index_binary_gpu_to_cpu(
+        const faiss::IndexBinary* gpu_index) {
+    if (auto ii = dynamic_cast<const GpuIndexBinaryFlat*>(gpu_index)) {
+        IndexBinaryFlat* ret = new IndexBinaryFlat();
+        ii->copyTo(ret);
+        return ret;
+    } else {
+        FAISS_THROW_MSG("cannot clone this type of index");
+    }
+}
+
+faiss::IndexBinary* index_binary_cpu_to_gpu(
+        GpuResourcesProvider* provider,
+        int device,
+        const faiss::IndexBinary* index,
+        const GpuClonerOptions* options) {
+    if (auto ii = dynamic_cast<const IndexBinaryFlat*>(index)) {
+        GpuIndexBinaryFlatConfig config;
+        config.device = device;
+        if (options) {
+            config.use_raft = options->use_raft;
+        }
+        return new GpuIndexBinaryFlat(provider, ii, config);
+    } else {
+        FAISS_THROW_MSG("cannot clone this type of index");
+    }
+}
+
+faiss::IndexBinary* index_binary_cpu_to_gpu_multiple(
+        std::vector<GpuResourcesProvider*>& provider,
+        std::vector<int>& devices,
+        const faiss::IndexBinary* index,
+        const GpuMultipleClonerOptions* options) {
+    GpuMultipleClonerOptions defaults;
+    FAISS_THROW_IF_NOT(devices.size() == provider.size());
+    int n = devices.size();
+    if (n == 1) {
+        return index_binary_cpu_to_gpu(provider[0], devices[0], index, options);
+    }
+    if (!options) {
+        options = &defaults;
+    }
+    if (options->shard) {
+        auto* fi = dynamic_cast<const IndexBinaryFlat*>(index);
+        FAISS_THROW_IF_NOT_MSG(fi, "only flat index cloning supported");
+        IndexBinaryShards* ret = new IndexBinaryShards(true, true);
+        for (int i = 0; i < n; i++) {
+            IndexBinaryFlat fig(fi->d);
+            size_t i0 = i * fi->ntotal / n;
+            size_t i1 = (i + 1) * fi->ntotal / n;
+            fig.add(i1 - i0, fi->xb.data() + i0 * fi->code_size);
+            ret->addIndex(index_binary_cpu_to_gpu(
+                    provider[i], devices[i], &fig, options));
+        }
+        ret->own_indices = true;
+        return ret;
+    } else { // replicas
+        IndexBinaryReplicas* ret = new IndexBinaryReplicas(true);
+        for (int i = 0; i < n; i++) {
+            ret->addIndex(index_binary_cpu_to_gpu(
+                    provider[i], devices[i], index, options));
+        }
+        ret->own_indices = true;
+        return ret;
+    }
+}
+
 } // namespace gpu
 } // namespace faiss

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

@@ -11,10 +11,12 @@
 
 #include <faiss/Clustering.h>
 #include <faiss/Index.h>
+#include <faiss/IndexBinary.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 {
 
@@ -95,5 +97,25 @@ struct GpuProgressiveDimIndexFactory : ProgressiveDimIndexFactory {
     virtual ~GpuProgressiveDimIndexFactory() override;
 };
 
+/*********************************************
+ * Cloning binary indexes
+ *********************************************/
+
+faiss::IndexBinary* index_binary_gpu_to_cpu(
+        const faiss::IndexBinary* gpu_index);
+
+/// converts any CPU index that can be converted to GPU
+faiss::IndexBinary* index_binary_cpu_to_gpu(
+        GpuResourcesProvider* provider,
+        int device,
+        const faiss::IndexBinary* index,
+        const GpuClonerOptions* options = nullptr);
+
+faiss::IndexBinary* index_binary_cpu_to_gpu_multiple(
+        std::vector<GpuResourcesProvider*>& provider,
+        std::vector<int>& devices,
+        const faiss::IndexBinary* index,
+        const GpuMultipleClonerOptions* options = nullptr);
+
 } // namespace gpu
 } // namespace faiss

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

@@ -36,6 +36,19 @@ struct GpuClonerOptions {
 
     /// Set verbose options on the index
     bool verbose = false;
+
+    /// use the RAFT implementation
+#if defined USE_NVIDIA_RAFT
+    bool use_raft = true;
+#else
+    bool use_raft = false;
+#endif
+
+    /// This flag controls the CPU fallback logic for coarse quantizer
+    /// component of the index. When set to false (default), the cloner will
+    /// throw an exception for indices not implemented on GPU. When set to
+    /// true, it will fallback to a CPU implementation.
+    bool allowCpuCoarseQuantizer = false;
 };
 
 struct GpuMultipleClonerOptions : public GpuClonerOptions {

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

@@ -9,6 +9,7 @@
 
 #include <faiss/Index.h>
 
+#pragma GCC visibility push(default)
 namespace faiss {
 namespace gpu {
 
@@ -28,44 +29,24 @@ enum class IndicesDataType {
 
 /// 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),
-              device(-1) {}
-
     //
     // Search parameters
     //
 
     /// Search parameter: distance metric
-    faiss::MetricType metric;
+    faiss::MetricType metric = METRIC_L2;
 
     /// Search parameter: distance metric argument (if applicable)
     /// For metric == METRIC_Lp, this is the p-value
-    float metricArg;
+    float metricArg = 0;
 
     /// Search parameter: return k nearest neighbors
     /// If the value provided is -1, then we report all pairwise distances
     /// without top-k filtering
-    int k;
+    int k = 0;
 
     /// Vector dimensionality
-    int dims;
+    int dims = 0;
 
     //
     // Vectors being queried
@@ -74,14 +55,14 @@ struct GpuDistanceParams {
     /// 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;
-    idx_t numVectors;
+    const void* vectors = nullptr;
+    DistanceDataType vectorType = DistanceDataType::F32;
+    bool vectorsRowMajor = true;
+    idx_t numVectors = 0;
 
     /// Precomputed L2 norms for each vector in `vectors`, which can be
     /// optionally provided in advance to speed computation for METRIC_L2
-    const float* vectorNorms;
+    const float* vectorNorms = nullptr;
 
     //
     // The query vectors (i.e., find k-nearest neighbors in `vectors` for each
@@ -91,10 +72,10 @@ struct GpuDistanceParams {
     /// 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;
-    idx_t numQueries;
+    const void* queries = nullptr;
+    DistanceDataType queryType = DistanceDataType::F32;
+    bool queriesRowMajor = true;
+    idx_t numQueries = 0;
 
     //
     // Output results
@@ -103,16 +84,16 @@ struct GpuDistanceParams {
     /// 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;
+    float* outDistances = nullptr;
 
     /// Do we only care about the indices reported, rather than the output
     /// distances? Not used if k == -1 (all pairwise distances)
-    bool ignoreOutDistances;
+    bool ignoreOutDistances = false;
 
     /// 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;
+    IndicesDataType outIndicesType = IndicesDataType::I64;
+    void* outIndices = nullptr;
 
     //
     // Execution information
@@ -123,9 +104,17 @@ struct GpuDistanceParams {
     /// (via cudaGetDevice/cudaSetDevice) is used
     /// Otherwise, an integer 0 <= device < numDevices indicates the device for
     /// execution
-    int device;
+    int device = -1;
+
+    /// Should the index dispatch down to RAFT?
+    /// TODO: change default to true if RAFT is enabled
+    bool use_raft = false;
 };
 
+/// A function that determines whether RAFT should be used based on various
+/// conditions (such as unsupported architecture)
+bool should_use_raft(GpuDistanceParams args);
+
 /// 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).
@@ -140,6 +129,24 @@ struct GpuDistanceParams {
 /// nearest neighbors with respect to the given metric
 void bfKnn(GpuResourcesProvider* resources, const GpuDistanceParams& args);
 
+// bfKnn which takes two extra parameters to control the maximum GPU
+// memory allowed for vectors and queries, the latter including the
+// memory required for the results.
+// If 0, the corresponding input must fit into GPU memory.
+// If greater than 0, the function will use at most this much GPU
+// memory (in bytes) for vectors and queries respectively.
+// Vectors are broken up into chunks of size vectorsMemoryLimit,
+// and queries are broken up into chunks of size queriesMemoryLimit.
+// The tiles resulting from the product of the query and vector
+// chunks are processed sequentially on the GPU.
+// Only supported for row major matrices and k > 0. The input that
+// needs sharding must reside on the CPU.
+void bfKnn_tiling(
+        GpuResourcesProvider* resources,
+        const GpuDistanceParams& args,
+        size_t vectorsMemoryLimit,
+        size_t queriesMemoryLimit);
+
 /// Deprecated legacy implementation
 void bruteForceKnn(
         GpuResourcesProvider* resources,
@@ -167,3 +174,4 @@ void bruteForceKnn(
 
 } // namespace gpu
 } // namespace faiss
+#pragma GCC visibility pop

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

@@ -15,7 +15,7 @@
 /// Assertions
 ///
 
-#ifdef __CUDA_ARCH__
+#if defined(__CUDA_ARCH__) || defined(USE_AMD_ROCM)
 #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)

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

@@ -4,6 +4,21 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 
 #pragma once
 
@@ -14,17 +29,26 @@ namespace faiss {
 namespace gpu {
 
 struct GpuIndexConfig {
-    inline GpuIndexConfig() : device(0), memorySpace(MemorySpace::Device) {}
-
     /// GPU device on which the index is resident
-    int device;
+    int device = 0;
 
     /// 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;
+    MemorySpace memorySpace = MemorySpace::Device;
+
+    /// Should the index dispatch down to RAFT?
+#if defined USE_NVIDIA_RAFT
+    bool use_raft = true;
+#else
+    bool use_raft = false;
+#endif
 };
 
+/// A centralized function that determines whether RAFT should
+/// be used based on various conditions (such as unsupported architecture)
+bool should_use_raft(GpuIndexConfig config_);
+
 class GpuIndex : public faiss::Index {
    public:
     GpuIndex(
@@ -60,19 +84,14 @@ class GpuIndex : public faiss::Index {
 
     /// `x` and `labels` can be resident on the CPU or any GPU; copies are
     /// performed as needed
-    void assign(
-            idx_t n,
-            const float* x,
-            idx_t* labels,
-            // faiss::Index has idx_t for k
-            idx_t k = 1) const override;
+    void assign(idx_t n, const float* x, idx_t* labels, 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(
             idx_t n,
             const float* x,
-            // faiss::Index has idx_t for k
             idx_t k,
             float* distances,
             idx_t* labels,
@@ -83,7 +102,6 @@ class GpuIndex : public faiss::Index {
     void search_and_reconstruct(
             idx_t n,
             const float* x,
-            // faiss::Index has idx_t for k
             idx_t k,
             float* distances,
             idx_t* labels,