AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

msccl: enable basic collective trace (#959)

Browse Source 
To avoid increasing number of kernels, colltrace is only enabled with
RCCL_MSCCL_FORCE_FULLOPS=1
This commit is contained in:
Wenkai Du
2023-11-08 20:14:28 -08:00
committed by GitHub
parent 8e0258a73d
commit 5a800e00cd
5 changed files with 69 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/misc/msccl/msccl_lifecycle.cc
+19
View File
@@ -305,6 +305,7 @@ static ncclResult_t mscclSetSavedSchedulerParam(
param->p.nRanks = comm->nRanks;
param->comm = comm;
param->stream = stream;
param->p.opCount = comm->opCount;
return ncclSuccess;
}
@@ -322,9 +323,27 @@ static ncclResult_t mscclSaveCountsAndDispls(struct mscclSavedSchedulerParam* pa
return ncclSuccess;
}
const char *mscclFuncNames[] = {
"mscclFuncReduce",
"mscclFuncBroadcast",
"mscclFuncAllReduce",
"mscclFuncReduceScatter",
"mscclFuncAllGather",
"mscclFuncSend",
"mscclFuncRecv",
"mscclFuncGather",
"mscclFuncScatter",
"mscclFuncAllToAll",
"mscclFuncAllToAllv",
};
static ncclResult_t mscclRunSavedParams() {
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
for (auto& param : threadLocalStatus.savedSchedulerParams) {
INFO(NCCL_COLL,"%s: opCount %lx sendbuff %p recvbuff %p count %zi datatype %d op %d root %d comm %p [nranks=%d] stream %p",
mscclFuncNames[param.p.func], param.p.opCount, param.p.sendBuff, param.p.recvBuff, param.p.count,
param.p.dataType, param.p.op, param.p.root, param.comm, param.p.nRanks, param.stream);
NCCLCHECK(mscclRunAlgo(
param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
src/misc/msccl/msccl_setup.cc
+12 -10
View File
@@ -391,6 +391,16 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
ncclDevRedOpFull opFull = {};
NCCLCHECK(hostToDevRedOp(&opFull, op, dataType, comm));
uint32_t fnIndex = (opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol;
uint8_t fullOpMask = (1<<MSCCL_RECV_COPY_SEND) |
(1<<MSCCL_RECV_REDUCE_SEND) |
(1<<MSCCL_RECV_REDUCE_COPY_SEND) |
(1<<MSCCL_RECV_REDUCE_COPY) |
(1<<MSCCL_LOCAL_COPY);
//check if need full ops msccl kernel
if ((hostAlgo->typeMask & fullOpMask) || rcclParamMscclForceFullOps())
fnIndex += sizeof(mscclKernelEntries)/sizeof(void *)/2;
mscclWork work;
work.syncFlags = status.syncFlags;
work.scratchBuffer = status.scratchBuffer;
@@ -403,7 +413,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
work.maxAllowedCount = status.maxAllowedCount;
work.hasReduce = hostAlgo->hasReduce;
work.redOpArgIsPtr = opFull.scalarArgIsPtr;
INFO(NCCL_COLL, "MSCCL: typeMask %x Setup Kernel finished", hostAlgo->typeMask);
work.fnIndex = fnIndex;
INFO(NCCL_COLL, "MSCCL: typeMask %x fnIndex %d Setup Kernel finished", hostAlgo->typeMask, fnIndex);
uint32_t workFifoIdxMask = status.workFifoDepth - 1;
uint32_t workFifoSent = status.workFifoSent;
@@ -428,15 +439,6 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
struct mscclWork *workPtr = status.workFifo + (workFifoSent & workFifoIdxMask);
void *args[3] = {&comm->devComm, &devAlgo, &workPtr};
uint32_t fnIndex = (opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol;
uint8_t fullOpMask = (1<<MSCCL_RECV_COPY_SEND) |
(1<<MSCCL_RECV_REDUCE_SEND) |
(1<<MSCCL_RECV_REDUCE_COPY_SEND) |
(1<<MSCCL_RECV_REDUCE_COPY) |
(1<<MSCCL_LOCAL_COPY);
//check if need full ops msccl kernel
if ((hostAlgo->typeMask & fullOpMask) || rcclParamMscclForceFullOps())
fnIndex += sizeof(mscclKernelEntries)/sizeof(void *)/2;
void *func = mscclKernelEntries[fnIndex];
if (enableDoneEvent) {
CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, comm->doneEvent, 0));