Implement rocshmem_g API and optimize memory usage

- Implement `rocshmem_g` API
- Free up memory space allocated for `rocshmem_g` and atomic operations' return values

src/ipc/backend_ipc.cpp

@@ -72,10 +72,6 @@ IPCBackend::IPCBackend(MPI_Comm comm)
    */
   assert(num_pes == ipcImpl.shm_size);
 
-  auto *bp{ipc_backend_proxy.get()};
-
-  bp->heap_ptr = &heap;
-
   /* Initialize the host interface */
   host_interface = std::make_shared<HostInterface>(hdp_proxy_.get(),
                                                  thread_comm,
@@ -85,10 +81,6 @@ IPCBackend::IPCBackend(MPI_Comm comm)
 
   ROCSHMEM_HOST_CTX_DEFAULT.ctx_opaque = default_host_ctx.get();
 
-  init_g_ret(&heap, thread_comm, MAX_NUM_BLOCKS, &bp->g_ret);
-
-  allocate_atomic_region(&bp->atomic_ret, MAX_NUM_BLOCKS);
-
   setup_team_world();
 
   init_wrk_sync_buffer();
@@ -107,17 +99,6 @@ IPCBackend::IPCBackend(MPI_Comm comm)
 }
 
 IPCBackend::~IPCBackend() {
-  /*
-   * Validate that a handle was passed that is not a nullptr.
-   */
-  auto *bp{ipc_backend_proxy.get()};
-  assert(bp);
-
-  /*
-   * Free the atomic_ret array.
-   */
-  CHECK_HIP(hipFree(bp->atomic_ret->atomic_base_ptr));
-
   /**
    * Destroy teams infrastructure
    * and team world

src/ipc/backend_ipc.hpp

@@ -27,7 +27,6 @@
 #include "../containers/free_list_impl.hpp"
 #include "../hdp_proxy.hpp"
 #include "../memory/hip_allocator.hpp"
-#include "ipc_backend_proxy.hpp"
 #include "../context_incl.hpp"
 #include "ipc_context_proxy.hpp"
 #include "../ipc_policy.hpp"
@@ -117,11 +116,6 @@ class IPCBackend : public Backend {
    */
   char** get_wrk_sync_bases() { return Wrk_Sync_buffer_bases_; }
 
-  /**
-   * @brief Handle to device memory fields.
-   */
-  IPCBackendProxyT ipc_backend_proxy{};
-
   /**
    * @brief The host-facing interface that will be used
    * by all contexts of the IPCBackend

src/ipc/context_ipc_device.cpp

@@ -42,11 +42,7 @@ __host__ IPCContext::IPCContext(Backend *b)
   ipcImpl_.ipc_bases = b->ipcImpl.ipc_bases;
   ipcImpl_.shm_size = b->ipcImpl.shm_size;
 
-  auto *bp{backend->ipc_backend_proxy.get()};
-
   barrier_sync = backend->barrier_sync;
-  g_ret = bp->g_ret;
-  atomic_base_ptr = bp->atomic_ret->atomic_base_ptr;
   fence_pool = backend->fence_pool;
   Wrk_Sync_buffer_bases_ = backend->get_wrk_sync_bases();
 

src/ipc/context_ipc_device.hpp

@@ -212,10 +212,6 @@ class IPCContext : public Context {
   //context class has IpcImpl object (ipcImpl_)
   IpcImpl *ipcImpl{nullptr};
 
-  uint64_t* atomic_base_ptr{nullptr};
-
-  char* g_ret;
-
   //internal functions used by collective operations
   template <typename T>
   __device__ void internal_broadcast(T *dest, const T *source, int nelems, int pe_root,

src/ipc/context_ipc_tmpl_device.hpp

@@ -55,6 +55,7 @@ __device__ void IPCContext::put_nbi(T *dest, const T *source, size_t nelems,
 template <typename T>
 __device__ T IPCContext::g(const T *source, int pe) {
   T ret;
+  getmem(&ret, source, sizeof(T), pe);
   return ret;
 }
 

src/ipc/ipc_backend_proxy.hpp

@@ -1,69 +0,0 @@
-/******************************************************************************
- * Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to
- * deal in the Software without restriction, including without limitation the
- * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
- * sell copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
- * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
- * IN THE SOFTWARE.
- *****************************************************************************/
-
-#ifndef LIBRARY_SRC_IPC_BACKEND_PROXY_HPP_
-#define LIBRARY_SRC_IPC_BACKEND_PROXY_HPP_
-
-#include "../device_proxy.hpp"
-#include "../atomic_return.hpp"
-
-namespace rocshmem {
-
-struct IPCBackendRegister {
-  char *g_ret{nullptr};
-  atomic_ret_t *atomic_ret{nullptr};
-  SymmetricHeap *heap_ptr{nullptr};
-};
-
-template <typename ALLOCATOR>
-class IPCBackendProxy {
-  using ProxyT = DeviceProxy<ALLOCATOR, IPCBackendRegister>;
-
- public:
-  /*
-   * Placement new the memory which is allocated by proxy_
-   */
-  IPCBackendProxy() { new (proxy_.get()) IPCBackendRegister(); }
-
-  /*
-   * Since placement new is called in the constructor, then
-   * delete must be called manually.
-   */
-  ~IPCBackendProxy() { proxy_.get()->~IPCBackendRegister(); }
-
-  /*
-   * @brief Provide access to the memory referenced by the proxy
-   */
-  __host__ __device__ IPCBackendRegister *get() { return proxy_.get(); }
-
- private:
-  /*
-   * @brief Memory managed by the lifetime of this object
-   */
-  ProxyT proxy_{};
-};
-
-using IPCBackendProxyT = IPCBackendProxy<HIPHostAllocator>;
-
-}  // namespace rocshmem
-
-#endif  // #define LIBRARY_SRC_IPC_BACKEND_PROXY_HPP_