20 #ifndef MXNET_COMMON_CUDA_NVTX_H_
21 #define MXNET_COMMON_CUDA_NVTX_H_
23 #if MXNET_USE_CUDA && MXNET_USE_NVTX
25 #include <cuda_runtime.h>
26 #include <nvToolsExtCuda.h>
37 explicit NVTXDuration(
const char* name) noexcept : range_id_(0), name_(name) {}
40 range_id_ = nvtxRangeStartA(name_);
44 nvtxRangeEnd(range_id_);
48 nvtxRangeId_t range_id_;
56 static const uint32_t kRed = 0xFF0000;
57 static const uint32_t kGreen = 0x00FF00;
58 static const uint32_t kBlue = 0x0000FF;
59 static const uint32_t kYellow = 0xB58900;
60 static const uint32_t kOrange = 0xCB4B16;
61 static const uint32_t kRed1 = 0xDC322F;
62 static const uint32_t kMagenta = 0xD33682;
63 static const uint32_t kViolet = 0x6C71C4;
64 static const uint32_t kBlue1 = 0x268BD2;
65 static const uint32_t kCyan = 0x2AA198;
66 static const uint32_t kGreen1 = 0x859900;
68 static void gpuRangeStart(
const uint32_t rgb,
const std::string& range_name) {
69 nvtxEventAttributes_t att;
70 att.version = NVTX_VERSION;
71 att.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
72 att.colorType = NVTX_COLOR_ARGB;
73 att.color = rgb | 0xff000000;
74 att.messageType = NVTX_MESSAGE_TYPE_ASCII;
75 att.message.ascii = range_name.c_str();
76 nvtxRangePushEx(&att);
80 static uint32_t nameToColor(
const std::string& range_name,
int prefix_len) {
81 static std::vector<uint32_t> colors{
82 kRed, kGreen, kBlue, kYellow, kOrange, kRed1, kMagenta, kViolet, kBlue1, kCyan, kGreen1};
83 std::string s(range_name, 0, prefix_len);
84 std::hash<std::string> hash_fn;
85 return colors[hash_fn(s) % colors.size()];
89 static uint32_t nameToColor(
const std::string& range_name) {
90 return nameToColor(range_name, range_name.size());
93 static void gpuRangeStop() {
102 #endif // MXNET_UDE_CUDA && MXNET_USE_NVTX
103 #endif // MXNET_COMMON_CUDA_NVTX_H_