Reduce memory footprint by batching reverse edge computation

bkarsin · bkarsin · commit 2b41a0e05eb3 · 2025-01-15T11:51:48.000Z
diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp
@@ -59,6 +59,8 @@ struct index_params : cuvs::neighbors::index_params {
   float batch_base = 2;
   /** Size of candidate queue structure - should be (2^x)-1 */
   uint32_t queue_size = 127;
+  /** Max batchsize of reverse edge processing (reduces memory footprint) */
+  uint32_t reverse_batchsize = 1000000;
 };
 
 /**
diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh
@@ -104,11 +104,12 @@ void batched_insert_vamana(
       "to 1.0");
     max_batchsize = (int)dataset.extent(0);
   }
-  int insert_iters = (int)(params.vamana_iters);
-  double base      = (double)(params.batch_base);
-  float alpha      = (float)(params.alpha);
-  int visited_size = params.visited_size;
-  int queue_size   = params.queue_size;
+  int insert_iters  = (int)(params.vamana_iters);
+  double base       = (double)(params.batch_base);
+  float alpha       = (float)(params.alpha);
+  int visited_size  = params.visited_size;
+  int queue_size    = params.queue_size;
+  int reverse_batch = params.reverse_batchsize;
 
   if ((visited_size & (visited_size - 1)) != 0) {
     RAFT_LOG_WARN("visited_size must be a power of 2, rounding up.");
@@ -152,39 +153,20 @@ void batched_insert_vamana(
   std::vector<IdxT> insert_order;
   create_insert_permutation<IdxT>(insert_order, (uint32_t)N);
 
-  // Memory needed to sort reverse edges - potentially large memory footprint
-  auto edge_dest =
-    raft::make_device_mdarray<IdxT>(res,
-                                    raft::resource::get_large_workspace_resource(res),
-                                    raft::make_extents<int64_t>(max_batchsize, degree));
-  auto edge_src =
-    raft::make_device_mdarray<IdxT>(res,
-                                    raft::resource::get_large_workspace_resource(res),
-                                    raft::make_extents<int64_t>(max_batchsize, degree));
-
-  size_t temp_storage_bytes = max_batchsize * degree * (2 * sizeof(IdxT));
-  RAFT_LOG_DEBUG("Temp storage needed for sorting (bytes): %lu", temp_storage_bytes);
-  auto temp_sort_storage =
-    raft::make_device_mdarray<IdxT>(res,
-                                    raft::resource::get_large_workspace_resource(res),
-                                    raft::make_extents<int64_t>(2 * max_batchsize, degree));
-
   // Calculate the shared memory sizes of each kernel
   int search_smem_sort_size = 0;
   int prune_smem_sort_size  = 0;
   SELECT_SMEM_SIZES(degree, visited_size);  // Sets above 2 variables to appropriate sizes
 
   // Total dynamic shared memory used by GreedySearch
-  int align_padding = ((((dim-1)/16)+1)*16) - dim;
-  int search_smem_total_size =
-    static_cast<int>(search_smem_sort_size + (dim+align_padding) * sizeof(T) + 
-		     visited_size * sizeof(Node<accT>) +
-                     degree * sizeof(int) + queue_size * sizeof(DistPair<IdxT, accT>));
+  int align_padding          = ((((dim - 1) / 16) + 1) * 16) - dim;
+  int search_smem_total_size = static_cast<int>(
+    search_smem_sort_size + (dim + align_padding) * sizeof(T) + visited_size * sizeof(Node<accT>) +
+    degree * sizeof(int) + queue_size * sizeof(DistPair<IdxT, accT>));
 
   // Total dynamic shared memory size needed by both RobustPrune calls
-  int prune_smem_total_size =
-                    prune_smem_sort_size + (dim+align_padding) * sizeof(T) 
-                    + (degree + visited_size) * sizeof(DistPair<IdxT, accT>);
+  int prune_smem_total_size = prune_smem_sort_size + (dim + align_padding) * sizeof(T) +
+                              (degree + visited_size) * sizeof(DistPair<IdxT, accT>);
 
   RAFT_LOG_DEBUG("Dynamic shared memory usage (bytes): GreedySearch: %d, RobustPrune: %d",
                  search_smem_total_size,
@@ -255,6 +237,15 @@ void batched_insert_vamana(
       int total_edges;
       raft::copy(&total_edges, d_total_edges.data_handle(), 1, stream);
 
+      auto edge_dest =
+        raft::make_device_mdarray<IdxT>(res,
+                                        raft::resource::get_large_workspace_resource(res),
+                                        raft::make_extents<int64_t>(total_edges));
+      auto edge_src =
+        raft::make_device_mdarray<IdxT>(res,
+                                        raft::resource::get_large_workspace_resource(res),
+                                        raft::make_extents<int64_t>(total_edges));
+
       // Create reverse edge list
       create_reverse_edge_list<accT, IdxT>
         <<<num_blocks, blockD, 0, stream>>>(query_list_ptr.data_handle(),
@@ -263,6 +254,24 @@ void batched_insert_vamana(
                                             edge_src.data_handle(),
                                             edge_dest.data_handle());
 
+      void* d_temp_storage      = nullptr;
+      size_t temp_storage_bytes = 0;
+
+      cub::DeviceMergeSort::SortPairs(d_temp_storage,
+                                      temp_storage_bytes,
+                                      edge_dest.data_handle(),
+                                      edge_src.data_handle(),
+                                      total_edges,
+                                      CmpEdge<IdxT>(),
+                                      stream);
+
+      RAFT_LOG_DEBUG("Temp storage needed for sorting (bytes): %lu", temp_storage_bytes);
+
+      auto temp_sort_storage = raft::make_device_mdarray<IdxT>(
+        res,
+        raft::resource::get_large_workspace_resource(res),
+        raft::make_extents<int64_t>(temp_storage_bytes / sizeof(IdxT)));
+
       // Sort to group reverse edges by destination
       cub::DeviceMergeSort::SortPairs(temp_sort_storage.data_handle(),
                                       temp_storage_bytes,
@@ -285,61 +294,72 @@ void batched_insert_vamana(
       thrust::unique_by_key(
         edge_dest_vec.begin(), edge_dest_vec.end(), unique_indices.data_handle());
 
-      // Allocate reverse QueryCandidate list based on number of unique destinations
-      // TODO - Do this in batches to reduce memory footprint / support larger datasets
-      auto reverse_list_ptr = raft::make_device_mdarray<QueryCandidates<IdxT, accT>>(
-        res,
-        raft::resource::get_large_workspace_resource(res),
-        raft::make_extents<int64_t>(unique_dests));
-      auto rev_ids =
-        raft::make_device_mdarray<IdxT>(res,
-                                        raft::resource::get_large_workspace_resource(res),
-                                        raft::make_extents<int64_t>(unique_dests, visited_size));
-      auto rev_dists =
-        raft::make_device_mdarray<accT>(res,
-                                        raft::resource::get_large_workspace_resource(res),
-                                        raft::make_extents<int64_t>(unique_dests, visited_size));
-
-      QueryCandidates<IdxT, accT>* reverse_list =
-        static_cast<QueryCandidates<IdxT, accT>*>(reverse_list_ptr.data_handle());
-
-      init_query_candidate_list<IdxT, accT><<<256, blockD, 0, stream>>>(reverse_list,
-                                                                        rev_ids.data_handle(),
-                                                                        rev_dists.data_handle(),
-                                                                        (int)unique_dests,
-                                                                        visited_size);
-
-      // May need more blocks for reverse list
-      num_blocks = min(maxBlocks, unique_dests);
-
-      // Populate reverse list ids and candidate lists from edge_src and edge_dest
-      populate_reverse_list_struct<T, accT, IdxT>
-        <<<num_blocks, blockD, 0, stream>>>(reverse_list,
-                                            edge_src.data_handle(),
-                                            edge_dest.data_handle(),
-                                            unique_indices.data_handle(),
-                                            unique_dests,
-                                            total_edges,
-                                            dataset.extent(0));
-
-      // Recompute distances (avoided keeping it during sorting)
-      recompute_reverse_dists<T, accT, IdxT>
-        <<<num_blocks, blockD, 0, stream>>>(reverse_list, dataset, unique_dests, metric);
-
-      // Call 2nd RobustPrune on reverse query_list
-      RobustPruneKernel<T, accT, IdxT>
-        <<<num_blocks, blockD, prune_smem_total_size, stream>>>(d_graph.view(),
-                                                                raft::make_const_mdspan(dataset),
-                                                                reverse_list_ptr.data_handle(),
-                                                                unique_dests,
-                                                                visited_size,
-                                                                metric,
-                                                                alpha,
-                                                                prune_smem_sort_size);
-
-      // Write new edge lists to graph
-      write_graph_edges_kernel<accT, IdxT><<<num_blocks, blockD, 0, stream>>>(
-        d_graph.view(), reverse_list_ptr.data_handle(), degree, unique_dests);
+      edge_dest_vec.clear();
+      edge_dest_vec.shrink_to_fit();
+
+      // Batch execution of reverse edge creation/application
+      for (int rev_start = 0; rev_start < (int)unique_dests; rev_start += reverse_batch) {
+        if (rev_start + reverse_batch > (int)unique_dests) {
+          reverse_batch = (int)unique_dests - rev_start;
+        }
+
+        // Allocate reverse QueryCandidate list based on number of unique destinations
+        auto reverse_list_ptr = raft::make_device_mdarray<QueryCandidates<IdxT, accT>>(
+          res,
+          raft::resource::get_large_workspace_resource(res),
+          raft::make_extents<int64_t>(reverse_batch));
+        auto rev_ids =
+          raft::make_device_mdarray<IdxT>(res,
+                                          raft::resource::get_large_workspace_resource(res),
+                                          raft::make_extents<int64_t>(reverse_batch, visited_size));
+        auto rev_dists =
+          raft::make_device_mdarray<accT>(res,
+                                          raft::resource::get_large_workspace_resource(res),
+                                          raft::make_extents<int64_t>(reverse_batch, visited_size));
+
+        QueryCandidates<IdxT, accT>* reverse_list =
+          static_cast<QueryCandidates<IdxT, accT>*>(reverse_list_ptr.data_handle());
+
+        init_query_candidate_list<IdxT, accT><<<256, blockD, 0, stream>>>(reverse_list,
+                                                                          rev_ids.data_handle(),
+                                                                          rev_dists.data_handle(),
+                                                                          (int)reverse_batch,
+                                                                          visited_size);
+
+        // May need more blocks for reverse list
+        num_blocks = min(maxBlocks, reverse_batch);
+
+        // Populate reverse list ids and candidate lists from edge_src and edge_dest
+        populate_reverse_list_struct<T, accT, IdxT>
+          <<<num_blocks, blockD, 0, stream>>>(reverse_list,
+                                              edge_src.data_handle(),
+                                              edge_dest.data_handle(),
+                                              unique_indices.data_handle(),
+                                              unique_dests,
+                                              total_edges,
+                                              dataset.extent(0),
+                                              rev_start,
+                                              reverse_batch);
+
+        // Recompute distances (avoided keeping it during sorting)
+        recompute_reverse_dists<T, accT, IdxT>
+          <<<num_blocks, blockD, 0, stream>>>(reverse_list, dataset, reverse_batch, metric);
+
+        // Call 2nd RobustPrune on reverse query_list
+        RobustPruneKernel<T, accT, IdxT>
+          <<<num_blocks, blockD, prune_smem_total_size, stream>>>(d_graph.view(),
+                                                                  raft::make_const_mdspan(dataset),
+                                                                  reverse_list_ptr.data_handle(),
+                                                                  reverse_batch,
+                                                                  visited_size,
+                                                                  metric,
+                                                                  alpha,
+                                                                  prune_smem_sort_size);
+
+        // Write new edge lists to graph
+        write_graph_edges_kernel<accT, IdxT><<<num_blocks, blockD, 0, stream>>>(
+          d_graph.view(), reverse_list_ptr.data_handle(), degree, reverse_batch);
+      }
 
       start += step_size;
       step_size *= base;
diff --git a/cpp/src/neighbors/detail/vamana/vamana_structs.cuh b/cpp/src/neighbors/detail/vamana/vamana_structs.cuh
@@ -420,22 +420,24 @@ __global__ void populate_reverse_list_struct(QueryCandidates<IdxT, accT>* revers
                                              int* unique_indices,
                                              int unique_dests,
                                              int total_edges,
-                                             int N)
+                                             int N,
+                                             int rev_start,
+                                             int reverse_batch)
 {
-  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < unique_dests;
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < reverse_batch;
        i += blockDim.x * gridDim.x) {
-    reverse_list[i].queryId = edge_dest[unique_indices[i]];
-    if (i == unique_dests - 1) {
-      reverse_list[i].size = total_edges - unique_indices[i];
+    reverse_list[i].queryId = edge_dest[unique_indices[i + rev_start]];
+    if (rev_start + i == unique_dests - 1) {
+      reverse_list[i].size = total_edges - unique_indices[i + rev_start];
     } else {
-      reverse_list[i].size = unique_indices[i + 1] - unique_indices[i];
+      reverse_list[i].size = unique_indices[i + rev_start + 1] - unique_indices[i + rev_start];
     }
     if (reverse_list[i].size > reverse_list[i].maxSize) {
       reverse_list[i].size = reverse_list[i].maxSize;
     }
 
     for (int j = 0; j < reverse_list[i].size; j++) {
-      reverse_list[i].ids[j] = edge_src[unique_indices[i] + j];
+      reverse_list[i].ids[j] = edge_src[unique_indices[i + rev_start] + j];
     }
     for (int j = reverse_list[i].size; j < reverse_list[i].maxSize; j++) {
       reverse_list[i].ids[j]   = raft::upper_bound<IdxT>();
diff --git a/cpp/test/neighbors/ann_vamana.cuh b/cpp/test/neighbors/ann_vamana.cuh
@@ -64,6 +64,7 @@ struct AnnVamanaInputs {
   double max_fraction;
   cuvs::distance::DistanceType metric;
   bool host_dataset;
+  int reverse_batchsize;
 
   // cagra search params
   int n_queries;
@@ -131,10 +132,11 @@ class AnnVamanaTest : public ::testing::TestWithParam<AnnVamanaInputs> {
   void testVamana()
   {
     vamana::index_params index_params;
-    index_params.metric       = ps.metric;
-    index_params.graph_degree = ps.graph_degree;
-    index_params.visited_size = ps.visited_size;
-    index_params.max_fraction = ps.max_fraction;
+    index_params.metric            = ps.metric;
+    index_params.graph_degree      = ps.graph_degree;
+    index_params.visited_size      = ps.visited_size;
+    index_params.max_fraction      = ps.max_fraction;
+    index_params.reverse_batchsize = ps.reverse_batchsize;
 
     auto database_view = raft::make_device_matrix_view<const DataT, int64_t>(
       (const DataT*)database.data(), ps.n_rows, ps.dim);
@@ -264,12 +266,13 @@ inline std::vector<AnnVamanaInputs> generate_inputs()
 {
   std::vector<AnnVamanaInputs> inputs = raft::util::itertools::product<AnnVamanaInputs>(
     {1000},
-    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024}, 
-    {32},                                    // graph degree
-    {64, 128, 256},                          // visited_size
+    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},
+    {32},            // graph degree
+    {64, 128, 256},  // visited_size
     {0.06, 0.1},
     {cuvs::distance::DistanceType::L2Expanded},
     {false},
+    {100, 1000000},
     {100},
     {10},
     {cagra::search_algo::AUTO},
@@ -278,14 +281,15 @@ inline std::vector<AnnVamanaInputs> generate_inputs()
     {1},
     {0.2});
 
-  std::vector<AnnVamanaInputs> inputs2 =
-    raft::util::itertools::product<AnnVamanaInputs>({1000},
-    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},  
+  std::vector<AnnVamanaInputs> inputs2 = raft::util::itertools::product<AnnVamanaInputs>(
+    {1000},
+    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},
     {64},             // graph degree
     {128, 256, 512},  // visited_size
     {0.06, 0.1},
     {cuvs::distance::DistanceType::L2Expanded},
     {false},
+    {100, 1000000},
     {100},
     {10},
     {cagra::search_algo::AUTO},
@@ -295,14 +299,15 @@ inline std::vector<AnnVamanaInputs> generate_inputs()
     {0.2});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
 
-  inputs2 =
-    raft::util::itertools::product<AnnVamanaInputs>({1000},
-    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024}, 
+  inputs2 = raft::util::itertools::product<AnnVamanaInputs>(
+    {1000},
+    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},
     {128},       // graph degree
     {256, 512},  // visited_size
     {0.06, 0.1},
     {cuvs::distance::DistanceType::L2Expanded},
     {false},
+    {100, 1000000},
     {100},
     {10},
     {cagra::search_algo::AUTO},
@@ -312,14 +317,15 @@ inline std::vector<AnnVamanaInputs> generate_inputs()
     {0.2});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
 
-  inputs2 =
-    raft::util::itertools::product<AnnVamanaInputs>({1000},
-    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024}, 
+  inputs2 = raft::util::itertools::product<AnnVamanaInputs>(
+    {1000},
+    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},
     {256},        // graph degree
     {512, 1024},  // visited_size
     {0.06, 0.1},
     {cuvs::distance::DistanceType::L2Expanded},
     {false},
+    {100, 1000000},
     {100},
     {10},
     {cagra::search_algo::AUTO},
@@ -328,7 +334,6 @@ inline std::vector<AnnVamanaInputs> generate_inputs()
     {1},
     {0.2});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
-  
 
   return inputs;
 }
