Enable CAGRA index building without adding dataset to the index (#2126)

We can build a CAGRA graph even for datasets that do not fit GPU mem. The IVF-PQ build method only requires that the temporary IVF-PQ index (that we use for creating the knn-graph) fits the GPU. But once the CAGRA graph is constructed, we try to initialize the CAGRA index which would copy the dataset to device memory. 

This PR adds a build flag, that would disable copying the dataset into the index. That enables building CAGRA graph for large datasets, and also allows users to customize allocator used for storing the data.

Authors:
  - Tamas Bela Feher (https://github.com/tfeher)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #2126

cpp/bench/ann/src/raft/raft_cagra_wrapper.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -159,7 +159,8 @@ void RaftCagra<T, IdxT>::build(const T* dataset, size_t nrow, cudaStream_t strea
                                                     index_params_.nn_descent_params,
                                                     index_params_.ivf_pq_refine_rate,
                                                     index_params_.ivf_pq_build_params,
-                                                    index_params_.ivf_pq_search_params)));
+                                                    index_params_.ivf_pq_search_params,
+                                                    false)));
 
   handle_.stream_wait(stream);  // RAFT stream -> bench stream
 }

cpp/include/raft/neighbors/cagra_types.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -145,7 +145,7 @@ struct index : ann::index {
   /** Total length of the index (number of vectors). */
   [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT
   {
-    return dataset_view_.extent(0);
+    return dataset_view_.extent(0) ? dataset_view_.extent(0) : graph_view_.extent(0);
   }
 
   /** Dimensionality of the data. */

cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -296,7 +296,8 @@ index<T, IdxT> build(
   std::optional<experimental::nn_descent::index_params> nn_descent_params = std::nullopt,
   std::optional<float> refine_rate                                        = std::nullopt,
   std::optional<ivf_pq::index_params> pq_build_params                     = std::nullopt,
-  std::optional<ivf_pq::search_params> search_params                      = std::nullopt)
+  std::optional<ivf_pq::search_params> search_params                      = std::nullopt,
+  bool construct_index_with_dataset                                       = true)
 {
   size_t intermediate_degree = params.intermediate_graph_degree;
   size_t graph_degree        = params.graph_degree;
@@ -334,12 +335,22 @@ index<T, IdxT> build(
 
   auto cagra_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), graph_degree);
 
+  RAFT_LOG_INFO("optimizing graph");
   optimize<IdxT>(res, knn_graph->view(), cagra_graph.view());
 
   // free intermediate graph before trying to create the index
   knn_graph.reset();
 
+  RAFT_LOG_INFO("Graph optimized, creating index");
   // Construct an index from dataset and optimized knn graph.
-  return index<T, IdxT>(res, params.metric, dataset, raft::make_const_mdspan(cagra_graph.view()));
+  if (construct_index_with_dataset) {
+    return index<T, IdxT>(res, params.metric, dataset, raft::make_const_mdspan(cagra_graph.view()));
+  } else {
+    // We just add the graph. User is expected to update dataset separately. This branch is used
+    // if user needs special control of memory allocations for the dataset.
+    index<T, IdxT> idx(res, params.metric);
+    idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view()));
+    return idx;
+  }
 }
 }  // namespace raft::neighbors::cagra::detail

cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh

@@ -65,8 +65,11 @@ void serialize(raft::resources const& res,
   serialize_scalar(res, os, index_.metric());
   serialize_mdspan(res, os, index_.graph());
 
+  include_dataset &= (index_.dataset().extent(0) > 0);
+
   serialize_scalar(res, os, include_dataset);
   if (include_dataset) {
+    RAFT_LOG_INFO("Saving CAGRA index with dataset");
     auto dataset = index_.dataset();
     // Remove padding before saving the dataset
     auto host_dataset = make_host_matrix<T, int64_t>(dataset.extent(0), dataset.extent(1));
@@ -80,6 +83,8 @@ void serialize(raft::resources const& res,
                                     resource::get_cuda_stream(res)));
     resource::sync_stream(res);
     serialize_mdspan(res, os, host_dataset.view());
+  } else {
+    RAFT_LOG_INFO("Saving CAGRA index WITHOUT dataset");
   }
 }