faiss 0.6.1 → 0.6.3

data/vendor/faiss/faiss/factory_tools.cpp

@@ -14,8 +14,10 @@
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexEDEN.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIDMap.h>
+#include <faiss/IndexIVFEDEN.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQFastScan.h>
 #include <faiss/IndexIVFRaBitQ.h>
@@ -33,16 +35,26 @@ namespace {
 const std::map<faiss::ScalarQuantizer::QuantizerType, std::string> sq_types = {
         {faiss::ScalarQuantizer::QT_8bit, "SQ8"},
         {faiss::ScalarQuantizer::QT_4bit, "SQ4"},
+        // QT_8bit_uniform and QT_4bit_uniform have no round-trippable
+        // index_factory string; the names below are synthetic identifiers used
+        // for telemetry logging only.
+        {faiss::ScalarQuantizer::QT_8bit_uniform, "SQ8u"},
+        {faiss::ScalarQuantizer::QT_4bit_uniform, "SQ4u"},
         {faiss::ScalarQuantizer::QT_6bit, "SQ6"},
         {faiss::ScalarQuantizer::QT_fp16, "SQfp16"},
         {faiss::ScalarQuantizer::QT_bf16, "SQbf16"},
         {faiss::ScalarQuantizer::QT_8bit_direct_signed, "SQ8_direct_signed"},
         {faiss::ScalarQuantizer::QT_8bit_direct, "SQ8_direct"},
+        {faiss::ScalarQuantizer::QT_0bit, "SQ0"},
         {faiss::ScalarQuantizer::QT_1bit_tqmse, "SQtqmse1"},
         {faiss::ScalarQuantizer::QT_2bit_tqmse, "SQtqmse2"},
         {faiss::ScalarQuantizer::QT_3bit_tqmse, "SQtqmse3"},
         {faiss::ScalarQuantizer::QT_4bit_tqmse, "SQtqmse4"},
         {faiss::ScalarQuantizer::QT_8bit_tqmse, "SQtqmse8"},
+        {faiss::ScalarQuantizer::QT_2bit_tq, "SQtq2"},
+        {faiss::ScalarQuantizer::QT_3bit_tq, "SQtq3"},
+        {faiss::ScalarQuantizer::QT_4bit_tq, "SQtq4"},
+        {faiss::ScalarQuantizer::QT_5bit_tq, "SQtq5"},
 };
 
 int get_hnsw_M(const faiss::IndexHNSW* index) {
@@ -61,6 +73,19 @@ int get_hnsw_M(const faiss::IndexBinaryHNSW* index) {
     return 0;
 }
 
+std::string eden_factory_string(
+        const faiss::ScalarQuantizer& sq,
+        faiss::EDENScaleType scale_type) {
+    std::string result = "EDEN";
+    if (sq.bits != 1) {
+        result += std::to_string(sq.bits);
+    }
+    if (scale_type == faiss::EDENScaleType_BIASED) {
+        result += "BIASED";
+    }
+    return result;
+}
+
 } // namespace
 
 // Reference for reverse_index_factory:
@@ -110,8 +135,17 @@ std::string reverse_index_factory(const faiss::Index* index) {
                                 ivf_index)) {
             return prefix + ",PQ" + std::to_string(ivfpqfs_index->pq.M) + "x" +
                     std::to_string(ivfpqfs_index->pq.nbits) + "fs";
-        } else if (dynamic_cast<const faiss::IndexIVFRaBitQ*>(ivf_index)) {
-            return prefix + ",RaBitQ";
+        } else if (
+                const faiss::IndexIVFRaBitQ* ivfrabitq_index =
+                        dynamic_cast<const faiss::IndexIVFRaBitQ*>(ivf_index)) {
+            size_t nb_bits = ivfrabitq_index->rabitq.nb_bits;
+            return prefix + ",RaBitQ" +
+                    (nb_bits == 1 ? "" : std::to_string(nb_bits));
+        } else if (
+                const faiss::IndexIVFEDEN* ivf_eden =
+                        dynamic_cast<const faiss::IndexIVFEDEN*>(ivf_index)) {
+            return prefix + "," +
+                    eden_factory_string(ivf_eden->sq, ivf_eden->scale_type);
         }
     } else if (
             const faiss::IndexPreTransform* pretransform_index =
@@ -180,12 +214,25 @@ std::string reverse_index_factory(const faiss::Index* index) {
             const faiss::IndexScalarQuantizer* sq_index =
                     dynamic_cast<const faiss::IndexScalarQuantizer*>(index)) {
         return sq_types.at(sq_index->sq.qtype);
+    } else if (
+            // IndexIDMap2 inherits IndexIDMap — check subclass first.
+            const faiss::IndexIDMap2* idmap2 =
+                    dynamic_cast<const faiss::IndexIDMap2*>(index)) {
+        return std::string("IDMap2,") + reverse_index_factory(idmap2->index);
     } else if (
             const faiss::IndexIDMap* idmap =
                     dynamic_cast<const faiss::IndexIDMap*>(index)) {
         return std::string("IDMap,") + reverse_index_factory(idmap->index);
-    } else if (dynamic_cast<const faiss::IndexRaBitQ*>(index)) {
-        return "RaBitQ";
+    } else if (
+            const faiss::IndexRaBitQ* rabitq_index =
+                    dynamic_cast<const faiss::IndexRaBitQ*>(index)) {
+        size_t nb_bits = rabitq_index->rabitq.nb_bits;
+        return "RaBitQ" +
+                (nb_bits == 1 ? std::string() : std::to_string(nb_bits));
+    } else if (
+            const faiss::IndexEDEN* eden =
+                    dynamic_cast<const faiss::IndexEDEN*>(index)) {
+        return eden_factory_string(eden->sq, eden->scale_type);
     }
     // Avoid runtime error, just return empty string for logging.
     return "";

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

@@ -14,7 +14,7 @@
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexFlat.h>
-#if defined USE_NVIDIA_CUVS
+#if defined(USE_NVIDIA_CUVS) && !defined(FAISS_CUVS_NO_CAGRA)
 #include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexHNSW.h>
 #endif
@@ -28,7 +28,7 @@
 #include <faiss/MetaIndexes.h>
 #include <faiss/gpu/GpuIndex.h>
 #include <faiss/gpu/GpuIndexBinaryFlat.h>
-#if defined USE_NVIDIA_CUVS
+#if defined(USE_NVIDIA_CUVS) && !defined(FAISS_CUVS_NO_CAGRA)
 #include <faiss/gpu/GpuIndexBinaryCagra.h>
 #include <faiss/gpu/GpuIndexCagra.h>
 #endif
@@ -92,7 +92,7 @@ Index* ToCPUCloner::clone_Index(const Index* index) {
         // (inverse op of ToGpuClonerMultiple)
 
     }
-#if defined USE_NVIDIA_CUVS
+#if defined(USE_NVIDIA_CUVS) && !defined(FAISS_CUVS_NO_CAGRA)
     else if (auto icg = dynamic_cast<const GpuIndexCagra*>(index)) {
         IndexHNSWCagra* res = new IndexHNSWCagra();
         if (icg->get_numeric_type() != faiss::NumericType::Float32) {
@@ -150,8 +150,8 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         GpuIndexFlatConfig config;
         config.device = device;
         config.useFloat16 = true;
-        FAISS_THROW_IF_NOT_MSG(
-                !use_cuvs, "this type of index is not implemented for cuVS");
+        FAISS_THROW_IF_MSG(
+                use_cuvs, "this type of index is not implemented for cuVS");
         GpuIndexFlat* gif = new GpuIndexFlat(
                 provider, index->d, index->metric_type, config);
         // transfer data by blocks
@@ -187,8 +187,8 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.device = device;
         config.indicesOptions = indicesOptions;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
-        FAISS_THROW_IF_NOT_MSG(
-                !use_cuvs, "this type of index is not implemented for cuVS");
+        config.use_cuvs = use_cuvs;
+        config.allowCpuCoarseQuantizer = allowCpuCoarseQuantizer;
 
         GpuIndexIVFScalarQuantizer* res = new GpuIndexIVFScalarQuantizer(
                 provider,
@@ -233,7 +233,7 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
 
         return res;
     }
-#if defined USE_NVIDIA_CUVS
+#if defined(USE_NVIDIA_CUVS) && !defined(FAISS_CUVS_NO_CAGRA)
     else if (auto icg = dynamic_cast<const faiss::IndexHNSWCagra*>(index)) {
         GpuIndexCagraConfig config;
         config.device = device;
@@ -533,7 +533,7 @@ faiss::IndexBinary* index_binary_gpu_to_cpu(
         ii->copyTo(ret);
         return ret;
     }
-#if defined USE_NVIDIA_CUVS
+#if defined(USE_NVIDIA_CUVS) && !defined(FAISS_CUVS_NO_CAGRA)
     else if (auto ii = dynamic_cast<const GpuIndexBinaryCagra*>(gpu_index)) {
         IndexBinaryHNSWCagra* ret = new IndexBinaryHNSWCagra();
         ii->copyTo(ret);
@@ -549,13 +549,13 @@ faiss::IndexBinary* index_binary_cpu_to_gpu(
         GpuResourcesProvider* provider,
         int device,
         const faiss::IndexBinary* index,
-        const GpuClonerOptions* options) {
+        const GpuClonerOptions* /*options*/) {
     if (auto ii = dynamic_cast<const IndexBinaryFlat*>(index)) {
         GpuIndexBinaryFlatConfig config;
         config.device = device;
         return new GpuIndexBinaryFlat(provider, ii, config);
     }
-#if defined USE_NVIDIA_CUVS
+#if defined(USE_NVIDIA_CUVS) && !defined(FAISS_CUVS_NO_CAGRA)
     else if (
             auto ii = dynamic_cast<const faiss::IndexBinaryHNSWCagra*>(index)) {
         GpuIndexCagraConfig config;

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

@@ -177,6 +177,33 @@ class GpuIndex : public faiss::Index {
         }
     };
 
+    virtual void addImplPrecomputed_(
+            idx_t n,
+            const float* x,
+            const idx_t* ids,
+            const idx_t* precomputed_idx);
+
+    virtual void addImplPrecomputed_ex_(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            const idx_t* ids,
+            const idx_t* precomputed_idx);
+
+    /// Handles paged adds if the add set is too large, passes to
+    /// addImpl_ to actually perform the add for the current page
+    void addPaged_(
+            idx_t n,
+            const float* x,
+            const idx_t* ids,
+            const idx_t* precomputed_idx = nullptr);
+    void addPaged_ex_(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            const idx_t* ids,
+            const idx_t* precomputed_idx = nullptr);
+
     /// Overridden to actually perform the search
     /// All data is guaranteed to be resident on our device
     virtual void searchImpl_(
@@ -209,22 +236,18 @@ class GpuIndex : public faiss::Index {
     }
 
    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_(idx_t n, const float* x, const idx_t* ids);
-    void addPaged_ex_(
-            idx_t n,
-            const void* x,
-            NumericType numeric_type,
-            const idx_t* ids);
-
     /// Calls addImpl_ for a single page of GPU-resident data
-    void addPage_(idx_t n, const float* x, const idx_t* ids);
+    void addPage_(
+            idx_t n,
+            const float* x,
+            const idx_t* ids,
+            const idx_t* precomputed_idx = nullptr);
     void addPage_ex_(
             idx_t n,
             const void* x,
             NumericType numeric_type,
-            const idx_t* ids);
+            const idx_t* ids,
+            const idx_t* precomputed_idx = nullptr);
 
     /// Calls searchImpl_ for a single page of GPU-resident data
     void searchNonPaged_(

data/vendor/faiss/faiss/gpu/GpuIndexCagra.h

@@ -284,6 +284,46 @@ struct GpuIndexCagra : public GpuIndex {
     /// in the index instance
     void copyTo(faiss::IndexHNSWCagra* index) const;
 
+    /// Train CAGRA using multiple GPUs by sharding the dataset.
+    /// Uses cuVS native SNMG (single-node multi-GPU) CAGRA build.
+    /// Each device builds one shard in parallel via OpenMP.
+    /// Float32 only. After training, call copyTo() to produce a CPU
+    /// IndexHNSWCagra with full HNSW upper levels.
+    /// The training data pointer must remain valid until copyTo() completes.
+    /// stitch_mode: 0=CPU HNSW (Approach C), 1=GPU brute-force (Approach B)
+    void trainMultiGpu(
+            idx_t n,
+            const float* x,
+            std::vector<GpuResourcesProvider*>& providers,
+            std::vector<int>& devices,
+            idx_t stitch_per_shard = 0,
+            int stitch_k = 2,
+            int stitch_mode = 0);
+
+    /// Build a unified CAGRA graph using cuVS all_neighbors
+    /// (multi-GPU kNN graph construction with overlapping clusters) followed
+    /// by cagra::optimize (graph pruning). Produces a single unified graph
+    /// without stitching. The training data pointer must remain valid until
+    /// copyTo() completes.
+    /// build_algo: 0=NN-descent (default), 1=brute-force, 2=IVF-PQ
+    /// refinement_rate: IVF-PQ refinement multiplier (only used when
+    /// build_algo==2). Higher values trade build time for recall; the cuVS
+    /// default is 2.0.
+    /// ivfpq_search_batch: cap the IVF-PQ search `max_internal_batch_size` used
+    /// during the all_neighbors kNN build (build_algo==2). 0 = cuVS default
+    /// (128*1024), which can OOM at 100M; a smaller value (e.g. 8192) bounds
+    /// the GPU search workspace with no effect on results (recall-neutral).
+    void trainAllNeighbors(
+            idx_t n,
+            const float* x,
+            std::vector<int>& devices,
+            int n_clusters = 0,
+            int overlap_factor = 0,
+            bool multi_gpu_optimize = false,
+            int build_algo = 0,
+            float refinement_rate = 2.0f,
+            int ivfpq_search_batch = 0);
+
     void reset() override;
 
     std::vector<idx_t> get_knngraph() const;
@@ -317,6 +357,8 @@ struct GpuIndexCagra : public GpuIndex {
             idx_t* labels,
             const SearchParameters* search_params) const override;
 
+    void copyToMultiGpu_(faiss::IndexHNSWCagra* index) const;
+
     /// Our configuration options
     const GpuIndexCagraConfig cagraConfig_;
 
@@ -329,6 +371,11 @@ struct GpuIndexCagra : public GpuIndex {
             std::shared_ptr<CuvsCagra<half>>,
             std::shared_ptr<CuvsCagra<int8_t>>>
             index_;
+
+    /// Multi-GPU state: populated by trainMultiGpu(), used by copyTo()
+    std::vector<idx_t> merged_knngraph_;
+    idx_t merged_knngraph_degree_ = 0;
+    const float* multi_gpu_dataset_ = nullptr;
 };
 
 } // namespace gpu

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

@@ -98,6 +98,17 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface {
     /// debugging purposes.
     virtual std::vector<idx_t> getListIndices(idx_t listId) const;
 
+    /// Add vectors with precomputed IVF list assignments.
+    /// Same naming and argument order as IndexIVF::add_core.
+    /// `x`, `xids`, and `precomputed_idx` can be resident on the CPU or any
+    /// GPU; copies are performed as needed.
+    void add_core(
+            idx_t n,
+            const float* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx,
+            void* inverted_list_context = nullptr);
+
     void search_preassigned(
             idx_t n,
             const float* x,
@@ -133,6 +144,12 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface {
     /// Called from GpuIndex for add/add_with_ids
     void addImpl_(idx_t n, const float* x, const idx_t* ids) override;
 
+    void addImplPrecomputed_(
+            idx_t n,
+            const float* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx) override;
+
     /// Called from GpuIndex for search
     void searchImpl_(
             idx_t n,

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

@@ -93,6 +93,22 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
     /// Validates index SQ parameters
     void verifySQSettings_() const;
 
+    /// Whether this index configuration can use cuVS IVF-SQ
+    bool shouldUseCuvs_() const;
+
+    /// Construct the backing IVF implementation
+    void setIndex_(
+            GpuResources* resources,
+            int dim,
+            idx_t nlist,
+            faiss::MetricType metric,
+            float metricArg,
+            bool useResidual,
+            faiss::ScalarQuantizer* scalarQ,
+            bool interleavedLayout,
+            IndicesOptions indicesOptions,
+            MemorySpace space);
+
     /// Called from train to handle SQ residual training
     void trainResiduals_(idx_t n, const float* x);
 

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

@@ -33,7 +33,8 @@
 
 #if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <cuda/memory_resource>
+#include <optional>
 #endif
 
 namespace faiss {
@@ -163,7 +164,7 @@ struct AllocRequest : public AllocInfo {
     size_t size = 0;
 
 #if defined USE_NVIDIA_CUVS
-    rmm::mr::device_memory_resource* mr = nullptr;
+    std::optional<cuda::mr::any_resource<cuda::mr::device_accessible>> mr;
 #endif
 };
 

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

@@ -92,8 +92,8 @@ std::string allocsToString(const std::unordered_map<void*, AllocRequest>& map) {
 StandardGpuResourcesImpl::StandardGpuResourcesImpl()
         :
 #if defined USE_NVIDIA_CUVS
-          mmr_(new rmm::mr::managed_memory_resource),
-          pmr_(new rmm::mr::pinned_host_memory_resource),
+          mmr_{},
+          pmr_{},
 #endif
           pinnedMemAlloc_(nullptr),
           pinnedMemAllocSize_(0),
@@ -164,7 +164,7 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
 
     if (pinnedMemAlloc_) {
 #if defined USE_NVIDIA_CUVS
-        pmr_->deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
+        pmr_.deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
 #else
         auto err = cudaFreeHost(pinnedMemAlloc_);
         FAISS_ASSERT_FMT(
@@ -350,7 +350,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
         // pinned memory allocation
         if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
             try {
-                pinnedMemAlloc_ = pmr_->allocate_sync(pinnedMemSize_);
+                pinnedMemAlloc_ = pmr_.allocate_sync(pinnedMemSize_);
             } catch (const std::bad_alloc& rmm_ex) {
                 FAISS_THROW_MSG("CUDA memory allocation error");
             }
@@ -546,10 +546,9 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
     } else if (adjReq.space == MemorySpace::Device) {
 #if defined USE_NVIDIA_CUVS
         try {
-            rmm::mr::device_memory_resource* current_mr =
-                    rmm::mr::get_per_device_resource(
-                            rmm::cuda_device_id{adjReq.device});
-            p = current_mr->allocate(adjReq.stream, adjReq.size);
+            auto current_mr = rmm::mr::get_per_device_resource_ref(
+                    rmm::cuda_device_id{adjReq.device});
+            p = current_mr.allocate(adjReq.stream, adjReq.size);
             adjReq.mr = current_mr;
         } catch (const std::bad_alloc& rmm_ex) {
             FAISS_THROW_MSG("CUDA memory allocation error");
@@ -562,7 +561,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // FIXME: as of CUDA 11, a memory allocation error appears to be
             // presented via cudaGetLastError as well, and needs to be
             // cleared. Just call the function to clear it
-            cudaGetLastError();
+            (void)cudaGetLastError();
 
             std::stringstream ss;
             ss << "StandardGpuResources: alloc fail " << adjReq.toString()
@@ -584,8 +583,8 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // TODO: change this to use the current device resource once RMM has
             // a way to retrieve a "guaranteed" managed memory resource for a
             // device.
-            p = mmr_->allocate(adjReq.stream, adjReq.size);
-            adjReq.mr = mmr_.get();
+            p = mmr_.allocate(adjReq.stream, adjReq.size);
+            adjReq.mr = mmr_;
         } catch (const std::bad_alloc& rmm_ex) {
             FAISS_THROW_MSG("CUDA memory allocation error");
         }
@@ -596,7 +595,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // FIXME: as of CUDA 11, a memory allocation error appears to be
             // presented via cudaGetLastError as well, and needs to be cleared.
             // Just call the function to clear it
-            cudaGetLastError();
+            (void)cudaGetLastError();
 
             std::stringstream ss;
             ss << "StandardGpuResources: alloc fail " << adjReq.toString()

data/vendor/faiss/faiss/gpu/StandardGpuResources.h

@@ -25,7 +25,7 @@
 
 #if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/pinned_host_memory_resource.hpp>
 #endif
 
@@ -171,10 +171,10 @@ class StandardGpuResourcesImpl : public GpuResources {
      */
 
     // managed_memory_resource
-    std::unique_ptr<rmm::mr::device_memory_resource> mmr_;
+    rmm::mr::managed_memory_resource mmr_;
 
     // pinned_host_memory_resource
-    std::unique_ptr<rmm::mr::pinned_host_memory_resource> pmr_;
+    rmm::mr::pinned_host_memory_resource pmr_;
 #endif
 
     /// Pinned memory allocation for use with this GPU

data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp

@@ -42,7 +42,7 @@ using namespace faiss::gpu;
 int main(int argc, char** argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-    cudaProfilerStop();
+    CUDA_VERIFY(cudaProfilerStop());
 
     auto seed = FLAGS_seed != -1 ? FLAGS_seed : time(nullptr);
     printf("using seed %ld\n", seed);

data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp

@@ -31,7 +31,7 @@ DEFINE_bool(reserve_memory, false, "whether or not to pre-reserve memory");
 int main(int argc, char** argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-    cudaProfilerStop();
+    CUDA_VERIFY(cudaProfilerStop());
 
     int dim = FLAGS_dim;
     int numCentroids = FLAGS_centroids;
@@ -72,7 +72,7 @@ int main(int argc, char** argv) {
         }
     }
 
-    cudaDeviceSynchronize();
+    CUDA_VERIFY(cudaDeviceSynchronize());
     CUDA_VERIFY(cudaProfilerStart());
 
     float totalGpuTime = 0.0f;