From 1ca817d0c4c793dcb46b57e371ddbfc5dc3a70bb Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Thu, 6 Aug 2026 09:45:29 -0700 Subject: [PATCH] Add device-memory overload for cagra::helpers::optimize `cagra::helpers::optimize` currently accepts host matrices only, so a caller that already holds its k-NN graph in device memory -- for example the output of `all_neighbors::build` -- must copy the graph to host, optimize, and copy back. The device implementation already exists. `graph::optimize` is templated on both mdspan accessors and `make_reverse_graph_gpu` has an `is_device_accessible` fast path (graph_core.cuh:826), and `batch_load_iterator` switches to a zero-copy `kPassthrough` mode for device accessors. That code is unreachable today because `detail::optimize` in cagra_build.cuh erases the caller's accessor: using g_accessor_internal = raft::host_device_accessor, raft::memory_type::host>; so `graph::optimize` is only ever instantiated with host accessors. With host accessors the reverse-graph phase degrades into `graph_degree` separate host column gathers, each with its own H2D copy and a full stream synchronisation. This change: - propagates the caller's memory types through `detail::optimize` instead of erasing them, and makes `new_graph` accessor-generic there and in `cagra::optimize`; - adds a `device_matrix_view` overload of `cagra::helpers::optimize` to the public API; - exposes the existing `guarantee_connectivity` flag on the public overloads, which previously could not be reached from outside. Both public overloads keep their existing signatures via a defaulted argument, so this is source compatible. Measured on 100M x 129d vectors (graph degree 32, intermediate degree 32) on 8x H100, as part of a multi-GPU CAGRA build: the optimize step goes from 119.1s to 1.65s, a 72x reduction, with recall unchanged. At that scale it takes the whole build->serialize pipeline from 8.0 to 5.4 minutes. Tests: adds a device-to-device case and a case asserting the device overload produces the same graph as the host overload for the same input. --- cpp/include/cuvs/neighbors/cagra.hpp | 37 ++++++++++- cpp/src/neighbors/cagra.cuh | 7 ++- cpp/src/neighbors/cagra_optimize.cu | 13 +++- .../neighbors/detail/cagra/cagra_build.cuh | 26 +++++--- .../ann_cagra/test_optimize_uint32_t.cu | 63 +++++++++++++++++++ 5 files changed, 134 insertions(+), 12 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 43ae7a6235..a243808d09 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -4656,10 +4656,45 @@ std::pair cagra_build_mem_usage(raft::resources const& res, * @param[in] handle RAFT resources * @param[in] knn_graph Input KNN graph on host [n_rows, k_in] * @param[out] new_graph Output CAGRA graph on host [n_rows, k_out] + * @param[in] guarantee_connectivity Run the MST pass so the pruned graph is guaranteed + * to be connected */ void optimize(raft::resources const& handle, raft::host_matrix_view knn_graph, - raft::host_matrix_view new_graph); + raft::host_matrix_view new_graph, + bool guarantee_connectivity = false); + +/** + * @brief Optimize a KNN graph into a CAGRA graph without leaving device memory. + * + * Same as the host overload, but both graphs stay in device memory. Pruning, + * reverse-graph construction and the final merge all run on device with no host + * staging. This avoids copying the graph out and back, and avoids serialising the + * reverse-graph phase into `graph_degree` separate host gathers, each with its own + * H2D copy and stream synchronisation. + * + * Prefer this overload when the k-NN graph is already on device, for example the + * output of `all_neighbors::build`. + * + * Usage example: + * @code{.cpp} + * raft::resources res; + * auto d_knn = raft::make_device_matrix(res, N, K_in); + * // Fill d_knn with the KNN graph + * auto d_out = raft::make_device_matrix(res, N, K_out); + * cuvs::neighbors::cagra::helpers::optimize(res, d_knn.view(), d_out.view()); + * @endcode + * + * @param[in] handle RAFT resources + * @param[in] knn_graph Input KNN graph on device [n_rows, k_in] + * @param[out] new_graph Output CAGRA graph on device [n_rows, k_out] + * @param[in] guarantee_connectivity Run the MST pass so the pruned graph is guaranteed + * to be connected + */ +void optimize(raft::resources const& handle, + raft::device_matrix_view knn_graph, + raft::device_matrix_view new_graph, + bool guarantee_connectivity = false); } // namespace helpers } // namespace cagra diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 80e2f2a07e..dfce8c19e1 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -264,15 +264,18 @@ void sort_knn_graph( * @param[in] res raft resources * @param[in] knn_graph a matrix view (host or device) of the input knn graph [n_rows, * knn_graph_degree] - * @param[out] new_graph a host matrix view of the optimized knn graph [n_rows, graph_degree] + * @param[out] new_graph a matrix view (host or device) of the optimized knn graph [n_rows, + * graph_degree] */ template , raft::memory_type::host>, + typename n_accessor = raft::host_device_accessor, raft::memory_type::host>> void optimize( raft::resources const& res, raft::mdspan, raft::row_major, g_accessor> knn_graph, - raft::host_matrix_view new_graph, + raft::mdspan, raft::row_major, n_accessor> new_graph, const bool guarantee_connectivity = false) { detail::optimize(res, knn_graph, new_graph, guarantee_connectivity); diff --git a/cpp/src/neighbors/cagra_optimize.cu b/cpp/src/neighbors/cagra_optimize.cu index 951b850aab..9b54b8c418 100644 --- a/cpp/src/neighbors/cagra_optimize.cu +++ b/cpp/src/neighbors/cagra_optimize.cu @@ -10,9 +10,18 @@ namespace cuvs::neighbors::cagra::helpers { void optimize(raft::resources const& handle, raft::host_matrix_view knn_graph, - raft::host_matrix_view new_graph) + raft::host_matrix_view new_graph, + bool guarantee_connectivity) { - cuvs::neighbors::cagra::optimize(handle, knn_graph, new_graph); + cuvs::neighbors::cagra::optimize(handle, knn_graph, new_graph, guarantee_connectivity); +} + +void optimize(raft::resources const& handle, + raft::device_matrix_view knn_graph, + raft::device_matrix_view new_graph, + bool guarantee_connectivity) +{ + cuvs::neighbors::cagra::optimize(handle, knn_graph, new_graph, guarantee_connectivity); } } // namespace cuvs::neighbors::cagra::helpers diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index c06f9b12e3..4bfe3f0429 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1941,22 +1941,34 @@ void build_knn_graph( template , raft::memory_type::host>, + typename n_accessor = raft::host_device_accessor, raft::memory_type::host>> void optimize( raft::resources const& res, raft::mdspan, raft::row_major, g_accessor> knn_graph, - raft::host_matrix_view new_graph, + raft::mdspan, raft::row_major, n_accessor> new_graph, const bool guarantee_connectivity = false) { using internal_IdxT = typename std::make_unsigned::type; - auto new_graph_internal = raft::make_host_matrix_view( - reinterpret_cast(new_graph.data_handle()), - new_graph.extent(0), - new_graph.extent(1)); - + // Propagate the caller's memory types rather than erasing them to + // `memory_type::host`. Erasing them makes the device-resident branches + // inside graph::optimize unreachable -- the `is_device_accessible` path in + // make_reverse_graph_gpu, and the zero-copy `kPassthrough` mode of + // batch_load_iterator -- so a caller that already holds the graph on device + // is forced through host staging. using g_accessor_internal = - raft::host_device_accessor, raft::memory_type::host>; + raft::host_device_accessor, g_accessor::mem_type>; + using n_accessor_internal = + raft::host_device_accessor, n_accessor::mem_type>; + + auto new_graph_internal = + raft::mdspan, raft::row_major, n_accessor_internal>( + reinterpret_cast(new_graph.data_handle()), + new_graph.extent(0), + new_graph.extent(1)); + auto knn_graph_internal = raft::mdspan, raft::row_major, g_accessor_internal>( reinterpret_cast(knn_graph.data_handle()), diff --git a/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu index 633d876ea3..e1d7dbc43d 100644 --- a/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu @@ -5,7 +5,10 @@ #include #include +#include +#include #include +#include #include // This test targets public API exposure and basic invariants only (shapes, in-range indices). @@ -53,4 +56,64 @@ TEST(CagraOptimize, HostToHostOptimizesGraph) } } +TEST(CagraOptimize, DeviceToDeviceOptimizesGraph) +{ + raft::resources res; + + constexpr int64_t num_rows = 8; + constexpr int64_t kin = 8; + constexpr int64_t kout = 4; + + auto knn_graph_h = make_ring_knn_host(num_rows, kin); + auto knn_graph_d = raft::make_device_matrix(res, num_rows, kin); + raft::copy(res, knn_graph_d.view(), raft::make_const_mdspan(knn_graph_h.view())); + + auto optimized_graph_d = raft::make_device_matrix(res, num_rows, kout); + + cuvs::neighbors::cagra::helpers::optimize(res, knn_graph_d.view(), optimized_graph_d.view()); + + ASSERT_EQ(optimized_graph_d.extent(0), num_rows); + ASSERT_EQ(optimized_graph_d.extent(1), kout); + + auto optimized_graph_h = raft::make_host_matrix(num_rows, kout); + raft::copy(res, optimized_graph_h.view(), raft::make_const_mdspan(optimized_graph_d.view())); + raft::resource::sync_stream(res); + + for (int64_t i = 0; i < num_rows; ++i) { + for (int64_t j = 0; j < kout; ++j) { + EXPECT_LT(optimized_graph_h(i, j), static_cast(num_rows)); + } + } +} + +// The device and host overloads must agree: same input, same optimized graph. +TEST(CagraOptimize, DeviceMatchesHost) +{ + raft::resources res; + + constexpr int64_t num_rows = 64; + constexpr int64_t kin = 16; + constexpr int64_t kout = 8; + + auto knn_graph_h = make_ring_knn_host(num_rows, kin); + + auto expected_h = raft::make_host_matrix(num_rows, kout); + cuvs::neighbors::cagra::helpers::optimize(res, knn_graph_h.view(), expected_h.view()); + + auto knn_graph_d = raft::make_device_matrix(res, num_rows, kin); + raft::copy(res, knn_graph_d.view(), raft::make_const_mdspan(knn_graph_h.view())); + auto actual_d = raft::make_device_matrix(res, num_rows, kout); + cuvs::neighbors::cagra::helpers::optimize(res, knn_graph_d.view(), actual_d.view()); + + auto actual_h = raft::make_host_matrix(num_rows, kout); + raft::copy(res, actual_h.view(), raft::make_const_mdspan(actual_d.view())); + raft::resource::sync_stream(res); + + for (int64_t i = 0; i < num_rows; ++i) { + for (int64_t j = 0; j < kout; ++j) { + EXPECT_EQ(actual_h(i, j), expected_h(i, j)) << "mismatch at (" << i << ", " << j << ")"; + } + } +} + } // namespace