From 461b289cfbe33711a132d5a73980bddab9fe7ce6 Mon Sep 17 00:00:00 2001
From: Tianyu Liu <kingcrimsontianyu@gmail.com>
Date: Sun, 27 Oct 2024 23:13:32 -0400
Subject: [PATCH] Improve implementation a bit

---
 cpp/include/kvikio/utils.hpp | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/cpp/include/kvikio/utils.hpp b/cpp/include/kvikio/utils.hpp
index ed32da58f4..f6c73c4e28 100644
--- 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