Transfer files from RAD repository

src/CMakeLists.txt

@@ -0,0 +1,69 @@
+###############################################################################
+# 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.
+###############################################################################
+
+###############################################################################
+# ADD ROCSHMEM TARGET FOR FILES IN CURRENT DIRECTORY
+###############################################################################
+target_sources(
+  ${PROJECT_NAME}
+  PRIVATE
+    atomic_return.cpp
+    backend_bc.cpp
+    context_host.cpp
+    context_device.cpp
+    ipc_policy.cpp
+    mpi_init_singleton.cpp
+    roc_shmem_gpu.cpp
+    roc_shmem.cpp
+    team.cpp
+    team_tracker.cpp
+    util.cpp
+    wf_coal_policy.cpp
+)
+
+target_compile_options(
+  ${PROJECT_NAME}
+  PUBLIC
+    -fgpu-rdc
+#   xnack allows address translation fault recovery
+#   required option for managed heap configs
+#    -mxnack
+)
+
+#target_link_options(
+  #${PROJECT_NAME}
+  #PUBLIC
+    #--hip-link
+#)
+
+###############################################################################
+# ROCSHMEM TARGET FOR BACKENDS
+###############################################################################
+IF (USE_GPU_IB)
+add_subdirectory(gpu_ib)
+ELSE()
+add_subdirectory(reverse_offload)
+ENDIF()
+add_subdirectory(containers)
+add_subdirectory(host)
+add_subdirectory(memory)
+add_subdirectory(sync)

src/atomic_return.cpp

@@ -0,0 +1,84 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "atomic_return.hpp"
+
+#include <cassert>
+
+#include "constants.hpp"
+
+namespace rocshmem {
+
+void allocate_atomic_region(atomic_ret_t** atomic_ret, int num_wg) {
+  atomic_ret_t* tmp_ret{nullptr};
+  /*
+   * Allocate device-side control struct for the atomic return region.
+   */
+  CHECK_HIP(
+      hipMalloc(reinterpret_cast<void**>(&tmp_ret), sizeof(atomic_ret_t)));
+
+  /*
+   * Allocate fine-grained device-side memory for the atomic return
+   * region.
+   */
+  size_t size_bytes{max_nb_atomic * num_wg * sizeof(uint64_t)};
+#ifdef HIP_SUPPORTS_MALLOC_UNCACHED
+  CHECK_HIP(
+      hipExtMallocWithFlags(reinterpret_cast<void**>(&tmp_ret->atomic_base_ptr),
+                            size_bytes, hipDeviceMallocUncached));
+#else
+  CHECK_HIP(
+      hipExtMallocWithFlags(reinterpret_cast<void**>(&tmp_ret->atomic_base_ptr),
+                            size_bytes, hipDeviceMallocFinegrained));
+#endif
+
+  /*
+   * Zero-initialize the entire atomic return region.
+   */
+  memset(tmp_ret->atomic_base_ptr, 0, size_bytes);
+
+  *atomic_ret = tmp_ret;
+}
+
+void init_g_ret(SymmetricHeap* heap_handle, MPI_Comm thread_comm, int num_wg,
+                char** g_ret) {
+  /*
+   * Create space on the symmetric heap
+   */
+  void* ptr{nullptr};
+  size_t size_bytes{sizeof(int64_t) * MAX_WG_SIZE * num_wg};
+  heap_handle->malloc(&ptr, size_bytes);
+  assert(ptr);
+
+  /*
+   * Assign g_ret the output of the malloc
+   */
+  *g_ret = reinterpret_cast<char*>(ptr);
+
+  /*
+   * Make sure that all processing elements have done this before
+   * continuing.
+   */
+  MPI_Barrier(thread_comm);
+}
+
+}  // namespace rocshmem

src/atomic_return.hpp

@@ -0,0 +1,49 @@
+/******************************************************************************
+ * 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_ATOMIC_RETURN_HPP_
+#define LIBRARY_SRC_ATOMIC_RETURN_HPP_
+
+#include <hip/hip_runtime.h>
+#include <mpi.h>
+
+#include "memory/symmetric_heap.hpp"
+#include "util.hpp"
+
+namespace rocshmem {
+
+const int max_nb_atomic = 4096;
+
+struct atomic_ret_t {
+  uint64_t* atomic_base_ptr;
+  uint32_t atomic_lkey;
+  uint64_t atomic_counter;
+};
+
+void allocate_atomic_region(atomic_ret_t** atomic_ret, int num_wg);
+
+void init_g_ret(SymmetricHeap* heap_handle, MPI_Comm thread_comm, int num_wg,
+                char** g_ret);
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_ATOMIC_RETURN_HPP_

src/backend_bc.cpp

@@ -0,0 +1,219 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "backend_bc.hpp"
+
+#include "backend_type.hpp"
+#include "context_incl.hpp"
+
+#ifndef USE_GPU_IB
+#include "reverse_offload/backend_ro.hpp"
+#else
+#include "gpu_ib/backend_ib.hpp"
+#endif
+
+namespace rocshmem {
+
+Backend::Backend() {
+  int num_cus{};
+  if (hipDeviceGetAttribute(&num_cus, hipDeviceAttributeMultiprocessorCount,
+                            0)) {
+    abort();
+  }
+
+  CHECK_HIP(hipGetDevice(&hip_dev_id));
+
+  /*
+   * Initialize 'print_lock' global and copy to the device memory space.
+   */
+  CHECK_HIP(hipMalloc(&print_lock, sizeof(*print_lock)));
+  *print_lock = 0;
+
+  int* print_lock_addr{nullptr};
+  CHECK_HIP(hipGetSymbolAddress(reinterpret_cast<void**>(&print_lock_addr),
+                                HIP_SYMBOL(print_lock)));
+
+  CHECK_HIP(hipMemcpy(print_lock_addr, &print_lock, sizeof(print_lock),
+                      hipMemcpyDefault));
+
+  /*
+   * Copy this Backend object to 'backend_device_proxy' global in the
+   * device memory space to provide a device-side handle to Backend.
+   */
+  int* device_backend_proxy_addr{nullptr};
+  CHECK_HIP(
+      hipGetSymbolAddress(reinterpret_cast<void**>(&device_backend_proxy_addr),
+                          HIP_SYMBOL(device_backend_proxy)));
+
+  Backend* this_temp_addr{this};
+  CHECK_HIP(hipMemcpy(device_backend_proxy_addr, &this_temp_addr, sizeof(this),
+                      hipMemcpyDefault));
+
+  CHECK_HIP(
+      hipHostMalloc(reinterpret_cast<void**>(&done_init), sizeof(uint8_t)));
+
+  /*
+   * Notify other threads that Backend has been initialized.
+   */
+  *done_init = 0;
+}
+
+void Backend::track_ctx(Context* ctx) {
+  /**
+   * TODO: Don't track CTX_PRIVATE when we support it
+   * since destroying CTX_PRIVATE is the user's
+   * responsibility.
+   */
+  list_of_ctxs.push_back(ctx);
+}
+
+void Backend::untrack_ctx(Context* ctx) {
+  /* Get an iterator to this ctx in the vector */
+  std::vector<Context*>::iterator it =
+      std::find(list_of_ctxs.begin(), list_of_ctxs.end(), ctx);
+  assert(it != list_of_ctxs.end());
+
+  /* Remove the ctx from the vector */
+  list_of_ctxs.erase(it);
+}
+
+void Backend::destroy_remaining_ctxs() {
+  while (!list_of_ctxs.empty()) {
+    ctx_destroy(list_of_ctxs.back());
+    list_of_ctxs.pop_back();
+  }
+}
+
+Backend::~Backend() { CHECK_HIP(hipFree(print_lock)); }
+
+void Backend::dump_stats() {
+  printf("PE %d\n", my_pe);
+
+  const auto& device_stats{globalStats};
+  printf("DEVICE STATS\n");
+  printf("Puts (Blocking/P/Nbi) %llu/%llu/%llu\n",
+         device_stats.getStat(NUM_PUT), device_stats.getStat(NUM_P),
+         device_stats.getStat(NUM_PUT_NBI));
+  printf("WG_Puts (Blocking/Nbi) %llu/%llu\n", device_stats.getStat(NUM_PUT_WG),
+         device_stats.getStat(NUM_PUT_NBI_WG));
+  printf("WAVE_Puts (Blocking/Nbi) %llu/%llu\n",
+         device_stats.getStat(NUM_PUT_WAVE),
+         device_stats.getStat(NUM_PUT_NBI_WAVE));
+  printf("Gets (Blocking/G/Nbi) %llu/%llu/%llu\n",
+         device_stats.getStat(NUM_GET), device_stats.getStat(NUM_G),
+         device_stats.getStat(NUM_GET_NBI));
+  printf("WG_Gets (Blocking/Nbi) %llu/%llu\n", device_stats.getStat(NUM_GET_WG),
+         device_stats.getStat(NUM_GET_NBI_WG));
+  printf("WAVE_Gets (Blocking/Nbi) %llu/%llu\n",
+         device_stats.getStat(NUM_GET_WAVE),
+         device_stats.getStat(NUM_GET_NBI_WAVE));
+  printf("Fences %llu\n", device_stats.getStat(NUM_FENCE));
+  printf("Quiets %llu\n", device_stats.getStat(NUM_QUIET));
+  printf("ToAll %llu\n", device_stats.getStat(NUM_TO_ALL));
+  printf("BarrierAll %llu\n", device_stats.getStat(NUM_BARRIER_ALL));
+  printf("Wait Until %llu\n", device_stats.getStat(NUM_WAIT_UNTIL));
+  printf("Wait Until Any %llu\n", device_stats.getStat(NUM_WAIT_UNTIL_ANY));
+  printf("Wait Until All %llu\n", device_stats.getStat(NUM_WAIT_UNTIL_ALL));
+  printf("Wait Until Some %llu\n", device_stats.getStat(NUM_WAIT_UNTIL_SOME));
+  printf("Wait Until All Vector %llu\n",
+         device_stats.getStat(NUM_WAIT_UNTIL_ALL_VECTOR));
+  printf("Wait Until Any Vector %llu\n",
+         device_stats.getStat(NUM_WAIT_UNTIL_ANY_VECTOR));
+  printf("Wait Until Some Vector %llu\n",
+         device_stats.getStat(NUM_WAIT_UNTIL_SOME_VECTOR));
+  printf("Finalizes %llu\n", device_stats.getStat(NUM_FINALIZE));
+  printf("Coalesced %llu\n", device_stats.getStat(NUM_MSG_COAL));
+  printf("Atomic_FAdd %llu\n", device_stats.getStat(NUM_ATOMIC_FADD));
+  printf("Atomic_FCswap %llu\n", device_stats.getStat(NUM_ATOMIC_FCSWAP));
+  printf("Atomic_FInc %llu\n", device_stats.getStat(NUM_ATOMIC_FINC));
+  printf("Atomic_Fetch %llu\n", device_stats.getStat(NUM_ATOMIC_FETCH));
+  printf("Atomic_Add %llu\n", device_stats.getStat(NUM_ATOMIC_ADD));
+  printf("Atomic_Set %llu\n", device_stats.getStat(NUM_ATOMIC_SET));
+  printf("Atomic_Cswap %llu\n", device_stats.getStat(NUM_ATOMIC_CSWAP));
+  printf("Atomic_Inc %llu\n", device_stats.getStat(NUM_ATOMIC_INC));
+  printf("Tests %llu\n", device_stats.getStat(NUM_TEST));
+  printf("SHMEM_PTR %llu\n", device_stats.getStat(NUM_SHMEM_PTR));
+  printf("SyncAll %llu\n", device_stats.getStat(NUM_SYNC_ALL));
+
+  const auto& host_stats{globalHostStats};
+  printf("HOST STATS\n");
+  printf("Puts (Blocking/P/Nbi) %llu/%llu/%llu\n",
+         host_stats.getStat(NUM_HOST_PUT), host_stats.getStat(NUM_HOST_P),
+         host_stats.getStat(NUM_HOST_PUT_NBI));
+  printf("Gets (Blocking/G/Nbi) (%llu/%llu/%llu)\n",
+         host_stats.getStat(NUM_HOST_GET), host_stats.getStat(NUM_HOST_G),
+         host_stats.getStat(NUM_HOST_GET_NBI));
+  printf("Fences %llu\n", host_stats.getStat(NUM_HOST_FENCE));
+  printf("Quiets %llu\n", host_stats.getStat(NUM_HOST_QUIET));
+  printf("ToAll %llu\n", host_stats.getStat(NUM_HOST_TO_ALL));
+  printf("BarrierAll %llu\n", host_stats.getStat(NUM_HOST_BARRIER_ALL));
+  printf("Wait Until %llu\n", host_stats.getStat(NUM_HOST_WAIT_UNTIL));
+  printf("Wait Until Any %llu\n", host_stats.getStat(NUM_HOST_WAIT_UNTIL_ANY));
+  printf("Wait Until All %llu\n", host_stats.getStat(NUM_HOST_WAIT_UNTIL_ALL));
+  printf("Wait Until Some %llu\n",
+         host_stats.getStat(NUM_HOST_WAIT_UNTIL_SOME));
+  printf("Wait Until All Vector %llu\n",
+         host_stats.getStat(NUM_HOST_WAIT_UNTIL_ALL_VECTOR));
+  printf("Wait Until Any Vector %llu\n",
+         host_stats.getStat(NUM_HOST_WAIT_UNTIL_ANY_VECTOR));
+  printf("Wait Until Some Vector %llu\n",
+         host_stats.getStat(NUM_HOST_WAIT_UNTIL_SOME_VECTOR));
+  printf("Finalizes %llu\n", host_stats.getStat(NUM_HOST_FINALIZE));
+  printf("Atomic_FAdd %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_FADD));
+  printf("Atomic_FCswap %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_FCSWAP));
+  printf("Atomic_FInc %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_FINC));
+  printf("Atomic_Fetch %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_FETCH));
+  printf("Atomic_Add %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_ADD));
+  printf("Atomic_Set %llu\n", host_stats.getStat(NUM_ATOMIC_SET));
+  printf("Atomic_Cswap %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_CSWAP));
+  printf("Atomic_Inc %llu\n", host_stats.getStat(NUM_HOST_ATOMIC_INC));
+  printf("Tests %llu\n", host_stats.getStat(NUM_HOST_TEST));
+  printf("SHMEM_PTR %llu\n", host_stats.getStat(NUM_HOST_SHMEM_PTR));
+  printf("SyncAll %llu\n", host_stats.getStat(NUM_HOST_SYNC_ALL));
+
+  dump_backend_stats();
+}
+
+void Backend::reset_stats() {
+  globalStats.resetStats();
+  globalHostStats.resetStats();
+
+  reset_backend_stats();
+}
+
+__device__ bool Backend::create_ctx(int64_t option, roc_shmem_ctx_t* ctx) {
+#ifndef USE_GPU_IB
+  return static_cast<ROBackend*>(this)->create_ctx(option, ctx);
+#else
+  return static_cast<GPUIBBackend*>(this)->create_ctx(option, ctx);
+#endif
+}
+
+__device__ void Backend::destroy_ctx(roc_shmem_ctx_t* ctx) {
+#ifndef USE_GPU_IB
+  static_cast<ROBackend*>(this)->destroy_ctx(ctx);
+#else
+  static_cast<GPUIBBackend*>(this)->destroy_ctx(ctx);
+#endif
+}
+
+}  // namespace rocshmem

src/backend_bc.hpp

@@ -0,0 +1,314 @@
+/******************************************************************************
+ * 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_BACKEND_BC_HPP_
+#define LIBRARY_SRC_BACKEND_BC_HPP_
+
+/**
+ * @file backend_bc.hpp
+ * Defines the Backend base class.
+ *
+ * The backend base class sets up most of the host-side library resources.
+ * It is the top-level interface for these resources.
+ */
+
+#include <mpi.h>
+
+#include <vector>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "roc_shmem/roc_shmem.hpp"
+#include "backend_type.hpp"
+#include "ipc_policy.hpp"
+#include "memory/slab_heap.hpp"
+#include "memory/symmetric_heap.hpp"
+#include "stats.hpp"
+#include "team_tracker.hpp"
+
+namespace rocshmem {
+
+class Team;
+class TeamInfo;
+
+/**
+ * @class Backend backend.hpp
+ * @brief Container class for the persistent state used by the library.
+ *
+ * Backend is populated by host-side initialization and allocation calls.
+ * It uses this state to populate Context objects which the GPU may use to
+ * perform networking operations.
+ *
+ * The roc_shmem.cpp implementation file wraps many the Backend public
+ * members to implement the library's public API.
+ */
+class Backend {
+ public:
+  /**
+   * @brief Constructor.
+   *
+   * @note Implementation may reduce the number of workgroups if the
+   * number exceeds hardware limits.
+   */
+  explicit Backend();
+
+  /**
+   * @brief Destructor.
+   */
+  virtual ~Backend();
+
+  __device__ bool create_ctx(int64_t option, roc_shmem_ctx_t* ctx);
+  __device__ void destroy_ctx(roc_shmem_ctx_t* ctx);
+
+  /**
+   * @brief Create a new team object and initialize it.
+   *
+   * @param[in] parent_team Pointer to the parrent team object.
+   * @param[in] team_info_wrt_parent TeamInfo object wrt parent team.
+   * @param[in] team_info_wrt_world TeamInfo object wrt TEAM_WORLD.
+   * @param[in] num_pes Number of PEs in this team.
+   * @param[in] my_pe_in_new_team Index of this PE in the new team.
+   * @param[in] team_comm MPI communicator for this team.
+   *
+   * @param[out] new_team pointer to the new team.
+   */
+  virtual void create_new_team(Team* parent_team,
+                               TeamInfo* team_info_wrt_parent,
+                               TeamInfo* team_info_wrt_world, int num_pes,
+                               int my_pe_in_new_team, MPI_Comm team_comm,
+                               roc_shmem_team_t* new_team) = 0;
+
+  /**
+   * @brief Destruct a team
+   *
+   * @param[in] team Handle to the team to destroy.
+   */
+  virtual void team_destroy(roc_shmem_team_t team) = 0;
+
+  /**
+   * @brief Reports processing element number id.
+   *
+   * @return Unique numeric identifier for each processing element.
+   */
+  __host__ __device__ int getMyPE() const { return my_pe; }
+
+  /**
+   * @brief Reports number of processing elements.
+   *
+   * @return Number of active processing elements tracked by library.
+   */
+  __host__ __device__ int getNumPEs() const { return num_pes; }
+
+  /**
+   * @brief Allocates and initializes device-side library state.
+   *
+   * Preallocates a single private context for this workgroup (thread-block)
+   * and binds it to the WGState instance.
+   *
+   * The code below carves the allocation out of the "extern __shared__"
+   * partition and then builds the context object in that memory.
+   *
+   * @return void
+   */
+  __device__ void create_wg_state();
+
+  /**
+   * @brief Frees device-side library resources.
+   *
+   * @return void
+   */
+  __device__ void finalize_wg_state();
+
+  /**
+   * @brief Dumps statistics for public API invocations.
+   *
+   * @note Implementation may dump additional statistics from backend
+   * derived classes when calling this function. If so, the method,
+   * dump_backend_stats, will be used as the interface for the
+   * additional statistics.
+   */
+  void dump_stats();
+
+  /**
+   * @brief Resets statistics for public API invocations.
+   *
+   * @note Implementation may reset additional statistics from backend
+   * derived classes when calling this function. If so, the method,
+   * reset_backend_stats, will be used as the interface for the
+   * additional statistics.
+   */
+  void reset_stats();
+
+  /**
+   * @brief Abort the application.
+   *
+   * @param[in] status Exit code.
+   *
+   * @return void.
+   *
+   * @note This routine terminates the entire application.
+   */
+  virtual void global_exit(int status) = 0;
+
+  /**
+   * @brief Creates a new OpenSHMEM context.
+   *
+   * @param[in] options Options for context creation
+   * @param[in] ctx     Address of the pointer to the new context
+   *
+   * @return Zero on success, nonzero otherwise.
+   */
+  virtual void ctx_create(int64_t options, void** ctx) = 0;
+
+  /**
+   * @brief Destroys a context.
+   *
+   * @param[in] ctx Context handle.
+   *
+   * @return void.
+   */
+  virtual void ctx_destroy(Context* ctx) = 0;
+
+  /**
+   * @brief High level device stats that do not depend on backend type.
+   */
+  ROCStats globalStats{};
+
+  /**
+   * @brief High level host stats that do not depend on backend type.
+   */
+  ROCHostStats globalHostStats{};
+
+  /**
+   * @brief Number of processing elements running in job.
+   *
+   * @todo Change to size_t.
+   */
+  int num_pes{0};
+
+  /**
+   * @brief Unique numeric identifier ranging from 0 (inclusive) to
+   * num_pes (exclusive) [0 ... num_pes).
+   *
+   * @todo Change to size_t and set invalid entry to max size.
+   */
+  int my_pe{-1};
+
+  /**
+   * @brief indicate when init is done on the CPU. Non-blocking init is only
+   * available with GPU-IB
+   */
+  uint8_t* done_init{nullptr};
+
+  /**
+   * @todo document where this is used and try to coalesce this into another
+   * class
+   */
+  MPI_Comm thread_comm{};
+
+  /**
+   * @brief Object contains the interface and internal data structures
+   * needed to allocate/free memory on the symmetric heap.
+   */
+  SymmetricHeap heap{};
+
+  /**
+   * @brief Object contains the interface and internal data structures
+   * needed to allocate/free memory on the slab heap.
+   */
+  SlabHeap slab{};
+
+  /**
+   * @brief Determines which device to launch device kernels onto.
+   *
+   * Multi-device nodes can specify which one they would like to use.
+   */
+  int hip_dev_id{0};
+
+  /**
+   * @brief Add ctx from the list of user-created ctxs
+   */
+  void track_ctx(Context* ctx);
+
+  /**
+   * @brief Remove ctx from the list of user-created ctxs
+   */
+  void untrack_ctx(Context* ctx);
+
+  /**
+   * @brief Remove all ctxs from the list of user-created ctxs
+   */
+  void destroy_remaining_ctxs();
+
+  /**
+   * @brief Compile-time configuration policy for intra-node shared memory
+   * accesses.
+   *
+   * The configuration option "USE_IPC" can be enabled to allow shared
+   * memory accesses to the symmetric heap from processing elements
+   * co-located on the same node.
+   */
+  IpcImpl ipcImpl{};
+
+  /**
+   * @brief Maintains information about teams
+   */
+  TeamTracker team_tracker{};
+
+ protected:
+  /**
+   * @brief Required to support static inheritance for device calls.
+   *
+   * The Context DISPATCH implementation requires this member.
+   * The implementation needs to know the derived class type to
+   * issue a static_cast.
+   *
+   * GPU devices do not support virtual functions. Therefore, we cannot
+   * rely on the normal inheritance mechanism to tailor behavior for
+   * derived backend types.
+   */
+  BackendType type{BackendType::GPU_IB_BACKEND};
+
+  /**
+   * @brief Dumps derived class statistics.
+   */
+  virtual void dump_backend_stats() = 0;
+
+  /**
+   * @brief Resets derived class statistics.
+   */
+  virtual void reset_backend_stats() = 0;
+
+ private:
+  /**
+   * @brief List of ctxs created by the user.
+   */
+  std::vector<Context*> list_of_ctxs{};
+};
+
+/**
+ * @brief Global handle used by the device to access the backend.
+ */
+extern __constant__ Backend* device_backend_proxy;
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_BACKEND_BC_HPP_

src/backend_type.hpp

@@ -0,0 +1,132 @@
+/******************************************************************************
+ * 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_BACKEND_TYPE_HPP_
+#define LIBRARY_SRC_BACKEND_TYPE_HPP_
+
+/**
+ * @file backend_type.hpp
+ * Defines the Backend derived class types and contains the DISPATCH macros.
+ *
+ * The type information is required to be known at compile time because
+ * we use static dispatch to produce compile time polymorphism.
+ *
+ * The device cannot use runtime polymorphism because calls through virtual
+ * functions are not supported at this time.
+ */
+
+#include "config.h"  // NOLINT(build/include_subdir)
+
+namespace rocshmem {
+
+/**
+ * @brief Enumerates the Backend derived classes.
+ *
+ * @note Derived classes which use Backend as a base class must add
+ * themselves to this enum class to support static polymorphism.
+ */
+enum class BackendType { RO_BACKEND, GPU_IB_BACKEND };
+
+/**
+ * @brief Helper macro for some dispatch calls
+ */
+#define PAIR(A, B) A, B
+
+/**
+ * @brief Device static dispatch method call.
+ */
+#ifndef USE_GPU_IB
+#define DISPATCH(Func)                  \
+  static_cast<ROContext *>(this)->Func;
+#else
+#define DISPATCH(Func)                     \
+  dev_mtx_.lock();                         \
+  static_cast<GPUIBContext *>(this)->Func; \
+  dev_mtx_.unlock();
+#endif
+
+/**
+ * @brief Device static dispatch method call with a return value.
+ */
+#ifndef USE_GPU_IB
+#define DISPATCH_RET(Func)                        \
+  auto ret_val{0};                                \
+  ret_val = static_cast<ROContext *>(this)->Func; \
+  return ret_val;
+#else
+#define DISPATCH_RET(Func)                           \
+  dev_mtx_.lock();                                   \
+  auto ret_val{0};                                   \
+  ret_val = static_cast<GPUIBContext *>(this)->Func; \
+  dev_mtx_.unlock();                                 \
+  return ret_val;
+#endif
+/**
+ * @brief Device static dispatch method call with a return type of pointer.
+ */
+#ifndef USE_GPU_IB
+#define DISPATCH_RET_PTR(Func)                    \
+  void *ret_val{nullptr};                         \
+  ret_val = static_cast<ROContext *>(this)->Func; \
+  return ret_val;
+#else
+#define DISPATCH_RET_PTR(Func)                       \
+  void *ret_val{nullptr};                            \
+  ret_val = static_cast<GPUIBContext *>(this)->Func; \
+  return ret_val;
+#endif
+
+/**
+ * @brief Host static dispatch method call.
+ *
+ * @note There is no need to lock-unlock on host since we are using
+ * MPI_THREAD_MULTIPLE (for RMA and AMO operations) and the ordering and
+ * threading semantics of collectives in OpenSHMEM match those of MPI.
+ */
+#ifndef USE_GPU_IB
+#define HOST_DISPATCH(Func) static_cast<ROHostContext *>(this)->Func;
+#else
+#define HOST_DISPATCH(Func) static_cast<GPUIBHostContext *>(this)->Func;
+#endif
+/**
+ * @brief Host static dispatch method call with return value.
+ *
+ * @note There is no need to lock-unlock on host since we are using
+ * MPI_THREAD_MULTIPLE (for RMA and AMO operations) and the ordering and
+ * threading semantics of collectives in OpenSHMEM match those of MPI.
+ */
+
+#ifndef USE_GPU_IB
+#define HOST_DISPATCH_RET(Func)                       \
+  auto ret_val{0};                                    \
+  ret_val = static_cast<ROHostContext *>(this)->Func; \
+  return ret_val;
+#else
+#define HOST_DISPATCH_RET(Func)                          \
+  auto ret_val{0};                                       \
+  ret_val = static_cast<GPUIBHostContext *>(this)->Func; \
+  return ret_val;
+#endif
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_BACKEND_TYPE_HPP_

src/constants.hpp

@@ -0,0 +1,62 @@
+/******************************************************************************
+ * 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_CONSTANTS_HPP_
+#define LIBRARY_SRC_CONSTANTS_HPP_
+
+/**
+ * @file constants.hpp
+ *
+ * @brief Contains global constants for ROCSHMEM library
+ */
+
+namespace rocshmem {
+
+/**
+ * @brief Minimum object alignment for symmetric heap.
+ *
+ * @note Cache line size on most systems is either 64 or 128.
+ */
+inline const unsigned ALIGNMENT{128};
+
+/**
+ * @brief Constant number which holds maximum workgroup size.
+ *
+ * @todo Remove this member from this class. It belongs in a class
+ * that specifically holds device hardware information. If this
+ * device class existed, we could consolidate the various flavours of
+ * the Instinct cards into their own groups and then set these
+ * hard-coded fields by querying the rocm runtime during our library
+ * initialization.
+ */
+inline const unsigned MAX_WG_SIZE{1024};
+
+/**
+ * @brief Constant number which holds the wavefront size
+ *
+ * @note Wavefront size on most systems is either 32 or 64.
+ */
+inline const unsigned WF_SIZE{64};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONSTANTS_HPP_

src/containers/CMakeLists.txt

@@ -0,0 +1,31 @@
+###############################################################################
+# 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.
+###############################################################################
+
+###############################################################################
+# ADD ROCSHMEM TARGET FOR FILES IN CURRENT DIRECTORY
+###############################################################################
+target_sources(
+  ${PROJECT_NAME}
+  PRIVATE
+    share_strategy.cpp
+    strategies.cpp
+)

src/containers/array.hpp

@@ -0,0 +1,199 @@
+/******************************************************************************
+ * 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_CONTAINERS_ARRAY_HPP_
+#define LIBRARY_SRC_CONTAINERS_ARRAY_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include "index_strategy.hpp"
+#include "strategies.hpp"
+#include "memory_allocator.hpp"
+
+namespace rocshmem {
+
+template <typename TYPE>
+class Array {
+ public:
+  /**
+   * @brief
+   *
+   * @param[in] index
+   *
+   * @return
+   */
+  __host__ __device__ TYPE& operator[](size_t index);
+
+  /**
+   * @brief
+   *
+   * @param[in] index
+   *
+   * @return
+   */
+  __host__ __device__ const TYPE& operator[](size_t index) const;
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __host__ __device__ size_t size() const;
+
+  /**
+   * @brief
+   *
+   * @param[in] v
+   * @param[in] start_index
+   * @param[in] length
+   *
+   * @return void
+   */
+  __host__ void fill(TYPE v, size_t start_index, size_t length);
+
+  /**
+   * @brief
+   *
+   * @param[in] v
+   * @param[in] start_index
+   * @param[in] length
+   *
+   * @return void
+   */
+  __device__ void fill(TYPE v, size_t start_index, size_t length);
+
+  /**
+   * @brief
+   *
+   * @return void
+   */
+  __host__ __device__ void zero();
+
+  /**
+   * @brief
+   *
+   * @param[in] other
+   * @param[in] start_index
+   * @param[in] length
+   *
+   * @return void
+   */
+  __host__ void copy(const Array* other, size_t start_index,
+                     size_t length);  // NOLINT
+
+  /**
+   * @brief
+   *
+   * @param[in] other
+   * @param[in] start_index
+   * @param[in] length
+   *
+   * @return void
+   */
+  __device__ void copy(const Array* other, size_t start_index,  // NOLINT
+                       size_t length);
+
+  /**
+   * @brief
+   */
+  Array() = default;
+
+  /**
+   * @brief
+   *
+   * @param[in] array_size
+   * @param[in] allocator
+   */
+  Array(size_t array_size, MemoryAllocator allocator,
+        const ObjectStrategy* strategy);
+
+  /**
+   * @brief
+   */
+  ~Array();
+
+  /**
+   * @brief
+   *
+   * @param[in] other
+   */
+  Array(const Array& other);
+
+  /**
+   * @brief
+   *
+   * @param[in] rhs
+   *
+   * @return
+   */
+  Array& operator=(const Array& rhs);
+
+  /**
+   * @brief
+   */
+  __device__ void zero_thread_dump();
+
+  /**
+   * @brief
+   */
+  __device__ void any_thread_dump();
+
+ private:
+  /**
+   * @brief
+   */
+  __device__ void _dump();
+
+  /**
+   * @brief
+   *
+   * @note _allocator is required to be declared before '_array' in this
+   * class because _allocator deallocates _array in class destructor.
+   */
+  MemoryAllocator _allocator{};
+
+ protected:
+  /**
+   * @brief
+   */
+  TYPE* _array{nullptr};
+
+  /**
+   * @brief
+   */
+  size_t _size{0};
+
+ private:
+  /**
+   * @brief
+   */
+  IndexStrategy _dev_idx{};
+
+  /**
+   * @brief
+   */
+  const size_t dump_num_data_per_line{8};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_ARRAY_HPP_

src/containers/array_impl.hpp

@@ -0,0 +1,186 @@
+/******************************************************************************
+ * 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_CONTAINERS_ARRAY_IMPL_HPP_
+#define LIBRARY_SRC_CONTAINERS_ARRAY_IMPL_HPP_
+
+#include <cassert>
+
+#include "array.hpp"
+
+namespace rocshmem {
+
+extern __constant__ int* GLOBAL_DEVICE_PRINT_LOCK;
+
+template <typename TYPE>
+Array<TYPE>::Array(size_t array_size, MemoryAllocator allocator,
+                   const ObjectStrategy* strategy)
+    : _allocator(allocator),
+      _size(array_size),
+      _dev_idx(strategy->index_strategy_two) {
+  _allocator.allocate(reinterpret_cast<void**>(&_array),
+                      array_size * sizeof(TYPE));
+}
+
+template <typename TYPE>
+Array<TYPE>::~Array() {
+  if (_array) {
+    _allocator.deallocate(_array);
+  }
+}
+
+template <typename TYPE>
+Array<TYPE>::Array(const Array<TYPE>& other) {
+  _size = other._size;
+
+  _allocator.allocate(reinterpret_cast<void**>(&_array), _size * sizeof(TYPE));
+
+  memcpy(_array, other._array, _size * sizeof(TYPE));
+}
+
+template <typename TYPE>
+Array<TYPE>& Array<TYPE>::operator=(const Array<TYPE>& rhs) {
+  if (this == &rhs) {
+    return *this;
+  }
+
+  if (_array) {
+    _allocator.deallocate(_array);
+  }
+
+  _size = rhs._size;
+
+  _allocator = rhs._allocator;
+
+  _allocator.allocate(reinterpret_cast<void**>(&_array), _size * sizeof(TYPE));
+  assert(_array);
+
+  memcpy(_array, rhs._array, _size * sizeof(TYPE));
+
+  return *this;
+}
+
+template <typename TYPE>
+__host__ __device__ TYPE& Array<TYPE>::operator[](size_t index) {
+  assert(index < _size);
+  return _array[index];
+}
+
+template <typename TYPE>
+__host__ __device__ const TYPE& Array<TYPE>::operator[](size_t index) const {
+  assert(index < _size);
+  return _array[index];
+}
+
+template <typename TYPE>
+__host__ __device__ size_t Array<TYPE>::size() const {
+  return _size;
+}
+
+template <typename TYPE>
+__host__ void Array<TYPE>::fill(TYPE v, size_t start_index, size_t length) {
+  for (size_t i = start_index; i < start_index + length; i++) {
+    (*this)[i] = v;
+  }
+}
+
+template <typename TYPE>
+__device__ void Array<TYPE>::fill(TYPE v, size_t start_index, size_t length) {
+  size_t i = start_index + _dev_idx.start();
+  while ((i < start_index + length) && (i < _size)) {
+    (*this)[i] = v;
+    i = _dev_idx.next(i);
+  }
+}
+
+template <typename TYPE>
+__host__ __device__ void Array<TYPE>::zero() {
+  fill({0, 0}, 0, size());
+}
+
+template <typename TYPE>
+__host__ void Array<TYPE>::copy(const Array<TYPE>* other,
+                                size_t start_index,  // NOLINT
+                                size_t length) {
+  assert(other);
+  for (size_t i = start_index; i < start_index + length; i++) {
+    (*this)[i] = (*other)[i];
+  }
+}
+
+template <typename TYPE>
+__device__ void Array<TYPE>::copy(const Array<TYPE>* other,  // NOLINT
+                                  size_t start_index, size_t length) {
+  assert(other);
+  for (size_t i = start_index + _dev_idx.start(); i < start_index + length;
+       i = _dev_idx.next(i)) {
+    (*this)[i] = (*other)[i];
+  }
+}
+
+template <typename TYPE>
+__device__ void Array<TYPE>::zero_thread_dump() {
+  Identity id{};
+  if (id.global_thread_id() == 0) {
+    _dump();
+  }
+}
+
+template <typename TYPE>
+__device__ void Array<TYPE>::any_thread_dump() {
+  Identity id{};
+  for (int i = 0; i < __AMDGCN_WAVEFRONT_SIZE; i++) {
+    if ((id.local_thread_id() % __AMDGCN_WAVEFRONT_SIZE) == i) {
+      while (atomicCAS(GLOBAL_DEVICE_PRINT_LOCK, 0, 1) == 1) {
+      }
+
+      printf("(thread %lu)\n", id.global_thread_id());
+      _dump();
+
+      *GLOBAL_DEVICE_PRINT_LOCK = 0;
+    }
+  }
+}
+
+template <typename TYPE>
+__device__ void Array<TYPE>::_dump() {
+  Thread_Contiguous_Block_Agnostic idx(_size);
+  for (size_t i = idx.start(); i < idx.end(); i = idx.next(i)) {
+    /*
+     * Limit the number of printed elements per line.
+     */
+    if (i % dump_num_data_per_line == 0) {
+      printf("\n");
+      printf("%10lu ", i);
+    }
+
+    /*
+     * Print the data for this index.
+     */
+    dump((*this)[i]);
+  }
+  printf("\n");
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_ARRAY_IMPL_HPP_

src/containers/forward_list.hpp

@@ -0,0 +1,500 @@
+/******************************************************************************
+ * 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_CONTAINERS_FORWARD_LIST_HPP_
+#define LIBRARY_SRC_CONTAINERS_FORWARD_LIST_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include <initializer_list>
+#include <ostream>
+#include <string>
+
+#include "../src/device_proxy.hpp"
+#include "../src/memory/hip_allocator.hpp"
+#include "../sync/abql_block_mutex.hpp"
+
+namespace rocshmem {
+
+/*****************************************************************************
+ ******************************* FORWARD LIST ********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC = HIPAllocator>
+class ForwardList {
+  friend class ForwardListTestFixture;
+
+  struct Node {
+    TYPE data;
+    Node* next;
+  };
+
+  template <bool CONST>
+  class Iterator;
+
+ public:
+  using iterator = Iterator<false>;
+  using const_iterator = Iterator<true>;
+
+  /**
+   * @brief Default constructor
+   */
+  ForwardList() = default;
+
+  explicit ForwardList(const ALLOC& alloc);
+
+  /**
+   * @brief Fill constructor
+   */
+  explicit ForwardList(size_t n, const ALLOC& alloc = ALLOC());
+
+  explicit ForwardList(size_t n, const TYPE& val, const ALLOC& alloc = ALLOC());
+
+  /**
+   * @brief Range constructor
+   */
+  template <class InputIterator>
+  ForwardList(InputIterator first, InputIterator last,
+              const ALLOC& alloc = ALLOC());
+
+  /**
+   * @brief Copy constructor
+   */
+  ForwardList(const ForwardList& fwdlst);
+
+  ForwardList(const ForwardList& fwdlst, const ALLOC& alloc);
+
+  /**
+   * @brief Move constructor
+   */
+  ForwardList(ForwardList&& fwdlst);
+
+  ForwardList(ForwardList&& fwdlst, const ALLOC& alloc);
+
+  /**
+   * @brief Initializer list constructor
+   */
+  ForwardList(std::initializer_list<TYPE> il, const ALLOC& alloc = ALLOC());
+
+  /**
+   * @brief Destructor
+   */
+  ~ForwardList();
+
+  /**
+   * @brief Copy assignment operator
+   */
+  ForwardList& operator=(const ForwardList& fwdlst);
+
+  /**
+   * @brief Move assignment operator
+   */
+  ForwardList& operator=(ForwardList&& fwdlst);
+
+  /**
+   * @brief Initializer list assignment
+   */
+  ForwardList& operator=(std::initializer_list<TYPE> il);
+
+  /**
+   * @brief Returns iterator pointing to position before first element.
+   *
+   * The iterator returned shall not be dereferenced:
+   * It is meant to be used as an argument for member functions
+   * emplace_after, insert_after, erase_after or splice_after, to specify
+   * the beginning of the sequence as the location where the action is
+   * performed.
+   */
+  iterator before_begin() noexcept;
+
+  /**
+   * @brief Returns iterator pointing to first element in ForwardList.
+   *
+   * Notice that, unlike member ForwardList::front, which returns a
+   * reference to the first element, this function returns a forward iterator
+   * pointing to it.
+   *
+   * If the container is empty, the returned iterator value shall not be
+   * dereferenced.
+   */
+  iterator begin() noexcept;
+
+  const_iterator begin() const noexcept;
+
+  /**
+   * @brief Returns an iterator to past-the-end element in ForwardList.
+   */
+  iterator end() noexcept;
+
+  const_iterator end() const noexcept;
+
+  /**
+   * @brief Returns const_iterator pointing to position before first element.
+   */
+  const_iterator cbefore_begin() const noexcept;
+
+  /**
+   * @brief Returns const_iterator pointing to first element.
+   */
+  const_iterator cbegin() const noexcept;
+
+  /**
+   * @brief Returns const_iterator pointing to past-the-end element.
+   */
+  const_iterator cend() const noexcept;
+
+  /**
+   * @brief Returns bool indicating if ForwardList container is empty.
+   */
+  bool empty() const noexcept;
+
+  /**
+   * @brief Returns maximum number of elements that ForwardList can hold.
+   */
+  size_t max_size() const noexcept;
+
+  /**
+   * @brief Returns reference to first element in the ForwardList.
+   */
+  ForwardList& front();
+
+  const ForwardList& front() const;
+
+  /**
+   * @brief Assigns new values, replacing current contents, and modify size.
+   *
+   * In range version, the new contents are elements constructed from each
+   * of the elements in the range between first and last, in the same order.
+   */
+  template <class InputIterator>
+  void assign(InputIterator first, InputIterator last);
+
+  /**
+   * @brief Assigns new values, replacing current contents, and modify size.
+   *
+   * In fill version, the new contents are n elements, each initialized to
+   * a copy of val.
+   */
+  void assign(size_t n, const TYPE& val);
+
+  /**
+   * @brief Assigns new values, replacing current contents, and modify size.
+   *
+   * In initializer list version, the new contents are copies of the values
+   * passed as initializer list, in the same order.
+   */
+  void assign(std::initializer_list<TYPE> il);
+
+  /**
+   * @brief Inserts new element at beginning of the ForwardList.
+   *
+   * The element goes into the container right before its current first
+   * element. This new element is constructed in place using args as the
+   * arguments for its construction.
+   */
+  template <class... Args>
+  void emplace_front(Args&&... args);
+
+  /**
+   * @brief Inserts new element at the beginning of the ForwardList.
+   *
+   * The element goes into the container right before its current first
+   * element. The content of val is copied (or moved) to the inserted
+   * element.
+   */
+  void push_front(const TYPE& val);
+
+  void push_front(TYPE&& val);
+
+  /**
+   * @brief Removes first element in ForwardList, reducing its size by one.
+   */
+  void pop_front();
+
+  /**
+   * @brief Inserts a new element after element at position.
+   *
+   * This new element is constructed in place using args as the arguments
+   * for its construction.
+   */
+  template <class... Args>
+  iterator emplace_after(const_iterator position, Args&&... args);
+
+  /**
+   * @brief Inserts new elements after the element at position.
+   */
+  iterator insert_after(const_iterator position, const TYPE& val);
+
+  iterator insert_after(const_iterator position, TYPE&& val);
+
+  iterator insert_after(const_iterator position, size_t n, const TYPE& val);
+
+  template <class InputIterator>
+  iterator insert_after(const_iterator position, InputIterator first,
+                        InputIterator last);
+
+  iterator insert_after(const_iterator position,
+                        std::initializer_list<TYPE> il);
+
+  /**
+   * @brief Removes either a single element or a range of elements.
+   */
+  iterator erase_after(const_iterator position);
+
+  iterator erase_after(const_iterator position, const_iterator last);
+
+  /**
+   * @brief Exchanges content by the content of fwdlst.
+   *
+   * Sizes may differ.
+   */
+  void swap(ForwardList& fwdlst);
+
+  /**
+   * @brief Resizes the container to contain n elements.
+   */
+  void resize(size_t n);
+
+  void resize(size_t n, const TYPE& val);
+
+  /**
+   * @brief Removes all elements, all leaves container with size 0.
+   */
+  void clear() noexcept;
+
+  /**
+   * @brief Transfers all elements of fwdlist into container.
+   */
+  void splice_after(const_iterator position, ForwardList& fwdlst);  // NOLINT
+
+  void splice_after(const_iterator position, ForwardList&& fwdlst);
+
+  /**
+   * @brief Transfers only elements pointed by from fwdlist into container.
+   */
+  void splice_after(const_iterator position, ForwardList& fwdlst,  // NOLINT
+                    const_iterator i);
+
+  void splice_after(const_iterator position, ForwardList&& fwdlst,
+                    const_iterator i);
+
+  /**
+   * @brief Transfers the range (first,last) from fwdlist into container.
+   */
+  void splice_after(const_iterator position, ForwardList& fwdlst,  // NOLINT
+                    const_iterator first, const_iterator last);
+
+  void splice_after(const_iterator position, ForwardList&& fwdlst,  // NOLINT
+                    const_iterator first, const_iterator last);
+
+  /**
+   * @brief Removes all the elements that compare equal to val.
+   */
+  void remove(const TYPE& val);
+
+  /**
+   * @brief Removes all elements for which Predicate pred returns true.
+   */
+  template <class Predicate>
+  void remove_if(Predicate pred);
+
+  /**
+   * @brief Removes all but first element from consecutive group equal elems.
+   */
+  void unique();
+
+  /**
+   * @brief Takes comparison func that determines "uniqueness" of elem.
+   */
+  template <class BinaryPredicate>
+  void unique(BinaryPredicate binary_pred);
+
+  /**
+   * @brief Merges ForwardList by transferring all elements.
+   *
+   * Both containers shall already be ordered before calling merge.
+   */
+  void merge(ForwardList& fwdlst);  // NOLINT
+
+  void merge(ForwardList&& fwdlst);  // NOLINT
+
+  /**
+   * @brief Same as merge, but take specific predicate to perform comparison.
+   */
+  template <class Compare>
+  void merge(ForwardList& fwdlst, Compare comp);  // NOLINT
+
+  template <class Compare>
+  void merge(ForwardList&& fwdlst, Compare comp);  // NOLINT
+
+  /**
+   * @brief Sorts elements in ForwardList, altering position in container.
+   */
+  void sort();
+
+  template <class Compare>
+  void sort(Compare comp);  // NOLINT
+
+  /**
+   * @brief Reverses order of elements in the ForwardList container.
+   */
+  void reverse() noexcept;
+
+  /**
+   * @brief Returns a copy of the allocator object associated with container.
+   */
+  ALLOC
+  get_allocator() const noexcept;
+
+ private:
+  /**
+   * @brief Internal memory allocator used to create list nodes.
+   */
+  MemoryAllocator allocator_{};
+
+  /**
+   * @brief First element in the list.
+   */
+  Node* head_{nullptr};
+
+  /**
+   * @brief Last element in the list.
+   */
+  Node* tail_{nullptr};
+
+  /**
+   * @brief Size of the list.
+   */
+  size_t size_{0};
+};
+
+/*****************************************************************************
+ ********************************* ITERATOR **********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+class ForwardList<TYPE, ALLOC>::Iterator {
+  using NodeT = typename ForwardList<TYPE, ALLOC>::Node;
+
+ public:
+  using iterator_category = std::forward_iterator_tag;
+  using difference_type = std::ptrdiff_t;
+  using value_type = TYPE;
+  using pointer = TYPE*;
+  using reference = TYPE&;
+
+  Iterator(NodeT* ptr);
+
+  Iterator operator++();
+
+  Iterator operator++(int);
+
+  template <bool Q = CONST>
+  typename std::enable_if<Q, const TYPE&>::type operator*();
+
+  template <bool Q = CONST>
+  typename std::enable_if<!Q, TYPE&>::type operator*();
+
+  pointer operator->();
+
+  template <typename ITER_TYPE>
+  friend bool operator==(ITER_TYPE& a, ITER_TYPE& b);  // NOLINT
+
+  template <typename ITER_TYPE>
+  friend bool operator==(ITER_TYPE& a, ITER_TYPE b);  // NOLINT
+
+  template <typename ITER_TYPE>
+  friend bool operator!=(ITER_TYPE& a, ITER_TYPE& b);  // NOLINT
+
+  template <typename ITER_TYPE>
+  friend bool operator!=(ITER_TYPE& a, ITER_TYPE b);  // NOLINT
+
+ private:
+  NodeT* node_ptr_;
+};
+
+/*****************************************************************************
+ ******************************* COMPARATORS *********************************
+ *****************************************************************************/
+
+template <class TYPE, class ALLOC>
+bool operator==(const ForwardList<TYPE, ALLOC>& lhs,
+                const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+bool operator==(const ForwardList<TYPE, ALLOC>& lhs, const std::string rhs);
+
+template <class TYPE, class ALLOC>
+bool operator==(const std::string lhs, const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+bool operator!=(const ForwardList<TYPE, ALLOC>& lhs,
+                const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+bool operator<(const ForwardList<TYPE, ALLOC>& lhs,
+               const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+bool operator<=(const ForwardList<TYPE, ALLOC>& lhs,
+                const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+bool operator>(const ForwardList<TYPE, ALLOC>& lhs,
+               const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+bool operator>=(const ForwardList<TYPE, ALLOC>& lhs,
+                const ForwardList<TYPE, ALLOC>& rhs);
+
+template <class TYPE, class ALLOC>
+void swap(ForwardList<TYPE, ALLOC>& x, ForwardList<TYPE, ALLOC>& y);  // NOLINT
+
+/*****************************************************************************
+ ******************************** STRINGIFIERS *******************************
+ *****************************************************************************/
+
+template <class TYPE>
+std::ostream& operator<<(std::ostream& os, ForwardList<TYPE> const& list);
+
+template <class TYPE, class ALLOC>
+std::string to_string(const ForwardList<TYPE, ALLOC>& list);
+
+/*****************************************************************************
+ *********************************** PROXY ***********************************
+ *****************************************************************************/
+
+template <typename ALLOC, typename TYPE>
+class ForwardListProxy {
+  using ProxyT = DeviceProxy<ALLOC, ForwardList<TYPE, ALLOC>, 1>;
+
+ public:
+  __host__ __device__ ForwardList<TYPE, ALLOC>* get() { return proxy_.get(); }
+
+ private:
+  ProxyT proxy_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_FORWARD_LIST_HPP_

src/containers/forward_list_impl.hpp

@@ -0,0 +1,233 @@
+/******************************************************************************
+ * 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_CONTAINERS_FORWARD_LIST_IMPL_HPP_
+#define LIBRARY_SRC_CONTAINERS_FORWARD_LIST_IMPL_HPP_
+
+#include <algorithm>
+#include <cassert>
+#include <string>
+
+#include "forward_list.hpp"
+
+namespace rocshmem {
+
+/*****************************************************************************
+ ******************************* FORWARD LIST ********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC>
+ForwardList<TYPE, ALLOC>::~ForwardList() {
+  while (head_ != nullptr) {
+    Node* temp{head_};
+    head_ = head_->next;
+    allocator_.deallocate(temp);
+  }
+  tail_ = nullptr;
+}
+
+template <typename TYPE, typename ALLOC>
+ForwardList<TYPE, ALLOC>::ForwardList(const ForwardList& fwdlst) {}
+
+template <typename TYPE, typename ALLOC>
+ForwardList<TYPE, ALLOC>::ForwardList(const ForwardList& fwdlst,
+                                      const ALLOC& alloc) {}
+
+template <typename TYPE, typename ALLOC>
+ForwardList<TYPE, ALLOC>::ForwardList(std::initializer_list<TYPE> il,
+                                      const ALLOC& allocator) {
+  allocator_ = allocator;
+  assign(il);
+}
+
+template <typename TYPE, typename ALLOC>
+typename ForwardList<TYPE, ALLOC>::iterator
+ForwardList<TYPE, ALLOC>::begin() noexcept {
+  iterator i(head_);
+  return i;
+}
+
+template <typename TYPE, typename ALLOC>
+typename ForwardList<TYPE, ALLOC>::const_iterator
+ForwardList<TYPE, ALLOC>::begin() const noexcept {
+  const_iterator i(head_);
+  return i;
+}
+
+template <typename TYPE, typename ALLOC>
+typename ForwardList<TYPE, ALLOC>::iterator
+ForwardList<TYPE, ALLOC>::end() noexcept {
+  iterator i(tail_);
+  return i;
+}
+
+template <typename TYPE, typename ALLOC>
+typename ForwardList<TYPE, ALLOC>::const_iterator
+ForwardList<TYPE, ALLOC>::end() const noexcept {
+  const_iterator i(tail_);
+  return i;
+}
+
+template <typename TYPE, typename ALLOC>
+void ForwardList<TYPE, ALLOC>::assign(std::initializer_list<TYPE> il) {
+  resize(il.size());
+  std::copy_n(il.begin(), il.size(), begin());
+}
+
+template <typename TYPE, typename ALLOC>
+void ForwardList<TYPE, ALLOC>::resize(size_t n) {}
+
+template <typename TYPE, typename ALLOC>
+void ForwardList<TYPE, ALLOC>::resize(size_t n, const TYPE& val) {}
+
+template <typename TYPE, typename ALLOC>
+void ForwardList<TYPE, ALLOC>::clear() noexcept {}
+
+/*****************************************************************************
+ ********************************* ITERATOR **********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+ForwardList<TYPE, ALLOC>::Iterator<CONST>::Iterator(NodeT* ptr)
+    : node_ptr_(ptr) {}
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+typename ForwardList<TYPE, ALLOC>::template Iterator<CONST>
+ForwardList<TYPE, ALLOC>::Iterator<CONST>::operator++() {
+  if (node_ptr_) {
+    node_ptr_ = node_ptr_->next;
+  }
+  return *this;
+}
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+typename ForwardList<TYPE, ALLOC>::template Iterator<CONST>
+ForwardList<TYPE, ALLOC>::Iterator<CONST>::operator++(int) {
+  Iterator iterator = *this;
+  ++*this;
+  return iterator;
+}
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+template <bool Q>
+typename std::enable_if<!Q, TYPE&>::type
+ForwardList<TYPE, ALLOC>::Iterator<CONST>::operator*() {
+  return node_ptr_->data;
+}
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+template <bool Q>
+typename std::enable_if<Q, const TYPE&>::type
+ForwardList<TYPE, ALLOC>::Iterator<CONST>::operator*() {
+  return node_ptr_->data;
+}
+
+template <typename TYPE, typename ALLOC>
+template <bool CONST>
+TYPE* ForwardList<TYPE, ALLOC>::Iterator<CONST>::operator->() {
+  return &node_ptr_->data;
+}
+
+template <typename ITER_TYPE>
+bool operator==(ITER_TYPE& a, ITER_TYPE& b) {  // NOLINT
+  return a.node_ptr_ == b.node_ptr_;
+}
+
+template <typename ITER_TYPE>
+bool operator==(ITER_TYPE& a, ITER_TYPE b) {  // NOLINT
+  return a.node_ptr_ == b.node_ptr_;
+}
+
+template <typename ITER_TYPE>
+bool operator!=(ITER_TYPE& a, ITER_TYPE& b) {  // NOLINT
+  return a.node_ptr_ != b.node_ptr_;
+}
+
+template <typename ITER_TYPE>
+bool operator!=(ITER_TYPE& a, ITER_TYPE b) {  // NOLINT
+  return a.node_ptr_ != b.node_ptr_;
+}
+
+/*****************************************************************************
+ ******************************* COMPARATORS *********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC>
+bool operator==(const ForwardList<TYPE, ALLOC>& lhs,
+                const ForwardList<TYPE, ALLOC>& rhs) {
+  return false;
+}
+
+template <typename TYPE, typename ALLOC>
+bool operator==(const ForwardList<TYPE, ALLOC>& lhs, const std::string rhs) {
+  return false;
+}
+
+template <typename TYPE, typename ALLOC>
+bool operator==(const std::string lhs, const ForwardList<TYPE, ALLOC>& rhs) {
+  return false;
+}
+
+template <typename TYPE, typename ALLOC>
+bool operator!=(const ForwardList<TYPE, ALLOC>& lhs,
+                const ForwardList<TYPE, ALLOC>& rhs) {
+  return false;
+}
+
+template <typename TYPE, typename ALLOC>
+bool operator!=(const ForwardList<TYPE, ALLOC>& lhs, const std::string rhs) {
+  return false;
+}
+
+template <typename TYPE, typename ALLOC>
+bool operator!=(const std::string lhs, const ForwardList<TYPE, ALLOC>& rhs) {
+  return false;
+}
+
+/*****************************************************************************
+ ******************************** STRINGIFIERS *******************************
+ *****************************************************************************/
+
+template <typename TYPE>
+std::ostream& operator<<(std::ostream& os, ForwardList<TYPE> const& list) {
+  using CItr = typename ForwardList<TYPE>::const_iterator;
+  for (CItr ci = list.begin(); ci != list.end(); ++ci) {
+    os << *ci << " ";
+  }
+  return os;
+}
+
+template <typename TYPE, typename ALLOC>
+std::string to_string(const ForwardList<TYPE, ALLOC>& list) {
+  std::stringstream ss;
+  ss << list;
+  return ss.str();
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_FORWARD_LIST_IMPL_HPP_

src/containers/free_list.hpp

@@ -0,0 +1,278 @@
+/******************************************************************************
+ * 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_CONTAINERS_FREE_LIST_HPP_
+#define LIBRARY_SRC_CONTAINERS_FREE_LIST_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include "../memory/hip_allocator.hpp"
+#include "../sync/abql_block_mutex.hpp"
+
+namespace rocshmem {
+
+// Forward declaration of the proxy.
+template <typename ALLOCATOR, typename TYPE>
+class FreeListProxy;
+
+/*****************************************************************************
+ ******************************* FREE LIST ***********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC = HIPDefaultFinegrainedAllocator>
+class FreeList {
+  friend class FreeListProxy<ALLOC, TYPE>;
+
+  using MutexProxyType = ABQLBlockMutexProxy<ALLOC>;
+  using MutexType = ABQLBlockMutex;
+
+  struct Node {
+    TYPE data;
+    Node* next{nullptr};
+  };
+
+  struct PopBackResult {
+    TYPE value;
+    bool success;
+  };
+
+  /**
+   * @brief A lock guard for ticket-based locks that follows the design of
+   * `std::lock_guard`.
+   *
+   * @tparam MUTEX The type of the ticket-based mutex to lock.
+   */
+  template <typename MUTEX>
+  struct TicketLockGuard {
+    /**
+     * @brief Constructs the `TicketLockGuard` and locks the mutex.
+     *
+     * @param m Mutex to take ownership of.
+     */
+    __device__ explicit TicketLockGuard(MUTEX& m) : mutex_{m} {
+      ticket_ = mutex_.lock();
+      __threadfence();
+    }
+
+    /**
+     * @brief Lock guards are not copyable
+     */
+    __device__ TicketLockGuard(const TicketLockGuard&) = delete;
+
+    /**
+     * @brief Lock guards are not moveable
+     */
+    __device__ TicketLockGuard(TicketLockGuard&&) = delete;
+
+    /**
+     * @brief Destructor the unlocks the mutex.
+     */
+    __device__ ~TicketLockGuard() {
+      __threadfence();
+      mutex_.unlock(ticket_);
+    }
+
+   private:
+    using TicketT = uint64_t;
+    MUTEX& mutex_;
+    TicketT ticket_;
+  };
+
+ public:
+  /**
+   * @brief Construct a new Free List object
+   *
+   * @param alloc Allocator to use for free list nodes allocations.
+   */
+  explicit FreeList(const ALLOC& alloc = ALLOC());
+
+  /**
+   * @brief Constructs a FreeList object with contents from a range of elements
+   * defined by [`first`, `last`).
+   *
+   * @tparam InputIt Iterator type of the elements to store in the free-list.
+   * @param first First element in the range defining the input elements.
+   * @param last Element after last that defines the input elements range.
+   * @param alloc Allocator to use for allocating internal structures of the
+   * free list.
+   */
+  template <class InputIt>
+  FreeList(InputIt first, InputIt last, const ALLOC& alloc = ALLOC());
+
+  /**
+   * @brief Pushes a range of elements defined by [`first`, `last`).
+   *
+   * @tparam InputIt Iterator type of the elements to store in the free-list.
+   * @param first First element in the range defining the input elements.
+   * @param last Element after last that defines the input elements range.
+   *
+   * @return @c true if success, or @c false otherwise.
+   */
+  template <class InputIt>
+  __host__ bool push_back_range(InputIt first, InputIt last);
+
+  /**
+   * @brief Destroy the Free List object
+   */
+  ~FreeList();
+
+  /**
+   * @brief  Inserts new element at the end of the FreeList.
+   *
+   * The element goes into the container right after its last
+   * element. The content of val is copied (or moved) to the inserted
+   * element.
+   *
+   * @note Host-side API is not thread safe.
+   *
+   * @param val The value to insert in the FreeList.
+   * @return @c true if the operation succeed, and @c false otherwise.
+   */
+  __device__ bool push_back(const TYPE& val);
+
+  /// @copydoc bool FreeList<TYPE, ALLOC>::push_back(const TYPE&)
+  __device__ bool push_back(TYPE&& val);
+
+  /// @copydoc bool FreeList<TYPE, ALLOC>::push_back(const TYPE&)
+  __host__ bool push_back(const TYPE& val);
+
+  /// @copydoc bool FreeList<TYPE, ALLOC>::push_back(const TYPE&)
+  __host__ bool push_back(TYPE&& val);
+
+  /**
+   * @brief Removes the first element in FreeList, reducing its size by one.
+   *
+   * @return An object with two fields `value` and `success`. `success` is a
+   * boolean indicating if the operation succeeded, and if the operation
+   * succeeded, the `value` field contains the popped value.
+   */
+  __device__ PopBackResult pop_front();
+
+ private:
+  /**
+   * @brief Deallocates all memory that was dynamically allocated using the free
+   * list.
+   */
+  void deallocate_all_nodes();
+
+  /**
+   * @brief Allocates a node using the host-side allocator or the recycled
+   * pointers on the device.
+   *
+   * @note The device-side API assumes the data structure is protected by a
+   * lock.
+   *
+   * @return A pointer to the allocated node.
+   */
+  __host__ Node* allocate_node() {
+    Node* node;
+    allocator_.allocate((void**)(&node), sizeof(Node));
+    return node;
+  };
+
+  /// @copydoc FreeList<TYPE, ALLOC>::Node* FreeList<TYPE,
+  /// ALLOC>::allocate_node(const TYPE&)
+  __device__ Node* allocate_node() {
+    Node* node = deallocated_nodes_;
+    if (node != nullptr) {
+      deallocated_nodes_ = node->next;
+    }
+    return node;
+  };
+
+  /**
+   * @brief Appends a node to the tail of the free list.
+   *
+   * @param node Node to append to the list.
+   */
+  __host__ __device__ void insert_node_at_tail(Node* node) {
+    if (tail_ != nullptr) {
+      tail_->next = node;
+    }
+    tail_ = node;
+
+    // if the list is empty, set the head_ to the first node
+    if (head_ == nullptr) {
+      head_ = node;
+    }
+  }
+
+  /**
+   * @brief Device-side node deallocation that inserts the deallocated node into
+   * a linked-list of pointers to reuse on the device.
+   *
+   * @note Assumes structure is protected by a lock.
+   */
+  __device__ void deallocate_node(Node* node) {
+    // append the node to the head of the linked list
+    node->next = deallocated_nodes_;
+    deallocated_nodes_ = node;
+  };
+
+  /**
+   * @brief Internal memory allocator used to create list nodes.
+   */
+  MemoryAllocator allocator_{};
+
+  /**
+   * @brief First element in the list.
+   */
+  Node* head_{};
+
+  /**
+   * @brief Last element in the list.
+   */
+  Node* tail_{};
+
+  /**
+   * @brief A linked-list of deallocated nodes.
+   */
+  Node* deallocated_nodes_{};
+
+  /**
+   * @brief Mutex protecting the free-list mutations.
+   */
+  MutexProxyType mutex_;
+};
+
+template <typename ALLOCATOR, typename TYPE>
+class FreeListProxy {
+  using FreeListT = FreeList<TYPE, ALLOCATOR>;
+  using ProxyT = DeviceProxy<ALLOCATOR, FreeListT>;
+
+ public:
+  __host__ __device__ FreeListT* get() { return proxy_.get(); }
+
+  FreeListProxy() { new (proxy_.get()) FreeListT(); }
+
+  ~FreeListProxy() {
+    auto free_list = proxy_.get();
+    free_list->deallocate_all_nodes();
+    free_list->~FreeList();
+  }
+
+ private:
+  ProxyT proxy_{};
+};
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_FREE_LIST_HPP_

src/containers/free_list_impl.hpp

@@ -0,0 +1,141 @@
+/******************************************************************************
+ * 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_CONTAINERS_FREE_LIST_IMPL_HPP_
+#define LIBRARY_SRC_CONTAINERS_FREE_LIST_IMPL_HPP_
+
+#include "free_list.hpp"
+
+namespace rocshmem {
+
+/*****************************************************************************
+ ******************************* FREE LIST ***********************************
+ *****************************************************************************/
+
+template <typename TYPE, typename ALLOC>
+FreeList<TYPE, ALLOC>::~FreeList() {}
+
+template <typename TYPE, typename ALLOC>
+void FreeList<TYPE, ALLOC>::deallocate_all_nodes() {
+  // Deallocate any existing nodes
+  while (head_ != nullptr) {
+    auto temp = head_;
+    head_ = temp->next;
+    allocator_.deallocate(temp);
+  }
+
+  // The tail no longer points to any nodes
+  tail_ = nullptr;
+
+  // Deallocate all recycled nodes
+  while (deallocated_nodes_ != nullptr) {
+    auto temp = deallocated_nodes_;
+    deallocated_nodes_ = temp->next;
+    allocator_.deallocate(temp);
+  }
+}
+
+template <typename TYPE, typename ALLOC>
+FreeList<TYPE, ALLOC>::FreeList(const ALLOC& alloc)
+    : allocator_{alloc},
+      head_{nullptr},
+      tail_{nullptr},
+      deallocated_nodes_{nullptr} {}
+
+template <typename TYPE, typename ALLOC>
+template <class InputIt>
+bool FreeList<TYPE, ALLOC>::push_back_range(InputIt first, InputIt last) {
+  for (auto iter = first; iter != last; iter++) {
+    auto key = *iter;
+    const bool result = push_back(key);
+    if (!result) {
+      return false;
+    }
+  }
+  return true;
+}
+
+template <typename TYPE, typename ALLOC>
+__device__ bool FreeList<TYPE, ALLOC>::push_back(const TYPE& val) {
+  TicketLockGuard<MutexType> guard(*mutex_.get());
+  auto node = allocate_node();
+  if (node == nullptr) {
+    return false;
+  }
+
+  node->data = val;
+  node->next = nullptr;
+
+  insert_node_at_tail(node);
+
+  return true;
+}
+
+template <typename TYPE, typename ALLOC>
+__device__ bool FreeList<TYPE, ALLOC>::push_back(TYPE&& val) {
+  return push_back(std::forward<const TYPE>(val));
+}
+
+template <typename TYPE, typename ALLOC>
+__host__ bool FreeList<TYPE, ALLOC>::push_back(const TYPE& val) {
+  auto node = allocate_node();
+  if (node == nullptr) {
+    return false;
+  }
+
+  node->data = val;
+  node->next = nullptr;
+  insert_node_at_tail(node);
+
+  return true;
+}
+template <typename TYPE, typename ALLOC>
+__host__ bool FreeList<TYPE, ALLOC>::push_back(TYPE&& val) {
+  return push_back(std::forward<const TYPE>(val));
+}
+
+template <typename TYPE, typename ALLOC>
+__device__ typename FreeList<TYPE, ALLOC>::PopBackResult
+FreeList<TYPE, ALLOC>::pop_front() {
+  TicketLockGuard<MutexType> guard(*mutex_.get());
+
+  if (head_ == nullptr) {
+    return {{}, false};
+  }
+  auto last_node = head_;
+  head_ = head_->next;
+
+  // if we removed all nodes, we should reset the tail
+  if (head_ == nullptr) {
+    tail_ = nullptr;
+  }
+
+  TYPE result{last_node->data};
+
+  deallocate_node(last_node);
+
+  return {result, true};
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_FREE_LIST_IMPL_HPP_

src/containers/helper_macros.hpp

@@ -0,0 +1,52 @@
+/******************************************************************************
+ * 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_CONTAINERS_HELPER_MACROS_HPP_
+#define LIBRARY_SRC_CONTAINERS_HELPER_MACROS_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <iostream>
+
+#include "roc_shmem/roc_shmem.hpp"
+
+#define BARRIER() rocshmem::roc_shmem_wg_barrier_all()
+#define RANK rocshmem::roc_shmem_my_pe()
+#define NPES rocshmem::roc_shmem_n_pes()
+
+#define PE_BITS ((uint64_t)ceil(log(NPES) / log(2)))
+#define PE_OF(X) ((X) >> (64 - PE_BITS))
+
+#define _printf \
+  if (RANK == 0) printf
+#define _cout \
+  if (RANK == 0) std::cout
+#define _cerr \
+  if (RANK == 0) std::cerr
+
+#define GIBI 1073741824L
+#define MEBI 1048576
+
+#endif  // LIBRARY_SRC_CONTAINERS_HELPER_MACROS_HPP_

src/containers/index_strategy.hpp

@@ -0,0 +1,639 @@
+/******************************************************************************
+ * 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_CONTAINERS_INDEX_STRATEGY_HPP_
+#define LIBRARY_SRC_CONTAINERS_INDEX_STRATEGY_HPP_
+
+/**
+ * @file index_strategy.hpp
+ *
+ * @section
+ * Q: What is an indexing strategy?
+ *
+ * A: The indexing strategy is a scheme used by __device__ code to access
+ *    a container's raw memory region.
+ *    The indexing strategy behaves like a simple forward iterator (in normal
+ *    STL code).
+ *
+ * Q: What are the inputs and outputs of an indexing strategy?
+ *
+ * A: [INPUT]  an indexing strategy needs block and grid information
+ *    [INPUT]  an indexing strategy needs to know the current index
+ *    [INPUT]  an indexing strategy needs to know number of container elements
+ *    [OUTPUT] an indexing strategy returns an index
+ *
+ * Q: What are the names of the strategies?
+ *
+ * A: Thread_Contiguous_Block_Agnostic
+ *    Thread_Discontiguous_Block_Discontiguous
+ *    Thread_Discontiguous_Block_Contiguous
+ *
+ * Q: What is the Thread_Contiguous_Block_Agnostic strategy?
+ *
+ * A: Thread_Contiguous_Block_Agnostic keeps a thread's memory accesses
+ *    contiguous. Only one thread may access the container; the container
+ *    is private to the thread (thread-private).
+ *
+ * Q: What is the Thread_Discontiguous_Block_Discontiguous strategy?
+ *
+ * A: Thread_Discontiguous_Block_Discontiguous does not keep a thread
+ *    block's memory accesses contiguous. The container memory is accessed
+ *    by multiple thread blocks and the accesses by different thread blocks
+ *    are interleaved.
+ *
+ *    For example:
+ *      assume grid_dim {8, 1, 1} block_dim {4, 1, 1}
+ *      assume container._size = 72
+ *
+ *      In table below, '#' denotes thread block #'s accesses.
+ *           0   1   2   3   4   5   6   7   8   9   10  11
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      0  | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      1  | 3 | 3 | 3 | 3 | 4 | 4 | 4 | 4 | 5 | 5 | 5 | 5 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      2  | 6 | 6 | 6 | 6 | 7 | 7 | 7 | 7 | 0 | 0 | 0 | 0 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      3  | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3 | 3 | 3 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      4  | 4 | 4 | 4 | 4 | 5 | 5 | 5 | 5 | 6 | 6 | 6 | 6 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      5  | 7 | 7 | 7 | 7 | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *
+ *      thread_id_00: grid{0, 0, 0} block{0, 0, 0}
+ *      thread_id_00.2d_accesses := { {0,0},
+ *                                    {2,8},
+ *                                    {5,4} }
+ *      thread_id_00.1d_accesses := { 0, 32, 64 }
+ *
+ *      thread_id_14: grid{3, 0, 0} block{2, 0, 0}
+ *      thread_id_14.2d_accesses := { {1,2},
+ *                                    {3,9} }
+ *      thread_id_14.1d_accesses := { 14, 46 }
+ *
+ * Q: What is the Thread_Discontiguous_Block_Contiguous strategy?
+ *
+ * A: Thread_Discontiguous_Block_Contiguous does keep a thread block's memory
+ *    accesses contiguous. The container memory is accessed by multiple
+ *    thread blocks and the accesses by different thread blocks are
+ *    restricted by artificial boundaries called 'tiles'.
+ *
+ *    For example:
+ *      assume grid_dim {8, 1, 1} block_dim {4, 1, 1}
+ *      assume container._size = 72
+ *
+ *      In table below, '#' denotes thread block #'s accesses.
+ *           0   1   2   3   4   5   6   7   8   9   10  11
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      0  | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 1 | 1 | 1 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      1  | 1 | 1 | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 2 | 2 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      2  | 2 | 2 | 2 | 3 | 3 | 3 | 3 | 3 | 3 | 3 | 3 | 3 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      3  | 4 | 4 | 4 | 4 | 4 | 4 | 4 | 4 | 4 | 5 | 5 | 5 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      4  | 5 | 5 | 5 | 5 | 5 | 5 | 6 | 6 | 6 | 6 | 6 | 6 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *      5  | 6 | 6 | 6 | 7 | 7 | 7 | 7 | 7 | 7 | 7 | 7 | 7 |
+ *         +---+---+---+---+---+---+---+---+---+---+---+---+
+ *
+ * Q: How do I use the strategies?
+ *
+ * A: You may use a strategy like this:
+ *
+ *      for (size_t i = _dev_idx.start();
+ *           i < _dev_idx.end();
+ *           i = _dev_idx.next(i)) {
+ *          container[i] = ...
+ *      }
+ */
+
+#include <hip/hip_runtime.h>
+
+#include <algorithm>
+
+namespace rocshmem {
+
+class Identity {
+ public:
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t block_size() const {
+    return hipBlockDim_x * hipBlockDim_y * hipBlockDim_z;
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t grid_size() const {
+    return hipGridDim_x * hipGridDim_y * hipGridDim_z;
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t block_id() const {
+    return hipBlockIdx_x + hipBlockIdx_y * hipGridDim_x +
+           hipBlockIdx_z * hipGridDim_x * hipGridDim_y;
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t local_thread_id() const {
+    return hipThreadIdx_x + hipThreadIdx_y * hipBlockDim_x +
+           hipThreadIdx_z * hipBlockDim_x * hipBlockDim_y;
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t global_thread_id() const {
+    return local_thread_id() + block_id() * block_size();
+  }
+};
+
+class Thread_Contiguous_Block_Agnostic {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Thread_Contiguous_Block_Agnostic() = default;
+
+  /**
+   * @brief
+   *
+   * @param container_elems
+   */
+  __host__ __device__ Thread_Contiguous_Block_Agnostic(size_t container_elems)
+      : _container_elems(container_elems) {}
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() { return 0; }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(_container_elems);
+    return _container_elems;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) { return current_index + 1; }
+
+ private:
+  /**
+   * @brief
+   */
+  size_t _container_elems{0};
+};
+
+class Thread_Discontiguous_Block_Discontiguous {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Thread_Discontiguous_Block_Discontiguous() = default;
+
+  /**
+   * @brief
+   *
+   * @param container_elems
+   */
+  __host__ __device__ explicit Thread_Discontiguous_Block_Discontiguous(
+      size_t container_elems)
+      : _container_elems(container_elems) {}
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() { return _id.global_thread_id(); }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(_container_elems);
+    return _container_elems;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    return current_index + _id.grid_size() * _id.block_size();
+  }
+
+ private:
+  size_t _container_elems{0};
+
+  Identity _id{};
+};
+
+class Thread_Discontiguous_Block_Contiguous {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Thread_Discontiguous_Block_Contiguous() = default;
+
+  /**
+   * @brief
+   *
+   * @param container_elems
+   */
+  __host__ __device__ explicit Thread_Discontiguous_Block_Contiguous(
+      size_t container_elems)
+      : _container_elems(container_elems) {
+    size_t left_over = _container_elems % _id.grid_size();
+
+    _tile_offset = _id.block_id() * (_container_elems / _id.grid_size()) +
+                   min(_id.block_id(), left_over);
+
+    _tile_num_elems = _container_elems / _id.grid_size();
+    if (_id.block_id() < left_over) {
+      _tile_num_elems++;
+    }
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() {
+    assert(_container_elems);
+    return _tile_offset + _id.local_thread_id();
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(_container_elems);
+    return _tile_offset + _tile_num_elems;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    assert(_container_elems);
+    return current_index + _id.block_size();
+  }
+
+ private:
+  size_t _container_elems{0};
+  size_t _tile_offset{0};
+  size_t _tile_num_elems{0};
+
+  Identity _id{};
+};
+
+class Thread_Discontiguous_Block_Private {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Thread_Discontiguous_Block_Private() = default;
+
+  /**
+   * @brief
+   *
+   * @param container_elems
+   */
+  __host__ __device__ Thread_Discontiguous_Block_Private(size_t container_elems)
+      : _container_elems(container_elems) {}
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() { return _id.local_thread_id(); }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(_container_elems);
+    return _container_elems;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    return current_index + _id.block_size();
+  }
+
+ private:
+  size_t _container_elems{0};
+
+  Identity _id{};
+};
+
+class Matrix_Private {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Matrix_Private() = default;
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() { return _id.global_thread_id(); }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(false);
+    return 0;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    assert(false);
+    return 0;
+  }
+
+ private:
+  Identity _id{};
+};
+
+class Matrix_Block {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Matrix_Block() = default;
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() { return _id.block_id(); }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(false);
+    return 0;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    assert(false);
+    return 0;
+  }
+
+ private:
+  Identity _id{};
+};
+
+class Matrix_Device {
+ public:
+  /**
+   * @brief
+   */
+  __host__ __device__ Matrix_Device() = default;
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() { return 0; }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    assert(false);
+    return 0;
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    assert(false);
+    return 0;
+  }
+};
+
+enum class IndexStrategyEnum {
+  TCBA = 0,
+  TDBD = 1,
+  TDBC = 2,
+  TDBP = 3,
+  MP = 4,
+  MB = 5,
+  MD = 6,
+  UNSET = 7
+};
+
+class IndexStrategy {
+ public:
+  __host__ __device__ IndexStrategy() = default;
+
+  __host__ __device__ IndexStrategy(IndexStrategyEnum ise) : _ise(ise) {}
+
+  __host__ __device__ IndexStrategy(IndexStrategyEnum ise,
+                                    size_t container_elems)
+      : _ise(ise),
+        _tcba(container_elems),
+        _tdbd(container_elems),
+        _tdbc(container_elems),
+        _tdbp(container_elems) {}
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t start() {
+    switch (_ise) {
+      case IndexStrategyEnum::TCBA:
+        return _tcba.start();
+      case IndexStrategyEnum::TDBD:
+        return _tdbd.start();
+      case IndexStrategyEnum::TDBC:
+        return _tdbc.start();
+      case IndexStrategyEnum::TDBP:
+        return _tdbp.start();
+      case IndexStrategyEnum::MP:
+        return _mp.start();
+      case IndexStrategyEnum::MB:
+        return _mb.start();
+      case IndexStrategyEnum::MD:
+        return _md.start();
+      case IndexStrategyEnum::UNSET:
+        assert(false);
+        return 0;
+    }
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ size_t end() {
+    switch (_ise) {
+      case IndexStrategyEnum::TCBA:
+        return _tcba.end();
+      case IndexStrategyEnum::TDBD:
+        return _tdbd.end();
+      case IndexStrategyEnum::TDBC:
+        return _tdbc.end();
+      case IndexStrategyEnum::TDBP:
+        return _tdbp.end();
+      case IndexStrategyEnum::MP:
+        return _mp.end();
+      case IndexStrategyEnum::MB:
+        return _mb.end();
+      case IndexStrategyEnum::MD:
+        return _md.end();
+      case IndexStrategyEnum::UNSET:
+        assert(false);
+        return 0;
+    }
+  }
+
+  /**
+   * @brief
+   *
+   * @param current_index
+   *
+   * @return
+   */
+  __device__ size_t next(size_t current_index) {
+    switch (_ise) {
+      case IndexStrategyEnum::TCBA:
+        return _tcba.next(current_index);
+      case IndexStrategyEnum::TDBD:
+        return _tdbd.next(current_index);
+      case IndexStrategyEnum::TDBC:
+        return _tdbc.next(current_index);
+      case IndexStrategyEnum::TDBP:
+        return _tdbp.next(current_index);
+      case IndexStrategyEnum::MP:
+        return _mp.next(current_index);
+      case IndexStrategyEnum::MB:
+        return _mb.next(current_index);
+      case IndexStrategyEnum::MD:
+        return _md.next(current_index);
+      case IndexStrategyEnum::UNSET:
+        assert(false);
+        return 0;
+    }
+  }
+
+ private:
+  IndexStrategyEnum _ise{IndexStrategyEnum::UNSET};
+  Thread_Contiguous_Block_Agnostic _tcba{};
+  Thread_Discontiguous_Block_Discontiguous _tdbd{};
+  Thread_Discontiguous_Block_Contiguous _tdbc{};
+  Thread_Discontiguous_Block_Private _tdbp{};
+  Matrix_Private _mp{};
+  Matrix_Block _mb{};
+  Matrix_Device _md{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_INDEX_STRATEGY_HPP_

src/containers/matrix.hpp

@@ -0,0 +1,167 @@
+/******************************************************************************
+ * 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_CONTAINERS_MATRIX_HPP_
+#define LIBRARY_SRC_CONTAINERS_MATRIX_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include <vector>
+
+#include "index_strategy.hpp"
+#include "memory_allocator.hpp"
+
+namespace rocshmem {
+
+template <typename T>
+struct requires_internal_allocator {
+  static const bool value{false};
+};
+
+template <typename T>
+struct one_generic_constructor_parameter {
+  static const bool value{false};
+};
+
+template <typename TYPE>
+class Matrix {
+ public:
+  /**
+   * @brief
+   *
+   * @param[in] row_index
+   * @param[in] col_index
+   *
+   * @return
+   */
+  __host__ __device__ TYPE* access(size_t row_index, size_t col_index);
+
+  /**
+   * @brief
+   *
+   * @param[in] col_index
+   *
+   * @return
+   */
+  __device__ TYPE* access(size_t col_index);
+
+  /**
+   * @brief
+   *
+   * @param[in] col_index
+   *
+   * @return
+   */
+  __device__ TYPE* access();
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __host__ __device__ size_t rows() const;
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __host__ __device__ size_t columns() const;
+
+ protected:
+  /**
+   * @brief
+   */
+  Matrix() = default;
+
+  /**
+   * @brief
+   *
+   * @param[in] number_rows
+   * @param[in] number_columns
+   * @param[in] allocator
+   * @param[in] strategy
+   */
+  Matrix(size_t number_rows, size_t number_columns, MemoryAllocator allocator,
+         const IndexStrategy index_strat);
+
+  /**
+   * @brief
+   *
+   * @param[in] number_rows
+   * @param[in] number_columns
+   * @param[in] TYPE_constructor_param
+   * @param[in] allocator
+   * @param[in] strategy
+   */
+  Matrix(size_t number_rows, size_t number_columns,
+         size_t TYPE_constructor_param, MemoryAllocator allocator,
+         const IndexStrategy index_strat);
+
+  /**
+   * @brief
+   *
+   * @param[in] number_rows
+   * @param[in] number_columns
+   * @param[in] TYPE_constructor_param
+   * @param[in] allocator
+   * @param[in] strategy
+   */
+  Matrix(size_t number_rows, size_t number_columns,
+         std::vector<size_t> TYPE_constructor_param, MemoryAllocator allocator,
+         const IndexStrategy index_strat);
+
+  /**
+   * @brief
+   */
+  ~Matrix();
+
+ private:
+  /**
+   * @brief
+   */
+  size_t _number_rows{0};
+
+  /**
+   * @brief
+   */
+  size_t _number_columns{0};
+
+  /**
+   * @brief
+   */
+  MemoryAllocator _allocator{};
+
+  /**
+   * @brief
+   */
+  TYPE** _flat_c_array{nullptr};
+
+  /**
+   * @brief
+   */
+  IndexStrategy _dev_idx{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_MATRIX_HPP_

src/containers/matrix_impl.hpp

@@ -0,0 +1,209 @@
+/******************************************************************************
+ * 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_CONTAINERS_MATRIX_IMPL_HPP_
+#define LIBRARY_SRC_CONTAINERS_MATRIX_IMPL_HPP_
+
+#include <vector>
+
+#include "matrix.hpp"
+
+namespace rocshmem {
+
+template <typename TYPE>
+Matrix<TYPE>::Matrix(size_t number_rows, size_t number_columns,
+                     MemoryAllocator allocator, const IndexStrategy strategy)
+    : _number_rows(number_rows),
+      _number_columns(number_columns),
+      _allocator(allocator),
+      _dev_idx(strategy) {
+  /*
+   * Allocate the flattened c-array which contains the top-level
+   * TYPE pointers.
+   */
+  size_t flat_c_array_dimensions = _number_rows * _number_columns;
+  size_t size_bytes = flat_c_array_dimensions * sizeof(TYPE*);
+  _allocator.allocate(reinterpret_cast<void**>(&_flat_c_array), size_bytes);
+
+  /*
+   * Iterate through the flattened c-array and initialize each pointer
+   * with a valid TYPE.
+   */
+  for (size_t i = 0; i < flat_c_array_dimensions; i++) {
+    _allocator.allocate(reinterpret_cast<void**>(&_flat_c_array[i]),
+                        sizeof(TYPE));
+
+    if constexpr (one_generic_constructor_parameter<TYPE>::value &&
+                  requires_internal_allocator<TYPE>::value) {
+      /*
+       * Do not invoke constructor since type traits do not match.
+       */
+    } else if constexpr (requires_internal_allocator<TYPE>::value) {
+      /*
+       * Construct the TYPE with placement new.
+       */
+      new (_flat_c_array[i]) TYPE(allocator, strategy);
+    }
+  }
+}
+
+template <typename TYPE>
+Matrix<TYPE>::Matrix(size_t number_rows, size_t number_columns,
+                     size_t TYPE_constructor_param, MemoryAllocator allocator,
+                     const IndexStrategy strategy)
+    : _number_rows(number_rows),
+      _number_columns(number_columns),
+      _allocator(allocator),
+      _dev_idx(strategy) {
+  /*
+   * Allocate the flattened c-array which contains the top-level
+   * TYPE pointers.
+   */
+  size_t flat_c_array_dimensions = _number_rows * _number_columns;
+  size_t size_bytes = flat_c_array_dimensions * sizeof(TYPE*);
+  _allocator.allocate(reinterpret_cast<void**>(&_flat_c_array), size_bytes);
+
+  /*
+   * Iterate through the flattened c-array and initialize each pointer
+   * with a valid TYPE.
+   */
+  for (size_t i = 0; i < flat_c_array_dimensions; i++) {
+    _allocator.allocate(reinterpret_cast<void**>(&_flat_c_array[i]),
+                        sizeof(TYPE));
+
+    /*
+     * If type traits match, construct the TYPE with placement new.
+     */
+    if constexpr (one_generic_constructor_parameter<TYPE>::value &&
+                  requires_internal_allocator<TYPE>::value) {
+      new (_flat_c_array[i]) TYPE(TYPE_constructor_param, allocator, strategy);
+    } else if constexpr (one_generic_constructor_parameter<TYPE>::value) {
+      new (_flat_c_array[i]) TYPE(TYPE_constructor_param, strategy);
+    }
+  }
+}
+
+template <typename TYPE>
+Matrix<TYPE>::Matrix(size_t number_rows, size_t number_columns,
+                     std::vector<size_t> TYPE_constructor_param,
+                     MemoryAllocator allocator, const IndexStrategy strategy)
+    : _number_rows(number_rows),
+      _number_columns(number_columns),
+      _allocator(allocator),
+      _dev_idx(strategy->index_strategy_three) {
+  /*
+   * Allocate the flattened c-array which contains the top-level
+   * TYPE pointers.
+   */
+  size_t flat_c_array_dimensions{_number_rows * _number_columns};
+  size_t size_bytes{flat_c_array_dimensions * sizeof(TYPE*)};
+  _allocator.allocate(reinterpret_cast<void**>(&_flat_c_array), size_bytes);
+
+  /*
+   * Check the TYPE_constructor_param vector to see if it has enough
+   * entries to fully initialize the matrix.
+   */
+  assert(TYPE_constructor_param.size() == flat_c_array_dimensions);
+
+  /*
+   * Iterate through the flattened c-array and initialize each pointer
+   * with a valid TYPE.
+   */
+  for (size_t i = 0; i < flat_c_array_dimensions; i++) {
+    _allocator.allocate(reinterpret_cast<void**>(&_flat_c_array[i]),
+                        sizeof(TYPE));
+
+    /*
+     * If type traits match, construct the TYPE with placement new.
+     */
+    if constexpr (one_generic_constructor_parameter<TYPE>::value &&
+                  requires_internal_allocator<TYPE>::value) {
+      new (_flat_c_array[i])
+          TYPE(TYPE_constructor_param[i], allocator, strategy);
+    } else if constexpr (one_generic_constructor_parameter<TYPE>::value) {
+      new (_flat_c_array[i]) TYPE(TYPE_constructor_param[i], strategy);
+    }
+  }
+}
+
+template <typename TYPE>
+Matrix<TYPE>::~Matrix() {
+  if (_flat_c_array) {
+    /*
+     * Free internal TYPE instances.
+     */
+    size_t flat_c_array_dimensions = _number_rows * _number_columns;
+    for (size_t i = 0; i < flat_c_array_dimensions; i++) {
+      if (_flat_c_array[i]) {
+        _flat_c_array[i]->~TYPE();
+        _allocator.deallocate(_flat_c_array[i]);
+      }
+    }
+
+    /*
+     * Free top-level flat c-array.
+     */
+    _allocator.deallocate(_flat_c_array);
+  }
+}
+
+template <typename TYPE>
+__host__ __device__ TYPE* Matrix<TYPE>::access(size_t row_index,
+                                               size_t col_index) {
+  assert(row_index < _number_rows);
+  assert(col_index < _number_columns);
+  size_t offset = row_index * _number_columns + col_index;
+  return _flat_c_array[offset];
+}
+
+template <typename TYPE>
+__device__ TYPE* Matrix<TYPE>::access(size_t col_index) {
+  auto row_index = _dev_idx.start();
+  assert(row_index < _number_rows);
+  assert(col_index < _number_columns);
+  size_t offset = row_index * _number_columns + col_index;
+  return _flat_c_array[offset];
+}
+
+template <typename TYPE>
+__device__ TYPE* Matrix<TYPE>::access() {
+  auto row_index = 0;
+  auto col_index = _dev_idx.start();
+  assert(row_index < _number_rows);
+  assert(col_index < _number_columns);
+  size_t offset = row_index * _number_columns + col_index;
+  return _flat_c_array[offset];
+}
+
+template <typename TYPE>
+__host__ __device__ size_t Matrix<TYPE>::rows() const {
+  return _number_rows;
+}
+
+template <typename TYPE>
+__host__ __device__ size_t Matrix<TYPE>::columns() const {
+  return _number_columns;
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_MATRIX_IMPL_HPP_

src/containers/share_strategy.cpp

@@ -0,0 +1,230 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "share_strategy.hpp"
+
+namespace rocshmem {
+
+/*
+ * @brief The warp (wave-front) size.
+ */
+static constexpr uint64_t WARP_SIZE = __AMDGCN_WAVEFRONT_SIZE;
+
+__device__ uint64_t Block::lane_id() {
+  /*
+   * amd_detail/device_functions.h::__lane_id()
+   */
+  return __builtin_amdgcn_mbcnt_hi(-1, __builtin_amdgcn_mbcnt_lo(-1, 0));
+}
+
+__device__ uint64_t Block::number_active_lanes() {
+  /*
+   * The __ballot(1) built-in instruction conducts an active lane roll
+   * call storing the roll call result in a bit vector. Using its index
+   * within the warp, each active lane contributes a '1' to the bit vector
+   * Inactive lanes cannot contribute to the bit vector so their lane
+   * values contain a '0'.
+   *
+   * For example, assume lanes 2, 5, and 62 active out of 64 lanes (0...63).
+   * The result of the ballot would be the following:
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 1 0 0 1 0 ... 1 0]
+   *
+   * The __ballot(1) result is fed to __popcll.
+   *
+   * The __popcll instruction conducts a population count; it checks
+   * the number of '1' values in a bit vector.
+   *
+   * Using the previous example, the result would be a 3 (since indices
+   * 2, 5, and 62) contain a '1'.
+   */
+  return __popcll(__ballot(1));
+}
+
+__device__ bool Block::is_lowest_active_lane() {
+  return active_logical_lane_id() == 0;
+}
+
+__device__ uint64_t Block::lowest_active_lane() {
+  /*
+   * The __ballot(1) built-in instruction conducts an active lane roll
+   * call storing the roll call result in a bit vector. Using its index
+   * within the warp, each active lane contributes a '1' to the bit vector
+   * Inactive lanes cannot contribute to the bit vector so their lane
+   * values contain a '0'.
+   *
+   * For example, assume lanes 2, 5, and 62 active out of 64 lanes (0...63).
+   * The result of the ballot would be the following:
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 1 0 0 1 0 ... 1 0]
+   *
+   * The __ballot(1) result is fed to __ffsll.
+   *
+   * The __ffsll instruction finds the index of the least significant
+   * bit set to '1' (in the input bit vector).
+   *
+   * In the previous example, the return result would be 2 (since index
+   * 2 is the first bit with a '1').
+   *
+   * The '- 1' at the end is necessary because the return index is 1-based
+   * instead of 0-based ([1...64] vs [0...63]).
+   */
+  return __ffsll(__ballot(1)) - 1;
+}
+
+__device__ uint64_t Block::active_logical_lane_id() {
+  /*
+   * The __ballot(1) built-in instruction conducts an active lane roll
+   * call storing the roll call result in a bit vector. Using its index
+   * within the warp, each active lane contributes a '1' to the bit vector
+   * Inactive lanes cannot contribute to the bit vector so their lane
+   * values contain a '0'.
+   *
+   * For example, assume lanes 2, 5, and 62 active out of 64 lanes (0...63).
+   * The result of the ballot would be the following:
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 1 0 0 1 0 ... 1 0]
+   */
+  uint64_t ballot = __ballot(1);
+
+  /*
+   * The physical_lane_id is the warp lane index of the thread executing
+   * this code. The word 'physical' here denotes that the lane_id will
+   * be the actual hardware lane_id.
+   */
+  uint64_t my_physical_lane_id = lane_id();
+
+  /*
+   * Create a full bitset for subsequent operations.
+   *
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [1 1 1 1 1 1 1 ... 1 1]
+   */
+  uint64_t all_ones_mask = -1;
+
+  /*
+   * Left-shift to zero-out the mask elements up to our lane_id.
+   *
+   * As an example, assume our lane_id is '5':
+   *
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 0 0 0 1 1 ... 1 1]
+   */
+  uint64_t lane_mask = all_ones_mask << my_physical_lane_id;
+
+  /*
+   * Invert the lane_mask.
+   *
+   * Continue with lane_id '5' example:
+   *
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [1 1 1 1 1 0 0 ... 0 0]
+   */
+  uint64_t inverted_mask = ~lane_mask;
+
+  /*
+   * Bit-wise And the inverted_mask and the ballot.
+   *
+   * The result contains a bitset with all active_lanes preceding this
+   * thread in the ballot (all active threads with lower lane_ids).
+   *
+   * Continue with lane_id '5' example:
+   *
+   * ballot
+   * ------------------------------
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 1 0 0 1 0 ... 1 0]
+   *
+   * inverted_mask
+   * ------------------------------
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [1 1 1 1 1 0 0 ... 0 0]
+   *
+   * lower_active_lanes
+   * ------------------------------
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 1 0 0 0 0 ... 0 0]
+   */
+  uint64_t lower_active_lanes = ballot & inverted_mask;
+
+  /*
+   * Conduct a population count on lower_active_lanes.
+   *
+   * The result gives an index into our logical_lane_id.
+   *
+   * Continue with lane_id '5' example:
+   *
+   * lower_active_lanes
+   * ------------------------------
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   * value: [0 0 1 0 0 0 0 ... 0 0]
+   *
+   * my_logical_lane_id
+   * ------------------------------
+   *        [- - X - - - - ... - -] <- population_count = 1
+   *
+   * index:  0 0 0 0 0 0 0 ... 6 6
+   *         0 1 2 3 4 5 6 ... 2 3
+   *        [- - 0 - - 1 - ... 2 -] <- my_logical_lane_id = 1
+   */
+  uint64_t my_logical_lane_id = __popcll(lower_active_lanes);
+
+  return my_logical_lane_id;
+}
+
+__device__ uint64_t Block::broadcast_up(uint64_t fetch_value) {
+  for (unsigned i = 0; i < WARP_SIZE; i++) {
+    uint64_t temp = __shfl_up(fetch_value, i);
+    if (temp) {
+      fetch_value = temp;
+    }
+  }
+  return fetch_value;
+}
+
+__device__ void ShareStrategy::syncthreads() const {
+  switch (_sse) {
+    case ShareStrategyEnum::PRIVATE:
+      return;
+    case ShareStrategyEnum::BLOCK:
+      __syncthreads();
+      return;
+    case ShareStrategyEnum::DEVICE:
+      assert(false);
+      return;
+    case ShareStrategyEnum::UNUSED:
+      assert(false);
+      return;
+  }
+}
+
+}  // namespace rocshmem

src/containers/share_strategy.hpp

@@ -0,0 +1,216 @@
+/******************************************************************************
+ * 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_CONTAINERS_SHARE_STRATEGY_HPP_
+#define LIBRARY_SRC_CONTAINERS_SHARE_STRATEGY_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include "index_strategy.hpp"
+
+namespace rocshmem {
+
+class Global {
+ public:
+  /**
+   * @brief
+   *
+   * @param out_ptr
+   *
+   * @return
+   */
+  template <typename T>
+  __host__ T fetch_incr(T *out_ptr) {}
+};
+
+class Grid {
+ public:
+  /**
+   * @brief
+   *
+   * @param out_ptr
+   *
+   * @return
+   */
+  template <typename T>
+  __device__ T fetch_incr(T *out_ptr) {}
+};
+
+class Block {
+ private:
+  friend class BitwiseDeviceMethods;
+
+ public:
+  /**
+   * @brief
+   *
+   * @param out_ptr
+   *
+   * @return
+   */
+  template <typename T>
+  __device__ T fetch_incr(T *out_ptr) {
+    T fetch_value = 0;
+
+    auto num_active_lanes = number_active_lanes();
+
+    if (is_lowest_active_lane()) {
+      fetch_value = atomicAdd(out_ptr, num_active_lanes);
+    }
+
+    fetch_value = broadcast_up(fetch_value);
+
+    return fetch_value + active_logical_lane_id();
+  }
+
+ private:
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ uint64_t lane_id();
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ uint64_t number_active_lanes();
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ bool is_lowest_active_lane();
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ uint64_t lowest_active_lane();
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ uint64_t active_logical_lane_id();
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ uint64_t broadcast_up(uint64_t fetch_value);
+};
+
+class Private {
+ public:
+  /**
+   * @brief
+   *
+   * @param out_ptr
+   *
+   * @return
+   */
+  template <typename T>
+  __device__ T fetch_incr(T *out_ptr) {
+    auto orig_value = *out_ptr;
+    *out_ptr = orig_value + 1;
+    return orig_value;
+  }
+};
+
+enum class ShareStrategyEnum { PRIVATE = 0, BLOCK = 1, DEVICE = 2, UNUSED = 3 };
+
+class ShareStrategy {
+ public:
+  __host__ __device__ ShareStrategy() = default;
+
+  __host__ __device__ ShareStrategy(ShareStrategyEnum sse) : _sse(sse) {}
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ void syncthreads() const;
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  template <typename T>
+  __device__ T fetch_incr(T *out_ptr, size_t my_pe, size_t num_pes) {
+    T value = 0;
+    switch (_sse) {
+      case ShareStrategyEnum::PRIVATE:
+        return _private.fetch_incr(out_ptr);
+      case ShareStrategyEnum::BLOCK:
+        for (size_t i = 0; i < num_pes; i++) {
+          if (i == my_pe) {
+            value = _block.fetch_incr(out_ptr);
+          }
+        }
+        return value;
+      case ShareStrategyEnum::DEVICE:
+        assert(false);
+        return 0;
+      case ShareStrategyEnum::UNUSED:
+        assert(false);
+        return 0;
+    }
+  }
+
+  /**
+   * @brief
+   *
+   * @return
+   */
+  __device__ bool uses_designated_send_thread() {
+    switch (_sse) {
+      case ShareStrategyEnum::PRIVATE:
+        return false;
+      case ShareStrategyEnum::BLOCK:
+        return true;
+      case ShareStrategyEnum::DEVICE:
+        assert(false);
+        return false;
+      case ShareStrategyEnum::UNUSED:
+        assert(false);
+        return false;
+    }
+  }
+
+ private:
+  Private _private{};
+  Block _block{};
+  ShareStrategyEnum _sse{ShareStrategyEnum::UNUSED};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_SHARE_STRATEGY_HPP_

src/containers/strategies.cpp

@@ -0,0 +1,37 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "strategies.hpp"
+
+namespace rocshmem {
+
+DefaultObjectStrategy* DefaultObjectStrategy::_instance = nullptr;
+
+DefaultObjectStrategy* DefaultObjectStrategy::instance() {
+  if (!_instance) {
+    _instance = new DefaultObjectStrategy();
+    return _instance;
+  }
+  return _instance;
+}
+
+}  // namespace rocshmem

src/containers/strategies.hpp

@@ -0,0 +1,63 @@
+/******************************************************************************
+ * 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_CONTAINERS_STRATEGIES_HPP_
+#define LIBRARY_SRC_CONTAINERS_STRATEGIES_HPP_
+
+#include "index_strategy.hpp"
+#include "share_strategy.hpp"
+
+namespace rocshmem {
+
+struct ObjectStrategy {
+  ShareStrategy share_strategy{};
+  IndexStrategy index_strategy_one{};
+  IndexStrategy index_strategy_two{};
+  IndexStrategy index_strategy_three{};
+  IndexStrategy index_strategy_four{};
+};
+
+class DefaultObjectStrategy {
+ private:
+  static DefaultObjectStrategy* _instance;
+
+  DefaultObjectStrategy() {
+    _os.share_strategy = ShareStrategy(ShareStrategyEnum::PRIVATE);
+    _os.index_strategy_one = IndexStrategy(IndexStrategyEnum::TDBD);
+    _os.index_strategy_two = IndexStrategy(IndexStrategyEnum::TDBD);
+    _os.index_strategy_three = IndexStrategy(IndexStrategyEnum::TDBD);
+    _os.index_strategy_four = IndexStrategy(IndexStrategyEnum::TDBD);
+  }
+
+  ObjectStrategy _os{};
+
+ public:
+  static DefaultObjectStrategy* instance();
+
+  const ObjectStrategy* get() { return &_os; }
+};
+
+struct Strategies {};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTAINERS_STRATEGIES_HPP_

src/context.hpp

@@ -0,0 +1,468 @@
+/******************************************************************************
+ * 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_CONTEXT_HPP_
+#define LIBRARY_SRC_CONTEXT_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include "backend_type.hpp"
+#include "fence_policy.hpp"
+#include "host/host.hpp"
+#include "ipc_policy.hpp"
+#include "stats.hpp"
+#include "sync/spin_ebo_block_mutex.hpp"
+#include "wf_coal_policy.hpp"
+
+namespace rocshmem {
+
+class Backend;
+
+/**
+ * @file context.hpp
+ * @brief Context class corresponds directly to an OpenSHMEM context.
+ *
+ * GPUs perform networking operations on a context that is created by the
+ * application programmer or a "default context" managed by the runtime.
+ *
+ * Contexts can be allocated in shared memory, in which case they are private
+ * to the creating workgroup, or they can be allocated in global memory, in
+ * which case they are shareable across workgroups.
+ *
+ * This is an 'abstract' class, as much as there is such a thing on a GPU.
+ * It uses 'type' to dispatch to a derived class for most of the interesting
+ * behavior.
+ */
+class Context {
+ public:
+  __host__ Context(Backend* handle, bool shareable);
+
+  __device__ Context(Backend* handle, bool shareable);
+
+  /*
+   * Dispatch functions to get runtime polymorphism without 'virtual' or
+   * function pointers. Each one of these guys will use 'type' to
+   * static_cast themselves and dispatch to the appropriate derived class.
+   * It's basically doing part of what the 'virtual' keyword does, so when
+   * we get that working in ROCm it will be super easy to adapt to it by
+   * just removing the dispatch implementations.
+   *
+   * No comments for these guys since its basically the same as in the
+   * roc_shmem.hpp public header.
+   */
+
+  /**************************************************************************
+   ***************************** DEVICE METHODS *****************************
+   *************************************************************************/
+  template <typename T>
+  __device__ void wait_until(T* ptr, roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __device__ void wait_until_all(T* ptr, size_t nelems,
+                                 const int *status,
+                                 roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __device__ size_t wait_until_any(T* ptr, size_t nelems,
+                                   const int *status,
+                                   roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __device__ size_t wait_until_some(T* ptr, size_t nelems,
+                                    size_t* indices,
+                                    const int *status,
+                                    roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __device__ void wait_until_all_vector(T* ptr, size_t nelems,
+                                        const int *status,
+                                        roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __device__ size_t wait_until_any_vector(T* ptr, size_t nelems,
+                                          const int *status,
+                                          roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __device__ size_t wait_until_some_vector(T* ptr, size_t nelems,
+                                           size_t* indices,
+                                           const int *status,
+                                           roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __device__ int test(T* ptr, roc_shmem_cmps cmp, T val);
+
+  __device__ void threadfence_system();
+
+  __device__ void ctx_create();
+
+  __device__ void ctx_destroy();
+
+  __device__ void putmem(void* dest, const void* source, size_t nelems, int pe);
+
+  __device__ void getmem(void* dest, const void* source, size_t nelems, int pe);
+
+  __device__ void putmem_nbi(void* dest, const void* source, size_t nelems,
+                             int pe);
+
+  __device__ void getmem_nbi(void* dest, const void* source, size_t size,
+                             int pe);
+
+  __device__ void fence();
+
+  __device__ void fence(int pe);
+
+  __device__ void quiet();
+
+  __device__ void* shmem_ptr(const void* dest, int pe);
+
+  __device__ void barrier_all();
+
+  __device__ void sync_all();
+
+  __device__ void sync(roc_shmem_team_t team);
+
+  template <typename T>
+  __device__ T amo_fetch(void* dst, T value, T cond, int pe, uint8_t atomic_op);
+
+  template <typename T>
+  __device__ void amo_add(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_set(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_swap(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_and(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_and(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_or(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_or(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_xor(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_xor(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_cas(void* dst, T value, T cond, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_add(void* dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_cas(void* dst, T value, T cond, int pe);
+
+  template <typename T>
+  __device__ void p(T* dest, T value, int pe);
+
+  template <typename T>
+  __device__ T g(T* source, int pe);
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __device__ void to_all(T* dest, const T* source, int nreduce, int PE_start,
+                         int logPE_stride, int PE_size, T* pWrk,
+                         long* pSync);  // NOLINT(runtime/int)
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __device__ void to_all(roc_shmem_team_t team, T* dest, const T* source,
+                         int nreduce);
+
+  template <typename T>
+  __device__ void put(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_nbi(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_nbi(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void alltoall(roc_shmem_team_t team, T* dest, const T* source,
+                           int nelems);
+
+  template <typename T>
+  __device__ void fcollect(roc_shmem_team_t team, T* dest, const T* source,
+                           int nelems);
+
+  template <typename T>
+  __device__ void broadcast(roc_shmem_team_t team, T* dest, const T* source,
+                            int nelems, int pe_root);
+
+  template <typename T>
+  __device__ void broadcast(T* dest, const T* source, int nelems, int pe_root,
+                            int pe_start, int log_pe_stride, int pe_size,
+                            long* p_sync);  // NOLINT(runtime/int)
+
+  __device__ void putmem_wg(void* dest, const void* source, size_t nelems,
+                            int pe);
+
+  __device__ void getmem_wg(void* dest, const void* source, size_t nelems,
+                            int pe);
+
+  __device__ void putmem_nbi_wg(void* dest, const void* source, size_t nelems,
+                                int pe);
+
+  __device__ void getmem_nbi_wg(void* dest, const void* source, size_t size,
+                                int pe);
+
+  __device__ void putmem_wave(void* dest, const void* source, size_t nelems,
+                              int pe);
+
+  __device__ void getmem_wave(void* dest, const void* source, size_t nelems,
+                              int pe);
+
+  __device__ void putmem_nbi_wave(void* dest, const void* source, size_t nelems,
+                                  int pe);
+
+  __device__ void getmem_nbi_wave(void* dest, const void* source, size_t size,
+                                  int pe);
+
+  template <typename T>
+  __device__ void put_wg(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_nbi_wg(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_wg(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_nbi_wg(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_wave(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_nbi_wave(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_wave(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_nbi_wave(T* dest, const T* source, size_t nelems, int pe);
+
+  /**************************************************************************
+   ****************************** HOST METHODS ******************************
+   *************************************************************************/
+  template <typename T>
+  __host__ void p(T* dest, T value, int pe);
+
+  template <typename T>
+  __host__ T g(const T* source, int pe);
+
+  template <typename T>
+  __host__ void put(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __host__ void get(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __host__ void put_nbi(T* dest, const T* source, size_t nelems, int pe);
+
+  template <typename T>
+  __host__ void get_nbi(T* dest, const T* source, size_t nelems, int pe);
+
+  __host__ void putmem(void* dest, const void* source, size_t nelems, int pe);
+
+  __host__ void getmem(void* dest, const void* source, size_t nelems, int pe);
+
+  __host__ void putmem_nbi(void* dest, const void* source, size_t nelems,
+                           int pe);
+
+  __host__ void getmem_nbi(void* dest, const void* source, size_t size, int pe);
+
+  template <typename T>
+  __host__ void amo_add(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ void amo_set(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ T amo_swap(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_and(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ void amo_and(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_or(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ void amo_or(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_xor(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ void amo_xor(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ void amo_cas(void* dst, T value, T cond, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_add(void* dst, T value, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_cas(void* dst, T value, T cond, int pe);
+
+  __host__ void fence();
+
+  __host__ void quiet();
+
+  __host__ void barrier_all();
+
+  __host__ void sync_all();
+
+  template <typename T>
+  __host__ void broadcast(T* dest, const T* source, int nelems, int pe_root,
+                          int pe_start, int log_pe_stride, int pe_size,
+                          long* p_sync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __host__ void broadcast(roc_shmem_team_t team, T* dest, const T* source,
+                          int nelems, int pe_root);
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all(T* dest, const T* source, int nreduce, int PE_start,
+                       int logPE_stride, int PE_size, T* pWrk,
+                       long* pSync);  // NOLINT(runtime/int)
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all(roc_shmem_team_t team, T* dest, const T* source,
+                       int nreduce);
+
+  template <typename T>
+  __host__ void wait_until(T* ptr, roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ void wait_until_all(T* ptr, size_t nelems,
+                               const int *status,
+                               roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ size_t wait_until_any(T* ptr, size_t nelems,
+                                 const int *status,
+                                 roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ size_t wait_until_some(T* ptr, size_t nelems,
+                                  size_t* indices,
+                                  const int *status,
+                                  roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ void wait_until_all_vector(T* ptr, size_t nelems,
+                                      const int *status,
+                                      roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __host__ size_t wait_until_any_vector(T* ptr, size_t nelems,
+                                        const int *status,
+                                        roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __host__ size_t wait_until_some_vector(T* ptr, size_t nelems,
+                                         size_t* indices,
+                                         const int *status,
+                                         roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __host__ int test(T* ptr, roc_shmem_cmps cmp, T val);
+
+ public:
+  /**
+   * @brief Set the fence policy using a runtime option
+   *
+   * @param[in] options interpreted as a bitfield using bitwise operations
+   */
+  __device__ void setFence(long options) { fence_ = Fence(options); }
+
+  /**************************************************************************
+   ***************************** PUBLIC MEMBERS *****************************
+   *************************************************************************/
+  /**
+   * @brief Duplicated local copy of backend's num_pes
+   */
+  int num_pes{0};
+
+  /**
+   * @brief Duplicated local copy of backend's my_pe
+   */
+  int my_pe{-1};
+
+  /**
+   * @brief Stats common to all types of device contexts.
+   */
+  ROCStats ctxStats{};
+
+  /**
+   * @brief Stats common to all types of host contexts.
+   */
+  ROCHostStats ctxHostStats{};
+
+  /**
+   * @brief Lock to prevent data races on shared data
+   */
+  SpinEBOBlockMutex dev_mtx_{};
+
+ protected:
+  /**************************************************************************
+   ***************************** POLICY MEMBERS *****************************
+   *************************************************************************/
+
+  /**
+   * @brief Coalesce policy for 'multi' configuration builds
+   */
+  WavefrontCoalescer wf_coal_{};
+
+  /**
+   * @brief Controls fence behavior in device code
+   */
+  Fence fence_{};
+
+ public:
+  /**
+   * @brief Inter-Process Communication (IPC) interface for context class
+   *
+   * This member is an interface to allow intra-node interprocess
+   * communication through shared memory.
+   */
+  IpcImpl ipcImpl_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTEXT_HPP_

src/context_device.cpp

@@ -0,0 +1,252 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "backend_bc.hpp"
+#include "context_incl.hpp"
+#include "util.hpp"
+
+namespace rocshmem {
+
+__device__ Context::Context(Backend* handle, bool shareable)
+    : num_pes(handle->getNumPEs()),
+      my_pe(handle->getMyPE()),
+      fence_(shareable),
+      dev_mtx_(shareable) {
+  /*
+   * Device-side context constructor is a work-group collective, so make
+   * sure all the members have their default values before returning.
+   *
+   * Each thread is essentially initializing the same thing right over the
+   * top of each other for all the default values in context.hh (and the
+   * initializer list). It's not incorrect, but it is weird and probably
+   * wasteful.
+   *
+   * TODO: Might consider refactoring so that constructor is always called
+   * from a single thread, and the parallel portion of initialization can be
+   * a separate function. This requires reworking all the derived classes
+   * since their constructors actually make use of all the threads to boost
+   * performance.
+   */
+  __syncthreads();
+}
+
+/******************************************************************************
+ ********************** CONTEXT DISPATCH IMPLEMENTATIONS **********************
+ *****************************************************************************/
+
+__device__ void Context::threadfence_system() {
+  DISPATCH(threadfence_system());
+}
+
+__device__ void Context::ctx_create() {
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_CREATE);
+  }
+
+  DISPATCH(ctx_create());
+}
+
+__device__ void Context::ctx_destroy() {
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_FINALIZE);
+    device_backend_proxy->globalStats.accumulateStats(ctxStats);
+  }
+
+  DISPATCH(ctx_destroy());
+}
+
+__device__ void Context::putmem(void* dest, const void* source, size_t nelems,
+                                int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT);
+
+  DISPATCH(putmem(dest, source, nelems, pe));
+}
+
+__device__ void Context::getmem(void* dest, const void* source, size_t nelems,
+                                int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET);
+
+  DISPATCH(getmem(dest, source, nelems, pe));
+}
+
+__device__ void Context::putmem_nbi(void* dest, const void* source,
+                                    size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_NBI);
+
+  DISPATCH(putmem_nbi(dest, source, nelems, pe));
+}
+
+__device__ void Context::getmem_nbi(void* dest, const void* source, size_t size,
+                                    int pe) {
+  if (size == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_NBI);
+
+  DISPATCH(getmem_nbi(dest, source, size, pe));
+}
+
+__device__ void Context::fence() {
+  ctxStats.incStat(NUM_FENCE);
+
+  DISPATCH(fence());
+}
+
+__device__ void Context::fence(int pe) {
+  ctxStats.incStat(NUM_FENCE);
+
+  DISPATCH(fence(pe));
+}
+
+__device__ void Context::quiet() {
+  ctxStats.incStat(NUM_QUIET);
+
+  DISPATCH(quiet());
+}
+
+__device__ void* Context::shmem_ptr(const void* dest, int pe) {
+  ctxStats.incStat(NUM_SHMEM_PTR);
+
+  DISPATCH_RET_PTR(shmem_ptr(dest, pe));
+}
+
+__device__ void Context::barrier_all() {
+  ctxStats.incStat(NUM_BARRIER_ALL);
+
+  DISPATCH(barrier_all());
+}
+
+__device__ void Context::sync_all() {
+  ctxStats.incStat(NUM_SYNC_ALL);
+
+  DISPATCH(sync_all());
+}
+
+__device__ void Context::sync(roc_shmem_team_t team) {
+  ctxStats.incStat(NUM_SYNC_ALL);
+
+  DISPATCH(sync(team));
+}
+
+__device__ void Context::putmem_wg(void* dest, const void* source,
+                                   size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_WG);
+
+  DISPATCH(putmem_wg(dest, source, nelems, pe));
+}
+
+__device__ void Context::getmem_wg(void* dest, const void* source,
+                                   size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_WG);
+
+  DISPATCH(getmem_wg(dest, source, nelems, pe));
+}
+
+__device__ void Context::putmem_nbi_wg(void* dest, const void* source,
+                                       size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_NBI_WG);
+
+  DISPATCH(putmem_nbi_wg(dest, source, nelems, pe));
+}
+
+__device__ void Context::getmem_nbi_wg(void* dest, const void* source,
+                                       size_t size, int pe) {
+  if (size == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_NBI_WG);
+
+  DISPATCH(getmem_nbi_wg(dest, source, size, pe));
+}
+
+__device__ void Context::putmem_wave(void* dest, const void* source,
+                                     size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_WAVE);
+
+  DISPATCH(putmem_wave(dest, source, nelems, pe));
+}
+
+__device__ void Context::getmem_wave(void* dest, const void* source,
+                                     size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_WAVE);
+
+  DISPATCH(getmem_wave(dest, source, nelems, pe));
+}
+
+__device__ void Context::putmem_nbi_wave(void* dest, const void* source,
+                                         size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_NBI_WAVE);
+
+  DISPATCH(putmem_nbi_wave(dest, source, nelems, pe));
+}
+
+__device__ void Context::getmem_nbi_wave(void* dest, const void* source,
+                                         size_t size, int pe) {
+  if (size == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_NBI_WAVE);
+
+  DISPATCH(getmem_nbi_wave(dest, source, size, pe));
+}
+
+}  // namespace rocshmem

src/context_host.cpp

@@ -0,0 +1,107 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "backend_bc.hpp"
+#include "context_incl.hpp"
+
+namespace rocshmem {
+
+__host__ Context::Context(Backend* handle, bool shareable)
+    : num_pes(handle->getNumPEs()),
+      my_pe(handle->getMyPE()),
+      fence_(shareable),
+      dev_mtx_(shareable) {}
+
+/******************************************************************************
+ ********************** CONTEXT DISPATCH IMPLEMENTATIONS **********************
+ *****************************************************************************/
+
+__host__ void Context::putmem(void* dest, const void* source, size_t nelems,
+                              int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_PUT);
+
+  HOST_DISPATCH(putmem(dest, source, nelems, pe));
+}
+
+__host__ void Context::getmem(void* dest, const void* source, size_t nelems,
+                              int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_GET);
+
+  HOST_DISPATCH(getmem(dest, source, nelems, pe));
+}
+
+__host__ void Context::putmem_nbi(void* dest, const void* source, size_t nelems,
+                                  int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_PUT_NBI);
+
+  HOST_DISPATCH(putmem_nbi(dest, source, nelems, pe));
+}
+
+__host__ void Context::getmem_nbi(void* dest, const void* source, size_t nelems,
+                                  int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_GET_NBI);
+
+  HOST_DISPATCH(getmem_nbi(dest, source, nelems, pe));
+}
+
+__host__ void Context::fence() {
+  ctxHostStats.incStat(NUM_HOST_FENCE);
+
+  HOST_DISPATCH(fence());
+}
+
+__host__ void Context::quiet() {
+  ctxHostStats.incStat(NUM_HOST_QUIET);
+
+  HOST_DISPATCH(quiet());
+}
+
+__host__ void Context::sync_all() {
+  ctxHostStats.incStat(NUM_HOST_SYNC_ALL);
+
+  HOST_DISPATCH(sync_all());
+}
+
+__host__ void Context::barrier_all() {
+  ctxHostStats.incStat(NUM_HOST_BARRIER_ALL);
+
+  HOST_DISPATCH(barrier_all());
+}
+
+}  // namespace rocshmem

src/context_incl.hpp

@@ -0,0 +1,37 @@
+/******************************************************************************
+ * 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_CONTEXT_INCL_HPP_
+#define LIBRARY_SRC_CONTEXT_INCL_HPP_
+
+#include "context.hpp"
+#include "context_tmpl_device.hpp"
+#include "context_tmpl_host.hpp"
+#ifdef USE_GPU_IB
+#include "gpu_ib/context_ib_device.hpp"
+#include "gpu_ib/context_ib_host.hpp"
+#else
+#include "reverse_offload/context_ro_device.hpp"
+#include "reverse_offload/context_ro_host.hpp"
+#endif
+
+#endif  // LIBRARY_SRC_CONTEXT_INCL_HPP_

src/context_tmpl_device.hpp

@@ -0,0 +1,559 @@
+/******************************************************************************
+ * 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_CONTEXT_TMPL_DEVICE_HPP_
+#define LIBRARY_SRC_CONTEXT_TMPL_DEVICE_HPP_
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "backend_type.hpp"
+#ifdef USE_GPU_IB
+#include "gpu_ib/context_ib_device.hpp"
+#else
+#include "reverse_offload/context_ro_device.hpp"
+#endif
+
+namespace rocshmem {
+
+/*
+ * Context dispatch implementations for the template functions. Needs to
+ * be in a header and not cpp because it is a template.
+ */
+template <typename T>
+__device__ void Context::p(T *dest, T value, int pe) {
+  ctxStats.incStat(NUM_P);
+
+  /*
+   * TODO: Need to handle _p a bit differently for coalescing, since the
+   * owner of a coalesced message needs val from all absorbed messages.
+   */
+  DISPATCH(p(dest, value, pe));
+}
+
+template <typename T>
+__device__ T Context::g(T *source, int pe) {
+  ctxStats.incStat(NUM_G);
+
+  /*
+   * TODO: Need to handle _g a bit differently for coalescing, since the
+   * owner of a coalesced message needs val from all absorbed messages.
+   */
+  DISPATCH_RET(g(source, pe));
+}
+
+// The only way to get multi-arg templates to feed into a macro
+template <typename T, ROC_SHMEM_OP Op>
+__device__ void Context::to_all(T *dest, const T *source, int nreduce,
+                                int PE_start, int logPE_stride, int PE_size,
+                                T *pWrk,
+                                long *pSync) {  // NOLINT(runtime/int)
+  if (nreduce == 0) {
+    return;
+  }
+
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_TO_ALL);
+  }
+
+  DISPATCH(to_all<PAIR(T, Op)>(dest, source, nreduce, PE_start, logPE_stride,
+                               PE_size, pWrk, pSync));
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__device__ void Context::to_all(roc_shmem_team_t team, T *dest, const T *source,
+                                int nreduce) {
+  if (nreduce == 0) {
+    return;
+  }
+
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_TO_ALL);
+  }
+
+  DISPATCH(to_all<PAIR(T, Op)>(team, dest, source, nreduce));
+}
+
+template <typename T>
+__device__ void Context::put(T *dest, const T *source, size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT);
+
+  DISPATCH(put(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::put_nbi(T *dest, const T *source, size_t nelems,
+                                 int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_NBI);
+
+  DISPATCH(put_nbi(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::get(T *dest, const T *source, size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET);
+
+  DISPATCH(get(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::get_nbi(T *dest, const T *source, size_t nelems,
+                                 int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_NBI);
+
+  DISPATCH(get_nbi(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::alltoall(roc_shmem_team_t team, T *dest,
+                                  const T *source, int nelems) {
+  if (nelems == 0) {
+    return;
+  }
+
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_ALLTOALL);
+  }
+
+  DISPATCH(alltoall<T>(team, dest, source, nelems));
+}
+
+template <typename T>
+__device__ void Context::fcollect(roc_shmem_team_t team, T *dest,
+                                  const T *source, int nelems) {
+  if (nelems == 0) {
+    return;
+  }
+
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_FCOLLECT);
+  }
+
+  DISPATCH(fcollect<T>(team, dest, source, nelems));
+}
+
+template <typename T>
+__device__ void Context::broadcast(roc_shmem_team_t team, T *dest,
+                                   const T *source, int nelems, int pe_root) {
+  if (nelems == 0) {
+    return;
+  }
+
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_BROADCAST);
+  }
+
+  DISPATCH(broadcast<T>(team, dest, source, nelems, pe_root));
+}
+
+template <typename T>
+__device__ void Context::broadcast(T *dest, const T *source, int nelems,
+                                   int pe_root, int pe_start, int log_pe_stride,
+                                   int pe_size,
+                                   long *p_sync) {  // NOLINT(runtime/int)
+  if (nelems == 0) {
+    return;
+  }
+
+  if (is_thread_zero_in_block()) {
+    ctxStats.incStat(NUM_BROADCAST);
+  }
+
+  DISPATCH(broadcast<T>(dest, source, nelems, pe_root, pe_start, log_pe_stride,
+                        pe_size, p_sync));
+}
+
+template <typename T>
+__device__ __forceinline__ void Context::wait_until(T *ptr, roc_shmem_cmps cmp,
+                                                    T val) {
+  while (!test(ptr, cmp, val)) {
+  }
+}
+
+__device__ __forceinline__ size_t status_entry(size_t nelems,
+                                               const int *status) {
+  size_t i{0};
+  while (i < nelems) {
+    if (status[i] == 0) {
+      return i;
+    }
+    i++;
+  }
+  return i;
+}
+
+template <typename T>
+__device__ __forceinline__
+size_t Context::wait_until_any(T* ptr, size_t nelems,
+                               const int *status,
+                               roc_shmem_cmps cmp, T val) {
+  // zero nelems error condition
+  if (!nelems) {
+    return SIZE_MAX;
+  }
+
+  size_t pos{status_entry(nelems, status)};
+
+  // invalid (empty) status array error condition
+  if (pos == nelems) {
+    return SIZE_MAX;
+  }
+
+  while (true) {
+    for (size_t i{pos}; i < nelems; i++) {
+      // skip entries marked with non-zero status
+      if (status[i]) {
+        continue;
+      }
+      if (test(ptr + i, cmp, val)) {
+        return i;
+      }
+    }
+  }
+}
+
+template <typename T>
+__device__ __forceinline__
+void Context::wait_until_all(T* ptr, size_t nelems,
+                             const int *status,
+                             roc_shmem_cmps cmp, T val) {
+  // zero nelems error condition
+  if (!nelems) {
+    return;
+  }
+
+  size_t pos{status_entry(nelems, status)};
+
+  // invalid (empty) status array error condition
+  if (pos == nelems) {
+    return;
+  }
+
+  for (size_t i{pos}; i < nelems; i++) {
+    if (status[i]) {
+      continue;
+    }
+    while (!test(ptr + i, cmp, val)) {
+    }
+  }
+}
+
+template <typename T>
+__device__ __forceinline__
+size_t Context::wait_until_some(T* ptr, size_t nelems,
+                              size_t* indices,
+                              const int *status,
+                              roc_shmem_cmps cmp, T val) {
+  // zero nelems error condition
+  if (!nelems) {
+    return 0;
+  }
+
+  size_t pos{status_entry(nelems, status)};
+
+  // invalid (empty) status array error condition
+  if (pos == nelems) {
+    return 0;
+  }
+
+  bool done {false};
+  size_t ncompleted {0};
+  while (!done) {
+    for (size_t i{pos}; i < nelems; i++) {
+      // skip entries marked with non-zero status
+      if (status[i]) {
+        continue;
+      }
+      if (test(ptr + i, cmp, val)) {
+        done = true;
+        indices[ncompleted] = i;
+        ncompleted++;
+      }
+    }
+  }
+  return ncompleted;
+}
+
+template <typename T>
+__device__ __forceinline__
+void Context::wait_until_all_vector(T* ptr, size_t nelems,
+                                    const int *status,
+                                    roc_shmem_cmps cmp, T* vals) {
+  ;
+}
+
+template <typename T>
+__device__ __forceinline__
+size_t Context::wait_until_any_vector(T* ptr, size_t nelems,
+                                      const int *status,
+                                      roc_shmem_cmps cmp, T* vals) {
+  return 0;
+}
+
+template <typename T>
+__device__ __forceinline__
+size_t Context::wait_until_some_vector(T* ptr, size_t nelems,
+                                     size_t* indices,
+                                     const int *status,
+                                     roc_shmem_cmps cmp, T* vals) {
+  return 0;
+}
+
+template <typename T>
+__device__ __forceinline__ int Context::test(T *ptr, roc_shmem_cmps cmp,
+                                             T val) {
+  int ret = 0;
+  volatile T *vol_ptr = reinterpret_cast<T *>(ptr);
+  switch (cmp) {
+    case ROC_SHMEM_CMP_EQ:
+      if (uncached_load(vol_ptr) == val) {
+        ret = 1;
+      }
+      break;
+    case ROC_SHMEM_CMP_NE:
+      if (uncached_load(vol_ptr) != val) {
+        ret = 1;
+      }
+      break;
+    case ROC_SHMEM_CMP_GT:
+      if (uncached_load(vol_ptr) > val) {
+        ret = 1;
+      }
+      break;
+    case ROC_SHMEM_CMP_GE:
+      if (uncached_load(vol_ptr) >= val) {
+        ret = 1;
+      }
+      break;
+    case ROC_SHMEM_CMP_LT:
+      if (uncached_load(vol_ptr) < val) {
+        ret = 1;
+      }
+      break;
+    case ROC_SHMEM_CMP_LE:
+      if (uncached_load(vol_ptr) <= val) {
+        ret = 1;
+      }
+      break;
+    default:
+      break;
+  }
+  return ret;
+}
+
+template <typename T>
+__device__ void Context::put_wg(T *dest, const T *source, size_t nelems,
+                                int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_WG);
+
+  DISPATCH(put_wg(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::put_nbi_wg(T *dest, const T *source, size_t nelems,
+                                    int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_NBI_WG);
+
+  DISPATCH(put_nbi_wg(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::get_wg(T *dest, const T *source, size_t nelems,
+                                int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_WG);
+
+  DISPATCH(get_wg(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::get_nbi_wg(T *dest, const T *source, size_t nelems,
+                                    int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_NBI_WG);
+
+  DISPATCH(get_nbi_wg(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::put_wave(T *dest, const T *source, size_t nelems,
+                                  int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_WAVE);
+
+  DISPATCH(put_wave(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::put_nbi_wave(T *dest, const T *source, size_t nelems,
+                                      int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_PUT_NBI_WAVE);
+
+  DISPATCH(put_nbi_wave(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::get_wave(T *dest, const T *source, size_t nelems,
+                                  int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_WAVE);
+
+  DISPATCH(get_wave(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ void Context::get_nbi_wave(T *dest, const T *source, size_t nelems,
+                                      int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxStats.incStat(NUM_GET_NBI_WAVE);
+
+  DISPATCH(get_nbi_wave(dest, source, nelems, pe));
+}
+
+template <typename T>
+__device__ T Context::amo_fetch_add(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_FADD);
+
+  DISPATCH_RET(amo_fetch_add(dst, value, pe));
+}
+
+template <typename T>
+__device__ void Context::amo_add(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_ADD);
+
+  DISPATCH(amo_add(dst, value, pe));
+}
+
+template <typename T>
+__device__ void Context::amo_set(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_SET);
+
+  DISPATCH(amo_set(dst, value, pe));
+}
+
+template <typename T>
+__device__ T Context::amo_swap(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_SWAP);
+
+  DISPATCH_RET(amo_swap(dst, value, pe));
+}
+
+template <typename T>
+__device__ T Context::amo_fetch_and(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_FETCH_AND);
+
+  DISPATCH_RET(amo_fetch_and(dst, value, pe));
+}
+
+template <typename T>
+__device__ void Context::amo_and(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_AND);
+
+  DISPATCH(amo_and(dst, value, pe));
+}
+
+template <typename T>
+__device__ T Context::amo_fetch_or(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_FETCH_OR);
+
+  DISPATCH_RET(amo_fetch_or(dst, value, pe));
+}
+
+template <typename T>
+__device__ void Context::amo_or(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_OR);
+
+  DISPATCH(amo_or(dst, value, pe));
+}
+
+template <typename T>
+__device__ T Context::amo_fetch_xor(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_FETCH_XOR);
+
+  DISPATCH_RET(amo_fetch_xor(dst, value, pe));
+}
+
+template <typename T>
+__device__ void Context::amo_xor(void *dst, T value, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_XOR);
+
+  DISPATCH(amo_xor(dst, value, pe));
+}
+
+template <typename T>
+__device__ T Context::amo_fetch_cas(void *dst, T value, T cond, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_FCSWAP);
+
+  DISPATCH_RET(amo_fetch_cas(dst, value, cond, pe));
+}
+
+template <typename T>
+__device__ void Context::amo_cas(void *dst, T value, T cond, int pe) {
+  ctxStats.incStat(NUM_ATOMIC_CSWAP);
+
+  DISPATCH(amo_cas(dst, value, cond, pe));
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTEXT_TMPL_DEVICE_HPP_

src/context_tmpl_host.hpp

@@ -0,0 +1,306 @@
+/******************************************************************************
+ * 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_CONTEXT_TMPL_HOST_HPP_
+#define LIBRARY_SRC_CONTEXT_TMPL_HOST_HPP_
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "backend_type.hpp"
+#ifdef USE_GPU_IB
+#include "gpu_ib/context_ib_host.hpp"
+#else
+#include "reverse_offload/context_ro_host.hpp"
+#endif
+namespace rocshmem {
+
+template <typename T>
+__host__ void Context::p(T *dest, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_P);
+
+  HOST_DISPATCH(p(dest, value, pe));
+}
+
+template <typename T>
+__host__ T Context::g(const T *source, int pe) {
+  ctxHostStats.incStat(NUM_HOST_G);
+
+  HOST_DISPATCH_RET(g(source, pe));
+}
+
+template <typename T>
+__host__ void Context::put(T *dest, const T *source, size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_PUT);
+
+  HOST_DISPATCH(put(dest, source, nelems, pe));
+}
+
+template <typename T>
+__host__ void Context::get(T *dest, const T *source, size_t nelems, int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_GET);
+
+  HOST_DISPATCH(get(dest, source, nelems, pe));
+}
+
+template <typename T>
+__host__ void Context::put_nbi(T *dest, const T *source, size_t nelems,
+                               int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_PUT_NBI);
+
+  HOST_DISPATCH(put_nbi(dest, source, nelems, pe));
+}
+
+template <typename T>
+__host__ void Context::get_nbi(T *dest, const T *source, size_t nelems,
+                               int pe) {
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_GET_NBI);
+
+  HOST_DISPATCH(get_nbi(dest, source, nelems, pe));
+}
+
+template <typename T>
+__host__ T Context::amo_fetch_add(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_FADD);
+
+  HOST_DISPATCH_RET(amo_fetch_add(dst, value, pe));
+}
+
+template <typename T>
+__host__ void Context::amo_add(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_ADD);
+
+  HOST_DISPATCH(amo_add(dst, value, pe));
+}
+
+template <typename T>
+__host__ void Context::amo_set(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_ADD);
+
+  HOST_DISPATCH(amo_set(dst, value, pe));
+}
+
+template <typename T>
+__host__ T Context::amo_swap(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_ADD);
+
+  HOST_DISPATCH_RET(amo_swap(dst, value, pe));
+}
+
+template <typename T>
+__host__ T Context::amo_fetch_and(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_FETCH_AND);
+
+  HOST_DISPATCH_RET(amo_fetch_and(dst, value, pe));
+}
+
+template <typename T>
+__host__ void Context::amo_and(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_AND);
+
+  HOST_DISPATCH(amo_and(dst, value, pe));
+}
+
+template <typename T>
+__host__ T Context::amo_fetch_or(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_FETCH_OR);
+
+  HOST_DISPATCH_RET(amo_fetch_or(dst, value, pe));
+}
+
+template <typename T>
+__host__ void Context::amo_or(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_OR);
+
+  HOST_DISPATCH(amo_or(dst, value, pe));
+}
+
+template <typename T>
+__host__ T Context::amo_fetch_xor(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_FETCH_XOR);
+
+  HOST_DISPATCH_RET(amo_fetch_xor(dst, value, pe));
+}
+
+template <typename T>
+__host__ void Context::amo_xor(void *dst, T value, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_XOR);
+
+  HOST_DISPATCH(amo_xor(dst, value, pe));
+}
+
+template <typename T>
+__host__ T Context::amo_fetch_cas(void *dst, T value, T cond, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_FCSWAP);
+
+  HOST_DISPATCH_RET(amo_fetch_cas(dst, value, cond, pe));
+}
+
+template <typename T>
+__host__ void Context::amo_cas(void *dst, T value, T cond, int pe) {
+  ctxHostStats.incStat(NUM_HOST_ATOMIC_CSWAP);
+
+  HOST_DISPATCH(amo_cas(dst, value, cond, pe));
+}
+
+template <typename T>
+__host__ void Context::broadcast(T *dest, const T *source, int nelems,
+                                 int pe_root, int pe_start, int log_pe_stride,
+                                 int pe_size,
+                                 long *p_sync) {  // NOLINT(runtime/int)
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_BROADCAST);
+
+  HOST_DISPATCH(broadcast<T>(dest, source, nelems, pe_root, pe_start,
+                             log_pe_stride, pe_size, p_sync));
+}
+
+template <typename T>
+__host__ void Context::broadcast(roc_shmem_team_t team, T *dest,
+                                 const T *source, int nelems,
+                                 int pe_root) {  // NOLINT(runtime/int)
+  if (nelems == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_BROADCAST);
+
+  HOST_DISPATCH(broadcast<T>(team, dest, source, nelems, pe_root));
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void Context::to_all(T *dest, const T *source, int nreduce,
+                              int PE_start, int logPE_stride, int PE_size,
+                              T *pWrk,
+                              long *pSync) {  // NOLINT(runtime/int)
+  if (nreduce == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_TO_ALL);
+
+  HOST_DISPATCH(to_all<PAIR(T, Op)>(dest, source, nreduce, PE_start,
+                                    logPE_stride, PE_size, pWrk, pSync));
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void Context::to_all(roc_shmem_team_t team, T *dest, const T *source,
+                              int nreduce) {  // NOLINT(runtime/int)
+  if (nreduce == 0) {
+    return;
+  }
+
+  ctxHostStats.incStat(NUM_HOST_TO_ALL);
+
+  HOST_DISPATCH(to_all<PAIR(T, Op)>(team, dest, source, nreduce));
+}
+
+template <typename T>
+__host__ void Context::wait_until(T *ptr, roc_shmem_cmps cmp, T val) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL);
+
+  HOST_DISPATCH(wait_until<T>(ptr, cmp, val));
+}
+
+template <typename T>
+__host__ size_t Context::wait_until_any(T* ptr, size_t nelems,
+                                        const int* status,
+                                        roc_shmem_cmps cmp, T val) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL_ANY);
+
+  return HOST_DISPATCH(wait_until_any<T>(ptr, nelems, status, cmp, val));
+}
+
+template <typename T>
+__host__ void Context::wait_until_all(T* ptr, size_t nelems,
+                                      const int* status,
+                                      roc_shmem_cmps cmp, T val) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL_ALL);
+
+  HOST_DISPATCH(wait_until_all<T>(ptr, nelems, status, cmp, val));
+}
+
+template <typename T>
+__host__ size_t Context::wait_until_some(T* ptr, size_t nelems,
+                                         size_t* indices,
+                                         const int* status,
+                                         roc_shmem_cmps cmp, T val) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL_SOME);
+
+  HOST_DISPATCH_RET(wait_until_some<T>(ptr, nelems, indices, status, cmp, val));
+}
+
+template <typename T>
+__host__ void Context::wait_until_all_vector(T* ptr, size_t nelems,
+                                             const int *status,
+                                             roc_shmem_cmps cmp, T* vals) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL_ALL_VECTOR);
+
+  HOST_DISPATCH(wait_until_all_vector<T>(ptr, nelems, status, cmp, vals));
+}
+
+template <typename T>
+__host__ size_t Context::wait_until_any_vector(T* ptr, size_t nelems,
+                                               const int *status,
+                                               roc_shmem_cmps cmp, T* vals) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL_ANY_VECTOR);
+
+  HOST_DISPATCH_RET(wait_until_any_vector<T>(ptr, nelems, status, cmp, vals));
+}
+
+template <typename T>
+__host__ size_t Context::wait_until_some_vector(T* ptr, size_t nelems,
+                                                size_t* indices,
+                                                const int *status,
+                                                roc_shmem_cmps cmp, T* vals) {
+  ctxHostStats.incStat(NUM_HOST_WAIT_UNTIL_SOME_VECTOR);
+
+  HOST_DISPATCH_RET(wait_until_some_vector<T>(ptr, nelems, indices, status, cmp, vals));
+}
+
+template <typename T>
+__host__ int Context::test(T *ptr, roc_shmem_cmps cmp, T val) {
+  ctxHostStats.incStat(NUM_HOST_TEST);
+
+  HOST_DISPATCH_RET(test<T>(ptr, cmp, val));
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_CONTEXT_TMPL_HOST_HPP_

src/device_proxy.hpp

@@ -0,0 +1,108 @@
+/******************************************************************************
+ * 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_DEVICE_PROXY_HPP_
+#define LIBRARY_SRC_DEVICE_PROXY_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include <memory>
+#include <utility>
+
+namespace rocshmem {
+
+template <typename ALLOCATOR, typename T, size_t SIZE_IN = 1>
+class DeviceProxy {
+ public:
+  DeviceProxy() {
+    /*
+     * Allocate memory and verify that the allocation worked.
+     */
+    T* temp{nullptr};
+    allocator_.allocate(reinterpret_cast<void**>(&temp), SIZE_BYTES_);
+    assert(temp);
+
+    /*
+     * Default memory provided by the allocation to recognizable bytes.
+     */
+    memset(static_cast<void*>(temp), 0xBC, SIZE_BYTES_);
+
+    /*
+     * Pass the memory into a unique ptr for tracking.
+     */
+    std::unique_ptr<T, Deleter> up{temp};
+    up_ = std::move(up);
+
+    /*
+     * Set a c-style ptr for access by the device.
+     */
+    ptr_ = up_.get();
+  }
+
+  /**
+   * @brief Return internal storage tracked by the Proxy.
+   *
+   * @note Do not try to free this memory yourself. The proxy maintains
+   * the lifetime of the data itself.
+   */
+  __host__ __device__ T* get() { return ptr_; }
+
+ private:
+  /**
+   * @brief Internal Deleter functor is required by up_ member
+   */
+  class Deleter {
+   public:
+    void operator()(void* x) { a_.deallocate(x); }
+
+   private:
+    ALLOCATOR a_;
+  };
+
+  /**
+   * @brief Externally provided allocator type.
+   */
+  ALLOCATOR allocator_{};
+
+  /**
+   * @brief Unique pointer for tracking the proxy.
+   */
+  std::unique_ptr<T, Deleter> up_{nullptr};
+
+  /**
+   * @brief A handle to access the internal memory.
+   *
+   * In general, device code cannot access standard library routines
+   * like std::unique_ptr::get(). Circumvent this problem by caching
+   * the pointer manually in this class.
+   */
+  T* ptr_{nullptr};
+
+  /**
+   * @brief The allocation size for the internal memory
+   */
+  static constexpr size_t SIZE_BYTES_{sizeof(T) * SIZE_IN};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_DEVICE_PROXY_HPP_

src/fence_policy.hpp

@@ -0,0 +1,76 @@
+/******************************************************************************
+ * 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_FENCE_POLICY_HPP_
+#define LIBRARY_SRC_FENCE_POLICY_HPP_
+
+#include "roc_shmem/roc_shmem.hpp"
+
+namespace rocshmem {
+
+/**
+ * @brief Controls the behavior of device code which may need to stall
+ */
+class Fence {
+ public:
+  /**
+   * Secondary constructor
+   */
+  __host__ __device__ Fence() = default;
+
+  /**
+   * Primary constructor
+   *
+   * @param[in] options interpreted as a bitfield using bitwise operations
+   */
+  __host__ __device__ Fence(long option) {
+    if (option & ROC_SHMEM_CTX_NOSTORE) {
+      flush_ = false;
+    }
+  }
+
+  /**
+   * @brief Wait for outstanding memory operations to complete
+   *
+   * This can be useful when code needs guarantees about visibility
+   * before moving past the flush.
+   *
+   * @return void
+   */
+  __device__ void flush() {
+    if (flush_) {
+      __threadfence();
+    }
+  }
+
+ private:
+  /**
+   * @brief Used to toggle flushes behavior on and off
+   *
+   * @note By default, flushing is enabled.
+   */
+  bool flush_{true};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_FENCE_POLICY_HPP_

src/gpu_ib/CMakeLists.txt

@@ -0,0 +1,45 @@
+###############################################################################
+# 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.
+###############################################################################
+
+###############################################################################
+# ADD ROCSHMEM TARGET FOR FILES IN CURRENT DIRECTORY
+###############################################################################
+target_sources(
+  ${PROJECT_NAME}
+  PRIVATE
+    backend_ib.cpp
+    connection.cpp
+    connection_policy.cpp
+    context_ib_device.cpp
+    context_ib_device_coll.cpp
+    context_ib_host.cpp
+    debug.cpp
+    dynamic_connection.cpp
+    endian.cpp
+    gpu_ib_team.cpp
+    network_policy.cpp
+    qe_dumper.cpp
+    queue_pair.cpp
+    reliable_connection.cpp
+    segment_builder.cpp
+    thread_policy.cpp
+)

src/gpu_ib/backend_ib.cpp

@@ -0,0 +1,494 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "backend_ib.hpp"
+
+#include <endian.h>
+#include <mpi.h>
+#include <unistd.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <mutex>  // NOLINT(build/c++11)
+
+#include "roc_shmem/roc_shmem.hpp"
+#include "../backend_type.hpp"
+#include "../context_incl.hpp"
+#include "gpu_ib_team.hpp"
+#include "queue_pair.hpp"
+#include "../host/host.hpp"
+
+namespace rocshmem {
+
+#define NET_CHECK(cmd)                                       \
+  {                                                          \
+    if (cmd != MPI_SUCCESS) {                                \
+      fprintf(stderr, "Unrecoverable error: MPI Failure\n"); \
+      abort();                                               \
+    }                                                        \
+  }
+
+extern roc_shmem_ctx_t ROC_SHMEM_HOST_CTX_DEFAULT;
+
+roc_shmem_team_t get_external_team(GPUIBTeam *team) {
+  return reinterpret_cast<roc_shmem_team_t>(team);
+}
+
+int get_ls_non_zero_bit(char *bitmask, int mask_length) {
+  int position = -1;
+
+  for (int bit_i = 0; bit_i < mask_length; bit_i++) {
+    int byte_i = bit_i / CHAR_BIT;
+    if (bitmask[byte_i] & (1 << (bit_i % CHAR_BIT))) {
+      position = bit_i;
+      break;
+    }
+  }
+
+  return position;
+}
+
+GPUIBBackend::GPUIBBackend(MPI_Comm comm) : Backend() {
+  if (auto maximum_num_contexts_str = getenv("ROC_SHMEM_MAX_NUM_CONTEXTS")) {
+    std::stringstream sstream(maximum_num_contexts_str);
+    sstream >> maximum_num_contexts_;
+  }
+  num_blocks_ = maximum_num_contexts_;
+
+  init_mpi_once(comm);
+
+  type = BackendType::GPU_IB_BACKEND;
+
+  NET_CHECK(MPI_Comm_dup(backend_comm, &gpu_ib_comm_world));
+  NET_CHECK(MPI_Comm_size(gpu_ib_comm_world, &num_pes));
+  NET_CHECK(MPI_Comm_rank(gpu_ib_comm_world, &my_pe));
+
+  /* Initialize the host interface */
+  host_interface =
+      new HostInterface(hdp_proxy_.get(), gpu_ib_comm_world, &heap);
+
+  /*
+   * Construct default host context independently of the
+   * default device context (done in the async thread)
+   * so that host operations can execute regardless of
+   * device operations.
+   */
+  setup_default_host_ctx();
+
+  setup_team_world();
+
+  roc_shmem_collective_init();
+
+  teams_init();
+
+  // MPI_Comm_dup(gpu_ib_comm_world, &thread_comm);
+  thread_comm = gpu_ib_comm_world;
+
+  NET_CHECK(MPI_Barrier(gpu_ib_comm_world));
+
+  worker_thread_exit = false;
+
+#ifdef USE_HOST_SIDE_HDP_FLUSH
+  hdp_gpu_cpu_flush_flag_ =
+      static_cast<unsigned int *>(roc_shmem_malloc(sizeof(unsigned int)));
+  hdp_policy->set_flush_polling_ptr(hdp_gpu_cpu_flush_flag_);
+  hdp_flush_worker_thread = std::thread(&GPUIBBackend::hdp_flush_poll, this);
+
+  // We can now initialize and set the HDP window in the host interface
+  host_interface->create_hdp_window();
+
+#endif
+
+  // commenting out the async  thread as there is some issues with ROCm
+  // this makes the CPU init blocking
+  // async_thread_ = thread_spawn(this);
+  thread_func_internal(this);
+}
+
+__device__ bool GPUIBBackend::create_ctx(int64_t options,
+                                         roc_shmem_ctx_t *ctx) {
+  GPUIBContext *ctx_;
+
+  auto pop_result = ctx_free_list.get()->pop_front();
+  if (!pop_result.success) {
+    return false;
+  }
+  ctx_ = pop_result.value;
+
+  ctx_->dev_mtx_.shareable_ = static_cast<bool>(options >> 3);
+  ctx->ctx_opaque = ctx_;
+  return true;
+}
+
+void GPUIBBackend::ctx_create(int64_t options, void **ctx) {
+  GPUIBHostContext *new_ctx = nullptr;
+
+  new_ctx = new GPUIBHostContext(this, options);
+
+  *ctx = new_ctx;
+}
+
+GPUIBHostContext *get_internal_gpu_ib_ctx(Context *ctx) {
+  return reinterpret_cast<GPUIBHostContext *>(ctx);
+}
+
+void GPUIBBackend::ctx_destroy(Context *ctx) {
+  GPUIBHostContext *gpu_ib_host_ctx = get_internal_gpu_ib_ctx(ctx);
+  delete gpu_ib_host_ctx;
+}
+
+__device__ void GPUIBBackend::destroy_ctx(roc_shmem_ctx_t *ctx) {
+  ctx_free_list.get()->push_back(static_cast<GPUIBContext *>(ctx->ctx_opaque));
+}
+
+GPUIBBackend::~GPUIBBackend() {
+  // need to get this back once ROCm is fixed
+  // async_thread_.join();
+
+  worker_thread_exit = true;
+
+#ifdef USE_HOST_SIDE_HDP_FLUSH
+  hdp_flush_worker_thread.join();
+  hdp_policy->set_flush_polling_ptr(nullptr);
+  roc_shmem_free(hdp_gpu_cpu_flush_flag_);
+#endif
+
+  /**
+   * Destroy teams infrastructure
+   * and team world
+   */
+  teams_destroy();
+  auto *team_world{team_tracker.get_team_world()};
+  team_world->~Team();
+  CHECK_HIP(hipFree(team_world));
+
+  delete default_host_ctx_;
+
+  NET_CHECK(MPI_Comm_free(&gpu_ib_comm_world));
+
+  CHECK_HIP(hipFree(default_ctx_->device_qp_proxy));
+  CHECK_HIP(hipFree(default_ctx_));
+  default_ctx_ = nullptr;
+
+  delete host_interface;
+  host_interface = nullptr;
+
+  networkImpl.networkHostFinalize();
+
+  CHECK_HIP(hipFree(ctx_array));
+}
+
+__host__ void GPUIBBackend::global_exit(int status) {
+  MPI_Abort(gpu_ib_comm_world, status);
+}
+
+void GPUIBBackend::create_new_team([[maybe_unused]] Team *parent_team,
+                                   TeamInfo *team_info_wrt_parent,
+                                   TeamInfo *team_info_wrt_world, int num_pes,
+                                   int my_pe_in_new_team, MPI_Comm team_comm,
+                                   roc_shmem_team_t *new_team) {
+  /**
+   * Read the bit mask and find out a common index into
+   * the pool of available work arrays.
+   */
+  NET_CHECK(MPI_Allreduce(pool_bitmask_, reduced_bitmask_, bitmask_size_,
+                          MPI_CHAR, MPI_BAND, team_comm));
+
+  /* Pick the least significant non-zero bit (logical layout) in the reduced
+   * bitmask */
+  auto max_num_teams{team_tracker.get_max_num_teams()};
+  int common_index = get_ls_non_zero_bit(reduced_bitmask_, max_num_teams);
+  if (common_index < 0) {
+    /* No team available */
+    abort();
+  }
+
+  /* Mark the team as taken (by unsetting the bit in the pool bitmask) */
+  int byte = common_index / CHAR_BIT;
+  pool_bitmask_[byte] &= ~(1 << (common_index % CHAR_BIT));
+
+  /**
+   * Allocate device-side memory for team_world and
+   * construct a GPU_IB team in it
+   */
+  GPUIBTeam *new_team_obj;
+  CHECK_HIP(hipMalloc(&new_team_obj, sizeof(GPUIBTeam)));
+  new (new_team_obj)
+      GPUIBTeam(this, team_info_wrt_parent, team_info_wrt_world, num_pes,
+                my_pe_in_new_team, team_comm, common_index);
+
+  *new_team = get_external_team(new_team_obj);
+}
+
+void GPUIBBackend::team_destroy(roc_shmem_team_t team) {
+  GPUIBTeam *team_obj = get_internal_gpu_ib_team(team);
+
+  /* Mark the pool as available */
+  int bit = team_obj->pool_index_;
+  int byte_i = bit / CHAR_BIT;
+  pool_bitmask_[byte_i] |= 1 << (bit % CHAR_BIT);
+
+  team_obj->~GPUIBTeam();
+  CHECK_HIP(hipFree(team_obj));
+}
+
+void GPUIBBackend::dump_backend_stats() {
+  networkImpl.dump_backend_stats(&globalStats);
+}
+
+void GPUIBBackend::reset_backend_stats() { networkImpl.reset_backend_stats(); }
+
+void GPUIBBackend::initialize_ipc() {
+  ipcImpl.ipcHostInit(my_pe, heap.get_heap_bases(), thread_comm);
+}
+
+void GPUIBBackend::initialize_network() { networkImpl.networkHostSetup(this); }
+
+void GPUIBBackend::setup_default_host_ctx() {
+  default_host_ctx_ = new GPUIBHostContext(this, 0);
+  ROC_SHMEM_HOST_CTX_DEFAULT.ctx_opaque = default_host_ctx_;
+}
+
+void GPUIBBackend::setup_ctxs() {
+  /*
+   * Allocate device-side memory for all context and construct an
+   * InfiniBand context in it.
+   */
+  CHECK_HIP(
+      hipMalloc(&ctx_array, sizeof(GPUIBContext) * maximum_num_contexts_));
+  for (int i = 0; i < maximum_num_contexts_; i++) {
+    new (&ctx_array[i]) GPUIBContext(this, false, i);
+    ctx_free_list.get()->push_back(ctx_array + i);
+  }
+}
+
+void GPUIBBackend::setup_default_ctx() {
+  /*
+   * Allocate device-side memory for default context and construct an
+   * InfiniBand context in it.
+   */
+  CHECK_HIP(hipMalloc(&default_ctx_, sizeof(GPUIBContext)));
+  new (default_ctx_) GPUIBContext(this, true, 0);
+
+  /*
+   * Set the ROC_SHMEM_CTX_DEFAULT in constant memory.
+   */
+  int *symbol_address;
+  CHECK_HIP(hipGetSymbolAddress(reinterpret_cast<void **>(&symbol_address),
+                                HIP_SYMBOL(ROC_SHMEM_CTX_DEFAULT)));
+
+  roc_shmem_ctx_t ctx_default_host{default_ctx_, nullptr};
+
+  hipStream_t stream;
+  CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+  CHECK_HIP(hipMemcpyAsync(symbol_address, &ctx_default_host,
+                           sizeof(roc_shmem_ctx_t), hipMemcpyDefault, stream));
+  CHECK_HIP(hipStreamSynchronize(stream));
+  CHECK_HIP(hipStreamDestroy(stream));
+}
+
+void GPUIBBackend::setup_team_world() {
+  TeamInfo *team_info_wrt_parent, *team_info_wrt_world;
+
+  /**
+   * Allocate device-side memory for team_world and construct a
+   * GPU_IB team in it.
+   */
+  CHECK_HIP(hipMalloc(&team_info_wrt_parent, sizeof(TeamInfo)));
+  CHECK_HIP(hipMalloc(&team_info_wrt_world, sizeof(TeamInfo)));
+
+  new (team_info_wrt_parent) TeamInfo(nullptr, 0, 1, num_pes);
+  new (team_info_wrt_world) TeamInfo(nullptr, 0, 1, num_pes);
+
+  MPI_Comm team_world_comm;
+  NET_CHECK(MPI_Comm_dup(gpu_ib_comm_world, &team_world_comm));
+
+  GPUIBTeam *team_world{nullptr};
+  CHECK_HIP(hipMalloc(&team_world, sizeof(GPUIBTeam)));
+  new (team_world) GPUIBTeam(this, team_info_wrt_parent, team_info_wrt_world,
+                             num_pes, my_pe, team_world_comm, 0);
+  team_tracker.set_team_world(team_world);
+
+  /**
+   * Copy the address to ROC_SHMEM_TEAM_WORLD.
+   */
+  ROC_SHMEM_TEAM_WORLD = reinterpret_cast<roc_shmem_team_t>(team_world);
+}
+
+void GPUIBBackend::init_mpi_once(MPI_Comm comm) {
+  static std::mutex init_mutex;
+  const std::lock_guard<std::mutex> lock(init_mutex);
+
+  int init_done = 0;
+  NET_CHECK(MPI_Initialized(&init_done));
+  if (init_done == 0) {
+    int provided;
+    NET_CHECK(
+        MPI_Init_thread(nullptr, nullptr, MPI_THREAD_SINGLE, &provided));
+  }
+  if (comm == MPI_COMM_NULL) {
+    NET_CHECK(MPI_Comm_dup(MPI_COMM_WORLD, &backend_comm));
+  } else {
+    NET_CHECK(MPI_Comm_dup(comm, &backend_comm));
+  }
+}
+
+std::thread GPUIBBackend::thread_spawn(GPUIBBackend *b) {
+  return std::thread(&GPUIBBackend::thread_func_internal, this, b);
+}
+
+void GPUIBBackend::thread_func_internal(GPUIBBackend *b) {
+  CHECK_HIP(hipSetDevice(hip_dev_id));
+
+  b->initialize_ipc();
+  b->initialize_network();
+  b->setup_ctxs();
+  b->setup_default_ctx();
+  *(b->done_init) = 1;
+}
+
+#ifdef USE_HOST_SIDE_HDP_FLUSH
+void GPUIBBackend::hdp_flush_poll() {
+  while (!worker_thread_exit) {
+    if (hdp_policy->has_active_flush_request()) {
+      hdp_policy->hdp_flush();
+      hdp_policy->clear_active_flush_flag();
+    }
+  }
+}
+#endif
+
+void GPUIBBackend::teams_init() {
+  /**
+   * Allocate pools for the teams sync and work arrary from the SHEAP.
+   */
+  auto max_num_teams{team_tracker.get_max_num_teams()};
+  barrier_pSync_pool = reinterpret_cast<long *>(roc_shmem_malloc(
+      sizeof(long) * ROC_SHMEM_BARRIER_SYNC_SIZE * max_num_teams));
+  reduce_pSync_pool = reinterpret_cast<long *>(roc_shmem_malloc(
+      sizeof(long) * ROC_SHMEM_REDUCE_SYNC_SIZE * max_num_teams));
+  bcast_pSync_pool = reinterpret_cast<long *>(roc_shmem_malloc(
+      sizeof(long) * ROC_SHMEM_BCAST_SYNC_SIZE * max_num_teams));
+  alltoall_pSync_pool = reinterpret_cast<long *>(roc_shmem_malloc(
+      sizeof(long) * ROC_SHMEM_ALLTOALL_SYNC_SIZE * max_num_teams));
+
+  /* Accommodating for largest possible data type for pWrk */
+  pWrk_pool = roc_shmem_malloc(
+      sizeof(double) * ROC_SHMEM_REDUCE_MIN_WRKDATA_SIZE * max_num_teams);
+  pAta_pool = roc_shmem_malloc(sizeof(double) * ROC_SHMEM_ATA_MAX_WRKDATA_SIZE *
+                               max_num_teams);
+
+  /**
+   * Initialize the sync arrays in the pool with default values.
+   */
+  long *barrier_pSync, *reduce_pSync, *bcast_pSync, *alltoall_pSync;
+  for (int team_i = 0; team_i < max_num_teams; team_i++) {
+    barrier_pSync = reinterpret_cast<long *>(
+        &barrier_pSync_pool[team_i * ROC_SHMEM_BARRIER_SYNC_SIZE]);
+    reduce_pSync = reinterpret_cast<long *>(
+        &reduce_pSync_pool[team_i * ROC_SHMEM_REDUCE_SYNC_SIZE]);
+    bcast_pSync = reinterpret_cast<long *>(
+        &bcast_pSync_pool[team_i * ROC_SHMEM_BCAST_SYNC_SIZE]);
+    alltoall_pSync = reinterpret_cast<long *>(
+        &alltoall_pSync_pool[team_i * ROC_SHMEM_ALLTOALL_SYNC_SIZE]);
+
+    for (int i = 0; i < ROC_SHMEM_BARRIER_SYNC_SIZE; i++) {
+      barrier_pSync[i] = ROC_SHMEM_SYNC_VALUE;
+    }
+    for (int i = 0; i < ROC_SHMEM_REDUCE_SYNC_SIZE; i++) {
+      reduce_pSync[i] = ROC_SHMEM_SYNC_VALUE;
+    }
+    for (int i = 0; i < ROC_SHMEM_BCAST_SYNC_SIZE; i++) {
+      bcast_pSync[i] = ROC_SHMEM_SYNC_VALUE;
+    }
+    for (int i = 0; i < ROC_SHMEM_ALLTOALL_SYNC_SIZE; i++) {
+      alltoall_pSync[i] = ROC_SHMEM_SYNC_VALUE;
+    }
+  }
+
+  /**
+   * Initialize bit mask
+   *
+   * Logical:
+   * MSB..........................................................................LSB
+   * Physical: MSB...1st least significant 8 bits...LSB  MSB...2nd least
+   * signifant 8 bits...LSB
+   *
+   * Description shows only a 2-byte long mask but idea extends to any
+   * arbitrary size.
+   */
+  bitmask_size_ = (max_num_teams % CHAR_BIT) ? (max_num_teams / CHAR_BIT + 1)
+                                             : (max_num_teams / CHAR_BIT);
+  pool_bitmask_ = reinterpret_cast<char *>(malloc(bitmask_size_));
+  reduced_bitmask_ = reinterpret_cast<char *>(malloc(bitmask_size_));
+
+  memset(pool_bitmask_, 0, bitmask_size_);
+  memset(reduced_bitmask_, 0, bitmask_size_);
+  /* Set all to available except the 0th one (reserved for TEAM_WORLD) */
+  for (int bit_i = 1; bit_i < max_num_teams; bit_i++) {
+    int byte_i = bit_i / CHAR_BIT;
+
+    pool_bitmask_[byte_i] |= 1 << (bit_i % CHAR_BIT);
+  }
+
+  /**
+   * Make sure that all processing elements have done this before
+   * continuing.
+   */
+  NET_CHECK(MPI_Barrier(gpu_ib_comm_world));
+}
+
+void GPUIBBackend::teams_destroy() {
+  roc_shmem_free(barrier_pSync_pool);
+  roc_shmem_free(reduce_pSync_pool);
+  roc_shmem_free(bcast_pSync_pool);
+  roc_shmem_free(alltoall_pSync_pool);
+  roc_shmem_free(pWrk_pool);
+  roc_shmem_free(pAta_pool);
+
+  free(pool_bitmask_);
+  free(reduced_bitmask_);
+}
+
+void GPUIBBackend::roc_shmem_collective_init() {
+  /*
+   * Allocate heap space for barrier_sync
+   */
+  size_t one_sync_size_bytes{sizeof(*barrier_sync)};
+  size_t sync_size_bytes{one_sync_size_bytes * ROC_SHMEM_BARRIER_SYNC_SIZE};
+  heap.malloc(reinterpret_cast<void **>(&barrier_sync), sync_size_bytes);
+
+  /*
+   * Initialize the barrier synchronization array with default values.
+   */
+  for (int i = 0; i < num_pes; i++) {
+    barrier_sync[i] = ROC_SHMEM_SYNC_VALUE;
+  }
+
+  /*
+   * Make sure that all processing elements have done this before
+   * continuing.
+   */
+  NET_CHECK(MPI_Barrier(gpu_ib_comm_world));
+}
+
+}  // namespace rocshmem

src/gpu_ib/backend_ib.hpp

@@ -0,0 +1,351 @@
+/******************************************************************************
+ * 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_GPU_IB_BACKEND_IB_HPP_
+#define LIBRARY_SRC_GPU_IB_BACKEND_IB_HPP_
+
+#include "../backend_bc.hpp"
+#include "../containers/free_list_impl.hpp"
+#include "network_policy.hpp"
+#include "../hdp_policy.hpp"
+#include "../hdp_proxy.hpp"
+#include "../memory/hip_allocator.hpp"
+
+namespace rocshmem {
+
+class HostInterface;
+
+/**
+ * @class GPUIBBackend backend.hpp
+ * @brief InfiniBand specific backend.
+ *
+ * The InfiniBand (GPUIB) backend enables the device to enqueue network
+ * requests to InfiniBand queues (with minimal host intervention). The setup
+ * requires some effort from the host, but the device is able to craft
+ * InfiniBand requests and send them on its own.
+ */
+class GPUIBBackend : public Backend {
+ public:
+  /**
+   * @copydoc Backend::Backend(unsigned)
+   */
+  explicit GPUIBBackend(MPI_Comm comm);
+
+  /**
+   * @copydoc Backend::~Backend()
+   */
+  virtual ~GPUIBBackend();
+
+  /**
+   * @brief Abort the application.
+   *
+   * @param[in] status Exit code.
+   *
+   * @return void.
+   *
+   * @note This routine terminates the entire application.
+   */
+  void global_exit(int status) override;
+
+  /**
+   * @copydoc Backend::create_new_team
+   */
+  void create_new_team(Team *parent_team, TeamInfo *team_info_wrt_parent,
+                       TeamInfo *team_info_wrt_world, int num_pes,
+                       int my_pe_in_new_team, MPI_Comm team_comm,
+                       roc_shmem_team_t *new_team) override;
+
+  /**
+   * @copydoc Backend::team_destroy(roc_shmem_team_t)
+   */
+  void team_destroy(roc_shmem_team_t team) override;
+
+  /**
+   * @copydoc Backend::ctx_create
+   */
+  void ctx_create(int64_t options, void **ctx) override;
+
+  __device__ bool create_ctx(int64_t options, roc_shmem_ctx_t *ctx);
+
+  /**
+   * @copydoc Backend::ctx_destroy
+   */
+  void ctx_destroy(Context *ctx) override;
+
+  /**
+   * @copydoc Backend::ctx_destroy
+   */
+  __device__ void destroy_ctx(roc_shmem_ctx_t *ctx);
+
+ protected:
+  /**
+   * @copydoc Backend::dump_backend_stats()
+   */
+  void dump_backend_stats() override;
+
+  /**
+   * @copydoc Backend::reset_backend_stats()
+   */
+  void reset_backend_stats() override;
+
+  /**
+   * @brief spawn a new thread to perform the rest of initialization
+   */
+  std::thread thread_spawn(GPUIBBackend *b);
+
+  /**
+   * @brief overheads for helper thread to run
+   *
+   * @param[in] the thread needs access to the class
+   *
+   * @return void
+   */
+  void thread_func_internal(GPUIBBackend *b);
+
+  /**
+   * @brief initialize MPI.
+   *
+   * GPUIB relies on MPI just to exchange the connection information.
+   *
+   * todo: remove the dependency on MPI and make it generic to PMI-X or just
+   * to OpenSHMEM to have support for both CPU and GPU
+   */
+  void init_mpi_once(MPI_Comm comm);
+
+  /**
+   * @brief init the network support
+   */
+  void initialize_network();
+
+  /**
+   * @brief Invokes the IPC policy class initialization method.
+   *
+   * This method delegates Inter Process Communication (IPC)
+   * initialization to the appropriate policy class. The initialization
+   * needs to be exposed to the Backed due to initialization ordering
+   * constraints. (The symmetric heaps needs to be allocated and
+   * initialized before this method can be called.)
+   *
+   * The policy class encapsulates what the initialization process so
+   * refer to that class for more details.
+   */
+  void initialize_ipc();
+
+  /**
+   * @brief Allocate and initialize the ROC_SHMEM_CTX_DEFAULT variable.
+   *
+   * @todo The default_ctx member looks unused after it is copied into
+   * the ROC_SHMEM_CTX_DEFAULT variable.
+   */
+  void setup_default_ctx();
+  void setup_ctxs();
+
+  /**
+   * @brief Allocate and initialize the default context for host
+   * operations.
+   */
+  void setup_default_host_ctx();
+
+  /**
+   * @brief Allocate and initialize team world.
+   */
+  void setup_team_world();
+
+  /**
+   * @brief Initialize the resources required to support teams
+   */
+  void teams_init();
+
+  /**
+   * @brief Destruct the resources required to support teams
+   */
+  void teams_destroy();
+
+  /**
+   * @brief Allocate and initialize barrier operation addresses on
+   * symmetric heap.
+   *
+   * When this method completes, the barrier_sync member will be available
+   * for use.
+   */
+  void roc_shmem_collective_init();
+
+#ifdef USE_HOST_SIDE_HDP_FLUSH
+  /**
+   * @brief A service thread routine that flushes the hdp cache on behalf of the
+   * GPU.
+   */
+  void hdp_flush_poll();
+
+  /**
+   * @brief Workers used to poll on the device hdp flush request.
+   */
+  std::thread hdp_flush_worker_thread{};
+#endif
+
+  /**
+   * @brief Signals to the worker threads to exist
+   */
+  std::atomic<bool> worker_thread_exit{false};
+
+ public:
+  /**
+   * @brief The host-facing interface that will be used
+   * by all contexts of the GPUIBBackend
+   */
+  HostInterface *host_interface{nullptr};
+
+  /**
+   * @brief Handle for raw memory for barrier sync
+   */
+  long *barrier_pSync_pool{nullptr};
+
+  /**
+   * @brief Handle for raw memory for reduce sync
+   */
+  long *reduce_pSync_pool{nullptr};
+
+  /**
+   * @brief Handle for raw memory for broadcast sync
+   */
+  long *bcast_pSync_pool{nullptr};
+
+  /**
+   * @brief Handle for raw memory for alltoall sync
+   */
+  long *alltoall_pSync_pool{nullptr};
+
+  /**
+   * @brief Handle for raw memory for work
+   */
+  void *pWrk_pool{nullptr};
+
+  /**
+   * @brief Handle for raw memory for alltoall
+   */
+  void *pAta_pool{nullptr};
+
+  /**
+   * @brief ROC_SHMEM's copy of MPI_COMM_WORLD (for interoperability
+   * with orthogonal MPI usage in an MPI+ROC_SHMEM program).
+   */
+  MPI_Comm gpu_ib_comm_world{};
+  MPI_Comm backend_comm{};
+
+  /**
+   * @brief Holds number of blocks used in library
+   */
+  size_t num_blocks_{1};
+
+ private:
+  /**
+   * @brief Allocates cacheable, device memory for the hdp policy.
+   *
+   * @note Internal data ownership is managed by the proxy
+   */
+  HdpProxy<HIPAllocator> hdp_proxy_{};
+
+ public:
+  /**
+   * @brief Policy choice for two HDP implementations.
+   *
+   * @todo Combine HDP related stuff together into a class with a
+   * reasonable interface. The functionality does not need to exist in
+   * multiple pieces in the Backend and QueuePair classes. The hdp_rkey,
+   * hdp_addresses, and hdp_policy fields should all live in the class.
+   */
+  HdpPolicy *hdp_policy{hdp_proxy_.get()};
+
+  /**
+   * @brief Scratchpad for the internal barrier algorithms.
+   */
+  int64_t *barrier_sync{nullptr};
+
+  /**
+   * @brief Compile-time configuration policy for network (IB)
+   *
+   *
+   * The configuration option "USE_SINGLE_NODE" can be enabled to not build
+   * with network support.
+   */
+  NetworkImpl networkImpl{};
+
+ private:
+  /**
+   * @brief An array of @ref ROContexts that backs the context FreeList.
+   */
+  GPUIBContext *ctx_array{nullptr};
+
+  /**
+   * @brief A free-list containing contexts.
+   */
+  FreeListProxy<HIPAllocator, GPUIBContext *> ctx_free_list{};
+
+  /**
+   * @brief Holds maximum number of contexts used in library
+   */
+  size_t maximum_num_contexts_{1024};
+
+  /**
+   * @brief The bitmask representing the availability of teams in the pool
+   */
+  char *pool_bitmask_{nullptr};
+
+  /**
+   * @brief Bitmask to store the reduced result of bitmasks on pariticipating
+   * PEs
+   *
+   * With no thread-safety for this bitmask, multithreaded creation of teams is
+   * not supported.
+   */
+  char *reduced_bitmask_{nullptr};
+
+  /**
+   * @brief Size of the bitmask
+   */
+  int bitmask_size_{-1};
+
+  /**
+   * @brief a helper thread to perform the initialization (non-blocking init)
+   */
+  std::thread async_thread_{};
+
+  /**
+   * @brief Holds a copy of the default context (see OpenSHMEM
+   * specification).
+   *
+   * @todo Remove this member from the backend class. There is another
+   * copy stored in ROC_SHMEM_CTX_DEFAULT.
+   */
+  GPUIBContext *default_ctx_{nullptr};
+
+  /**
+   * @brief Holds a copy of the default context for host functions
+   */
+  GPUIBHostContext *default_host_ctx_{nullptr};
+
+  unsigned int* hdp_gpu_cpu_flush_flag_;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_BACKEND_IB_HPP_

src/gpu_ib/connect_tests.cpp

@@ -0,0 +1,54 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include <gtest/gtest.h>
+#include <mpi.h>
+
+#include "dynamic_connection.hpp"
+#include "reliable_connection.hpp"
+
+namespace rocshmem {
+
+// test with different use_ib_hca
+// test with different heap size
+// test with different sleep
+// test with different sq_size
+
+TEST(DynamicConnect, ToNothing) {
+  DynamicConnection connect;
+  connect.construct_init(1);
+}
+
+// test with different num_dcis
+// test with different num_dcts
+
+TEST(ReliableConnect, ToNothing) {
+  ReliableConnection connect;
+  connect.construct_init(1);
+}
+
+int main(int argc, char **argv) {
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
+
+}  // namespace rocshmem

src/gpu_ib/connection.cpp

@@ -0,0 +1,431 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "connection.hpp"
+
+#include <mpi.h>
+
+#include <mutex>  // NOLINT(build/c++11)
+#include <vector>
+
+#include "backend_ib.hpp"
+#include "queue_pair.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+int Connection::use_gpu_mem = 0;
+int Connection::coherent_cq = 0;
+
+Connection::Connection(GPUIBBackend* b, int k) : backend(b), key_offset(k) {
+  char* value = nullptr;
+
+  if ((value = getenv("ROC_SHMEM_USE_IB_HCA"))) {
+    requested_dev = value;
+  }
+
+  if ((value = getenv("ROC_SHMEM_SQ_SIZE"))) {
+    sq_size = atoi(value);
+  }
+
+  if ((value = getenv("ROC_SHMEM_USE_CQ_GPU_MEM")) != nullptr) {
+    cq_use_gpu_mem = atoi(value);
+  }
+
+  if ((value = getenv("ROC_SHMEM_USE_SQ_GPU_MEM")) != nullptr) {
+    sq_use_gpu_mem = atoi(value);
+  }
+}
+
+Connection::~Connection() { delete ib_state; }
+
+void Connection::reg_mr(void* ptr, size_t size, ibv_mr** mr, bool managed) {
+  int access = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE |
+               IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_ATOMIC;
+  if (managed) {
+    access |= IBV_ACCESS_ON_DEMAND;
+  }
+
+  *mr = ibv_reg_mr(ib_state->pd, ptr, size, access);
+
+  if (*mr == nullptr) {
+    abort();
+  }
+}
+
+unsigned Connection::total_number_connections() {
+  int connections;
+  get_remote_conn(&connections);
+  return backend->num_blocks_ * connections;
+}
+
+void Connection::initialize(int num_block) {
+  allocate_dynamic_members(num_block);
+
+  int ib_devices{0};
+  dev_list = ibv_get_device_list(&ib_devices);
+  if (dev_list == nullptr) {
+    abort();
+  }
+
+  struct ibv_device* ib_dev = dev_list[0];
+  if (requested_dev != nullptr) {
+    for (int i = 0; i < ib_devices; i++) {
+      const char* select_dev = ibv_get_device_name(dev_list[i]);
+      if (strstr(select_dev, requested_dev) != nullptr) {
+        ib_dev = dev_list[i];
+        break;
+      }
+    }
+  }
+
+  uint8_t port = 1;
+  ib_init(ib_dev, port);
+
+  int hip_dev_id = 0;
+  CHECK_HIP(hipGetDevice(&hip_dev_id));
+
+  int ib_fork_err = ibv_fork_init();
+  if (ib_fork_err != 0) printf("error: ibv_fork_init failed \n");
+
+  sq_post_dv = static_cast<sq_post_dv_t*>(
+      malloc(sizeof(sq_post_dv_t) * total_number_connections()));
+
+  if (sq_post_dv == nullptr) {
+    abort();
+  }
+
+  create_qps(port, backend->my_pe, &ib_state->portinfo);
+  initialize_1(port, num_block);
+
+  MPI_Barrier(backend->thread_comm);
+  free_dynamic_members();
+}
+
+void Connection::finalize() {
+  ibv_free_device_list(dev_list);
+
+  int ret = ibv_dereg_mr(backend->networkImpl.heap_mr);
+  if (ret) {
+    abort();
+  }
+  // comment until rocm 4.5
+  // ibv_dereg_mr(backend->networkImpl.hdp_mr);
+  ibv_dereg_mr(backend->networkImpl.mr);
+}
+
+void Connection::ib_init(struct ibv_device* ib_dev, uint8_t port) {
+  ib_state = new ib_state_t;
+  if (!ib_state) {
+    abort();
+  }
+
+  ib_state->context = ibv_open_device(ib_dev);
+  if (!ib_state->context) {
+    delete ib_state;
+    abort();
+  }
+
+  ib_state->pd = ibv_alloc_pd(ib_state->context);
+  if (!ib_state->pd) {
+    delete ib_state;
+    abort();
+  }
+
+  ibv_parent_domain_init_attr pattr;
+  init_parent_domain_attr(&pattr);
+  ib_state->pd = ibv_alloc_parent_domain(ib_state->context, &pattr);
+
+  ibv_query_port(ib_state->context, port, &ib_state->portinfo);
+}
+
+template <typename StateType>
+void Connection::try_to_modify_qp(ibv_qp* qp, StateType state) {
+  ibv_modify_qp(qp, &state.exp_qp_attr, state.exp_attr_mask);
+}
+
+void Connection::init_qp_status(ibv_qp* qp, uint8_t port) {
+  try_to_modify_qp<InitQPState>(qp, initqp(port));
+}
+
+/**
+ * rtr stands for 'ready to receive'
+ */
+void Connection::change_status_rtr(ibv_qp* qp, dest_info_t* dest,
+                                     uint8_t port) {
+  try_to_modify_qp<RtrState>(qp, rtr(dest, port));
+}
+
+/**
+ * rts stands for 'ready to send'
+ */
+void Connection::change_status_rts(ibv_qp* qp, dest_info_t* dest) {
+  try_to_modify_qp<RtsState>(qp, rts(dest));
+}
+
+void Connection::create_qps(uint8_t port, int my_rank,
+                              ibv_port_attr* ib_port_att) {
+  create_qps_1();
+
+  ibv_qp_cap cap{};
+  cap.max_send_wr = sq_size;
+  cap.max_send_sge = 1;
+  cap.max_inline_data = 4;
+
+  QPInitAttr qp_init_attr = qpattr(cap);
+
+  size_t qp_size = total_number_connections();
+  cqs.resize(qp_size);
+  qps.resize(qp_size);
+
+  int cqe = qp_init_attr.attr.cap.max_send_wr;
+  for (auto& entry : cqs) {
+    entry = create_cq(ib_state->context, ib_state->pd, cqe);
+    if (!entry) {
+      abort();
+    }
+  }
+
+  create_qps_2(port, my_rank, ib_port_att);
+
+  for (int i = 0; i < qps.size(); i++) {
+    qps[i] =
+        create_qp(ib_state->pd, ib_state->context, &qp_init_attr.attr, cqs[i]);
+    if (!qps[i]) {
+      abort();
+    }
+
+    create_qps_3(port, qps[i], i, ib_port_att);
+  }
+}
+
+void Connection::initialize_gpu_policy(ConnectionImpl** conn,
+                                         uint32_t* heap_rkey) {
+  CHECK_HIP(hipMalloc(reinterpret_cast<void**>(conn), sizeof(ConnectionImpl)));
+  new (*conn) ConnectionImpl(this, heap_rkey);
+}
+
+/*
+ * Create and write the rdma segment to the SQ
+ */
+void Connection::set_rdma_seg(mlx5_wqe_raddr_seg* rdma, uint64_t address,
+                              uint32_t rkey) {
+  rdma->raddr = htobe64(address);
+  rdma->rkey = htobe32(rkey);
+}
+
+/*
+ * Retrieve the address of a SQ.
+ * We used this address to write the WQE directly to the SQ.
+ */
+uint64_t* Connection::get_address_sq(int i) {
+  mlx5dv_obj mlx_obj;
+  mlx5dv_qp qp_out;
+
+  mlx_obj.qp.in = qps[i];
+  mlx_obj.qp.out = &qp_out;
+
+  mlx5dv_init_obj(&mlx_obj, MLX5DV_OBJ_QP);
+
+  return reinterpret_cast<uint64_t*>(qp_out.sq.buf);
+}
+
+void* Connection::buf_alloc([[maybe_unused]] struct ibv_pd* pd,
+                            [[maybe_unused]] void* pd_context, size_t size,
+                            [[maybe_unused]] size_t alignment,
+                            [[maybe_unused]] uint64_t resource_type) {
+  if (use_gpu_mem) {
+    void* dev_ptr;
+    if (coherent_cq == 1) {
+#if defined USE_COHERENT_HEAP || defined USE_CACHED_HEAP
+      CHECK_HIP(hipMalloc(reinterpret_cast<void**>(&dev_ptr), size));
+#else
+  #ifdef HIP_SUPPORTS_MALLOC_UNCACHED
+        CHECK_HIP(hipExtMallocWithFlags(reinterpret_cast<void**>(&dev_ptr), size,
+                                        hipDeviceMallocUncached));
+  #else
+        CHECK_HIP(hipExtMallocWithFlags(reinterpret_cast<void**>(&dev_ptr), size,
+                                        hipDeviceMallocFinegrained));
+  #endif
+#endif
+    } else {
+#ifdef HIP_SUPPORTS_MALLOC_UNCACHED
+      CHECK_HIP(hipExtMallocWithFlags(reinterpret_cast<void**>(&dev_ptr), size,
+                                      hipDeviceMallocUncached));
+#else
+      CHECK_HIP(hipExtMallocWithFlags(reinterpret_cast<void**>(&dev_ptr), size,
+                                      hipDeviceMallocFinegrained));
+#endif
+
+    }
+    memset(dev_ptr, 0, size);
+    return dev_ptr;
+  }
+  return IBV_ALLOCATOR_USE_DEFAULT;
+}
+
+void Connection::buf_release([[maybe_unused]] struct ibv_pd* pd,
+                             [[maybe_unused]] void* pd_context, void* ptr,
+                             [[maybe_unused]] uint64_t resource_type) {
+  if (use_gpu_mem) {
+    CHECK_HIP(hipFree(ptr));
+  } else {
+    free(ptr);
+  }
+}
+
+void Connection::init_parent_domain_attr(ibv_parent_domain_init_attr* attr1) {
+  attr1->pd = ib_state->pd;
+  attr1->td = nullptr;
+  attr1->comp_mask = IBV_PARENT_DOMAIN_INIT_ATTR_ALLOCATORS;
+  attr1->alloc = Connection::buf_alloc;
+  attr1->free = Connection::buf_release;
+  attr1->pd_context = nullptr;
+}
+
+ibv_cq* Connection::create_cq(ibv_context* context, ibv_pd* pd, int cqe) {
+  use_gpu_mem = cq_use_gpu_mem;
+
+  ibv_cq_init_attr_ex cq_attr;
+  memset(&cq_attr, 0, sizeof(ibv_cq_init_attr_ex));
+  cq_attr.cqe = cqe;
+  cq_attr.cq_context = nullptr;
+  cq_attr.channel = nullptr;
+  cq_attr.comp_vector = 0;
+  cq_attr.flags = 0;  // see ibv_exp_cq_create_flags
+  cq_attr.comp_mask = IBV_CQ_INIT_ATTR_MASK_PD;
+  cq_attr.parent_domain = pd;
+
+  coherent_cq = 1;
+  ibv_cq_ex* cq = ibv_create_cq_ex(context, &cq_attr);
+  coherent_cq = 0;
+  if (!cq) {
+    printf("error in ibv_create_cq_ex: %d %s\n", errno, strerror(errno));
+    return nullptr;
+  }
+  return ibv_cq_ex_to_cq(cq);
+}
+
+void Connection::init_gpu_qp_from_connection(QueuePair* gpu_qp,
+                                               int conn_num) {
+  int hip_dev_id = 0;
+  CHECK_HIP(hipGetDevice(&hip_dev_id));
+  use_gpu_mem = cq_use_gpu_mem;
+
+  mlx5dv_cq cq_out;
+  mlx5dv_obj mlx_obj;
+  mlx_obj.cq.in = cqs[conn_num];
+  mlx_obj.cq.out = &cq_out;
+
+  mlx5dv_init_obj(&mlx_obj, MLX5DV_OBJ_CQ);
+  gpu_qp->cq_log_size = log2(cq_out.cqe_cnt);
+  gpu_qp->cq_size = cq_out.cqe_cnt;
+
+  void* gpu_ptr = nullptr;
+  if (use_gpu_mem) {
+    gpu_qp->current_cq_q = reinterpret_cast<mlx5_cqe64*>(cq_out.buf);
+  } else {
+    rocm_memory_lock_to_fine_grain(reinterpret_cast<void*>(cq_out.buf),
+                                   cq_out.cqe_cnt * 64, &gpu_ptr, hip_dev_id);
+    gpu_qp->current_cq_q = reinterpret_cast<mlx5_cqe64*>(gpu_ptr);
+  }
+  gpu_qp->current_cq_q_H = reinterpret_cast<mlx5_cqe64*>(cq_out.buf);
+
+  rocm_memory_lock_to_fine_grain(reinterpret_cast<void*>(cq_out.dbrec), 64,
+                                 &gpu_ptr, hip_dev_id);
+
+  gpu_qp->dbrec_cq = reinterpret_cast<volatile uint32_t*>(gpu_ptr);
+
+  use_gpu_mem = sq_use_gpu_mem;
+
+  mlx5dv_qp qp_out;
+  mlx_obj.qp.in = qps[conn_num];
+  mlx_obj.qp.out = &qp_out;
+
+  mlx5dv_init_obj(&mlx_obj, MLX5DV_OBJ_QP);
+
+  gpu_qp->max_nwqe = (qp_out.sq.wqe_cnt);
+
+  volatile uint32_t* dbrec_send = qp_out.dbrec + 1;
+
+  if (use_gpu_mem) {
+    gpu_qp->current_sq = reinterpret_cast<uint64_t*>(qp_out.sq.buf);
+    gpu_qp->dbrec_send = reinterpret_cast<volatile uint32_t*>(dbrec_send);
+  } else {
+    gpu_ptr = nullptr;
+    rocm_memory_lock_to_fine_grain(reinterpret_cast<void*>(qp_out.sq.buf),
+                                   qp_out.sq.wqe_cnt * 64, &gpu_ptr,
+                                   hip_dev_id);
+
+    gpu_qp->current_sq = reinterpret_cast<uint64_t*>(gpu_ptr);
+
+    rocm_memory_lock_to_fine_grain(
+        reinterpret_cast<void*>(const_cast<uint32_t*>(dbrec_send)), 32,
+        &gpu_ptr, hip_dev_id);
+
+    gpu_qp->dbrec_send = reinterpret_cast<volatile uint32_t*>(gpu_ptr);
+  }
+
+  gpu_qp->current_sq_H = reinterpret_cast<uint64_t*>(qp_out.sq.buf);
+
+  gpu_qp->setDBval(*(reinterpret_cast<uint64_t*>(qp_out.sq.buf)));
+
+  rocm_memory_lock_to_fine_grain(qp_out.bf.reg, qp_out.bf.size * 2, &gpu_ptr,
+                                 hip_dev_id);
+
+  gpu_qp->db.ptr = reinterpret_cast<uint64_t*>(gpu_ptr);
+
+  uint32_t* sq = reinterpret_cast<uint32_t*>(qp_out.sq.buf);
+  uint32_t ctrl_qp_sq = (reinterpret_cast<uint32_t*>(sq))[1];
+  gpu_qp->ctrl_qp_sq = ctrl_qp_sq & 0xFFFFFF;
+  gpu_qp->ctrl_sig = (reinterpret_cast<uint64_t*>(sq))[1];
+  gpu_qp->rkey = (reinterpret_cast<uint32_t*>(sq))[6 + key_offset];
+  gpu_qp->lkey = (reinterpret_cast<uint32_t*>(sq))[9 + key_offset];
+}
+
+ibv_qp* Connection::create_qp(ibv_pd* pd, ibv_context* context,
+                              ibv_qp_init_attr_ex* qp_attr, ibv_cq* cq) {
+  use_gpu_mem = sq_use_gpu_mem;
+
+  ibv_qp* qp = nullptr;
+
+  assert(pd);
+  assert(context);
+  assert(qp_attr);
+
+  qp_attr->send_cq = cq;
+  qp_attr->recv_cq = cq;
+  qp_attr->pd = pd;
+
+  qp_attr->comp_mask = IBV_QP_INIT_ATTR_PD;
+
+  qp = create_qp_0(context, qp_attr);
+
+  if (!qp) {
+    printf("***** error ibv_create_qp failed %d m %m \n", errno, errno);
+    ibv_destroy_cq(cq);
+  }
+
+  return qp;
+}
+
+}  // namespace rocshmem

src/gpu_ib/connection.hpp

@@ -0,0 +1,259 @@
+/******************************************************************************
+ * 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_GPU_IB_CONNECTION_HPP_
+#define LIBRARY_SRC_GPU_IB_CONNECTION_HPP_
+
+#include <infiniband/verbs.h>
+
+extern "C" {
+#include <infiniband/mlx5dv.h>
+}
+
+#include <vector>
+
+#include "roc_shmem/roc_shmem.hpp"
+#include "connection_policy.hpp"
+
+namespace rocshmem {
+
+class GPUIBBackend;
+class QueuePair;
+
+class Connection {
+ protected:
+  typedef struct ib_state {
+    struct ibv_context* context;
+    struct ibv_pd* pd;
+    struct ibv_mr* mr;
+    struct ibv_port_attr portinfo;
+  } ib_state_t;
+
+  typedef struct dest_info {
+    int lid;
+    int qpn;
+    int psn;
+    union ibv_gid gid;
+  } dest_info_t;
+
+  typedef struct heap_info {
+    void* base_heap;
+    uint32_t rkey;
+  } heap_info_t;
+
+  struct sq_post_dv_t {
+    uint64_t segments[16];
+    uint32_t current_sq;
+    uint16_t wqe_idx;
+  };
+
+  class State {
+   public:
+    ibv_qp_attr exp_qp_attr{};
+    uint64_t exp_attr_mask{};
+  };
+
+  class InitQPState : public State {
+   public:
+    InitQPState() {
+      exp_qp_attr.qp_state = IBV_QPS_INIT;
+      exp_qp_attr.qp_access_flags =
+          IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_LOCAL_WRITE |
+          IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_ATOMIC;
+
+      exp_attr_mask = IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT;
+    }
+  };
+
+  class RtrState : public State {
+   public:
+    RtrState() {
+      exp_qp_attr.qp_state = IBV_QPS_RTR;
+      exp_qp_attr.path_mtu = IBV_MTU_4096;
+      exp_qp_attr.ah_attr.sl = 1;
+      exp_qp_attr.max_dest_rd_atomic = 1;
+      exp_qp_attr.min_rnr_timer = 12;
+
+      exp_attr_mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU;
+    }
+  };
+
+  class RtsState : public State {
+   public:
+    RtsState() {
+      exp_qp_attr.qp_state = IBV_QPS_RTS;
+      exp_qp_attr.timeout = 14;
+      exp_qp_attr.retry_cnt = 7;
+      exp_qp_attr.rnr_retry = 7;
+      exp_qp_attr.max_rd_atomic = 1;
+
+      exp_attr_mask = IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT |
+                      IBV_QP_RNR_RETRY | IBV_QP_MAX_QP_RD_ATOMIC;
+    }
+  };
+
+  class QPInitAttr {
+   public:
+    explicit QPInitAttr(ibv_qp_cap cap) {
+      attr.cap = cap;
+      attr.sq_sig_all = 0;
+    }
+    ibv_qp_init_attr_ex attr{};
+  };
+
+ public:
+  Connection(GPUIBBackend* backend, int key_offset);
+
+  virtual ~Connection();
+
+  void initialize(int num_block);
+
+  void finalize();
+
+  virtual void post_wqes() = 0;
+
+  void reg_mr(void* ptr, size_t size, ibv_mr** mr, bool is_managed);
+
+  virtual void get_remote_conn(int* remote_conn) = 0;
+
+  unsigned total_number_connections();
+
+  virtual void initialize_rkey_handle(uint32_t** heap_rkey_handle,
+                                        ibv_mr* mr) = 0;
+
+  virtual void free_rkey_handle(uint32_t* heap_rkey_handle) = 0;
+
+  void initialize_gpu_policy(ConnectionImpl** conn, uint32_t* heap_rkey);
+
+  /*
+   * Populate a QueuePair for use on the GPU from the internal IB state.
+   */
+  void init_gpu_qp_from_connection(QueuePair* qp, int conn_num);
+
+ protected:
+  Connection() = default;
+
+  virtual InitQPState initqp(uint8_t port) = 0;
+
+  virtual RtrState rtr(dest_info_t* dest, uint8_t port) = 0;
+
+  virtual RtsState rts(dest_info_t* dest) = 0;
+
+  virtual QPInitAttr qpattr(ibv_qp_cap cap) = 0;
+
+  void init_qp_status(ibv_qp* qp, uint8_t port);
+
+  void change_status_rtr(ibv_qp* qp, dest_info_t* dest, uint8_t port);
+
+  void change_status_rts(ibv_qp* qp, dest_info_t* dest);
+
+  void create_qps(uint8_t port, int my_rank, ibv_port_attr* ib_port_att);
+
+  template <typename T>
+  void try_to_modify_qp(ibv_qp* qp, T state);
+
+  virtual void create_qps_1() = 0;
+
+  virtual void create_qps_2(int port, int my_rank,
+                              ibv_port_attr* ib_port_att) = 0;
+
+  virtual void create_qps_3(int port, ibv_qp* qp, int offset,
+                              ibv_port_attr* ib_port_att) = 0;
+
+  virtual ibv_qp* create_qp_0(ibv_context* context,
+                              ibv_qp_init_attr_ex* qp_attr) = 0;
+
+  virtual void allocate_dynamic_members(int num_block) = 0;
+
+  virtual void free_dynamic_members() = 0;
+
+  virtual void initialize_1(int port, int num_block) = 0;
+
+  virtual void initialize_wr_fields(ibv_send_wr* wr, ibv_ah* ah,
+                                    int dc_key) = 0;
+
+  virtual int get_sq_dv_offset(int pe_idx, int num_qps, int wg_idx) = 0;
+
+  void set_sq_dv(int num_block, int wg_idx, int pe_idx);
+
+  /*
+   * ibv interface functions must be static.
+   */
+  static void* buf_alloc(ibv_pd* pd, void* pd_context, size_t size,
+                         size_t alignment, uint64_t resource_type);
+
+  static void buf_release(ibv_pd* pd, void* pd_context, void* ptr,
+                          uint64_t resource_type);
+
+  void init_parent_domain_attr(ibv_parent_domain_init_attr* attr);
+
+  void set_rdma_seg(mlx5_wqe_raddr_seg* rdma, uint64_t address, uint32_t rkey);
+
+  uint64_t* get_address_sq(int i);
+
+  ibv_cq* create_cq(ibv_context* context, ibv_pd* pd, int cqe);
+
+  ibv_qp* create_qp(ibv_pd* pd, ibv_context* context,
+                    ibv_qp_init_attr_ex* qp_attr, ibv_cq* rcq);
+
+  /*
+   * TODO: Remove this eventually. Goal is to have backend delegate
+   * connection stuff to this class, while this class knows nothing about
+   * GPUs or backends.
+   */
+  GPUIBBackend* backend{nullptr};
+
+  uint32_t sq_size{1024};
+
+  ib_state_t* ib_state{nullptr};
+
+  const int key_offset{0};
+
+  sq_post_dv_t* sq_post_dv{nullptr};
+
+  std::vector<ibv_cq*> cqs;
+
+  std::vector<ibv_qp*> qps;
+
+  uint64_t counter_wqe{0};
+
+  static int use_gpu_mem;
+
+  static int coherent_cq;
+
+  int cq_use_gpu_mem{1};
+
+  int sq_use_gpu_mem{1};
+
+ private:
+  void init_shmem_handle();
+
+  void ib_init(ibv_device* ib_dev, uint8_t port);
+
+  char* requested_dev{nullptr};
+
+  ibv_device** dev_list{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_CONNECTION_HPP_

src/gpu_ib/connection_policy.cpp

@@ -0,0 +1,81 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "connection_policy.hpp"
+
+#include <infiniband/mlx5dv.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "dynamic_connection.hpp"
+#include "queue_pair.hpp"
+
+#ifdef DEBUG
+#define HIP_ENABLE_PRINTF 1
+#endif
+
+namespace rocshmem {
+
+RCConnectionImpl::RCConnectionImpl([[maybe_unused]] Connection* conn,
+                                   [[maybe_unused]] uint32_t* _vec_rkey) {}
+
+DCConnectionImpl::DCConnectionImpl(Connection* conn, uint32_t* _vec_rkey)
+    : vec_dct_num(static_cast<DynamicConnection*>(conn)->get_vec_dct_num()),
+      vec_rkey(_vec_rkey),
+      vec_lids(static_cast<DynamicConnection*>(conn)->get_vec_lids()) {}
+
+__device__ uint32_t RCConnectionImpl::getNumWqesImpl([
+    [maybe_unused]] uint8_t opcode) {
+  return 1;
+}
+
+__device__ uint32_t DCConnectionImpl::getNumWqesImpl(uint8_t opcode) {
+  // FIXME: We assume all threads in wave are performing ATOMIC ops.
+  // While this might be common, we do not have such restriction
+  // so need to be fixed.
+  // Since OFED 5.2, a DC segments uses 48bytes - so with or without
+  // atomic we need 2 wqes.
+  // return 2;
+  return (opcode == MLX5_OPCODE_ATOMIC_FA || opcode == MLX5_OPCODE_ATOMIC_CS)
+             ? 2
+             : 1;
+}
+
+__device__ bool RCConnectionImpl::updateConnectionSegmentImpl(
+    [[maybe_unused]] ib_mlx5_base_av_t* wqe, [[maybe_unused]] int pe) {
+  return false;
+}
+
+__device__ bool DCConnectionImpl::updateConnectionSegmentImpl(
+    ib_mlx5_base_av_t* wqe, int pe) {
+  wqe->dqp_dct = vec_dct_num[pe];
+  wqe->rlid = vec_lids[pe];
+  return true;
+}
+
+__device__ void RCConnectionImpl::setRkeyImpl([[maybe_unused]] uint32_t* rkey,
+                                              [[maybe_unused]] int pe) {}
+
+__device__ void DCConnectionImpl::setRkeyImpl(uint32_t* rkey, int pe) {
+  *rkey = vec_rkey[pe];
+}
+
+}  // namespace rocshmem

src/gpu_ib/connection_policy.hpp

@@ -0,0 +1,130 @@
+/******************************************************************************
+ * 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_GPU_IB_CONNECTION_POLICY_HPP_
+#define LIBRARY_SRC_GPU_IB_CONNECTION_POLICY_HPP_
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "infiniband_structs.hpp"
+
+namespace rocshmem {
+
+/*
+ * CRTP base class for connection type
+ */
+template <typename Derived>
+class ConnectionBase {
+ public:
+  /*
+   * Control segment WQE offset imposed by this connection type.
+   */
+  __device__ int wqeCntrlOffset() {
+    return static_cast<Derived*>(this)->wqeCntrlOffsetImpl();
+  }
+
+  /*
+   * Whether or not we need to force PE-level divergence when posting for
+   * this connection type.
+   */
+  __device__ bool forcePostDivergence() {
+    return static_cast<Derived*>(this)->forcePostDivergenceImpl();
+  }
+
+  /*
+   * Number of WQEs produced by this connection type for the given opcode.
+   */
+  __device__ uint32_t getNumWqes(uint8_t opcode) {
+    return static_cast<Derived*>(this)->getNumWqesImpl(opcode);
+  }
+
+  /*
+   * Updates the connection-specific segment in the SQ.
+   */
+  __device__ bool updateConnectionSegment(ib_mlx5_base_av_t* wqe, int pe) {
+    return static_cast<Derived*>(this)->updateConnectionSegmentImpl(wqe, pe);
+  }
+
+  /*
+   * Set the rkey based on this connection type.
+   */
+  __device__ void setRkey(uint32_t* rkey, int pe) {
+    static_cast<Derived*>(this)->setRkeyImpl(rkey, pe);
+  }
+};
+
+class Connection;
+
+/*
+ * Connection policy corresponding to an RC connection type.
+ */
+class RCConnectionImpl : public ConnectionBase<RCConnectionImpl> {
+ public:
+  RCConnectionImpl(Connection* conn, uint32_t* _vec_rkey);
+
+  __device__ int wqeCntrlOffsetImpl() { return 0; }
+
+  __device__ bool forcePostDivergenceImpl() { return true; }
+
+  __device__ uint32_t getNumWqesImpl(uint8_t opcode);
+
+  __device__ bool updateConnectionSegmentImpl(ib_mlx5_base_av_t* wqe, int pe);
+
+  __device__ void setRkeyImpl(uint32_t* rkey, int pe);
+};
+
+/*
+ * Connection policy corresponding to a DC connection type.
+ */
+class DCConnectionImpl : public ConnectionBase<DCConnectionImpl> {
+ public:
+  DCConnectionImpl(Connection* conn, uint32_t* _vec_rkey);
+
+  __device__ int wqeCntrlOffsetImpl() { return 1; }
+
+  __device__ bool forcePostDivergenceImpl() { return false; }
+
+  __device__ uint32_t getNumWqesImpl(uint8_t opcode);
+
+  __device__ bool updateConnectionSegmentImpl(ib_mlx5_base_av_t* wqe, int pe);
+
+  __device__ void setRkeyImpl(uint32_t* rkey, int pe);
+
+ private:
+  uint32_t* vec_dct_num{nullptr};
+
+  uint32_t* vec_rkey{nullptr};
+
+  uint16_t* vec_lids{nullptr};
+};
+
+/*
+ * Select which one of our connection policies to use at compile time.
+ */
+#ifdef USE_DC
+typedef DCConnectionImpl ConnectionImpl;
+#else
+typedef RCConnectionImpl ConnectionImpl;
+#endif
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_CONNECTION_POLICY_HPP_

src/gpu_ib/context_ib_device.cpp

@@ -0,0 +1,371 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "context_ib_device.hpp"
+
+#include <hip/hip_runtime.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "roc_shmem/roc_shmem.hpp"
+#include "../backend_type.hpp"
+#include "../context_incl.hpp"
+#include "backend_ib.hpp"
+#include "queue_pair.hpp"
+
+namespace rocshmem {
+
+__host__ GPUIBContext::GPUIBContext(Backend *backend, bool option, int idx)
+    : Context(backend, option) {
+  GPUIBBackend *b{static_cast<GPUIBBackend *>(backend)};
+  ctx_idx = idx;
+  networkImpl = b->networkImpl;
+  base_heap = b->heap.get_heap_bases().data();
+  networkImpl.networkHostInit(this, idx);
+
+  barrier_sync = b->barrier_sync;
+  ipcImpl_.ipc_bases = b->ipcImpl.ipc_bases;
+  ipcImpl_.shm_size = b->ipcImpl.shm_size;
+}
+
+__device__ void GPUIBContext::ctx_create() {
+  /* Nothing to do in the GPU_IB backend */
+  return;
+}
+
+/*
+ * TODO(bpotter): these will go in a policy class based on DC/RC.
+ * I am not completely sure at this point what else is needed in said class,
+ * so just leave them up here for now.
+ */
+__device__ __host__ QueuePair *GPUIBContext::getQueuePair(int pe) {
+  return networkImpl.getQueuePair(device_qp_proxy, pe);
+}
+
+__device__ __host__ int GPUIBContext::getNumQueuePairs() {
+  return networkImpl.getNumQueuePairs();
+}
+
+__device__ __host__ int GPUIBContext::getNumDest() {
+  return networkImpl.getNumDest();
+}
+
+__device__ void GPUIBContext::fence() {
+#ifdef USE_SINGLE_NODE
+  threadfence_system();
+#else
+
+  for (int k = 0; k < getNumDest(); k++) {
+    getQueuePair(k)->fence(k);
+  }
+
+  fence_.flush();
+#endif
+}
+
+__device__ void GPUIBContext::fence(int pe) {
+#ifdef USE_SINGLE_NODE
+  threadfence_system();
+#else
+  getQueuePair(pe)->fence(pe);
+  fence_.flush();
+#endif
+}
+
+__device__ void GPUIBContext::putmem_nbi(void *dest, const void *source,
+                                         size_t nelems, int pe) {
+  uint64_t L_offset = reinterpret_cast<char *>(dest) - base_heap[my_pe];
+
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy(ipcImpl_.ipc_bases[local_pe] + L_offset,
+                     const_cast<void *>(source), nelems);
+  } else {
+    bool must_send_message = wf_coal_.coalesce(pe, source, dest, &nelems);
+    if (!must_send_message) {
+      return;
+    }
+
+    auto *qp = getQueuePair(pe);
+    qp->put_nbi<THREAD>(base_heap[pe] + L_offset, source, nelems, pe, true);
+  }
+}
+
+__device__ void GPUIBContext::getmem_nbi(void *dest, const void *source,
+                                         size_t nelems, int pe) {
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy(dest, ipcImpl_.ipc_bases[local_pe] + L_offset, nelems);
+  } else {
+    bool must_send_message = wf_coal_.coalesce(pe, source, dest, &nelems);
+    if (!must_send_message) {
+      return;
+    }
+
+    auto *qp = getQueuePair(pe);
+    qp->get_nbi<THREAD>(base_heap[pe] + L_offset, dest, nelems, pe, true);
+  }
+}
+
+__device__ void GPUIBContext::quiet() {
+#ifdef USE_SINGLE_NODE
+  threadfence_system();
+  for (int pe = 0; pe < ipcImpl_.shm_size; pe++) {
+    if (pe != my_pe) {
+      ipcImpl_.zero_byte_read(pe);
+    }
+  }
+#else
+
+  for (int k = 0; k < getNumDest(); k++) {
+    getQueuePair(k)->quiet_single_heavy<THREAD>(k);
+  }
+  fence_.flush();
+
+#endif
+}
+
+__device__ void *GPUIBContext::shmem_ptr(const void *dest, int pe) {
+  void *ret = nullptr;
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    void *dst = const_cast<void *>(dest);
+    uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
+    int local_pe = pe % ipcImpl_.shm_size;
+    ret = ipcImpl_.ipc_bases[local_pe] + L_offset;
+  }
+  return ret;
+}
+
+__device__ void GPUIBContext::threadfence_system() {
+  int thread_id = get_flat_block_id();
+
+  if (thread_id % WF_SIZE == lowerID()) {
+#ifdef USE_SINGLE_NODE
+    // Flush current PE HDP
+    HdpPolicy::hdp_flush(
+        reinterpret_cast<unsigned int *>(networkImpl.hdp_address));
+
+    // Flush the rest of the HDPs
+    for (int pe = 0; pe < ipcImpl_.shm_size; pe++) {
+      auto target_address = networkImpl.hdp_address;
+      const int value = HdpPolicy::HDP_FLUSH_VAL;
+      if (pe != my_pe) {
+        const int value = HdpPolicy::HDP_FLUSH_VAL;
+        auto mapped_address =
+            shmem_ptr(reinterpret_cast<void *>(target_address), pe);
+        __hip_atomic_store(static_cast<int *>(mapped_address), value,
+                           __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);
+      }
+    }
+#else
+    getQueuePair(my_pe)->hdp_policy->flushCoherency();
+#endif
+  }
+
+  __threadfence_system();
+}
+
+__device__ void GPUIBContext::getmem(void *dest, const void *source,
+                                     size_t nelems, int pe) {
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy(dest, ipcImpl_.ipc_bases[local_pe] + L_offset, nelems);
+  } else {
+    bool must_send_message = wf_coal_.coalesce(pe, source, dest, &nelems);
+    if (!must_send_message) {
+      return;
+    }
+    auto *qp = getQueuePair(pe);
+    qp->get_nbi_cqe<THREAD>(base_heap[pe] + L_offset, dest, nelems, pe, true);
+    qp->quiet_single<THREAD>();
+  }
+  fence_.flush();
+}
+
+__device__ void GPUIBContext::putmem(void *dest, const void *source,
+                                     size_t nelems, int pe) {
+  uint64_t L_offset = reinterpret_cast<char *>(dest) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy(ipcImpl_.ipc_bases[local_pe] + L_offset,
+                     const_cast<void *>(source), nelems);
+
+    threadfence_system();
+    ipcImpl_.zero_byte_read(pe);
+  } else {
+    bool must_send_message = wf_coal_.coalesce(pe, source, dest, &nelems);
+    if (!must_send_message) {
+      return;
+    }
+    auto *qp = getQueuePair(pe);
+    qp->put_nbi_cqe<THREAD>(base_heap[pe] + L_offset, source, nelems, pe, true);
+    qp->quiet_single<THREAD>();
+  }
+  fence_.flush();
+}
+
+/******************************************************************************
+ ************************ WORKGROUP/WAVE-LEVEL RMA API ************************
+ *****************************************************************************/
+__device__ void GPUIBContext::putmem_nbi_wg(void *dest, const void *source,
+                                            size_t nelems, int pe) {
+  uint64_t L_offset = reinterpret_cast<char *>(dest) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wg(ipcImpl_.ipc_bases[local_pe] + L_offset,
+                        const_cast<void *>(source), nelems);
+  } else {
+    if (is_thread_zero_in_block()) {
+      auto *qp = getQueuePair(pe);
+      qp->put_nbi<WG>(base_heap[pe] + L_offset, source, nelems, pe, true);
+    }
+  }
+  __syncthreads();
+}
+
+__device__ void GPUIBContext::putmem_nbi_wave(void *dest, const void *source,
+                                              size_t nelems, int pe) {
+  uint64_t L_offset = reinterpret_cast<char *>(dest) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wave(ipcImpl_.ipc_bases[local_pe] + L_offset,
+                          const_cast<void *>(source), nelems);
+  } else {
+    if (is_thread_zero_in_wave()) {
+      auto *qp = getQueuePair(pe);
+      qp->put_nbi<WAVE>(base_heap[pe] + L_offset, source, nelems, pe, true);
+    }
+  }
+}
+
+__device__ void GPUIBContext::putmem_wg(void *dest, const void *source,
+                                        size_t nelems, int pe) {
+  uint64_t L_offset = reinterpret_cast<char *>(dest) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wg(ipcImpl_.ipc_bases[local_pe] + L_offset,
+                        const_cast<void *>(source), nelems);
+    __syncthreads();
+    threadfence_system();
+    ipcImpl_.zero_byte_read(pe);
+  } else {
+    auto *qp = getQueuePair(pe);
+    if (is_thread_zero_in_block()) {
+      qp->put_nbi_cqe<WG>(base_heap[pe] + L_offset, source, nelems, pe, true);
+    }
+    qp->quiet_single<WG>();
+  }
+  __syncthreads();
+  fence_.flush();
+}
+
+__device__ void GPUIBContext::putmem_wave(void *dest, const void *source,
+                                          size_t nelems, int pe) {
+  uint64_t L_offset = reinterpret_cast<char *>(dest) - base_heap[my_pe];
+  auto *qp = getQueuePair(pe);
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wave(ipcImpl_.ipc_bases[local_pe] + L_offset,
+                          const_cast<void *>(source), nelems);
+    threadfence_system();
+    ipcImpl_.zero_byte_read(pe);
+  } else {
+    if (is_thread_zero_in_wave()) {
+      qp->put_nbi_cqe<WAVE>(base_heap[pe] + L_offset, source, nelems, pe, true);
+    }
+    qp->quiet_single<WAVE>();
+  }
+  fence_.flush();
+}
+
+__device__ void GPUIBContext::getmem_wg(void *dest, const void *source,
+                                        size_t nelems, int pe) {
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - base_heap[my_pe];
+  auto *qp = getQueuePair(pe);
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wg(dest, ipcImpl_.ipc_bases[local_pe] + L_offset, nelems);
+  } else {
+    if (is_thread_zero_in_block()) {
+      qp->get_nbi_cqe<WG>(base_heap[pe] + L_offset, dest, nelems, pe, true);
+    }
+    qp->quiet_single<WG>();
+  }
+  __syncthreads();
+  fence_.flush();
+}
+
+__device__ void GPUIBContext::getmem_wave(void *dest, const void *source,
+                                          size_t nelems, int pe) {
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - base_heap[my_pe];
+  auto *qp = getQueuePair(pe);
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wave(dest, ipcImpl_.ipc_bases[local_pe] + L_offset,
+                          nelems);
+  } else {
+    if (is_thread_zero_in_wave()) {
+      qp->get_nbi_cqe<WAVE>(base_heap[pe] + L_offset, dest, nelems, pe, true);
+    }
+    qp->quiet_single<WAVE>();
+  }
+  fence_.flush();
+}
+
+__device__ void GPUIBContext::getmem_nbi_wg(void *dest, const void *source,
+                                            size_t nelems, int pe) {
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wg(dest, ipcImpl_.ipc_bases[local_pe] + L_offset, nelems);
+  } else {
+    if (is_thread_zero_in_block()) {
+      auto *qp = getQueuePair(pe);
+      qp->get_nbi<WG>(base_heap[pe] + L_offset, dest, nelems, pe, true);
+    }
+  }
+  __syncthreads();
+}
+
+__device__ void GPUIBContext::getmem_nbi_wave(void *dest, const void *source,
+                                              size_t nelems, int pe) {
+  const char *src_typed = reinterpret_cast<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - base_heap[my_pe];
+  if (ipcImpl_.isIpcAvailable(my_pe, pe)) {
+    int local_pe = pe % ipcImpl_.shm_size;
+    ipcImpl_.ipcCopy_wave(dest, ipcImpl_.ipc_bases[local_pe] + L_offset,
+                          nelems);
+  } else {
+    if (is_thread_zero_in_wave()) {
+      auto *qp = getQueuePair(pe);
+      qp->get_nbi<WAVE>(base_heap[pe] + L_offset, dest, nelems, pe, true);
+    }
+  }
+}
+
+}  // namespace rocshmem

src/gpu_ib/context_ib_device.hpp

@@ -0,0 +1,304 @@
+/******************************************************************************
+ * 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_GPU_IB_CONTEXT_IB_DEVICE_HPP_
+#define LIBRARY_SRC_GPU_IB_CONTEXT_IB_DEVICE_HPP_
+
+#include "../context.hpp"
+#include "memory_builder_policy.hpp"
+#include "network_policy.hpp"
+
+namespace rocshmem {
+
+class QueuePair;
+
+class GPUIBContext : public Context {
+ public:
+  __host__ GPUIBContext(Backend *b, bool option, int idx);
+
+  __device__ __host__ QueuePair *getQueuePair(int pe);
+
+  __device__ __host__ int getNumQueuePairs();
+
+  __device__ __host__ int getNumDest();
+
+  __device__ __attribute__((noinline)) void threadfence_system();
+
+  __device__ void ctx_create();
+
+  __device__ void ctx_destroy();
+
+  __device__ void putmem(void *dest, const void *source, size_t nelems, int pe);
+
+  __device__ void getmem(void *dest, const void *source, size_t nelems, int pe);
+
+  __device__ void putmem_nbi(void *dest, const void *source, size_t nelems,
+                             int pe);
+
+  __device__ void getmem_nbi(void *dest, const void *source, size_t size,
+                             int pe);
+
+  __device__ void fence();
+
+  __device__ void fence(int pe);
+
+  __device__ void quiet();
+
+  __device__ void *shmem_ptr(const void *dest, int pe);
+
+  __device__ void barrier_all();
+
+  __device__ void sync_all();
+
+  __device__ void sync(roc_shmem_team_t team);
+
+  template <typename T>
+  __device__ void amo_add(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_set(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_swap(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_and(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_and(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_or(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_or(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_xor(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_xor(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ void amo_cas(void *dst, T value, T cond, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_add(void *dst, T value, int pe);
+
+  template <typename T>
+  __device__ T amo_fetch_cas(void *dst, T value, T cond, int pe);
+
+  template <typename T>
+  __device__ void p(T *dest, T value, int pe);
+
+  template <typename T>
+  __device__ T g(const T *source, int pe);
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __device__ void to_all(T *dest, const T *source, int nreduce, int PE_start,
+                         int logPE_stride, int PE_size, T *pWrk,
+                         long *pSync);  // NOLINT(runtime/int)
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __device__ void to_all(roc_shmem_team_t team, T *dest, const T *source,
+                         int nreduce);
+
+  template <typename T>
+  __device__ void put(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_nbi(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_nbi(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void broadcast(roc_shmem_team_t team, T *dest, const T *source,
+                            int nelems, int pe_root);
+
+  template <typename T>
+  __device__ void broadcast(T *dest, const T *source, int nelems, int pe_root,
+                            int pe_start, int log_pe_stride, int pe_size,
+                            long *p_sync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __device__ void alltoall(roc_shmem_team_t team, T *dest, const T *source,
+                           int nelems);
+
+  template <typename T>
+  __device__ void alltoall_broadcast(roc_shmem_team_t team, T *dest,
+                                     const T *source, int nelems);
+
+  template <typename T>
+  __device__ void alltoall_brucks(roc_shmem_team_t team, T *dest,
+                                  const T *source, int nelems);
+
+  template <typename T>
+  __device__ void alltoall_gcen(roc_shmem_team_t team, T *dest, const T *source,
+                                int nelems);
+
+  template <typename T>
+  __device__ void alltoall_gcen2(roc_shmem_team_t team, T *dest,
+                                 const T *source, int nelems);
+
+  template <typename T>
+  __device__ void fcollect(roc_shmem_team_t team, T *dest, const T *source,
+                           int nelems);
+
+  template <typename T>
+  __device__ void fcollect_broadcast(roc_shmem_team_t team, T *dest,
+                                     const T *source, int nelems);
+
+  template <typename T>
+  __device__ void fcollect_brucks(roc_shmem_team_t team, T *dest,
+                                  const T *source, int nelems);
+
+  template <typename T>
+  __device__ void fcollect_gcen(roc_shmem_team_t team, T *dest, const T *source,
+                                int nelems);
+
+  template <typename T>
+  __device__ void fcollect_gcen2(roc_shmem_team_t team, T *dest,
+                                 const T *source, int nelems);
+
+  __device__ void putmem_wg(void *dest, const void *source, size_t nelems,
+                            int pe);
+
+  __device__ void getmem_wg(void *dest, const void *source, size_t nelems,
+                            int pe);
+
+  __device__ void putmem_nbi_wg(void *dest, const void *source, size_t nelems,
+                                int pe);
+
+  __device__ void getmem_nbi_wg(void *dest, const void *source, size_t size,
+                                int pe);
+
+  __device__ void putmem_wave(void *dest, const void *source, size_t nelems,
+                              int pe);
+
+  __device__ void getmem_wave(void *dest, const void *source, size_t nelems,
+                              int pe);
+
+  __device__ void putmem_nbi_wave(void *dest, const void *source, size_t nelems,
+                                  int pe);
+
+  __device__ void getmem_nbi_wave(void *dest, const void *source, size_t size,
+                                  int pe);
+
+  template <typename T>
+  __device__ void put_wg(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_nbi_wg(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_wg(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_nbi_wg(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_wave(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void put_nbi_wave(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_wave(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __device__ void get_nbi_wave(T *dest, const T *source, size_t nelems, int pe);
+
+ private:
+  template <typename T, ROC_SHMEM_OP Op>
+  __device__ void internal_direct_allreduce(
+      T *dst, const T *src, int nelems, int PE_start, int logPE_stride,
+      int PE_size, T *pWrk,
+      long *pSync);  // NOLINT(runtime/int)
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __device__ void internal_ring_allreduce(T *dst, const T *src, int nelems,
+                                          int PE_start, int logPE_stride,
+                                          int PE_size, T *pWrk,
+                                          long *pSync,  // NOLINT(runtime/int)
+                                          int n_seg, int seg_size,
+                                          int chunk_size);
+
+  template <typename T>
+  __device__ void internal_put_broadcast(T *dst, const T *src, int nelems,
+                                         int pe_root, int PE_start,
+                                         int logPE_stride, int PE_size,
+                                         long *pSync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __device__ void internal_get_broadcast(T *dst, const T *src, int nelems,
+                                         int pe_root,
+                                         long *pSync);  // NOLINT(runtime/int)
+
+  __device__ void internal_direct_barrier(int pe, int PE_start, int stride,
+                                          int n_pes, int64_t *pSync);
+
+  __device__ void internal_atomic_barrier(int pe, int PE_start, int stride,
+                                          int n_pes, int64_t *pSync);
+
+  __device__ void internal_sync(int pe, int PE_start, int stride, int PE_size,
+                                int64_t *pSync);
+
+  __device__ void quiet_single(int cq_num);
+
+ public:
+  /*
+   * Collection of queue pairs that are currently checked out by this
+   * context from GPUIBBackend.
+   */
+  // FIXME: keep it private and destroy in destructor for better
+  // encapsulation.
+  QueuePair *device_qp_proxy{nullptr};
+
+  /*
+   * Array of char * pointers corresponding to the heap base pointers VA for
+   * each PE that we can communicate with.
+   */
+  char *const *base_heap{nullptr};
+
+  /*
+   * Buffer used to store the results of a *_g operation. These ops do not
+   * provide a destination buffer, so the runtime must manage one.
+   */
+  char *g_ret{nullptr};
+
+  NetworkImpl networkImpl{};
+
+  /*
+   * Temporary scratchpad memory used by internal barrier algorithms.
+   */
+  int64_t *barrier_sync{nullptr};
+
+  int ctx_idx;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_CONTEXT_IB_DEVICE_HPP_

src/gpu_ib/context_ib_device_coll.cpp

@@ -0,0 +1,122 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "roc_shmem/roc_shmem.hpp"
+#include "../context_incl.hpp"
+#include "context_ib_tmpl_device.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+__device__ void GPUIBContext::internal_direct_barrier(int pe, int PE_start,
+                                                      int stride, int n_pes,
+                                                      int64_t *pSync) {
+  int64_t flag_val = 1;
+  if (pe == PE_start) {
+    // Go through all PE offsets (except current offset = 0)
+    // and wait until they all reach
+    for (size_t i = 1; i < n_pes; i++) {
+      wait_until(&pSync[i], ROC_SHMEM_CMP_EQ, flag_val);
+      pSync[i] = ROC_SHMEM_SYNC_VALUE;
+    }
+    threadfence_system();
+    // Announce to other PEs that all have reached
+    for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) {
+      put_nbi(&pSync[0], &flag_val, 1, j);
+    }
+
+  } else {
+    // Mark current PE offset as reached
+    size_t pe_offset = (pe - PE_start) / stride;
+    put_nbi(&pSync[pe_offset], &flag_val, 1, PE_start);
+    wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val);
+    pSync[0] = ROC_SHMEM_SYNC_VALUE;
+    threadfence_system();
+  }
+}
+
+__device__ void GPUIBContext::internal_atomic_barrier(int pe, int PE_start,
+                                                      int stride, int n_pes,
+                                                      int64_t *pSync) {
+  int64_t flag_val = 1;
+  if (pe == PE_start) {
+    wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, (int64_t)(n_pes - 1));
+    pSync[0] = ROC_SHMEM_SYNC_VALUE;
+    threadfence_system();
+    for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) {
+      put_nbi(&pSync[0], &flag_val, 1, j);
+    }
+  } else {
+    amo_add<int64_t>(&pSync[0], flag_val, PE_start);
+    wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val);
+    pSync[0] = ROC_SHMEM_SYNC_VALUE;
+    threadfence_system();
+  }
+}
+
+// Uses PE values that are relative to world
+__device__ void GPUIBContext::internal_sync(int pe, int PE_start, int stride,
+                                            int PE_size, int64_t *pSync) {
+  __syncthreads();
+  if (is_thread_zero_in_block()) {
+    if (PE_size < 64) {
+      internal_direct_barrier(pe, PE_start, stride, PE_size, pSync);
+    } else {
+      internal_atomic_barrier(pe, PE_start, stride, PE_size, pSync);
+    }
+  }
+  __threadfence();
+  __syncthreads();
+}
+
+__device__ void GPUIBContext::sync(roc_shmem_team_t team) {
+  GPUIBTeam *team_obj = reinterpret_cast<GPUIBTeam *>(team);
+
+  double dbl_log_pe_stride = team_obj->tinfo_wrt_world->log_stride;
+  int log_pe_stride = static_cast<int>(dbl_log_pe_stride);
+  /**
+   * Ensure that the stride is a multiple of 2 for GPU_IB.
+   * TODO: enable GPU_IB to work with non-powers-of-2 strides
+   * and remove this assert.
+   */
+  assert((dbl_log_pe_stride - log_pe_stride) == 0);
+
+  int pe = team_obj->my_pe_in_world;
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_stride = (1 << log_pe_stride);
+  int pe_size = team_obj->num_pes;
+  internal_sync(pe, pe_start, pe_stride, pe_size, barrier_sync);
+}
+
+__device__ void GPUIBContext::sync_all() {
+  internal_sync(my_pe, 0, 1, num_pes, barrier_sync);
+}
+
+__device__ void GPUIBContext::barrier_all() {
+  if (is_thread_zero_in_block()) {
+    quiet();
+  }
+  sync_all();
+  __syncthreads();
+}
+
+}  // namespace rocshmem

src/gpu_ib/context_ib_host.cpp

@@ -0,0 +1,85 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "context_ib_host.hpp"
+
+#include <mpi.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "../backend_type.hpp"
+#include "../context_incl.hpp"
+#include "backend_ib.hpp"
+#include "../host/host.hpp"
+
+namespace rocshmem {
+
+__host__ GPUIBHostContext::GPUIBHostContext(Backend *backend,
+                                            [[maybe_unused]] int64_t options)
+    : Context(backend, true) {
+  GPUIBBackend *b{static_cast<GPUIBBackend *>(backend)};
+
+  host_interface = b->host_interface;
+
+  context_window_info = host_interface->acquire_window_context();
+}
+
+__host__ GPUIBHostContext::~GPUIBHostContext() {
+  host_interface->release_window_context(context_window_info);
+}
+
+__host__ void GPUIBHostContext::putmem_nbi(void *dest, const void *source,
+                                           size_t nelems, int pe) {
+  host_interface->putmem_nbi(dest, source, nelems, pe, context_window_info);
+}
+
+__host__ void GPUIBHostContext::getmem_nbi(void *dest, const void *source,
+                                           size_t nelems, int pe) {
+  host_interface->getmem_nbi(dest, source, nelems, pe, context_window_info);
+}
+
+__host__ void GPUIBHostContext::putmem(void *dest, const void *source,
+                                       size_t nelems, int pe) {
+  host_interface->putmem(dest, source, nelems, pe, context_window_info);
+}
+
+__host__ void GPUIBHostContext::getmem(void *dest, const void *source,
+                                       size_t nelems, int pe) {
+  host_interface->getmem(dest, source, nelems, pe, context_window_info);
+}
+
+__host__ void GPUIBHostContext::fence() {
+  host_interface->fence(context_window_info);
+}
+
+__host__ void GPUIBHostContext::quiet() {
+  host_interface->quiet(context_window_info);
+}
+
+__host__ void GPUIBHostContext::sync_all() {
+  host_interface->sync_all(context_window_info);
+}
+
+__host__ void GPUIBHostContext::barrier_all() {
+  host_interface->barrier_all(context_window_info);
+}
+
+}  // namespace rocshmem

src/gpu_ib/context_ib_host.hpp

@@ -0,0 +1,149 @@
+/******************************************************************************
+ * 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_GPU_IB_CONTEXT_IB_HOST_HPP_
+#define LIBRARY_SRC_GPU_IB_CONTEXT_IB_HOST_HPP_
+
+#include "../context.hpp"
+
+namespace rocshmem {
+
+class GPUIBHostContext : public Context {
+ public:
+  __host__ GPUIBHostContext(Backend *b, int64_t options);
+
+  __host__ ~GPUIBHostContext();
+
+  template <typename T>
+  __host__ void p(T *dest, T value, int pe);
+
+  template <typename T>
+  __host__ T g(const T *source, int pe);
+
+  template <typename T>
+  __host__ void put(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __host__ void get(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __host__ void put_nbi(T *dest, const T *source, size_t nelems, int pe);
+
+  template <typename T>
+  __host__ void get_nbi(T *dest, const T *source, size_t nelems, int pe);
+
+  __host__ void putmem(void *dest, const void *source, size_t nelems, int pe);
+
+  __host__ void getmem(void *dest, const void *source, size_t nelems, int pe);
+
+  __host__ void putmem_nbi(void *dest, const void *source, size_t nelems,
+                           int pe);
+
+  __host__ void getmem_nbi(void *dest, const void *source, size_t size, int pe);
+
+  template <typename T>
+  __host__ void amo_add(void *dst, T value, int pe);
+
+  template <typename T>
+  __host__ void amo_cas(void *dst, T value, T cond, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_add(void *dst, T value, int pe);
+
+  template <typename T>
+  __host__ T amo_fetch_cas(void *dst, T value, T cond, int pe);
+
+  __host__ void fence();
+
+  __host__ void quiet();
+
+  __host__ void barrier_all();
+
+  __host__ void sync_all();
+
+  template <typename T>
+  __host__ void broadcast(T *dest, const T *source, int nelems, int pe_root,
+                          int pe_start, int log_pe_stride, int pe_size,
+                          long *p_sync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __host__ void broadcast(roc_shmem_team_t team, T *dest, const T *source,
+                          int nelems, int pe_root);
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all(T *dest, const T *source, int nreduce, int pe_start,
+                       int log_pe_stride, int pe_size, T *p_wrk,
+                       long *p_sync);  // NOLINT(runtime/int)
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all(roc_shmem_team_t team, T *dest, const T *source,
+                       int nreduce);
+
+  template <typename T>
+  __host__ void wait_until(T *ptr, roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ size_t wait_until_any(T* ptr, size_t nelems,
+                                 const int *status,
+                                 roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ void wait_until_all(T* ptr, size_t nelems,
+                               const int *status,
+                               roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ size_t wait_until_some(T* ptr, size_t nelems,
+                                size_t* indices,
+                                const int *status,
+                                roc_shmem_cmps cmp, T val);
+
+  template <typename T>
+  __host__ void wait_until_all_vector(T* ptr, size_t nelems,
+                                      const int *status,
+                                      roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __host__ size_t wait_until_any_vector(T* ptr, size_t nelems,
+                                        const int *status,
+                                        roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __host__ size_t wait_until_some_vector(T* ptr, size_t nelems,
+                                         size_t* indices,
+                                         const int *status,
+                                         roc_shmem_cmps cmp, T* vals);
+
+  template <typename T>
+  __host__ int test(T *ptr, roc_shmem_cmps cmp, T val);
+
+ public:
+  /* Pointer to the backend's host interface */
+  HostInterface *host_interface{nullptr};
+
+  /* An MPI Window implements a context */
+  WindowInfo *context_window_info{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_CONTEXT_IB_HOST_HPP_

src/gpu_ib/context_ib_tmpl_host.hpp

@@ -0,0 +1,173 @@
+/******************************************************************************
+ * 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_GPU_IB_CONTEXT_IB_TMPL_HOST_HPP_
+#define LIBRARY_SRC_GPU_IB_CONTEXT_IB_TMPL_HOST_HPP_
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "../host/host_templates.hpp"
+
+namespace rocshmem {
+
+template <typename T>
+__host__ void GPUIBHostContext::p(T *dest, T value, int pe) {
+  host_interface->p<T>(dest, value, pe, context_window_info);
+}
+
+template <typename T>
+__host__ T GPUIBHostContext::g(const T *source, int pe) {
+  return host_interface->g<T>(source, pe, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::put(T *dest, const T *source, size_t nelems,
+                                    int pe) {
+  host_interface->put<T>(dest, source, nelems, pe, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::get(T *dest, const T *source, size_t nelems,
+                                    int pe) {
+  host_interface->get<T>(dest, source, nelems, pe, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::put_nbi(T *dest, const T *source, size_t nelems,
+                                        int pe) {
+  host_interface->put_nbi<T>(dest, source, nelems, pe, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::get_nbi(T *dest, const T *source, size_t nelems,
+                                        int pe) {
+  host_interface->get_nbi<T>(dest, source, nelems, pe, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::amo_add(void *dst, T value, int pe) {
+  host_interface->amo_add(dst, value, pe, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::amo_cas(void *dst, T value, T cond, int pe) {
+  host_interface->amo_cas(dst, value, cond, pe, context_window_info);
+}
+
+template <typename T>
+__host__ T GPUIBHostContext::amo_fetch_add(void *dst, T value, int pe) {
+  return host_interface->amo_fetch_add(dst, value, pe, context_window_info);
+}
+
+template <typename T>
+__host__ T GPUIBHostContext::amo_fetch_cas(void *dst, T value, T cond, int pe) {
+  return host_interface->amo_fetch_cas(dst, value, cond, pe,
+                                       context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::broadcast(
+    T *dest, const T *source, int nelems, int pe_root, int pe_start,
+    int log_pe_stride, int pe_size,
+    long *p_sync) {  // NOLINT(runtime/int)
+  host_interface->broadcast<T>(dest, source, nelems, pe_root, pe_start,
+                               log_pe_stride, pe_size, p_sync);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::broadcast(roc_shmem_team_t team, T *dest,
+                                          const T *source, int nelems,
+                                          int pe_root) {
+  host_interface->broadcast<T>(team, dest, source, nelems, pe_root);
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void GPUIBHostContext::to_all(T *dest, const T *source, int nreduce,
+                                       int pe_start, int log_pe_stride,
+                                       int pe_size, T *p_wrk,
+                                       long *p_sync) {  // NOLINT(runtime/int)
+  host_interface->to_all<T, Op>(dest, source, nreduce, pe_start, log_pe_stride,
+                                pe_size, p_wrk, p_sync);
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void GPUIBHostContext::to_all(roc_shmem_team_t team, T *dest,
+                                       const T *source, int nreduce) {
+  host_interface->to_all<T, Op>(team, dest, source, nreduce);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::wait_until(T *ptr, roc_shmem_cmps cmp, T val) {
+  host_interface->wait_until<T>(ptr, cmp, val, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::wait_until_all(T *ptr, size_t nelems,
+                                               const int* status,
+                                               roc_shmem_cmps cmp, T val) {
+  host_interface->wait_until_all<T>(ptr, nelems, status, cmp, val, context_window_info);
+}
+
+template <typename T>
+__host__ size_t GPUIBHostContext::wait_until_any(T *ptr, size_t nelems,
+                                               const int* status,
+                                               roc_shmem_cmps cmp, T val) {
+  return host_interface->wait_until_any<T>(ptr, nelems, status, cmp, val, context_window_info);
+}
+
+template <typename T>
+__host__ size_t GPUIBHostContext::wait_until_some(T *ptr, size_t nelems,
+                                                 size_t* indices,
+                                                 const int* status,
+                                                 roc_shmem_cmps cmp, T val) {
+  return host_interface->wait_until_some<T>(ptr, nelems, indices, status, cmp, val, context_window_info);
+}
+
+template <typename T>
+__host__ void GPUIBHostContext::wait_until_all_vector(T *ptr, size_t nelems,
+                                                      const int* status,
+                                                      roc_shmem_cmps cmp, T* vals) {
+  host_interface->wait_until_all_vector<T>(ptr, nelems, status, cmp, vals, context_window_info);
+}
+
+template <typename T>
+__host__ size_t GPUIBHostContext::wait_until_any_vector(T *ptr, size_t nelems,
+                                                        const int* status,
+                                                        roc_shmem_cmps cmp, T* vals) {
+  return host_interface->wait_until_any_vector<T>(ptr, nelems, status, cmp, vals, context_window_info);
+}
+
+template <typename T>
+__host__ size_t GPUIBHostContext::wait_until_some_vector(T *ptr, size_t nelems,
+                                                         size_t* indices,
+                                                         const int* status,
+                                                         roc_shmem_cmps cmp, T* vals) {
+  return host_interface->wait_until_some_vector<T>(ptr, nelems, indices, status, cmp, vals, context_window_info);
+}
+
+template <typename T>
+__host__ int GPUIBHostContext::test(T *ptr, roc_shmem_cmps cmp, T val) {
+  return host_interface->test<T>(ptr, cmp, val, context_window_info);
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_CONTEXT_IB_TMPL_HOST_HPP_

src/gpu_ib/debug.cpp

@@ -0,0 +1,39 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "roc_shmem/debug.hpp"
+
+#include "qe_dumper.hpp"
+
+namespace rocshmem {
+
+void debug_print_cq(int dest_pe, int src_wg, int cqe_index) {
+  QeDumper dumper(dest_pe, src_wg, cqe_index);
+  dumper.dump_cq();
+}
+
+void debug_print_sq(int dest_pe, int src_wg, int wqe_index) {
+  QeDumper dumper(dest_pe, src_wg, wqe_index);
+  dumper.dump_sq();
+}
+
+}  // namespace rocshmem

src/gpu_ib/dynamic_connection.cpp

@@ -0,0 +1,381 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "dynamic_connection.hpp"
+
+#include <mpi.h>
+
+#include "backend_ib.hpp"
+
+namespace rocshmem {
+
+DynamicConnection::DynamicConnection(GPUIBBackend* b) : Connection(b, 4) {
+  char* value = nullptr;
+
+  if ((value = getenv("ROC_SHMEM_NUM_DCIs"))) {
+    num_dcis = atoi(value);
+  }
+
+  if ((value = getenv("ROC_SHMEM_NUM_DCT"))) {
+    num_dct = atoi(value);
+  }
+}
+
+DynamicConnection::~DynamicConnection() {
+  CHECK_HIP(hipFree(vec_lids));
+  CHECK_HIP(hipFree(vec_dct_num));
+}
+
+ibv_qp_init_attr_ex DynamicConnection::dct_qp_init_attr(
+    ibv_cq* cq, ibv_srq* srq, [[maybe_unused]] uint8_t port) const {
+  ibv_qp_init_attr_ex attr{};
+
+  attr.comp_mask = IBV_QP_INIT_ATTR_PD;
+  attr.pd = ib_state->pd;
+  attr.recv_cq = cq;
+  attr.send_cq = cq;
+  attr.srq = srq;
+  attr.qp_type = IBV_QPT_DRIVER;
+
+  return attr;
+}
+
+mlx5dv_qp_init_attr DynamicConnection::dct_dv_init_attr() {
+  mlx5dv_qp_init_attr dv_attr{};
+  dv_attr.comp_mask = MLX5DV_QP_INIT_ATTR_MASK_DC;
+  dv_attr.dc_init_attr.dc_type = MLX5DV_DCTYPE_DCT;
+  dv_attr.dc_init_attr.dct_access_key = DC_IB_KEY;
+
+  return dv_attr;
+}
+
+Connection::InitQPState DynamicConnection::initqp(uint8_t port) {
+  InitQPState initqp{};
+
+  initqp.exp_qp_attr.port_num = port;
+  initqp.exp_qp_attr.pkey_index = 0;
+  initqp.exp_qp_attr.qp_access_flags = 0;
+
+  return initqp;
+}
+
+Connection::RtrState DynamicConnection::rtr([[maybe_unused]] dest_info_t* dest,
+                                            uint8_t port) {
+  RtrState rtr{};
+
+  rtr.exp_qp_attr.ah_attr.is_global = 1;
+  rtr.exp_qp_attr.ah_attr.port_num = port;
+
+  rtr.exp_qp_attr.max_dest_rd_atomic = 0;
+  rtr.exp_qp_attr.min_rnr_timer = 0;
+
+  return rtr;
+}
+
+Connection::RtsState DynamicConnection::rts([
+    [maybe_unused]] dest_info_t* dest) {
+  RtsState rts{};
+  rts.exp_attr_mask |= IBV_QP_SQ_PSN;
+  return rts;
+}
+
+void DynamicConnection::connect_dci(ibv_qp* qp, uint8_t port) {
+  init_qp_status(qp, port);
+  change_status_rtr(qp, nullptr, port);
+  change_status_rts(qp, nullptr);
+}
+
+/*
+ * create a DCT and get is to ready state
+ */
+void DynamicConnection::create_dct(int32_t* dct_num, ibv_cq* cq, ibv_srq* srq,
+                                     uint8_t port) {
+  auto init_attr = dct_qp_init_attr(cq, srq, port);
+  auto dv_attr = dct_dv_init_attr();
+  auto dct = mlx5dv_create_qp(ib_state->context, &init_attr, &dv_attr);
+
+  if (dct == nullptr) {
+    printf("Failed to create dct \n");
+    abort();
+  }
+
+  ibv_qp_attr qp_attr{};
+  qp_attr.qp_state = IBV_QPS_INIT;
+  qp_attr.port_num = port;
+  qp_attr.qp_access_flags = IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_LOCAL_WRITE |
+                            IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_ATOMIC;
+
+  int attr_mask =
+      IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS;
+
+  int ret = ibv_modify_qp(dct, &qp_attr, attr_mask);
+  if (ret) {
+    abort();
+  }
+
+  qp_attr.qp_state = IBV_QPS_RTR;
+  qp_attr.path_mtu = IBV_MTU_4096;
+  qp_attr.min_rnr_timer = 7;
+  qp_attr.ah_attr.is_global = 1;
+  qp_attr.ah_attr.grh.hop_limit = 1;
+  qp_attr.ah_attr.grh.traffic_class = 0;
+  qp_attr.ah_attr.grh.sgid_index = 0;
+  qp_attr.ah_attr.port_num = port;
+
+  attr_mask = IBV_QP_STATE | IBV_QP_MIN_RNR_TIMER | IBV_QP_AV | IBV_QP_PATH_MTU;
+
+  ret = ibv_modify_qp(dct, &qp_attr, attr_mask);
+  if (ret) {
+    abort();
+  }
+
+  *dct_num = dct->qp_num;
+}
+
+/*
+ * @brief create a qp (DCI qp) using DEVX
+ */
+ibv_qp* DynamicConnection::create_qp_0(ibv_context* context,
+                                       ibv_qp_init_attr_ex* qp_attr) {
+  ibv_qp* qp;
+  qp_attr->qp_type = IBV_QPT_DRIVER;
+
+  mlx5dv_qp_init_attr dv_attr{};
+  dv_attr.comp_mask = MLX5DV_QP_INIT_ATTR_MASK_DC;
+  dv_attr.dc_init_attr.dc_type = MLX5DV_DCTYPE_DCI;
+  dv_attr.dc_init_attr.dct_access_key = DC_IB_KEY;
+
+  qp = mlx5dv_create_qp(context, qp_attr, &dv_attr);
+
+  return qp;
+}
+
+void DynamicConnection::create_qps_1() {
+  ibv_srq_init_attr srq_init_attr{};
+  srq_init_attr.attr.max_wr = 1;
+  srq_init_attr.attr.max_sge = 1;
+
+  srq = ibv_create_srq(ib_state->pd, &srq_init_attr);
+  if (!srq) {
+    abort();
+  }
+
+  dct_cq = ibv_create_cq(ib_state->context, 100, nullptr, nullptr, 0);
+  if (!dct_cq) {
+    abort();
+  }
+}
+
+void DynamicConnection::create_qps_2(int port, int my_rank,
+                                       ibv_port_attr* ib_port_att) {
+  for (int i = 0; i < num_dct; i++) {
+    int32_t dct_num;
+    create_dct(&dct_num, dct_cq, srq, port);
+    dct_num = htobe32(dct_num);
+    dcts_num[my_rank * num_dct + i] = dct_num;
+  }
+  lids[my_rank] = htobe16(ib_port_att->lid);
+}
+
+void DynamicConnection::create_qps_3(
+    int port, ibv_qp* qp, [[maybe_unused]] int offset,
+    [[maybe_unused]] ibv_port_attr* ib_port_att) {
+  return connect_dci(qp, port);
+}
+
+void DynamicConnection::get_remote_conn(int* remote_conn) {
+  *remote_conn = num_dcis;
+}
+
+void DynamicConnection::allocate_dynamic_members([
+    [maybe_unused]] int num_wg) {
+  size_t num_pes_size_bytes = sizeof(uint16_t) * backend->num_pes;
+  lids = reinterpret_cast<uint16_t*>(malloc(num_pes_size_bytes));
+  if (lids == nullptr) {
+    abort();
+  }
+
+  size_t num_dcts = num_dct * backend->num_pes;
+  size_t num_dcts_size_bytes = sizeof(uint32_t) * num_dcts;
+  dcts_num = reinterpret_cast<uint32_t*>(malloc(num_dcts_size_bytes));
+  if (dcts_num == nullptr) {
+    abort();
+  }
+}
+
+/*
+ * get the wqe_av information from tyhe ibv_ah
+ * rely on DEVX to extract the AV. We use the AV to create
+ * the DC segment
+ */
+void DynamicConnection::dc_get_av(ibv_ah* ah, mlx5_wqe_av* mlx5_av) {
+  mlx5dv_obj dv;
+  mlx5dv_ah dah;
+
+  dv.ah.in = ah;
+  dv.ah.out = &dah;
+  mlx5dv_init_obj(&dv, MLX5DV_OBJ_AH);
+
+  memcpy(mlx5_av, dah.av, sizeof(mlx5_wqe_av));
+}
+
+void DynamicConnection::free_dynamic_members() {
+  free(lids);
+  free(dcts_num);
+}
+
+void DynamicConnection::initialize_1(int port, [[maybe_unused]] int num_wg) {
+  MPI_Allgather(MPI_IN_PLACE, sizeof(int32_t) * num_dct, MPI_CHAR, dcts_num,
+                sizeof(int32_t) * num_dct, MPI_CHAR, backend->thread_comm);
+
+  MPI_Allgather(MPI_IN_PLACE, sizeof(int16_t), MPI_CHAR, lids, sizeof(int16_t),
+                MPI_CHAR, backend->thread_comm);
+
+  hipStream_t stream;
+  CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+  CHECK_HIP(hipMalloc(reinterpret_cast<void**>(&vec_dct_num),
+                      sizeof(int32_t) * num_dct * backend->num_pes));
+
+  CHECK_HIP(hipMemcpyAsync(vec_dct_num, dcts_num,
+                           sizeof(int32_t) * num_dct * backend->num_pes,
+                           hipMemcpyHostToDevice, stream));
+
+  CHECK_HIP(hipMalloc(reinterpret_cast<void**>(&vec_lids),
+                      sizeof(int16_t) * backend->num_pes));
+
+  CHECK_HIP(hipMemcpyAsync(vec_lids, lids, sizeof(int16_t) * backend->num_pes,
+                           hipMemcpyHostToDevice, stream));
+
+  struct ibv_ah_attr ah_attr;
+  memset(&ah_attr, 0, sizeof(ah_attr));
+
+  ah_attr.is_global = 1;
+  ah_attr.dlid = ib_state->portinfo.lid;
+  ah_attr.sl = 1;
+  ah_attr.src_path_bits = 0;
+  ah_attr.port_num = port;
+
+  ah = ibv_create_ah(ib_state->pd, &ah_attr);
+  if (ah == nullptr) {
+    abort();
+  }
+
+  dc_get_av(ah, &mlx5_av);
+
+  CHECK_HIP(hipStreamSynchronize(stream));
+  CHECK_HIP(hipStreamDestroy(stream));
+}
+
+void DynamicConnection::initialize_rkey_handle(uint32_t** heap_rkey_handle,
+                                                 ibv_mr* mr) {
+  CHECK_HIP(hipMalloc(heap_rkey_handle, sizeof(uint32_t) * backend->num_pes));
+  (*heap_rkey_handle)[backend->my_pe] = htobe32(mr->rkey);
+}
+
+void DynamicConnection::free_rkey_handle(uint32_t* heap_rkey_handle) {
+  CHECK_HIP(hipFree(heap_rkey_handle));
+}
+
+Connection::QPInitAttr DynamicConnection::qpattr(ibv_qp_cap cap) {
+  QPInitAttr qpattr(cap);
+  return qpattr;
+}
+
+/*
+ * Create and write the DC segment to SQ.
+ * We get all the info needed from the mlx5_wqe_av that we extract from ibv_ah.
+ */
+void DynamicConnection::set_dgram_seg(mlx5_wqe_datagram_seg* dc_seg,
+                                      uint64_t dc_key, uint32_t dct_num,
+                                      uint8_t ext, mlx5_wqe_av* mlx5_av) {
+  dc_seg->av.key.dc_key = htobe64(dc_key);
+  dc_seg->av.dqp_dct = htobe32(((uint32_t)ext << 31 | dct_num));
+  dc_seg->av.stat_rate_sl = mlx5_av->stat_rate_sl;
+  dc_seg->av.fl_mlid = mlx5_av->fl_mlid;
+  dc_seg->av.rlid = mlx5_av->rlid;
+}
+
+/*
+ * create a DC wqe and post it to the SQ
+ * we rely on mlx5dv functions to ceate the ctrl and data
+ * segments but we use our own function to write teh DC and rdma segments
+ */
+void DynamicConnection::post_dv_dc_wqe(int remote_conn) {
+  mlx5_wqe_ctrl_seg* ctrl;
+  mlx5_wqe_datagram_seg* dc_seg;
+  mlx5_wqe_raddr_seg* rdma;
+  mlx5_wqe_data_seg* data;
+
+  for (int i = 0; i < remote_conn; i++) {
+    uint64_t* ptr = get_address_sq(i);
+
+    const uint32_t nb_post = 4 * sq_size;
+    for (uint16_t index = 0; index < nb_post; index++) {
+      uint8_t op_mod = 0;
+      uint8_t op_code = 8;
+      uint32_t qp_num = qps[i]->qp_num;
+      uint8_t fm_ce_se = 0;
+      uint8_t ds = 4;
+      ctrl = reinterpret_cast<mlx5_wqe_ctrl_seg*>(ptr);
+      mlx5dv_set_ctrl_seg(ctrl, index, op_code, op_mod, qp_num, fm_ce_se, ds, 0,
+                          0);
+      ptr = ptr + 2;
+
+      uint32_t dct_num = dcts_num[i];
+      uint8_t ext = 1;
+      dc_seg = reinterpret_cast<mlx5_wqe_datagram_seg*>(ptr);
+      set_dgram_seg(dc_seg, (uint64_t)DC_IB_KEY, dct_num, ext, &mlx5_av);
+      ptr = ptr + 2;
+
+      uint64_t address = 0;
+      uint32_t rkey = 0;
+      rdma = reinterpret_cast<mlx5_wqe_raddr_seg*>(ptr);
+      set_rdma_seg(rdma, address, rkey);
+      ptr = ptr + 2;
+
+      uint32_t lkey = backend->networkImpl.heap_mr->lkey;
+      data = reinterpret_cast<mlx5_wqe_data_seg*>(ptr);
+      mlx5dv_set_data_seg(data, 1, lkey, 0);
+      ptr = ptr + 2;
+    }
+  }
+}
+
+// TODO(bpotter): remove redundancies with the other derived class
+void DynamicConnection::post_wqes() {
+  int remote_conn;
+  get_remote_conn(&remote_conn);
+  remote_conn *= backend->num_blocks_;
+  post_dv_dc_wqe(remote_conn);
+}
+
+void DynamicConnection::initialize_wr_fields([[maybe_unused]] ibv_send_wr* wr,
+                                             [[maybe_unused]] ibv_ah* ah,
+                                             [[maybe_unused]] int dc_key) {}
+
+int DynamicConnection::get_sq_dv_offset([[maybe_unused]] int pe_idx,
+                                        [[maybe_unused]] int num_qps,
+                                        int wg_idx) {
+  return wg_idx;
+}
+
+}  // namespace rocshmem

src/gpu_ib/dynamic_connection.hpp

@@ -0,0 +1,122 @@
+/******************************************************************************
+ * 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_GPU_IB_DYNAMIC_CONNECTION_HPP_
+#define LIBRARY_SRC_GPU_IB_DYNAMIC_CONNECTION_HPP_
+
+#include "connection.hpp"
+
+namespace rocshmem {
+
+class DynamicConnection : public Connection {
+ public:
+  explicit DynamicConnection(GPUIBBackend* backend);
+
+  ~DynamicConnection() override;
+
+  void get_remote_conn(int* remote_conn) override;
+
+  void post_wqes() override;
+
+  void initialize_rkey_handle(uint32_t** heap_rkey_handle,
+                                ibv_mr* mr) override;
+
+  void free_rkey_handle(uint32_t* heap_rkey_handle) override;
+
+  uint32_t* get_vec_dct_num() const { return vec_dct_num; }
+
+  uint16_t* get_vec_lids() const { return vec_lids; }
+
+ private:
+  InitQPState initqp(uint8_t port) override;
+
+  RtrState rtr(dest_info_t* dest, uint8_t port) override;
+
+  RtsState rts(dest_info_t* dest) override;
+
+  QPInitAttr qpattr(ibv_qp_cap cap) override;
+
+  void connect_dci(ibv_qp* qp, uint8_t port);
+
+  void create_dct(int32_t* dct_num, ibv_cq* cq, ibv_srq* srq, uint8_t port);
+
+  ibv_qp_init_attr_ex dct_qp_init_attr(ibv_cq* cq, ibv_srq* srq,
+                                       uint8_t port) const;
+
+  mlx5dv_qp_init_attr dct_dv_init_attr();
+
+  void dc_get_av(ibv_ah* ah, mlx5_wqe_av* mlx5_av);
+
+  void set_dgram_seg(mlx5_wqe_datagram_seg* dc_seg, uint64_t dc_key,
+                     uint32_t dct_num, uint8_t ext, mlx5_wqe_av* av);
+
+  void set_data_seg(mlx5_wqe_data_seg* data_seg, uint32_t lkey);
+
+  void post_dv_dc_wqe(int remote_conn);
+
+  void create_qps_1() override;
+
+  void create_qps_2(int port, int my_rank,
+                      ibv_port_attr* ib_port_att) override;
+
+  void create_qps_3(int port, ibv_qp* qp, int offset,
+                      ibv_port_attr* ib_port_att) override;
+
+  ibv_qp* create_qp_0(ibv_context* context,
+                      ibv_qp_init_attr_ex* qp_attr) override;
+
+  void allocate_dynamic_members(int num_wg) override;
+
+  void free_dynamic_members() override;
+
+  void initialize_1(int port, int num_wg) override;
+
+  void initialize_wr_fields(ibv_send_wr* wr, ibv_ah* ah, int dc_key) override;
+
+  int get_sq_dv_offset(int pe_idx, int32_t num_qps, int wg_idx) override;
+
+  int num_dcis{1};
+
+  int num_dct{1};
+
+  static constexpr int DC_IB_KEY{0x1ee7a330};
+
+  uint32_t* dcts_num{nullptr};
+
+  uint16_t* lids{nullptr};
+
+  mlx5_wqe_av mlx5_av{};
+
+  ibv_ah* ah{nullptr};
+
+  ibv_srq* srq{nullptr};
+
+  ibv_cq* dct_cq{nullptr};
+
+  uint32_t* vec_dct_num{nullptr};
+
+  uint16_t* vec_lids{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_DYNAMIC_CONNECTION_HPP_

src/gpu_ib/endian.cpp

@@ -0,0 +1,79 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "endian.hpp"
+
+namespace rocshmem {
+
+template <typename T>
+__device__ void swap_endian_store(T *dst, const T val) {
+  typedef union U {
+    T val;
+    uint8_t bytes[sizeof(T)];
+  } union_type;
+  union_type src;
+  union_type dst_tmp;
+
+  src.val = val;
+  std::reverse_copy(src.bytes, src.bytes + sizeof(T), dst_tmp.bytes);
+  *dst = dst_tmp.val;
+}
+
+template <>
+__device__ void swap_endian_store(uint64_t *dst, const uint64_t val) {
+  uint64_t new_val = ((val << 8) & 0xFF00FF00FF00FF00ULL) |
+                     ((val >> 8) & 0x00FF00FF00FF00FFULL);
+
+  new_val = ((new_val << 16) & 0xFFFF0000FFFF0000ULL) |
+            ((new_val >> 16) & 0x0000FFFF0000FFFFULL);
+
+  *dst = (new_val << 32) | (new_val >> 32);
+}
+
+template <>
+__device__ void swap_endian_store(int64_t *dst, const int64_t val) {
+  swap_endian_store(reinterpret_cast<uint64_t *>(dst), (const uint64_t)val);
+}
+
+template <>
+__device__ void swap_endian_store(uint32_t *dst, const uint32_t val) {
+  uint32_t new_val = ((val << 8) & 0xFF00FF00) | ((val >> 8) & 0xFF00FF);
+
+  *dst = (new_val << 16) | (new_val >> 16);
+}
+
+template <>
+__device__ void swap_endian_store(int32_t *dst, const int32_t val) {
+  swap_endian_store(reinterpret_cast<uint32_t *>(dst), (const uint32_t)val);
+}
+
+template <>
+__device__ void swap_endian_store(uint16_t *dst, const uint16_t val) {
+  *dst = ((val << 8) & 0xFF00) | ((val >> 8) & 0x00FF);
+}
+
+template <>
+__device__ void swap_endian_store(int16_t *dst, const int16_t val) {
+  swap_endian_store(reinterpret_cast<uint16_t *>(dst), (const uint16_t)val);
+}
+
+}  // namespace rocshmem

src/gpu_ib/endian.hpp

@@ -0,0 +1,53 @@
+/******************************************************************************
+ * 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_GPU_IB_ENDIAN_HPP_
+#define LIBRARY_SRC_GPU_IB_ENDIAN_HPP_
+
+#include <hip/hip_runtime.h>
+
+namespace rocshmem {
+
+template <typename T>
+__device__ void swap_endian_store(T *dst, const T val);
+
+template <>
+__device__ void swap_endian_store(uint64_t *dst, const uint64_t val);
+
+template <>
+__device__ void swap_endian_store(int64_t *dst, const int64_t val);
+
+template <>
+__device__ void swap_endian_store(uint32_t *dst, const uint32_t val);
+
+template <>
+__device__ void swap_endian_store(int32_t *dst, const int32_t val);
+
+template <>
+__device__ void swap_endian_store(uint16_t *dst, const uint16_t val);
+
+template <>
+__device__ void swap_endian_store(int16_t *dst, const int16_t val);
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_ENDIAN_HPP_

src/gpu_ib/gpu_ib_team.cpp

@@ -0,0 +1,56 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "gpu_ib_team.hpp"
+
+#include "../backend_type.hpp"
+#include "backend_ib.hpp"
+
+namespace rocshmem {
+
+GPUIBTeam::GPUIBTeam(Backend *backend, TeamInfo *team_info_parent,
+                     TeamInfo *team_info_world, int num_pes, int my_pe,
+                     MPI_Comm mpi_comm, int pool_index)
+    : Team(backend, team_info_parent, team_info_world, num_pes, my_pe,
+           mpi_comm) {
+  type = BackendType::GPU_IB_BACKEND;
+  const GPUIBBackend *b = static_cast<const GPUIBBackend *>(backend);
+
+  pool_index_ = pool_index;
+
+  barrier_pSync =
+      &(b->barrier_pSync_pool[pool_index * ROC_SHMEM_BARRIER_SYNC_SIZE]);
+  reduce_pSync =
+      &(b->reduce_pSync_pool[pool_index * ROC_SHMEM_REDUCE_SYNC_SIZE]);
+  bcast_pSync = &(b->bcast_pSync_pool[pool_index * ROC_SHMEM_BCAST_SYNC_SIZE]);
+  alltoall_pSync =
+      &(b->alltoall_pSync_pool[pool_index * ROC_SHMEM_ALLTOALL_SYNC_SIZE]);
+
+  pWrk = reinterpret_cast<char *>(b->pWrk_pool) +
+         ROC_SHMEM_REDUCE_MIN_WRKDATA_SIZE * sizeof(double) * pool_index;
+  pAta = reinterpret_cast<char *>(b->pAta_pool) +
+         ROC_SHMEM_ATA_MAX_WRKDATA_SIZE * sizeof(double) * pool_index;
+}
+
+GPUIBTeam::~GPUIBTeam() {}
+
+}  // namespace rocshmem

src/gpu_ib/gpu_ib_team.hpp

@@ -0,0 +1,50 @@
+/******************************************************************************
+ * 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_GPU_IB_GPU_IB_TEAM_HPP_
+#define LIBRARY_SRC_GPU_IB_GPU_IB_TEAM_HPP_
+
+#include "../team.hpp"
+
+namespace rocshmem {
+
+class GPUIBTeam : public Team {
+ public:
+  GPUIBTeam(Backend* handle, TeamInfo* team_info_wrt_parent,
+            TeamInfo* team_info_wrt_world, int num_pes, int my_pe,
+            MPI_Comm team_comm, int pool_index);
+
+  virtual ~GPUIBTeam();
+
+  long* barrier_pSync{nullptr};
+  long* reduce_pSync{nullptr};
+  long* bcast_pSync{nullptr};
+  long* alltoall_pSync{nullptr};
+  void* pWrk{nullptr};
+  void* pAta{nullptr};
+
+  int pool_index_{-1};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_GPU_IB_TEAM_HPP_

src/gpu_ib/infiniband_structs.hpp

@@ -0,0 +1,49 @@
+/******************************************************************************
+ * 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_GPU_IB_INFINIBAND_STRUCTS_HPP_
+#define LIBRARY_SRC_GPU_IB_INFINIBAND_STRUCTS_HPP_
+
+#include <infiniband/mlx5dv.h>
+
+namespace rocshmem {
+
+typedef struct ib_mlx5_base_av {
+  uint64_t dc_key;
+  uint32_t dqp_dct;
+  uint8_t stat_rate_sl;
+  uint8_t fl_mlid;
+  uint16_t rlid;
+} ib_mlx5_base_av_t;
+
+union mlx5_segment {
+  mlx5_wqe_ctrl_seg ctrl_seg;
+  mlx5_wqe_raddr_seg raddr_seg;
+  mlx5_wqe_atomic_seg atomic_seg;
+  mlx5_wqe_data_seg data_seg;
+  mlx5_wqe_inl_data_seg inl_data_seg;
+  ib_mlx5_base_av_t base_av;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_INFINIBAND_STRUCTS_HPP_

src/gpu_ib/memory_builder_policy.hpp

@@ -0,0 +1,78 @@
+/******************************************************************************
+ * 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_GPU_IB_MEMORY_BUILDER_POLICY_HPP_
+#define LIBRARY_SRC_GPU_IB_MEMORY_BUILDER_POLICY_HPP_
+
+#include <hip/hip_runtime.h>
+
+#include <utility>
+
+namespace rocshmem {
+
+class GPUIBContext;
+
+class MemoryBuilderPolicyWrapper {
+ public:
+  __device__ MemoryBuilderPolicyWrapper() = default;
+
+  __device__ ~MemoryBuilderPolicyWrapper() {
+    if (wrapped_policy_) {
+      delete wrapped_policy_;
+    }
+  }
+
+  template <typename T>
+  __device__ MemoryBuilderPolicyWrapper(T&& policy)
+      : wrapped_policy_(new Wrapper<T>(std::forward<T>(policy))) {}
+
+  __device__ void operator()(GPUIBContext* context) {
+    return (*wrapped_policy_)(context);
+  }
+
+ private:
+  class PolicyBase {
+   public:
+    __device__ virtual void operator()(GPUIBContext* context) = 0;
+
+    __device__ virtual ~PolicyBase() {}
+  };
+
+  template <typename T>
+  class Wrapper : public PolicyBase {
+   public:
+    __device__ Wrapper(const T& t) : wrapped_policy_(t) {}
+
+    __device__ void operator()(GPUIBContext* context) override {
+      return wrapped_policy_(context);
+    }
+
+   private:
+    T wrapped_policy_;
+  };
+
+  PolicyBase* wrapped_policy_;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_MEMORY_BUILDER_POLICY_HPP_

src/gpu_ib/network_policy.cpp

@@ -0,0 +1,500 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "network_policy.hpp"
+
+#include <mpi.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "../atomic_return.hpp"
+#include "../context_incl.hpp"
+#include "backend_ib.hpp"
+#include "connection.hpp"
+#include "dynamic_connection.hpp"
+#include "queue_pair.hpp"
+#include "reliable_connection.hpp"
+
+namespace rocshmem {
+
+void NetworkOnImpl::dump_backend_stats(ROCStats *globalStats) {
+  /*
+   * TODO(bpotter): Refactor this into the Stats class to remove the ifdef
+   */
+#ifdef PROFILE
+  int statblocks = connection->total_number_connections();
+
+  uint64_t cycles_ring_sq_db = 0;
+  uint64_t cycles_update_wqe = 0;
+  uint64_t cycles_poll_cq = 0;
+  uint64_t cycles_next_cq = 0;
+  uint64_t cycles_init = gpu_qps[statblocks - 1].profiler.getStat(INIT);
+  uint64_t cycles_finalize = gpu_qps[statblocks - 1].profiler.getStat(FINALIZE);
+
+  uint64_t total_quiet_count = 0;
+  uint64_t total_db_count = 0;
+  uint64_t total_wqe_count = 0;
+
+  for (int i = 0; i < statblocks; i++) {
+    cycles_ring_sq_db += gpu_qps[i].profiler.getStat(RING_SQ_DB);
+    cycles_update_wqe += gpu_qps[i].profiler.getStat(UPDATE_WQE);
+    cycles_poll_cq += gpu_qps[i].profiler.getStat(POLL_CQ);
+    cycles_next_cq += gpu_qps[i].profiler.getStat(NEXT_CQ);
+    total_quiet_count += gpu_qps[i].profiler.getStat(QUIET_COUNT);
+    total_db_count += gpu_qps[i].profiler.getStat(DB_COUNT);
+    total_wqe_count += gpu_qps[i].profiler.getStat(WQE_COUNT);
+  }
+
+  double us_ring_sq_db = cycles_ring_sq_db / gpu_clock_freq_mhz;
+  double us_update_wqe = cycles_update_wqe / gpu_clock_freq_mhz;
+  double us_poll_cq = cycles_poll_cq / gpu_clock_freq_mhz;
+  double us_next_cq = cycles_next_cq / gpu_clock_freq_mhz;
+  double us_init = cycles_init / gpu_clock_freq_mhz;
+  double us_finalize = cycles_finalize / gpu_clock_freq_mhz;
+
+  const int FIELD_WIDTH = 20;
+  const int FLOAT_PRECISION = 2;
+
+  printf("Counts: Internal Quiets %lu DB Rings %lu WQE Posts %lu\n",
+         total_quiet_count, total_db_count, total_wqe_count);
+
+  printf("\n%*s%*s%*s%*s%*s%*s\n", FIELD_WIDTH + 1, "Init (us)",
+         FIELD_WIDTH + 1, "Finalize (us)", FIELD_WIDTH + 1, "Ring SQ DB (us)",
+         FIELD_WIDTH + 1, "Update WQE (us)", FIELD_WIDTH + 1, "Poll CQ (us)",
+         FIELD_WIDTH + 1, "Next CQ (us)");
+
+  uint64_t totalFinalize = globalStats->getStat(NUM_FINALIZE);
+  printf("%*.*f %*.*f %*.*f %*.*f %*.*f %*.*f\n", FIELD_WIDTH, FLOAT_PRECISION,
+         us_init / totalFinalize, FIELD_WIDTH, FLOAT_PRECISION,
+         us_finalize / totalFinalize, FIELD_WIDTH, FLOAT_PRECISION,
+         us_ring_sq_db / total_db_count, FIELD_WIDTH, FLOAT_PRECISION,
+         us_update_wqe / total_wqe_count, FIELD_WIDTH, FLOAT_PRECISION,
+         us_poll_cq / total_quiet_count, FIELD_WIDTH, FLOAT_PRECISION,
+         us_next_cq / total_quiet_count);
+#endif
+}
+
+void NetworkOnImpl::reset_backend_stats() {
+  int statblocks = connection->total_number_connections();
+
+  for (size_t i = 0; i < statblocks; i++) {
+    gpu_qps[i].profiler.resetStats();
+  }
+}
+
+void NetworkOnImpl::exchange_hdp_info(HdpPolicy *hdp_policy,
+                                      MPI_Comm thread_comm) {
+  /*
+   * Using Connection class, register the host-side hdp flush address
+   * with the InfiniBand network.
+   */
+  connection->reg_mr(hdp_policy->get_hdp_flush_ptr(), 32, &hdp_mr, false);
+
+  /*
+   * Allocate device-side memory for the remote HDP keys.
+   */
+  CHECK_HIP(hipMalloc(reinterpret_cast<void **>(&hdp_rkey),
+                      num_pes * sizeof(uint32_t)));
+
+  /*
+   * Allocate device-side memory for the remote HDP addresses.
+   */
+  CHECK_HIP(hipMalloc(reinterpret_cast<void **>(&hdp_address),
+                      num_pes * sizeof(uintptr_t)));
+
+  /*
+   * Allocate host-side memory to exchange hdp keys using MPI_Allgather.
+   */
+  uint32_t *host_hdp_cpy =
+      reinterpret_cast<uint32_t *>(malloc(num_pes * sizeof(uint32_t)));
+  if (host_hdp_cpy == nullptr) {
+    abort();
+  }
+
+  /*
+   * Allocate host-side memory to exchange hdp addresses using
+   * MPI_Allgather.
+   */
+  uint32_t **host_hdp_address_cpy =
+      reinterpret_cast<uint32_t **>(malloc(num_pes * sizeof(uint32_t *)));
+  if (host_hdp_address_cpy == nullptr) {
+    free(host_hdp_cpy);
+    abort();
+  }
+
+  /*
+   * This processing element writes its personal HDP key and HDP address
+   * into the host-side arrays which were just allocated.
+   */
+  int my_rank = my_pe;
+  host_hdp_cpy[my_rank] = htobe32(hdp_mr->rkey);
+  host_hdp_address_cpy[my_rank] = hdp_policy->get_hdp_flush_ptr();
+
+  /*
+   * Do all-to-all exchange of our HDP key with other processing elements.
+   */
+  MPI_Allgather(MPI_IN_PLACE, sizeof(uint32_t), MPI_CHAR, host_hdp_cpy,
+                sizeof(uint32_t), MPI_CHAR, thread_comm);
+
+  /*
+   * Do all-to-all exchange of our HDP address with other processing
+   * elements.
+   */
+  MPI_Allgather(MPI_IN_PLACE, sizeof(uintptr_t), MPI_CHAR, host_hdp_address_cpy,
+                sizeof(uint32_t *), MPI_CHAR, thread_comm);
+
+  /*
+   * Copy the recently exchanged HDP keys to device memory.
+   */
+  hipStream_t stream;
+  CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+  CHECK_HIP(hipMemcpyAsync(hdp_rkey, host_hdp_cpy, num_pes * sizeof(uint32_t),
+                           hipMemcpyHostToDevice, stream));
+
+  /*
+   * Copy the recently exchanged HDP addresses to device memory.
+   */
+  CHECK_HIP(hipMemcpyAsync(hdp_address, host_hdp_address_cpy,
+                           num_pes * sizeof(uint32_t *), hipMemcpyHostToDevice,
+                           stream));
+  CHECK_HIP(hipStreamSynchronize(stream));
+  CHECK_HIP(hipStreamDestroy(stream));
+
+  /*
+   * Free the host-side resources used to exchange HDP resources
+   * between processing elements.
+   */
+  free(host_hdp_cpy);
+  free(host_hdp_address_cpy);
+}
+
+void NetworkOnImpl::setup_atomic_region() {
+  /*
+   * Allocate fine-grained device-side memory for the atomic return
+   * region.
+   */
+  allocate_atomic_region(&atomic_ret, num_blocks);
+
+  /*
+   * Register the atomic return region on the InfiniBand network.
+   */
+  connection->reg_mr(atomic_ret->atomic_base_ptr,
+                     sizeof(uint64_t) * max_nb_atomic * num_blocks, &mr, false);
+
+  /*
+   * Set member variable from class.
+   */
+  atomic_ret->atomic_lkey = htobe32(mr->lkey);
+}
+
+void NetworkOnImpl::heap_memory_rkey(char *local_heap_base, size_t heap_size,
+                                     MPI_Comm thread_comm, bool is_managed) {
+  /*
+   * Allocate host-side memory to hold remote keys for all processing
+   * elements.
+   */
+  const size_t rkeys_size = sizeof(uint32_t) * num_pes;
+  uint32_t *host_rkey_cpy = reinterpret_cast<uint32_t *>(malloc(rkeys_size));
+  if (host_rkey_cpy == nullptr) {
+    abort();
+  }
+
+  /*
+   * Using the Connection class, register the symmetric heap with the
+   * InfiniBand network.
+   */
+  void *base_heap = local_heap_base;
+  connection->reg_mr(base_heap, heap_size, &heap_mr, is_managed);
+
+  /*
+   * Using the memory region from the prior heap memory registration,
+   * allocate and initialize some device-side memory to hold the remote
+   * keys for the symmetric heap base.
+   *
+   * Only the device-side memory entry for this processing element will be
+   * updated with the key for the heap memory region.
+   */
+  connection->initialize_rkey_handle(&heap_rkey, heap_mr);
+
+  /*
+   * Copy the device-side heap base remote key array to the host-side
+   * heap base remote key array.
+   */
+  hipStream_t stream;
+  CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+  CHECK_HIP(hipMemcpyAsync(host_rkey_cpy, heap_rkey, rkeys_size,
+                           hipMemcpyDeviceToHost, stream));
+  CHECK_HIP(hipStreamSynchronize(stream));
+
+  /*
+   * Do all-to-all exchange of symmetric heap base remote key between the
+   * processing elements.
+   */
+  MPI_Allgather(MPI_IN_PLACE, sizeof(uint32_t), MPI_CHAR, host_rkey_cpy,
+                sizeof(uint32_t), MPI_CHAR, thread_comm);
+
+  /*
+   * Copy the recently updated host-side heap base remote key array back
+   * to the device-side memory.
+   */
+  CHECK_HIP(hipMemcpyAsync(heap_rkey, host_rkey_cpy, rkeys_size,
+                           hipMemcpyHostToDevice, stream));
+  CHECK_HIP(hipStreamSynchronize(stream));
+  CHECK_HIP(hipStreamDestroy(stream));
+
+  /*
+   * Free the host-side resources used to do the processing element
+   * exchange of keys and addresses for the symmetric heap base.
+   */
+  free(host_rkey_cpy);
+
+  /*
+   * Initialize this member variable to hold the InfiniBand memory
+   * region's local key.
+   */
+  lkey = heap_mr->lkey;
+}
+
+void NetworkOnImpl::setup_gpu_qps(GPUIBBackend *B) {
+  /*
+   * Determine how many connections are needed.
+   * The number of connections depends on the connection type and the
+   * number of workgroups.
+   */
+  int connections;
+  connection->get_remote_conn(&connections);
+  connections *= num_blocks;
+
+  /*
+   * Allocate device-side memory for the queue pairs.
+   */
+  CHECK_HIP(hipMalloc(&gpu_qps, sizeof(QueuePair) * connections));
+
+  /*
+   * For every connection, initialize the QueuePair.
+   */
+  for (int i = 0; i < connections; i++) {
+    new (&gpu_qps[i]) QueuePair(B);
+    connection->init_gpu_qp_from_connection(&gpu_qps[i], i);
+  }
+}
+
+void NetworkOnImpl::roc_shmem_g_init(SymmetricHeap *heap_handle,
+                                     MPI_Comm thread_comm) {
+  init_g_ret(heap_handle, thread_comm, num_blocks, &g_ret);
+}
+
+__host__ void NetworkOnImpl::networkHostSetup(GPUIBBackend *B) {
+  num_pes = B->num_pes;
+  my_pe = B->my_pe;
+  num_blocks = B->num_blocks_;
+
+#ifdef USE_DC
+  connection = new DynamicConnection(B);
+#else
+  connection = new ReliableConnection(B);
+#endif
+
+  connection->initialize(B->num_blocks_);
+  exchange_hdp_info(B->hdp_policy, B->thread_comm);
+
+  const auto &heap_bases{B->heap.get_heap_bases()};
+  heap_memory_rkey(heap_bases[my_pe], B->heap.get_size(), B->thread_comm,
+                   B->heap.is_managed());
+  // The earliest we can allow the main thread to launch a kernel to
+  // avoid potential deadlock
+  network_init_done = true;
+
+  setup_atomic_region();
+
+  connection->initialize_gpu_policy(&connection_policy, heap_rkey);
+
+  roc_shmem_g_init(&B->heap, B->thread_comm);
+
+  connection->post_wqes();
+
+  setup_gpu_qps(B);
+}
+
+__host__ void NetworkOnImpl::networkHostFinalize() {
+  CHECK_HIP(hipFree(hdp_rkey));
+  hdp_rkey = nullptr;
+
+  CHECK_HIP(hipFree(hdp_address));
+  hdp_address = nullptr;
+
+  CHECK_HIP(hipFree(atomic_ret));
+  atomic_ret = nullptr;
+
+  CHECK_HIP(hipFree(gpu_qps));
+  gpu_qps = nullptr;
+
+  CHECK_HIP(hipFree(connection_policy));
+  connection_policy = nullptr;
+
+  connection->free_rkey_handle(heap_rkey);
+
+  connection->finalize();
+  delete connection;
+  connection = nullptr;
+}
+
+__host__ void NetworkOnImpl::networkHostInit(GPUIBContext *ctx, int buffer_id) {
+  int remote_conn = getNumQueuePairs();
+
+  CHECK_HIP(hipMalloc(&ctx->device_qp_proxy, remote_conn * sizeof(QueuePair)));
+
+  for (int i = 0; i < getNumQueuePairs(); i++) {
+    /*
+     * RC gpu_qp is actually [NUM_PE][NUM_BLOCK] qps but is flattened.
+     * Each num_pe entry contains num_block QPs connected to that PE.
+     * For RC, we need to iterate gpu_qp[i][buffer_id] to collect a
+     * single QP for each connected PE in order to build context.
+     * For DC, NUM_PE = 1 so can just use buffer_id directly.
+     */
+    int offset = num_blocks * i + buffer_id;
+    new (ctx->getQueuePair(i)) QueuePair(gpu_qps[offset]);
+
+    auto *qp = ctx->getQueuePair(i);
+    qp->global_qp = &gpu_qps[offset];
+    qp->num_cqs = getNumQueuePairs();
+    qp->atomic_ret.atomic_base_ptr =
+        &atomic_ret->atomic_base_ptr[max_nb_atomic * buffer_id];
+    qp->base_heap = ctx->base_heap;
+  }
+  ctx->g_ret = g_ret;
+}
+
+__device__ void NetworkOnImpl::networkGpuInit(GPUIBContext *ctx,
+                                              int buffer_id) {
+  for (int i = 0; i < getNumQueuePairs(); i++) {
+    int offset = num_blocks * i + buffer_id;
+
+    auto *qp = ctx->getQueuePair(i);
+    new (qp) QueuePair(gpu_qps[offset]);
+
+    qp->global_qp = &gpu_qps[offset];
+    qp->num_cqs = getNumQueuePairs();
+    qp->atomic_ret.atomic_base_ptr =
+        &atomic_ret->atomic_base_ptr[max_nb_atomic * buffer_id];
+    qp->base_heap = ctx->base_heap;
+  }
+  ctx->g_ret = g_ret;
+}
+
+__device__ __host__ QueuePair *NetworkOnImpl::getQueuePair(QueuePair *qp_handle,
+                                                           int pe) {
+#ifdef USE_DC
+  return qp_handle;
+#else
+  return &qp_handle[pe];
+#endif
+}
+
+__device__ __host__ int NetworkOnImpl::getNumQueuePairs() {
+#ifdef USE_DC
+  return 1;
+#else
+  return num_pes;
+#endif
+}
+
+void NetworkOffImpl::networkHostSetup(GPUIBBackend *B) {
+  num_pes = B->num_pes;
+  my_pe = B->my_pe;
+  num_blocks = B->num_blocks_;
+
+  exchange_hdp_info(B->hdp_policy, B->thread_comm);
+}
+void NetworkOffImpl::exchange_hdp_info(HdpPolicy *hdp_policy,
+                                       MPI_Comm thread_comm) {
+#ifdef USE_SINGLE_NODE
+  // We are using the symmetric heap for the HDP flush ptr
+  hdp_address = reinterpret_cast<uintptr_t *>(hdp_policy->get_hdp_flush_ptr());
+#else
+  /*
+   * Allocate device-side memory for the remote HDP addresses.
+   */
+  CHECK_HIP(hipMalloc(reinterpret_cast<void **>(&hdp_address),
+                      num_pes * sizeof(uintptr_t)));
+
+  /*
+   * Allocate host-side memory to exchange hdp keys using MPI_Allgather.
+   */
+  uint32_t *host_hdp_cpy =
+      reinterpret_cast<uint32_t *>(malloc(num_pes * sizeof(uint32_t)));
+  if (host_hdp_cpy == nullptr) {
+    abort();
+  }
+
+  /*
+   * Allocate host-side memory to exchange hdp addresses using
+   * MPI_Allgather.
+   */
+  uint32_t **host_hdp_address_cpy =
+      reinterpret_cast<uint32_t **>(malloc(num_pes * sizeof(uint32_t *)));
+  if (host_hdp_address_cpy == nullptr) {
+    free(host_hdp_cpy);
+    abort();
+  }
+
+  /*
+   * This processing element writes its personal HDP address
+   * into the host-side array which were just allocated.
+   */
+  int my_rank = my_pe;
+  host_hdp_address_cpy[my_rank] = hdp_policy->get_hdp_flush_ptr();
+
+  /*
+   * Do all-to-all exchange of our HDP address with other processing
+   * elements.
+   */
+  MPI_Allgather(MPI_IN_PLACE, sizeof(uintptr_t), MPI_CHAR, host_hdp_address_cpy,
+                sizeof(uint32_t *), MPI_CHAR, thread_comm);
+
+  /*
+   * Copy the recently exchanged HDP addresses to device memory.
+   */
+  hipStream_t stream;
+  CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+  CHECK_HIP(hipMemcpyAsync(hdp_address, host_hdp_address_cpy,
+                           num_pes * sizeof(uint32_t *), hipMemcpyHostToDevice,
+                           stream));
+  CHECK_HIP(hipStreamSynchronize(stream));
+  CHECK_HIP(hipStreamDestroy(stream));
+
+  /*
+   * Free the host-side resources used to exchange HDP resources
+   * between processing elements.
+   */
+  free(host_hdp_cpy);
+  free(host_hdp_address_cpy);
+#endif
+}
+
+void NetworkOffImpl::networkHostFinalize() {
+#ifndef USE_SINGLE_NODE
+  CHECK_HIP(hipFree(hdp_address));
+#endif
+  hdp_address = nullptr;
+}
+
+}  // namespace rocshmem

src/gpu_ib/network_policy.hpp

@@ -0,0 +1,357 @@
+/******************************************************************************
+ * 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_GPU_IB_NETWORK_POLICY_HPP_
+#define LIBRARY_SRC_GPU_IB_NETWORK_POLICY_HPP_
+
+#include <hip/hip_runtime.h>
+#include <mpi.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "roc_shmem/roc_shmem.hpp"
+#include "connection_policy.hpp"
+#include "queue_pair.hpp"
+#include "../hdp_policy.hpp"
+#include "../memory/symmetric_heap.hpp"
+#include "../stats.hpp"
+#include "../util.hpp"
+
+struct ibv_mr;
+struct hdp_reg_t;
+
+namespace rocshmem {
+
+struct atomic_ret_t;
+class GPUIBBackend;
+class GPUIBContext;
+class GPUIBHostContext;
+class Connection;
+
+class NetworkOnImpl {
+ public:
+  void dump_backend_stats(ROCStats *globalStats);
+
+  void reset_backend_stats();
+
+  /**
+   * @brief setup the network resources and initialization for the
+   * GPUIBBackend
+   */
+  __host__ void networkHostSetup(GPUIBBackend *B);
+
+  /**
+   * @brief deallocate and close the network resources
+   */
+  __host__ void networkHostFinalize();
+
+  /**
+   * @brief initialize the network resources for each context
+   */
+  __host__ void networkHostInit(GPUIBContext *ctx, int buffer_id);
+
+  /**
+   * @brief initialize the network resources for each context on GPU side
+   */
+  __device__ void networkGpuInit(GPUIBContext *ctx, int buffer_id);
+
+  /**
+   * @brief returns the QP for the targeted pe
+   */
+  __device__ __host__ QueuePair *getQueuePair(QueuePair *qp, int pe);
+
+  /**
+   * @brief returns the numbers of QPs used per the calling PE
+   */
+  __device__ __host__ int getNumQueuePairs();
+
+  /**
+   * @brief returns the number of PEs accessible via network
+   */
+  __device__ __host__ int getNumDest() { return num_pes; }
+
+  static uint32_t externSharedBytes(int num_pes) {
+    int remote_conn{1};
+#ifndef USE_DC
+    remote_conn = num_pes;
+#endif
+    return remote_conn * sizeof(QueuePair);
+  }
+
+ protected:
+  /**
+   * @brief flag to indicated that the helper thread reach this milestone
+   */
+  volatile bool network_init_done{false};
+
+  void heap_memory_rkey(char *local_heap_base, size_t heap_size,
+                        MPI_Comm thread_comm, bool is_managed);
+
+  /**
+   * @brief Exchange HDP information between all processing elements.
+   *
+   * Each device has a Host Data Path (HDP) associated with it must be
+   * manually controlled when using fine-grained memory accesses. (The
+   * symmetric heap is allocated with fine-grained memory to support both
+   * host memory accesses and device memory accesses.) The HDP can be
+   * cleared by accessing an address on the device. These addresses must be
+   * shared across the network (to support updates on remote accesses).
+   *
+   * These HDPs are visible to the network by registering them as
+   * InfiniBand memory regions. Every memory region has a remote key
+   * which needs to be shared across the network (to access the memory
+   * region).
+   *
+   * This method is responsible to allocating and initializing the
+   * library's HDP device-side memory and running the all-to-all exchange
+   * to share both the keys and addresses.
+   *
+   * @todo Implement HDP policy class methods to hide most of this
+   * method. The guts should be encapsulated in the policy class and
+   * not exposed here in the backend. Within the policy class methods,
+   * create helper function to improve code reuse regarding the many
+   * data transfers.
+   */
+  void exchange_hdp_info(HdpPolicy *hdp_policy, MPI_Comm thread_comm);
+
+  /**
+   * @brief Allocate and initialize the atomic region.
+   *
+   * The atomic region is used by the atomic operations which have return
+   * values. The library user does not need to provide an address for the
+   * return value so we are forced to do it on their behalf.
+   *
+   * The atomic_ret member is initialized upon completion of this method.
+   */
+  void setup_atomic_region();
+
+  /**
+   * @brief Allocate and initialize device-side queue pair objects.
+   *
+   * Upon completion, the gpu_qps member will be initialized.
+   */
+  void setup_gpu_qps(GPUIBBackend *B);
+
+  /**
+   * @brief Allocate and initialize device-side memory that will be used for
+   * the return of g shmem ops (eg: shmem_int_g)
+   */
+  void roc_shmem_g_init(SymmetricHeap *heap_handle, MPI_Comm thread_comm);
+
+  /**
+   * @brief The backend delegates some InfiniBand connection setup to
+   * the Connection class.
+   */
+  Connection *connection{nullptr};
+
+ public:
+  /**
+   * @brief Number of PEs. Get directly from the GPUIBBackend.
+   */
+  int num_pes{0};
+
+  /**
+   * @brief This PE's rank.
+   */
+  int my_pe{-1};
+
+  /**
+   * @brief Number of WG that will be performing communication
+   */
+  int num_blocks{0};
+
+  /**
+   * @brief Holds InfiniBand remote keys for HDP memory regions.
+   *
+   * The member holds a C-array allocation for remote keys (from
+   * InfiniBand memory registrations) for remote HDP registers. The C-array
+   * has one entry for each processing element (indexed by processing
+   * element ID).
+   *
+   * @todo Remove duplication between the backend class and the QueuePair
+   * class. QueuePair stores a copy of this member too. The backend
+   * class does not do much besides initialize this data structure and
+   * hold it until the QueuePair can consume it.
+   */
+  uint32_t *hdp_rkey{nullptr};
+
+  /**
+   * @brief Holds HDP register addresses for each processing element.
+   *
+   * The Host Data Path (HDP) addresses are used to clear a buffer
+   * which interferes with memory visibility of accesses to fine-grained
+   * allocations.
+   *
+   * The member holds a C-array allocation for the register addresses.
+   * The C-array has one entry for each processing element (indexed by
+   * processing element ID).
+   *
+   * @todo Remove duplication between the backend class and the QueuePair
+   * class. QueuePair stores a copy of this member too. The backend
+   * class does not do much besides initialize this data structure and
+   * hold it until the QueuePair can consume it.
+   */
+  uintptr_t *hdp_address{nullptr};
+
+  /**
+   * @brief Handle for the HDP memory region.
+   */
+  ibv_mr *hdp_mr{nullptr};
+
+  /**
+   * @brief Set of QueuePairs used by device to do networking.
+   *
+   * The member is used during Context creation.
+   *
+   * @todo What we really need here is a collection of Contexts that can
+   * either be copied into LDS or used directly by the GPU depending on
+   * what type of context it is (shareable, serialized, or private).
+   * No need to pool up QueuePairs, they can just be managed by their
+   * owning Context. Should then consider pushing into base class since
+   * it's not gpu-ib specific.
+   */
+  QueuePair *gpu_qps{nullptr};
+
+  /**
+   * @brief C-array of symmetric heap base pointers.
+   *
+   * A C-array of char* pointers corresponding to the heap base pointers
+   * virtual address for each processing element that we can communicate
+   * with.
+   */
+  uint32_t *heap_rkey{nullptr};
+
+  /**
+   * @brief Handle for the symmetric heap memory region.
+   */
+  ibv_mr *heap_mr{nullptr};
+
+  /**
+   * @brief Local key for the symmetric heap memory region.
+   */
+  uint32_t lkey{0};
+
+  /**
+   * @brief Control struct for atomic memory region.
+   *
+   * The atomic region is used by the atomic operations which have return
+   * values. The library user does not need to provide an address for the
+   * return value so we are forced to do it on their behalf.
+   */
+  atomic_ret_t *atomic_ret{nullptr};
+
+  /**
+   * @brief Handle for the atomic memory region.
+   *
+   * @todo Provide more descriptive variable name.
+   */
+  ibv_mr *mr{nullptr};
+
+  /**
+   * @brief Buffer used to store the results of a *_g operation.
+   *
+   * These operations do not provide a destination buffer so the runtime
+   * must manage one.
+   */
+  char *g_ret{nullptr};
+
+  /**
+   * @brief Compile-time configuration policy for InfiniBand connections.
+   *
+   * The configuration option "USE_DC" can be enabled to create
+   * Dynamic connection types. By default, Reliable connections are
+   * created.
+   */
+  ConnectionImpl *connection_policy{nullptr};
+};
+
+// clang-format off
+NOWARN(-Wunused-parameter,
+class NetworkOffImpl {
+ public:
+  void dump_backend_stats(ROCStats *globalStats) { }
+
+  void reset_backend_stats() { }
+
+  __host__ void networkHostSetup(GPUIBBackend *B);
+
+  __host__ void exchange_hdp_info(HdpPolicy *hdp_policy, MPI_Comm thread_comm);
+
+  __host__ void networkHostFinalize();
+
+  __host__ void networkHostInit(GPUIBContext *ctx, int buffer_id) {}
+
+  __device__ void networkGpuInit(GPUIBContext *ctx, int buffer_id) {}
+
+  __device__ __host__ QueuePair *getQueuePair(QueuePair *qp, int pe) {
+    return nullptr;
+  }
+
+  __device__ __host__ int getNumQueuePairs() { return 0; }
+
+  __device__ __host__ int getNumDest() { return 0; }
+
+  static uint32_t externSharedBytes(int num_pes) { return 0; }
+
+ public:
+  int num_pes{0};
+
+  int my_pe{-1};
+
+  int num_blocks{0};
+
+  uint32_t *hdp_rkey{nullptr};
+
+  uintptr_t *hdp_address{nullptr};
+
+  ibv_mr *hdp_mr{nullptr};
+
+  QueuePair *gpu_qps{nullptr};
+
+  uint32_t *heap_rkey{nullptr};
+
+  ibv_mr *heap_mr{nullptr};
+
+  uint32_t lkey{0};
+
+  atomic_ret_t *atomic_ret{nullptr};
+
+  ibv_mr *mr{nullptr};
+
+  char *g_ret{nullptr};
+
+  ConnectionImpl *connection_policy{nullptr};
+};
+)
+// clang-format on
+
+/*
+ * Select which one of our IPC policies to use at compile time.
+ */
+#ifdef USE_SINGLE_NODE
+typedef NetworkOffImpl NetworkImpl;
+#else
+typedef NetworkOnImpl NetworkImpl;
+#endif
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_NETWORK_POLICY_HPP_

src/gpu_ib/qe_dumper.cpp

@@ -0,0 +1,79 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "qe_dumper.hpp"
+
+namespace rocshmem {
+
+QeDumper::QeDumper(int dest_pe, int src_wg, int index)
+    : dest_pe_(dest_pe), src_wg_(src_wg), index_(index) {
+  void* temp = malloc(sizeof(GPUIBBackend*));
+  gpu_backend_ = static_cast<GPUIBBackend*>(temp);
+
+  GPUIBBackend* device_backend_proxy_address;
+  CHECK_HIP(hipGetSymbolAddress(
+      reinterpret_cast<void**>(&device_backend_proxy_address),
+      HIP_SYMBOL(device_backend_proxy)));
+
+  CHECK_HIP(hipMemcpy(&gpu_backend_, device_backend_proxy_address,
+                      sizeof(GPUIBBackend*), hipMemcpyDeviceToHost));
+
+  int qp_offset = gpu_backend_->num_blocks_ * dest_pe_ + src_wg_;
+
+  qp_ = &(gpu_backend_->networkImpl.gpu_qps[qp_offset]);
+}
+
+QeDumper::~QeDumper() {
+  /*if (gpu_backend_) {
+      free(gpu_backend_);
+  }*/
+}
+
+void QeDumper::dump_cq() {
+  type_ = "CQ";
+
+  auto* raw_cqe = &(qp_->current_cq_q_H[index_]);
+  raw_u64_ = reinterpret_cast<uint64_t*>(raw_cqe);
+
+  dump_uint64_(8);
+}
+
+void QeDumper::dump_sq() {
+  type_ = "SQ";
+
+  auto* raw_sqe = &(qp_->current_sq_H[index_ * 8]);
+  raw_u64_ = reinterpret_cast<uint64_t*>(raw_sqe);
+
+  dump_uint64_(8);
+}
+
+void QeDumper::dump_uint64_(size_t num_elems) const {
+  printf("%s(%d, %d, %d) *** = ", type_.c_str(), dest_pe_, src_wg_, index_);
+
+  for (size_t i = 0; i < num_elems; i++) {
+    printf(" %lx ", raw_u64_[i]);
+  }
+
+  printf("done %s\n", type_.c_str());
+}
+
+}  // namespace rocshmem

src/gpu_ib/qe_dumper.hpp

@@ -0,0 +1,66 @@
+/******************************************************************************
+ * 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_GPU_IB_QE_DUMPER_HPP_
+#define LIBRARY_SRC_GPU_IB_QE_DUMPER_HPP_
+
+#include <hip/hip_runtime.h>
+#include <infiniband/mlx5dv.h>
+
+#include <string>
+
+#include "backend_ib.hpp"
+#include "queue_pair.hpp"
+
+namespace rocshmem {
+
+class QeDumper {
+ public:
+  QeDumper(int dest_pe, int src_wg, int index);
+
+  ~QeDumper();
+
+  void dump_cq();
+
+  void dump_sq();
+
+ private:
+  void dump_uint64_(size_t num_elems) const;
+
+  int dest_pe_{-1};
+
+  int src_wg_{-1};
+
+  int index_{-1};
+
+  GPUIBBackend* gpu_backend_{nullptr};
+
+  std::string type_{};
+
+  QueuePair* qp_{nullptr};
+
+  uint64_t* raw_u64_{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_QE_DUMPER_HPP_

src/gpu_ib/queue_pair.cpp

@@ -0,0 +1,437 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "queue_pair.hpp"
+
+#include <hip/hip_runtime.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "backend_ib.hpp"
+#include "endian.hpp"
+#include "segment_builder.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+QueuePair::QueuePair(GPUIBBackend *backend)
+    : hdp_policy(backend->hdp_policy),
+      connection_policy(*backend->networkImpl.connection_policy) {
+  hdp_rkey = backend->networkImpl.hdp_rkey;
+  hdp_address = backend->networkImpl.hdp_address;
+
+  atomic_ret.atomic_lkey = backend->networkImpl.atomic_ret->atomic_lkey;
+  atomic_ret.atomic_counter = 0;
+}
+
+__device__ QueuePair::~QueuePair() {
+  uint64_t start = profiler.startTimer();
+
+  global_qp->sq_counter = sq_counter;
+  global_qp->local_sq_cnt = local_sq_cnt;
+  global_qp->cq_consumer_counter = cq_consumer_counter;
+  global_qp->current_sq = current_sq;
+  global_qp->current_cq_q = current_cq_q;
+  global_qp->sq_overflow = sq_overflow;
+  global_qp->quiet_counter = quiet_counter;
+  profiler.endTimer(start, FINALIZE);
+
+  global_qp->profiler.accumulateStats(profiler);
+
+  __syncthreads();
+}
+
+__device__ uint8_t QueuePair::get_cq_error_syndrome(mlx5_cqe64 *cqe_entry) {
+  mlx5_err_cqe *cqe_err = reinterpret_cast<mlx5_err_cqe *>(cqe_entry);
+  return cqe_err->syndrome;
+}
+
+__device__ void QueuePair::ring_doorbell(uint64_t db_val) {
+  swap_endian_store(const_cast<uint32_t *>(dbrec_send),
+                    reinterpret_cast<uint32_t>(sq_counter));
+  STORE(db.ptr, db_val);
+  db.uint ^= 256;
+}
+
+__device__ void QueuePair::set_completion_flag_on_wqe(int num_wqes) {
+  uint64_t *wqe = &current_sq[8 * ((sq_counter - num_wqes) % max_nwqe)];
+  uint8_t *wqe_ce = reinterpret_cast<uint8_t *>(wqe) + 11;
+  *wqe_ce = 8;
+}
+
+template <>
+__device__ void QueuePair::update_wqe_ce_single<false>(int num_wqes) {
+  if (sq_counter % max_nwqe == (max_nwqe - 2)) {
+    set_completion_flag_on_wqe(num_wqes);
+    quiet_counter++;
+  }
+}
+
+template <>
+__device__ void QueuePair::update_wqe_ce_single<true>(int num_wqes) {
+  set_completion_flag_on_wqe(num_wqes);
+  quiet_counter++;
+}
+
+template <>
+__device__ void QueuePair::update_wqe_ce_thread<false>(int num_wqes) {}
+
+template <>
+__device__ void QueuePair::update_wqe_ce_thread<true>(int num_wqes) {
+  set_completion_flag_on_wqe(num_wqes);
+  atomicAdd(&quiet_counter, 1);
+}
+
+__device__ void QueuePair::compute_db_val_opcode(uint64_t *db_val,
+                                                 uint16_t dbrec_val,
+                                                 uint8_t opcode) {
+  uint64_t opcode64 = opcode;
+  opcode64 = opcode64 << 24 & 0x000000FFFF000000;
+
+  uint64_t dbrec = dbrec_val << 8;
+  dbrec = dbrec & 0x0000000000FFFF00;
+
+  uint64_t val = *db_val;
+  val = val & 0xFFFFFFFFFF0000FF;
+
+  *db_val = val | dbrec | opcode64;
+}
+
+template <class level>
+__device__ void QueuePair::quiet_internal() {
+  /*
+   * If there are nothing to quiet, just return early.
+   */
+  uint32_t quiet_val = quiet_counter;
+  if (!quiet_val) {
+    return;
+  }
+
+  profiler.incStat(QUIET_COUNT);
+  uint64_t start = profiler.startTimer();
+
+  /*
+   * Generate a pointer to the completion queue entry.
+   */
+  cq_consumer_counter = cq_consumer_counter + quiet_val - 1;
+  uint32_t indx = (cq_consumer_counter % cq_size);
+  mlx5_cqe64 *cqe_entry = &current_cq_q[indx];
+
+  /*
+   * Access the op_own value in the completion queue entry.
+   */
+  int val_ld = uncached_load_ubyte(&(cqe_entry->op_own));
+  uint8_t val_op_own = val_ld;
+
+  /*
+   * If the completion queue entry is not valid, wait for it to become so.
+   */
+  while (!((val_op_own & 0x1) == ((cq_consumer_counter >> cq_log_size) & 1)) ||
+         ((val_op_own) >> 4) == 0xF) {
+    val_ld = uncached_load_ubyte(&(cqe_entry->op_own));
+    val_op_own = val_ld;
+  }
+
+  /*
+   * Grab the opcode from the op_own field and report if it is an error.
+   */
+  uint8_t opcode = val_op_own >> 4;
+  if (opcode != 0) {
+    uint8_t syndrome = get_cq_error_syndrome(cqe_entry);
+    mlx5_err_cqe *cqe_err = reinterpret_cast<mlx5_err_cqe *>(cqe_entry);
+    GPU_DPRINTF("QUIET ERROR: signature %d opcode_qpn %llx wqe_cnt %llx \n",
+                syndrome, cqe_err->s_wqe_opcode_qpn, cqe_err->wqe_counter);
+  }
+
+  /*
+   * Decrement the quiet count by the amount determined at the beginning
+   * of this method.
+   *
+   * bpotter - There are two areas of concern in this method for me.
+   * 1) In multithreaded builds, we may need to make this method a critical
+   * section to prevent data races on these variables.
+   *
+   * 2) Is there a data race in the API if a one remote process calls quiet
+   * while another process continues adding events? Is it ever possible for
+   * a quiet to complete, but the quiet_counter decrement here is not set
+   * to zero?
+   */
+  level L;
+  L.decQuietCounter(&quiet_counter, quiet_val);
+
+  profiler.endTimer(start, POLL_CQ);
+  start = profiler.startTimer();
+
+  /*
+   * Increment the trailing index counter which tracks our spot in the
+   * completion queue.
+   */
+  cq_consumer_counter++;
+  swap_endian_store(const_cast<uint32_t *>(dbrec_cq), cq_consumer_counter);
+
+  profiler.endTimer(start, NEXT_CQ);
+}
+
+template <class level>
+__device__ void QueuePair::quiet_single() {
+  level L;
+  L.quiet(this);
+}
+
+template <class level>
+__device__ void QueuePair::quiet_single_heavy(int pe) {
+  level L;
+  L.quiet_heavy(this, pe);
+}
+
+template <class level, bool cqe>
+__device__ void QueuePair::update_posted_wqe_generic(
+    int pe, int32_t size, uintptr_t *laddr, uintptr_t *raddr, uint8_t opcode,
+    int64_t atomic_data, int64_t atomic_cmp, bool ring_db,
+    uint64_t atomic_ret_pos, bool zero_byte_rd) {
+  uint64_t start = profiler.startTimer();
+
+  level L;
+  L.postLock(this, pe);
+  uint32_t num_wqes = connection_policy.getNumWqes(opcode);
+
+  // Get the index for my thread's put in the SQ.
+  uint64_t my_sq_counter = L.threadAtomicAdd(&sq_counter, num_wqes);
+  uint64_t my_sq_index = my_sq_counter % max_nwqe;
+
+  // 16-bit little endian version of the SQ index needed to build the cntrl
+  // segment in the WQE.
+  uint16_t le_sq_counter;
+  uint16_t sq_counter_u16 = my_sq_counter;
+  swap_endian_store(&le_sq_counter, sq_counter_u16);
+
+  bool flag = sq_overflow;
+  uint32_t lkey_in_stack_frame = lkey;
+  uint32_t rkey_in_stack_frame = rkey;
+  uint32_t ctrl_qp_sq_in_stack_frame = ctrl_qp_sq;
+  uint64_t ctrl_sig_in_stack_frame = ctrl_sig;
+
+  connection_policy.setRkey(&rkey_in_stack_frame, pe);
+
+  if (opcode == MLX5_OPCODE_RDMA_WRITE && !size) {
+    rkey_in_stack_frame = hdp_rkey[pe];
+    size = 4;
+  }
+
+  /*
+   * Build out all the segments required for my WQE(s) based on the
+   * operation, starting at my_sq_index into the SQ. SegmentBuilder will
+   * keep track of placing the segments in the correct location.
+   */
+  SegmentBuilder seg_build(my_sq_index, current_sq);
+  seg_build.update_cntrl_seg(opcode, le_sq_counter, ctrl_qp_sq_in_stack_frame,
+                             ctrl_sig_in_stack_frame, &connection_policy,
+                             zero_byte_rd);
+  seg_build.update_connection_seg(pe, &connection_policy);
+  seg_build.update_rdma_seg(raddr, rkey_in_stack_frame);
+
+  if (opcode == MLX5_OPCODE_ATOMIC_FA || opcode == MLX5_OPCODE_ATOMIC_CS) {
+    seg_build.update_atomic_data_seg(atomic_data, atomic_cmp);
+    size = 8;
+    lkey_in_stack_frame = atomic_ret.atomic_lkey;
+    laddr = &atomic_ret.atomic_base_ptr[atomic_ret_pos];
+  }
+
+  if (size <= inline_threshold && opcode == MLX5_OPCODE_RDMA_WRITE) {
+    seg_build.update_inl_data_seg(laddr, size);
+  } else {
+    seg_build.update_data_seg(laddr, size, lkey_in_stack_frame);
+  }
+
+  profiler.incStat(WQE_COUNT);
+  profiler.endTimer(start, UPDATE_WQE);
+  start = profiler.startTimer();
+
+  L.template finishPost<cqe>(this, ring_db, num_wqes, pe, le_sq_counter,
+                             opcode);
+
+  profiler.incStat(DB_COUNT);
+  profiler.endTimer(start, RING_SQ_DB);
+}
+
+/******************************************************************************
+ ****************************** SHMEM INTERFACE *******************************
+ *****************************************************************************/
+template <class level>
+__device__ void QueuePair::put_nbi(void *dest, const void *source,
+                                   size_t nelems, int pe, bool db_ring) {
+  uintptr_t *src = reinterpret_cast<uintptr_t *>(const_cast<void *>(source));
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(dest);
+
+  update_posted_wqe_generic<level, false>(
+      pe, nelems, src, dst, MLX5_OPCODE_RDMA_WRITE, 0, 0, db_ring, 0);
+}
+
+template <class level>
+__device__ void QueuePair::put_nbi_cqe(void *dest, const void *source,
+                                       size_t nelems, int pe, bool db_ring) {
+  uintptr_t *src = reinterpret_cast<uintptr_t *>(const_cast<void *>(source));
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(dest);
+
+  update_posted_wqe_generic<level, true>(
+      pe, nelems, src, dst, MLX5_OPCODE_RDMA_WRITE, 0, 0, db_ring, 0);
+}
+
+template <class level>
+__device__ void QueuePair::get_nbi(void *dest, const void *source,
+                                   size_t nelems, int pe, bool db_ring) {
+  uintptr_t *src = reinterpret_cast<uintptr_t *>(const_cast<void *>(source));
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(dest);
+
+  update_posted_wqe_generic<level, false>(
+      pe, nelems, src, dst, MLX5_OPCODE_RDMA_READ, 0, 0, db_ring, 0);
+}
+
+template <class level>
+__device__ void QueuePair::get_nbi_cqe(void *dest, const void *source,
+                                       size_t nelems, int pe, bool db_ring) {
+  uintptr_t *src = reinterpret_cast<uintptr_t *>(const_cast<void *>(source));
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(dest);
+
+  update_posted_wqe_generic<level, true>(
+      pe, nelems, src, dst, MLX5_OPCODE_RDMA_READ, 0, 0, db_ring, 0);
+}
+
+template <class level>
+__device__ void QueuePair::zero_b_rd(int pe) {
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(base_heap[pe]);
+
+  update_posted_wqe_generic<level, true>(pe, 0, nullptr, dst,
+                                         MLX5_OPCODE_RDMA_READ, 0, 0, true, 0,
+                                         true);  // enable 0_byte read op
+}
+
+__device__ int64_t QueuePair::atomic_fetch(void *dest, int64_t value,
+                                           int64_t cond, int pe, bool db_ring,
+                                           uint8_t atomic_op) {
+  THREAD TH;
+  uint64_t pos = TH.threadAtomicAdd(
+      reinterpret_cast<unsigned long long *>(/* NOLINT(runtime/int) */
+                                             &atomic_ret.atomic_counter));
+
+  pos = pos % max_nb_atomic;
+
+  int64_t *atomic_base_ptr =
+      reinterpret_cast<int64_t *>(atomic_ret.atomic_base_ptr);
+
+  int64_t *load_address = &atomic_base_ptr[pos];
+
+  *load_address = -100;
+
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(dest);
+
+  update_posted_wqe_generic<THREAD, true>(pe, sizeof(int64_t), nullptr, dst,
+                                          atomic_op, value, cond, db_ring, pos);
+  quiet_single<THREAD>();
+
+  while (uncached_load(load_address) == -100) {
+  }
+
+  int64_t ret = *load_address;
+
+  __threadfence();
+
+  return ret;
+}
+
+__device__ void QueuePair::atomic_nofetch(void *dest, int64_t value,
+                                          int64_t cond, int pe, bool db_ring,
+                                          uint8_t atomic_op) {
+  THREAD TH;
+  uint64_t pos = TH.threadAtomicAdd(
+      reinterpret_cast<unsigned long long *>(/* NOLINT(runtime/int) */
+                                             &atomic_ret.atomic_counter));
+  pos = pos % max_nb_atomic;
+  uintptr_t *dst = reinterpret_cast<uintptr_t *>(dest);
+
+  update_posted_wqe_generic<THREAD, true>(pe, sizeof(int64_t), nullptr, dst,
+                                          atomic_op, value, cond, db_ring, pos);
+
+  quiet_single<THREAD>();
+}
+
+__device__ void QueuePair::fence(int pe) {
+  // TODO(khamidou): should this be replaced by a zero_byte_rd?
+  // FIXME: the relaxed ordering requires an intervening read to order
+  // prior operations.
+  auto remote_hdp_uncast = hdp_address[pe];
+  uintptr_t *remote_hdp = reinterpret_cast<uintptr_t *>(remote_hdp_uncast);
+  update_posted_wqe_generic<THREAD, true>(
+      pe, 0, nullptr, remote_hdp, MLX5_OPCODE_RDMA_WRITE, 0, 0, true, 0);
+}
+
+__device__ void QueuePair::waitCQSpace(int num_msgs) {
+  // We cannot post more outstanding requests than the completion queue
+  // size.  Force a quiet if we are out of space.
+  if ((quiet_counter + num_msgs) >= cq_size) {
+    GPU_DPRINTF(
+        "*** inside post_cq forcing flush: outstanding %d "
+        "adding %d cq_size %d\n",
+        quiet_counter, num_msgs, cq_size);
+
+    // TODO(khamidou): More targeted flush would be better here.
+    quiet_single<THREAD>();
+  }
+}
+
+__device__ void QueuePair::waitSQSpace(int num_msgs) {
+  // We cannot post more outstanding requests than the Send queue
+  // size.  Force a quiet if we are out of space.
+  local_sq_cnt += num_msgs;
+  int div = local_sq_cnt / max_nwqe;
+
+  if (div > 0) {
+    GPU_DPRINTF(
+        "*** inside waitSQSpace forcing flush to overrun the SQ"
+        " sq_counter %d  adding %d quiet_conter %d \n",
+        sq_counter, num_msgs, max_nwqe, quiet_counter);
+
+    quiet_single<THREAD>();
+    local_sq_cnt = local_sq_cnt % max_nwqe;
+  }
+}
+
+void QueuePair::setDBval(uint64_t val) { db_val = val; }
+
+#define THREAD_LEVEL_GEN(T)                                                 \
+  template __device__ void QueuePair::put_nbi<T>(                           \
+      void *dest, const void *source, size_t nelems, int pe, bool db_ring); \
+  template __device__ void QueuePair::put_nbi_cqe<T>(                       \
+      void *dest, const void *source, size_t nelems, int pe, bool db_ring); \
+  template __device__ void QueuePair::get_nbi<T>(                           \
+      void *dest, const void *source, size_t nelems, int pe, bool db_ring); \
+  template __device__ void QueuePair::get_nbi_cqe<T>(                       \
+      void *dest, const void *source, size_t nelems, int pe, bool db_ring); \
+  template __device__ void QueuePair::zero_b_rd<T>(int pe);                 \
+  template __device__ void QueuePair::quiet_single<T>();                    \
+  template __device__ void QueuePair::quiet_single_heavy<T>(int pe);        \
+  template __device__ void QueuePair::quiet_internal<T>();
+
+THREAD_LEVEL_GEN(THREAD)
+THREAD_LEVEL_GEN(WG)
+THREAD_LEVEL_GEN(WAVE)
+
+}  // namespace rocshmem

src/gpu_ib/queue_pair.hpp

@@ -0,0 +1,431 @@
+/******************************************************************************
+ * 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_GPU_IB_QUEUE_PAIR_HPP_
+#define LIBRARY_SRC_GPU_IB_QUEUE_PAIR_HPP_
+
+/**
+ * @file queue_pair.hpp
+ *
+ * @section DESCRIPTION
+ * An IB QueuePair (SQ and CQ) that the device can use to perform network
+ * operations. Most important ROC_SHMEM operations are performed by this
+ * class.
+ */
+
+#include <infiniband/mlx5dv.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "../atomic_return.hpp"
+#include "connection_policy.hpp"
+#include "thread_policy.hpp"
+#include "../hdp_policy.hpp"
+#include "../stats.hpp"
+
+namespace rocshmem {
+
+class GPUIBBackend;
+
+enum gpu_ib_stats {
+  RING_SQ_DB = 0,
+  UPDATE_WQE,
+  POLL_CQ,
+  NEXT_CQ,
+  QUIET_COUNT,
+  DB_COUNT,
+  WQE_COUNT,
+  MEM_WAIT,
+  INIT,
+  FINALIZE,
+  GPU_IB_NUM_STATS
+};
+
+typedef union db_reg {
+  uint64_t *ptr;
+  uintptr_t uint;
+} db_reg_t;
+
+class QueuePair {
+ public:
+  /**
+   * @brief Constructor.
+   *
+   * @param[in] backend Backend needed for member access.
+   */
+  explicit QueuePair(GPUIBBackend *backend);
+
+  /**
+   * @brief Destructor.
+   */
+  __device__ ~QueuePair();
+
+  /**
+   * @brief Inspect completion queue and possibly wait for free space.
+   *
+   * @param[in] num_msgs Number of entries needing space in completion queue.
+   */
+  __device__ void waitCQSpace(int num_msgs);
+
+  /**
+   * @brief Inspect send queue and possibly wait for free space.
+   *
+   * @param[in] num_msgs Number of entries needing space in send queue.
+   */
+  __device__ void waitSQSpace(int num_msgs);
+
+  /**
+   * @brief Create and enqueue a non-blocking put work queue entry (wqe).
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   * @param[in] dest Destination address for data transmission.
+   * @param[in] source Source address for data transmission.
+   * @param[in] nelems Size in bytes of data transmission.
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] db_ring Denotes whether send queue door bell should be rung.
+   */
+  template <class level>
+  __device__ void put_nbi(void *dest, const void *source, size_t nelems, int pe,
+                          bool db_ring);
+
+  /**
+   * @brief Create and enqueue a non-blocking put work queue entry (wqe).
+   *
+   * @note This variant differs from put_nbi by requesting that a completion
+   * queue entry is generated in the completion queue.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   * @param[in] dest Destination address for data transmission.
+   * @param[in] source Source address for data transmission.
+   * @param[in] nelems Size in bytes of data transmission.
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] db_ring Denotes whether send queue door bell should be rung.
+   */
+  template <class level>
+  __device__ void put_nbi_cqe(void *dest, const void *source, size_t nelems,
+                              int pe, bool db_ring);
+
+  /**
+   * @brief Consume a completion queue entry from this queue pair's
+   * completion queue.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   */
+  template <class level>
+  __device__ void quiet_single();
+
+  /**
+   * @brief Send a zero-byte read to enforce ordering and then consume
+   * a completion queue entry from this queue pair's completion queue.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   * @param[in] pe Processing element id to send the zero_b_rd.
+   */
+  template <class level>
+  __device__ void quiet_single_heavy(int pe);
+
+  /**
+   * @brief Create and enqueue a HDP flush work queue entry on the remote PE.
+   *
+   * @param[in] pe Processing element id to send the HDP flush operation.
+   *
+   * TODO(@khamidou): does this require a zero_b_rd to enforce write ordering
+   * The HDP flush is itself a write. Could this write be reordered with
+   * respect to other write on the network and arrive out-of-order?
+   */
+  __device__ void fence(int pe);
+
+  /**
+   * @brief Create and enqueue a non-blocking get work queue entry (wqe).
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   * @param[in] dest Destination address for data transmission.
+   * @param[in] source Source address for data transmission.
+   * @param[in] nelems Size in bytes of data transmission.
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] db_ring Denotes whether send queue door bell should be rung.
+   */
+  template <class level>
+  __device__ void get_nbi(void *dest, const void *source, size_t nelems, int pe,
+                          bool db_ring);
+
+  /**
+   * @brief Create and enqueue a non-blocking get work queue entry (wqe).
+   *
+   * @note This variant differs from get_nbi by requesting that a completion
+   * queue entry is generated in the completion queue.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   * @param[in] dest Destination address for data transmission.
+   * @param[in] source Source address for data transmission.
+   * @param[in] nelems Size in bytes of data transmission.
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] db_ring Denotes whether send queue door bell should be rung.
+   */
+  template <class level>
+  __device__ void get_nbi_cqe(void *dest, const void *source, size_t nelems,
+                              int pe, bool db_ring);
+
+  /**
+   * @brief Create and enqueue a zero-byte read to enforce write ordering.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   * @param[in] pe Processing element id to send the zero_b_rd.
+   */
+  template <class level>
+  __device__ void zero_b_rd(int pe);
+
+  /**
+   * @brief Create and enqueue an atomic fetch work queue entry (wqe).
+   *
+   * @param[in] dest Destination address for data transmission.
+   * @param[in] value Data value for the atomic operation.
+   * @param[in] cond Used in atomic comparisons.
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] db_ring Denotes whether send queue door bell should be rung.
+   * @param[in] atomic_op The atomic operation to perform.
+   *
+   * @return An atomic value
+   */
+  __device__ int64_t atomic_fetch(void *dest, int64_t value, int64_t cond,
+                                  int pe, bool db_ring, uint8_t atomic_op);
+
+  /**
+   * @brief Create and enqueue an atomic fetch work queue entry (wqe).
+   *
+   * @param[in] dest Destination address for data transmission.
+   * @param[in] value Data value for the atomic operation.
+   * @param[in] cond Used in atomic comparisons.
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] db_ring Denotes whether send queue door bell should be rung.
+   * @param[in] atomic_op The atomic operation to perform.
+   */
+  __device__ void atomic_nofetch(void *dest, int64_t value, int64_t cond,
+                                 int pe, bool db_ring, uint8_t atomic_op);
+
+  /**
+   * @brief Helper method to set the doorbell's value.
+   *
+   * @param[in] val Desired value for the doorbell.
+   */
+  void setDBval(uint64_t val);
+
+ protected:
+  /**
+   * @brief Helper method to build work requests for the send queue.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   * @tparam cqe Flag to optionally generate cqes.
+   *
+   * @param[in] pe Destination processing element of data transmission.
+   * @param[in] size Size in bytes of data transmission.
+   * @param[in] laddr Local address.
+   * @param[in] raddr Remote address.
+   * @param[in] opcode Operation to be performed.
+   * @param[in] atomic_data An atomic data value to be used.
+   * @param[in] atomic_cmp An atomic comparison operation to be performed.
+   * @param[in] ring_db Boolean denoting if doorbell should be rung.
+   * @param[in] atomic_ret_pos Index into atomic return structure.
+   * @param[in] zero_byte_rd Boolean if zero byte read should be used.
+   */
+  template <class level, bool cqe>
+  __device__ __attribute__((noinline)) void update_posted_wqe_generic(
+      int pe, int32_t size, uintptr_t *laddr, uintptr_t *raddr, uint8_t opcode,
+      int64_t atomic_data, int64_t atomic_cmp, bool ring_db,
+      uint64_t atomic_ret_pos, bool zero_byte_rd = false);
+
+  /**
+   * @brief Helper method to drain completion queue entries.
+   *
+   * @tparam level Implements specific behaviors for thread, warp, block access.
+   *
+   */
+  template <class level>
+  __device__ __attribute__((noinline)) void quiet_internal();
+
+  /**
+   * @brief Helper method to compute doorbell value opcode which is used to
+   * ring the doorbell.
+   *
+   * @param[in,out] db_val
+   * @param[in] dbrec_val
+   * @param[in] opcode
+   */
+  __device__ void compute_db_val_opcode(uint64_t *db_val, uint16_t dbrec_val,
+                                        uint8_t opcode);
+
+  /**
+   * @brief Helper method that sets the field in a work queue entry to
+   * generate a completion entry in the completion queue.
+   *
+   * @param num_wqes Number of work entries this completion entry represents.
+   */
+  __device__ void set_completion_flag_on_wqe(int num_wqes);
+
+  /**
+   * @brief Helper method to update fields for the work queue entry.
+   *
+   * @tparam cqe Flag to optionally generate cqes.
+   *
+   * @note Single variant is meant to be callable by a block leader.
+   */
+  template <bool cqe>
+  __device__ void update_wqe_ce_single(int num_wqes);
+
+  /**
+   * @brief Helper method to update fields for the work queue entry.
+   *
+   * @tparam cqe Flag to optionally generate cqes.
+   *
+   * @note Thread variant is meant to be callable by multiple threads.
+   */
+  template <bool cqe>
+  __device__ void update_wqe_ce_thread(int num_wqes);
+
+  /**
+   * @brief Helper method to ring the doorbell
+   *
+   * @param[in] db_val Doorbell value is written by method.
+   */
+  __device__ void ring_doorbell(uint64_t db_val);
+
+  /**
+   * @brief Helper method to extract syndrome field from cqe.
+   *
+   * @param[in] cq_entry Completion queue entry.
+   */
+  __device__ uint8_t get_cq_error_syndrome(mlx5_cqe64 *cq_entry);
+
+ private:
+  const int inline_threshold{8};
+
+  /* TODO(bpotter): Most of these should be private/protected */
+ public:
+#ifdef PROFILE
+  typedef Stats<GPU_IB_NUM_STATS> GPUIBStats;
+#else
+  typedef NullStats<GPU_IB_NUM_STATS> GPUIBStats;
+#endif
+
+  /*
+   * Pointer to the hardware doorbell register for the QP.
+   */
+  db_reg_t db{};
+
+  /*
+   * Base pointer of this QP's SQ
+   * TODO(bpotter): Use the correct struct type for this.
+   */
+  uint64_t *current_sq{nullptr};
+  uint64_t *current_sq_H{nullptr};
+
+  /*
+   * Base pointer of this QP's CQ
+   */
+  mlx5_cqe64 *current_cq_q{nullptr};
+  mlx5_cqe64 *current_cq_q_H{nullptr};
+
+  /*
+   * Pointer to the doorbell record for this SQ.
+   */
+  volatile uint32_t *dbrec_send{nullptr};
+
+  /*
+   * Pointer to the doorbell record for the CQ.
+   */
+  volatile uint32_t *dbrec_cq{nullptr};
+
+  uint32_t *hdp_rkey{nullptr};
+
+  uintptr_t *hdp_address{nullptr};
+
+  HdpPolicy *hdp_policy{};
+
+  atomic_ret_t atomic_ret{};
+
+  ThreadImpl threadImpl{};
+
+  ConnectionImpl connection_policy;
+
+  char *const *base_heap{nullptr};
+  /*
+   * Current index into the SQ (non-modulo size).
+   */
+  uint32_t sq_counter{0};
+  uint32_t local_sq_cnt{0};
+
+  /*
+   * Number of outstanding messages on this QP that need to be completed
+   * during a quiet operation.
+   */
+  uint32_t quiet_counter{0};
+
+  int num_cqs{0};
+
+  /*
+   * Current index into the SQ (non-module size).
+   */
+  uint32_t cq_consumer_counter{0};
+
+  uint16_t cq_log_size{0};
+
+  uint16_t cq_size{0};
+
+  uint32_t ctrl_qp_sq{0};
+
+  uint64_t ctrl_sig{0};
+
+  uint32_t rkey{0};
+
+  uint32_t lkey{0};
+
+  GPUIBStats profiler{};
+
+  uint16_t max_nwqe{0};
+
+  bool sq_overflow{0};
+
+  uint64_t db_val{};
+  /*
+   * Pointer to the QP in global memory that this QP is copied from.  When
+   * this QP is destroyed, the dynamic (indicies, stats, etc) in the
+   * global_qp are updated.
+   */
+  QueuePair *global_qp{nullptr};
+
+  friend SingleThreadImpl;
+  friend MultiThreadImpl;
+  friend THREAD;
+  friend WG;
+  friend WAVE;
+  friend RCConnectionImpl;
+  friend DCConnectionImpl;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_QUEUE_PAIR_HPP_

src/gpu_ib/reliable_connection.cpp

@@ -0,0 +1,201 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "reliable_connection.hpp"
+
+#include <mpi.h>
+
+#include "backend_ib.hpp"
+
+namespace rocshmem {
+
+ReliableConnection::ReliableConnection(GPUIBBackend* b) : Connection(b, 0) {}
+
+ReliableConnection::~ReliableConnection() {}
+
+Connection::InitQPState ReliableConnection::initqp(uint8_t port) {
+  InitQPState init{};
+
+  init.exp_qp_attr.qp_access_flags =
+      IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_LOCAL_WRITE |
+      IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_ATOMIC;
+  init.exp_qp_attr.port_num = port;
+
+  init.exp_attr_mask |= IBV_QP_ACCESS_FLAGS;
+
+  return init;
+}
+
+Connection::RtrState ReliableConnection::rtr(dest_info_t* dest, uint8_t port) {
+  RtrState rtr{};
+
+  rtr.exp_qp_attr.dest_qp_num = dest->qpn;
+  rtr.exp_qp_attr.rq_psn = dest->psn;
+  rtr.exp_qp_attr.ah_attr.port_num = port;
+  if (ib_state->portinfo.link_layer == IBV_LINK_LAYER_INFINIBAND) {
+    rtr.exp_qp_attr.ah_attr.dlid = dest->lid;
+  } else {
+    rtr.exp_qp_attr.ah_attr.is_global = 1;
+    rtr.exp_qp_attr.ah_attr.grh.dgid = dest->gid;
+    rtr.exp_qp_attr.ah_attr.grh.sgid_index = 0;
+    rtr.exp_qp_attr.ah_attr.grh.hop_limit = 1;
+  }
+
+  rtr.exp_attr_mask |= IBV_QP_DEST_QPN | IBV_QP_RQ_PSN |
+                       IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER;
+
+  return rtr;
+}
+
+Connection::RtsState ReliableConnection::rts(dest_info_t* dest) {
+  RtsState rts{};
+
+  rts.exp_qp_attr.sq_psn = dest->psn;
+
+  rts.exp_attr_mask |= IBV_QP_SQ_PSN;
+
+  return rts;
+}
+
+ibv_qp* ReliableConnection::create_qp_0(ibv_context* context,
+                                        ibv_qp_init_attr_ex* qp_attr) {
+  return ibv_create_qp_ex(context, qp_attr);
+}
+
+void ReliableConnection::create_qps_1() { }
+
+void ReliableConnection::create_qps_2(int port, int my_rank,
+                                        ibv_port_attr* ib_port_att) { }
+
+void ReliableConnection::create_qps_3(int port, ibv_qp* qp, int offset,
+                                        ibv_port_attr* ib_port_att) {
+  init_qp_status(qp, port);
+
+  all_qp[offset].lid = ib_port_att->lid;
+  all_qp[offset].qpn = qp->qp_num;
+  all_qp[offset].psn = 0;
+  union ibv_gid gid;
+  ibv_query_gid(ib_state->context, port, 0, &gid);
+  all_qp[offset].gid = gid;
+}
+
+void ReliableConnection::get_remote_conn(int* remote_conn) {
+  *remote_conn = backend->num_pes;
+}
+
+void ReliableConnection::allocate_dynamic_members(int num_blocks) {
+  all_qp.resize(backend->num_pes * num_blocks);
+}
+
+void ReliableConnection::free_dynamic_members() {
+}
+
+void ReliableConnection::initialize_1(int port, int num_blocks) {
+  MPI_Alltoall(MPI_IN_PLACE, sizeof(dest_info_t) * num_blocks, MPI_CHAR,
+               all_qp.data(), sizeof(dest_info_t) * num_blocks, MPI_CHAR,
+               backend->thread_comm);
+
+  for (int i = 0; i < qps.size(); i++) {
+    change_status_rtr(qps[i], &all_qp[i], port);
+  }
+
+  MPI_Barrier(backend->thread_comm);
+
+  for (int i = 0; i < qps.size(); i++) {
+    change_status_rts(qps[i], &all_qp[i]);
+  }
+}
+
+void ReliableConnection::initialize_rkey_handle(uint32_t** heap_rkey_handle,
+                                                  ibv_mr* mr) {
+  CHECK_HIP(
+      hipHostMalloc(heap_rkey_handle, sizeof(uint32_t) * backend->num_pes));
+  (*heap_rkey_handle)[backend->my_pe] = mr->rkey;
+}
+
+void ReliableConnection::free_rkey_handle(uint32_t* heap_rkey_handle) {
+  CHECK_HIP(hipHostFree(heap_rkey_handle));
+}
+
+Connection::QPInitAttr ReliableConnection::qpattr(ibv_qp_cap cap) {
+  QPInitAttr qpattr(cap);
+  qpattr.attr.qp_type = IBV_QPT_RC;
+  return qpattr;
+}
+
+void ReliableConnection::post_dv_rc_wqe(int remote_conn) {
+  mlx5_wqe_ctrl_seg* ctrl;
+  mlx5_wqe_raddr_seg* rdma;
+  mlx5_wqe_data_seg* data;
+
+  for (int i = 0; i < remote_conn; i++) {
+    int num_blocks = backend->num_blocks_;
+    for (int j = 0; j < num_blocks; j++) {
+      int qp_index = i * num_blocks + j;
+      uint64_t* ptr = get_address_sq(qp_index);
+
+      const uint16_t nb_post = 1;  // 4 * sq_size;
+      for (uint16_t index = 0; index < nb_post; index++) {
+        uint8_t op_mod = 0;
+        uint8_t op_code = 8;
+        uint32_t qp_num = qps[qp_index]->qp_num;
+        uint8_t fm_ce_se = 0;
+        uint8_t ds = 3;
+        ctrl = reinterpret_cast<mlx5_wqe_ctrl_seg*>(ptr);
+        mlx5dv_set_ctrl_seg(ctrl, index, op_code, op_mod, qp_num, fm_ce_se, ds,
+                            0, 0);
+        ptr = ptr + 2;
+
+        rdma = reinterpret_cast<mlx5_wqe_raddr_seg*>(ptr);
+        const auto& heap_bases = backend->heap.get_heap_bases();
+        auto temp = heap_bases[(backend->my_pe + 1) % 2];
+        uint64_t r_address = reinterpret_cast<uint64_t>(temp);
+        uint32_t rkey = backend->networkImpl.heap_rkey[i];
+        set_rdma_seg(rdma, r_address, rkey);
+        ptr = ptr + 2;
+
+        data = reinterpret_cast<mlx5_wqe_data_seg*>(ptr);
+        uint32_t lkey = backend->networkImpl.heap_mr->lkey;
+        temp = heap_bases[backend->my_pe];
+        uint64_t address = reinterpret_cast<uint64_t>(temp);
+        mlx5dv_set_data_seg(data, 1, lkey, address);
+        ptr = ptr + 4;
+      }
+    }
+  }
+}
+
+// TODO(bpotter): remove redundancies with the other derived class
+void ReliableConnection::post_wqes() {
+  int remote_conn;
+  get_remote_conn(&remote_conn);
+  post_dv_rc_wqe(remote_conn);
+}
+
+void ReliableConnection::initialize_wr_fields(ibv_send_wr* wr, ibv_ah* ah,
+                                              int dc_key) {}
+
+int ReliableConnection::get_sq_dv_offset(int pe_idx, int num_qps, int wg_idx) {
+  return pe_idx * num_qps + wg_idx;
+}
+
+}  // namespace rocshmem

src/gpu_ib/reliable_connection.hpp

@@ -0,0 +1,84 @@
+/******************************************************************************
+ * 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_GPU_IB_RELIABLE_CONNECTION_HPP_
+#define LIBRARY_SRC_GPU_IB_RELIABLE_CONNECTION_HPP_
+
+#include <vector>
+
+#include "connection.hpp"
+
+namespace rocshmem {
+
+class ReliableConnection : public Connection {
+ public:
+  explicit ReliableConnection(GPUIBBackend* backend);
+
+  ~ReliableConnection() override;
+
+  void get_remote_conn(int* remote_conn) override;
+
+  void post_wqes() override;
+
+  void initialize_rkey_handle(uint32_t** heap_rkey_handle,
+                                ibv_mr* mr) override;
+
+  void free_rkey_handle(uint32_t* heap_rkey_handle) override;
+
+ private:
+  InitQPState initqp(uint8_t port) override;
+
+  RtrState rtr(dest_info_t* dest, uint8_t port) override;
+
+  RtsState rts(dest_info_t* dest) override;
+
+  QPInitAttr qpattr(ibv_qp_cap cap) override;
+
+  void create_qps_1() override;
+
+  void create_qps_2(int port, int my_rank,
+                      ibv_port_attr* ib_port_att) override;
+
+  void create_qps_3(int port, ibv_qp* qp, int offset,
+                      ibv_port_attr* ib_port_att) override;
+
+  ibv_qp* create_qp_0(ibv_context* context,
+                      ibv_qp_init_attr_ex* qp_attr) override;
+
+  void allocate_dynamic_members(int num_wg) override;
+
+  void free_dynamic_members() override;
+
+  void initialize_1(int port, int num_wg) override;
+
+  void initialize_wr_fields(ibv_send_wr* wr, ibv_ah* ah, int dc_key) override;
+
+  int get_sq_dv_offset(int pe_idx, int num_qps, int wg_idx) override;
+
+  std::vector<dest_info_t> all_qp;
+
+  void post_dv_rc_wqe(int remote_conn);
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_RELIABLE_CONNECTION_HPP_

src/gpu_ib/segment_builder.cpp

@@ -0,0 +1,130 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "segment_builder.hpp"
+
+#include "endian.hpp"
+
+namespace rocshmem {
+
+__device__ SegmentBuilder::SegmentBuilder(uint64_t wqe_idx, void *base) {
+  mlx5_segment *base_ptr = static_cast<mlx5_segment *>(base);
+  size_t segment_offset = SEGMENTS_PER_WQE * wqe_idx;
+  seg_ptr = &base_ptr[segment_offset];
+}
+
+__device__ void SegmentBuilder::update_cntrl_seg(
+    uint8_t opcode, uint16_t wqe_idx, uint32_t ctrl_qp_sq, uint64_t ctrl_sig,
+    ConnectionImpl *connection_policy, bool zero_byte_rd) {
+  mlx5_wqe_ctrl_seg *ctrl_seg = &seg_ptr->ctrl_seg;
+
+  ctrl_seg->opmod_idx_opcode = (opcode << 24) | (wqe_idx << 8);
+
+  uint32_t DS = 2;
+  if (zero_byte_rd == false) {
+    DS = (opcode == MLX5_OPCODE_RDMA_WRITE || opcode == MLX5_OPCODE_RDMA_READ)
+             ? 3
+             : 4;
+  }
+
+  DS += connection_policy->wqeCntrlOffset();
+
+  ctrl_seg->qpn_ds = (DS << 24) | ctrl_qp_sq;
+
+  ctrl_seg->signature = ctrl_sig;
+
+  ctrl_seg->fm_ce_se = ctrl_sig >> 24;
+
+  ctrl_seg->imm = ctrl_sig >> 32;
+
+  seg_ptr++;
+}
+
+__device__ void SegmentBuilder::update_atomic_data_seg(uint64_t atomic_data,
+                                                       uint64_t atomic_cmp) {
+  mlx5_wqe_atomic_seg *atomic_seg = &seg_ptr->atomic_seg;
+
+  swap_endian_store(reinterpret_cast<uint64_t *>(&atomic_seg->swap_add),
+                    atomic_data);
+
+  swap_endian_store(reinterpret_cast<uint64_t *>(&atomic_seg->compare),
+                    atomic_cmp);
+
+  seg_ptr++;
+}
+
+__device__ void SegmentBuilder::update_rdma_seg(uintptr_t *raddr,
+                                                uint32_t rkey) {
+  mlx5_wqe_raddr_seg *raddr_seg = &seg_ptr->raddr_seg;
+
+  raddr_seg->rkey = rkey;
+
+  swap_endian_store(reinterpret_cast<uint64_t *>(&raddr_seg->raddr),
+                    reinterpret_cast<uint64_t>(raddr));
+
+  seg_ptr++;
+}
+
+__device__ void SegmentBuilder::update_data_seg(uintptr_t *laddr, int32_t size,
+                                                uint32_t lkey) {
+  if (laddr == nullptr) {
+    return;
+  }
+
+  mlx5_wqe_data_seg *data_seg = &seg_ptr->data_seg;
+
+  data_seg->lkey = lkey;
+
+  swap_endian_store(&data_seg->byte_count, size & 0x7FFFFFFFU);
+
+  swap_endian_store(reinterpret_cast<uint64_t *>(&data_seg->addr),
+                    reinterpret_cast<uint64_t>(laddr));
+
+  seg_ptr++;
+}
+
+__device__ void SegmentBuilder::update_inl_data_seg(uintptr_t *laddr,
+                                                    int32_t size) {
+  mlx5_wqe_inl_data_seg *inl_data_seg = &seg_ptr->inl_data_seg;
+
+  swap_endian_store(&inl_data_seg->byte_count, (size & 0x3FF) | 0x80000000);
+
+  // Assume fence HDP flush
+  // TODO(khamidou): Rework fence interface to avoid this
+  if (!laddr) {
+    uint8_t flush_val = 1;
+    memcpy(inl_data_seg + 1, &flush_val, sizeof(flush_val));
+  } else {
+    memcpy(inl_data_seg + 1, laddr, size);
+  }
+
+  seg_ptr++;
+}
+
+__device__ void SegmentBuilder::update_connection_seg(
+    int pe, ConnectionImpl *conn_policy) {
+  if (conn_policy->updateConnectionSegmentImpl(&seg_ptr->base_av, pe)) {
+    seg_ptr++;
+  }
+}
+
+}  // namespace rocshmem

src/gpu_ib/segment_builder.hpp

@@ -0,0 +1,64 @@
+/******************************************************************************
+ * 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_GPU_IB_SEGMENT_BUILDER_HPP_
+#define LIBRARY_SRC_GPU_IB_SEGMENT_BUILDER_HPP_
+
+#include <infiniband/mlx5dv.h>
+
+#include "connection_policy.hpp"
+#include "infiniband_structs.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+class SegmentBuilder {
+ public:
+  __device__ SegmentBuilder(uint64_t wqe_idx, void *base);
+
+  __device__ void update_cntrl_seg(uint8_t opcode, uint16_t wqe_idx,
+                                   uint32_t ctrl_qp_sq, uint64_t ctrl_sig,
+                                   ConnectionImpl *connection_policy,
+                                   bool zero_byte_rd);
+
+  __device__ void update_connection_seg(int pe,
+                                        ConnectionImpl *connection_policy);
+
+  __device__ void update_atomic_data_seg(uint64_t atomic_data,
+                                         uint64_t atomic_cmp);
+
+  __device__ void update_rdma_seg(uintptr_t *raddr, uint32_t rkey);
+
+  __device__ void update_inl_data_seg(uintptr_t *laddr, int32_t size);
+
+  __device__ void update_data_seg(uintptr_t *laddr, int32_t size,
+                                  uint32_t lkey);
+
+ private:
+  const int SEGMENTS_PER_WQE = 4;
+
+  mlx5_segment *seg_ptr;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_SEGMENT_BUILDER_HPP_

src/gpu_ib/thread_policy.cpp

@@ -0,0 +1,358 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "thread_policy.hpp"
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "queue_pair.hpp"
+
+namespace rocshmem {
+
+__device__ void SingleThreadImpl::quiet(QueuePair *handle) {
+  handle->quiet_internal<THREAD>();
+}
+
+__device__ void SingleThreadImpl::quiet_heavy(QueuePair *handle, int pe) {
+  handle->zero_b_rd<THREAD>(pe);
+  handle->quiet_internal<THREAD>();
+}
+
+__device__ void MultiThreadImpl::quiet(QueuePair *handle) {
+  int thread_id = get_flat_block_id();
+  /*
+   * Each WF selects one thread to perform the quiet.  Only one thread
+   * per WG is allowed to do a quiet at once to avoid races with the CQ.
+   */
+  if (thread_id % WF_SIZE == lowerID()) {
+    while (atomicCAS(&(handle->threadImpl.cq_lock), 0, 1) == 1) {
+    }
+    handle->quiet_internal<THREAD>();
+    __threadfence();
+    handle->threadImpl.cq_lock = 0;
+  }
+}
+
+__device__ void MultiThreadImpl::quiet_heavy(QueuePair *handle, int pe) {
+  int thread_id = get_flat_block_id();
+  /*
+   * Each WF selects one thread to perform the quiet.  Only one thread
+   * per WG is allowed to do a quiet at once to avoid races with the CQ.
+   */
+  if (thread_id % WF_SIZE == lowerID()) {
+    // zero_byte read
+    handle->zero_b_rd<THREAD>(pe);
+
+    while (atomicCAS(&(handle->threadImpl.cq_lock), 0, 1) == 1) {
+    }
+    handle->quiet_internal<THREAD>();
+    __threadfence();
+    handle->threadImpl.cq_lock = 0;
+  }
+}
+
+__device__ void WG::quiet(QueuePair *handle) { handle->quiet_internal<WG>(); }
+
+__device__ void WG::quiet_heavy(QueuePair *handle, int pe) {
+  handle->zero_b_rd<THREAD>(pe);
+  handle->quiet_internal<WG>();
+}
+
+__device__ void WAVE::quiet(QueuePair *handle) {
+  int thread_id = get_flat_block_id();
+  /*
+   * Each WF selects one thread to perform the quiet.  Only one thread
+   * per WG is allowed to do a quiet at once to avoid races with the CQ.
+   */
+
+  if (thread_id % WF_SIZE == 0) {
+    while (atomicCAS(&(handle->threadImpl.cq_lock), 0, 1) == 1) {
+    }
+    handle->quiet_internal<WAVE>();
+    __threadfence();
+    handle->threadImpl.cq_lock = 0;
+  }
+}
+
+__device__ void WAVE::quiet_heavy(QueuePair *handle, int pe) {
+  int thread_id = get_flat_block_id();
+  /*
+   * Each WF selects one thread to perform the quiet.  Only one thread
+   * per WG is allowed to do a quiet at once to avoid races with the CQ.
+   */
+  if (thread_id % WF_SIZE == 0) {
+    // post a zero-byte read
+    handle->zero_b_rd<THREAD>(pe);
+    while (atomicCAS(&(handle->threadImpl.cq_lock), 0, 1) == 1) {
+    }
+    handle->quiet_internal<WAVE>();
+    __threadfence();
+    handle->threadImpl.cq_lock = 0;
+  }
+}
+
+__device__ void SingleThreadImpl::decQuietCounter(uint32_t *quiet_counter,
+                                                  int num) {
+  *quiet_counter -= num;
+}
+
+__device__ void MultiThreadImpl::decQuietCounter(uint32_t *quiet_counter,
+                                                 int num) {
+  atomicSub(quiet_counter, num);
+}
+
+__device__ void WG::decQuietCounter(uint32_t *quiet_counter, int num) {
+  *quiet_counter -= num;
+}
+
+__device__ void WAVE::decQuietCounter(uint32_t *quiet_counter, int num) {
+  *quiet_counter -= num;
+}
+
+template <bool cqe>
+__device__ void SingleThreadImpl::finishPost(QueuePair *handle, bool ring_db,
+                                             int num_wqes, int pe,
+                                             uint16_t le_sq_counter,
+                                             uint8_t opcode) {
+  if (ring_db) {
+    uint64_t db_val = handle->db_val;
+    handle->compute_db_val_opcode(&db_val, le_sq_counter, opcode);
+    handle->update_wqe_ce_single<cqe>(num_wqes);
+    handle->ring_doorbell(db_val);
+  }
+}
+
+template <bool cqe>
+__device__ void MultiThreadImpl::finishPost(QueuePair *handle, bool ring_db,
+                                            int num_wqes, int pe,
+                                            uint16_t le_sq_counter,
+                                            uint8_t opcode) {
+  /*
+   * For RC, we can't allow a wave to have different PEs in it, else the
+   * doorbell ringing logic will not work.  This little for loop forces
+   * control flow divergence based on the PE.  It works well for small
+   * numbers of PEs, but we might want a different solution for large
+   * numbers.
+   */
+  if (handle->connection_policy.forcePostDivergence()) {
+    for (int i = 0; i < handle->num_cqs; i++) {
+      if (i != pe) {
+        continue;
+      }
+
+      finishPost_internal<cqe>(handle, ring_db, num_wqes, pe, le_sq_counter,
+                               opcode);
+    }
+  } else {
+    finishPost_internal<cqe>(handle, ring_db, num_wqes, pe, le_sq_counter,
+                             opcode);
+  }
+}
+
+template <bool cqe>
+__device__ void MultiThreadImpl::finishPost_internal(QueuePair *handle,
+                                                     bool ring_db, int num_wqes,
+                                                     int pe,
+                                                     uint16_t le_sq_counter,
+                                                     uint8_t opcode) {
+  /*
+   * Assuming here that postLock locks out all wavefronts in this WG but
+   * one, and that this will select a single thread in the wavefront.
+   */
+  __threadfence();
+  if (get_flat_block_id() % WF_SIZE == lowerID()) {
+    if (ring_db) {
+      uint64_t db_val =
+          handle->current_sq[8 * ((handle->sq_counter - num_wqes) %
+                                  handle->max_nwqe)];
+      handle->update_wqe_ce_thread<true>(num_wqes);
+      handle->ring_doorbell(db_val);
+    }
+
+    handle->threadImpl.sq_lock = 0;
+  }
+}
+
+template <bool cqe>
+__device__ void WG::finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                               int pe, uint16_t le_sq_counter, uint8_t opcode) {
+  if (ring_db) {
+    uint64_t db_val = handle->current_sq[8 * ((handle->sq_counter - num_wqes) %
+                                              handle->max_nwqe)];
+    handle->update_wqe_ce_single<cqe>(num_wqes);
+    handle->ring_doorbell(db_val);
+  }
+}
+
+template <bool cqe>
+__device__ void WAVE::finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                                 int pe, uint16_t le_sq_counter,
+                                 uint8_t opcode) {
+  if (ring_db) {
+    uint64_t db_val = handle->current_sq[8 * ((handle->sq_counter - num_wqes) %
+                                              handle->max_nwqe)];
+    handle->update_wqe_ce_thread<cqe>(num_wqes);
+    handle->ring_doorbell(db_val);
+  }
+  handle->threadImpl.sq_lock = 0;
+}
+
+__device__ void SingleThreadImpl::postLock(QueuePair *handle, int pe) {
+  handle->hdp_policy->hdp_flush();
+  // handle->waitCQSpace(1);
+  handle->waitSQSpace(1);
+}
+
+__device__ void MultiThreadImpl::postLock_internal(QueuePair *handle) {
+  int thread_id = get_flat_block_id();
+  int active_threads = wave_SZ();
+
+  if (thread_id % WF_SIZE == lowerID()) {
+    handle->hdp_policy->hdp_flush();
+    /*
+     * Don't let more than one wave in this WG go any further or a
+     * horrible variety of impossible to debug race conditions can occur.
+     */
+    while (atomicCAS(&(handle->threadImpl.sq_lock), 0, 1) == 1) {
+    }
+
+    /*
+     * This is a tiny bit over-aggressive as it assumes that all of the
+     * active_threads are going to the same PE when calculating whether
+     * we are full.
+     */
+    // handle->waitCQSpace(active_threads);
+    handle->waitSQSpace(active_threads);
+  }
+
+  /*
+   * Double check we've got the same exec mask (assuming divergence after
+   * the previous if.
+   */
+  if (active_threads != wave_SZ()) {
+    __builtin_trap();
+  }
+}
+
+__device__ void MultiThreadImpl::postLock(QueuePair *handle, int pe) {
+  /*
+   * For RC, we can't allow a wave to have different PEs in it, else the
+   * doorbell ringing logic will not work.  This little for loop forces
+   * control flow divergence based on the PE.  It works well for small
+   * numbers of PEs, but we might want a different solution for large
+   * numbers.
+   */
+  if (handle->connection_policy.forcePostDivergence()) {
+    for (int i = 0; i < handle->num_cqs; i++) {
+      if (i != pe) {
+        continue;
+      }
+
+      postLock_internal(handle);
+    }
+  } else {
+    postLock_internal(handle);
+  }
+}
+
+__device__ void WG::postLock(QueuePair *handle, int pe) {
+  handle->hdp_policy->hdp_flush();
+  // handle->waitCQSpace(1);
+  handle->waitSQSpace(1);
+}
+
+__device__ void WAVE::postLock(QueuePair *handle, int pe) {
+  handle->hdp_policy->hdp_flush();
+
+  /*
+   * Don't let more than one wave in this WG go any further or a horrible
+   * variety of impossible to debug race conditions can occur.
+   */
+  while (atomicCAS(&(handle->threadImpl.sq_lock), 0, 1) == 1) {
+  }
+
+  /*
+   * This is a tiny bit over-aggressive as it assumes that all of the
+   * active_threads are going to the same PE when calculating whether
+   * we are full.
+   */
+  // handle->waitCQSpace(1);
+  handle->waitSQSpace(1);
+}
+
+template <typename T>
+__device__ T SingleThreadImpl::threadAtomicAdd(T *val, T value) {
+  T old_val = *val;
+  *val += value;
+  return old_val;
+}
+
+template <typename T>
+__device__ T MultiThreadImpl::threadAtomicAdd(T *val, T value) {
+  return atomicAdd(val, value);
+}
+
+template <typename T>
+__device__ T WG::threadAtomicAdd(T *val, T value) {
+  T old_val = *val;
+  *val += value;
+  return old_val;
+}
+
+template <typename T>
+__device__ T WAVE::threadAtomicAdd(T *val, T value) {
+  return atomicAdd(val, value);
+}
+
+#define TYPE_GEN(T)                                                            \
+  template __device__ T SingleThreadImpl::threadAtomicAdd<T>(T * val,          \
+                                                             T value);         \
+  template __device__ T MultiThreadImpl::threadAtomicAdd<T>(T * val, T value); \
+  template __device__ T WG::threadAtomicAdd<T>(T * val, T value);              \
+  template __device__ T WAVE::threadAtomicAdd<T>(T * val, T value);
+
+TYPE_GEN(float)
+TYPE_GEN(double)
+TYPE_GEN(int)
+TYPE_GEN(unsigned int)
+TYPE_GEN(unsigned long long)  // NOLINT(runtime/int)
+
+#define TYPE_BOOL(T)                                                \
+  template __device__ void SingleThreadImpl::finishPost<T>(         \
+      QueuePair * handle, bool ring_db, int num_wqes, int pe,       \
+      uint16_t le_sq_counter, uint8_t opcode);                      \
+  template __device__ void MultiThreadImpl::finishPost<T>(          \
+      QueuePair * handle, bool ring_db, int num_wqes, int pe,       \
+      uint16_t le_sq_counter, uint8_t opcode);                      \
+  template __device__ void WG::finishPost<T>(                       \
+      QueuePair * handle, bool ring_db, int num_wqes, int pe,       \
+      uint16_t le_sq_counter, uint8_t opcode);                      \
+  template __device__ void WAVE::finishPost<T>(                     \
+      QueuePair * handle, bool ring_db, int num_wqes, int pe,       \
+      uint16_t le_sq_counter, uint8_t opcode);                      \
+  template __device__ void MultiThreadImpl::finishPost_internal<T>( \
+      QueuePair * handle, bool ring_db, int num_wqes, int pe,       \
+      uint16_t le_sq_counter, uint8_t opcode);
+
+TYPE_BOOL(true)
+TYPE_BOOL(false)
+
+}  // namespace rocshmem

src/gpu_ib/thread_policy.hpp

@@ -0,0 +1,176 @@
+/******************************************************************************
+ * 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_GPU_IB_THREAD_POLICY_HPP_
+#define LIBRARY_SRC_GPU_IB_THREAD_POLICY_HPP_
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "../util.hpp"
+
+namespace rocshmem {
+
+class QueuePair;
+
+/*
+ * GPU single-thread policy class. Only a single work-item per work-group
+ * is allowed to call into a ROC_SHMEM function (unless it is specifically
+ * called out as a collective API. This thread policy is the fastest but
+ * is not as flexible.
+ */
+class SingleThreadImpl {
+ public:
+  uint32_t cq_lock = 0;
+  uint32_t sq_lock = 0;
+
+  __device__ void quiet(QueuePair *handle);
+
+  __device__ void quiet_heavy(QueuePair *handle, int pe);
+
+  __device__ void decQuietCounter(uint32_t *quiet_counter, int num);
+
+  template <bool cqe>
+  __device__ void finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                             int pe, uint16_t le_sq_counter, uint8_t opcode);
+
+  __device__ void postLock(QueuePair *handle, int pe);
+
+  template <typename T>
+  __device__ T threadAtomicAdd(T *val, T value = 1);
+};
+
+/*
+ * GPU multi-thread policy class. Multiple work-items per work-group are
+ * allowed to call into a ROC_SHMEM function.  A bit slower than its
+ * single-thread counterpart but it enables a much more flexible user-facing
+ * API.
+ */
+class MultiThreadImpl {
+  /*
+   * Per-wg locks for the CQ and the SQ, respectively.
+   */
+  template <bool cqe>
+  __device__ void finishPost_internal(QueuePair *handle, bool ring_db,
+                                      int num_wqes, int pe,
+                                      uint16_t le_sq_counter, uint8_t opcode);
+
+  __device__ void postLock_internal(QueuePair *handle);
+
+ public:
+  uint32_t cq_lock = 0;
+  uint32_t sq_lock = 0;
+
+  __device__ void quiet(QueuePair *handle);
+
+  __device__ void quiet_heavy(QueuePair *handle, int pe);
+
+  __device__ void decQuietCounter(uint32_t *quiet_counter, int num);
+
+  template <bool cqe>
+  __device__ void finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                             int pe, uint16_t le_sq_counter, uint8_t opcode);
+
+  __device__ void postLock(QueuePair *handle, int pe);
+
+  template <typename T>
+  __device__ T threadAtomicAdd(T *val, T value = 1);
+};
+
+/*
+ * Select which one of our thread policies to use at compile time.
+ */
+#ifdef USE_THREADS
+typedef MultiThreadImpl ThreadImpl;
+#else
+typedef SingleThreadImpl ThreadImpl;
+#endif
+
+class THREAD {
+ public:
+  ThreadImpl threadImpl;
+
+  __device__ void quiet(QueuePair *handle) { threadImpl.quiet(handle); }
+
+  __device__ void quiet_heavy(QueuePair *handle, int pe) {
+    threadImpl.quiet_heavy(handle, pe);
+  }
+
+  __device__ void decQuietCounter(uint32_t *quiet_counter, int num) {
+    threadImpl.decQuietCounter(quiet_counter, num);
+  }
+
+  template <bool cqe>
+  __device__ void finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                             int pe, uint16_t le_sq_counter, uint8_t opcode) {
+    threadImpl.finishPost<cqe>(handle, ring_db, num_wqes, pe, le_sq_counter,
+                               opcode);
+  }
+
+  __device__ void postLock(QueuePair *handle, int pe) {
+    threadImpl.postLock(handle, pe);
+  }
+
+  template <typename T>
+  __device__ T threadAtomicAdd(T *val, T value = 1) {
+    T tmp = threadImpl.threadAtomicAdd(val, value);
+    return tmp;
+  }
+};
+
+class WAVE {
+ public:
+  __device__ void quiet(QueuePair *handle);
+
+  __device__ void quiet_heavy(QueuePair *handle, int pe);
+
+  __device__ void decQuietCounter(uint32_t *quiet_counter, int num);
+
+  template <bool cqe>
+  __device__ void finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                             int pe, uint16_t le_sq_counter, uint8_t opcode);
+
+  __device__ void postLock(QueuePair *handle, int pe);
+
+  template <typename T>
+  __device__ T threadAtomicAdd(T *val, T value = 1);
+};
+
+class WG {
+ public:
+  __device__ void quiet(QueuePair *handle);
+
+  __device__ void quiet_heavy(QueuePair *handle, int pe);
+
+  __device__ void decQuietCounter(uint32_t *quiet_counter, int num);
+
+  template <bool cqe>
+  __device__ void finishPost(QueuePair *handle, bool ring_db, int num_wqes,
+                             int pe, uint16_t le_sq_counter, uint8_t opcode);
+
+  __device__ void postLock(QueuePair *handle, int pe);
+
+  template <typename T>
+  __device__ T threadAtomicAdd(T *val, T value = 1);
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_GPU_IB_THREAD_POLICY_HPP_

src/hdp_policy.hpp

@@ -0,0 +1,197 @@
+/******************************************************************************
+ * 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_HDP_POLICY_HPP_
+#define LIBRARY_SRC_HDP_POLICY_HPP_
+
+#include <hip/hip_runtime.h>
+#include <hsa/hsa_ext_amd.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "memory/hip_allocator.hpp"
+#include "util.hpp"
+
+namespace rocshmem {
+
+class HdpHostSideFlushRocmPolicy {
+ public:
+  HdpHostSideFlushRocmPolicy() { set_hdp_flush_ptr(); }
+  ~HdpHostSideFlushRocmPolicy() {}
+
+  /**
+   * @brief Host-side API to set the flush polling ptr
+   * @param hdp_gpu_cpu_flush_flag A pointer allocated on the symmetric heap
+   */
+  __host__ void set_flush_polling_ptr(unsigned int* hdp_gpu_cpu_flush_flag) {
+    hdp_gpu_cpu_flush_flag_ = hdp_gpu_cpu_flush_flag;
+    if (hdp_gpu_cpu_flush_flag_ != nullptr) {
+      *hdp_gpu_cpu_flush_flag_ =
+          static_cast<std::underlying_type_t<hdp_poll_flag>>(
+              hdp_poll_flag::NO_FLUSH);
+    }
+  }
+
+  /**
+   * @brief Host-side API that performs the flush.
+   */
+  __host__ void hdp_flush() { *hdp_flush_ptr_ = HDP_FLUSH_VAL; }
+
+  /**
+   * @brief Device-side API that signal to the CPU spinning thread to do the
+   * flush on behalf of the GPU.
+   *
+   * @param hdp_ptr A pointer to use as the HDP flush ptr
+   */
+  __device__ static void hdp_flush(unsigned int* hdp_ptr) {
+    STORE(hdp_ptr, static_cast<unsigned int>(hdp_poll_flag::FLUSH));
+  }
+
+  /**
+   * @brief Device-side API that signal to the CPU spinning thread to do the
+   * flush on behalf of the GPU.
+   */
+  __device__ void hdp_flush() { hdp_flush(hdp_gpu_cpu_flush_flag_); }
+
+  __device__ void flushCoherency() { hdp_flush(); }
+
+  /**
+   * @brief Get the hdp flush signal flag.
+   *
+   * @return A pointer to an object that stores a signal for the CPU thread to
+   * perform the flush on the host.
+   */
+  __host__ unsigned int* get_hdp_flush_ptr() const {
+    return hdp_gpu_cpu_flush_flag_;
+  }
+
+  /**
+   * @brief Check if there is an active request to flush the hdp cache.
+   *
+   * @return A boolean indicating whether there is a flush request.
+   */
+  __host__ bool has_active_flush_request() const {
+    if (hdp_gpu_cpu_flush_flag_ == nullptr) {
+      return false;
+    }
+    auto device_flag_value = __hip_atomic_load(
+        hdp_gpu_cpu_flush_flag_, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM);
+    const auto flush_value = static_cast<std::underlying_type_t<hdp_poll_flag>>(
+        hdp_poll_flag::FLUSH);
+    return device_flag_value == flush_value;
+  }
+
+  /**
+   * @brief Clears the flag that signals hdp cache flush.
+   */
+  __host__ void clear_active_flush_flag() {
+    const auto no_flush_value =
+        static_cast<std::underlying_type_t<hdp_poll_flag>>(
+            hdp_poll_flag::NO_FLUSH);
+    __hip_atomic_store(hdp_gpu_cpu_flush_flag_, no_flush_value,
+                       __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM);
+  }
+
+  static const int HDP_FLUSH_VAL{0x01};
+
+  enum class hdp_poll_flag : unsigned int {
+    NO_FLUSH = 0x0,
+    FLUSH = 0x1,
+  };
+
+ private:
+  void set_hdp_flush_ptr() {
+    int hip_dev_id{};
+    CHECK_HIP(hipGetDevice(&hip_dev_id));
+    CHECK_HIP(hipDeviceGetAttribute(reinterpret_cast<int*>(&hdp_flush_ptr_),
+                                    hipDeviceAttributeHdpMemFlushCntl,
+                                    hip_dev_id));
+  }
+
+  unsigned int* hdp_flush_ptr_{nullptr};
+  unsigned int* hdp_gpu_cpu_flush_flag_{nullptr};
+};
+
+class HdpDeviceSideFlushRocmPolicy {
+ public:
+  HdpDeviceSideFlushRocmPolicy() { set_hdp_flush_ptr(); }
+  ~HdpDeviceSideFlushRocmPolicy() {}
+
+  /**
+   * @brief Flush the HDP by setting the flush control signal
+   */
+  __host__ __device__ void hdp_flush() {
+    STORE(hdp_flush_ptr_, static_cast<unsigned int>(HDP_FLUSH_VAL));
+  }
+  __device__ void flushCoherency() { hdp_flush(); }
+
+  /**
+   * @brief Get the hdp flush signal flag.
+   *
+   * @return A pointer to the HDP flush control signal
+   */
+  __host__ unsigned int* get_hdp_flush_ptr() const { return hdp_flush_ptr_; }
+
+  static const int HDP_FLUSH_VAL{0x01};
+
+ private:
+  void set_hdp_flush_ptr() {
+    int hip_dev_id{};
+    CHECK_HIP(hipGetDevice(&hip_dev_id));
+    CHECK_HIP(hipDeviceGetAttribute(reinterpret_cast<int*>(&hdp_flush_ptr_),
+                                    hipDeviceAttributeHdpMemFlushCntl,
+                                    hip_dev_id));
+  }
+
+  unsigned int* hdp_flush_ptr_{nullptr};
+};
+
+class NoHdpPolicy {
+ public:
+  NoHdpPolicy() = default;
+
+  __host__ void hdp_flush() {}
+
+  __host__ unsigned int* get_hdp_flush_ptr() const { return nullptr; }
+
+  __device__ void hdp_flush() {}
+
+  __device__ void flushCoherency() { __roc_flush(); }
+};
+
+/*
+ * Select which one of our HDP policies to use at compile time.
+ */
+#ifdef USE_COHERENT_HEAP
+typedef NoHdpPolicy HdpPolicy;
+#else
+// Only when we are using the IB conduit, we have to use a polling thread to
+// flush the HDP cache on the GPU's behalf.
+#ifdef USE_HOST_SIDE_HDP_FLUSH
+typedef HdpHostSideFlushRocmPolicy HdpPolicy;
+#else
+typedef HdpDeviceSideFlushRocmPolicy HdpPolicy;
+#endif
+#endif
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_HDP_POLICY_HPP_

src/hdp_proxy.hpp

@@ -0,0 +1,65 @@
+/******************************************************************************
+ * 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_HDP_PROXY_HPP_
+#define LIBRARY_SRC_HDP_PROXY_HPP_
+
+#include "device_proxy.hpp"
+#include "hdp_policy.hpp"
+
+namespace rocshmem {
+
+template <typename ALLOCATOR>
+class HdpProxy {
+  using HdpProxyT = DeviceProxy<ALLOCATOR, HdpPolicy>;
+
+ public:
+  /*
+   * Placement new the memory which is allocated by proxy_
+   */
+  HdpProxy() {
+    new (proxy_.get()) HdpPolicy();
+  }
+
+  /*
+   * Since placement new is called in the constructor, then
+   * delete must be called manually.
+   */
+  ~HdpProxy() {
+    proxy_.get()->~HdpPolicy();
+  }
+
+  /*
+   * @brief Provide access to the memory referenced by the proxy
+   */
+  __host__ __device__ HdpPolicy* get() { return proxy_.get(); }
+
+ private:
+  /*
+   * @brief Memory managed by the lifetime of this object
+   */
+  HdpProxyT proxy_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_HDP_PROXY_HPP_

src/host/CMakeLists.txt

@@ -0,0 +1,30 @@
+###############################################################################
+# 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.
+###############################################################################
+
+###############################################################################
+# ADD ROCSHMEM TARGET FOR FILES IN CURRENT DIRECTORY
+###############################################################################
+target_sources(
+  ${PROJECT_NAME}
+  PRIVATE
+    host.cpp
+)

src/host/host.cpp

@@ -0,0 +1,255 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "host.hpp"
+
+#include <mpi.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "host_helpers.hpp"
+#include "../memory/window_info.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+__host__ HostContextWindowInfo::HostContextWindowInfo(MPI_Comm comm_world,
+                                                      SymmetricHeap* heap) {
+  window_info_ =
+      new WindowInfo(comm_world, heap->get_local_heap_base(), heap->get_size());
+}
+
+__host__ HostContextWindowInfo::~HostContextWindowInfo() {
+  delete window_info_;
+}
+
+WindowInfo* HostInterface::acquire_window_context() {
+  auto index{find_avail_pool_entry()};
+
+  HostContextWindowInfo* acquired_win_info = host_window_context_pool_[index];
+
+  acquired_win_info->mark_unavail();
+
+  return acquired_win_info->get();
+}
+
+__host__ void HostInterface::release_window_context(WindowInfo* window_info) {
+  auto index{find_win_info_in_pool(window_info)};
+
+  host_window_context_pool_[index]->mark_avail();
+}
+
+int HostInterface::find_avail_pool_entry() {
+  for (int i{0}; i < max_num_ctxs_; i++) {
+    if (host_window_context_pool_[i]->is_avail()) {
+      return i;
+    }
+  }
+  /* Entry should have been available; consider this as an error. */
+  assert(false);
+  return -1;
+}
+
+int HostInterface::find_win_info_in_pool(WindowInfo* window_info) {
+  for (int i{0}; i < max_num_ctxs_; i++) {
+    if (host_window_context_pool_[i]->is_avail()) {
+      continue;
+    }
+    if (window_info == host_window_context_pool_[i]->get()) {
+      return i;
+    }
+  }
+  /* Entry should have been present; consider this as an error. */
+  assert(false);
+  return -1;
+}
+
+__host__ HostInterface::HostInterface(HdpPolicy* hdp_policy,
+                                      MPI_Comm roc_shmem_comm,
+                                      SymmetricHeap* heap) {
+  /*
+   * Duplicate a communicator from roc_shem's comm
+   * world for the host interface
+   */
+  MPI_Comm_dup(roc_shmem_comm, &host_comm_world_);
+  MPI_Comm_rank(host_comm_world_, &my_pe_);
+  MPI_Comm_rank(host_comm_world_, &num_pes_);
+
+  /*
+   * Create an MPI window on the HDP so that it can be flushed
+   * by remote PEs for host-facing functions
+   */
+  hdp_policy_ = hdp_policy;
+
+  /*
+   * Allocate and initialize pool of windows for contexts
+   */
+  char* value{nullptr};
+  if ((value = getenv("ROC_SHMEM_MAX_NUM_HOST_CONTEXTS"))) {
+    max_num_ctxs_ = atoi(value);
+  }
+
+  size_t pool_size = max_num_ctxs_ * sizeof(HostContextWindowInfo*);
+  host_window_context_pool_ =
+      reinterpret_cast<HostContextWindowInfo**>(malloc(pool_size));
+
+  for (int ctx_i = 0; ctx_i < max_num_ctxs_; ctx_i++) {
+    host_window_context_pool_[ctx_i] =
+        new HostContextWindowInfo(host_comm_world_, heap);
+  }
+
+#if !defined(USE_COHERENT_HEAP) && !defined(USE_SINGLE_NODE)
+  // The single node implementation needs a different path since
+  // the HDP flush pointers are allocated on the symmetric heap
+  // and we need to wait for other initialization to happen before
+  // calling `get_hdp_flush_ptr`.
+  create_hdp_window();
+#endif  // defined(USE_COHERENT_HEAP) && !defined(USE_SINGLE_NODE)
+}
+
+#ifndef USE_COHERENT_HEAP
+__host__ void HostInterface::create_hdp_window() {
+  MPI_Win_create(hdp_policy_->get_hdp_flush_ptr(),
+                 sizeof(unsigned int), /* size of window */
+                 sizeof(unsigned int), /* displacement */
+                 MPI_INFO_NULL, host_comm_world_, &hdp_win);
+
+  /*
+   * Start a shared access epoch on windows of all ranks,
+   * and let the library there is no need to check for
+   * lock exclusivity during operations on this window
+   * (MPI_MODE_NOCHECK).
+   */
+  MPI_Win_lock_all(MPI_MODE_NOCHECK, hdp_win);
+}
+#endif  // USE_COHERENT_HEAP
+
+__host__ HostInterface::~HostInterface() {
+#ifndef USE_COHERENT_HEAP
+  MPI_Win_unlock_all(hdp_win);
+
+  MPI_Win_free(&hdp_win);
+#endif  // USE_COHERENT_HEAP
+
+  /* Detroy the pool of contexts */
+  for (int ctx_i = 0; ctx_i < max_num_ctxs_; ctx_i++) {
+    delete host_window_context_pool_[ctx_i];
+  }
+
+  free(host_window_context_pool_);
+
+  MPI_Comm_free(&host_comm_world_);
+}
+
+__host__ void HostInterface::putmem_nbi(void* dest, const void* source,
+                                        size_t nelems, int pe,
+                                        WindowInfo* window_info) {
+  initiate_put(dest, source, nelems, pe, window_info);
+}
+
+__host__ void HostInterface::getmem_nbi(void* dest, const void* source,
+                                        size_t nelems, int pe,
+                                        WindowInfo* window_info) {
+  initiate_get(dest, source, nelems, pe, window_info);
+}
+
+__host__ void HostInterface::putmem(void* dest, const void* source,
+                                    size_t nelems, int pe,
+                                    WindowInfo* window_info) {
+  initiate_put(dest, source, nelems, pe, window_info);
+
+  MPI_Win_flush_local(pe, window_info->get_win());
+}
+
+__host__ void HostInterface::getmem(void* dest, const void* source,
+                                    size_t nelems, int pe,
+                                    WindowInfo* window_info) {
+  initiate_get(dest, source, nelems, pe, window_info);
+
+  MPI_Win_flush_local(pe, window_info->get_win());
+
+  /*
+   * Flush local HDP to ensure that the NIC's write
+   * of the fetched data is visible in device memory
+   */
+  hdp_policy_->hdp_flush();
+}
+
+__host__ void HostInterface::fence(WindowInfo* window_info) {
+  complete_all(window_info->get_win());
+
+  /*
+   * Flush my HDP and the HDPs of remote GPUs.
+   * The HDP is a write-combining (WC) write-through
+   * cache. But, even after the WC buffer is full and
+   * the data is passed to the Data Fabric (DF), DF
+   * can still reorder the writes. A flush ensures
+   * that writes after the flush are written only
+   * after those before the flush.
+   */
+  hdp_policy_->hdp_flush();
+  flush_remote_hdps();
+
+  return;
+}
+
+__host__ void HostInterface::quiet(WindowInfo* window_info) {
+  complete_all(window_info->get_win());
+
+  /* Same explanation as in fence */
+  hdp_policy_->hdp_flush();
+  flush_remote_hdps();
+
+  return;
+}
+
+__host__ void HostInterface::sync_all(WindowInfo* window_info) {
+  MPI_Win_sync(window_info->get_win());
+
+  hdp_policy_->hdp_flush();
+  /*
+   * No need to flush remote
+   * HDPs here since all PEs are
+   * participating.
+   */
+
+  MPI_Barrier(host_comm_world_);
+
+  return;
+}
+
+__host__ void HostInterface::barrier_all(WindowInfo* window_info) {
+  complete_all(window_info->get_win());
+
+  /*
+   * Flush my HDP cache so remote NICs will
+   * see the latest values in device memory
+   */
+  hdp_policy_->hdp_flush();
+
+  MPI_Barrier(host_comm_world_);
+}
+
+__host__ void HostInterface::barrier_for_sync() {
+  MPI_Barrier(host_comm_world_);
+}
+
+}  // namespace rocshmem

src/host/host.hpp

@@ -0,0 +1,406 @@
+/******************************************************************************
+ * 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_HOST_HOST_HPP_
+#define LIBRARY_SRC_HOST_HOST_HPP_
+
+/**
+ * @file host.hpp
+ * Defines the HostInterface class.
+ *
+ * The file contains the HostInterface class that defines all the
+ * host-facing functions that will be used by all host contexts of
+ * any backend type.
+ */
+
+#include <mpi.h>
+
+#include <map>
+
+#include "roc_shmem/roc_shmem.hpp"
+#include "../hdp_policy.hpp"
+#include "../memory/symmetric_heap.hpp"
+#include "../memory/window_info.hpp"
+
+namespace rocshmem {
+
+class HostContextWindowInfo {
+ public:
+  /**
+   * @brief Constructor with default members
+   */
+  HostContextWindowInfo() = default;
+
+  /**
+   * @brief Constructor with initialized members
+   *
+   * @param[in] team pointer used to track team info
+   * @param[in] team_info information about participating PEs
+   */
+  HostContextWindowInfo(MPI_Comm comm_world, SymmetricHeap* heap);
+
+  /**
+   * @brief Destructor
+   */
+  __host__ ~HostContextWindowInfo();
+
+  /**
+   * @brief Retrieve a pointer to the internal WindowInfo
+   *
+   * @return WindowInfo pointer
+   */
+  WindowInfo* get() { return window_info_; }
+
+  /**
+   * @brief Mark the window info as avaialable (not allocated)
+   */
+  void mark_avail() { avail_ = true; }
+
+  /**
+   * @brief Mark the window info as unavaialble (allocated)
+   */
+  void mark_unavail() { avail_ = false; }
+
+  /**
+   * @brief Check if the window info has been allocated
+   */
+  bool is_avail() { return avail_; }
+
+ private:
+  /**
+   * @brief Flag to state whether or not this window is available to be assigned
+   * to a Host Context
+   */
+  bool avail_{true};
+
+  /**
+   * @brief Pointer to the WindowInfo object that manages the MPI Window for
+   * this context
+   */
+  WindowInfo* window_info_{nullptr};
+};
+
+class HostInterface {
+ public:
+  /**
+   * @brief Primary constructor
+   */
+  __host__ HostInterface(HdpPolicy* hdp_policy, MPI_Comm roc_shmem_comm,
+                         SymmetricHeap* heap);
+
+  /**
+   * @brief Destructor
+   */
+  __host__ ~HostInterface();
+
+  /**
+   * @brief Accessor for copy of comm world
+   *
+   * @return MPI_Comm containing host comm world
+   */
+  MPI_Comm get_comm_world() { return host_comm_world_; }
+
+  /**
+   * @brief Get a window context from the pool
+   *
+   * @return Pointer to the WindowInfo in the allocated one from the pool
+   */
+  WindowInfo* acquire_window_context();
+
+  /**
+   * @brief Return a window context back to the pool
+   */
+  void release_window_context(WindowInfo* window_info);
+
+  /**************************************************************************
+   ***************************** HOST FUNCTIONS *****************************
+   *************************************************************************/
+  template <typename T>
+  __host__ void p(T* dest, T value, int pe, WindowInfo* window_info);
+
+  template <typename T>
+  __host__ T g(const T* source, int pe, WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void put(T* dest, const T* source, size_t nelems, int pe,
+                    WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void get(T* dest, const T* source, size_t nelems, int pe,
+                    WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void put_nbi(T* dest, const T* source, size_t nelems, int pe,
+                        WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void get_nbi(T* dest, const T* source, size_t nelems, int pe,
+                        WindowInfo* window_info);
+
+  __host__ void putmem(void* dest, const void* source, size_t nelems, int pe,
+                       WindowInfo* window_info);
+
+  __host__ void getmem(void* dest, const void* source, size_t nelems, int pe,
+                       WindowInfo* window_info);
+
+  __host__ void putmem_nbi(void* dest, const void* source, size_t nelems,
+                           int pe, WindowInfo* window_info);
+
+  __host__ void getmem_nbi(void* dest, const void* source, size_t size, int pe,
+                           WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void amo_add(void* dst, T value, int pe, WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void amo_cas(void* dst, T value, T cond, int pe,
+                        WindowInfo* window_info);
+
+  template <typename T>
+  __host__ T amo_fetch_add(void* dst, T value, int pe, WindowInfo* window_info);
+
+  template <typename T>
+  __host__ T amo_fetch_cas(void* dst, T value, T cond, int pe,
+                           WindowInfo* window_info);
+
+  __host__ void fence(WindowInfo* window_info);
+
+  __host__ void quiet(WindowInfo* window_info);
+
+  __host__ void barrier_all(WindowInfo* window_info);
+
+  __host__ void barrier_for_sync();
+
+  __host__ void sync_all(WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void broadcast(T* dest, const T* source, int nelems, int pe_root,
+                          int pe_start, int log_pe_stride, int pe_size,
+                          long* p_sync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __host__ void broadcast(roc_shmem_team_t team, T* dest, const T* source,
+                          int nelems, int pe_root);
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all(T* dest, const T* source, int nreduce, int pe_start,
+                       int log_pe_stride, int pe_size, T* p_wrk,
+                       long* p_sync);  // NOLINT(runtime/int)
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all(roc_shmem_team_t team, T* dest, const T* source,
+                       int nreduce);
+
+  template <typename T>
+  __host__ void wait_until(T* ptr, roc_shmem_cmps cmp, T val,
+                           WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void wait_until_all(T* ptr, size_t nelems, const int* status,
+                               roc_shmem_cmps cmp, T val,
+                               WindowInfo* window_info);
+
+  template <typename T>
+  __host__ size_t wait_until_any(T* ptr, size_t nelems, const int* status,
+                                 roc_shmem_cmps cmp, T val,
+                                 WindowInfo* window_info);
+
+  template <typename T>
+  __host__ size_t wait_until_some(T* ptr, size_t nelems, size_t* indices,
+                                  const int* status, roc_shmem_cmps cmp, T val,
+                                  WindowInfo* window_info);
+
+  template <typename T>
+  __host__ void wait_until_all_vector(T* ptr, size_t nelems, const int* status,
+                                      roc_shmem_cmps cmp, T* vals,
+                                      WindowInfo* window_info);
+
+  template <typename T>
+  __host__ size_t wait_until_any_vector(T* ptr, size_t nelems,
+                                        const int* status, roc_shmem_cmps cmp,
+                                        T* vals, WindowInfo* window_info);
+
+  template <typename T>
+  __host__ size_t wait_until_some_vector(T* ptr, size_t nelems, size_t* indices,
+                                         const int* status, roc_shmem_cmps cmp,
+                                         T* vals, WindowInfo* window_info);
+
+  template <typename T>
+  __host__ int test(T* ptr, roc_shmem_cmps cmp, T val, WindowInfo* window_info);
+
+#ifndef USE_COHERENT_HEAP
+  __host__ void create_hdp_window();
+#endif // USE_COHERENT_HEAP
+
+ private:
+  /**************************************************************************
+   **************************** INTERNAL METHODS ****************************
+   *************************************************************************/
+  __host__ void flush_remote_hdps() {
+#ifndef USE_COHERENT_HEAP
+    unsigned flush_val{HdpPolicy::HDP_FLUSH_VAL};
+    for (size_t i{0}; i < num_pes_; i++) {
+      if (i == my_pe_) {
+        continue;
+      }
+      MPI_Put(&flush_val, 1, MPI_UNSIGNED, i, 0, 1, MPI_UNSIGNED, hdp_win);
+    }
+    MPI_Win_flush_all(hdp_win);
+#endif // USE_COHERENT_HEAP
+  }
+
+  __host__ void flush_remote_hdp(int pe) {
+#ifndef USE_COHERENT_HEAP
+    unsigned flush_val{HdpPolicy::HDP_FLUSH_VAL};
+    MPI_Put(&flush_val, 1, MPI_UNSIGNED, pe, 0, 1, MPI_UNSIGNED, hdp_win);
+    MPI_Win_flush(pe, hdp_win);
+#endif // USE_COHERENT_HEAP
+  }
+
+  __host__ void initiate_put(void* dest, const void* source, size_t nelems,
+                             int pe, WindowInfo* window_info);
+
+  __host__ void initiate_get(void* dest, const void* source, size_t nelems,
+                             int pe, WindowInfo* window_info);
+
+  __host__ void complete_all(MPI_Win win);
+
+  __host__ MPI_Aint compute_offset(const void* dest, void* win_start,
+                                   void* win_end);
+
+  __host__ MPI_Comm get_mpi_comm(int pe_start, int log_pe_stride, int pe_size);
+
+  __host__ MPI_Op get_mpi_op(ROC_SHMEM_OP Op);
+
+  template <typename T>
+  __host__ MPI_Datatype get_mpi_type();
+
+  template <typename T>
+  __host__ int compare(roc_shmem_cmps cmp, T input_val, T target_val);
+
+  template <typename T>
+  __host__ int test_and_compare(MPI_Aint offset, MPI_Datatype mpi_type,
+                                roc_shmem_cmps cmp, T val, MPI_Win win);
+
+  template <typename T, ROC_SHMEM_OP Op>
+  __host__ void to_all_internal(MPI_Comm mpi_comm, T* dest, const T* source,
+                                int nreduce);
+
+  template <typename T>
+  __host__ void broadcast_internal(MPI_Comm mpi_comm, T* dest, const T* source,
+                                   int nelems, int pe_root);
+
+  /**************************************************************************
+   **************************** INTERNAL MEMBERS ****************************
+   *************************************************************************/
+  /**
+   * @brief Duplicate to the Backend's hdp policy pointer
+   */
+  HdpPolicy* hdp_policy_{nullptr};
+
+  /**
+   * @brief Global MPI communicator for those host API
+   */
+  MPI_Comm host_comm_world_{};
+
+  /**
+   * @brief Duplicate of this processing element's id within global rank
+   */
+  int my_pe_{-1};
+
+  /**
+   * @brief Duplicate of global number of processing elements
+   */
+  int num_pes_{0};
+
+#ifndef USE_COHERENT_HEAP
+  /**
+   * @brief MPI window for hdp flushing
+   */
+  MPI_Win hdp_win;
+#endif  // USE_COHERENT_HEAP
+
+  /**
+   * @brief Max number of contexts for the application
+   */
+  int max_num_ctxs_{40};
+
+  /**
+   * @brief Pool of HostContexWindowInfos
+   */
+  HostContextWindowInfo** host_window_context_pool_{nullptr};
+
+  int find_win_info_in_pool(WindowInfo* window_info);
+
+  int find_avail_pool_entry();
+
+  /*
+   * @brief Used by comm_map map for active sets.
+   *
+   * This data structure that stores the parameters defining the active
+   * set of PEs in a collective. This struct also serves as a key
+   * into the comm_map map.
+   */
+  class ActiveSetKey {
+   public:
+    /**
+     * @brief Primary constructor
+     */
+    ActiveSetKey(int pe_start, int log_pe_stride, int pe_size)
+        : pe_start_(pe_start),
+          log_pe_stride_(log_pe_stride),
+          pe_size_(pe_size) {}
+
+    bool operator<(const ActiveSetKey& key) const {
+      return pe_start_ < key.pe_start_ ||
+             (pe_start_ == key.pe_start_ &&
+              log_pe_stride_ < key.log_pe_stride_) ||
+             (pe_start_ == key.pe_start_ &&
+              log_pe_stride_ == key.log_pe_stride_ && pe_size_ < key.pe_size_);
+    }
+
+   private:
+    /**
+     * @brief Records start location in (logical) active set bitmask
+     */
+    int pe_start_{-1};
+
+    /**
+     * @brief Records stride in (logical) active set bitmask
+     */
+    int log_pe_stride_{-1};
+
+    /**
+     * @brief Records (logical) active set bitmask size
+     */
+    int pe_size_{-1};
+  };
+
+  /*
+   * @brief Map of active set descriptors to MPI communicators
+   */
+  std::map<ActiveSetKey, MPI_Comm> comm_map{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_HOST_HOST_HPP_

src/host/host_helpers.hpp

@@ -0,0 +1,92 @@
+/******************************************************************************
+ * 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_HOST_HOST_HELPERS_HPP_
+#define LIBRARY_SRC_HOST_HOST_HELPERS_HPP_
+
+#include "host.hpp"
+#include "../memory/window_info.hpp"
+
+namespace rocshmem {
+
+__host__ inline MPI_Aint HostInterface::compute_offset(
+    const void* dest, void* win_start, [[maybe_unused]] void* win_end) {
+  assert((reinterpret_cast<char*>(const_cast<void*>(dest)) >=
+          reinterpret_cast<char*>(win_start)) &&
+         (reinterpret_cast<char*>(const_cast<void*>(dest)) <
+          reinterpret_cast<char*>(win_end)));
+
+  MPI_Aint dest_disp{};
+  MPI_Aint start_disp{};
+
+  MPI_Get_address(dest, &dest_disp);
+  MPI_Get_address(win_start, &start_disp);
+
+  return MPI_Aint_diff(dest_disp, start_disp);
+}
+
+__host__ inline void HostInterface::complete_all(MPI_Win win) {
+  MPI_Win_flush_all(win); /* RMA operations */
+  MPI_Win_sync(win);      /* memory stores */
+}
+
+__host__ inline void HostInterface::initiate_put(void* dest, const void* source,
+                                                 size_t nelems, int pe,
+                                                 WindowInfo* window_info) {
+  MPI_Win win{window_info->get_win()};
+  void* win_start{window_info->get_start()};
+  void* win_end{window_info->get_end()};
+
+  /* Calculate offset of remote dest from base address of window */
+  MPI_Aint offset{compute_offset(dest, win_start, win_end)};
+
+  /*
+   * Current semantics of our API restrict the buffers
+   * passed in to be on the symmetric heap only. So,
+   * flush the HDP since the GPU may have written the
+   * latest value to the source buffer and we want the
+   * NIC to DMA read the latest value instead of the
+   * value that may have been cached in the HDP.
+   */
+  hdp_policy_->hdp_flush();
+
+  /* Offload remote write operation to MPI */
+  MPI_Put(source, nelems, MPI_CHAR, pe, offset, nelems, MPI_CHAR, win);
+}
+
+__host__ inline void HostInterface::initiate_get(void* dest, const void* source,
+                                                 size_t nelems, int pe,
+                                                 WindowInfo* window_info) {
+  MPI_Win win{window_info->get_win()};
+  void* win_start{window_info->get_start()};
+  void* win_end{window_info->get_end()};
+
+  /* Calculate offset of remote source from base address of window */
+  MPI_Aint offset = compute_offset(source, win_start, win_end);
+
+  /* Offload remote fetch operation to MPI */
+  MPI_Get(dest, nelems, MPI_CHAR, pe, offset, nelems, MPI_CHAR, win);
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_HOST_HOST_HELPERS_HPP_

src/host/host_templates.hpp

@@ -0,0 +1,645 @@
+/******************************************************************************
+ * 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_HOST_HOST_TEMPLATES_HPP_
+#define LIBRARY_SRC_HOST_HOST_TEMPLATES_HPP_
+
+#include <utility>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "host_helpers.hpp"
+#include "../memory/window_info.hpp"
+#include "../team.hpp"
+
+namespace rocshmem {
+
+template <typename T>
+__host__ void HostInterface::p(T* dest, T value, int pe,
+                               WindowInfo* window_info) {
+  DPRINTF("Function: host_p\n");
+  putmem(dest, &value, sizeof(T), pe, window_info);
+}
+
+template <typename T>
+__host__ void HostInterface::put(T* dest, const T* source, size_t nelems,
+                                 int pe, WindowInfo* window_info) {
+  DPRINTF("Function: host_put\n");
+  putmem(dest, source, sizeof(T) * nelems, pe, window_info);
+}
+
+template <typename T>
+__host__ void HostInterface::put_nbi(T* dest, const T* source, size_t nelems,
+                                     int pe, WindowInfo* window_info) {
+  DPRINTF("Function: host_put_nbi\n");
+  putmem_nbi(dest, source, sizeof(T) * nelems, pe, window_info);
+}
+
+template <typename T>
+__host__ T HostInterface::g(const T* source, int pe, WindowInfo* window_info) {
+  DPRINTF("Function: host_g\n");
+
+  T ret{};
+
+  /*
+   * We don't call getmem directly here
+   * since it flushes the local HDP. We
+   * don't need the flush because the
+   * destination buffer is on the CPU.
+   */
+  getmem_nbi(&ret, source, sizeof(T), pe, window_info);
+
+  MPI_Win_flush_local(pe, window_info->get_win());
+
+  return ret;
+}
+
+template <typename T>
+__host__ void HostInterface::get(T* dest, const T* source, size_t nelems,
+                                 int pe, WindowInfo* window_info) {
+  DPRINTF("Function: host_get\n");
+  getmem(dest, source, sizeof(T) * nelems, pe, window_info);
+}
+
+template <typename T>
+__host__ void HostInterface::get_nbi(T* dest, const T* source, size_t nelems,
+                                     int pe, WindowInfo* window_info) {
+  DPRINTF("Function: host_get_nbi\n");
+  getmem_nbi(dest, source, sizeof(T) * nelems, pe, window_info);
+}
+
+__host__ MPI_Comm HostInterface::get_mpi_comm(int pe_start, int log_pe_stride,
+                                              int pe_size) {
+  MPI_Comm active_set_comm{};
+
+  /*
+   * First, check to see if the active set is the same as COMM_WORLD
+   */
+  int comm_world_size{-1};
+  MPI_Comm_size(host_comm_world_, &comm_world_size);
+
+  if (pe_start == 0 && log_pe_stride == 0 && pe_size == comm_world_size) {
+    /*
+     * Use the host interface's copy of MPI_COMM_WORLD
+     * TODO: replace with a per-context copy of MPI_COMM_WORLD when we
+     * have multiple contexts
+     */
+    active_set_comm = host_comm_world_;
+    return active_set_comm;
+  }
+
+  /*
+   * Then, check to see if we had already created a communicator for
+   * this active set
+   */
+  ActiveSetKey key(pe_start, log_pe_stride, pe_size);
+
+  auto it{comm_map.find(key)};
+  if (it != comm_map.end()) {
+    DPRINTF("Using cached communicator\n");
+    return it->second;
+  }
+
+  /*
+   * If there is not one cached, create a new one (expensive)
+   */
+  int active_set_ranks[pe_size];  // NOLINT
+  int stride{1 << log_pe_stride};
+  active_set_ranks[0] = pe_start;
+
+  for (int i{1}; i < pe_size; i++) {
+    active_set_ranks[i] = active_set_ranks[i - 1] + stride;
+  }
+
+  MPI_Group comm_world_group{};
+  MPI_Group active_set_group{};
+
+  MPI_Comm_group(host_comm_world_, &comm_world_group);
+
+  MPI_Group_incl(comm_world_group, pe_size, active_set_ranks,
+                 &active_set_group);
+
+  MPI_Comm_create_group(host_comm_world_, active_set_group, 0,
+                        &active_set_comm);
+
+  /*
+   * Cache the new communicator
+   */
+  DPRINTF("Created a new communicator. Now caching it\n");
+  comm_map.insert(std::pair<ActiveSetKey, MPI_Comm>(key, active_set_comm));
+
+  return active_set_comm;
+}
+
+template <typename T>
+__host__ void HostInterface::broadcast_internal(MPI_Comm mpi_comm, T* dest,
+                                                const T* source, int nelems,
+                                                int pe_root) {
+  DPRINTF("Function: host_broadcast_internal\n");
+
+  /*
+   * Choose the right pointer for my buffer depending
+   * on whether or not I am the root.
+   */
+  int active_set_rank{-1};
+  void* buffer{nullptr};
+  MPI_Comm_rank(mpi_comm, &active_set_rank);
+  if (pe_root == active_set_rank) {
+    buffer = const_cast<T*>(source);
+  } else {
+    buffer = const_cast<T*>(dest);
+  }
+
+  /*
+   * Flush my HDP so that the NIC does not read stale values
+   */
+  hdp_policy_->hdp_flush();
+
+  /*
+   * Offload the broadcast to MPI
+   */
+  MPI_Bcast(buffer, nelems * sizeof(T), MPI_CHAR, pe_root, mpi_comm);
+
+  return;
+}
+
+template <typename T>
+__host__ void HostInterface::broadcast(T* dest, const T* source, int nelems,
+                                       int pe_root, int pe_start,
+                                       int log_pe_stride, int pe_size,
+                                       [[maybe_unused]] long* p_sync) {
+  DPRINTF("Function: host_broadcast\n");
+
+  /*
+   * Get an MPI communicator for active set of PEs
+   * Note: pe_root is w.r.t the active set, hence
+   * the MPI communicator contains the root as well.
+   */
+  MPI_Comm mpi_comm{get_mpi_comm(pe_start, log_pe_stride, pe_size)};
+
+  broadcast_internal<T>(mpi_comm, dest, source, nelems, pe_root);
+
+  return;
+}
+
+template <typename T>
+__host__ void HostInterface::broadcast(roc_shmem_team_t team, T* dest,
+                                       const T* source, int nelems,
+                                       int pe_root) {
+  DPRINTF("Function: Team-based host_broadcast\n");
+
+  /*
+   * Get the MPI communicator of this team
+   */
+  Team* team_obj{get_internal_team(team)};
+  MPI_Comm mpi_comm{team_obj->mpi_comm};
+
+  broadcast_internal<T>(mpi_comm, dest, source, nelems, pe_root);
+
+  return;
+}
+
+__host__ inline MPI_Op HostInterface::get_mpi_op(ROC_SHMEM_OP Op) {
+  switch (Op) {
+    case ROC_SHMEM_SUM:
+      return MPI_SUM;
+    case ROC_SHMEM_MAX:
+      return MPI_MAX;
+    case ROC_SHMEM_MIN:
+      return MPI_MIN;
+    case ROC_SHMEM_PROD:
+      return MPI_PROD;
+    case ROC_SHMEM_AND:
+      return MPI_BAND;
+    case ROC_SHMEM_OR:
+      return MPI_BOR;
+    case ROC_SHMEM_XOR:
+      return MPI_BXOR;
+    default:
+      fprintf(stderr, "Unknown ROC_SHMEM op MPI conversion %d\n", Op);
+      abort();
+      return 0;
+  }
+}
+
+template <typename T>
+__host__ inline MPI_Datatype HostInterface::get_mpi_type() {
+  fprintf(stderr, "Unknown or unimplemented datatype \n");
+}
+
+#define GET_MPI_TYPE(T, MPI_T)                                    \
+  template <>                                                     \
+  __host__ inline MPI_Datatype HostInterface::get_mpi_type<T>() { \
+    return MPI_T;                                                 \
+  }
+
+GET_MPI_TYPE(int, MPI_INT)
+GET_MPI_TYPE(unsigned int, MPI_UNSIGNED)
+GET_MPI_TYPE(short, MPI_SHORT)
+GET_MPI_TYPE(unsigned short, MPI_UNSIGNED_SHORT)
+GET_MPI_TYPE(long, MPI_LONG)
+GET_MPI_TYPE(unsigned long, MPI_UNSIGNED_LONG)
+GET_MPI_TYPE(long long, MPI_LONG_LONG)
+GET_MPI_TYPE(unsigned long long, MPI_UNSIGNED_LONG_LONG)
+GET_MPI_TYPE(float, MPI_FLOAT)
+GET_MPI_TYPE(double, MPI_DOUBLE)
+GET_MPI_TYPE(char, MPI_CHAR)
+GET_MPI_TYPE(signed char, MPI_SIGNED_CHAR)
+GET_MPI_TYPE(unsigned char, MPI_UNSIGNED_CHAR)
+
+template <typename T>
+__host__ void HostInterface::amo_add(void* dst, T value, int pe,
+                                     WindowInfo* window_info) {
+  /*
+   * Most MPI implementations tend to use active messages to implement
+   * MPI_Accumulate. So, to eliminate the potential involvement of the
+   * target PE, we instead use fetch_add and disregard the return value.
+   */
+  [[maybe_unused]] T ret{amo_fetch_add(dst, value, pe, window_info)};
+}
+
+template <typename T>
+__host__ void HostInterface::amo_cas(void* dst, T value, T cond, int pe,
+                                     WindowInfo* window_info) {
+  /* Perform the compare and swap and disregard the return value */
+  [[maybe_unused]] T ret{amo_fetch_cas(dst, value, cond, pe, window_info)};
+}
+
+template <typename T>
+__host__ T HostInterface::amo_fetch_add(void* dst, T value, int pe,
+                                        WindowInfo* window_info) {
+  /* Calculate offset of remote dest from base address of window */
+  MPI_Aint offset{
+      compute_offset(dst, window_info->get_start(), window_info->get_end())};
+
+  /*
+   * Flush the HDP of the remote PE so that the NIC does not
+   * read stale values
+   */
+  flush_remote_hdp(pe);
+
+  /* Offload remote fetch and op operation to MPI */
+  T ret{};
+  MPI_Win win{window_info->get_win()};
+  MPI_Datatype mpi_type{get_mpi_type<T>()};
+  MPI_Fetch_and_op(&value, &ret, mpi_type, pe, offset, MPI_SUM, win);
+
+  MPI_Win_flush_local(pe, win);
+
+  return ret;
+}
+
+template <typename T>
+__host__ T HostInterface::amo_fetch_cas(void* dst, T value, T cond, int pe,
+                                        WindowInfo* window_info) {
+  /* Calculate offset of remote dest from base address of window */
+  MPI_Aint offset{
+      compute_offset(dst, window_info->get_start(), window_info->get_end())};
+
+  /*
+   * Flush the HDP of the remote PE so that the NIC does not
+   * read stale values
+   */
+  flush_remote_hdp(pe);
+
+  /* Offload remote compare and swap operation to MPI */
+  T ret{};
+  MPI_Win win{window_info->get_win()};
+  MPI_Datatype mpi_type{get_mpi_type<T>()};
+  MPI_Compare_and_swap(&value, &cond, &ret, mpi_type, pe, offset, win);
+
+  MPI_Win_flush_local(pe, win);
+
+  return ret;
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void HostInterface::to_all_internal(MPI_Comm mpi_comm, T* dest,
+                                             const T* source, int nreduce) {
+  DPRINTF("Function: host_to_all_internal\n");
+
+  MPI_Op mpi_op{get_mpi_op(Op)};
+
+  MPI_Datatype mpi_type{get_mpi_type<T>()};
+
+  void* send_buf{const_cast<T*>(source)};
+  void* recv_buf{const_cast<T*>(dest)};
+
+  /*
+   * Flush my HDP so that the NIC does not read stale values
+   */
+  hdp_policy_->hdp_flush();
+
+  /*
+   * Offload the allreduce to MPI
+   */
+  MPI_Allreduce((dest == source) ? MPI_IN_PLACE : send_buf, recv_buf, nreduce,
+                mpi_type, mpi_op, mpi_comm);
+
+  return;
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void HostInterface::to_all(T* dest, const T* source, int nreduce,
+                                    int pe_start, int log_pe_stride,
+                                    int pe_size, [[maybe_unused]] T* p_wrk,
+                                    [[maybe_unused]] long* p_sync) {
+  DPRINTF("Function: host_to_all\n");
+
+  /*
+   * Get an MPI communicator for active set of PEs
+   * Note: pe_root is w.r.t. the active set, hence
+   * the MPI communicator contains the root as well.
+   */
+  MPI_Comm mpi_comm{get_mpi_comm(pe_start, log_pe_stride, pe_size)};
+
+  to_all_internal<T, Op>(mpi_comm, dest, source, nreduce);
+
+  return;
+}
+
+template <typename T, ROC_SHMEM_OP Op>
+__host__ void HostInterface::to_all(roc_shmem_team_t team, T* dest,
+                                    const T* source, int nreduce) {
+  DPRINTF("Function: Team-based host_to_all\n");
+
+  /*
+   * Get the MPI communicator of this team
+   */
+  Team* team_obj{get_internal_team(team)};
+  MPI_Comm mpi_comm{team_obj->mpi_comm};
+
+  to_all_internal<T, Op>(mpi_comm, dest, source, nreduce);
+
+  return;
+}
+
+template <typename T>
+__host__ inline int HostInterface::compare(roc_shmem_cmps cmp, T input_val,
+                                           T target_val) {
+  int cond_satisfied{0};
+
+  switch (cmp) {
+    case ROC_SHMEM_CMP_EQ:
+      cond_satisfied = (input_val == target_val) ? 1 : 0;
+      break;
+    case ROC_SHMEM_CMP_NE:
+      cond_satisfied = (input_val != target_val) ? 1 : 0;
+      break;
+    case ROC_SHMEM_CMP_GT:
+      cond_satisfied = (input_val > target_val) ? 1 : 0;
+      break;
+    case ROC_SHMEM_CMP_GE:
+      cond_satisfied = (input_val >= target_val) ? 1 : 0;
+      break;
+    case ROC_SHMEM_CMP_LT:
+      cond_satisfied = (input_val < target_val) ? 1 : 0;
+      break;
+    case ROC_SHMEM_CMP_LE:
+      cond_satisfied = (input_val <= target_val) ? 1 : 0;
+      break;
+    default:
+      assert(cmp >= ROC_SHMEM_CMP_EQ && cmp <= ROC_SHMEM_CMP_LE);
+      break;
+  }
+
+  return cond_satisfied;
+}
+
+template <typename T>
+__host__ inline int HostInterface::test_and_compare(MPI_Aint offset,
+                                                    MPI_Datatype mpi_type,
+                                                    roc_shmem_cmps cmp, T val,
+                                                    MPI_Win win) {
+  T fetched_val{};
+
+  /*
+   * Flush the HDP so that the CPU doesn't read stale values
+   */
+  hdp_policy_->hdp_flush();
+
+  MPI_Fetch_and_op(nullptr,  // because no operation happening here
+                   &fetched_val, mpi_type, my_pe_, offset, MPI_NO_OP, win);
+  MPI_Win_flush_local(my_pe_, win);
+
+  /*
+   * Compare based on the operation
+   */
+  return compare(cmp, fetched_val, val);
+}
+
+template <typename T>
+__host__ void HostInterface::wait_until(T* ptr, roc_shmem_cmps cmp, T val,
+                                        WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until\n");
+
+  /*
+   * Find the offset of this memory in the window
+   */
+  MPI_Aint offset{
+      compute_offset(ptr, window_info->get_start(), window_info->get_end())};
+
+  MPI_Datatype mpi_type{get_mpi_type<T>()};
+  MPI_Win win{window_info->get_win()};
+
+  /*
+   * Continuously read the ptr atomically until it satisfies the condition
+   */
+  while (1) {
+    int cond_satisfied{test_and_compare(offset, mpi_type, cmp, val, win)};
+
+    if (cond_satisfied) {
+      break;
+    }
+  }
+}
+
+__host__ size_t status_entry(size_t nelems,
+                             const int *status,
+                             bool* done_flags) {
+  size_t i{0};
+  size_t pos{SIZE_MAX};
+  while (i < nelems) {
+    if (status[i]) {
+      done_flags[i] = 1;
+    } else {
+      pos = min(i, pos);
+    }
+    i++;
+  }
+  return pos;
+}
+
+__host__ size_t status_entry(size_t nelems,
+                             const int *status) {
+  size_t i{0};
+  while (i < nelems) {
+    if (status[i] == 0) {
+      return i;
+    }
+    i++;
+  }
+  return i;
+}
+
+template <typename T>
+__host__ size_t HostInterface::wait_until_any(T* ptr, size_t nelems,
+                                              const int *status,
+                                              roc_shmem_cmps cmp, T val,
+                                              WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until_any\n");
+
+  // zero nelems error condition
+  if (!nelems) {
+    return SIZE_MAX;
+  }
+
+  size_t pos{status_entry(nelems, status)};
+
+  // invalid (empty) status array error condition
+  if (pos == nelems) {
+    return SIZE_MAX;
+  }
+
+  while (true) {
+    for (size_t i{pos}; i < nelems; i++) {
+      // skip entries marked with non-zero status
+      if (status[i]) {
+        continue;
+      }
+      if (test(ptr + i, cmp, val, window_info)) {
+        return i;
+      }
+    }
+  }
+}
+
+template <typename T>
+__host__ void HostInterface::wait_until_all(T* ptr, size_t nelems,
+                                            const int *status,
+                                            roc_shmem_cmps cmp, T val,
+                                            WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until_all\n");
+
+  // zero nelems error condition
+  if (!nelems) {
+    return;
+  }
+
+  size_t pos{status_entry(nelems, status)};
+
+  // invalid (empty) status array error condition
+  if (pos == nelems) {
+    return;
+  }
+
+  for (size_t i{pos}; i < nelems; i++) {
+    if (status[i]) {
+      continue;
+    }
+    while (!test(ptr + i, cmp, val, window_info)) {
+    }
+  }
+}
+
+template <typename T>
+__host__ size_t HostInterface::wait_until_some(T* ptr, size_t nelems,
+                                             size_t* indices,
+                                             const int *status,
+                                             roc_shmem_cmps cmp, T val,
+                                             WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until_some\n");
+
+  // zero nelems error condition
+  if (!nelems) {
+    return 0;
+  }
+
+  size_t pos{status_entry(nelems, status)};
+
+  // invalid (empty) status array error condition
+  if (pos == nelems) {
+    return 0;
+  }
+
+  bool done {false};
+  size_t ncompleted {0};
+  while (!done) {
+    for (size_t i{pos}; i < nelems; i++) {
+      // skip entries marked with non-zero status
+      if (status[i]) {
+        continue;
+      }
+      if (test(ptr + i, cmp, val, window_info)) {
+        done = true;
+        indices[ncompleted] = i;
+        ncompleted++;
+      }
+    }
+  }
+  return ncompleted;
+}
+
+template <typename T>
+__host__ void HostInterface::wait_until_all_vector(T* ptr, size_t nelems,
+                                                   const int *status,
+                                                   roc_shmem_cmps cmp, T* vals,
+                                                   WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until_all_vector\n");
+}
+
+template <typename T>
+__host__ size_t HostInterface::wait_until_any_vector(T* ptr, size_t nelems,
+                                                     const int *status,
+                                                     roc_shmem_cmps cmp, T* vals,
+                                                     WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until_any_vector\n");
+  return 0;
+}
+
+template <typename T>
+__host__ size_t HostInterface::wait_until_some_vector(T* ptr, size_t nelems,
+                                                    size_t* indices,
+                                                    const int *status,
+                                                    roc_shmem_cmps cmp, T* vals,
+                                                    WindowInfo* window_info) {
+  DPRINTF("Function: host_wait_until_some_vector\n");
+  return 0;
+}
+
+template <typename T>
+__host__ int HostInterface::test(T* ptr, roc_shmem_cmps cmp, T val,
+                                 WindowInfo* window_info) {
+  DPRINTF("Function: host_test\n");
+
+  /*
+   * Find the offset of this memory in the window
+   */
+  MPI_Aint offset{
+      compute_offset(ptr, window_info->get_start(), window_info->get_end())};
+
+  MPI_Datatype mpi_type{get_mpi_type<T>()};
+
+  return test_and_compare(offset, mpi_type, cmp, val, window_info->get_win());
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_HOST_HOST_TEMPLATES_HPP_

src/ipc_policy.cpp

@@ -0,0 +1,125 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "ipc_policy.hpp"
+
+#include <mpi.h>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "backend_bc.hpp"
+#include "context_incl.hpp"
+#include "util.hpp"
+
+namespace rocshmem {
+
+__host__ void IpcOnImpl::ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
+                                     MPI_Comm thread_comm) {
+  /*
+   * Create an MPI communicator that deals only with local processes.
+   */
+  MPI_Comm shmcomm;
+  MPI_Comm_split_type(thread_comm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL,
+                      &shmcomm);
+
+  /*
+   * Figure out how many local process there are.
+   */
+  int Shm_size;
+  MPI_Comm_size(shmcomm, &Shm_size);
+  shm_size = Shm_size;
+
+  /*
+   * Figure out how this process' rank among local processes.
+   */
+  int shm_rank;
+  MPI_Comm_rank(shmcomm, &shm_rank);
+
+  /*
+   * Allocate a host-side c-array to hold the IPC handles.
+   */
+  void *ipc_mem_handle_uncast = malloc(shm_size * sizeof(hipIpcMemHandle_t));
+  hipIpcMemHandle_t *vec_ipc_handle =
+      reinterpret_cast<hipIpcMemHandle_t *>(ipc_mem_handle_uncast);
+
+  /*
+   * Call into the hip runtime to get an IPC handle for my symmetric
+   * heap and store that IPC handle into the host-side c-array which was
+   * just allocated.
+   */
+  char *base_heap = heap_bases[my_pe];
+  CHECK_HIP(hipIpcGetMemHandle(&vec_ipc_handle[shm_rank], base_heap));
+
+  /*
+   * Do an all-to-all exchange with each local processing element to
+   * share the symmetric heap IPC handles.
+   */
+  MPI_Allgather(MPI_IN_PLACE, sizeof(hipIpcMemHandle_t), MPI_CHAR,
+                vec_ipc_handle, sizeof(hipIpcMemHandle_t), MPI_CHAR, shmcomm);
+
+  /*
+   * Allocate device-side array to hold the IPC symmetric heap base
+   * addresses.
+   */
+  char **ipc_base;
+  CHECK_HIP(hipMalloc(reinterpret_cast<void **>(&ipc_base),
+                      shm_size * sizeof(char **)));
+
+  /*
+   * For all local processing elements, initialize the device-side array
+   * with the IPC symmetric heap base addresses.
+   */
+  for (size_t i = 0; i < shm_size; i++) {
+    if (i != shm_rank) {
+      void **ipc_base_uncast = reinterpret_cast<void **>(&ipc_base[i]);
+      CHECK_HIP(hipIpcOpenMemHandle(ipc_base_uncast, vec_ipc_handle[i],
+                                    hipIpcMemLazyEnablePeerAccess));
+      // TODO(bpotter): add some error checking here if happens to fail
+    } else {
+      ipc_base[i] = base_heap;
+    }
+  }
+
+  /*
+   * Set member variables used by subsequent method calls.
+   */
+  ipc_bases = ipc_base;
+
+  /*
+   * Free the host-side memory used to exchange the symmetric heap base
+   * addresses.
+   */
+  free(vec_ipc_handle);
+}
+
+__device__ void IpcOnImpl::ipcCopy(void *dst, void *src, size_t size) {
+  memcpy(dst, src, size);
+}
+
+__device__ void IpcOnImpl::ipcCopy_wave(void *dst, void *src, size_t size) {
+  memcpy_wave(dst, src, size);
+}
+
+__device__ void IpcOnImpl::ipcCopy_wg(void *dst, void *src, size_t size) {
+  memcpy_wg(dst, src, size);
+}
+
+}  // namespace rocshmem

src/ipc_policy.hpp

@@ -0,0 +1,166 @@
+/******************************************************************************
+ * 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_POLICY_HPP_
+#define LIBRARY_SRC_IPC_POLICY_HPP_
+
+#include <hip/hip_runtime.h>
+#include <mpi.h>
+
+#include <atomic>
+#include <vector>
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "memory/hip_allocator.hpp"
+#include "util.hpp"
+
+namespace rocshmem {
+
+class Backend;
+class Context;
+
+class IpcOnImpl {
+  using HEAP_BASES_T = std::vector<char *, StdAllocatorHIP<char *>>;
+
+ public:
+  uint32_t shm_size{0};
+
+  char **ipc_bases{nullptr};
+
+  __host__ void ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
+                            MPI_Comm thread_comm);
+
+  __device__ bool isIpcAvailable(int my_pe, int target_pe) {
+    return my_pe / shm_size == target_pe / shm_size;
+  }
+  __device__ void ipcGpuInit(Backend *gpu_backend, Context *ctx, int thread_id);
+
+  __device__ void ipcCopy(void *dst, void *src, size_t size);
+
+  __device__ void ipcCopy_wg(void *dst, void *src, size_t size);
+
+  __device__ void ipcCopy_wave(void *dst, void *src, size_t size);
+
+  __device__ void ipcFence() { __threadfence(); }
+
+  template <typename T>
+  __device__ T ipcAMOFetchAdd(T *val, T value) {
+    return __hip_atomic_fetch_add(val, value, __ATOMIC_RELAXED,
+                                  __HIP_MEMORY_SCOPE_AGENT);
+  }
+
+  template <typename T>
+  __device__ T ipcAMOFetchCas(T *val, T cond, T value) {
+    __hip_atomic_compare_exchange_strong(val, &cond, value, __ATOMIC_RELAXED,
+                                         __ATOMIC_RELAXED,
+                                         __HIP_MEMORY_SCOPE_AGENT);
+    return cond;
+  }
+
+  template <typename T>
+  __device__ void ipcAMOAdd(T *val, T value) {
+    __hip_atomic_fetch_add(val, value, __ATOMIC_RELAXED,
+                           __HIP_MEMORY_SCOPE_AGENT);
+  }
+
+  template <typename T>
+  __device__ void ipcAMOCas(T *val, T cond, T value) {
+    __hip_atomic_compare_exchange_strong(val, &cond, value, __ATOMIC_RELAXED,
+                                         __ATOMIC_RELAXED,
+                                         __HIP_MEMORY_SCOPE_AGENT);
+  }
+
+  template <typename T>
+  __device__ void ipcAMOSet(T *val, T value) {
+    __hip_atomic_store(val, value, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
+  }
+
+  __device__ void zero_byte_read(int pe) {
+    int local_pe = pe % shm_size;
+    uint32_t *pe_ipc_base = reinterpret_cast<uint32_t *>(ipc_bases[local_pe]);
+    volatile uint32_t read_value = __hip_atomic_load(
+        pe_ipc_base, __ATOMIC_SEQ_CST, __HIP_MEMORY_SCOPE_SYSTEM);
+  }
+};
+
+// clang-format off
+NOWARN(-Wunused-parameter,
+class IpcOffImpl {
+  using HEAP_BASES_T = std::vector<char *, StdAllocatorHIP<char *>>;
+
+ public:
+  uint32_t shm_size{0};
+
+  char **ipc_bases{nullptr};
+
+  __host__ void ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
+                            MPI_Comm thread_comm) {}
+
+  __device__ bool isIpcAvailable(int my_pe, int target_pe) { return false; }
+
+  __device__ void ipcGpuInit(Backend *roc_shmem_handle, Context *ctx,
+                             int thread_id) {}
+
+  __device__ void ipcCopy(void *dst, void *src, size_t size) {}
+
+  __device__ void ipcCopy_wg(void *dst, void *src, size_t size) {}
+
+  __device__ void ipcCopy_wave(void *dst, void *src, size_t size) {}
+
+  __device__ void ipcFence() {}
+
+  template <typename T>
+  __device__ T ipcAMOFetchAdd(T *val, T value) {
+    return T();
+  }
+
+  template <typename T>
+  __device__ T ipcAMOFetchCas(T *val, T cond, T value) {
+    return T();
+  }
+
+  template <typename T>
+  __device__ void ipcAMOAdd(T *val, T value) {}
+
+  template <typename T>
+  __device__ void ipcAMOSet(T *val, T value) {}
+
+  template <typename T>
+  __device__ void ipcAMOCas(T *val, T cond, T value) {}
+
+  __device__ void zero_byte_read(int pe) {}
+};
+)
+// clang-format on
+
+/*
+ * Select which one of our IPC policies to use at compile time.
+ */
+#ifdef USE_IPC
+typedef IpcOnImpl IpcImpl;
+#else
+typedef IpcOffImpl IpcImpl;
+#endif
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_IPC_POLICY_HPP_

src/memory/CMakeLists.txt

@@ -0,0 +1,32 @@
+###############################################################################
+# 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.
+###############################################################################
+
+###############################################################################
+# ADD ROCSHMEM TARGET FOR FILES IN CURRENT DIRECTORY
+###############################################################################
+target_sources(
+  ${PROJECT_NAME}
+  PRIVATE
+    single_heap.cpp
+    slab_heap.cpp
+    memory_allocator.cpp
+)

src/memory/address_record.hpp

@@ -0,0 +1,133 @@
+/******************************************************************************
+ * 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_MEMORY_ADDRESS_RECORD_HPP_
+#define LIBRARY_SRC_MEMORY_ADDRESS_RECORD_HPP_
+
+#include <cassert>
+#include <utility>
+
+/**
+ * @file address_record.hpp
+ *
+ * @brief Contains an address record class
+ *
+ * The address record is intended to be used by allocators which need to
+ * keep track of raw pointers and the corresponding allocation size.
+ */
+
+namespace rocshmem {
+
+class AddressRecord {
+  /**
+   * @brief Helper type for address records
+   */
+  using AR_T = AddressRecord;
+
+ public:
+  /**
+   * @brief Default constructor
+   *
+   * @note This constructor generates an empty record. In some contexts,
+   * an empty record signifies an INVALID record and is treated
+   * accordingly.
+   */
+  AddressRecord() = default;
+
+  /**
+   * @brief Primary constructor type
+   *
+   * @param[in] raw pointer holding an address
+   * @param[in] size
+   */
+  AddressRecord(char* address, size_t size) : address_(address), size_(size) {}
+
+  /**
+   * @brief Accessor for address_
+   *
+   * @return raw pointer
+   */
+  char* get_address() { return address_; }
+
+  /**
+   * @brief Accessor for size_
+   *
+   * @return size
+   */
+  size_t get_size() { return size_; }
+
+  /**
+   * @brief Splits record in half
+   *
+   * @return pair containing initialized address records
+   */
+  std::pair<AR_T, AR_T> split() {
+    assert(address_);
+    assert(size_);
+    auto half_size{size_ >> 1};
+    assert(half_size);
+
+    AR_T e1{address_, half_size};
+    AR_T e2{address_ + half_size, half_size};
+
+    return {e1, e2};
+  }
+
+  /**
+   * @brief Combines record with another record
+   *
+   * @return an address record containing both input records
+   */
+  AR_T combine(AR_T other) {
+    assert(address_);
+    assert(other.get_address());
+    bool this_smaller{address_ < other.get_address()};
+
+    auto smaller_addr{this_smaller ? address_ : other.get_address()};
+    [[maybe_unused]] auto larger_addr{this_smaller ? other.get_address()
+                                                   : address_};
+
+    assert(size_ == other.get_size());
+    auto combined_size{size_ + size_};
+
+    assert(smaller_addr + size_ == larger_addr);
+
+    AR_T record{smaller_addr, combined_size};
+
+    return record;
+  }
+
+ private:
+  /**
+   * @brief raw memory pointer
+   */
+  char* address_{nullptr};
+
+  /**
+   * @brief size of address record
+   */
+  size_t size_{0};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_ADDRESS_RECORD_HPP_

src/memory/bin.hpp

@@ -0,0 +1,85 @@
+/******************************************************************************
+ * 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_MEMORY_BIN_HPP_
+#define LIBRARY_SRC_MEMORY_BIN_HPP_
+
+#include <cassert>
+#include <stack>
+
+/**
+ * @file bin.hpp
+ *
+ * @brief Simple container class
+ *
+ * The class wraps a stack and provides simple mutators. A different internal
+ * container class might be preferable to a stack.
+ */
+
+namespace rocshmem {
+
+template <typename T>
+class Bin {
+ public:
+  /**
+   * @brief Is stack empty?
+   *
+   * @return A boolean denoting stack emptiness
+   */
+  bool empty() { return stack_.empty(); }
+
+  /**
+   * @brief How many elements in stack?
+   *
+   * @return The number of elements in the stack
+   */
+  size_t size() { return stack_.size(); }
+
+  /**
+   * @brief Emplace an element in stack
+   *
+   * @param[in] An element
+   */
+  void put(T element) { stack_.emplace(element); }
+
+  /**
+   * @brief Retrieve an element from stack
+   *
+   * @return An element
+   */
+  T get() {
+    assert(stack_.size());
+    auto top = stack_.top();
+    stack_.pop();
+    return top;
+  }
+
+ private:
+  /**
+   * @brief Implementation container
+   */
+  std::stack<T> stack_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_BIN_HPP_

src/memory/binner.hpp

@@ -0,0 +1,246 @@
+/******************************************************************************
+ * 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_MEMORY_BINNER_HPP_
+#define LIBRARY_SRC_MEMORY_BINNER_HPP_
+
+#include <cassert>
+#include <climits>
+#include <iostream>
+#include <vector>
+
+#include "../constants.hpp"
+#include "bin.hpp"
+
+/**
+ * @file binner.hpp
+ *
+ * @brief Contains a builder for bins_ object in Pow2Bins.
+ */
+
+namespace rocshmem {
+
+/**
+ * @brief Helper to choose between clz variants
+ *
+ * @param[in] An input value
+ *
+ * @return __builtin_clz value
+ */
+template <typename T>
+inline int clz_fn([[maybe_unused]] T v) {
+  assert(false);
+  return 0;
+}
+
+template <>
+inline int clz_fn(unsigned v) {
+  return __builtin_clz(v);
+}
+
+template <>
+inline int clz_fn(unsigned long v) {
+  return __builtin_clzl(v);
+}
+
+template <>
+inline int clz_fn(unsigned long long v) {
+  return __builtin_clzll(v);
+}
+
+/**
+ * @brief Finds the bit position of the first one
+ *
+ * @param[in] An input value
+ *
+ * @return A bit position or UINT_MAX
+ */
+template <typename T>
+inline unsigned find_first_set_one(T v) {
+  if (v == 0) {
+    return UINT_MAX;
+  }
+  auto num_bytes{sizeof(T)};
+  auto num_bits{num_bytes * CHAR_BIT};
+  auto num_leading_zeroes{clz_fn(v)};
+  auto last_zero_bit_position{num_bits - num_leading_zeroes};
+  auto first_one_bit_position{last_zero_bit_position - 1};
+  return first_one_bit_position;
+}
+
+template <typename AR_T, typename BINS_T>
+class Binner {
+ public:
+  /**
+   * @brief Required for default construction of other objects
+   *
+   * @note Not intended for direct usage.
+   */
+  Binner() = default;
+
+  /**
+   * @brief Primary constructor type
+   *
+   * Constructs bins by inserting a set of bin into it.
+   *
+   * @param[in] Raw memory pointer to bins object
+   * @param[in] Raw memory pointer to symmetric heap memory.
+   * @param[in] Size of symmetric heap allocation.
+   */
+  Binner(BINS_T* bins, char* heap_ptr, size_t heap_size)
+      : bins_{bins}, heap_{heap_ptr, heap_size} {
+    auto bit_position_heap{find_first_set_one(heap_size)};
+    assert(bit_position_heap != UINT_MAX);
+
+    auto bit_position_align{find_first_set_one(ALIGNMENT)};
+    assert(bit_position_align != UINT_MAX);
+
+    assert(bit_position_heap >= bit_position_align);
+
+    auto difference_in_position{bit_position_heap - bit_position_align};
+    auto vector_size{difference_in_position + 1};
+
+    std::vector<size_t> v;
+    v.resize(vector_size);
+
+    auto pow2 = [&bit_position_align]() {
+      return std::pow(2, bit_position_align++);
+    };
+
+    std::generate(v.begin(), v.end(), pow2);
+
+    for (auto x : v) {
+      bins_->insert({x, Bin<AR_T>{}});
+    }
+  }
+
+  /**
+   * @brief Assign symmetric heap memory to bins object
+   */
+  void assign_heap_to_bins() {
+    assert(bins_);
+    assert(heap_.get_address());
+    assert(heap_.get_size());
+
+    ignore_unaligned_heap_memory();
+    assign_heap_memory_to_bins();
+  }
+
+  /**
+   * @brief Accessor for bins_ pointer.
+   *
+   * @return Raw pointer to bins_
+   *
+   * @note Used by unit-test test fixture
+   */
+  BINS_T* get_bins() { return bins_; }
+
+  /**
+   * @brief Dump the bins_ object to the standard out
+   */
+  void dump_bins() {
+    auto dump_bin = [](const auto& pair) {
+      auto [address, bin] = pair;
+      std::cout << "address " << address << " bin.size " << bin.size()
+                << std::endl;
+    };
+    std::for_each(bins_->begin(), bins_->end(), dump_bin);
+  }
+
+ private:
+  /**
+   * @brief Shift assignment point past unaligned memory.
+   */
+  void ignore_unaligned_heap_memory() {
+    static_assert(sizeof(char) == 1);
+
+    auto heap_ptr{heap_.get_address()};
+    size_t off_by_bytes{reinterpret_cast<size_t>(heap_ptr) % ALIGNMENT};
+
+    if (off_by_bytes) {
+      size_t required_offset{ALIGNMENT - off_by_bytes};
+      heap_ptr += required_offset;
+      auto heap_size{heap_.get_size() - required_offset};
+      heap_ = AR_T(heap_ptr, heap_size);
+    }
+  }
+
+  /**
+   * @brief Delegate method for top-level heap assignment.
+   */
+  void assign_heap_memory_to_bins() {
+    AR_T chunk = heap_;
+    while (chunk.get_address()) {
+      chunk = assign_heap_chunk(chunk);
+    }
+  }
+
+  /**
+   * @brief Chunk heap memory up into bin sized allocations
+   *
+   * The method starts with the largest chunks possible and proceeds
+   * to smaller chunks. It continues chunking until there are no more
+   * chunks which will fit into bins.
+   *
+   * @param[in] An address record
+   *
+   * @return An address record minus a chunk
+   */
+  AR_T assign_heap_chunk(AR_T address_record) {
+    char* memory_chunk{address_record.get_address()};
+    size_t memory_chunk_size{address_record.get_size()};
+
+    auto rit{bins_->rbegin()};
+    auto rend{bins_->rend()};
+    while (rit != rend) {
+      auto& [bin_memory_chunk_size, bin] = *rit;
+      if (memory_chunk_size >= bin_memory_chunk_size) {
+        break;
+      }
+      rit++;
+    }
+
+    if (rit == rend) {
+      return {nullptr, 0};
+    }
+
+    auto& [bin_memory_chunk_size, bin] = *rit;
+    bin.put({memory_chunk, bin_memory_chunk_size});
+
+    return {memory_chunk + bin_memory_chunk_size,
+            memory_chunk_size - bin_memory_chunk_size};
+  }
+
+  /**
+   * @brief Raw pointer to external bins_ object
+   */
+  BINS_T* bins_{nullptr};
+
+  /**
+   * @brief An address record holding heap object's information
+   */
+  AR_T heap_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_BINNER_HPP_

src/memory/dev_mono_linear.hpp

@@ -0,0 +1,144 @@
+/******************************************************************************
+ * 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_MEMORY_DEV_MONO_LINEAR_HPP_
+#define LIBRARY_SRC_MEMORY_DEV_MONO_LINEAR_HPP_
+
+#include <cassert>
+
+#include "shmem_allocator_strategy.hpp"
+#include "../util.hpp"
+
+/**
+ * @file dev_mono_linear.hpp
+ *
+ * @brief Contains an allocator strategy for the heap.
+ *
+ * This strategy returns memory chunks by monotonically increasing a pointer.
+ */
+
+namespace rocshmem {
+
+template <typename HM_T>
+class DevMonoLinear : public ShmemAllocatorStrategy {
+ public:
+  /**
+   * @brief Required for default construction of other objects
+   *
+   * @note Not intended for direct usage.
+   */
+  DevMonoLinear() = default;
+
+  /**
+   * @brief Primary constructor type
+   *
+   * @param[in] Raw pointer to heap memory type
+   */
+  explicit DevMonoLinear(HM_T* heap_mem)
+      : heap_mem_{heap_mem}, current_ptr_{heap_mem_->get_ptr()} {}
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   */
+  void alloc(char** ptr, size_t request_size) override {
+    assert(ptr);
+    *ptr = nullptr;
+
+    if (!request_size) {
+      return;
+    }
+
+    char* heap_end{heap_mem_->get_ptr() + heap_mem_->get_size()};
+
+    if (current_ptr_ + request_size < heap_end) {
+      *ptr = current_ptr_;
+      current_ptr_ += request_size;
+    }
+  }
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   */
+  __device__ void alloc(char** ptr, size_t request_size) override {
+    if (is_thread_zero_in_block()) {
+      assert(ptr);
+      *ptr = nullptr;
+
+      if (!request_size) {
+        return;
+      }
+
+      char* heap_end{heap_mem_->get_ptr() + heap_mem_->get_size()};
+
+      if (current_ptr_ + request_size < heap_end) {
+        *ptr = current_ptr_;
+        current_ptr_ += request_size;
+      }
+    }
+  }
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * Released memory ignored.
+   *
+   * @param[in] Raw pointer to heap memory
+   */
+  __host__ void free([[maybe_unused]] char* ptr) override {}
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * Released memory ignored.
+   *
+   * @param[in] Raw pointer to heap memory
+   */
+  __device__ void free([[maybe_unused]] char* ptr) override {}
+
+  /**
+   * @brief Return pointer to monotonic linear ptr
+   *
+   * @return Raw pointer to heap memory.
+   */
+  __host__ __device__ char* current() { return current_ptr_; }
+
+ private:
+  /**
+   * The pointer is used to keep reference to heap memory.
+   */
+  HM_T* heap_mem_{nullptr};
+
+  /**
+   * The pointer is used to track the next allocation point.
+   */
+  char* current_ptr_{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_DEV_MONO_LINEAR_HPP_

src/memory/heap_memory.hpp

@@ -0,0 +1,133 @@
+/******************************************************************************
+ * 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_MEMORY_HEAP_MEMORY_HPP_
+#define LIBRARY_SRC_MEMORY_HEAP_MEMORY_HPP_
+
+#include <cassert>
+#include <memory>
+#include <utility>
+
+/**
+ * @file heap_memory.hpp
+ *
+ * @brief Contains heap memory class
+ *
+ * @note The heap memory class owns the symmetric heap memory allocation
+ */
+
+namespace rocshmem {
+
+template <typename ALLOCATOR>
+class HeapMemory {
+ public:
+  /**
+   * @brief Primary constructor type
+   *
+   * Uses default heap size specified in class body.
+   */
+  HeapMemory() : HeapMemory(gibibyte_) {}
+
+  /**
+   * @brief Secondary constructor type
+   *
+   * @param[in] User-specified size used as heap size
+   */
+  explicit HeapMemory(size_t size) : size_{size} {
+    char* temp;
+    allocator_.allocate(reinterpret_cast<void**>(&temp), size_);
+    assert(temp);
+    std::unique_ptr<char, Deleter> up{temp};
+    up_ = std::move(up);
+
+    /*
+     * Set a c-style ptr for access by the device.
+     */
+    ptr_ = up_.get();
+  }
+
+  /**
+   * @brief Accessor for heap ptr
+   *
+   * @return Raw memory pointer
+   */
+  __host__ __device__ char* get_ptr() { return ptr_; }
+
+  /**
+   * @brief Accessor for heap size
+   *
+   * @return Heap size
+   */
+  __host__ __device__ size_t get_size() { return size_; }
+
+  /**
+   * @brief Returns is the heap is allocated with managed memory
+   *
+   * @return bool
+   */
+  bool is_managed() { return allocator_.is_managed(); }
+
+ private:
+  /**
+   * @brief Template type member with allocate and deallocate methods.
+   */
+  ALLOCATOR allocator_{};
+
+  /**
+   * @brief Wrap deallocator into a functor for up_ template.
+   */
+  class Deleter {
+   public:
+    void operator()(void* x) { a_.deallocate(x); }
+
+   private:
+    ALLOCATOR a_;
+  };
+
+  /**
+   * @brief Owning pointer to heap memory.
+   */
+  std::unique_ptr<char, Deleter> up_{nullptr};
+
+  /**
+   * @brief Named constant for a gibibyte.
+   */
+  static constexpr size_t gibibyte_{1 << 30};
+
+  /**
+   * @brief Size of heap memory.
+   */
+  size_t size_{gibibyte_};
+
+  /**
+   * @brief A handle to access the internal memory.
+   *
+   * In general, device code cannot access standard library routines
+   * like std::unique_ptr::get(). Circumvent this problem by caching
+   * the pointer manually in this class.
+   */
+  char* ptr_{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_HEAP_MEMORY_HPP_

src/memory/heap_type.hpp

@@ -0,0 +1,55 @@
+/******************************************************************************
+ * 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_MEMORY_HEAP_TYPE_HPP_
+#define LIBRARY_SRC_MEMORY_HEAP_TYPE_HPP_
+
+#include "config.h"  // NOLINT(build/include_subdir)
+#include "hip_allocator.hpp"
+
+/**
+ * @file heap_type.hpp
+ *
+ * @brief Contains the type of heap memory used to allocate the symmetric heap
+ *
+ * The heap type in this file is used by other classes to select the allocation
+ * policy for the symmetric heap. The heap type choices depend on whether
+ * the heap is cacheable in hardware and if the heap is a managed memory type.
+ */
+
+namespace rocshmem {
+
+#if defined USE_MANAGED_HEAP
+using HEAP_T = HeapMemory<HIPAllocatorManaged>;
+#elif defined USE_COHERENT_HEAP || defined USE_CACHED_HEAP
+using HEAP_T = HeapMemory<HIPAllocator>;
+#elif defined USE_HOST_HEAP
+using HEAP_T = HeapMemory<HostAllocator>;
+#elif defined USE_HIP_HOST_HEAP
+using HEAP_T = HeapMemory<HIPHostAllocator>;
+#else
+using HEAP_T = HeapMemory<HIPDefaultFinegrainedAllocator>;
+#endif
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_HEAP_TYPE_HPP_

src/memory/hip_allocator.hpp

@@ -0,0 +1,139 @@
+/******************************************************************************
+ * 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_MEMORY_HIP_ALLOCATOR_HPP_
+#define LIBRARY_SRC_MEMORY_HIP_ALLOCATOR_HPP_
+
+/**
+ * @file hip_allocator.hpp
+ *
+ * @brief Contains HIP wrapper class for memory allocator
+ */
+
+#include <hip/hip_runtime_api.h>
+
+#include <cstdlib>
+#include <limits>
+
+#include "memory_allocator.hpp"
+
+// `hipDeviceMallocUncached` was introduced at ROCm 5.5
+#if HIP_VERSION_MAJOR >= 5 && HIP_VERSION_MINOR >= 5
+#define HIP_SUPPORTS_MALLOC_UNCACHED
+#endif
+namespace rocshmem {
+
+class HIPAllocator : public MemoryAllocator {
+ public:
+  HIPAllocator() : MemoryAllocator(hipMalloc, hipFree) {}
+};
+
+class HIPAllocatorFinegrained : public MemoryAllocator {
+ public:
+  HIPAllocatorFinegrained()
+      : MemoryAllocator(hipExtMallocWithFlags, hipFree,
+                        hipDeviceMallocFinegrained) {}
+};
+
+#ifdef HIP_SUPPORTS_MALLOC_UNCACHED
+class HIPAllocatorUncached : public MemoryAllocator {
+ public:
+  HIPAllocatorUncached()
+      : MemoryAllocator(hipExtMallocWithFlags, hipFree,
+                        hipDeviceMallocUncached) {}
+};
+// The default fine-grained coherence allocator is the uncached allocator
+using HIPDefaultFinegrainedAllocator = HIPAllocatorUncached;
+#else
+// The default fine-grained coherence allocator is the fine-grained allocator
+using HIPDefaultFinegrainedAllocator = HIPAllocatorFinegrained;
+#endif
+
+class HIPAllocatorManaged : public MemoryAllocator {
+ public:
+  HIPAllocatorManaged()
+      : MemoryAllocator(hipMallocManaged, hipFree, hipMemAttachHost) {
+    _managed = true;
+  }
+};
+
+class HIPHostAllocator : public MemoryAllocator {
+ public:
+  HIPHostAllocator()
+      : MemoryAllocator(hipHostMalloc, hipFree, hipHostMallocCoherent) {}
+};
+
+class HostAllocator : public MemoryAllocator {
+ public:
+  HostAllocator() : MemoryAllocator(std::malloc, std::free) {}
+};
+
+class PosixAligned64Allocator : public MemoryAllocator {
+ public:
+  PosixAligned64Allocator() : MemoryAllocator(posix_memalign, std::free, 64) {}
+};
+
+template <class T>
+class StdAllocatorHIP {
+ public:
+  typedef T value_type;
+
+  StdAllocatorHIP() = default;
+
+  template <class U>
+  constexpr StdAllocatorHIP(const StdAllocatorHIP<U>&) noexcept {}
+
+  [[nodiscard]] T* allocate(size_t n) {
+    if (n > std::numeric_limits<size_t>::max() / sizeof(T)) {
+      throw std::bad_array_new_length();
+    }
+
+    T* p{nullptr};
+    allocator_.allocate(reinterpret_cast<void**>(&p), n * sizeof(T));
+    if (p) {
+      return p;
+    }
+
+    throw std::bad_alloc();
+  }
+
+  void deallocate(T* p, [[maybe_unused]] size_t n) noexcept {
+    allocator_.deallocate(p);
+  }
+
+ private:
+  HIPDefaultFinegrainedAllocator allocator_{};
+};
+
+template <class T, class U>
+bool operator==(const StdAllocatorHIP<T>&, const StdAllocatorHIP<U>&) {
+  return true;
+}
+
+template <class T, class U>
+bool operator!=(const StdAllocatorHIP<T>&, const StdAllocatorHIP<U>&) {
+  return false;
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_HIP_ALLOCATOR_HPP_

src/memory/memory_allocator.cpp

@@ -0,0 +1,99 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "memory_allocator.hpp"
+
+#include <cassert>
+
+#include "../util.hpp"
+
+namespace rocshmem {
+
+MemoryAllocator::MemoryAllocator(hipError_t (*hip_alloc_fn)(void**, size_t,
+                                                            unsigned),
+                                 hipError_t (*hip_free_fn)(void*),
+                                 unsigned flags)
+    : _hip_alloc_with_flags(hip_alloc_fn),
+      _hip_free(hip_free_fn),
+      _flags(flags) {}
+
+MemoryAllocator::MemoryAllocator(hipError_t (*hip_alloc_fn)(void**, size_t),
+                                 hipError_t (*hip_free_fn)(void*))
+    : _hip_alloc(hip_alloc_fn), _hip_free(hip_free_fn) {}
+
+MemoryAllocator::MemoryAllocator(std::function<void*(size_t)> alloc_fn,
+                                 std::function<void(void*)> free_fn)
+    : _alloc(alloc_fn), _free(free_fn) {}
+
+MemoryAllocator::MemoryAllocator(
+    std::function<int(void**, size_t, size_t)> posix_align_fn,
+    std::function<void(void*)> free_fn, size_t alignment)
+    : _alloc_posix_memalign(posix_align_fn),
+      _free(free_fn),
+      _alignment(alignment) {}
+
+void MemoryAllocator::allocate(void** void_ptr, size_t size) {
+  assert(void_ptr);
+
+  if (_alloc) {
+    *(reinterpret_cast<char**>(void_ptr)) =
+        reinterpret_cast<char*>(_alloc(size));
+    assert(*reinterpret_cast<char**>(void_ptr));
+    return;
+  }
+  if (_alloc_posix_memalign) {
+    assert(_alignment);
+    _alloc_posix_memalign(void_ptr, _alignment, size);
+    assert(*reinterpret_cast<char**>(void_ptr));
+    return;
+  }
+  if (_hip_alloc) {
+    CHECK_HIP(_hip_alloc(void_ptr, size));
+    return;
+  }
+  if (_hip_alloc_with_flags) {
+    CHECK_HIP(_hip_alloc_with_flags(void_ptr, size, _flags));
+    return;
+  }
+
+  assert(false);
+}
+
+void MemoryAllocator::deallocate(void* ptr) {
+  if (!ptr) {
+    return;
+  }
+  if (_free) {
+    _free(ptr);
+    return;
+  }
+  if (_hip_free) {
+    CHECK_HIP(_hip_free(ptr));
+    return;
+  }
+
+  assert(false);
+}
+
+bool MemoryAllocator::is_managed() { return _managed; }
+
+}  // namespace rocshmem

src/memory/memory_allocator.hpp

@@ -0,0 +1,168 @@
+/******************************************************************************
+ * 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_MEMORY_MEMORY_ALLOCATOR_HPP_
+#define LIBRARY_SRC_MEMORY_MEMORY_ALLOCATOR_HPP_
+
+/**
+ * @file memory_allocator.hpp
+ *
+ * @brief Wraps allocators in a uniform interface
+ */
+
+#include <hip/hip_runtime_api.h>
+
+#include <cstdlib>
+#include <functional>
+
+namespace rocshmem {
+
+class MemoryAllocator {
+ public:
+  /**
+   * @brief Required for default construction of other objects
+   *
+   * @note Not intended for direct usage.
+   */
+  MemoryAllocator() = default;
+
+  /**
+   * @brief Primary constructor
+   *
+   * @param[in] hip library allocation function
+   * @param[in] hip library free function
+   * @param[in] hip flags
+   *
+   * @note Finegrained flag is assumed by default
+   */
+  MemoryAllocator(hipError_t (*hip_alloc_fn)(void**, size_t, unsigned),
+                  hipError_t (*hip_free_fn)(void*), unsigned flags);
+
+  /**
+   * @brief Primary constructor
+   *
+   * @param[in] hip library allocation function
+   * @param[in] hip library free function
+   */
+  MemoryAllocator(hipError_t (*hip_alloc_fn)(void**, size_t),
+                  hipError_t (*hip_free_fn)(void*));
+
+  /**
+   * @brief Primary constructor
+   *
+   * @param[in] an allocation function
+   * @param[in] a free function
+   */
+  MemoryAllocator(std::function<void*(size_t)> alloc_fn,
+                  std::function<void(void*)> free_fn);
+
+  /**
+   * @brief Primary constructor
+   *
+   * @param[in] an allocation function
+   * @param[in] a free function
+   * @param[in] alignment for allocations
+   */
+  MemoryAllocator(std::function<int(void**, size_t, size_t)> posix_align_fn,
+                  std::function<void(void*)> free_fn, size_t alignment);
+
+  /**
+   * @brief Allocates memory
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   */
+  void allocate(void** ptr, size_t size);
+
+  /**
+   * @brief Deallocates memory
+   *
+   * @param[in] Address of raw pointer (&pointer_to_char)
+   */
+  void deallocate(void* ptr);
+
+  /**
+   * @brief Returns is memory is managed
+   *
+   * @return returns whether this memory is managed
+   */
+  bool is_managed();
+
+ protected:
+  /**
+   * @brief is this memory allocated using managed memory
+   */
+  bool _managed{false};
+
+ private:
+  /**
+   * @brief a posix allocator function for memory aligned accesses
+   */
+  std::function<int(void**, size_t, size_t)> _alloc_posix_memalign;
+
+  /**
+   * @brief alignment for posix_memalign allocations
+   */
+  size_t _alignment{0};
+
+  /**
+   * @brief a standard allocator function
+   */
+  std::function<void*(size_t)> _alloc{nullptr};
+
+  /**
+   * @brief a standard free function
+   */
+  std::function<void(void*)> _free{nullptr};
+
+  /**
+   * @brief a hip-specific allocator function
+   */
+  std::function<hipError_t(void**, size_t)> _hip_alloc{nullptr};
+
+  /**
+   * @brief a hip-specific allocator function
+   */
+  std::function<hipError_t(void**, size_t, unsigned)> _hip_alloc_with_flags{
+      nullptr};
+
+  /**
+   * @brief flags specified in hip_runtime_api.h
+   *
+   * Used as a runtime input parameter to the _hip_alloc_with_flags function.
+   */
+  unsigned _flags{0x0};
+
+  /**
+   * @brief a hip-specific free function
+   */
+  std::function<hipError_t(void*)> _hip_free{nullptr};
+
+  /**
+   * @brief a hip-specific return code
+   */
+  hipError_t _hip_return_value{hipSuccess};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_MEMORY_ALLOCATOR_HPP_

src/memory/notifier.hpp

@@ -0,0 +1,83 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+/**
+ * @file notifier.hpp
+ *
+ * @brief Contains the notification memory space for threads to communicate
+ * results with one another.
+ *
+ * Assume one thread does work on behalf of other threads (as a leader) and
+ * that work needs to be communicated to the other threads as a result.
+ * To expose the result, the threads need to share a memory space where the
+ * result can be written. The leader thread writes the result out to this
+ * memory space and all threads synchronize on it.
+ *
+ * This class allows the leader thread to notify other threads of the update.
+ */
+
+#ifndef LIBRARY_SRC_MEMORY_NOTIFIER_HPP_
+#define LIBRARY_SRC_MEMORY_NOTIFIER_HPP_
+
+#include "../device_proxy.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+class Notifier {
+ public:
+  __device__ uint64_t read() { return value_; }
+
+  __device__ void write(uint64_t val) {
+    if (is_thread_zero_in_block()) {
+      value_ = val;
+    }
+    publish();
+  }
+
+  __device__ void done() { __syncthreads(); }
+
+ private:
+  __device__ void publish() {
+    if (is_thread_zero_in_block()) {
+      __threadfence();
+    }
+    __syncthreads();
+  }
+
+  uint64_t value_{};
+};
+
+template <typename ALLOCATOR>
+class NotifierProxy {
+  using ProxyT = DeviceProxy<ALLOCATOR, Notifier, 1>;
+
+ public:
+  __host__ __device__ Notifier* get() { return proxy_.get(); }
+
+ private:
+  ProxyT proxy_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_NOTIFIER_HPP_

src/memory/pow2_bins.hpp

@@ -0,0 +1,353 @@
+/******************************************************************************
+ * 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_MEMORY_POW2_BINS_HPP_
+#define LIBRARY_SRC_MEMORY_POW2_BINS_HPP_
+
+#include <cassert>
+#include <map>
+
+#include "../constants.hpp"
+#include "bin.hpp"
+#include "binner.hpp"
+#include "shmem_allocator_strategy.hpp"
+
+/**
+ * @file pow2_bins.hpp
+ *
+ * @brief Contains an allocator strategy for the heap.
+ *
+ * This strategy returns memory chunks with power-of-two sizes.
+ */
+
+namespace rocshmem {
+
+template <typename AR_T, typename HM_T>
+class Pow2Bins : public ShmemAllocatorStrategy {
+  /**
+   * @brief Helper type for bins of address records
+   */
+  using BIN_T = Bin<AR_T>;
+
+  /**
+   * @brief Helper type for size to bin maps
+   */
+  using BINS_T = std::map<size_t, BIN_T>;
+
+  /**
+   * @brief Helper type for BINS_T iterator
+   */
+  using BINS_IT_T = typename BINS_T::iterator;
+
+  /**
+   * @brief Helper type for raw pointers to address record maps
+   */
+  using PROFFERED_T = std::map<char*, AR_T>;
+
+ public:
+  /**
+   * @brief Required for default construction of other objects
+   *
+   * @note Not intended for direct usage.
+   */
+  Pow2Bins() = default;
+
+  /**
+   * @brief Primary constructor type
+   *
+   * @param[in] Raw pointer to heap memory type
+   */
+  explicit Pow2Bins(HM_T* heap_mem)
+      : binner_{&bins_, heap_mem->get_ptr(), heap_mem->get_size()} {
+    binner_.assign_heap_to_bins();
+  }
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   */
+  void alloc(char** ptr, size_t request_size) override {
+    assert(ptr);
+    *ptr = nullptr;
+
+    if (!request_size) {
+      return;
+    }
+
+    auto do_request_using = [&](auto it) {
+      if (it != bins_.end()) {
+        AR_T record{retrieve_record_from_bin(it)};
+        /*
+         * Record retrieval may have generated INVALID record.
+         * If INVALID, do not mark record as proffered.
+         * Also notice, input pointer "ptr" may have nullptr value.
+         */
+        if (record.get_address()) {
+          emplace_in_proffered(record);
+        }
+        *ptr = record.get_address();
+      }
+    };
+
+    /*
+     * Try to find exact power-of-two matches.
+     */
+    do_request_using(bins_.find(request_size));
+    if (*ptr) {
+      return;
+    }
+
+    /*
+     * Round up to the nearest power-of-two size.
+     */
+    do_request_using(bins_.lower_bound(request_size));
+  }
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   *
+   * @note Not implemented
+   */
+  __device__ void alloc([[maybe_unused]] char** ptr,
+                        [[maybe_unused]] size_t request_size) override {}
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * Released memory is tracked by bookkeeping structures within this class.
+   *
+   * @param[in] Raw pointer to heap memory
+   *
+   * @note We do not attempt to combine records together into larger
+   * even if it is possible to do so. This may eventually exhaust the
+   * number of records in the larger bins.
+   */
+  void free(char* ptr) override {
+    auto record{retrieve_from_proffered(ptr)};
+    emplace_record_in_bin(record);
+  }
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * Released memory is tracked by bookkeeping structures within this class.
+   *
+   * @param[in] Raw pointer to heap memory
+   *
+   * @note Not implemented
+   */
+  __device__ void free([[maybe_unused]] char* ptr) override {}
+
+  /**
+   * @brief Sum of all proffered_ memory sizes
+   *
+   * @return memory size
+   */
+  size_t amount_proffered() {
+    size_t size{0};
+    auto sum_sizes = [&size](const auto& pair) {
+      auto [address, record] = pair;
+      size += record.get_size();
+    };
+    std::for_each(proffered_.begin(), proffered_.end(), sum_sizes);
+    return size;
+  }
+
+  /**
+   * @brief Accessor for bins_ internal data structure
+   *
+   * @return BINS_T raw pointer
+   *
+   * @note Used by unit-test test fixture
+   */
+  BINS_T* get_bins() { return &bins_; }
+
+ private:
+  /**
+   * @brief Store address record in appropriate bin
+   *
+   * @param[in] An address record
+   */
+  void emplace_record_in_bin(AR_T record) {
+    auto record_size{record.get_size()};
+
+    auto bins_it = bins_.find(record_size);
+    assert(bins_it != bins_.end());
+
+    auto& [IGNORE_BIN_SIZE, bin]{*bins_it};
+    bin.put(record);
+  }
+
+  /**
+   * @brief Retrieve address record from bin within a BINS_IT_T
+   *
+   * @param[in] a VALID FORWARD iterator containing a bin
+   *
+   * @return An address record
+   *
+   * @note May return an empty record indicating a retrieval failure.
+   * The failure must be resolved somewhere else within the allocation
+   * framework.
+   */
+  AR_T retrieve_record_from_bin(BINS_IT_T it) {
+    try_to_ensure_bin_has_records(it);
+
+    auto& [IGNORE, bin]{*it};
+
+    if (bin.empty()) {
+      return AR_T{};
+    }
+
+    return bin.get();
+  }
+
+  /**
+   * @brief Attempt to accommodate a subsequent record request
+   *
+   * The bin under iterator may not (or may) contain records when
+   * this method is invoked. This method tries to make at least one
+   * record available.
+   *
+   * @param[in] a VALID FORWARD iterator containing a bin
+   *
+   * @note BINS_T container must be sorted in ascending bin size for
+   * delegate method "split_larger_bin"
+   */
+  void try_to_ensure_bin_has_records(BINS_IT_T it) {
+    auto& [IGNORE, bin]{*it};
+
+    if (!bin.empty()) {
+      return;
+    }
+
+    split_larger_bin(it);
+  }
+
+  /**
+   * @brief Split records in next largest bin
+   *
+   * This method tries to create records within this bin by decomposing
+   * records in other bins (with larger bin sizes).
+   *
+   * This method may be called recursively if records are unavailable
+   * at each step until records are made available to the initial
+   * bin.
+   *
+   * @param[in] a VALID FORWARD iterator containing a bin
+   *
+   * @note BINS_T container must be sorted in ascending bin size
+   */
+  void split_larger_bin(BINS_IT_T it) {
+    assert(it != bins_.end());
+
+    auto next{std::next(it)};
+    if (next == bins_.end()) {
+      return;
+    }
+
+    auto& [IGNORE_NEXT, next_bin]{*next};
+    if (next_bin.empty()) {
+      split_larger_bin(next);
+    }
+
+    if (next_bin.empty()) {
+      return;
+    }
+
+    auto record{next_bin.get()};
+    auto [smaller, larger]{record.split()};
+
+    auto& [IGNORE, bin]{*it};
+    bin.put(larger);
+    bin.put(smaller);
+  }
+
+  /**
+   * @brief Store record in proffered_ data structure
+   *
+   * @param[in] An address record
+   */
+  void emplace_in_proffered(AR_T record) {
+    assert(record.get_address());
+    proffered_[record.get_address()] = record;
+  }
+
+  /**
+   * @brief Retrieve record from proffered_ data structure
+   *
+   * @param[in] raw memory pointer
+   *
+   * The memory pointer must match a pointer previously
+   * handed to the user.
+   *
+   * @return An address record
+   */
+  AR_T retrieve_from_proffered(char* ptr) {
+    assert(ptr);
+
+    auto prof_it{proffered_.find(ptr)};
+    assert(prof_it != proffered_.end());
+
+    auto [IGNORE_RECORD_ADDR, record]{*prof_it};
+    proffered_.erase(prof_it);
+
+    return record;
+  }
+
+  /**
+   * @brief Holds a SORTED set of bin objects
+   *
+   * The bin objects are sorted by ascending bin sizes. Each bin holds
+   * address records with matching sizes. The bin sizes increase by
+   * powers-of-two. There is a minimum bin size set by memory alignment
+   * constraints.
+   */
+  BINS_T bins_{};
+
+  /**
+   * @brief Initializes the bins_ object
+   *
+   * The binner functions as a bins_ builder. The binner access data in
+   * the heap memory to establish the number of bins and each bin's size.
+   *
+   * Additionally, binner_ will assign the heap memory to bins_.
+   */
+  Binner<AR_T, BINS_T> binner_{};
+
+  /**
+   * @brief Tracks all address records handed over to user.
+   *
+   * This member is required to track the sizes of memory allocations
+   * since the user is not required to provide a size_t field back
+   * through the "free" interface.
+   */
+  PROFFERED_T proffered_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_POW2_BINS_HPP_

src/memory/remote_heap_info.hpp

@@ -0,0 +1,223 @@
+/******************************************************************************
+ * 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_MEMORY_REMOTE_HEAP_INFO_HPP_
+#define LIBRARY_SRC_MEMORY_REMOTE_HEAP_INFO_HPP_
+
+#include <hip/hip_runtime_api.h>
+#include <mpi.h>
+
+#include <cassert>
+#include <vector>
+
+#include "hip_allocator.hpp"
+#include "window_info.hpp"
+
+/**
+ * @file remote_heap_info.hpp
+ *
+ * @brief Contains information about remote processing elements' heaps
+ */
+
+namespace rocshmem {
+
+class CommunicatorMPI {
+ public:
+  /**
+   * @brief Default constructor
+   */
+  CommunicatorMPI() = default;
+
+  /**
+   * @brief Primary constructor
+   */
+  CommunicatorMPI(char* heap_base, size_t heap_size) {
+    int initialized;
+    MPI_Initialized(&initialized);
+    if (!initialized) {
+      int provided;
+      MPI_Init_thread(nullptr, nullptr, MPI_THREAD_SINGLE, &provided);
+    }
+    MPI_Comm_rank(comm_, &my_pe_);
+    MPI_Comm_size(comm_, &num_pes_);
+
+    heap_window_info_ = WindowInfo(comm_, heap_base, heap_size);
+  }
+
+  /**
+   * @brief Destructor
+   */
+  ~CommunicatorMPI() {}
+
+  /**
+   * @brief Returns my processing element ID
+   */
+  int my_pe() { return my_pe_; }
+
+  /**
+   * @brief Returns number of processing elements
+   */
+  int num_pes() { return num_pes_; }
+
+  /**
+   * @brief Performs MPI_Barrier
+   */
+  void barrier() { MPI_Barrier(comm_); }
+
+  /**
+   * @brief Performs MPI_Allgather on recvbuf
+   */
+  void allgather(void* recvbuf) {
+    MPI_Allgather(MPI_IN_PLACE, sizeof(void*), MPI_CHAR, recvbuf, sizeof(void*),
+                  MPI_CHAR, comm_);
+  }
+
+  /**
+   * @brief Accessor method for heap_window_info_
+   */
+  WindowInfo* get_window_info() { return &heap_window_info_; }
+
+ private:
+  /**
+   * @brief Identifier for this processing element
+   */
+  MPI_Comm comm_{MPI_COMM_WORLD};
+
+  /**
+   * @brief Identifier for this processing element
+   */
+  int my_pe_{-1};
+
+  /**
+   * @brief The total number of processing elements
+   */
+  int num_pes_{-1};
+
+  /**
+   * @brief MPI window on the symmetric GPU heap
+   */
+  WindowInfo heap_window_info_{};
+};
+
+template <typename COMMUNICATOR_T>
+class RemoteHeapInfo {
+  /**
+   * @brief Helper type for heap_bases_ member
+   */
+  using HEAP_BASES_T = std::vector<char*, StdAllocatorHIP<char*>>;
+
+ public:
+  /**
+   * @brief Required for default construction of other objects
+   *
+   * @note Not intended for direct usage
+   */
+  RemoteHeapInfo() = default;
+
+  /**
+   * @brief Primary constructor
+   *
+   * @param[in] The identifier for this processing element
+   * @param[in] The total number of processing elements
+   */
+  RemoteHeapInfo(char* heap_ptr, size_t heap_size)
+      : communicator_{heap_ptr, heap_size} {
+    heap_bases_.resize(communicator_.num_pes());
+    for (auto& base : heap_bases_) {
+      base = nullptr;
+    }
+    heap_bases_[communicator_.my_pe()] = heap_ptr;
+    communicator_.allgather(heap_bases_.data());
+
+    device_heap_bases_ = heap_bases_.data();
+  }
+
+  /**
+   * @brief Invoke barrier on communicator
+   */
+  void barrier() { communicator_.barrier(); }
+
+  /**
+   * @brief Accessor for my processing element identifier
+   *
+   * @return My processing element identifier
+   */
+  int my_pe() { return communicator_.my_pe(); }
+
+  /**
+   * @brief Accessor for number of processing elements
+   *
+   * @return Number of processing elements
+   */
+  int num_pes() { return communicator_.num_pes(); }
+
+  /**
+   * @brief Accessor method for heap_window_info_
+   *
+   * @return Window info for Reverse Offload interface
+   */
+  auto get_window_info() { return communicator_.get_window_info(); }
+
+  /**
+   * @brief Accessor for heap bases
+   *
+   * @return Vector containing the addresses of the symmetric heap bases
+   */
+  __host__ const HEAP_BASES_T& get_heap_bases() { return heap_bases_; }
+
+  /**
+   * @brief Accessor for heap bases
+   *
+   * @return Vector containing the addresses of the symmetric heap bases
+   */
+  __device__ auto get_heap_bases() { return device_heap_bases_; }
+
+ private:
+  /**
+   * @brief Communicator implementation
+   */
+  COMMUNICATOR_T communicator_{};
+
+  /**
+   * @brief C-array of symmetric heap base pointers
+   *
+   * A vector of char* pointers corresponding to the heap_bases virtual
+   * address for all processing element that we can communicate with.
+   *
+   * @note Used by __host__ code
+   */
+  HEAP_BASES_T heap_bases_{};
+
+  /**
+   * @brief C-array of symmetric heap base pointers
+   *
+   * A vector of char* pointers corresponding to the heap_bases virtual
+   * address for all processing element that we can communicate with.
+   *
+   * @note Used by __device__ code
+   */
+  char** device_heap_bases_{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_REMOTE_HEAP_INFO_HPP_

src/memory/shmem_allocator_strategy.hpp

@@ -0,0 +1,90 @@
+/******************************************************************************
+ * 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_MEMORY_SHMEM_ALLOCATOR_STRATEGY_HPP_
+#define LIBRARY_SRC_MEMORY_SHMEM_ALLOCATOR_STRATEGY_HPP_
+
+/**
+ * @file shmem_allocator_strategy.hpp
+ *
+ * @brief Strategy design pattern for SHMEM symmetric heap allocator.
+ */
+
+namespace rocshmem {
+
+class ShmemAllocatorStrategy {
+ public:
+  /**
+   * @brief Default constructor
+   */
+  ShmemAllocatorStrategy() = default;
+
+  /**
+   * @brief Default destructor
+   */
+  virtual ~ShmemAllocatorStrategy() = default;
+
+  /**
+   * @brief Default copy construction
+   */
+  ShmemAllocatorStrategy(const ShmemAllocatorStrategy& other) = default;
+
+  /**
+   * @brief Disable copy assignment
+   */
+  ShmemAllocatorStrategy& operator=(const ShmemAllocatorStrategy& other) =
+      default;
+
+  /**
+   * @brief Allocates memory from the symmetric heap
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   */
+  virtual void alloc(char** ptr, size_t request_size) = 0;
+
+  /**
+   * @brief Allocates memory from the symmetric heap
+   *
+   * @param[in, out] Address of raw pointer (&pointer_to_char)
+   * @param[in] Size in bytes of memory allocation
+   */
+  __device__ virtual void alloc(char** ptr, size_t request_size) = 0;
+
+  /**
+   * @brief Frees memory from the symmetric heap
+   *
+   * @param[in] Raw pointer to symmetric heap memory
+   */
+  virtual void free(char* ptr) = 0;
+
+  /**
+   * @brief Frees memory from the symmetric heap
+   *
+   * @param[in] Raw pointer to symmetric heap memory
+   */
+  __device__ virtual void free(char* ptr) = 0;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_SHMEM_ALLOCATOR_STRATEGY_HPP_

src/memory/single_heap.cpp

@@ -0,0 +1,66 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "single_heap.hpp"
+
+#include <sstream>
+
+namespace rocshmem {
+
+SingleHeap::SingleHeap() {
+  if (auto heap_size_cstr = getenv("ROC_SHMEM_HEAP_SIZE")) {
+    std::stringstream sstream(heap_size_cstr);
+    size_t heap_size;
+    sstream >> heap_size;
+    heap_mem_ = HEAP_T{heap_size};
+    strat_ = STRAT_T{&heap_mem_};
+  }
+}
+
+void SingleHeap::malloc(void** ptr, size_t size) {
+  strat_.alloc(reinterpret_cast<char**>(ptr), size);
+}
+
+__device__ void SingleHeap::malloc(void** ptr, size_t size) {}
+
+void SingleHeap::free(void* ptr) {
+  if (!ptr) {
+    return;
+  }
+  strat_.free(reinterpret_cast<char*>(ptr));
+}
+
+__device__ void SingleHeap::free(void* ptr) {}
+
+void* SingleHeap::realloc(void* ptr, size_t size) { return nullptr; }
+
+void* SingleHeap::malign(size_t alignment, size_t size) { return nullptr; }
+
+char* SingleHeap::get_base_ptr() { return heap_mem_.get_ptr(); }
+
+size_t SingleHeap::get_size() { return heap_mem_.get_size(); }
+
+size_t SingleHeap::get_used() { return strat_.amount_proffered(); }
+
+size_t SingleHeap::get_avail() { return get_size() - get_used(); }
+
+}  // namespace rocshmem

src/memory/single_heap.hpp

@@ -0,0 +1,162 @@
+/******************************************************************************
+ * 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_MEMORY_SINGLE_HEAP_HPP_
+#define LIBRARY_SRC_MEMORY_SINGLE_HEAP_HPP_
+
+#include "address_record.hpp"
+#include "heap_memory.hpp"
+#include "heap_type.hpp"
+#include "pow2_bins.hpp"
+
+/**
+ * @file single_heap.hpp
+ *
+ * @brief Contains a single heap
+ *
+ * The single heap implements local processing element allocations. The
+ * symmetric heap delegates allocations to this class.
+ */
+
+namespace rocshmem {
+
+class SingleHeap {
+  /**
+   * @brief Helper type for address records
+   */
+  using AR_T = AddressRecord;
+
+  /**
+   * @brief Helper type for allocation strategy
+   */
+  using STRAT_T = Pow2Bins<AR_T, HEAP_T>;
+
+ public:
+  /**
+   * @brief Primary constructor
+   */
+  SingleHeap();
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in,out] A pointer to memory handle
+   * @param[in] Size in bytes of memory allocation
+   */
+  void malloc(void** ptr, size_t size);
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in,out] A pointer to memory handle
+   * @param[in] Size in bytes of memory allocation
+   *
+   * @note Not implemented
+   */
+  __device__ void malloc(void** ptr, size_t size);
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * @param[in] Raw pointer to heap memory
+   */
+  void free(void* ptr);
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * @param[in] Raw pointer to heap memory
+   *
+   * @note Not implemented
+   */
+  __device__ void free(void* ptr);
+
+  /**
+   * @brief
+   *
+   * @param[in]
+   * @param[in]
+   *
+   * @return
+   */
+  void* realloc(void* ptr, size_t size);
+
+  /**
+   * @brief
+   *
+   * @param[in]
+   * @param[in]
+   *
+   * @return
+   */
+  void* malign(size_t alignment, size_t size);
+
+  /**
+   * @brief Accessor for heap base ptr
+   *
+   * @return Pointer to base of my heap
+   */
+  char* get_base_ptr();
+
+  /**
+   * @brief Accessor for heap size
+   *
+   * @return Amount of bytes in heap
+   */
+  size_t get_size();
+
+  /**
+   * @brief Accessor for heap usage
+   *
+   * @return Amount of used bytes in heap
+   */
+  size_t get_used();
+
+  /**
+   * @brief Accessor for heap available
+   *
+   * @return Amount of available bytes in heap
+   */
+  size_t get_avail();
+
+  /**
+   * @brief Returns is the heap is allocated with managed memory
+   *
+   * @return bool
+   */
+  bool is_managed() { return heap_mem_.is_managed(); }
+
+ private:
+  /**
+   * @brief Heap memory object
+   */
+  HEAP_T heap_mem_{};
+
+  /**
+   * @brief Allocation strategy object
+   */
+  STRAT_T strat_{&heap_mem_};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_SINGLE_HEAP_HPP_

src/memory/slab_heap.cpp

@@ -0,0 +1,113 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "slab_heap.hpp"
+
+#include <sstream>
+
+#include "../util.hpp"
+
+namespace rocshmem {
+
+SlabHeap::SlabHeap() {
+  if (auto slab_size_cstr = getenv("ROC_SHMEM_SLAB_SIZE")) {
+    std::stringstream sstream(slab_size_cstr);
+    size_t slab_size;
+    sstream >> slab_size;
+    heap_mem_ = HEAP_T{slab_size};
+    strat_ = STRAT_T{&heap_mem_};
+  }
+}
+
+void SlabHeap::malloc(void** ptr, size_t size) {
+  strat_.alloc(reinterpret_cast<char**>(ptr), size);
+}
+
+__device__ void SlabHeap::malloc(void** ptr, size_t size) {
+  /*
+   * Grab the mutex from the proxy object which owns it.
+   */
+  auto mutex{mutex_.get()};
+
+  /*
+   * Take the ticketed lock.
+   *
+   * The lock is held jointly by all threads in the block.
+   */
+  auto ticket{mutex->lock()};
+
+  /*
+   * Perform allocation and verify it worked.
+   *
+   * Allocation should only be run by only one thread in the
+   * strategy code.
+   */
+  char** ptr_c{reinterpret_cast<char**>(ptr)};
+  strat_.alloc(ptr_c, size);
+  __threadfence();
+
+  /*
+   * The notifier works with uint64_t for the address broadcasts
+   * between threads (as type erasure for the pointer arithmetic).
+   */
+  uint64_t ptr_deref_u64{reinterpret_cast<uint64_t>(*ptr)};
+
+  /*
+   * Notify other threads in block about the allocation result.
+   */
+  auto notifier{notifier_.get()};
+  notifier->write(ptr_deref_u64);
+  uint64_t notification_u64{notifier->read()};
+  notifier->done();
+
+  /*
+   * Write to the ptr parameter (to return it back up the call stack).
+   */
+  char* read_value_c{reinterpret_cast<char*>(notification_u64)};
+  *ptr_c = read_value_c;
+
+  /*
+   * Release the lock with our ticket number.
+   */
+  mutex->unlock(ticket);
+}
+
+__host__ __device__ void SlabHeap::free(void* ptr) {
+  if (!ptr) {
+    return;
+  }
+  strat_.free(reinterpret_cast<char*>(ptr));
+}
+
+void* SlabHeap::realloc(void* ptr, size_t size) { return nullptr; }
+
+void* SlabHeap::malign(size_t alignment, size_t size) { return nullptr; }
+
+char* SlabHeap::get_base_ptr() { return heap_mem_.get_ptr(); }
+
+size_t SlabHeap::get_size() { return heap_mem_.get_size(); }
+
+size_t SlabHeap::get_used() { return strat_.current() - get_base_ptr(); }
+
+size_t SlabHeap::get_avail() { return get_size() - get_used(); }
+
+}  // namespace rocshmem

src/memory/slab_heap.hpp

@@ -0,0 +1,185 @@
+/******************************************************************************
+ * 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_MEMORY_SLAB_HEAP_HPP_
+#define LIBRARY_SRC_MEMORY_SLAB_HEAP_HPP_
+
+#include "dev_mono_linear.hpp"
+#include "heap_memory.hpp"
+#include "heap_type.hpp"
+#include "notifier.hpp"
+#include "../sync/abql_block_mutex.hpp"
+
+/**
+ * @file slab_heap.hpp
+ *
+ * @brief Contains a heap used to allocate library objects
+ *
+ * The slab heap is used internally by the library
+ */
+
+namespace rocshmem {
+
+class SlabHeap {
+  /**
+   * @brief Helper type for allocation strategy
+   */
+  using STRAT_T = DevMonoLinear<HEAP_T>;
+
+  /**
+   * @brief Helper type for notifier
+   */
+  using NOTIFIER_PROXY_T = NotifierProxy<HIPAllocator>;
+
+  /**
+   * @brief Helper type for notifier
+   */
+  using MUTEX_PROXY_T = ABQLBlockMutexProxy<HIPAllocator>;
+
+ public:
+  /**
+   * @brief Primary constructor
+   */
+  SlabHeap();
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in,out] A pointer to memory handle
+   * @param[in] Size in bytes of memory allocation
+   */
+  void malloc(void** ptr, size_t size);
+
+  /**
+   * @brief Allocates memory from the heap
+   *
+   * @param[in,out] A pointer to memory handle
+   * @param[in] Size in bytes of memory allocation
+   */
+  __device__ void malloc(void** ptr, size_t size);
+
+  /**
+   * @brief Frees memory from the heap
+   *
+   * @param[in] Raw pointer to heap memory
+   */
+  __host__ __device__ void free(void* ptr);
+
+  /**
+   * @brief
+   *
+   * @param[in]
+   * @param[in]
+   *
+   * @return
+   */
+  void* realloc(void* ptr, size_t size);
+
+  /**
+   * @brief
+   *
+   * @param[in]
+   * @param[in]
+   *
+   * @return
+   */
+  void* malign(size_t alignment, size_t size);
+
+  /**
+   * @brief Accessor for heap base ptr
+   *
+   * @return Pointer to base of my heap
+   */
+  char* get_base_ptr();
+
+  /**
+   * @brief Accessor for heap size
+   *
+   * @return Amount of bytes in heap
+   */
+  size_t get_size();
+
+  /**
+   * @brief Accessor for heap usage
+   *
+   * @return Amount of used bytes in heap
+   */
+  size_t get_used();
+
+  /**
+   * @brief Accessor for heap available
+   *
+   * @return Amount of available bytes in heap
+   */
+  size_t get_avail();
+
+ private:
+  /**
+   * @brief Heap memory object
+   */
+  HEAP_T heap_mem_{};
+
+  /**
+   * @brief Allocation strategy object
+   */
+  STRAT_T strat_{&heap_mem_};
+
+  /**
+   * @brief Notifier proxy to share information between threads.
+   *
+   * Need this object to share allocation information between the
+   * leader thread (that does allocation) and the follower threads
+   * (who need the allocation address).
+   */
+  NOTIFIER_PROXY_T notifier_{};
+
+  /**
+   * @brief Mutex to access the heap mutator methods.
+   */
+  MUTEX_PROXY_T mutex_;
+};
+
+template <typename ALLOCATOR>
+class SlabHeapProxy {
+  using ProxyT = DeviceProxy<ALLOCATOR, SlabHeap, 1>;
+
+ public:
+  /*
+   * Placement new the memory which is allocated by proxy_
+   */
+  SlabHeapProxy() { new (proxy_.get()) SlabHeap(); }
+
+  /*
+   * Since placement new is called in the constructor, then
+   * delete must be called manually.
+   */
+  ~SlabHeapProxy() { proxy_.get()->~SlabHeap(); }
+
+  __host__ __device__ SlabHeap* get() { return proxy_.get(); }
+
+ private:
+  ProxyT proxy_{};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_SLAB_HEAP_HPP_

src/memory/symmetric_heap.hpp

@@ -0,0 +1,129 @@
+/******************************************************************************
+ * 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_MEMORY_SYMMETRIC_HEAP_HPP_
+#define LIBRARY_SRC_MEMORY_SYMMETRIC_HEAP_HPP_
+
+/**
+ * @file symmetric_heap.hpp
+ *
+ * @brief Contains a symmetric heap
+ *
+ * Every processing element allocates a symmetric heap. (For details on
+ * the symmetric heap, refer to the OpenSHMEM specification.)
+ *
+ * The symmetric heap is allocated using fine-grained memory to allow
+ * both host access and device access to the memory space.
+ *
+ * The symmetric heaps are visible to network by registering them as
+ * InfiniBand memory regions. Every memory region has a remote key
+ * which needs to be shared across the network (to access the memory
+ * region).
+ */
+
+#include <hip/hip_runtime_api.h>
+
+#include "remote_heap_info.hpp"
+#include "single_heap.hpp"
+
+namespace rocshmem {
+
+class SymmetricHeap {
+  /**
+   * @brief Helper type for RemoteHeapInfo with MPI
+   */
+  using RemoteHeapInfoType = RemoteHeapInfo<CommunicatorMPI>;
+
+ public:
+  /**
+   * @brief Allocates heap memory and returns ptr to caller
+   *
+   * @param[in,out] A pointer to memory handle
+   * @param[in] Number of bytes of requested
+   */
+  void malloc(void** ptr, size_t size) { single_heap_.malloc(ptr, size); }
+
+  /**
+   * @brief Frees previously allocated network visible memory
+   *
+   * @param[in] Handle of previously allocated memory
+   */
+  void free(void* ptr) { single_heap_.free(ptr); }
+
+  /**
+   * @brief Accessor for local heap base
+   *
+   * @return Base address of the local symmetric heap
+   */
+  __host__ char* get_local_heap_base() { return single_heap_.get_base_ptr(); }
+
+  /**
+   * @brief Accessor method for heap size
+   */
+  auto get_size() { return single_heap_.get_size(); }
+
+  /**
+   * @brief Accessor method for heap_window_info_
+   */
+  auto get_window_info() { return remote_heap_info_.get_window_info(); }
+
+  /**
+   * @brief Accessor for heap bases
+   *
+   * @return Vector containing the addresses of the symmetric heap bases
+   */
+  __host__ const auto& get_heap_bases() {
+    return remote_heap_info_.get_heap_bases();
+  }
+
+  /**
+   * @brief Accessor for heap bases
+   *
+   * @return Vector containing the addresses of the symmetric heap bases
+   */
+  __device__ auto get_heap_bases() {
+    return remote_heap_info_.get_heap_bases();
+  }
+
+  /**
+   * @brief Returns is the heap is allocated with managed memory
+   *
+   * @return bool
+   */
+  bool is_managed() { return single_heap_.is_managed(); }
+
+ private:
+  /**
+   * @brief Processing element's implementation of heap
+   */
+  SingleHeap single_heap_{};
+
+  /**
+   * @brief Implementation of remote heaps
+   */
+  RemoteHeapInfoType remote_heap_info_{single_heap_.get_base_ptr(),
+                                       single_heap_.get_size()};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_SYMMETRIC_HEAP_HPP_

src/memory/window_info.hpp

@@ -0,0 +1,194 @@
+/******************************************************************************
+ * 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_MEMORY_WINDOW_INFO_HPP_
+#define LIBRARY_SRC_MEMORY_WINDOW_INFO_HPP_
+
+#include <mpi.h>
+
+#include <cassert>
+#include <memory>
+
+/**
+ * @file window_info.hpp
+ *
+ * @brief Contains information about symmetric heaps' windows
+ */
+
+namespace rocshmem {
+
+class WindowInfo {
+ public:
+  /**
+   * @brief Default constructor
+   */
+  WindowInfo() = default;
+
+  /**
+   * @brief Primary constructor
+   */
+  WindowInfo(MPI_Comm comm, void* start, size_t size)
+      : comm_{comm},
+        win_start_{start},
+        win_end_{reinterpret_cast<char*>(start) + size} {
+    up_win_ = std::unique_ptr<MPI_Win>(new MPI_Win);
+    MPI_Win_create(win_start_, size, 1, MPI_INFO_NULL, comm_, up_win_.get());
+    MPI_Win_lock_all(MPI_MODE_NOCHECK, *up_win_.get());
+  }
+
+  /**
+   * @brief Destructor
+   */
+  ~WindowInfo() {
+    if (up_win_) {
+      MPI_Win_unlock_all(*up_win_.get());
+      MPI_Win_free(up_win_.get());
+    }
+  }
+
+  /**
+   * @brief Copy constructor
+   *
+   * @note Disabled due to up_win_
+   */
+  WindowInfo(WindowInfo& other) = delete;  // NOLINT
+
+  /**
+   * @brief Const copy constructor
+   *
+   * @note Disabled due to up_win_
+   */
+  WindowInfo(const WindowInfo& other) = delete;
+
+  /**
+   * @brief Copy assignment
+   *
+   * @note Disabled due to up_win_
+   */
+  WindowInfo& operator=(WindowInfo other) = delete;
+
+  /**
+   * @brief Move constructor
+   */
+  WindowInfo(WindowInfo&& other) = default;
+
+  /**
+   * @brief Move assignment
+   */
+  WindowInfo& operator=(WindowInfo&& other) = default;
+
+  /**
+   * @brief Accessor for object in up_win_
+   *
+   * @return MPI_Win object
+   */
+  MPI_Win get_win() const { return *up_win_.get(); }
+
+  /**
+   * @brief Accessor for win_start_
+   *
+   * @return Raw start pointer
+   */
+  void* get_start() const { return win_start_; }
+
+  /**
+   * @brief Accessor for win_end_
+   *
+   * @return Raw end pointer
+   */
+  void* get_end() const { return win_end_; }
+
+  /**
+   * @brief Setter for object in up_win_
+   *
+   * @param[in] An MPI Window object
+   */
+  void set_win(MPI_Win win) { *up_win_.get() = win; }
+
+  /**
+   * @brief Setter for win_start_
+   *
+   * @param[in] Start raw pointer
+   */
+  void set_start(void* start) { win_start_ = start; }
+
+  /**
+   * @brief Setter for win_end_
+   *
+   * @param[in] End raw pointer
+   */
+  void set_end(void* end) { win_end_ = end; }
+
+  /**
+   * @brief Get offset between address and start of window
+   *
+   * @param[in] Address in raw pointer format
+   *
+   * @return Difference between dest and window start
+   */
+  MPI_Aint get_offset(const void* dest) {
+    assert(reinterpret_cast<char*>(const_cast<void*>(dest)) >=
+           reinterpret_cast<char*>(win_start_));
+    assert(reinterpret_cast<char*>(const_cast<void*>(dest)) >=
+           reinterpret_cast<char*>(win_start_));
+    assert(reinterpret_cast<char*>(const_cast<void*>(dest)) <
+           reinterpret_cast<char*>(win_end_));
+
+    MPI_Aint dest_disp;
+    MPI_Get_address(dest, &dest_disp);
+    MPI_Aint start_disp;
+    MPI_Get_address(win_start_, &start_disp);
+
+    return MPI_Aint_diff(dest_disp, start_disp);
+  }
+
+ private:
+  /**
+   * @brief MPI Communicator
+   */
+  MPI_Comm comm_{MPI_COMM_WORLD};
+
+  /**
+   * @brief Owning pointer to MPI_Win
+   *
+   * The pointer is used to track which object is responsible for
+   * releasing window resources during class destruction.
+   *
+   * This becomes an issue if objects of this class are move constructed
+   * or move assigned.
+   */
+  std::unique_ptr<MPI_Win> up_win_{nullptr};
+
+  /**
+   * @brief Raw pointer marking the start of window
+   */
+  void* win_start_{nullptr};
+
+  /**
+   * @brief Raw pointer marking the end of window
+   */
+  void* win_end_{nullptr};
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MEMORY_WINDOW_INFO_HPP_

src/mpi_init_singleton.cpp

@@ -0,0 +1,61 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "mpi_init_singleton.hpp"
+
+namespace rocshmem {
+
+MPIInitSingleton* MPIInitSingleton::instance{nullptr};
+
+MPIInitSingleton::MPIInitSingleton() {
+  MPI_Initialized(&pre_init_done);
+
+  if (!pre_init_done) {
+    int provided;
+    MPI_Init_thread(nullptr, nullptr, MPI_THREAD_SINGLE, &provided);
+  }
+
+  MPI_Comm_size(MPI_COMM_WORLD, &nprocs_);
+  MPI_Comm_rank(MPI_COMM_WORLD, &my_rank_);
+}
+
+MPIInitSingleton::~MPIInitSingleton() {
+  int finalized{0};
+  MPI_Finalized(&finalized);
+  if (!finalized && !pre_init_done) {
+    MPI_Finalize();
+  }
+}
+
+MPIInitSingleton* MPIInitSingleton::GetInstance() {
+  if (!instance) {
+    instance = new MPIInitSingleton();
+    return instance;
+  }
+  return instance;
+}
+
+int MPIInitSingleton::get_rank() { return my_rank_; }
+
+int MPIInitSingleton::get_nprocs() { return nprocs_; }
+
+}  // namespace rocshmem

src/mpi_init_singleton.hpp

@@ -0,0 +1,96 @@
+/******************************************************************************
+ * 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_MPI_INIT_SINGLETON_HPP_
+#define LIBRARY_SRC_MPI_INIT_SINGLETON_HPP_
+
+#include <mpi.h>
+
+#include <memory>
+
+/**
+ * @file mpi_init_singleton.hpp
+ *
+ * @brief Contains MPI library initialization code
+ */
+
+namespace rocshmem {
+
+class MPIInitSingleton {
+ private:
+  /**
+   * @brief Primary constructor
+   */
+  MPIInitSingleton();
+
+ public:
+  /**
+   * @brief Destructor
+   */
+  ~MPIInitSingleton();
+
+  /**
+   * @brief Invoke singleton construction or return handle
+   *
+   * @return Initialized handle to singleton
+   */
+  static MPIInitSingleton* GetInstance();
+
+  /**
+   * @brief Accessor for my COMM_WORLD rank identifier
+   *
+   * @return My COMM_WORLD rank identifier
+   */
+  int get_rank();
+
+  /**
+   * @brief Accessor for number or processes in COMM_WORLD
+   *
+   * @return Number of processes in COMM_WORLD
+   */
+  int get_nprocs();
+
+ private:
+  /**
+   * @brief My MPI rank identifier
+   */
+  int my_rank_{-1};
+
+  /**
+   * @brief Number of MPI processes
+   */
+  int nprocs_{-1};
+
+  /**
+   * @brief Was MPI initialized before ROCSHMEM_init call
+   */
+  int pre_init_done{0};
+
+  /**
+   * @brief Refers to global variable
+   */
+  static MPIInitSingleton* instance;
+};
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_MPI_INIT_SINGLETON_HPP_

src/reverse_offload/CMakeLists.txt

@@ -0,0 +1,35 @@
+###############################################################################
+# 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.
+###############################################################################
+
+###############################################################################
+# ADD ROCSHMEM TARGET FOR FILES IN CURRENT DIRECTORY
+###############################################################################
+target_sources(
+  ${PROJECT_NAME}
+  PRIVATE
+    backend_ro.cpp
+    context_ro_device.cpp
+    context_ro_host.cpp
+    mpi_transport.cpp
+    queue.cpp
+    ro_net_team.cpp
+)

src/reverse_offload/backend_proxy.hpp

@@ -0,0 +1,78 @@
+/******************************************************************************
+ * 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_REVERSE_OFFLOAD_BACKEND_PROXY_HPP_
+#define LIBRARY_SRC_REVERSE_OFFLOAD_BACKEND_PROXY_HPP_
+
+#include <atomic>
+
+#include "../device_proxy.hpp"
+#include "../stats.hpp"
+#include "queue.hpp"
+
+namespace rocshmem {
+
+struct BackendRegister {
+  ROStats *profiler{nullptr};
+  std::atomic<bool> worker_thread_exit{false};
+  bool *needs_quiet{nullptr};
+  bool *needs_blocking{nullptr};
+  char *g_ret{nullptr};
+  HdpPolicy *hdp_policy{nullptr};
+  WindowInfo **heap_window_info{nullptr};
+  atomic_ret_t *atomic_ret{nullptr};
+  SymmetricHeap *heap_ptr{nullptr};
+};
+
+template <typename ALLOCATOR>
+class BackendProxy {
+  using ProxyT = DeviceProxy<ALLOCATOR, BackendRegister>;
+
+ public:
+  /*
+   * Placement new the memory which is allocated by proxy_
+   */
+  BackendProxy() { new (proxy_.get()) BackendRegister(); }
+
+  /*
+   * Since placement new is called in the constructor, then
+   * delete must be called manually.
+   */
+  ~BackendProxy() { proxy_.get()->~BackendRegister(); }
+
+  /*
+   * @brief Provide access to the memory referenced by the proxy
+   */
+  __host__ __device__ BackendRegister *get() { return proxy_.get(); }
+
+ private:
+  /*
+   * @brief Memory managed by the lifetime of this object
+   */
+  ProxyT proxy_{};
+};
+
+using BackendProxyT = BackendProxy<HIPHostAllocator>;
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_REVERSE_OFFLOAD_BACKEND_PROXY_HPP_

src/reverse_offload/backend_ro.cpp

@@ -0,0 +1,288 @@
+/******************************************************************************
+ * 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.
+ *****************************************************************************/
+
+#include "backend_ro.hpp"
+
+#include <immintrin.h>
+#include <smmintrin.h>
+#include <unistd.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <iostream>
+#include <memory>
+#include <thread>  // NOLINT
+
+#include "roc_shmem/roc_shmem.hpp"
+#include "../atomic_return.hpp"
+#include "../backend_type.hpp"
+#include "../context_incl.hpp"
+#include "mpi_transport.hpp"
+#include "ro_net_team.hpp"
+#include "../util.hpp"
+
+namespace rocshmem {
+
+extern roc_shmem_ctx_t ROC_SHMEM_HOST_CTX_DEFAULT;
+
+ROBackend::ROBackend(MPI_Comm comm)
+    : profiler_proxy_(MAX_NUM_BLOCKS), Backend() {
+  type = BackendType::RO_BACKEND;
+
+  if (auto maximum_num_contexts_str = getenv("ROC_SHMEM_MAX_NUM_CONTEXTS")) {
+    std::stringstream sstream(maximum_num_contexts_str);
+    sstream >> maximum_num_contexts_;
+  }
+  poll_block_count_ = maximum_num_contexts_;
+
+  transport_ = new MPITransport(comm, &queue_);
+  num_pes = transport_->getNumPes();
+  my_pe = transport_->getMyPe();
+
+  auto *bp{backend_proxy.get()};
+
+  bp->hdp_policy = hdp_proxy_.get();
+
+  bp->profiler = profiler_proxy_.get();
+
+  bp->worker_thread_exit = false;
+
+  bp->heap_ptr = &heap;
+
+  ro_window_proxy_ = new WindowProxyT(&heap, transport_->get_world_comm());
+  bp->heap_window_info = ro_window_proxy_->get();
+
+  initIPC();
+
+  init_g_ret(&heap, transport_->get_world_comm(), MAX_NUM_BLOCKS, &bp->g_ret);
+
+  allocate_atomic_region(&bp->atomic_ret, MAX_NUM_BLOCKS);
+
+  transport_->initTransport(MAX_NUM_BLOCKS, &backend_proxy);
+
+  host_interface = transport_->host_interface;
+
+  default_host_ctx = std::make_unique<ROHostContext>(this, 0);
+
+  ROC_SHMEM_HOST_CTX_DEFAULT.ctx_opaque = default_host_ctx.get();
+
+  team_world_proxy_ = new ROTeamProxy<HIPAllocator>(
+      this, transport_->get_world_comm(), my_pe, num_pes);
+  team_tracker.set_team_world(team_world_proxy_->get());
+
+  ROC_SHMEM_TEAM_WORLD =
+      reinterpret_cast<roc_shmem_team_t>(team_world_proxy_->get());
+
+  default_block_handle_proxy_ = DefaultBlockHandleProxyT(
+      bp->g_ret, bp->atomic_ret, &queue_, &ipcImpl, hdp_proxy_.get());
+  default_context_proxy_ = DefaultContextProxyT(this);
+
+  block_handle_proxy_ = BlockHandleProxyT(bp->g_ret, bp->atomic_ret, &queue_,
+                                          &ipcImpl, hdp_proxy_.get());
+  setup_ctxs();
+
+  worker_thread = std::thread(&ROBackend::ro_net_poll, this);
+
+  *done_init = 1;
+}
+
+void ROBackend::setup_ctxs() {
+  CHECK_HIP(hipMalloc(&ctx_array, sizeof(ROContext) * maximum_num_contexts_));
+  for (int i = 0; i < maximum_num_contexts_; i++) {
+    new (&ctx_array[i]) ROContext(this, i);
+    ctx_free_list.get()->push_back(ctx_array + i);
+  }
+}
+
+ROBackend::~ROBackend() {
+  ro_net_free_runtime();
+  CHECK_HIP(hipFree(ctx_array));
+}
+
+__device__ bool ROBackend::create_ctx(int64_t options, roc_shmem_ctx_t *ctx) {
+  ROContext *ctx_;
+
+  auto pop_result = ctx_free_list.get()->pop_front();
+  if (!pop_result.success) {
+    return false;
+  }
+  ctx_ = pop_result.value;
+
+  ctx->ctx_opaque = ctx_;
+  return true;
+}
+
+__device__ void ROBackend::destroy_ctx(roc_shmem_ctx_t *ctx) {
+  ctx_free_list.get()->push_back(static_cast<ROContext *>(ctx->ctx_opaque));
+}
+
+void ROBackend::team_destroy(roc_shmem_team_t team) {
+  ROTeam *team_obj{get_internal_ro_team(team)};
+
+  team_obj->~ROTeam();
+  // CHECK_HIP(hipFree(team_obj));
+}
+
+void ROBackend::create_new_team(Team *parent_team,
+                                TeamInfo *team_info_wrt_parent,
+                                TeamInfo *team_info_wrt_world, int num_pes,
+                                int my_pe_in_new_team, MPI_Comm team_comm,
+                                roc_shmem_team_t *new_team) {
+  transport_->createNewTeam(this, parent_team, team_info_wrt_parent,
+                            team_info_wrt_world, num_pes, my_pe_in_new_team,
+                            team_comm, new_team);
+}
+
+void ROBackend::ctx_create(int64_t options, void **ctx) {
+  ROHostContext *new_ctx{nullptr};
+  new_ctx = new ROHostContext(this, options);
+  *ctx = new_ctx;
+}
+
+ROHostContext *get_internal_ro_net_ctx(Context *ctx) {
+  return reinterpret_cast<ROHostContext *>(ctx);
+}
+
+void ROBackend::ctx_destroy(Context *ctx) {
+  ROHostContext *ro_net_host_ctx{get_internal_ro_net_ctx(ctx)};
+  delete ro_net_host_ctx;
+}
+
+void ROBackend::reset_backend_stats() {
+  auto *bp{backend_proxy.get()};
+
+  for (size_t i{0}; i < MAX_NUM_BLOCKS; i++) {
+    bp->profiler[i].resetStats();
+  }
+}
+
+void ROBackend::dump_backend_stats() {
+  uint64_t total{0};
+  for (int i = 0; i < NUM_STATS; i++) {
+    total += globalStats.getStat(i);
+  }
+
+  uint64_t gpu_frequency_mhz{wallClk_freq_mhz()};
+
+  uint64_t us_wait_slot{0};
+  uint64_t us_pack{0};
+  uint64_t us_fence1{0};
+  uint64_t us_fence2{0};
+  uint64_t us_wait_host{0};
+
+  auto *bp{backend_proxy.get()};
+
+  for (size_t i{0}; i < MAX_NUM_BLOCKS; i++) {
+    // Average latency as perceived from a thread
+    const ROStats &prof{bp->profiler[i]};
+    us_wait_slot += prof.getStat(WAITING_ON_SLOT) / gpu_frequency_mhz;
+    us_pack += prof.getStat(PACK_QUEUE) / gpu_frequency_mhz;
+    us_fence1 += prof.getStat(THREAD_FENCE_1) / gpu_frequency_mhz;
+    us_fence2 += prof.getStat(THREAD_FENCE_2) / gpu_frequency_mhz;
+    us_wait_host += prof.getStat(WAITING_ON_HOST) / gpu_frequency_mhz;
+  }
+
+  constexpr int FIELD_WIDTH{20};
+  constexpr int FLOAT_PRECISION{2};
+
+  printf("%*s%*s%*s%*s%*s\n", FIELD_WIDTH + 1, "Wait On Slot (us)",
+         FIELD_WIDTH + 1, "Pack Queue (us)", FIELD_WIDTH + 1, "Fence 1 (us)",
+         FIELD_WIDTH + 1, "Fence 2 (us)", FIELD_WIDTH + 1, "Wait Host (us)");
+
+  printf("%*.*f %*.*f %*.*f %*.*f %*.*f\n\n", FIELD_WIDTH, FLOAT_PRECISION,
+         static_cast<double>(us_wait_slot) / total, FIELD_WIDTH,
+         FLOAT_PRECISION, static_cast<double>(us_pack) / total, FIELD_WIDTH,
+         FLOAT_PRECISION, static_cast<double>(us_fence1) / total, FIELD_WIDTH,
+         FLOAT_PRECISION, static_cast<double>(us_fence2) / total, FIELD_WIDTH,
+         FLOAT_PRECISION, static_cast<double>(us_wait_host) / total);
+}
+
+void ROBackend::ro_net_free_runtime() {
+  /*
+   * Validate that a handle was passed that is not a nullptr.
+   */
+  auto *bp{backend_proxy.get()};
+  assert(bp);
+
+  /*
+   * Set this flag to denote that the runtime is being torn down.
+   */
+  bp->worker_thread_exit = true;
+
+  /*
+   * Tear down the worker threads.
+   */
+  worker_thread.join();
+
+  /*
+   * Tear down the transport object.
+   */
+  while (!transport_->readyForFinalize()) {
+  }
+  transport_->finalizeTransport();
+
+  ro_window_proxy_->~WindowProxyT();
+  team_world_proxy_->~ROTeamProxy<HIPAllocator>();
+  transport_->~MPITransport();
+  /*
+   * Free the profiler statistics structure.
+   */
+  // CHECK_HIP(hipFree(bp->profiler));
+
+  /*
+   * Tear down team_world
+   */
+  auto *team_world{team_tracker.get_team_world()};
+  team_world->~Team();
+  // CHECK_HIP(hipFree(team_world));
+
+  /*
+   * Free the gpu_handle.
+   */
+  // CHECK_HIP(hipHostFree(bp));
+}
+
+void ROBackend::ro_net_poll() {
+  auto *bp{backend_proxy.get()};
+  while (!bp->worker_thread_exit) {
+    for (size_t i{0}; i < poll_block_count_; i++) {
+      int16_t request_count{0};
+      const int16_t max_count{64};
+      bool processed_req{true};
+      while (processed_req && (request_count < max_count)) {
+        processed_req = queue_.process(i, transport_);
+        request_count++;
+      }
+    }
+  }
+}
+
+void ROBackend::initIPC() {
+  const auto &heap_bases{heap.get_heap_bases()};
+
+  ipcImpl.ipcHostInit(transport_->getMyPe(), heap_bases,
+                      transport_->get_world_comm());
+}
+
+void ROBackend::global_exit(int status) { transport_->global_exit(status); }
+
+}  // namespace rocshmem