Fix workspace usage

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -204,10 +204,17 @@ auto calc_minibatch_size(const raft::resources& handle,
   // If we need to convert to MathT, space required for the converted batch.
   if (!needs_conversion) { mem_per_row += sizeof(MathT) * dim; }
 
-  // Heuristic: calculate the minibatch size in order to use at most 1GB of memory.
-  IdxT minibatch_size = (1 << 30) / mem_per_row;
-  minibatch_size      = 64 * raft::div_rounding_up_safe(minibatch_size, IdxT{64});
-  minibatch_size      = std::min<IdxT>(minibatch_size, n_rows);
+  // Heuristic: calculate the minibatch size in order to use at most 80% or 512MB workspace memory.
+  // We go below 1GB here as the allocation is mostly done in a single chunk which
+  // is problematic if e.g. a pool allocator manages its own chunks <= 1GB.
+  const auto free_ws_size = raft::resource::get_workspace_free_bytes(handle);
+  const auto available_ws_size =
+    std::min<size_t>((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29);
+
+  IdxT minibatch_size = std::max<IdxT>(IdxT{1}, static_cast<IdxT>(available_ws_size / mem_per_row));
+
+  minibatch_size = raft::round_down_safe<IdxT>(minibatch_size, IdxT{64});
+  minibatch_size = std::min<IdxT>(minibatch_size, n_rows);
   return std::make_tuple(minibatch_size, mem_per_row);
 }
 

cpp/src/neighbors/detail/cagra/graph_core.cuh

@@ -1626,11 +1626,12 @@ void prune_graph_gpu(
   namespace bli                       = cuvs::spatial::knn::detail::utils;
   auto [copy_stream, enable_prefetch] = bli::get_prefetch_stream(res);
   auto workspace_mr                   = raft::resource::get_workspace_resource_ref(res);
+  auto large_workspace_mr             = raft::resource::get_large_workspace_resource_ref(res);
 
   // Single-batch read-only iterator for the input graph (graph_size rows fit in one batch).
   bli::batch_load_iterator<
     raft::mdspan<IdxT, raft::matrix_extent<int64_t>, raft::row_major, AccessorKnnGraph>>
-    d_input_graph(res, knn_graph, graph_size, copy_stream, workspace_mr);
+    d_input_graph(res, knn_graph, graph_size, copy_stream, large_workspace_mr);
   auto input_view = (*d_input_graph).view();
 
   bli::batch_load_iterator<
@@ -1770,19 +1771,10 @@ void optimize(
     prune_graph_gpu<IdxT>(res, knn_graph, new_graph);
   }
 
-  // reverse graph creation will always use the GPU
-  // using default workspace resource for random access
-  // otherwise will be managed memory which is slow upon first access
-  auto d_rev_graph = raft::make_device_mdarray<IdxT>(res, raft::make_extents<int64_t>(0, 0));
-  try {
-    d_rev_graph = raft::make_device_mdarray<IdxT>(
-      res, default_ws_mr, raft::make_extents<int64_t>(graph_size, output_graph_degree));
-  } catch (const std::exception& e) {
-    RAFT_LOG_DEBUG(
-      "Failed to create device matrix for reverse graph, switching to large workspace resource");
-    d_rev_graph = raft::make_device_mdarray<IdxT>(
-      res, large_tmp_mr, raft::make_extents<int64_t>(graph_size, output_graph_degree));
-  }
+  // reverse graph creation will always use the GPU / large workspace resource
+  auto d_rev_graph = raft::make_device_mdarray<IdxT>(
+    res, large_tmp_mr, raft::make_extents<int64_t>(graph_size, output_graph_degree));
+
   // This should use the default workspace resource for random access / atomics
   auto d_rev_graph_count = raft::make_device_mdarray<uint32_t>(
     res, default_ws_mr, raft::make_extents<int64_t>(graph_size));

cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh

@@ -780,7 +780,13 @@ inline auto get_max_fine_batch_size(raft::resources const& res,
     return static_cast<uint64_t>(bs) *
            (other + (is_local_topk_feasible(k, n_probes, bs) ? buffers_fused : buffers_non_fused));
   };
-  auto max_ws_size = raft::resource::get_workspace_free_bytes(res);
+
+  // Use only 40% of the available workspace memory as the coarse batch allocations have not been
+  // allocated yet. Also limit the maximum to 1GB to avoid conflicts in case the workspace is shared
+  // with the large workspace.
+  const auto free_ws_size = raft::resource::get_workspace_free_bytes(res);
+  const auto max_ws_size =
+    std::min<size_t>((free_ws_size * size_t{40}) / size_t{100}, size_t{1} << 30);
   if (ws_size(max_batch_size) > max_ws_size) {
     uint32_t smaller_batch_size = raft::bound_by_power_of_two(max_batch_size);
     // gradually reduce the batch size until we fit into the max size limit.
@@ -837,11 +843,17 @@ inline auto get_max_coarse_batch_size(raft::resources const& res,
   // Transient allocations during coarse search (select_clusters): qc_distances + cluster_dists.
   auto transient_per_query = static_cast<size_t>(n_lists + n_probes) * qc_elem_size;
   auto total_per_query     = persistent_per_query + transient_per_query;
-  auto max_per_ws          = raft::resource::get_workspace_free_bytes(res) / total_per_query;
+
+  // Use only 40% of the available workspace memory as the fine batch allocations have not been
+  // allocated yet. Also limit the maximum to 1GB to avoid conflicts in case the workspace is shared
+  // with the large workspace.
+  const auto free_ws_size = raft::resource::get_workspace_free_bytes(res);
+  const auto available_ws_size =
+    std::min<size_t>((free_ws_size * size_t{40}) / size_t{100}, size_t{1} << 30);
+  auto max_per_ws = available_ws_size / total_per_query;
   return std::max<uint32_t>(
     1,
-    std::min<uint32_t>(max_per_ws / 2,
-                       std::min<uint32_t>(params.max_internal_batch_size, n_queries)));
+    std::min<uint32_t>(max_per_ws, std::min<uint32_t>(params.max_internal_batch_size, n_queries)));
 }
 
 template <typename T, typename IdxT>