merge commit

Author: jinsolp

.pre-commit-config.yaml

@@ -92,6 +92,12 @@ repos:
                 pass_filenames: false
                 language: python
                 additional_dependencies: [gitpython]
+              - id: cargo-fmt
+                name: cargo-fmt
+                entry: cargo fmt --manifest-path rust/Cargo.toml --all
+                pass_filenames: false
+                files: rust/.*
+                language: rust
       - repo: https://github.com/codespell-project/codespell
         rev: v2.2.2
         hooks:

cpp/cmake/modules/ConfigureCUDA.cmake

@@ -48,13 +48,8 @@ list(APPEND CUVS_CUDA_FLAGS --expt-extended-lambda --expt-relaxed-constexpr)
 list(APPEND CUVS_CXX_FLAGS "-DCUDA_API_PER_THREAD_DEFAULT_STREAM")
 list(APPEND CUVS_CUDA_FLAGS "-DCUDA_API_PER_THREAD_DEFAULT_STREAM")
 # make sure we produce smallest binary size
-list(APPEND CUVS_CUDA_FLAGS -Xfatbin=-compress-all)
-if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA"
-   AND (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9 AND CMAKE_CUDA_COMPILER_VERSION
-                                                                   VERSION_LESS 13.0)
-)
-  list(APPEND CUVS_CUDA_FLAGS -Xfatbin=--compress-level=3)
-endif()
+include(${rapids-cmake-dir}/cuda/enable_fatbin_compression.cmake)
+rapids_cuda_enable_fatbin_compression(VARIABLE CUVS_CUDA_FLAGS TUNE_FOR rapids)
 
 # Option to enable line info in CUDA device compilation to allow introspection when profiling /
 # memchecking

cpp/cmake/patches/faiss.diff

@@ -1,16 +1,9 @@
 {
     "packages" : {
       "faiss" : {
-        "version": "1.11.0",
+        "version": "1.12.0",
         "git_url": "https://github.com/facebookresearch/faiss.git",
-        "git_tag": "v1.11.0",
-        "patches" : [
-          {
-            "file" : "${current_json_dir}/faiss.diff",
-            "issue" : "Apply RAFT breaking changes",
-            "fixed_in" : ""
-          }
-        ]
+        "git_tag": "v1.12.0"
       }
     }
   }

cpp/cmake/patches/faiss_override.json

