From c8b0f2378e7eea965e8e7b82ee69420f1f89b961 Mon Sep 17 00:00:00 2001
From: avinashkethineedi <avinash.kethineedi@amd.com>
Date: Thu, 15 Aug 2024 13:17:44 -0700
Subject: [PATCH] Add gets and puts functionality to IPC context

---
 src/ipc/context_ipc_device.cpp | 53 ++++++++++++++++++++++++++++++++--
 src/ipc/context_ipc_device.hpp |  7 +++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/ipc/context_ipc_device.cpp b/src/ipc/context_ipc_device.cpp
index 97937fcd1e..7b9f3dd469 100644
--- a/src/ipc/context_ipc_device.cpp
+++ b/src/ipc/context_ipc_device.cpp
@@ -32,11 +32,19 @@
 
 #include "config.h"  // NOLINT(build/include_subdir)
 #include "roc_shmem/roc_shmem.hpp"
+#include "backend_ipc.hpp"
 
 namespace rocshmem {
 
 __host__ IPCContext::IPCContext(Backend *b)
     : Context(b, false) {
+  IPCBackend *backend{static_cast<IPCBackend *>(b)};
+  ipcImpl = &backend->ipcImpl;
+
+  auto *bp{backend->ipc_backend_proxy.get()};
+
+  g_ret = bp->g_ret;
+  atomic_base_ptr = bp->atomic_ret->atomic_base_ptr;
 }
 
 __device__ void IPCContext::threadfence_system() {
@@ -50,18 +58,32 @@ __device__ void IPCContext::ctx_destroy(){
 
 __device__ void IPCContext::putmem(void *dest, const void *source, size_t nelems,
                                   int pe) {
+  // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)
+  int local_pe = pe % ipcImpl->shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl->ipc_bases[my_pe];
+  ipcImpl->ipcCopy(ipcImpl->ipc_bases[local_pe] + L_offset,
+                   const_cast<void *>(source), nelems);
 }
 
 __device__ void IPCContext::getmem(void *dest, const void *source, size_t nelems,
                                   int pe) {
+  // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)
+  int local_pe = pe % ipcImpl->shm_size;
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset =
+      const_cast<char *>(src_typed) - ipcImpl->ipc_bases[my_pe];
+  ipcImpl->ipcCopy(dest, ipcImpl->ipc_bases[local_pe] + L_offset, nelems);
 }
 
 __device__ void IPCContext::putmem_nbi(void *dest, const void *source,
                                       size_t nelems, int pe) {
+  putmem(dest, source, nelems, pe);
 }
 
 __device__ void IPCContext::getmem_nbi(void *dest, const void *source,
                                       size_t nelems, int pe) {
+  getmem(dest, source, nelems, pe);
 }
 
 __device__ void IPCContext::fence() {
@@ -92,38 +114,65 @@ __device__ void IPCContext::sync(roc_shmem_team_t team) {
 
 __device__ void IPCContext::putmem_wg(void *dest, const void *source,
                                      size_t nelems, int pe) {
+  // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)
+  int local_pe = pe % ipcImpl->shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl->ipc_bases[my_pe];
+  ipcImpl->ipcCopy_wg(ipcImpl->ipc_bases[local_pe] + L_offset,
+                      const_cast<void *>(source), nelems);
   __syncthreads();
 }
 
 __device__ void IPCContext::getmem_wg(void *dest, const void *source,
                                      size_t nelems, int pe) {
+  // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)
+  int local_pe = pe % ipcImpl->shm_size;
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset =
+      const_cast<char *>(src_typed) - ipcImpl->ipc_bases[my_pe];
+  ipcImpl->ipcCopy_wg(dest, ipcImpl->ipc_bases[local_pe] + L_offset, nelems);
   __syncthreads();
 }
 
 __device__ void IPCContext::putmem_nbi_wg(void *dest, const void *source,
                                          size_t nelems, int pe) {
-  __syncthreads();
+  putmem_wg(dest, source, nelems, pe);
 }
 
 __device__ void IPCContext::getmem_nbi_wg(void *dest, const void *source,
                                          size_t nelems, int pe) {
-  __syncthreads();
+  getmem_wg(dest, source, nelems, pe);
 }
 
 __device__ void IPCContext::putmem_wave(void *dest, const void *source,
                                        size_t nelems, int pe) {
+  // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)
+  int local_pe = pe % ipcImpl->shm_size;
+  uint64_t L_offset =
+      reinterpret_cast<char *>(dest) - ipcImpl->ipc_bases[my_pe];
+  ipcImpl->ipcCopy_wave(ipcImpl->ipc_bases[local_pe] + L_offset,
+                        const_cast<void *>(source), nelems);
 }
 
 __device__ void IPCContext::getmem_wave(void *dest, const void *source,
                                        size_t nelems, int pe) {
+  // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)
+  int local_pe = pe % ipcImpl->shm_size;
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset =
+      const_cast<char *>(src_typed) - ipcImpl->ipc_bases[my_pe];
+  ipcImpl->ipcCopy_wave(dest, ipcImpl->ipc_bases[local_pe] + L_offset,
+                        nelems);
 }
 
 __device__ void IPCContext::putmem_nbi_wave(void *dest, const void *source,
                                            size_t nelems, int pe) {
+  putmem_wave(dest, source, nelems, pe);
 }
 
 __device__ void IPCContext::getmem_nbi_wave(void *dest, const void *source,
                                            size_t nelems, int pe) {
+  getmem_wave(dest, source, nelems, pe);
 }
 
 }  // namespace rocshmem
diff --git a/src/ipc/context_ipc_device.hpp b/src/ipc/context_ipc_device.hpp
index 492d63ca04..afadb9828c 100644
--- a/src/ipc/context_ipc_device.hpp
+++ b/src/ipc/context_ipc_device.hpp
@@ -234,6 +234,13 @@ class IPCContext : public Context {
   
  private:
 
+ //context class has IpcImpl object (ipcImpl_)
+ IpcImpl *ipcImpl{nullptr};
+
+ uint64_t* atomic_base_ptr{nullptr};
+
+ char* g_ret;
+
 };
 
 }  // namespace rocshmem