From 4c2fa05a2327fb2f388ffc599a37e1e07245a082 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Thu, 16 Nov 2023 07:57:25 -0800
Subject: [PATCH] msccl: allocate scratch as ext-scope fine-grained (#968)

[ROCm/rccl commit: 569d3f7d59754db703638959a0939bf037dda118]
---
 .../rccl/src/collectives/device/msccl_kernel_impl.h   | 11 +----------
 projects/rccl/src/misc/msccl/msccl_setup.cc           |  2 +-
 2 files changed, 2 insertions(+), 11 deletions(-)

diff --git a/projects/rccl/src/collectives/device/msccl_kernel_impl.h b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
index 99778488ed..bbaf9d5dac 100644
--- a/projects/rccl/src/collectives/device/msccl_kernel_impl.h
+++ b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
@@ -315,11 +315,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
           int16_t dependentStep = mscclShmem.mscclTB.dependentStep[dependentPointer+tid];
           uint64_t goalFlag = COMPUTE_FLAG(workIndex, iter, dependentStep);
           while (true){
-#if defined(__gfx942__)
-            uint64_t curFlag = __atomic_load_n(&(mscclFlags + dependentBid)->flag, __ATOMIC_ACQUIRE);
-#else
             uint64_t curFlag = __atomic_load_n(&(mscclFlags + dependentBid)->flag, __ATOMIC_RELAXED);
-#endif
             if (curFlag >= goalFlag && GET_WORKINDEX_FROM_FLAG(curFlag) == workIndex) break;
           }
         }
@@ -463,13 +459,8 @@ __device__ __forceinline__ void mscclRunInterpreter(
         else
           return;
       }
-      if (t->hasDependence && tid == nthreads-1){
-#if defined(__gfx942__)
-        __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), __ATOMIC_RELEASE);
-#else
+      if (t->hasDependence && tid == nthreads-1)
         __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), __ATOMIC_RELAXED);
-#endif
-      }
       step++;
     }
   }
diff --git a/projects/rccl/src/misc/msccl/msccl_setup.cc b/projects/rccl/src/misc/msccl/msccl_setup.cc
index 199c911a3d..606501af4e 100644
--- a/projects/rccl/src/misc/msccl/msccl_setup.cc
+++ b/projects/rccl/src/misc/msccl/msccl_setup.cc
@@ -73,7 +73,7 @@ ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream) {
   size_t sizeNeeded = (status.nBytes * (size_t)(hostAlgo->nScratchChunks)) / (size_t)(hostAlgo->nChunksPerLoop);
   if (sizeNeeded > status.scratchBufferSize){
     NCCLCHECK(ncclCudaFree(status.scratchBuffer));
-    NCCLCHECK(ncclCudaMalloc((char**)&status.scratchBuffer, sizeNeeded));
+    NCCLCHECK(ncclCudaMalloc((char**)&status.scratchBuffer, sizeNeeded, true));
     status.scratchBufferSize = sizeNeeded;
   }
   return ncclSuccess;