Improve implementation a bit

rapidsai · Oct 28, 2024 · 461b289 · 461b289
1 parent 3f3868a
commit 461b289
Showing 1 changed file with 13 additions and 8 deletions.
diff --git a/cpp/include/kvikio/utils.hpp b/cpp/include/kvikio/utils.hpp
@@ -291,16 +291,21 @@ struct libkvikio_domain {
 #define CONCAT_HELPER(x, y) x##y
 #define CONCAT(x, y)        CONCAT_HELPER(x, y)
 
+#define REGISTER_STRING(msg)                                             \
+  [&]() -> auto& {                                                       \
+    static nvtx3::registered_string_in<libkvikio_domain> a_reg_str{msg}; \
+    return a_reg_str;                                                    \
+  }()
+
 // Macro overloads of KVIKIO_NVTX_FUNC_RANGE
 #define KVIKIO_NVTX_FUNC_RANGE_1() NVTX3_FUNC_RANGE_IN(libkvikio_domain)
-#define KVIKIO_NVTX_FUNC_RANGE_2(msg, val)                                               \
-  static nvtx3::registered_string_in<libkvikio_domain> CONCAT(a_reg_str, __LINE__){msg}; \
-  nvtx3::scoped_range_in<libkvikio_domain> CONCAT(_kvikio_nvtx_range, __LINE__)          \
-  {                                                                                      \
-    nvtx3::event_attributes                                                              \
-    {                                                                                    \
-      CONCAT(a_reg_str, __LINE__), nvtx3::payload { convert_to_64bit(val) }              \
-    }                                                                                    \
+#define KVIKIO_NVTX_FUNC_RANGE_2(msg, val)                                      \
+  nvtx3::scoped_range_in<libkvikio_domain> CONCAT(_kvikio_nvtx_range, __LINE__) \
+  {                                                                             \
+    nvtx3::event_attributes                                                     \
+    {                                                                           \
+      REGISTER_STRING(msg), nvtx3::payload { convert_to_64bit(val) }            \
+    }                                                                           \
   }
 #define GET_KVIKIO_NVTX_FUNC_RANGE_MACRO(_1, _2, NAME, ...) NAME
 #endif