msccl: print stack and memory usage (#723)

* msccl: print stack and memory usage

* Update number of kernels calculation

src/include/msccl/msccl_lifecycle.h

@@ -32,4 +32,6 @@ ncclResult_t mscclGroupEnd();
 
 ncclResult_t mscclTeardown();
 
+size_t mscclKernMaxLocalSize();
+
 #endif

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;
 }
 

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<const void*>(mscclKernelEntries[i])), res, error);
+      if (attr.localSizeBytes > max) max = attr.localSizeBytes;
+    }
+  }
+
+error:
+  return (res != ncclSuccess) ? 0 : max;
+}