Skip to content

Refactor the NVTX and add the annotation for task begin and end #453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
eddy16112 wants to merge 7 commits into
base: main
Choose a base branch
from
Draft

Refactor the NVTX and add the annotation for task begin and end #453

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7cc34e6
new nvtx
eddy16112 Jun 4, 2026
56f2845
check
eddy16112 Jun 4, 2026
04eadfd
color
eddy16112 Jun 4, 2026
7efb702
fix
eddy16112 Jun 4, 2026
f654aa7
check
eddy16112 Jun 4, 2026
55c1283
fix
eddy16112 Jun 4, 2026
501083c
checkpoint
eddy16112 Jun 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -423,17 +423,11 @@ if(REALM_ENABLE_CUDA OR (REALM_ENABLE_HIP AND CMAKE_HIP_PLATFORM STREQUAL "nvidi
PATH_SUFFIXES include ../include extras/CUPTI/include ../extras/CUPTI/include
)

# Deal with the library name change in older versions of cmake.
#
# nvtx3 was added in cmake 3.25, and while later versions (at least 3.30+) claim that
# the CUDA::nvtx3 and CUDA::nvToolsExt targets are mutually exclusive, they don't
# appear to be so in earlier versions. So if we have nvtx3, prefer that, otherwise
# fall back to nvToolsExt.
# We use the header-only NVTX v3 C++ API, which is exposed via the
# CUDA::nvtx3 target (added in cmake 3.25). Require it; otherwise disable NVTX.
set(REALM_USE_NVTX ${REALM_ENABLE_NVTX})
if(TARGET CUDA::nvtx3)
list(APPEND REALM_LIBRARIES CUDA::nvtx3)
elseif(TARGET CUDA::nvToolsExt)
list(APPEND REALM_LIBRARIES CUDA::nvToolsExt)
else()
# Silently disable NVTX if we cannot find it
set(REALM_USE_NVTX OFF)
Expand Down
5 changes: 5 additions & 0 deletions src/realm/cuda/cuda_module.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1531,6 +1531,11 @@ namespace Realm {
tte = &it->second;
}

#ifdef REALM_USE_NVTX
std::string nvtx_msg = stringbuilder() << "task " << func_id;
nvtxUniqueRange nvtx_range(nvtx_category.get(), nvtx_msg.c_str());
#endif

if(tte->stream_aware_fnptr) {
// shouldn't be here without a valid stream
assert(ThreadLocal::current_gpu_stream != nullptr);
Expand Down
265 changes: 85 additions & 180 deletions src/realm/nvtx.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,258 +17,163 @@

#include "realm/nvtx.h"

#include <assert.h>
#include <functional>
#include <iostream>
#include <memory>
#include "realm/atomics.h"
#ifdef REALM_ON_WINDOWS
#include <processthreadsapi.h>
#else
#include <pthread.h>
#include <sys/syscall.h>
#include <unistd.h>
#endif

namespace Realm {

struct realm_nvtx_domain {
static constexpr char const *name{"Realm"};
};
// Definitions of the predefined category handles declared in the header.
NvtxCategory *nvtx_amsg = nullptr;
NvtxCategory *nvtx_bgwork = nullptr;
#ifdef REALM_USE_CUDA
NvtxCategory *nvtx_cuda = nullptr;
#endif
#ifdef REALM_USE_HIP
NvtxCategory *nvtx_hip = nullptr;
#endif
#ifdef REALM_USE_GASNET1
NvtxCategory *nvtx_gasnet1 = nullptr;
#endif
#ifdef REALM_USE_GASNETEX
NvtxCategory *nvtx_gasnetex = nullptr;
#endif
#ifdef REALM_USE_MPI
NvtxCategory *nvtx_mpi = nullptr;
#endif
#ifdef REALM_USE_OPENMP
NvtxCategory *nvtx_openmp = nullptr;
#endif
#ifdef REALM_USE_PYTHON
NvtxCategory *nvtx_python = nullptr;
#endif

struct nvtx_category_id_color {
// Static description of a predefined category: its id, default color, and the
// global handle to point at the created object. `slot` is a reference to one
// of the `nvtx_*` pointers defined above.
struct nvtx_category_def {
uint32_t id;
uint32_t color;
nvtx3::color color;
std::reference_wrapper<NvtxCategory *> slot;
};

static std::map<std::string, nvtx_category_id_color> nvtx_categories_predefined = {
{"amsg", {1, nvtx_color::red}},
{"bgwork", {2, nvtx_color::blue}},
static const std::map<std::string, nvtx_category_def> nvtx_categories_predefined = {
{"amsg", {1, nvtx_color::red, std::ref(nvtx_amsg)}},
{"bgwork", {2, nvtx_color::blue, std::ref(nvtx_bgwork)}},
#ifdef REALM_USE_CUDA
{"cuda", {100, nvtx_color::green}},
{"cuda", {100, nvtx_color::green, std::ref(nvtx_cuda)}},
#endif
#ifdef REALM_USE_HIP
{"hip", {101, nvtx_color::purple}},
{"hip", {101, nvtx_color::purple, std::ref(nvtx_hip)}},
#endif
#ifdef REALM_USE_GASNET1
{"gasnet1", {102, nvtx_color::lawn_green}},
{"gasnet1", {102, nvtx_color::lawn_green, std::ref(nvtx_gasnet1)}},
#endif
#ifdef REALM_USE_GASNETEX
{"gasnetex", {103, nvtx_color::cyan}},
{"gasnetex", {103, nvtx_color::cyan, std::ref(nvtx_gasnetex)}},
#endif
#ifdef REALM_USE_MPI
{"mpi", {104, nvtx_color::maroon}},
{"mpi", {104, nvtx_color::maroon, std::ref(nvtx_mpi)}},
#endif
#ifdef REALM_USE_OPENMP
{"openmp", {105, nvtx_color::navy}},
{"openmp", {105, nvtx_color::navy, std::ref(nvtx_openmp)}},
#endif
#ifdef REALM_USE_PYTHON
{"python", {106, nvtx_color::magenta}},
{"python", {106, nvtx_color::magenta, std::ref(nvtx_python)}},
#endif
};

thread_local std::map<std::string, NvtxCategory *> *nvtx_categories;

static nvtxDomainHandle_t nvtxRealmDomain = nullptr;
// Owns the category objects created at init; cleared at finalize. The global
// `nvtx_*` handles point into these.
static std::vector<std::unique_ptr<NvtxCategory>> nvtx_owned_categories;

static std::vector<std::string> enabled_nvtx_modules;

static inline NvtxCategory *find_category_by_name(const std::string &name)
{
std::map<std::string, NvtxCategory *>::iterator it = nvtx_categories->find(name);
if(it != nvtx_categories->end()) {
return it->second;
} else {
return nullptr;
}
}

////////////////////////////////////////////////////////////////////////
//
// class nvtxCategory

NvtxCategory::NvtxCategory(const std::string &category_name, uint32_t category_id,
uint32_t color)
: name(category_name)
{
// name the category
nvtxDomainNameCategoryA(nvtxRealmDomain, category_id, category_name.c_str());

// create nvtx event attribute and set values
memset(&nvtx_event, 0, NVTX_EVENT_ATTRIB_STRUCT_SIZE);
nvtx_event.version = NVTX_VERSION;
nvtx_event.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
nvtx_event.category = category_id;
nvtx_event.messageType = NVTX_MESSAGE_TYPE_ASCII;
nvtx_event.message.ascii = "";
nvtx_event.payloadType = NVTX_PAYLOAD_TYPE_INT32;
nvtx_event.payload.iValue = 0;
nvtx_event.colorType = NVTX_COLOR_ARGB;
nvtx_event.color = color;
}
static atomic<uint32_t> nvtx_proc_starting_category_id{1000};

////////////////////////////////////////////////////////////////////////
//
// class nvtxScopedRange

nvtxScopedRange::nvtxScopedRange(NvtxCategory *category, char const *message,
int32_t payload)
{
category->nvtx_event.message.ascii = message;
category->nvtx_event.payload.iValue = payload;
nvtxDomainRangePushEx(nvtxRealmDomain, &(category->nvtx_event));
}
// Palette of visually distinct colors handed out by nvtx_get_next_color().
static constexpr nvtx3::color nvtx_auto_color_palette[] = {
nvtx_color::red, nvtx_color::green, nvtx_color::blue,
nvtx_color::purple, nvtx_color::cyan, nvtx_color::maroon,
nvtx_color::navy, nvtx_color::magenta, nvtx_color::yellow,
nvtx_color::teal, nvtx_color::olive, nvtx_color::lawn_green,
};
static atomic<uint32_t> nvtx_next_color_index{0};

nvtxScopedRange::nvtxScopedRange(const std::string &name, char const *message,
int32_t payload)
// Create the category object, take ownership, and point its global handle at it.
static void create_category(const std::string &name, const nvtx_category_def &def)
{
NvtxCategory *category = find_category_by_name(name);
if(category) {
category->nvtx_event.message.ascii = message;
category->nvtx_event.payload.iValue = payload;
nvtxDomainRangePushEx(nvtxRealmDomain, &(category->nvtx_event));
}
nvtx_owned_categories.push_back(
std::make_unique<NvtxCategory>(name, def.id, def.color));
def.slot.get() = nvtx_owned_categories.back().get();
}

nvtxScopedRange::~nvtxScopedRange() { nvtxDomainRangePop(nvtxRealmDomain); }

void init_nvtx_thread(const char *thread_name)
{
#ifdef REALM_ON_WINDOWS
nvtxNameOsThread(GetCurrentThreadId(), thread_name)
nvtxNameOsThread(GetCurrentThreadId(), thread_name);
#else
nvtxNameOsThread(pthread_self(), thread_name);
// NVTX wants the OS-native (kernel) thread id, which on Linux is gettid(),
nvtxNameOsThread(static_cast<uint32_t>(syscall(SYS_gettid)), thread_name);
#endif
}

nvtx_categories = new std::map<std::string, NvtxCategory *>();
nvtx_categories->clear();
void init_nvtx(std::vector<std::string> &nvtx_modules)
{
enabled_nvtx_modules = nvtx_modules;

if(enabled_nvtx_modules.size() == 1 and enabled_nvtx_modules[0] == "all") {
if(enabled_nvtx_modules.size() == 1 && enabled_nvtx_modules[0] == "all") {
// handle -ll:nvtx_modules all
std::map<std::string, nvtx_category_id_color>::const_iterator it;
for(it = nvtx_categories_predefined.cbegin();
it != nvtx_categories_predefined.cend(); it++) {
nvtx_categories->insert(std::pair<std::string, NvtxCategory *>(
it->first, new NvtxCategory(it->first, it->second.id, it->second.color)));
for(const auto &entry : nvtx_categories_predefined) {
create_category(entry.first, entry.second);
}
} else {
for(const std::string &name : enabled_nvtx_modules) {
if(name == "all") {
std::cerr << "If all specified, then no other modules are needed." << std::endl;
abort();
}
std::map<std::string, nvtx_category_id_color>::const_iterator it =
std::map<std::string, nvtx_category_def>::const_iterator it =
nvtx_categories_predefined.find(name);
if(it == nvtx_categories_predefined.end()) {
std::cerr << "Unable to find specified nvtx module: " << name << std::endl;
abort();
}
nvtx_categories->insert(std::pair<std::string, NvtxCategory *>(
name, new NvtxCategory(name, it->second.id, it->second.color)));
create_category(it->first, it->second);
}
}
}

void finalize_nvtx_thread(void)
{
std::map<std::string, NvtxCategory *>::iterator it;
for(it = nvtx_categories->begin(); it != nvtx_categories->end(); it++) {
assert(it->second != nullptr);
delete it->second;
}
delete nvtx_categories;
}

void init_nvtx(std::vector<std::string> &nvtx_modules)
{
enabled_nvtx_modules = nvtx_modules;
nvtxInitialize(nullptr);
nvtxRealmDomain = nvtxDomainCreateA(realm_nvtx_domain::name);
init_nvtx_thread("MainThread");
}

void finalize_nvtx(void)
{
nvtxDomainDestroy(nvtxRealmDomain);
finalize_nvtx_thread();
}

void nvtx_range_push(NvtxCategory *category, const char *message, uint32_t color,
int32_t payload)
{
uint32_t origin_color;
if(color != nvtx_color::white) {
origin_color = category->nvtx_event.color;
category->nvtx_event.color = color;
}
category->nvtx_event.message.ascii = message;
category->nvtx_event.payload.iValue = payload;
nvtxDomainRangePushEx(nvtxRealmDomain, &(category->nvtx_event));
if(color != nvtx_color::white) {
category->nvtx_event.color = origin_color;
// The nvtx3 domain is intentionally never destroyed (see nvtx3 docs). Drop
// the owned categories and reset the global handles so none dangle.
nvtx_owned_categories.clear();
for(const auto &entry : nvtx_categories_predefined) {
entry.second.slot.get() = nullptr;
}
}

void nvtx_range_push(const std::string &name, const char *message, uint32_t color,
int32_t payload)
uint32_t nvtx_get_next_category_id(void)
{
NvtxCategory *category = find_category_by_name(name);
if(category) {
nvtx_range_push(category, message, color, payload);
}
return nvtx_proc_starting_category_id.fetch_add(1);
}

void nvtx_range_pop(void) { nvtxDomainRangePop(nvtxRealmDomain); }

nvtxRangeId_t nvtx_range_start(NvtxCategory *category, const char *message,
uint32_t color, int32_t payload)
nvtx3::color nvtx_get_next_color(void)
{
uint32_t origin_color;
if(color != nvtx_color::white) {
origin_color = category->nvtx_event.color;
category->nvtx_event.color = color;
}
category->nvtx_event.message.ascii = message;
category->nvtx_event.payload.iValue = payload;
nvtxRangeId_t id = nvtxDomainRangeStartEx(nvtxRealmDomain, &(category->nvtx_event));
if(color != nvtx_color::white) {
category->nvtx_event.color = origin_color;
}
return id;
}

nvtxRangeId_t nvtx_range_start(const std::string &name, const char *message,
uint32_t color, int32_t payload)
{
NvtxCategory *category = find_category_by_name(name);
if(category) {
return nvtx_range_start(category, message, color, payload);
} else {
return 0;
}
}

void nvtx_range_end(nvtxRangeId_t id) { nvtxDomainRangeEnd(nvtxRealmDomain, id); }

void nvtx_mark(NvtxCategory *category, const char *message, uint32_t color,
int32_t payload)
{
uint32_t origin_color;
if(color != nvtx_color::white) {
origin_color = category->nvtx_event.color;
category->nvtx_event.color = color;
}
category->nvtx_event.message.ascii = message;
category->nvtx_event.color = color;
category->nvtx_event.payload.iValue = payload;
nvtxDomainMarkEx(nvtxRealmDomain, &(category->nvtx_event));
if(color != nvtx_color::white) {
category->nvtx_event.color = origin_color;
}
}

void nvtx_mark(const std::string &name, const char *message, uint32_t color,
int32_t payload)
{
NvtxCategory *category = find_category_by_name(name);
if(category) {
nvtx_mark(category, message, color, payload);
}
constexpr size_t n =
sizeof(nvtx_auto_color_palette) / sizeof(nvtx_auto_color_palette[0]);
return nvtx_auto_color_palette[nvtx_next_color_index.fetch_add(1) % n];
}

}; // namespace Realm
Loading
Loading