sglang_v0.5.2/nvshmem_src/src/host/init/init_nvtx.cpp

114 lines
4.4 KiB
C++

/*
* Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
*
* See COPYRIGHT for license information
*/
#include "internal/host/nvshmem_nvtx.hpp" // for nvtxOpt_t, ALLOC_OPT
// IWYU pragma: no_include <nvtx3/nvtxDetail/nvtxImplCore.h>
#ifndef NVTX_DISABLE
#include <stdio.h> // for snprintf, NULL, printf
#include <stdlib.h> // for free
#include <string.h> // for strtok, strdup
#include <syscall.h> // for SYS_gettid
#include <unistd.h> // for syscall
#include "internal/host/debug.h" // for strcmp_case_insensi...
#include "internal/host/util.h" // for nvshmemi_options
#include "bootstrap_host_transport/env_defs_internal.h" // for nvshmemi_options_s
int nvshmem_nvtx_options = 0; // NVTX instrumentation is disabled by default
// mapping of environment variable options to internal ID
static const struct nvtxOptMap {
nvtxOpt_t opt;
const char *name;
const char *desc;
} options[] = {
{INIT_OPT, "init", "library setup"},
{ALLOC_OPT, "alloc", "memory management"},
{LAUNCH_OPT, "launch", "kernel launch routines"},
{COLL_OPT, "coll", "collective communications"},
{WAIT_OPT, "wait", "blocking point-to-point synchronization"},
{WAIT_ON_STREAM_OPT, "wait_on_stream", "point-to-point synchronization (on stream)"},
{TEST_OPT, "test", "non-blocking point-to-point synchronization"},
{MEMORDER_OPT, "memorder", "memory ordering (quiet, fence)"},
{QUIET_ON_STREAM_OPT, "quiet_on_stream", "nvshmemx_quiet_on_stream"},
{ATOMIC_FETCH_OPT, "atomic_fetch", "fetching atomic memory operations"},
{ATOMIC_SET_OPT, "atomic_set", "non-fetchong atomic memory operations"},
{RMA_BLOCKING_OPT, "rma_blocking", "blocking remote memory access operations"},
{RMA_NONBLOCKING_OPT, "rma_nonblocking", "non-blocking remote memory access operations"},
{PROXY_OPT, "proxy", "activity of the proxy thread"},
{DEFAULT_OPT, "common", "init,alloc,launch,coll,memorder,wait,atomic_fetch,rma_blocking,proxy"},
{ALL_OPT, "all", "all groups"},
{(nvtxOpt_t)0, "off", "disable all NVTX instrumentation"}};
static const int num_options = sizeof(options) / sizeof(options[0]);
static const int num_options_enable = num_options - 1;
void nvshmem_nvtx_set_thread_name(int pe, const char *suffix) {
char name[32];
if (suffix) {
snprintf(name, sizeof(name), "NVSHMEM PE %d %s", pe, suffix);
} else {
snprintf(name, sizeof(name), "NVSHMEM PE %d", pe);
}
nvtxNameOsThreadA(syscall(SYS_gettid), name);
}
static bool initialized = false;
void nvshmem_nvtx_init() {
if (initialized) {
return;
}
initialized = true;
// evaluate what should be instrumented from environment variable
if (nvshmemi_options.NVTX_provided) {
char *groups = strdup(nvshmemi_options.NVTX);
char seps[] = ",";
// iterate over all options from the environment variable
char *gname = strtok(groups, seps);
while (gname != NULL) {
// handle enable options
for (int i = 0; i < num_options_enable; ++i) {
if (strcmp_case_insensitive(gname, options[i].name) == 0) {
nvshmem_nvtx_options |= options[i].opt;
}
}
// handle disable options ('0' or last entry of options)
if (strcmp_case_insensitive(gname, "0") == 0 ||
strcmp_case_insensitive(gname, options[num_options - 1].name) == 0) {
nvshmem_nvtx_options = 0;
}
gname = strtok(NULL, seps);
}
if (groups != NULL) {
free(groups);
}
}
}
void nvshmem_nvtx_print_options() {
for (int i = 0; i < num_options; ++i) {
printf("\t %-20s: %s\n", options[i].name, options[i].desc);
}
}
#else /* !NVTX_DISABLE */
#include "internal/host/util.h"
void nvshmem_nvtx_init() {
if (nvshmemi_options.NVTX_provided &&
strcmp_case_insensitive(nvshmemi_options.NVTX, "off") != 0 &&
strcmp(nvshmemi_options.NVTX, "") != 0) {
printf(
"NVTX instrumentation is not available. Check your NVSHMEM "
"build. NVTX support requires C++11 or newer.\n");
}
}
#endif /* !NVTX_DISABLE */