diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 26ce110158..7fac255810 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/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(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((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); + + IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); + + minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); + minibatch_size = std::min(minibatch_size, n_rows); return std::make_tuple(minibatch_size, mem_per_row); } diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 0e2ad6d769..76992cc7d5 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/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, 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(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(res, raft::make_extents(0, 0)); - try { - d_rev_graph = raft::make_device_mdarray( - res, default_ws_mr, raft::make_extents(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( - res, large_tmp_mr, raft::make_extents(graph_size, output_graph_degree)); - } + // reverse graph creation will always use the GPU / large workspace resource + auto d_rev_graph = raft::make_device_mdarray( + res, large_tmp_mr, raft::make_extents(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( res, default_ws_mr, raft::make_extents(graph_size)); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 08fcd1f09a..1dd3d7663d 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/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(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((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(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((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( 1, - std::min(max_per_ws / 2, - std::min(params.max_internal_batch_size, n_queries))); + std::min(max_per_ws, std::min(params.max_internal_batch_size, n_queries))); } template