From 1538b908ac147bd0f550df47112c88c6cce0ed18 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Thu, 8 Feb 2024 12:09:56 -0800
Subject: [PATCH] msccl: use relaxed atomics on scratch buffer (#1075)

[ROCm/rccl commit: 5257c753c59ee708c96aa8c0eb296e6f9e2f98a1]
---
 projects/rccl/src/collectives/device/msccl_kernel_impl.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/projects/rccl/src/collectives/device/msccl_kernel_impl.h b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
index 3cbad93175..5e9b992e0e 100644
--- a/projects/rccl/src/collectives/device/msccl_kernel_impl.h
+++ b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
@@ -372,7 +372,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
           return;
       }
       if (t->hasDependence && tid == nthreads-1)
-        __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), (t->type == MSCCL_REDUCE || t->type == MSCCL_RECV) ? __ATOMIC_RELEASE : __ATOMIC_RELAXED);
+        __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), ((t->type == MSCCL_REDUCE || t->type == MSCCL_RECV) && (t->dstBuffer != MSCCL_SCRATCH_BUFFER)) ? __ATOMIC_RELEASE : __ATOMIC_RELAXED);
       step++;
     }
   }