[REVIEW] Add a public API for CAGRA graph optimize (rapidsai#1260)

This change exposes CAGRA's graph optimization functionality through a new public API in the `cuvs::neighbors::cagra::helpers` namespace. The optimize function allows users to optimize a custom KNN graphs to create CAGRA search graphs.

Key Changes:

- Added `cpp/include/cuvs/neighbors/cagra_optimize.hpp` with a public API declaration. As cagra_optimize assumes both input and output are host matrices, this PR exposes the API with host inputs only. Given this, users need to explicitly manage device-to-host transfers if KNN graph is not a RAFT host matrix
- Modified previously unused `cpp/src/neighbors/cagra_optimize.cu` with implementation that forwards to internal CAGRA optimize. Removed the version using device matrix for `knn_graph`, which generates runtime segfaults
- Added unit tests for the new API

This PR partially resolves rapidsai#456 and completes one of the steps for rapidsai#1146

Authors:
  - Rui Lan (https://github.com/abc99lr)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: rapidsai#1260

Author: abc99lr

cpp/include/cuvs/neighbors/cagra.hpp

@@ -2686,3 +2686,4 @@ auto distribute(const raft::resources& clique, const std::string& filename)
 }  // namespace cuvs::neighbors::cagra
 
 #include <cuvs/neighbors/cagra_index_wrapper.hpp>
+#include <cuvs/neighbors/cagra_optimize.hpp>

cpp/include/cuvs/neighbors/cagra_optimize.hpp

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/core/mdspan_types.hpp>
+#include <raft/core/resources.hpp>
+
+namespace cuvs::neighbors::cagra::helpers {
+
+/**
+ * @brief Optimize a KNN graph into a CAGRA graph.
+ *
+ * This function optimizes a k-NN graph to create a CAGRA graph.
+ * The input/output graphs must be on host memory.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   raft::resources res;
+ *   auto h_knn = raft::make_host_matrix<uint32_t, int64_t>(N, K_in);
+ *   // Fill h_knn with KNN graph
+ *   auto h_out = raft::make_host_matrix<uint32_t, int64_t>(N, K_out);
+ *   cuvs::neighbors::cagra::helpers::optimize(res, h_knn.view(), h_out.view());
+ * @endcode
+ *
+ * @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]
+ */
+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);
+
+}  // namespace cuvs::neighbors::cagra::helpers

cpp/src/neighbors/cagra_optimize.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,27 +15,15 @@
  */
 
 #include "cagra.cuh"
-#include <cuvs/neighbors/cagra.hpp>
+#include <cuvs/neighbors/cagra_optimize.hpp>
 
-namespace cuvs::neighbors::cagra {
+namespace cuvs::neighbors::cagra::helpers {
 
-void optimize(raft::resources const& handle,
-              raft::device_matrix_view<uint32_t, int64_t, raft::row_major> knn_graph,
-              raft::host_matrix_view<uint32_t, int64_t, raft::row_major> new_graph)
-{
-  cuvs::neighbors::cagra::optimize<
-    uint32_t,
-    raft::host_device_accessor<std::experimental::default_accessor<uint32_t>,
-                               raft::memory_type::device>>(handle, knn_graph, new_graph);
-}
 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)
 {
-  cuvs::neighbors::cagra::optimize<
-    uint32_t,
-    raft::host_device_accessor<std::experimental::default_accessor<uint32_t>,
-                               raft::memory_type::host>>(handle, knn_graph, new_graph);
+  cuvs::neighbors::cagra::optimize(handle, knn_graph, new_graph);
 }
 
-}  // namespace cuvs::neighbors::cagra
+}  // namespace cuvs::neighbors::cagra::helpers

cpp/tests/CMakeLists.txt

@@ -166,6 +166,11 @@ if(BUILD_TESTS)
     1 PERCENT 100
   )
 
+  ConfigureTest(
+    NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST PATH neighbors/ann_cagra/test_optimize_uint32_t.cu GPUS 1
+    PERCENT 100
+  )
+
   ConfigureTest(
     NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST PATH neighbors/ann_cagra/test_half_uint32_t.cu GPUS 1
     PERCENT 100

cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cuvs/neighbors/cagra.hpp>
+#include <gtest/gtest.h>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/core/resources.hpp>
+
+// This test targets public API exposure and basic invariants only (shapes, in-range indices).
+// Detailed optimization correctness is exercised by CAGRA build tests.
+
+namespace {
+
+using IdxT = uint32_t;
+
+// Helper to create a simple synthetic KNN graph (ring-like neighbors)
+auto make_ring_knn_host(int64_t num_rows, int64_t kin)
+{
+  auto knn_graph = raft::make_host_matrix<IdxT, int64_t>(num_rows, kin);
+  for (int64_t i = 0; i < num_rows; ++i) {
+    for (int64_t j = 0; j < kin; ++j) {
+      knn_graph(i, j) = static_cast<IdxT>((i + j + 1) % num_rows);
+    }
+  }
+  return knn_graph;
+}
+
+TEST(CagraOptimize, HostToHostOptimizesGraph)
+{
+  raft::resources res;
+
+  constexpr int64_t num_rows = 8;
+  constexpr int64_t kin      = 8;
+  constexpr int64_t kout     = 4;
+
+  auto knn_graph       = make_ring_knn_host(num_rows, kin);
+  auto optimized_graph = raft::make_host_matrix<IdxT, int64_t>(num_rows, kout);
+
+  // Test the optimize API
+  cuvs::neighbors::cagra::helpers::optimize(res, knn_graph.view(), optimized_graph.view());
+
+  // Check basic invariants
+  ASSERT_EQ(optimized_graph.extent(0), num_rows);
+  ASSERT_EQ(optimized_graph.extent(1), kout);
+
+  // Check that all neighbors are valid indices
+  for (int64_t i = 0; i < num_rows; ++i) {
+    for (int64_t j = 0; j < kout; ++j) {
+      EXPECT_LT(optimized_graph(i, j), static_cast<IdxT>(num_rows));
+    }
+  }
+}
+
+}  // namespace