Improve NVTX annotation for functions (#671)

## Background
Previously, to add NVTX function annotation for profiling, one of the following methods is used:
```cpp
namespace kvikio::detail {
// Method 1
void f() {
    KVIKIO_NVTX_FUNC_RANGE();
}

// Method 2
void f() {
    KVIKIO_NVTX_SCOPED_RANGE("detail::f()", 0);
}
} // end namespace kvikio::detail
```

Method 1 has the limitation that:
- `__func__` is used as the function name, which is unqualified and contains no information on namespace. Function overloads cannot be well handled.
- Users cannot specify NVTX payload or color.
- There is no way to annotate a lambda. Any lambda is simply named `operator()`.

Method 2 has the limitation that:
- Users have to manually enter the function name. Naming inconsistency has been observed across different source files.
- Users have to specify an NVTX payload value.

## This PR
This PR eliminates the limitations above to improve the annotation experience. More concretely:
- `KVIKIO_NVTX_FUNC_RANGE()` now uses fully qualified name with namespace, return type and parameter types. Function overloads can be properly handled.
- `KVIKIO_NVTX_FUNC_RANGE()` now works for lambdas. For example, for a lambda inside a function `kvikio::sample::foo()`, the name of the lambda will be: `kvikio::sample::foo::<lambda()>`
- `KVIKIO_NVTX_FUNC_RANGE()` now accepts two optional parameters: payload and color.
- `KVIKIO_NVTX_SCOPED_RANGE()` now only has one mandatory parameter, which is the message; both the payload and color have become optional.

With the flexible `KVIKIO_NVTX_FUNC_RANGE()` available, **this PR adds additional traces to many functions**.

Authors:
  - Tianyu Liu (https://github.com/kingcrimsontianyu)

Approvers:
  - Mads R. B. Kristensen (https://github.com/madsbk)

URL: #671

Author: kingcrimsontianyu

cpp/include/kvikio/nvtx.hpp

@@ -50,9 +50,33 @@ using nvtx_registered_string_type = nvtx3::registered_string_in<libkvikio_domain
   }(message)
 
 // Implementation of KVIKIO_NVTX_FUNC_RANGE()
-#define KVIKIO_NVTX_FUNC_RANGE_IMPL() NVTX3_FUNC_RANGE_IN(kvikio::libkvikio_domain)
+// todo: Although supported by many compilers, __PRETTY_FUNCTION__ is non-standard. Replacement may
+// be considered once reflection is standardized.
+#define KVIKIO_NVTX_FUNC_RANGE_IMPL_0() KVIKIO_NVTX_SCOPED_RANGE_IMPL_1(__PRETTY_FUNCTION__)
+#define KVIKIO_NVTX_FUNC_RANGE_IMPL_1(payload) \
+  KVIKIO_NVTX_SCOPED_RANGE_IMPL_2(__PRETTY_FUNCTION__, payload)
+#define KVIKIO_NVTX_FUNC_RANGE_IMPL_2(payload, color) \
+  KVIKIO_NVTX_SCOPED_RANGE_IMPL_3(__PRETTY_FUNCTION__, payload, color)
+#define KVIKIO_NVTX_FUNC_RANGE_SELECTOR(_0, _1, _2, NAME, ...) NAME
+// todo: Although supported by gcc and clang, ##__VA_ARGS__ is non-standard, and should be replaced
+// by __VA_OPT__ (since C++20) in the future.
+#define KVIKIO_NVTX_FUNC_RANGE_IMPL(...)                         \
+  KVIKIO_NVTX_FUNC_RANGE_SELECTOR(_0,                            \
+                                  ##__VA_ARGS__,                 \
+                                  KVIKIO_NVTX_FUNC_RANGE_IMPL_2, \
+                                  KVIKIO_NVTX_FUNC_RANGE_IMPL_1, \
+                                  KVIKIO_NVTX_FUNC_RANGE_IMPL_0) \
+  (__VA_ARGS__)
 
 // Implementation of KVIKIO_NVTX_SCOPED_RANGE(...)
+#define KVIKIO_NVTX_SCOPED_RANGE_IMPL_1(message)                             \
+  kvikio::nvtx_scoped_range_type KVIKIO_CONCAT(_kvikio_nvtx_range, __LINE__) \
+  {                                                                          \
+    nvtx3::event_attributes                                                  \
+    {                                                                        \
+      KVIKIO_REGISTER_STRING(message), kvikio::NvtxManager::default_color()  \
+    }                                                                        \
+  }
 #define KVIKIO_NVTX_SCOPED_RANGE_IMPL_3(message, payload_v, color)                                \
   kvikio::nvtx_scoped_range_type KVIKIO_CONCAT(_kvikio_nvtx_range, __LINE__)                      \
   {                                                                                               \
@@ -64,9 +88,11 @@ using nvtx_registered_string_type = nvtx3::registered_string_in<libkvikio_domain
 #define KVIKIO_NVTX_SCOPED_RANGE_IMPL_2(message, payload) \
   KVIKIO_NVTX_SCOPED_RANGE_IMPL_3(message, payload, kvikio::NvtxManager::default_color())
 #define KVIKIO_NVTX_SCOPED_RANGE_SELECTOR(_1, _2, _3, NAME, ...) NAME
-#define KVIKIO_NVTX_SCOPED_RANGE_IMPL(...)                                         \
-  KVIKIO_NVTX_SCOPED_RANGE_SELECTOR(                                               \
-    __VA_ARGS__, KVIKIO_NVTX_SCOPED_RANGE_IMPL_3, KVIKIO_NVTX_SCOPED_RANGE_IMPL_2) \
+#define KVIKIO_NVTX_SCOPED_RANGE_IMPL(...)                           \
+  KVIKIO_NVTX_SCOPED_RANGE_SELECTOR(__VA_ARGS__,                     \
+                                    KVIKIO_NVTX_SCOPED_RANGE_IMPL_3, \
+                                    KVIKIO_NVTX_SCOPED_RANGE_IMPL_2, \
+                                    KVIKIO_NVTX_SCOPED_RANGE_IMPL_1) \
   (__VA_ARGS__)
 
 // Implementation of KVIKIO_NVTX_MARKER(message, payload)
@@ -124,22 +150,39 @@ class NvtxManager {
 };
 
 /**
- * @brief Convenience macro for generating an NVTX range in the `libkvikio` domain
- * from the lifetime of a function.
+ * @brief Convenience macro for generating an NVTX range in the `libkvikio` domain from the lifetime
+ * of a function. Can be used inside a regular function or a lambda expression.
  *
- * Takes no argument. The name of the immediately enclosing function returned by `__func__` is used
- * as the message.
+ * The function name contains detailed information such as namespace, return type, parameter type,
+ * etc.
+ *
+ * @param payload (Optional) NVTX payload.
+ * @param color (Optional) NVTX color. If unspecified, a default NVTX color is used.
  *
  * Example:
  * ```
  * void some_function(){
- *    KVIKIO_NVTX_FUNC_RANGE();  // The name `some_function` is used as the message
+ *    // No argument
+ *    KVIKIO_NVTX_FUNC_RANGE();
+ *    ...
+ * }
+ *
+ * void some_function(){
+ *    // Specify payload
+ *    KVIKIO_NVTX_FUNC_RANGE(4096);
+ *    ...
+ * }
+ *
+ * void some_function(){
+ *    // Specify payload and color
+ *    auto const nvtx3::rgb color{0, 255, 0};
+ *    KVIKIO_NVTX_FUNC_RANGE(4096, color);
  *    ...
  * }
  * ```
  */
 #ifdef KVIKIO_CUDA_FOUND
-#define KVIKIO_NVTX_FUNC_RANGE() KVIKIO_NVTX_FUNC_RANGE_IMPL()
+#define KVIKIO_NVTX_FUNC_RANGE(...) KVIKIO_NVTX_FUNC_RANGE_IMPL(__VA_ARGS__)
 #else
 #define KVIKIO_NVTX_FUNC_RANGE(...) \
   do {                              \
@@ -152,7 +195,7 @@ class NvtxManager {
  *
  * @param message String literal for NVTX annotation. To improve profile-time performance, the
  * string literal is registered in NVTX.
- * @param payload NVTX payload.
+ * @param payload (Optional) NVTX payload.
  * @param color (Optional) NVTX color. If unspecified, a default NVTX color is used.
  *
  * Example:

cpp/include/kvikio/posix_io.hpp

@@ -183,7 +183,7 @@ std::size_t posix_device_io(int fd,
 template <PartialIO PartialIOStatus>
 std::size_t posix_host_read(int fd, void* buf, std::size_t size, std::size_t file_offset)
 {
-  KVIKIO_NVTX_SCOPED_RANGE("posix_host_read()", size);
+  KVIKIO_NVTX_FUNC_RANGE(size);
   return detail::posix_host_io<IOOperationType::READ, PartialIOStatus>(
     fd, buf, size, convert_size2off(file_offset));
 }
@@ -205,7 +205,7 @@ std::size_t posix_host_read(int fd, void* buf, std::size_t size, std::size_t fil
 template <PartialIO PartialIOStatus>
 std::size_t posix_host_write(int fd, void const* buf, std::size_t size, std::size_t file_offset)
 {
-  KVIKIO_NVTX_SCOPED_RANGE("posix_host_write()", size);
+  KVIKIO_NVTX_FUNC_RANGE(size);
   return detail::posix_host_io<IOOperationType::WRITE, PartialIOStatus>(
     fd, buf, size, convert_size2off(file_offset));
 }

cpp/include/kvikio/threadpool_wrapper.hpp

@@ -33,19 +33,26 @@ class thread_pool_wrapper : public pool_type {
    *
    * @param nthreads The number of threads to use.
    */
-  thread_pool_wrapper(unsigned int nthreads) : pool_type{nthreads, worker_thread_init_func} {}
+  thread_pool_wrapper(unsigned int nthreads) : pool_type{nthreads, worker_thread_init_func}
+  {
+    KVIKIO_NVTX_FUNC_RANGE();
+  }
 
   /**
    * @brief Reset the number of threads in the thread pool, and invoke a pre-defined initialization
    * function in each worker thread.
    *
    * @param nthreads The number of threads to use.
    */
-  void reset(unsigned int nthreads) { pool_type::reset(nthreads, worker_thread_init_func); }
+  void reset(unsigned int nthreads)
+  {
+    KVIKIO_NVTX_FUNC_RANGE();
+    pool_type::reset(nthreads, worker_thread_init_func);
+  }
 
  private:
   inline static std::function<void()> worker_thread_init_func{[] {
-    KVIKIO_NVTX_SCOPED_RANGE("worker thread init", 0, NvtxManager::default_color());
+    KVIKIO_NVTX_FUNC_RANGE();
     // Rename the worker thread in the thread pool to improve clarity from nsys-ui.
     // Note: This NVTX feature is currently not supported by nsys-ui.
     NvtxManager::rename_current_thread("thread pool");

cpp/src/bounce_buffer.cpp

@@ -20,28 +20,40 @@
 #include <kvikio/bounce_buffer.hpp>
 #include <kvikio/defaults.hpp>
 #include <kvikio/error.hpp>
+#include <kvikio/nvtx.hpp>
 #include <kvikio/shim/cuda.hpp>
 
 namespace kvikio {
 
 AllocRetain::Alloc::Alloc(AllocRetain* manager, void* alloc, std::size_t size)
   : _manager(manager), _alloc{alloc}, _size{size}
 {
+  KVIKIO_NVTX_FUNC_RANGE();
 }
 
-AllocRetain::Alloc::~Alloc() noexcept { _manager->put(_alloc, _size); }
+AllocRetain::Alloc::~Alloc() noexcept
+{
+  KVIKIO_NVTX_FUNC_RANGE();
+  _manager->put(_alloc, _size);
+}
 
-void* AllocRetain::Alloc::get() noexcept { return _alloc; }
+void* AllocRetain::Alloc::get() noexcept
+{
+  KVIKIO_NVTX_FUNC_RANGE();
+  return _alloc;
+}
 
 void* AllocRetain::Alloc::get(std::ptrdiff_t offset) noexcept
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   return static_cast<char*>(_alloc) + offset;
 }
 
 std::size_t AllocRetain::Alloc::size() noexcept { return _size; }
 
 std::size_t AllocRetain::_clear()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   std::size_t ret = _free_allocs.size() * _size;
   while (!_free_allocs.empty()) {
     CUDA_DRIVER_TRY(cudaAPI::instance().MemFreeHost(_free_allocs.top()));
@@ -52,6 +64,7 @@ std::size_t AllocRetain::_clear()
 
 void AllocRetain::_ensure_alloc_size()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   auto const bounce_buffer_size = defaults::bounce_buffer_size();
   if (_size != bounce_buffer_size) {
     _clear();
@@ -61,6 +74,7 @@ void AllocRetain::_ensure_alloc_size()
 
 AllocRetain::Alloc AllocRetain::get()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   std::lock_guard const lock(_mutex);
   _ensure_alloc_size();
 
@@ -81,6 +95,7 @@ AllocRetain::Alloc AllocRetain::get()
 
 void AllocRetain::put(void* alloc, std::size_t size)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   std::lock_guard const lock(_mutex);
   _ensure_alloc_size();
 
@@ -95,12 +110,14 @@ void AllocRetain::put(void* alloc, std::size_t size)
 
 std::size_t AllocRetain::clear()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   std::lock_guard const lock(_mutex);
   return _clear();
 }
 
 AllocRetain& AllocRetain::instance()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   static AllocRetain _instance;
   return _instance;
 }

cpp/src/buffer.cpp

@@ -21,6 +21,7 @@
 #include <kvikio/buffer.hpp>
 #include <kvikio/defaults.hpp>
 #include <kvikio/error.hpp>
+#include <kvikio/nvtx.hpp>
 #include <kvikio/shim/cufile.hpp>
 #include <kvikio/shim/cufile_h_wrapper.hpp>
 #include <kvikio/utils.hpp>
@@ -32,6 +33,7 @@ void buffer_register(void const* devPtr_base,
                      int flags,
                      std::vector<int> const& errors_to_ignore)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   if (defaults::is_compat_mode_preferred()) { return; }
   CUfileError_t status = cuFileAPI::instance().BufRegister(devPtr_base, size, flags);
   if (status.err != CU_FILE_SUCCESS) {
@@ -45,18 +47,21 @@ void buffer_register(void const* devPtr_base,
 
 void buffer_deregister(void const* devPtr_base)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   if (defaults::is_compat_mode_preferred()) { return; }
   CUFILE_TRY(cuFileAPI::instance().BufDeregister(devPtr_base));
 }
 
 void memory_register(void const* devPtr, int flags, std::vector<int> const& errors_to_ignore)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   auto [base, nbytes, offset] = get_alloc_info(devPtr);
   buffer_register(base, nbytes, flags, errors_to_ignore);
 }
 
 void memory_deregister(void const* devPtr)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   auto [base, nbytes, offset] = get_alloc_info(devPtr);
   buffer_deregister(base);
 }

cpp/src/compat_mode.cpp

@@ -22,13 +22,15 @@
 #include <kvikio/cufile/config.hpp>
 #include <kvikio/error.hpp>
 #include <kvikio/file_handle.hpp>
+#include <kvikio/nvtx.hpp>
 #include <kvikio/shim/cufile.hpp>
 
 namespace kvikio {
 
 namespace detail {
 CompatMode parse_compat_mode_str(std::string_view compat_mode_str)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   // Convert to lowercase
   std::string tmp{compat_mode_str};
   std::transform(
@@ -50,6 +52,7 @@ CompatMode parse_compat_mode_str(std::string_view compat_mode_str)
 
 CompatMode CompatModeManager::infer_compat_mode_if_auto(CompatMode compat_mode) noexcept
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   if (compat_mode == CompatMode::AUTO) {
     return is_cufile_available() ? CompatMode::OFF : CompatMode::ON;
   }
@@ -84,6 +87,7 @@ CompatModeManager::CompatModeManager(std::string const& file_path,
                                      CompatMode compat_mode_requested_v,
                                      FileHandle* file_handle)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   KVIKIO_EXPECT(file_handle != nullptr,
                 "The compatibility mode manager does not have a proper owning file handle.",
                 std::invalid_argument);
@@ -127,6 +131,7 @@ CompatModeManager::CompatModeManager(std::string const& file_path,
 
 void CompatModeManager::validate_compat_mode_for_async() const
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   if (!_is_compat_mode_preferred && _is_compat_mode_preferred_for_async &&
       _compat_mode_requested == CompatMode::OFF) {
     std::string err_msg;

cpp/src/defaults.cpp

@@ -33,6 +33,7 @@ namespace kvikio {
 template <>
 bool getenv_or(std::string_view env_var_name, bool default_val)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   auto const* env_val = std::getenv(env_var_name.data());
   if (env_val == nullptr) { return default_val; }
   try {
@@ -66,6 +67,7 @@ bool getenv_or(std::string_view env_var_name, bool default_val)
 template <>
 CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   auto* env_val = std::getenv(env_var_name.data());
   if (env_val == nullptr) { return default_val; }
   return detail::parse_compat_mode_str(env_val);
@@ -74,6 +76,7 @@ CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val)
 template <>
 std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> default_val)
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   auto* const env_val = std::getenv(env_var_name.data());
   if (env_val == nullptr) { return std::move(default_val); }
   std::string const int_str(env_val);
@@ -84,13 +87,15 @@ std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> defau
 
 unsigned int defaults::get_num_threads_from_env()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   int const ret = getenv_or("KVIKIO_NTHREADS", 1);
   KVIKIO_EXPECT(ret > 0, "KVIKIO_NTHREADS has to be a positive integer", std::invalid_argument);
   return ret;
 }
 
 defaults::defaults()
 {
+  KVIKIO_NVTX_FUNC_RANGE();
   // Determine the default value of `compat_mode`
   {
     _compat_mode = getenv_or("KVIKIO_COMPAT_MODE", CompatMode::AUTO);