[REVIEW] Add device-memory overload for cagra::helpers::optimize

cpp/include/cuvs/neighbors/cagra.hpp

@@ -4656,10 +4656,45 @@ std::pair<size_t, size_t> 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<uint32_t, int64_t, raft::row_major> knn_graph,
-              raft::host_matrix_view<uint32_t, int64_t, raft::row_major> new_graph);
+              raft::host_matrix_view<uint32_t, int64_t, raft::row_major> 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<uint32_t, int64_t>(res, N, K_in);
+ *   // Fill d_knn with the KNN graph
+ *   auto d_out = raft::make_device_matrix<uint32_t, int64_t>(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<uint32_t, int64_t, raft::row_major> knn_graph,
+              raft::device_matrix_view<uint32_t, int64_t, raft::row_major> new_graph,
+              bool guarantee_connectivity = false);
 
 }  // namespace helpers
 }  // namespace cagra

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 <typename IdxT = uint32_t,
           typename g_accessor =
+            raft::host_device_accessor<cuda::std::default_accessor<IdxT>, raft::memory_type::host>,
+          typename n_accessor =
             raft::host_device_accessor<cuda::std::default_accessor<IdxT>, raft::memory_type::host>>
 void optimize(
   raft::resources const& res,
   raft::mdspan<IdxT, raft::matrix_extent<int64_t>, raft::row_major, g_accessor> knn_graph,
-  raft::host_matrix_view<IdxT, int64_t, raft::row_major> new_graph,
+  raft::mdspan<IdxT, raft::matrix_extent<int64_t>, raft::row_major, n_accessor> new_graph,
   const bool guarantee_connectivity = false)
 {
   detail::optimize(res, knn_graph, new_graph, guarantee_connectivity);

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<uint32_t, int64_t, raft::row_major> knn_graph,
-              raft::host_matrix_view<uint32_t, int64_t, raft::row_major> new_graph)
+              raft::host_matrix_view<uint32_t, int64_t, raft::row_major> 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<uint32_t, int64_t, raft::row_major> knn_graph,
+              raft::device_matrix_view<uint32_t, int64_t, raft::row_major> new_graph,
+              bool guarantee_connectivity)
+{
+  cuvs::neighbors::cagra::optimize(handle, knn_graph, new_graph, guarantee_connectivity);
 }
 
 }  // namespace cuvs::neighbors::cagra::helpers

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

@@ -1941,22 +1941,34 @@ void build_knn_graph(
 
 template <typename IdxT = uint32_t,
           typename g_accessor =
+            raft::host_device_accessor<cuda::std::default_accessor<IdxT>, raft::memory_type::host>,
+          typename n_accessor =
             raft::host_device_accessor<cuda::std::default_accessor<IdxT>, raft::memory_type::host>>
 void optimize(
   raft::resources const& res,
   raft::mdspan<IdxT, raft::matrix_extent<int64_t>, raft::row_major, g_accessor> knn_graph,
-  raft::host_matrix_view<IdxT, int64_t, raft::row_major> new_graph,
+  raft::mdspan<IdxT, raft::matrix_extent<int64_t>, raft::row_major, n_accessor> new_graph,
   const bool guarantee_connectivity = false)
 {
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;
 
-  auto new_graph_internal = raft::make_host_matrix_view<internal_IdxT, int64_t>(
-    reinterpret_cast<internal_IdxT*>(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<cuda::std::default_accessor<internal_IdxT>, raft::memory_type::host>;
+    raft::host_device_accessor<cuda::std::default_accessor<internal_IdxT>, g_accessor::mem_type>;
+  using n_accessor_internal =
+    raft::host_device_accessor<cuda::std::default_accessor<internal_IdxT>, n_accessor::mem_type>;
+
+  auto new_graph_internal =
+    raft::mdspan<internal_IdxT, raft::matrix_extent<int64_t>, raft::row_major, n_accessor_internal>(
+      reinterpret_cast<internal_IdxT*>(new_graph.data_handle()),
+      new_graph.extent(0),
+      new_graph.extent(1));
+
   auto knn_graph_internal =
     raft::mdspan<internal_IdxT, raft::matrix_extent<int64_t>, raft::row_major, g_accessor_internal>(
       reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),

cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu

@@ -5,7 +5,10 @@
 
 #include <cuvs/neighbors/cagra.hpp>
 #include <gtest/gtest.h>
+#include <raft/core/copy.hpp>
+#include <raft/core/device_mdarray.hpp>
 #include <raft/core/host_mdspan.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 
 // 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<IdxT, int64_t>(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<IdxT, int64_t>(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<IdxT, int64_t>(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<IdxT>(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<IdxT, int64_t>(num_rows, kout);
+  cuvs::neighbors::cagra::helpers::optimize(res, knn_graph_h.view(), expected_h.view());
+
+  auto knn_graph_d = raft::make_device_matrix<IdxT, int64_t>(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<IdxT, int64_t>(res, num_rows, kout);
+  cuvs::neighbors::cagra::helpers::optimize(res, knn_graph_d.view(), actual_d.view());
+
+  auto actual_h = raft::make_host_matrix<IdxT, int64_t>(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