Skip to content

Expand NVTX annotations in CAGRA build and HNSW #711

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

Merged
rapids-bot merged 1 commit into from
Feb 20, 2025
Merged

Expand NVTX annotations in CAGRA build and HNSW #711

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
1 change: 1 addition & 0 deletions cpp/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -636,6 +636,7 @@ SECTIONS
if(CUVS_NVTX)
# This enables NVTX within the project with no option to disable it downstream.
target_link_libraries(cuvs PUBLIC CUDA::nvtx3)
target_compile_definitions(cuvs_objs PUBLIC NVTX_ENABLED)
target_compile_definitions(cuvs PUBLIC NVTX_ENABLED)

target_link_libraries(cuvs-cagra-search PUBLIC CUDA::nvtx3)
Expand Down
15 changes: 15 additions & 0 deletions cpp/src/neighbors/detail/cagra/graph_core.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#include <raft/core/resources.hpp>

// TODO: This shouldn't be invoking anything from spatial/knn
#include "../../../core/nvtx.hpp"
#include "../ann_utils.cuh"

#include <raft/util/bitonic_sort.cuh>
Expand Down Expand Up @@ -1086,6 +1087,8 @@ void optimize(
const uint64_t graph_size = new_graph.extent(0);
auto input_graph_ptr = knn_graph.data_handle();
auto output_graph_ptr = new_graph.data_handle();
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope(
"cagra::graph::optimize(%zu, %zu, %u)", graph_size, input_graph_degree, output_graph_degree);

// MST optimization
auto mst_graph = raft::make_host_matrix<IdxT, int64_t, raft::row_major>(0, 0);
Expand All @@ -1096,6 +1099,8 @@ void optimize(
mst_graph_num_edges_ptr[i] = 0;
}
if (guarantee_connectivity) {
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
"cagra::graph::optimize/check_connectivity");
mst_graph =
raft::make_host_matrix<IdxT, int64_t, raft::row_major>(graph_size, output_graph_degree);
RAFT_LOG_INFO("MST optimization is used to guarantee graph connectivity.");
Expand All @@ -1110,6 +1115,8 @@ void optimize(
}

{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
"cagra::graph::optimize/prune");
//
// Prune kNN graph
//
Expand Down Expand Up @@ -1270,6 +1277,8 @@ void optimize(
auto rev_graph_count = raft::make_host_vector<uint32_t, int64_t>(graph_size);

{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
"cagra::graph::optimize/reverse");
//
// Make reverse graph
//
Expand Down Expand Up @@ -1335,6 +1344,8 @@ void optimize(
}

{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
"cagra::graph::optimize/combine");
//
// Create search graphs from MST and pruned and reverse graphs
//
Expand Down Expand Up @@ -1435,6 +1446,8 @@ void optimize(

// Check number of incoming edges
{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
"cagra::graph::optimize/check_edges");
auto in_edge_count = raft::make_host_vector<uint32_t, int64_t>(graph_size);
auto in_edge_count_ptr = in_edge_count.data_handle();
#pragma omp parallel for
Expand Down Expand Up @@ -1477,6 +1490,8 @@ void optimize(

// Check duplication and out-of-range indices
{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
"cagra::graph::optimize/check_duplicates");
uint64_t num_dup = 0;
uint64_t num_oor = 0;
#pragma omp parallel for reduction(+ : num_dup) reduction(+ : num_oor)
Expand Down
18 changes: 10 additions & 8 deletions cpp/src/neighbors/hnsw.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,20 +15,22 @@
*/

#include "detail/hnsw.hpp"
#include "../core/nvtx.hpp"
#include <cstdint>
#include <cuvs/neighbors/hnsw.hpp>
#include <sys/types.h>

namespace cuvs::neighbors::hnsw {

#define CUVS_INST_HNSW_FROM_CAGRA(T) \
std::unique_ptr<index<T>> from_cagra( \
raft::resources const& res, \
const index_params& params, \
const cuvs::neighbors::cagra::index<T, uint32_t>& cagra_index, \
std::optional<raft::host_matrix_view<const T, int64_t, raft::row_major>> dataset) \
{ \
return detail::from_cagra<T>(res, params, cagra_index, dataset); \
#define CUVS_INST_HNSW_FROM_CAGRA(T) \
std::unique_ptr<index<T>> from_cagra( \
raft::resources const& res, \
const index_params& params, \
const cuvs::neighbors::cagra::index<T, uint32_t>& cagra_index, \
std::optional<raft::host_matrix_view<const T, int64_t, raft::row_major>> dataset) \
{ \
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope("hnsw::from_cagra"); \
return detail::from_cagra<T>(res, params, cagra_index, dataset); \
}

CUVS_INST_HNSW_FROM_CAGRA(float);
Expand Down