From 69135976d6487ec895b8a6ad48dccbf1edfacc1c Mon Sep 17 00:00:00 2001
From: corey-derochie-amd
 <161367113+corey-derochie-amd@users.noreply.github.com>
Date: Mon, 22 Jul 2024 15:46:47 -0600
Subject: [PATCH] Fix bug where the first collective call was using MSCCL
 instead of MSCCL++ (#1260)

---
 src/include/msccl/msccl_struct.h  |  1 +
 src/misc/msccl/msccl_lifecycle.cc | 10 ++++++++++
 src/misc/msccl/msccl_setup.cc     |  8 ++++++++
 3 files changed, 19 insertions(+)

diff --git a/src/include/msccl/msccl_struct.h b/src/include/msccl/msccl_struct.h
index f8a5d7e64a..351a557b11 100644
--- a/src/include/msccl/msccl_struct.h
+++ b/src/include/msccl/msccl_struct.h
@@ -165,6 +165,7 @@ struct mscclSavedSchedulerParam {
 };
 
 enum mscclCaptureStatus {
+  mscclUnknownCaptureStatus,
   mscclNoCapture,
   mscclNewCapture,
   mscclExistingCapture
diff --git a/src/misc/msccl/msccl_lifecycle.cc b/src/misc/msccl/msccl_lifecycle.cc
index 1d2b482f29..016137d126 100644
--- a/src/misc/msccl/msccl_lifecycle.cc
+++ b/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) {
diff --git a/src/misc/msccl/msccl_setup.cc b/src/misc/msccl/msccl_setup.cc
index 1068887067..7299820455 100644
--- a/src/misc/msccl/msccl_setup.cc
+++ b/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);