From 9cc9c3360bac062c4fbac63e58d060e6f2a68663 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Fri, 23 Apr 2021 16:36:19 -0700
Subject: [PATCH] Control collective trace from kernel separately (#356)

---
 src/init.cc | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/init.cc b/src/init.cc
index 6b90197f79..359dac875b 100644
--- a/src/init.cc
+++ b/src/init.cc
@@ -165,10 +165,14 @@ void NCCL_NO_OPTIMIZE commPoison(ncclComm_t comm) {
   comm->rank = comm->cudaDev = comm->busId = comm->nRanks = -1;
 }
 
+RCCL_PARAM(KernelCollEnable, "KERNEL_COLL_ENABLE", 0);
+RCCL_PARAM(KernelCollPrint, "KERNEL_COLL_PRINT", 0);
+
 #ifdef ENABLE_COLLTRACE
 void *ncclCommThreadMain(void *arg) {
   ncclComm_t comm = (ncclComm_t)arg;
   int head = comm->hostDevComm.collTraceHead;
+  bool kern_print = rcclParamKernelCollPrint() && rcclParamKernelCollEnable();
   #define MAX_NAME_LENGTH 32
   char* func_names = (char *)malloc(MAX_NAME_LENGTH*(FUNC_INDEX_P2P+1));
   for (int func = 0; func < NCCL_NUM_FUNCTIONS; func++) {
@@ -252,7 +256,7 @@ void *ncclCommThreadMain(void *arg) {
             break;
         }
       }
-      INFO(NCCL_COLL, "%s", line);
+      if (kern_print) INFO(NCCL_INIT, "%s", line);
       STORE(&(td->type), ncclCollTraceNotReady);
       head ++;
       head %= COLLTRACE_NUM_ITEMS;
@@ -416,7 +420,7 @@ static ncclResult_t commAlloc(ncclComm_t* comret, int ndev, int rank) {
   NCCLCHECK(ncclCudaHostCalloc(&comm->hostDevComm.collTrace, COLLTRACE_NUM_ITEMS));
   memset(comm->hostDevComm.collTrace, 0, sizeof(struct ncclCollTrace) * COLLTRACE_NUM_ITEMS);
   comm->hostDevComm.collTraceExit = comm->hostDevComm.collTraceHead = *comm->hostDevComm.collTraceTail = 0;
-  if ((ncclDebugLevel >= NCCL_LOG_INFO) && (ncclDebugMask & NCCL_COLL))
+  if ((ncclDebugLevel >= NCCL_LOG_INFO) && rcclParamKernelCollEnable())
     pthread_create(&comm->hostDevComm.collTraceThread, NULL, ncclCommThreadMain, (void *)comm);
   else
     comm->hostDevComm.collTraceThread = 0;