From 3eeaea3f00b0deae27a7435b544c9fcb8a6c747e Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Tue, 31 Oct 2023 14:46:57 -0700
Subject: [PATCH] msccl: use atomic to set dependency flags (#941)

[ROCm/rccl commit: a7400218a23273f9cd190db75c34c5b4d986c48b]
---
 .../rccl/src/collectives/device/msccl_kernel_impl.h  | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/projects/rccl/src/collectives/device/msccl_kernel_impl.h b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
index 44505d1a13..b439212439 100644
--- a/projects/rccl/src/collectives/device/msccl_kernel_impl.h
+++ b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
@@ -307,7 +307,11 @@ __device__ __forceinline__ void mscclRunInterpreter(
           int16_t dependentStep = mscclShmem.mscclTB.dependentStep[dependentPointer+tid];
           uint64_t goalFlag = COMPUTE_FLAG(workIndex, iter, dependentStep);
           while (true){
-            uint64_t curFlag = (mscclFlags + dependentBid)->flag;
+#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;
           }
         }
@@ -451,7 +455,11 @@ __device__ __forceinline__ void mscclRunInterpreter(
           return;
       }
       if (t->hasDependence && tid == nthreads-1){
-        mscclFlags[bid].flag = (uint64_t) COMPUTE_FLAG(workIndex, iter, step);
+#if defined(__gfx942__)
+        __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), __ATOMIC_RELEASE);
+#else
+        __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), __ATOMIC_RELAXED);
+#endif
       }
       step++;
     }