From fab20fe517174ff1eaa11beb6ca5b6e3c248c795 Mon Sep 17 00:00:00 2001
From: achirkin <achirkin@users.noreply.github.com>
Date: Thu, 20 Feb 2025 08:49:55 +0100
Subject: [PATCH] Expand NVTX annotations in CAGRA build and HNSW

---
 cpp/CMakeLists.txt                            |  1 +
 cpp/src/neighbors/detail/cagra/graph_core.cuh | 15 +++++++++++++++
 cpp/src/neighbors/hnsw.cpp                    | 18 ++++++++++--------
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 869854847..945a18157 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -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)
diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh
index 4adaa7a93..39037b745 100644
--- a/cpp/src/neighbors/detail/cagra/graph_core.cuh
+++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh
@@ -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>
@@ -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);
@@ -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.");
@@ -1110,6 +1115,8 @@ void optimize(
   }
 
   {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/prune");
     //
     // Prune kNN graph
     //
@@ -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
     //
@@ -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
     //
@@ -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
@@ -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)
diff --git a/cpp/src/neighbors/hnsw.cpp b/cpp/src/neighbors/hnsw.cpp
index f165176ec..bf274b0d9 100644
--- a/cpp/src/neighbors/hnsw.cpp
+++ b/cpp/src/neighbors/hnsw.cpp
@@ -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);