From 930ae6bb45dc40fc5e42d45e1f4032cdae3b0723 Mon Sep 17 00:00:00 2001
From: Vinay D <vinayd@nvidia.com>
Date: Tue, 5 Apr 2022 09:40:53 +0530
Subject: [PATCH] Updating cmake files to enable NVTX markers for cuML build

---
 cpp/CMakeLists.txt       | 5 +++++
 cpp/bench/CMakeLists.txt | 5 +++++
 cpp/test/CMakeLists.txt  | 5 +++++
 3 files changed, 15 insertions(+)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 9467423c78..3974274b3b 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -523,6 +523,11 @@ if(BUILD_CUML_CPP_LIBRARY)
       $<$<BOOL:${ENABLE_CUMLPRIMS_MG}>:cumlprims_mg::cumlprims_mg>
   )
 
+  if(NVTX)
+    target_link_libraries(${CUML_CPP_TARGET} PRIVATE nvToolsExt)
+    target_compile_options(${CUML_CPP_TARGET} PRIVATE "-DNVTX_ENABLED")
+  endif()
+
   # If we export the libdmlc symbols, they can lead to weird crashes with other
   # libraries that use libdmlc. This just hides the symbols internally.
   target_link_options(${CUML_CPP_TARGET} PRIVATE "-Wl,--exclude-libs,libdmlc.a")
diff --git a/cpp/bench/CMakeLists.txt b/cpp/bench/CMakeLists.txt
index 3647f0ed8a..82dbf8844f 100644
--- a/cpp/bench/CMakeLists.txt
+++ b/cpp/bench/CMakeLists.txt
@@ -104,6 +104,11 @@ if(BUILD_CUML_PRIMS_BENCH)
       raft::distance
   )
 
+  if(NVTX)
+    target_compile_options(${PRIMS_BENCH_TARGET} PRIVATE "-DNVTX_ENABLED")
+    target_link_libraries(${PRIMS_BENCH_TARGET} PRIVATE nvToolsExt)
+  endif()
+
   target_include_directories(${PRIMS_BENCH_TARGET}
     PRIVATE
       $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../src_prims>
diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt
index cc4b295bf0..4e35f66282 100644
--- a/cpp/test/CMakeLists.txt
+++ b/cpp/test/CMakeLists.txt
@@ -57,6 +57,11 @@ function(ConfigureTest)
     $<TARGET_NAME_IF_EXISTS:conda_env>
   )
 
+  if(NVTX)
+    target_compile_options(${TEST_NAME} PRIVATE "-DNVTX_ENABLED")
+    target_link_libraries(${TEST_NAME} PRIVATE nvToolsExt)
+  endif()
+
   target_compile_options(${TEST_NAME}
         PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:${CUML_CXX_FLAGS}>"
                 "$<$<COMPILE_LANGUAGE:CUDA>:${CUML_CUDA_FLAGS}>"