Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mnorris11 wants to merge 2 commits into
base: main
Choose a base branch
from
+123 −13
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dc3502f
Add device-memory overload for cagra::helpers::optimize
Aug 6, 2026
1fe0cff
Trim optimize() doc and inline comments
Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 31 additions & 2 deletions cpp/include/cuvs/neighbors/cagra.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4641,7 +4641,7 @@ std::pair<size_t, size_t> cagra_build_mem_usage(raft::resources const& res,
/**
* @brief Optimize a KNN graph into a CAGRA graph.
*
* This function optimizes a k-NN graph to create a CAGRA graph.
* This function optimizes a host-side k-NN graph to create a CAGRA graph.
* The input/output graphs must be on host memory.
*
* Usage example:
Expand All @@ -4656,10 +4656,39 @@ 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);

@mfoerste4 mfoerste4 Aug 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this parameter really required for your workflow? The mst path should be functionally correct, but is not optimized wrt. speed and memory consumption. I would rather not see it in the public API. @achirkin , what is your opinion on this?

@achirkin achirkin Aug 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this parameter will be a big problem since the default is false, but I agree unless you sure you need this feature it will be easier for us to support the API without it.

@cjnolet cjnolet Aug 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do absolutely want the guarantee_connectivity to be exposed to the public. The reason this parameter was introduced was not for testing, but rather because we came to realization that some types of neighborhood graphs, such as the RNG graph which underlines assumptions in Vamana, implicitly assume the graph is connected. This can obviously greatly impact the navigability since NSG-style graphs tend to start with an all-neighbors knn graph (which is not guaranteed to be connected).


/**
* @brief Optimize a KNN graph into a CAGRA graph.
*
* This function optimizes a device-side k-NN graph to create a CAGRA graph.
* The input/output graphs must be on device memory.
*
* 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
Expand Down
7 changes: 5 additions & 2 deletions cpp/src/neighbors/cagra.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
13 changes: 11 additions & 2 deletions cpp/src/neighbors/cagra_optimize.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
20 changes: 13 additions & 7 deletions cpp/src/neighbors/detail/cagra/cagra_build.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1941,22 +1941,28 @@ 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));

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()),
Expand Down
63 changes: 63 additions & 0 deletions cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down Expand Up @@ -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
Loading