From f3806f1e5d02aa17e9a647f56065b26a72ec0d2d Mon Sep 17 00:00:00 2001
From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com>
Date: Wed, 15 May 2024 06:46:16 +0200
Subject: [PATCH] Remove the shared state and the mutex from NVTX internals
 (#2310)

Until now, raft has stored a map of NVTX colors (annotation -> color) to avoid using the same color for different annotations and keep using the same color for the same annotations. This map is a shared state.
During an extensive ANN_BENCH throughput testing it has turned out that the mutex guarding the map can sometimes become a bottleneck when the number of concurrent threads is really large (>~ 256). This PR replaces the unordered map and the mutex guarding it with a deterministic hash value of the annotation instead (which is stateless).

**Pros:**
  - No shared state, no mutexes.
  - Assigns the same colors to the same annotations across program runs.

**Cons:**
  - Sometimes different annotations can have the same color (hash collisions).

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

Approvers:
  - Tamas Bela Feher (https://github.com/tfeher)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: https://github.com/rapidsai/raft/pull/2310
---
 cpp/include/raft/core/detail/nvtx.hpp | 40 +++++++++------------------
 1 file changed, 13 insertions(+), 27 deletions(-)

diff --git a/cpp/include/raft/core/detail/nvtx.hpp b/cpp/include/raft/core/detail/nvtx.hpp
index 82db75de84..253d8e5b93 100644
--- a/cpp/include/raft/core/detail/nvtx.hpp
+++ b/cpp/include/raft/core/detail/nvtx.hpp
@@ -24,23 +24,19 @@
 
 #include <cstdint>
 #include <cstdlib>
-#include <mutex>
+#include <limits>
 #include <string>
 #include <type_traits>
-#include <unordered_map>
 #include <vector>
 
 namespace raft::common::nvtx::detail {
 
 /**
- * @brief An internal struct to store associated state with the color
- * generator
+ * @brief An internal struct to to initialize 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_;
+struct color_gen {
+  /** This determines how many bits of the hash to use for the generator */
+  using hash_type = uint16_t;
   /** saturation */
   static inline constexpr float kS = 0.9f;
   /** value */
@@ -109,32 +105,22 @@ inline auto hsv2rgb(float h, float s, float v) -> uint32_t
 /**
  * @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.
+ * It calculates a hash of the passed string and uses the result to generate
+ * distinct yet deterministic colors.
  * 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
+ * @param tag a string used as an input to generate a distinct color.
  * @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;
+  auto x = static_cast<color_gen::hash_type>(std::hash<std::string>{}(tag));
+  auto u = std::numeric_limits<color_gen::hash_type>::max();
+  auto h = static_cast<float>(x) / static_cast<float>(u);
+  h += color_gen::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;
+  return hsv2rgb(h, color_gen::kS, color_gen::kV);
 }
 
 template <typename Domain, typename = Domain>