From 21eca7024d4e96c49d18730893a76a83d9b04334 Mon Sep 17 00:00:00 2001
From: Tianyu Liu <kingcrimsontianyu@gmail.com>
Date: Mon, 28 Oct 2024 00:04:23 -0400
Subject: [PATCH] Add NVTX marker

---
 cpp/include/kvikio/file_handle.hpp |  9 ++-------
 cpp/include/kvikio/utils.hpp       | 30 ++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/cpp/include/kvikio/file_handle.hpp b/cpp/include/kvikio/file_handle.hpp
index 4f871744e7..d5d3ade494 100644
--- a/cpp/include/kvikio/file_handle.hpp
+++ b/cpp/include/kvikio/file_handle.hpp
@@ -329,8 +329,6 @@ class FileHandle {
                    std::size_t devPtr_offset,
                    bool sync_default_stream = true)
   {
-    KVIKIO_NVTX_FUNC_RANGE("FileHandle::read()", size);
-
     if (_compat_mode) {
       return detail::posix_device_read(
         _fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
@@ -381,8 +379,6 @@ class FileHandle {
                     std::size_t devPtr_offset,
                     bool sync_default_stream = true)
   {
-    KVIKIO_NVTX_FUNC_RANGE("FileHandle::write()", size);
-
     _nbytes = 0;  // Invalidate the computed file size
 
     if (_compat_mode) {
@@ -438,8 +434,7 @@ class FileHandle {
                                  std::size_t gds_threshold = defaults::gds_threshold(),
                                  bool sync_default_stream  = true)
   {
-    KVIKIO_NVTX_FUNC_RANGE("FileHandle::pread()", size);
-
+    KVIKIO_NVTX_MARKER("FileHandle::pread()", size);
     if (is_host_memory(buf)) {
       auto op = [this](void* hostPtr_base,
                        std::size_t size,
@@ -516,7 +511,7 @@ class FileHandle {
                                   std::size_t gds_threshold = defaults::gds_threshold(),
                                   bool sync_default_stream  = true)
   {
-    KVIKIO_NVTX_FUNC_RANGE("FileHandle::pwrite()", size);
+    KVIKIO_NVTX_MARKER("FileHandle::pwrite()", size);
     if (is_host_memory(buf)) {
       auto op = [this](const void* hostPtr_base,
                        std::size_t size,
diff --git a/cpp/include/kvikio/utils.hpp b/cpp/include/kvikio/utils.hpp
index e9a5d13933..03d85dd5e7 100644
--- a/cpp/include/kvikio/utils.hpp
+++ b/cpp/include/kvikio/utils.hpp
@@ -310,6 +310,11 @@ struct libkvikio_domain {
     }                                                                           \
   }
 #define GET_KVIKIO_NVTX_FUNC_RANGE_MACRO(_1, _2, NAME, ...) NAME
+
+#define KVIKIO_NVTX_MARKER_IMPL(msg, val) \
+  nvtx3::mark_in<libkvikio_domain>(       \
+    nvtx3::event_attributes{REGISTER_STRING(msg), nvtx3::payload{convert_to_64bit(val)}})
+
 #endif
 
 /**
@@ -342,4 +347,29 @@ struct libkvikio_domain {
   } while (0)
 #endif
 
+/**
+ * @brief Convenience macro for generating an NVTX marker in the `libkvikio` domain to annotate a
+ * certain time point.
+ *
+ * Takes two arguments (message, payload). Use this macro to annotate asynchronous I/O operations,
+ * where the payload refers to the I/O size.
+ *
+ * Example:
+ * ```
+ * std::future<void> some_function(){
+ *     size_t io_size{2077};
+ *     KVIKIO_NVTX_MARKER("I/O operation", io_size);
+ *     perform_async_io_operation(io_size);
+ *     ...
+ * }
+ * ```
+ */
+#ifdef KVIKIO_CUDA_FOUND
+#define KVIKIO_NVTX_MARKER(message, payload) KVIKIO_NVTX_MARKER_IMPL(message, payload)
+#else
+#define KVIKIO_NVTX_MARKER(message, payload) \
+  do {                                       \
+  } while (0)
+#endif
+
 }  // namespace kvikio