From 69784a7423b2d2647fdae315338663e26d1e989e Mon Sep 17 00:00:00 2001
From: avinashkethineedi <avinash.kethineedi@amd.com>
Date: Thu, 3 Oct 2024 06:28:12 +0000
Subject: [PATCH] Add fence and quiet functionality

* Perform atomic stores to enforce memory ordering


[ROCm/rocshmem commit: 979aed105a88d3df0d8b8e23340c3058f6822d87]
---
 projects/rocshmem/src/ipc/backend_ipc.cpp        | 10 ++++++++++
 projects/rocshmem/src/ipc/backend_ipc.hpp        | 10 ++++++++++
 projects/rocshmem/src/ipc/context_ipc_device.cpp |  8 ++++++++
 projects/rocshmem/src/ipc/context_ipc_device.hpp |  7 +++++++
 4 files changed, 35 insertions(+)

diff --git a/projects/rocshmem/src/ipc/backend_ipc.cpp b/projects/rocshmem/src/ipc/backend_ipc.cpp
index 58caaf53f7..b09adee2ff 100644
--- a/projects/rocshmem/src/ipc/backend_ipc.cpp
+++ b/projects/rocshmem/src/ipc/backend_ipc.cpp
@@ -88,6 +88,8 @@ IPCBackend::IPCBackend(MPI_Comm comm)
 
   roc_shmem_collective_init();
 
+  setup_fence_buffer();
+
   teams_init();
 
   setup_ctxs();
@@ -287,6 +289,14 @@ void IPCBackend::teams_destroy() {
   free(reduced_bitmask_);
 }
 
+void IPCBackend::setup_fence_buffer() {
+  /*
+  * Allocate heap space for fence
+  */
+  fence_pool = reinterpret_cast<int *>(roc_shmem_malloc(
+      sizeof(int) * num_pes));
+}
+
 void IPCBackend::roc_shmem_collective_init() {
   /*
    * Allocate heap space for barrier_sync
diff --git a/projects/rocshmem/src/ipc/backend_ipc.hpp b/projects/rocshmem/src/ipc/backend_ipc.hpp
index cc385f16ee..7d00e6cf07 100644
--- a/projects/rocshmem/src/ipc/backend_ipc.hpp
+++ b/projects/rocshmem/src/ipc/backend_ipc.hpp
@@ -156,6 +156,11 @@ class IPCBackend : public Backend {
    */
   void *pAta_pool{nullptr};
 
+  /**
+   * @brief Handle for raw memory for fence/quiet
+  */
+  int *fence_pool{nullptr};
+
  protected:
    /**
    * @copydoc Backend::dump_backend_stats()
@@ -203,6 +208,11 @@ class IPCBackend : public Backend {
    */
   void roc_shmem_collective_init();
 
+  /**
+   * @brief Allocate buffer for fence/quiet operation
+   */
+  void setup_fence_buffer();
+
  private:
   /**
    * @brief Proxy for the default context
diff --git a/projects/rocshmem/src/ipc/context_ipc_device.cpp b/projects/rocshmem/src/ipc/context_ipc_device.cpp
index 6fbb362d73..036f5cf77b 100644
--- a/projects/rocshmem/src/ipc/context_ipc_device.cpp
+++ b/projects/rocshmem/src/ipc/context_ipc_device.cpp
@@ -47,6 +47,9 @@ __host__ IPCContext::IPCContext(Backend *b)
   barrier_sync = backend->barrier_sync;
   g_ret = bp->g_ret;
   atomic_base_ptr = bp->atomic_ret->atomic_base_ptr;
+  fence_pool = backend->fence_pool;
+
+  orders_.store = detail::atomic::rocshmem_memory_order::memory_order_seq_cst;
 }
 
 __device__ void IPCContext::threadfence_system() {
@@ -85,12 +88,17 @@ __device__ void IPCContext::getmem_nbi(void *dest, const void *source,
 }
 
 __device__ void IPCContext::fence() {
+  for (int i{0}; i < num_pes; i++) {
+    detail::atomic::store<int, detail::atomic::memory_scope_system>(&fence_pool[i], 1, orders_);
+  }
 }
 
 __device__ void IPCContext::fence(int pe) {
+  detail::atomic::store<int, detail::atomic::memory_scope_system>(&fence_pool[pe], 1, orders_);
 }
 
 __device__ void IPCContext::quiet() {
+  fence();
 }
 
 __device__ void *IPCContext::shmem_ptr(const void *dest, int pe) {
diff --git a/projects/rocshmem/src/ipc/context_ipc_device.hpp b/projects/rocshmem/src/ipc/context_ipc_device.hpp
index 4cc6d08d56..450bb2a3cc 100644
--- a/projects/rocshmem/src/ipc/context_ipc_device.hpp
+++ b/projects/rocshmem/src/ipc/context_ipc_device.hpp
@@ -24,6 +24,7 @@
 #define LIBRARY_SRC_IPC_CONTEXT_DEVICE_HPP_
 
 #include "../context.hpp"
+#include "../atomic.hpp"
 
 namespace rocshmem {
 
@@ -232,6 +233,12 @@ class IPCContext : public Context {
 
   //Temporary scratchpad memory used by internal barrier algorithms.
   int64_t *barrier_sync{nullptr};
+
+  //Struct defining memory ordering for atomic operations.
+  detail::atomic::rocshmem_memory_orders orders_{};
+
+  //Buffer to perform Atomic store to enforce memory ordering
+  int *fence_pool{nullptr};
 };
 
 }  // namespace rocshmem