From 3e66d42d954ac48118c8554dbbbe432a08e3bbdb Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Tue, 30 Jan 2024 08:22:50 -0800
Subject: [PATCH] msccl: ensure memory coherence after data receive (#1062)

[ROCm/rccl commit: f7550d83b83fe2685d7066d30c5323c42e1e9245]
---
 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 11894644a6..3cbad93175 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 ? __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) ? __ATOMIC_RELEASE : __ATOMIC_RELAXED);
       step++;
     }
   }