nvtx.hpp

/*
 * Copyright (c) 2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <cstdint>
 
#ifdef KVIKIO_CUDA_FOUND
#include <nvtx3/nvtx3.hpp>
#endif
 
#include <kvikio/shim/cuda.hpp>
#include <kvikio/utils.hpp>
 
namespace kvikio {
 
#ifdef KVIKIO_CUDA_FOUND
struct libkvikio_domain {
  static constexpr char const* name{"libkvikio"};
};
 
using nvtx_scoped_range_type      = nvtx3::scoped_range_in<libkvikio_domain>;
using nvtx_registered_string_type = nvtx3::registered_string_in<libkvikio_domain>;
 
// Macro to concatenate two tokens x and y.
#define KVIKIO_CONCAT_HELPER(x, y) x##y
#define KVIKIO_CONCAT(x, y)        KVIKIO_CONCAT_HELPER(x, y)
 
// Macro to create a static, registered string that will not have a name conflict with any
// registered string defined in the same scope.
#define KVIKIO_REGISTER_STRING(message)                              \
  [](const char* a_message) -> auto& {                               \
    static kvikio::nvtx_registered_string_type a_reg_str{a_message}; \
    return a_reg_str;                                                \
  }(message)
 
// Implementation of KVIKIO_NVTX_FUNC_RANGE()
#define KVIKIO_NVTX_FUNC_RANGE_IMPL() NVTX3_FUNC_RANGE_IN(kvikio::libkvikio_domain)
 
// Implementation of KVIKIO_NVTX_SCOPED_RANGE(...)
#define KVIKIO_NVTX_SCOPED_RANGE_IMPL_3(message, payload_v, color)                                \
  kvikio::nvtx_scoped_range_type KVIKIO_CONCAT(_kvikio_nvtx_range, __LINE__)                      \
  {                                                                                               \
    nvtx3::event_attributes                                                                       \
    {                                                                                             \
      KVIKIO_REGISTER_STRING(message), nvtx3::payload{kvikio::convert_to_64bit(payload_v)}, color \
    }                                                                                             \
  }
#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) \
  (__VA_ARGS__)
 
// Implementation of KVIKIO_NVTX_MARKER(message, payload)
#define KVIKIO_NVTX_MARKER_IMPL(message, payload_v)                 \
  nvtx3::mark_in<kvikio::libkvikio_domain>(nvtx3::event_attributes{ \
    KVIKIO_REGISTER_STRING(message), nvtx3::payload{kvikio::convert_to_64bit(payload_v)}})
 
#endif
 
#ifdef KVIKIO_CUDA_FOUND
using nvtx_color_type = nvtx3::color;
#else
using nvtx_color_type = int;
#endif
 
class NvtxManager {
 public:
  static NvtxManager& instance() noexcept;
 
  static const nvtx_color_type& default_color() noexcept;
 
  static const nvtx_color_type& get_color_by_index(std::uint64_t idx) noexcept;
 
  static void rename_current_thread(std::string_view new_name) noexcept;
 
  NvtxManager(NvtxManager const&)            = delete;
  NvtxManager& operator=(NvtxManager const&) = delete;
  NvtxManager(NvtxManager&&)                 = delete;
  NvtxManager& operator=(NvtxManager&&)      = delete;
 
 private:
  NvtxManager() = default;
};
 
#ifdef KVIKIO_CUDA_FOUND
#define KVIKIO_NVTX_FUNC_RANGE() KVIKIO_NVTX_FUNC_RANGE_IMPL()
#else
#define KVIKIO_NVTX_FUNC_RANGE(...) \
  do {                              \
  } while (0)
#endif
 
#ifdef KVIKIO_CUDA_FOUND
#define KVIKIO_NVTX_SCOPED_RANGE(...) KVIKIO_NVTX_SCOPED_RANGE_IMPL(__VA_ARGS__)
#else
#define KVIKIO_NVTX_SCOPED_RANGE(message, payload, ...) \
  do {                                                  \
  } while (0)
#endif
 
#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