diff --git a/src/include/msccl/msccl_lifecycle.h b/src/include/msccl/msccl_lifecycle.h index eb36b6def6..7735eeed99 100644 --- a/src/include/msccl/msccl_lifecycle.h +++ b/src/include/msccl/msccl_lifecycle.h @@ -32,4 +32,6 @@ ncclResult_t mscclGroupEnd(); ncclResult_t mscclTeardown(); +size_t mscclKernMaxLocalSize(); + #endif diff --git a/src/misc/msccl/msccl_lifecycle.cc b/src/misc/msccl/msccl_lifecycle.cc index 0e7308ec7e..d71c1310d2 100644 --- a/src/misc/msccl/msccl_lifecycle.cc +++ b/src/misc/msccl/msccl_lifecycle.cc @@ -172,7 +172,7 @@ ncclResult_t mscclInit(ncclComm_t comm) { mscclInitialized.store(true, std::memory_order_release); } - INFO(NCCL_INIT, "MSCCL: Initialization finished"); + INFO(NCCL_INIT, "MSCCL: Initialization finished, localSize %ld", mscclKernMaxLocalSize()); return ncclSuccess; } diff --git a/src/misc/msccl/msccl_setup.cc b/src/misc/msccl/msccl_setup.cc index ded6261587..21496f4e89 100644 --- a/src/misc/msccl/msccl_setup.cc +++ b/src/misc/msccl/msccl_setup.cc @@ -85,6 +85,7 @@ ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm) NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 0)); mscclClearIsCallerFlag(); + INFO(NCCL_INIT, "MSCCL: Setup connections finished, used %ld", allocTracker[comm->cudaDev].totalAllocSize); return ncclSuccess; } @@ -288,3 +289,20 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count status.lastStream = stream; return ncclSuccess; } + +// Determine the maximum kernel stack size of all MSCCL kernels +size_t mscclKernMaxLocalSize() { + ncclResult_t res = ncclSuccess; + int numMscclKerns = sizeof(mscclKernelEntries)/sizeof(void *); + hipFuncAttributes attr = {0}; + size_t max = 0; + for (int i = 0; i < numMscclKerns; i++) { + if (mscclKernelEntries[i] != nullptr) { + CUDACHECKGOTO(hipFuncGetAttributes(&attr, reinterpret_cast(mscclKernelEntries[i])), res, error); + if (attr.localSizeBytes > max) max = attr.localSizeBytes; + } + } + +error: + return (res != ncclSuccess) ? 0 : max; +}