NVTX range helpers

cpp/CMakeLists.txt

@@ -145,6 +145,7 @@ set(RAFT_LINK_LIBRARIES
         CUDA::cusolver
         CUDA::cudart
         CUDA::cusparse
+        $<$<BOOL:${NVTX}>:CUDA::nvToolsExt>
         rmm::rmm
         cuco::cuco
         )
@@ -153,6 +154,14 @@ target_link_libraries(raft INTERFACE ${RAFT_LINK_LIBRARIES})
 target_link_libraries(raft_distance PUBLIC ${RAFT_LINK_LIBRARIES})
 target_link_libraries(raft_nn PUBLIC ${RAFT_LINK_LIBRARIES} FAISS::FAISS)
 
+set(RAFT_COMPILE_DEFINITIONS
+        $<$<BOOL:${NVTX}>:NVTX_ENABLED>
+        )
+
+target_compile_definitions(raft INTERFACE ${RAFT_COMPILE_DEFINITIONS})
+target_compile_definitions(raft_distance PRIVATE ${RAFT_COMPILE_DEFINITIONS})
+target_compile_definitions(raft_nn PRIVATE ${RAFT_COMPILE_DEFINITIONS})
+
 target_compile_options(raft_distance
         PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:${RAFT_CXX_FLAGS}>"
         "$<$<COMPILE_LANGUAGE:CUDA>:${RAFT_CUDA_FLAGS}>"

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

@@ -0,0 +1,203 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <rmm/cuda_stream_view.hpp>
+
+namespace raft::common::nvtx::detail {
+
+#ifdef NVTX_ENABLED
+
+#include <cstdint>
+#include <cstdlib>
+#include <mutex>
+#include <nvToolsExt.h>
+#include <string>
+#include <type_traits>
+#include <unordered_map>
+
+/**
+ * @brief An internal struct to store associated state with the color
+ * generator
+ */
+struct color_gen_state {
+  /** collection of all tagged colors generated so far */
+  static inline std::unordered_map<std::string, uint32_t> all_colors_;
+  /** mutex for accessing the above map */
+  static inline std::mutex map_mutex_;
+  /** saturation */
+  static inline constexpr float kS = 0.9f;
+  /** value */
+  static inline constexpr float kV = 0.85f;
+  /** golden ratio */
+  static inline constexpr float kPhi = 1.61803f;
+  /** inverse golden ratio */
+  static inline constexpr float kInvPhi = 1.f / kPhi;
+};
+
+// all h, s, v are in range [0, 1]
+// Ref: http://en.wikipedia.org/wiki/HSL_and_HSV#Converting_to_RGB
+inline auto hsv2rgb(float h, float s, float v) -> uint32_t
+{
+  uint32_t out = 0xff000000u;
+  if (s <= 0.0f) { return out; }
+  // convert hue from [0, 1] range to [0, 360]
+  float h_deg = h * 360.f;
+  if (0.f > h_deg || h_deg >= 360.f) h_deg = 0.f;
+  h_deg /= 60.f;
+  int h_range = static_cast<int>(h_deg);
+  float h_mod = h_deg - h_range;
+  float x     = v * (1.f - s);
+  float y     = v * (1.f - (s * h_mod));
+  float z     = v * (1.f - (s * (1.f - h_mod)));
+  float r, g, b;
+  switch (h_range) {
+    case 0:
+      r = v;
+      g = z;
+      b = x;
+      break;
+    case 1:
+      r = y;
+      g = v;
+      b = x;
+      break;
+    case 2:
+      r = x;
+      g = v;
+      b = z;
+      break;
+    case 3:
+      r = x;
+      g = y;
+      b = v;
+      break;
+    case 4:
+      r = z;
+      g = x;
+      b = v;
+      break;
+    case 5:
+    default:
+      r = v;
+      g = x;
+      b = y;
+      break;
+  }
+  out |= (uint32_t(r * 256.f) << 16);
+  out |= (uint32_t(g * 256.f) << 8);
+  out |= uint32_t(b * 256.f);
+  return out;
+}
+
+/**
+ * @brief Helper method to generate 'visually distinct' colors.
+ * Inspired from https://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically/
+ * However, if an associated tag is passed, it will look up in its history for
+ * any generated color against this tag and if found, just returns it, else
+ * generates a new color, assigns a tag to it and stores it for future usage.
+ * Such a thing is very useful for nvtx markers where the ranges associated
+ * with a specific tag should ideally get the same color for the purpose of
+ * visualizing it on nsight-systems timeline.
+ * @param tag look for any previously generated colors with this tag or
+ * associate the currently generated color with it
+ * @return returns 32b RGB integer with alpha channel set of 0xff
+ */
+inline auto generate_next_color(const std::string& tag) -> uint32_t
+{
+  // std::unordered_map<std::string, uint32_t> color_gen_state::all_colors_;
+  // std::mutex color_gen_state::map_mutex_;
+
+  std::lock_guard<std::mutex> guard(color_gen_state::map_mutex_);
+  if (!tag.empty()) {
+    auto itr = color_gen_state::all_colors_.find(tag);
+    if (itr != color_gen_state::all_colors_.end()) { return itr->second; }
+  }
+  auto h = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
+  h += color_gen_state::kInvPhi;
+  if (h >= 1.f) h -= 1.f;
+  auto rgb = hsv2rgb(h, color_gen_state::kS, color_gen_state::kV);
+  if (!tag.empty()) { color_gen_state::all_colors_[tag] = rgb; }
+  return rgb;
+}
+
+template <typename Domain, typename = Domain>
+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;
+};
+
+template <typename Domain>
+struct domain_store<
+  Domain,
+  /* Check if there exists `Domain::name` */
+  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);
+};
+
+template <typename Domain>
+inline void push_range_name(const char* name)
+{
+  nvtxEventAttributes_t event_attrib = {0};
+  event_attrib.version               = NVTX_VERSION;
+  event_attrib.size                  = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+  event_attrib.colorType             = NVTX_COLOR_ARGB;
+  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);
+}
+
+template <typename Domain, typename... Args>
+inline void push_range(const char* format, Args... args)
+{
+  if constexpr (sizeof...(args) > 0) {
+    int length = std::snprintf(nullptr, 0, format, args...);
+    assert(length >= 0);
+    std::vector<char> buf(length + 1);
+    std::snprintf(buf.data(), length + 1, format, args...);
+    push_range_name<Domain>(buf.data());
+  } else {
+    push_range_name<Domain>(format);
+  }
+}
+
+template <typename Domain>
+inline void pop_range()
+{
+  nvtxDomainRangePop(domain_store<Domain>::kValue);
+}
+
+#else  // NVTX_ENABLED
+
+template <typename Domain, typename... Args>
+inline void push_range(const char* format, Args... args)
+{
+}
+
+template <typename Domain>
+inline void pop_range()
+{
+}
+
+#endif  // NVTX_ENABLED
+
+}  // namespace raft::common::nvtx::detail

