From e1e1ac6df6be452d3ca2e6a030ab44b85e437d09 Mon Sep 17 00:00:00 2001
From: avinashkethineedi <avinash.kethineedi@amd.com>
Date: Wed, 28 Aug 2024 08:30:46 -0700
Subject: [PATCH] Add atomics

* Add atomic_add, atomic_set, atomic_cas, atomic_fetch_add and atomic_fetch_cas to IPC backend
---
 src/ipc/context_ipc_tmpl_device.hpp | 44 +++++++++++++++++++++--------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/ipc/context_ipc_tmpl_device.hpp b/src/ipc/context_ipc_tmpl_device.hpp
index 28cd718f0a..5e697cafaf 100644
--- a/src/ipc/context_ipc_tmpl_device.hpp
+++ b/src/ipc/context_ipc_tmpl_device.hpp
@@ -70,13 +70,21 @@ __device__ void IPCContext::get_nbi(T *dest, const T *source, size_t nelems,
 
 // Atomics
 template <typename T>
-__device__ void IPCContext::amo_add(void *dst, T value, int pe) {
-  assert(false);
+__device__ void IPCContext::amo_add(void *dest, T value, int pe) {
+  int local_pe = pe % ipcImpl_.shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl_.ipc_bases[my_pe];
+  ipcImpl_.ipcAMOAdd(
+      reinterpret_cast<T *>(ipcImpl_.ipc_bases[local_pe] + L_offset), value);
 }
 
 template <typename T>
-__device__ void IPCContext::amo_set(void *dst, T value, int pe) {
-  assert(false);
+__device__ void IPCContext::amo_set(void *dest, T value, int pe) {
+  int local_pe = pe % ipcImpl_.shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl_.ipc_bases[my_pe];
+  ipcImpl_.ipcAMOSet(
+      reinterpret_cast<T *>(ipcImpl_.ipc_bases[local_pe] + L_offset), value);
 }
 
 template <typename T>
@@ -119,20 +127,32 @@ __device__ void IPCContext::amo_xor(void *dst, T value, int pe) {
 }
 
 template <typename T>
-__device__ void IPCContext::amo_cas(void *dst, T value, T cond, int pe) {
-  assert(false);
+__device__ void IPCContext::amo_cas(void *dest, T value, T cond, int pe) {
+  int local_pe = pe % ipcImpl_.shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl_.ipc_bases[my_pe];
+  ipcImpl_.ipcAMOCas(
+      reinterpret_cast<T *>(ipcImpl_.ipc_bases[local_pe] + L_offset), cond,
+      value);
 }
 
 template <typename T>
-__device__ T IPCContext::amo_fetch_add(void *dst, T value, int pe) {
-  assert(false);
-  return 0;
+__device__ T IPCContext::amo_fetch_add(void *dest, T value, int pe) {
+  int local_pe = pe % ipcImpl_.shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl_.ipc_bases[my_pe];
+  return ipcImpl_.ipcAMOFetchAdd(
+      reinterpret_cast<T *>(ipcImpl_.ipc_bases[local_pe] + L_offset), value);
 }
 
 template <typename T>
-__device__ T IPCContext::amo_fetch_cas(void *dst, T value, T cond, int pe) {
-  assert(false);
-  return 0;
+__device__ T IPCContext::amo_fetch_cas(void *dest, T value, T cond, int pe) {
+  int local_pe = pe % ipcImpl_.shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl_.ipc_bases[my_pe];
+  return ipcImpl_.ipcAMOFetchCas(
+      reinterpret_cast<T *>(ipcImpl_.ipc_bases[local_pe] + L_offset), cond,
+      value);
 }
   
 // Collectives