From aa605b0989d29e403bdc28745d2f85be61657942 Mon Sep 17 00:00:00 2001
From: Ziyue Yang <yzylivezh@hotmail.com>
Date: Sat, 25 Mar 2023 01:59:36 +0800
Subject: [PATCH] fix msccl stream usage (#717)

[ROCm/rccl commit: c8e33b1232986456acac3c56c8fab7c48b2e63a1]
---
 projects/rccl/src/misc/msccl/msccl_lifecycle.cc | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc
index 67177e2a19..0e7308ec7e 100644
--- a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc
+++ b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc
@@ -375,7 +375,11 @@ ncclResult_t mscclEnqueueCheck(
   switch (threadLocalStatus.groupStatus) {
     case mscclNoGroup:
       if (comm->mscclCompatible) {
-        CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        if (stream == (hipStream_t)0) {
+          captureStatus = hipStreamCaptureStatusNone;
+        } else {
+          CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        }
         if (captureStatus == hipStreamCaptureStatusNone) {
           NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
           if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
@@ -388,7 +392,11 @@ ncclResult_t mscclEnqueueCheck(
       break;
     case mscclGroupSupportedOp:
       if (comm->mscclCompatible) {
-        CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        if (stream == (hipStream_t)0) {
+          captureStatus = hipStreamCaptureStatusNone;
+        } else {
+          CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        }
         if (captureStatus == hipStreamCaptureStatusNone) {
           NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
           if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {