Allow enabling NVTX markers by downstream projects after install (#610)

Allow downstream projects enable `NVTX` option for the `raft::raft` target, if it hasn't been enabled at the install time.

Authors:
  - Artem M. Chirkin (https://github.com/achirkin)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)
  - Robert Maynard (https://github.com/robertmaynard)
  - AJ Schmidt (https://github.com/ajschmidt8)

URL: #610

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-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-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-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-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-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-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-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-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)
@@ -390,7 +410,8 @@ Imported Targets:
 
 ]=])
 
-set(code_string )
+set(code_string ${nvtx_export_string})
+
 if(RAFT_ENABLE_thrust_DEPENDENCY)
   string(APPEND code_string
   [=[

cpp/include/raft/common/detail/nvtx.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -140,7 +140,7 @@ struct domain_store {
   /* If `Domain::name` does not exist, this default instance is used and throws the error. */
   static_assert(sizeof(Domain) != sizeof(Domain),
                 "Type used to identify a domain must contain a static member 'char const* name'");
-  static inline nvtxDomainHandle_t const kValue = nullptr;
+  static inline auto value() -> const nvtxDomainHandle_t { return nullptr; }
 };
 
 template <typename Domain>
@@ -150,7 +150,12 @@ struct domain_store<
   std::enable_if_t<
     std::is_same<char const*, typename std::decay<decltype(Domain::name)>::type>::value,
     Domain>> {
-  static inline nvtxDomainHandle_t const kValue = nvtxDomainCreateA(Domain::name);
+  static inline auto value() -> const nvtxDomainHandle_t
+  {
+    // NB: static modifier ensures the domain is created only once
+    static const nvtxDomainHandle_t kValue = nvtxDomainCreateA(Domain::name);
+    return kValue;
+  }
 };
 
 template <typename Domain>
@@ -163,7 +168,7 @@ inline void push_range_name(const char* name)
   event_attrib.color                 = generate_next_color(name);
   event_attrib.messageType           = NVTX_MESSAGE_TYPE_ASCII;
   event_attrib.message.ascii         = name;
-  nvtxDomainRangePushEx(domain_store<Domain>::kValue, &event_attrib);
+  nvtxDomainRangePushEx(domain_store<Domain>::value(), &event_attrib);
 }
 
 template <typename Domain, typename... Args>
@@ -183,7 +188,7 @@ inline void push_range(const char* format, Args... args)
 template <typename Domain>
 inline void pop_range()
 {
-  nvtxDomainRangePop(domain_store<Domain>::kValue);
+  nvtxDomainRangePop(domain_store<Domain>::value());
 }
 
 #else  // NVTX_ENABLED