Allow enabling NVTX markers by downstream projects after install

build.sh

@@ -18,7 +18,7 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pyraft pylibraft docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
+VALIDARGS="clean libraft pyraft pylibraft docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --no-force-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
 HELP="$0 [<target> ...] [<flag> ...]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
@@ -41,7 +41,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    --allgpuarch     - build for all supported GPU architectures
    --buildfaiss     - build faiss statically into raft
    --install        - install cmake targets
-   --nvtx           - enable nvtx for profiling support
+   --no-force-nvtx  - disable nvtx (profiling markers), but allow enabling it in downstream projects
    --show_depr_warn - show cmake deprecation warnings
    -h               - print this text
 
@@ -70,7 +70,7 @@ ENABLE_thrust_DEPENDENCY=ON
 ENABLE_ucx_DEPENDENCY=OFF
 ENABLE_nccl_DEPENDENCY=OFF
 
-NVTX=OFF
+NVTX=ON
 CLEAN=0
 UNINSTALL=0
 DISABLE_DEPRECATION_WARNINGS=ON
@@ -157,8 +157,8 @@ fi
 if hasArg --buildfaiss; then
     BUILD_STATIC_FAISS=ON
 fi
-if hasArg --nvtx; then
-    NVTX=ON
+if hasArg --no-force-nvtx; then
+    NVTX=OFF
 fi
 if hasArg --show_depr_warn; then
     DISABLE_DEPRECATION_WARNINGS=OFF
@@ -219,7 +219,7 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -DCMAKE_CUDA_ARCHITECTURES=${RAFT_CMAKE_CUDA_ARCHITECTURES} \
           -DRAFT_COMPILE_LIBRARIES=${COMPILE_LIBRARIES} \
           -DRAFT_ENABLE_NN_DEPENDENCIES=${ENABLE_NN_DEPENDENCIES} \
-          -DNVTX=${NVTX} \
+          -DRAFT_NVTX=${NVTX} \
           -DDISABLE_DEPRECATION_WARNINGS=${DISABLE_DEPRECATION_WARNINGS} \
           -DBUILD_TESTS=${BUILD_TESTS} \
           -DBUILD_BENCH=${BUILD_BENCH} \

conda/recipes/libraft_distance/build.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-./build.sh libraft --install -v --allgpuarch --compile-dist
+./build.sh libraft --install -v --allgpuarch --compile-dist --no-force-nvtx

conda/recipes/libraft_headers/build.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-./build.sh libraft --install -v --allgpuarch
+./build.sh libraft --install -v --allgpuarch --no-force-nvtx

conda/recipes/libraft_nn/build.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-./build.sh libraft --install -v --allgpuarch --compile-nn
+./build.sh libraft --install -v --allgpuarch --compile-nn --no-force-nvtx

conda/recipes/pylibraft/build.sh

@@ -2,4 +2,4 @@
 #!/usr/bin/env bash
 
 # This assumes the script is executed from the root of the repo directory
-./build.sh pylibraft --install
+./build.sh pylibraft --install --no-force-nvtx

conda/recipes/pyraft/build.sh

@@ -3,4 +3,4 @@
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
 # This assumes the script is executed from the root of the repo directory
-./build.sh pyraft --install
+./build.sh pyraft --install --no-force-nvtx

cpp/CMakeLists.txt

@@ -59,7 +59,7 @@ option(CUDA_STATIC_RUNTIME "Statically link the CUDA runtime" OFF)
 option(DETECT_CONDA_ENV "Enable detection of conda environment for dependencies" ON)
 option(DISABLE_DEPRECATION_WARNINGS "Disable deprecaction warnings " ON)
 option(DISABLE_OPENMP "Disable OpenMP" OFF)
-option(NVTX "Enable nvtx markers" OFF)
+option(RAFT_NVTX "Enable nvtx markers" OFF)
 
 option(RAFT_COMPILE_LIBRARIES "Enable building raft shared library instantiations" ${BUILD_TESTS})
 option(RAFT_COMPILE_NN_LIBRARY "Enable building raft nearest neighbors shared library instantiations" OFF)
@@ -91,7 +91,7 @@ message(VERBOSE "RAFT: Disable depreaction warnings " ${DISABLE_DEPRECATION_WARN
 message(VERBOSE "RAFT: Disable OpenMP: ${DISABLE_OPENMP}")
 message(VERBOSE "RAFT: Enable kernel resource usage info: ${CUDA_ENABLE_KERNELINFO}")
 message(VERBOSE "RAFT: Enable lineinfo in nvcc: ${CUDA_ENABLE_LINEINFO}")
-message(VERBOSE "RAFT: Enable nvtx markers: ${NVTX}")
+message(VERBOSE "RAFT: Enable nvtx markers: ${RAFT_NVTX}")
 message(VERBOSE "RAFT: Statically link the CUDA runtime: ${CUDA_STATIC_RUNTIME}")
 
 # Set RMM logging level
@@ -167,7 +167,6 @@ target_include_directories(raft INTERFACE
 # Only CUDA libs, rmm, and mdspan should
 # be used in global target.
 target_link_libraries(raft INTERFACE
-        $<$<BOOL:${NVTX}>:CUDA::nvToolsExt>
         CUDA::cublas
         CUDA::curand
         CUDA::cusolver
@@ -178,7 +177,6 @@ target_link_libraries(raft INTERFACE
         $<$<BOOL:${RAFT_ENABLE_mdspan_DEPENDENCY}>:std::mdspan>
 )
 
-target_compile_definitions(raft INTERFACE $<$<BOOL:${NVTX}>:NVTX_ENABLED>)
 target_compile_features(raft INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
 
 if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_DIST_LIBRARY OR RAFT_COMPILE_NN_LIBRARY)
@@ -192,6 +190,28 @@ SECTIONS
 ]=])
 endif()
 
+##############################################################################
+# - NVTX support in raft -----------------------------------------------------
+
+if (RAFT_NVTX)
+  # This enables NVTX within the project with no option to disable it downstream.
+  target_link_libraries(raft INTERFACE CUDA::nvToolsExt)
+  target_compile_definitions(raft INTERFACE NVTX_ENABLED)
+else()
+  # Allow enable NVTX downstream if not set here.
+  # This creates a new option at build/install time, which is set by default to OFF,
+  # but can be enabled in the dependent project.
+  get_property(nvtx_option_help_string CACHE RAFT_NVTX PROPERTY HELPSTRING)
+  string(CONCAT nvtx_export_string
+  "option(RAFT_NVTX \"" ${nvtx_option_help_string} "\" OFF)"
+  [=[
+
+target_link_libraries(raft::raft INTERFACE $<$<BOOL:${RAFT_NVTX}>:CUDA::nvToolsExt>)
+target_compile_definitions(raft::raft INTERFACE $<$<BOOL:${RAFT_NVTX}>:NVTX_ENABLED>)
+
+  ]=])
+endif()
+
 ##############################################################################
 # - raft_distance ------------------------------------------------------------
 add_library(raft_distance INTERFACE)
@@ -381,7 +401,8 @@ Imported Targets:
 
 ]=])
 
-set(code_string )
+set(code_string ${nvtx_export_string})
+
 if(RAFT_ENABLE_thrust_DEPENDENCY)
   string(APPEND code_string
   [=[