cpp/include/raft/common/nvtx.hpp

@@ -0,0 +1,155 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "detail/nvtx.hpp"
+#include <optional>
+
+/**
+ * \section Usage
+ *
+ * To add NVTX ranges to your code, use the `nvtx::range` RAII object. A
+ * range begins when the object is created, and ends when the object is
+ * destroyed.
+ *
+ * The example below creates nested NVTX ranges. The range `fun_scope` spans
+ * the whole function, while the range `epoch_scope` spans an iteration
+ * (and appears 5 times in the timeline).
+ * \code{.cpp}
+ * #include <raft/common/nvtx.hpp>
+ * void some_function(int k){
+ *   // Begins a NVTX range with the messsage "some_function_{k}"
+ *   // The range ends when some_function() returns
+ *   common::nvtx::range fun_scope( r{"some_function_%d", k};
+ *
+ *   for(int i = 0; i < 5; i++){
+ *     common::nvtx::range epoch_scope{"epoch-%d", i};
+ *     // some logic inside the loop
+ *   }
+ * }
+ * \endcode
+ *
+ * \section Domains
+ *
+ * All NVTX ranges are assigned to domains. A domain defines a named timeline in
+ * the Nsight Systems view. By default, we put all ranges into a domain `domain::app`
+ * named "application". This is controlled by the template parameter `Domain`.
+ *
+ * The example below defines a domain and uses it in a function.
+ * \code{.cpp}
+ * #include <raft/common/nvtx.hpp>
+ *
+ * struct my_app_domain {
+ *   static constexpr char const* name{"my application"};
+ * }
+ *
+ * void some_function(int k){
+ *   // This NVTX range appears in the timeline named "my application" in Nsight Systems.
+ *   common::nvtx::range<my_app_domain> fun_scope( r{"some_function_%d", k};
+ *   // some logic inside the loop
+ * }
+ * \endcode
+ */
+namespace raft::common::nvtx {
+
+namespace domain {
+
+/** @brief The default NVTX domain. */
+struct app {
+  static constexpr char const* name{"application"};
+};
+
+/** @brief This NVTX domain is supposed to be used within raft.  */
+struct raft {
+  static constexpr char const* name{"raft"};
+};
+
+}  // namespace domain
+
+/**
+ * @brief Push a named NVTX range.
+ *
+ * @tparam Domain optional struct that defines the NVTX domain message;
+ *   You can create a new domain with a custom message as follows:
+ *   \code{.cpp}
+ *      struct custom_domain { static constexpr char const* name{"custom message"}; }
+ *   \endcode
+ *   NB: make sure to use the same domain for `push_range` and `pop_range`.
+ * @param format range name format (accepts printf-style arguments)
+ * @param args the arguments for the printf-style formatting
+ */
+template <typename Domain = domain::app, typename... Args>
+inline void push_range(const char* format, Args... args)
+{
+  detail::push_range<Domain, Args...>(format, args...);
+}
+
+/**
+ * @brief Pop the latest range.
+ *
+ * @tparam Domain optional struct that defines the NVTX domain message;
+ *   You can create a new domain with a custom message as follows:
+ *   \code{.cpp}
+ *      struct custom_domain { static constexpr char const* name{"custom message"}; }
+ *   \endcode
+ *   NB: make sure to use the same domain for `push_range` and `pop_range`.
+ */
+template <typename Domain = domain::app>
+inline void pop_range()
+{
+  detail::pop_range<Domain>();
+}
+
+/**
+ * @brief Push a named NVTX range that would be popped at the end of the object lifetime.
+ *
+ * Refer to \ref Usage for the usage examples.
+ *
+ * @tparam Domain optional struct that defines the NVTX domain message;
+ *   You can create a new domain with a custom message as follows:
+ *   \code{.cpp}
+ *      struct custom_domain { static constexpr char const* name{"custom message"}; }
+ *   \endcode
+ */
+template <typename Domain = domain::app>
+class range {
+ public:
+  /**
+   * Push a named NVTX range.
+   * At the end of the object lifetime, pop the range back.
+   *
+   * @param format range name format (accepts printf-style arguments)
+   * @param args the arguments for the printf-style formatting
+   */
+  template <typename... Args>
+  explicit range(const char* format, Args... args)
+  {
+    push_range<Domain, Args...>(format, args...);
+  }
+
+  ~range() { pop_range<Domain>(); }
+
+  /* This object is not meant to be touched. */
+  range(const range&) = delete;
+  range(range&&)      = delete;
+  auto operator=(const range&) -> range& = delete;
+  auto operator=(range&&) -> range&                = delete;
+  static auto operator new(std::size_t) -> void*   = delete;
+  static auto operator new[](std::size_t) -> void* = delete;
+};
+
+}  // namespace raft::common::nvtx