From 507c1ec7458efb32cbfac2edaca397a4afeecb4f Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Mon, 18 May 2026 13:48:15 +0200 Subject: [PATCH] Fix cagra::optimize modifying the state of raft::resources (#2103) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The new, improved version of CAGRA optimize() function makes use of raft CUDA stream pool resource. Setting the resource (`raft::resource::set_cuda_stream_pool`) affects the state of `raft::resources` and may lead to unintended consequences for a user. Using non-const `raft::resources` will be enforced at compile time in https://github.com/rapidsai/raft/pull/3005 , so this PR fixes the use case ahead of the upstream change by copying the resources handle. Authors: - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Malte Förster (https://github.com/mfoerste4) - Tamas Bela Feher (https://github.com/tfeher) URL: https://github.com/rapidsai/cuvs/pull/2103 --- cpp/src/neighbors/detail/cagra/graph_core.cuh | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 2c21018deb..0e2ad6d769 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -1706,7 +1706,7 @@ template , raft::memory_type::host>> void optimize( - raft::resources const& res, + raft::resources const& res_const, raft::mdspan, raft::row_major, AccessorKnnGraph> knn_graph, raft::mdspan, raft::row_major, AccessorOutputGraph> new_graph, const bool guarantee_connectivity = true, @@ -1715,6 +1715,12 @@ void optimize( RAFT_LOG_DEBUG( "# Pruning kNN graph (size=%lu, degree=%lu)\n", knn_graph.extent(0), knn_graph.extent(1)); + // TODO(achirkin): come up with a reasonable API to initialize a non-empty stream pool. + // raft::resource::set_cuda_stream_pool below modifies the resource, so it cannot be const. + // The optimize() is a heavy function, so copying the resource and creating a private stream pool + // is not a big overhead. + raft::resources res{res_const}; + // large temporary memory for large arrays, e.g. everything >= O(graph_size) auto large_tmp_mr = raft::resource::get_large_workspace_resource_ref(res); // temporary memory for small arrays, e.g. everything <= O(batchsize * graph_degree)