From 8334214b98fd91a4c75c24d1141f6cc7774755f7 Mon Sep 17 00:00:00 2001
From: Yiltan Hassan Temucin <yiltan.temucin@amd.com>
Date: Mon, 7 Oct 2024 06:06:42 -0700
Subject: [PATCH] added notifier->sync() when we are not using cooperative
 groups

updated scope bug


[ROCm/rocshmem commit: b1134e8633ecb1f9670bad7540b22857f1d87892]
---
 projects/rocshmem/src/ipc/context_ipc_device.hpp  | 12 ++++++++++++
 .../rocshmem/src/ipc/context_ipc_device_coll.cpp  | 15 ++++++---------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/projects/rocshmem/src/ipc/context_ipc_device.hpp b/projects/rocshmem/src/ipc/context_ipc_device.hpp
index 48fe5acbdf..8046fa2088 100644
--- a/projects/rocshmem/src/ipc/context_ipc_device.hpp
+++ b/projects/rocshmem/src/ipc/context_ipc_device.hpp
@@ -27,6 +27,13 @@
 #include "../atomic.hpp"
 #include "../team.hpp"
 
+#ifdef USE_COOPERATIVE_GROUPS
+#include <hip/hip_cooperative_groups.h>
+namespace cg = cooperative_groups;
+#else
+#include "../memory/notifier.hpp"
+#endif /* USE_COOPERATIVE_GROUPS */
+
 namespace rocshmem {
 
 class IPCContext : public Context {
@@ -235,6 +242,11 @@ class IPCContext : public Context {
   //Temporary scratchpad memory used by internal barrier algorithms.
   int64_t *barrier_sync{nullptr};
 
+#ifndef USE_COOPERATIVE_GROUPS
+  using NOTIFIER_PROXY_T = NotifierProxy<HIPAllocator, detail::atomic::memory_scope_agent>;
+  NOTIFIER_PROXY_T notifier_{};
+#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */
+
   //Struct defining memory ordering for atomic operations.
   detail::atomic::rocshmem_memory_orders orders_{};
 
diff --git a/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp b/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp
index 5e1f950c3b..02bb3ffcec 100644
--- a/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp
+++ b/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp
@@ -26,11 +26,6 @@
 #include "../util.hpp"
 #include "ipc_team.hpp"
 
-#ifdef USE_COOPERATIVE_GROUPS
-#include <hip/hip_cooperative_groups.h>
-namespace cg = cooperative_groups;
-#endif /* USE_COOPERATIVE_GROUPS */
-
 namespace rocshmem {
 
 __device__ void IPCContext::internal_direct_barrier(int pe, int PE_start,
@@ -92,11 +87,13 @@ __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride,
 #ifdef USE_COOPERATIVE_GROUPS
   cg::grid_group grid = cg::this_grid();
   grid.sync();
-  if (0 == grid.thread_rank()) {
+  if (0 == grid.thread_rank())
 #else
-  __syncthreads();
-  if (is_thread_zero_in_block()) {
+  auto notifier{notifier_.get()};
+  notifier->sync();
+  if (is_thread_zero_in_block() && is_block_zero_in_grid())
 #endif /* USE_COOPERATIVE_GROUPS */
+  {
     if (PE_size < 64) {
       internal_direct_barrier(pe, PE_start, stride, PE_size, pSync);
     } else {
@@ -107,7 +104,7 @@ __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride,
 #ifdef USE_COOPERATIVE_GROUPS
   grid.sync();
 #else
-  __syncthreads();
+  notifier->sync();
 #endif /* USE_COOPERATIVE_GROUPS */
 }