nvtx.hpp

 
#pragma once
 
#include <stdexcept>
#include <type_traits>
 
#include <nvtx3/nvtx3.hpp>
 
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf::detail {
 
template <typename T>
    requires std::is_integral_v<T>
[[nodiscard]] std::int64_t convert_to_64bit(T value) {
    if constexpr (std::numeric_limits<T>::max()
                  > std::numeric_limits<std::int64_t>::max())
    {
        if (value > std::numeric_limits<std::int64_t>::max()) {
            throw std::overflow_error(
                "convert_to_64bit(x): x too large to fit std::int64_t"
            );
        }
    }
    return rapidsmpf::safe_cast<std::int64_t>(value);
}
 
template <typename T>
    requires std::is_floating_point_v<T>
[[nodiscard]] double convert_to_64bit(T value) {
    return double(value);
}
}  // namespace rapidsmpf::detail
 
struct rapidsmpf_domain {
    static constexpr char const* name{"rapidsmpf"};  
};
 
// Macro to create a static, registered string that will not have a name conflict with any
// registered string defined in the same scope.
#define RAPIDSMPF_REGISTER_STRING(msg)                                         \
    [](char const* a_msg) -> auto& {                                           \
        static nvtx3::registered_string_in<rapidsmpf_domain> a_reg_str{a_msg}; \
        return a_reg_str;                                                      \
    }(msg)
 
// implement the func range macro with a value
#define RAPIDSMPF_NVTX_FUNC_RANGE_IMPL_WITH_VAL(val)              \
    static_assert(                                                \
        std::is_arithmetic_v<decltype(val)>,                      \
        "Value must be integral or floating point type"           \
    );                                                            \
    nvtx3::scoped_range_in<rapidsmpf_domain> RAPIDSMPF_CONCAT(    \
        _rapidsmpf_nvtx_range, __LINE__                           \
    ) {                                                           \
        nvtx3::event_attributes {                                 \
            RAPIDSMPF_REGISTER_STRING(__func__), nvtx3::payload { \
                rapidsmpf::detail::convert_to_64bit(val)          \
            }                                                     \
        }                                                         \
    }
 
// implement the func range macro without a value
#define RAPIDSMPF_NVTX_FUNC_RANGE_IMPL_WITHOUT_VAL() NVTX3_FUNC_RANGE_IN(rapidsmpf_domain)
 
// Macro selector for 0 vs 1 arguments
#define RAPIDSMPF_GET_MACRO_FUNC(_0, _1, NAME, ...) NAME
 
// unwrap the arguments and call the appropriate macro
#define RAPIDSMPF_NVTX_FUNC_RANGE_IMPL(...)                                                                                                          \
    RAPIDSMPF_GET_MACRO_FUNC(dummy __VA_OPT__(, ) __VA_ARGS__, RAPIDSMPF_NVTX_FUNC_RANGE_IMPL_WITH_VAL, RAPIDSMPF_NVTX_FUNC_RANGE_IMPL_WITHOUT_VAL)( \
        __VA_ARGS__                                                                                                                                  \
    )
 
#define RAPIDSMPF_NVTX_FUNC_RANGE(...) RAPIDSMPF_NVTX_FUNC_RANGE_IMPL(__VA_ARGS__)
 
// implement the scoped range macro with a value
#define RAPIDSMPF_NVTX_SCOPED_RANGE_IMPL_WITH_VAL(msg, val)    \
    nvtx3::scoped_range_in<rapidsmpf_domain> RAPIDSMPF_CONCAT( \
        _rapidsmpf_nvtx_range, __LINE__                        \
    ) {                                                        \
        nvtx3::event_attributes {                              \
            RAPIDSMPF_REGISTER_STRING(msg), nvtx3::payload {   \
                rapidsmpf::detail::convert_to_64bit(val)       \
            }                                                  \
        }                                                      \
    }
 
// implement the scoped range macro without a value
#define RAPIDSMPF_NVTX_SCOPED_RANGE_IMPL_WITHOUT_VAL(msg)      \
    nvtx3::scoped_range_in<rapidsmpf_domain> RAPIDSMPF_CONCAT( \
        _rapidsmpf_nvtx_range, __LINE__                        \
    ) {                                                        \
        nvtx3::event_attributes {                              \
            RAPIDSMPF_REGISTER_STRING(msg)                     \
        }                                                      \
    }
 
// Macro to detect number of arguments (1 or 2)
#define RAPIDSMPF_GET_MACRO(_1, _2, NAME, ...) NAME
 
// unwrap the arguments and call the appropriate macro
#define RAPIDSMPF_NVTX_SCOPED_RANGE_IMPL(...)        \
    RAPIDSMPF_GET_MACRO(                             \
        __VA_ARGS__,                                 \
        RAPIDSMPF_NVTX_SCOPED_RANGE_IMPL_WITH_VAL,   \
        RAPIDSMPF_NVTX_SCOPED_RANGE_IMPL_WITHOUT_VAL \
    )                                                \
    (__VA_ARGS__)
 
#define RAPIDSMPF_NVTX_SCOPED_RANGE(...) RAPIDSMPF_NVTX_SCOPED_RANGE_IMPL(__VA_ARGS__)
 
#if RAPIDSMPF_VERBOSE_INFO
#define RAPIDSMPF_NVTX_SCOPED_RANGE_VERBOSE(...) RAPIDSMPF_NVTX_SCOPED_RANGE(__VA_ARGS__)
#else
#define RAPIDSMPF_NVTX_SCOPED_RANGE_VERBOSE(...)
#endif
 
#define RAPIDSMPF_NVTX_MARKER_IMPL(msg, val)                     \
    nvtx3::mark_in<rapidsmpf_domain>(nvtx3::event_attributes{    \
        RAPIDSMPF_REGISTER_STRING(msg),                          \
        nvtx3::payload{rapidsmpf::detail::convert_to_64bit(val)} \
    })
 
#define RAPIDSMPF_NVTX_MARKER(message, payload) \
    RAPIDSMPF_NVTX_MARKER_IMPL(message, payload)
 
#if RAPIDSMPF_VERBOSE_INFO
#define RAPIDSMPF_NVTX_MARKER_VERBOSE(message, payload) \
    RAPIDSMPF_NVTX_MARKER_IMPL(message, payload)
#else
#define RAPIDSMPF_NVTX_MARKER_VERBOSE(message, payload)
#endif