Fix bug where the first collective call was using MSCCL instead of MSCCL++ (#1260)

src/include/msccl/msccl_struct.h

@@ -165,6 +165,7 @@ struct mscclSavedSchedulerParam {
 };
 
 enum mscclCaptureStatus {
+  mscclUnknownCaptureStatus,
   mscclNoCapture,
   mscclNewCapture,
   mscclExistingCapture

src/misc/msccl/msccl_lifecycle.cc

@@ -470,6 +470,11 @@ ncclResult_t mscclEnqueueCheck(
     case mscclNoGroup:
 #ifdef ENABLE_MSCCLPP
       if (comm->mscclppCompatible) {
+        if (threadLocalStatus.captureStatus == mscclUnknownCaptureStatus) {
+          INFO(NCCL_COLL, "MSCCL++: reading capture status");
+          NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
+        }
+
         /* check if one rank per GPU and graph mode is enabled */
         if ((threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
           if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold && (nBytes & 31) == 0) {
@@ -501,6 +506,11 @@ ncclResult_t mscclEnqueueCheck(
     case mscclGroupSupportedOp:
 #ifdef ENABLE_MSCCLPP
       if (comm->mscclppCompatible) {
+        if (threadLocalStatus.captureStatus == mscclUnknownCaptureStatus) {
+          INFO(NCCL_COLL, "MSCCL++: reading capture status");
+          NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
+        }
+
         /* check if one rank per GPU and graph mode is enabled */
         if ((threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
           if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold && (nBytes & 31) == 0) {

src/misc/msccl/msccl_setup.cc

@@ -187,6 +187,10 @@ ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm, hipStr
   mscclStatus& status = mscclGetStatus(comm->rank);
   mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
   mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs(comm->rank);
+  if (threadLocalStatus.captureStatus == mscclUnknownCaptureStatus) {
+    INFO(NCCL_NET, "mscclSetupProxy: reading capture status");
+    NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
+  }
   if (threadLocalStatus.captureStatus == mscclNoCapture) {
     INFO(NCCL_NET,"mscclSetupProxy: no capture\n");
     NCCLCHECK(mscclSetupProxyImpl(hostAlgo, comm));
@@ -476,6 +480,10 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
   work.fnIndex = fnIndex;
   INFO(NCCL_COLL, "MSCCL: typeMask %x fnIndex %d Setup Kernel finished", hostAlgo->typeMask, fnIndex);
 
+  if (threadLocalStatus.captureStatus == mscclUnknownCaptureStatus) {
+    INFO(NCCL_NET, "MSCCL: reading capture status");
+    NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
+  }
   mscclWorkFifoStatus* workFifoStatus = nullptr;
   if (threadLocalStatus.captureStatus == mscclNoCapture) {
     workFifoStatus = &(status.defaultWorkFifoStatus);