@@ -1975,6 +1975,9 @@ void serialize_to_hnswlib(
  * @note: When device memory is sufficient, the dataset attached to the returned index is allocated
  * in device memory by default; otherwise, host memory is used automatically.
  *
+ * @note: This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`), and
+ * attempting a logical merge here will throw an error.
+ *
  * Usage example:
  * @code{.cpp}
  *   using namespace raft::neighbors;

cpp/include/cuvs/neighbors/cagra.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -54,7 +54,6 @@ struct cuvsNNDescentIndexParams {
   size_t max_iterations;
   float termination_threshold;
   bool return_distances;
-  size_t n_clusters;
 };
 
 typedef struct cuvsNNDescentIndexParams* cuvsNNDescentIndexParams_t;

cpp/include/cuvs/neighbors/nn_descent.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -37,26 +37,24 @@ namespace cuvs::neighbors::nn_descent {
 
 /**
  * @brief Parameters used to build an nn-descent index
- *
- * `graph_degree`: For an input dataset of dimensions (N, D),
+ * - `graph_degree`: For an input dataset of dimensions (N, D),
  * determines the final dimensions of the all-neighbors knn graph
  * which turns out to be of dimensions (N, graph_degree)
- * `intermediate_graph_degree`: Internally, nn-descent builds an
+ * - `intermediate_graph_degree`: Internally, nn-descent builds an
  * all-neighbors knn graph of dimensions (N, intermediate_graph_degree)
  * before selecting the final `graph_degree` neighbors. It's recommended
  * that `intermediate_graph_degree` >= 1.5 * graph_degree
- * `max_iterations`: The number of iterations that nn-descent will refine
+ * - `max_iterations`: The number of iterations that nn-descent will refine
  * the graph for. More iterations produce a better quality graph at cost of performance
- * `termination_threshold`: The delta at which nn-descent will terminate its iterations
- *
+ * - `termination_threshold`: The delta at which nn-descent will terminate its iterations
+ * - `return_distances`: Boolean to decide whether to return distances array
  */
 struct index_params : cuvs::neighbors::index_params {
-  size_t graph_degree              = 64;      // Degree of output graph.
-  size_t intermediate_graph_degree = 128;     // Degree of input graph for pruning.
-  size_t max_iterations            = 20;      // Number of nn-descent iterations.
-  float termination_threshold      = 0.0001;  // Termination threshold of nn-descent.
-  bool return_distances            = true;    // return distances if true
-  size_t n_clusters                = 1;       // defaults to not using any batching
+  size_t graph_degree              = 64;
+  size_t intermediate_graph_degree = 128;
+  size_t max_iterations            = 20;
+  float termination_threshold      = 0.0001;
+  bool return_distances            = true;
 
   /** @brief Construct NN descent parameters for a specific kNN graph degree
    *
@@ -211,7 +209,10 @@ struct index : cuvs::neighbors::index {
  * @brief Build nn-descent Index with dataset in device memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
  *
  * Usage example:
  * @code{.cpp}
@@ -244,7 +245,10 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in host memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
  *
  * Usage example:
  * @code{.cpp}
@@ -279,7 +283,10 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in device memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
  *
  * Usage example:
  * @code{.cpp}
@@ -312,7 +319,10 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in host memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
  *
  * Usage example:
  * @code{.cpp}
@@ -347,7 +357,11 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in device memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
+ * - BitwiseHamming
  *
  * Usage example:
  * @code{.cpp}
@@ -380,7 +394,11 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in host memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
+ * - BitwiseHamming
  *
  * Usage example:
  * @code{.cpp}
@@ -415,7 +433,11 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in device memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
+ * - BitwiseHamming
  *
  * Usage example:
  * @code{.cpp}
@@ -448,7 +470,11 @@ auto build(raft::resources const& res,
  * @brief Build nn-descent Index with dataset in host memory
  *
  * The following distance metrics are supported:
- * - L2
+ * - L2Expanded
+ * - L2SqrtExpanded
+ * - CosineExpanded
+ * - InnerProduct
+ * - BitwiseHamming
  *
  * Usage example:
  * @code{.cpp}
@@ -479,6 +505,7 @@ auto build(raft::resources const& res,
            std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
              std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;
 
+/** @} */
 /**
  * @brief Test if we have enough GPU memory to run NN descent algorithm.
  *

cpp/include/cuvs/neighbors/nn_descent.hpp

@@ -246,6 +246,30 @@ cuvsError_t cuvsTieredIndexSearch(cuvsResources_t res,
 cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res,
                                   DLManagedTensor* new_vectors,
                                   cuvsTieredIndex_t index);
+/**
+ * @}
+ */
+
+/**
+ * @defgroup tiered_c_index_merge Tiered index merge
+ * @{
+ */
+/**
+ * @brief Merge multiple indices together into a single index
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] index_params Index parameters to use when merging
+ * @param[in] indices pointers to indices to merge together
+ * @param[in] num_indices the number of indices to merge
+ * @param[out] output_index the merged index
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res,
+                                 cuvsTieredIndexParams_t index_params,
+                                 cuvsTieredIndex_t* indices,
+                                 size_t num_indices,
+                                 cuvsTieredIndex_t output_index);
+
 /**
  * @}
  */

cpp/include/cuvs/neighbors/tiered_index.h

@@ -197,4 +197,34 @@ void search(raft::resources const& res,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
             const cuvs::neighbors::filtering::base_filter& sample_filter =
               cuvs::neighbors::filtering::none_sample_filter{});
+
+/** @brief Merge multiple tiered indices into a single index.
+ *
+ * This function merges multiple tiered indices into one, combining both the datasets and graph
+ * structures.
+ *
+ * @param[in] res
+ * @param[in] index_params configure the index building
+ * @param[in] indices A vector of pointers to the indices to merge. All indices should
+ *                    be of the same type, and have datasets with the same dimensionality
+ *
+ * @return A new tiered index containing the merged indices
+ */
+auto merge(raft::resources const& res,
+           const index_params<cagra::index_params>& index_params,
+           const std::vector<tiered_index::index<cagra::index<float, uint32_t>>*>& indices)
+  -> tiered_index::index<cagra::index<float, uint32_t>>;
+
+/** @copydoc merge */
+auto merge(raft::resources const& res,
+           const index_params<ivf_flat::index_params>& index_params,
+           const std::vector<tiered_index::index<ivf_flat::index<float, int64_t>>*>& indices)
+  -> tiered_index::index<ivf_flat::index<float, int64_t>>;
+
+/** @copydoc merge */
+auto merge(raft::resources const& res,
+           const index_params<ivf_pq::index_params>& index_params,
+           const std::vector<tiered_index::index<ivf_pq::typed_index<float, int64_t>>*>& indices)
+  -> tiered_index::index<ivf_pq::typed_index<float, int64_t>>;
+
 }  // namespace cuvs::neighbors::tiered_index