Update ext-profiler example

Sync ext-profiler example with 2.26.2.
This commit is contained in:
Giuseppe Congiu
2025-04-09 09:02:40 -07:00
committed by Sylvain Jeaugey
parent f44ac759fe
commit 145e67e707
13 changed files with 621 additions and 78 deletions
+99 -6
View File
@@ -58,6 +58,7 @@ __hidden double gettime(void) {
static pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
static pid_t pid;
static int* eActivationMaskPtr;
__hidden ncclResult_t exampleProfilerInit(void** context, int* eActivationMask) {
pthread_mutex_lock(&lock);
@@ -65,7 +66,7 @@ __hidden ncclResult_t exampleProfilerInit(void** context, int* eActivationMask)
// first thread initializes event mask, environment and detach pool
const char* str;
str = getenv("NCCL_PROFILE_EVENT_MASK");
__atomic_store_n(eActivationMask, str ? atoi(str) : defaultEActivationMask, __ATOMIC_RELAXED);
__atomic_store_n(eActivationMask, str ? atoi(str) : 0, __ATOMIC_RELAXED);
str = getenv("NCCL_PROFILE_GROUP_POOL_SIZE");
groupPoolSize = str ? atoi(str) : defaultGroupPoolSize;
@@ -100,6 +101,9 @@ __hidden ncclResult_t exampleProfilerInit(void** context, int* eActivationMask)
}
pthread_mutex_unlock(&lock);
// store pointer to activation mask globally
eActivationMaskPtr = eActivationMask;
// pre-allocate memory for event object pools in dedicated profiler context
struct context* ctx = (struct context *)calloc(1, sizeof(*ctx));
ctx->groupPool = (struct group *)calloc(groupPoolSize, sizeof(*ctx->groupPool));
@@ -199,8 +203,6 @@ __hidden ncclResult_t exampleProfilerStartEvent(void* context, void** eHandle, n
if (base->type == ncclProfileColl) {
struct collective* c = (struct collective *)base;
// reset event proxyOps & proxySteps
memset(c->send, 0, sizeof(struct proxyOp)*MAX_CHANNELS*MAX_OPS);
memset(c->recv, 0, sizeof(struct proxyOp)*MAX_CHANNELS*MAX_OPS);
memset(c->nProxyOps, 0, sizeof(int)*MAX_CHANNELS);
// release collective events in the group and return them to the collective pool
__atomic_fetch_add(&ctx->collPoolBase, 1, __ATOMIC_RELAXED);
@@ -252,7 +254,6 @@ __hidden ncclResult_t exampleProfilerStartEvent(void* context, void** eHandle, n
event->count = eDescr->coll.count;
event->root = eDescr->coll.root;
event->datatype = eDescr->coll.datatype;
event->trafficBytes = eDescr->coll.trafficBytes;
event->nMaxChannels = eDescr->coll.nMaxChannels;
event->nWarps = eDescr->coll.nWarps;
event->algo = eDescr->coll.algo;
@@ -373,7 +374,7 @@ __hidden ncclResult_t exampleProfilerStartEvent(void* context, void** eHandle, n
__atomic_fetch_add(&parent->base.refCount, 1, __ATOMIC_RELAXED);
debugEvent(event, "ProxyOpStart");
}
} else if (eDescr->type == ncclProfileProxyStep) {
} else if (eDescr->type == ncclProfileProxyStep) {
// the parent might be null if we run out of events
struct proxyOp* parent = (struct proxyOp *)eDescr->parentObj;
if (parent == NULL) return ncclSuccess;
@@ -385,8 +386,77 @@ __hidden ncclResult_t exampleProfilerStartEvent(void* context, void** eHandle, n
event->isSend = parent->isSend;
event->parent = parent;
event->startTs = gettime() - startTime;
event->nNetEvents = 0;
*eHandle = event;
debugEvent(event, "ProxyStepStart");
} else if (eDescr->type == ncclProfileKernelCh) {
struct taskEventBase* eventBase = (struct taskEventBase *)eDescr->parentObj;
if (eventBase == NULL) return ncclSuccess;
if (eventBase->type == ncclProfileColl) {
struct collective* parent = (struct collective *)eDescr->parentObj;
struct kernelCh* event = &parent->kernel[eDescr->kernelCh.channelId];
event->type = ncclProfileKernelCh;
event->channelId = eDescr->kernelCh.channelId;
event->parent = eventBase;
event->startTs = gettime() - startTime;
*eHandle = event;
__atomic_fetch_add(&parent->base.refCount, 1, __ATOMIC_RELAXED);
debugEvent(event, "KernelChStart");
} else { // ncclProfileP2p
struct p2p* parent = (struct p2p *)eDescr->parentObj;
struct kernelCh* event = &parent->kernel[eDescr->kernelCh.channelId];
event->type = ncclProfileKernelCh;
event->channelId = eDescr->kernelCh.channelId;
event->parent = eventBase;
event->startTs = gettime() - startTime;
*eHandle = event;
__atomic_fetch_add(&parent->base.refCount, 1, __ATOMIC_RELAXED);
debugEvent(event, "KernelChStart");
}
} else if (eDescr->type == ncclProfileNetPlugin) {
struct proxyStep* parent = (struct proxyStep *)eDescr->parentObj;
if (parent == NULL) return ncclSuccess;
int64_t pluginId = eDescr->netPlugin.id;
int64_t type = pluginId & NCCL_PROFILER_NET_TYPE_MASK;
int64_t ver = pluginId & NCCL_PROFILER_NET_VER_MASK;
if (type == NCCL_PROFILER_NET_TYPE_IB) {
if (ver == 1) {
ncclProfilerNetIbDescr_v1_t* descr = (ncclProfilerNetIbDescr_v1_t *)eDescr->netPlugin.data;
struct netPlugin* event = parent->net + __atomic_fetch_add(&parent->nNetEvents, 1, __ATOMIC_RELAXED);
event->type = ncclProfileNetPlugin;
event->pluginType = type;
event->pluginVer = ver;
if (descr->type == ncclProfileQp) {
event->pluginEvent = ncclProfileQp;
event->qp.device = descr->qp.device;
event->qp.wr_id = descr->qp.wr_id;
event->qp.opcode = descr->qp.opcode;
event->qp.qpNum = descr->qp.qpNum;
event->qp.length = descr->qp.length;
}
event->startTs = gettime() - startTime;
*eHandle = event;
debugEvent(event, "NetPluginStart");
}
} else if (type == NCCL_PROFILER_NET_TYPE_SOCK) {
if (ver == 1) {
ncclProfilerNetSockDescr_v1_t* descr = (ncclProfilerNetSockDescr_v1_t *)eDescr->netPlugin.data;
struct netPlugin* event = parent->net + __atomic_fetch_add(&parent->nNetEvents, 1, __ATOMIC_RELAXED);
event->type = ncclProfileNetPlugin;
event->pluginType = type;
event->pluginVer = ver;
if (descr->type == ncclProfileSocket) {
event->pluginEvent = ncclProfileSocket;
event->sock.fd = descr->sock.fd;
event->sock.op = descr->sock.op;
event->sock.length = descr->sock.length;
}
event->startTs = gettime() - startTime;
*eHandle = event;
debugEvent(event, "NetPluginStart");
}
}
}
return ncclSuccess;
}
@@ -445,6 +515,15 @@ void updateEvent(void* handle) {
struct proxyCtrl* event = (struct proxyCtrl *)handle;
event->stopTs = gettime() - startTime;
debugEvent(event, "ProxyCtrlStop");
} else if (type == ncclProfileKernelCh) {
struct kernelCh* event = (struct kernelCh *)handle;
event->stopTs = gettime() - startTime;
updateEvent(event->parent);
debugEvent(event, "KernelChStop");
} else if (type == ncclProfileNetPlugin) {
struct netPlugin* event = (struct netPlugin *)handle;
event->stopTs = gettime() - startTime;
debugEvent(event, "NetPluginStop");
}
}
@@ -506,7 +585,7 @@ __hidden ncclResult_t exampleProfilerRecordEventState(void* eHandle, ncclProfile
return ncclSuccess;
}
ncclProfiler_t ncclProfiler_v2 = {
ncclProfiler_t ncclProfiler_v3 = {
"Example-profiler",
exampleProfilerInit,
exampleProfilerStartEvent,
@@ -514,3 +593,17 @@ ncclProfiler_t ncclProfiler_v2 = {
exampleProfilerRecordEventState,
exampleProfilerFinalize,
};
int exampleProfilerStart(int eActivationMask) {
if (__atomic_load_n(&initialized, __ATOMIC_RELAXED)) {
__atomic_store_n(eActivationMaskPtr, eActivationMask, __ATOMIC_RELAXED);
}
return ncclSuccess;
}
int exampleProfilerStop(void) {
if (__atomic_load_n(&initialized, __ATOMIC_RELAXED)) {
__atomic_store_n(eActivationMaskPtr, 0, __ATOMIC_RELAXED);
}
return ncclSuccess;
}