Implement default RO context (#64)

* Allocate default context buffers and initialize queue for management

- Allocated the status flag, g return, and atomic return buffers for
  the default context.
- Initialized `AtomicWFQueueProxy` instances to manage these buffers
  efficiently for concurrent access.

* Update `BlockHandle` with default context buffers

* Add default context flag and update buffer retrieval functions

- Added a flag to distinguish the default context from other contexts.
- Modified return buffer functionns and `get_status_flag` function to accommodate
  the default context

* Add default context primitive tests

-  get, put, get_nbi, put_nbi, g, and p APIs.

src/reverse_offload/backend_ro.cpp

@@ -64,6 +64,10 @@ ROBackend::ROBackend(MPI_Comm comm)
 
   max_wg_size_ = device_props.maxThreadsPerBlock;
 
+  wf_size_ = device_props.warpSize;
+
+  setup_default_ctx_buffers();
+
   size_t num_buff_elems = maximum_num_contexts_ * max_wg_size_;
 
   g_ret_buffer_ = RetBufferProxyT(num_buff_elems);
@@ -113,7 +117,9 @@ ROBackend::ROBackend(MPI_Comm comm)
   default_block_handle_proxy_ = DefaultBlockHandleProxyT(
                                 g_ret_buffer_.get(),
                                 atomic_ret_buffer_.get(), &queue_,
-                                status_.get());
+                                status_.get(), default_ctx_status_.get(),
+                                default_ctx_g_ret_buffer_.get(),
+                                default_ctx_atomic_ret_buffer_.get());
 
   TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world;
 
@@ -137,6 +143,37 @@ void ROBackend::setup_ctxs() {
   }
 }
 
+void ROBackend::setup_default_ctx_buffers() {
+  if (auto maximum_wf_buffers_str = getenv("ROCSHMEM_MAX_WF_BUFFERS")) {
+    std::stringstream sstream(maximum_wf_buffers_str);
+    sstream >> max_wavefront_buffers_;
+  }
+
+  size_t num_buff_elems = max_wavefront_buffers_ * wf_size_;
+
+  g_ret_buffer_default_ctx_ = RetBufferProxyT(num_buff_elems);
+
+  atomic_ret_buffer_default_ctx_ = RetBufferProxyT(num_buff_elems);
+
+  status_default_ctx_ = StatusProxyT(num_buff_elems);
+
+  default_ctx_status_.get()->allocate_queue(max_wavefront_buffers_);
+  default_ctx_g_ret_buffer_.get()->allocate_queue(max_wavefront_buffers_);
+  default_ctx_atomic_ret_buffer_.get()->allocate_queue(max_wavefront_buffers_);
+
+
+  char* status = status_default_ctx_.get();
+  uint64_t* g_ret_buf = g_ret_buffer_default_ctx_.get();
+  uint64_t* atomic_ret_buf = atomic_ret_buffer_default_ctx_.get();
+
+  for (int i{0}; i < max_wavefront_buffers_; i++) {
+    int offset {i * wf_size_};
+    default_ctx_status_.get()->push(status + offset);
+    default_ctx_g_ret_buffer_.get()->push(g_ret_buf + offset);
+    default_ctx_atomic_ret_buffer_.get()->push(atomic_ret_buf + offset);
+  }
+}
+
 ROBackend::~ROBackend() {
   ro_net_free_runtime();
   CHECK_HIP(hipFree(ctx_array));

src/reverse_offload/backend_ro.hpp

@@ -178,6 +178,11 @@ class ROBackend : public Backend {
    */
   void setup_ctxs();
 
+  /**
+   * @brief Allocation and initialization of buffers for default context.
+   */
+  void setup_default_ctx_buffers();
+
   /**
    * @brief Handle for the transport class object.
    *
@@ -253,6 +258,32 @@ class ROBackend : public Backend {
    */
   FreeListProxy<HIPAllocator, ROContext *> ctx_free_list{};
 
+  /**
+   * @brief AtomicWFQueue containing status flag buffers for default context
+   */
+  AtomicWFQueueProxy<HIPAllocator, volatile char*> default_ctx_status_{};
+
+  /**
+   * @brief AtomicWFQueue containing rocshmem_g return buffers for default
+   * context
+   */
+  AtomicWFQueueProxy<HIPAllocator, uint64_t*> default_ctx_g_ret_buffer_{};
+
+  /**
+   * @brief AtomicWFQueue containing rocshmem return buffers for default
+   * context
+   */
+  AtomicWFQueueProxy<HIPAllocator, uint64_t*> default_ctx_atomic_ret_buffer_{};
+
+  /**
+   * @brief Maximum number of wavefront buffer arrays supported in the default
+   * context.
+   *
+   * This value determines the size of the status flag, rocshmem_g return, and
+   * rocshmem atomic return buffers.
+   */
+  size_t max_wavefront_buffers_{1024};
+
   /**
    * @brief Holds maximum number of contexts used in library
    */
@@ -260,9 +291,14 @@ class ROBackend : public Backend {
 
   /**
    * @brief Holds maximum threads per work-group
-  */
+   */
   int max_wg_size_{};
 
+  /**
+   * @brief Holds wavefront size
+   */
+  int wf_size_{};
+
   /**
    * @brief Holds the queue size for each context
   */
@@ -277,11 +313,13 @@ class ROBackend : public Backend {
    * @brief Return buffer for rocshmem_g API
    */
   RetBufferProxyT g_ret_buffer_;
+  RetBufferProxyT g_ret_buffer_default_ctx_;
 
   /**
    * @brief Return buffer for rocshmem atomic return APIs
    */
   RetBufferProxyT atomic_ret_buffer_;
+  RetBufferProxyT atomic_ret_buffer_default_ctx_;
 
   /**
    * This buffer is used by the GPU to wait on a blocking operation. The initial
@@ -291,6 +329,7 @@ class ROBackend : public Backend {
    * a separate status variable for each work-item in a RO Context
    */
   StatusProxyT status_;
+  StatusProxyT status_default_ctx_;
 };
 
 }  // namespace rocshmem

src/reverse_offload/block_handle.hpp

@@ -23,6 +23,7 @@
 #ifndef LIBRARY_SRC_REVERSE_OFFLOAD_BLOCK_HANDLE_HPP_
 #define LIBRARY_SRC_REVERSE_OFFLOAD_BLOCK_HANDLE_HPP_
 
+#include "../containers/atomic_wf_queue_impl.hpp"
 #include "../hdp_policy.hpp"
 #include "../ipc_policy.hpp"
 #include "profiler.hpp"
@@ -30,6 +31,9 @@
 
 namespace rocshmem {
 
+using AWF_Queue_statusT = AtomicWFQueue<volatile char*, HIPAllocator>;
+using AWF_Queue_ret_buffT = AtomicWFQueue<uint64_t*, HIPAllocator>;
+
 struct BlockHandle {
   ROStats profiler{};
   queue_element_t *queue{nullptr};
@@ -41,6 +45,9 @@ struct BlockHandle {
   void *g_ret{nullptr};
   void *atomic_ret{nullptr};
   volatile uint64_t lock{};
+  AWF_Queue_statusT *default_ctx_status{nullptr};
+  AWF_Queue_ret_buffT *default_ctx_g_ret{nullptr};
+  AWF_Queue_ret_buffT *default_ctx_atomic_ret{nullptr};
 };
 
 template <typename ALLOCATOR>
@@ -51,7 +58,11 @@ class DefaultBlockHandleProxy {
   DefaultBlockHandleProxy() = default;
 
   DefaultBlockHandleProxy(void *g_ret, void *atomic_ret, Queue *queue,
-                          volatile char *status, size_t num_elems = 1)
+                          volatile char *status,
+                          AWF_Queue_statusT *default_ctx_status,
+                          AWF_Queue_ret_buffT *default_ctx_g_ret,
+                          AWF_Queue_ret_buffT *default_ctx_atomic_ret,
+                          size_t num_elems = 1)
     : proxy_{num_elems} {
 
     // TODO(bpotter): create a default queue for this queue descriptor
@@ -67,6 +78,9 @@ class DefaultBlockHandleProxy {
     block_handle->g_ret = g_ret;
     block_handle->atomic_ret = atomic_ret;
     block_handle->lock = 0;
+    block_handle->default_ctx_status = default_ctx_status;
+    block_handle->default_ctx_g_ret = default_ctx_g_ret;
+    block_handle->default_ctx_atomic_ret = default_ctx_atomic_ret;
   }
 
   DefaultBlockHandleProxy(const DefaultBlockHandleProxy& other) = delete;

src/reverse_offload/context_proxy.hpp

@@ -46,7 +46,7 @@ class DefaultContextProxy {
                                size_t num_elems = 1)
   : constructed_{true}, proxy_{num_elems} {
     auto ctx{proxy_.get()};
-    new (ctx) ROContext(reinterpret_cast<Backend*>(backend), -1);
+    new (ctx) ROContext(reinterpret_cast<Backend*>(backend), -1, true);
     rocshmem_ctx_t local{ctx, tinfo};
     set_internal_ctx(&local);
   }

src/reverse_offload/context_ro_device.cpp

@@ -41,10 +41,14 @@
 
 namespace rocshmem {
 
-__host__ ROContext::ROContext(Backend *b, size_t block_id)
-    : Context(b, false) {
+__host__ ROContext::ROContext(Backend *b,
+                              size_t block_id,
+                              bool default_ctx)
+    : Context(b, false),
+      is_default_ctx{default_ctx} {
+
   ROBackend *backend{static_cast<ROBackend *>(b)};
-  if (block_id == -1) {
+  if (is_default_ctx) {
     block_handle = backend->default_block_handle_proxy_.get();
   } else {
     auto block_base{backend->block_handle_proxy_.get()};
@@ -72,7 +76,8 @@ __device__ void ROContext::putmem(void *dest, const void *source, size_t nelems,
     }
     build_queue_element(RO_NET_PUT, dest, const_cast<void *>(source), nelems,
                         pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
-                        ro_net_win_id, block_handle, true, get_status_flag());
+                        ro_net_win_id, block_handle, true, get_status_flag(),
+                        is_default_ctx);
   }
 }
 
@@ -91,7 +96,8 @@ __device__ void ROContext::getmem(void *dest, const void *source, size_t nelems,
     }
     build_queue_element(RO_NET_GET, dest, const_cast<void *>(source), nelems,
                         pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
-                        ro_net_win_id, block_handle, true, get_status_flag());
+                        ro_net_win_id, block_handle, true, get_status_flag(),
+                        is_default_ctx);
   }
 }
 
@@ -136,20 +142,20 @@ __device__ void ROContext::getmem_nbi(void *dest, const void *source,
 __device__ void ROContext::fence() {
   build_queue_element(RO_NET_FENCE, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
                       nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle,
-                      true, get_status_flag());
+                      true, get_status_flag(), is_default_ctx);
 }
 
 __device__ void ROContext::fence(int pe) {
   // TODO(khamidou): need to check if per pe has any special handling
   build_queue_element(RO_NET_FENCE, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
                       nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle,
-                      true, get_status_flag());
+                      true, get_status_flag(), is_default_ctx);
 }
 
 __device__ void ROContext::quiet() {
   build_queue_element(RO_NET_QUIET, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
                       nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle,
-                      true, get_status_flag());
+                      true, get_status_flag(), is_default_ctx);
 }
 
 __device__ void *ROContext::shmem_ptr(const void *dest, int pe) {
@@ -167,7 +173,7 @@ __device__ void ROContext::barrier_all() {
   if (is_thread_zero_in_block()) {
     build_queue_element(RO_NET_BARRIER, nullptr, nullptr, 0, 0, 0, 0, 0,
                         nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id,
-                        block_handle, true, get_status_flag());
+                        block_handle, true, get_status_flag(), is_default_ctx);
   }
   __syncthreads();
 }
@@ -177,7 +183,7 @@ __device__ void ROContext::barrier(rocshmem_team_t team) {
   if (is_thread_zero_in_block()) {
     build_queue_element(RO_NET_BARRIER, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
                         nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle,
-                        true, get_status_flag());
+                        true, get_status_flag(), is_default_ctx);
   }
   __syncthreads();
 }
@@ -186,7 +192,7 @@ __device__ void ROContext::sync_all() {
   if (is_thread_zero_in_block()) {
     build_queue_element(RO_NET_SYNC, nullptr, nullptr, 0, 0, 0, 0, 0,
                         nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id,
-                        block_handle, true, get_status_flag());
+                        block_handle, true, get_status_flag(), is_default_ctx);
   }
   __syncthreads();
 }
@@ -196,7 +202,7 @@ __device__ void ROContext::sync(rocshmem_team_t team) {
   if (is_thread_zero_in_block()) {
     build_queue_element(RO_NET_SYNC, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
                         nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle,
-                        true, get_status_flag());
+                        true, get_status_flag(), is_default_ctx);
   }
   __syncthreads();
 }
@@ -209,7 +215,7 @@ __device__ void ROContext::ctx_destroy() {
 
     build_queue_element(RO_NET_FINALIZE, nullptr, nullptr, 0, 0, 0, 0, 0,
                         nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id,
-                        block_handle, true, get_status_flag());
+                        block_handle, true, get_status_flag(), is_default_ctx);
 
     int buffer_id = ro_net_win_id;
     backend->queue_.descriptor(buffer_id)->write_index = block_handle->write_index;
@@ -233,7 +239,8 @@ __device__ void ROContext::putmem_wg(void *dest, const void *source,
     if (is_thread_zero_in_block()) {
       build_queue_element(RO_NET_PUT, dest, const_cast<void *>(source), nelems,
                           pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
-                          ro_net_win_id, block_handle, true, get_status_flag());
+                          ro_net_win_id, block_handle, true, get_status_flag(),
+                          is_default_ctx);
     }
   }
   __syncthreads();
@@ -251,7 +258,8 @@ __device__ void ROContext::getmem_wg(void *dest, const void *source,
     if (is_thread_zero_in_block()) {
       build_queue_element(RO_NET_GET, dest, const_cast<void *>(source), nelems,
                           pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
-                          ro_net_win_id, block_handle, true, get_status_flag());
+                          ro_net_win_id, block_handle, true, get_status_flag(),
+                          is_default_ctx);
     }
   }
   __syncthreads();
@@ -305,7 +313,8 @@ __device__ void ROContext::putmem_wave(void *dest, const void *source,
     if (is_thread_zero_in_wave()) {
       build_queue_element(RO_NET_PUT, dest, const_cast<void *>(source), nelems,
                           pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
-                          ro_net_win_id, block_handle, true, get_status_flag());
+                          ro_net_win_id, block_handle, true, get_status_flag(),
+                          is_default_ctx);
     }
   }
 }
@@ -323,7 +332,8 @@ __device__ void ROContext::getmem_wave(void *dest, const void *source,
     if (is_thread_zero_in_wave()) {
       build_queue_element(RO_NET_GET, dest, const_cast<void *>(source), nelems,
                           pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
-                          ro_net_win_id, block_handle, true, get_status_flag());
+                          ro_net_win_id, block_handle, true, get_status_flag(),
+                          is_default_ctx);
     }
   }
 }
@@ -592,7 +602,7 @@ __device__ void build_queue_element(
     ro_net_cmds type, void *dst, void *src, size_t size, int pe,
     int logPE_stride, int PE_size, int PE_root, void *pWrk, long *pSync,
     MPI_Comm team_comm, int ro_net_win_id, BlockHandle *handle,
-    bool blocking, volatile char *status, ROCSHMEM_OP op,
+    bool blocking, volatile char *status, bool default_ctx, ROCSHMEM_OP op,
     ro_net_types datatype) {
   auto write_slot{next_write_slot(handle)};
   auto queue_element = &handle->queue[write_slot];
@@ -661,25 +671,49 @@ __device__ void build_queue_element(
     *(queue_element->status) = 0;
     __threadfence();
   }
+  if (default_ctx) {
+    handle->default_ctx_status->enqueue(status);
+  }
 }
 
 __device__ uint64_t *ROContext::get_atomic_ret_buf() {
-  uint64_t *atomic_base_ptr{
-    reinterpret_cast<uint64_t*>(block_handle->atomic_ret)};
-  int thread_id{get_flat_block_id()};
-  return &atomic_base_ptr[thread_id];
+  uint64_t *buf_addr{nullptr};
+  if(is_default_ctx) {
+    buf_addr = block_handle->default_ctx_atomic_ret->dequeue();
+  }
+  else {
+    uint64_t *atomic_base_ptr{
+      reinterpret_cast<uint64_t*>(block_handle->atomic_ret)};
+    int thread_id{get_flat_block_id()};
+    buf_addr = &atomic_base_ptr[thread_id];
+  }
+  return buf_addr;
 }
 
 __device__ uint64_t *ROContext::get_g_ret_buf() {
-  uint64_t *g_ret{reinterpret_cast<uint64_t*>(block_handle->g_ret)};
-  int thread_id{get_flat_block_id()};
-  return &g_ret[thread_id];
+  uint64_t *buf_addr{nullptr};
+  if(is_default_ctx) {
+    buf_addr = block_handle->default_ctx_g_ret->dequeue();
+  }
+  else {
+    uint64_t *g_ret{reinterpret_cast<uint64_t*>(block_handle->g_ret)};
+    int thread_id{get_flat_block_id()};
+    buf_addr = &g_ret[thread_id];
+  }
+  return buf_addr;
 }
 
 __device__ volatile char *ROContext::get_status_flag() {
-  volatile char* status{block_handle->status};
-  int thread_id{get_flat_block_id()};
-  return &status[thread_id];
+  volatile char *status_addr{nullptr};
+  if(is_default_ctx) {
+    status_addr = block_handle->default_ctx_status->dequeue();
+  }
+  else {
+    volatile char* status{block_handle->status};
+    int thread_id{get_flat_block_id()};
+    status_addr = &status[thread_id];
+  }
+  return status_addr;
 }
 
 }  // namespace rocshmem

src/reverse_offload/context_ro_device.hpp

@@ -34,12 +34,12 @@ __device__ void build_queue_element(
     ro_net_cmds type, void *dst, void *src, size_t size, int pe,
     int logPE_stride, int PE_size, int PE_root, void *pWrk, long *pSync,
     MPI_Comm team_comm, int ro_net_win_id, BlockHandle *handle,
-    bool blocking, volatile char *status = nullptr, ROCSHMEM_OP op = ROCSHMEM_SUM,
-    ro_net_types datatype = RO_NET_INT);
+    bool blocking, volatile char *status = nullptr, bool default_ctx = false,
+    ROCSHMEM_OP op = ROCSHMEM_SUM, ro_net_types datatype = RO_NET_INT);
 
 class ROContext : public Context {
  public:
-  __host__ ROContext(Backend *b, size_t block_id = 0);
+  __host__ ROContext(Backend *b, size_t block_id = 0, bool default_ctx = false);
 
   __device__ void threadfence_system();
 
@@ -225,6 +225,8 @@ class ROContext : public Context {
   BlockHandle *block_handle{nullptr};
 
   int ro_net_win_id{-1};
+
+  bool is_default_ctx{false};
 };
 
 }  // namespace rocshmem

src/reverse_offload/context_ro_tmpl_device.hpp

@@ -121,7 +121,7 @@ __device__ int ROContext::reduce(rocshmem_team_t team, T *dest,
   build_queue_element(RO_NET_TEAM_REDUCE, dest, const_cast<T *>(source),
                       nreduce, 0, 0, 0, 0, nullptr, nullptr, team_obj->mpi_comm,
                       ro_net_win_id, block_handle, true, get_status_flag(),
-                      Op, GetROType<T>::Type);
+                      is_default_ctx, Op, GetROType<T>::Type);
 
   __syncthreads();
   return ROCSHMEM_SUCCESS;
@@ -150,7 +150,7 @@ __device__ void ROContext::p(T *dest, T value, int pe) {
   } else {
     build_queue_element(RO_NET_P, dest, &value, sizeof(T), pe, 0, 0, 0, nullptr,
                         nullptr, (MPI_Comm)NULL, ro_net_win_id,
-                        block_handle, true, get_status_flag());
+                        block_handle, true, get_status_flag(), is_default_ctx);
   }
 }
 
@@ -165,6 +165,9 @@ __device__ T ROContext::g(const T *source, int pe) {
   } else {
     auto dest{get_g_ret_buf()};
     get<T>(reinterpret_cast<T *>(dest), source, 1, pe);
+    if (is_default_ctx) {
+      block_handle->default_ctx_g_ret->enqueue(dest);
+    }
     return *(reinterpret_cast<T *>(dest));
   }
 }
@@ -190,9 +193,12 @@ __device__ T ROContext::amo_fetch_cas(void *dst, T value, T cond, int pe) {
                       value, pe, 0, 0, 0,
                       reinterpret_cast<void *>(static_cast<long long>(cond)),
                       nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle,
-                      true, get_status_flag(), ROCSHMEM_SUM,
+                      true, get_status_flag(), is_default_ctx, ROCSHMEM_SUM,
                       GetROType<T>::Type);
   __threadfence();
+  if (is_default_ctx) {
+    block_handle->default_ctx_atomic_ret->enqueue(source);
+  }
   return *source;
 }
 
@@ -207,8 +213,11 @@ __device__ T ROContext::amo_fetch_add(void *dst, T value, int pe) {
   build_queue_element(RO_NET_AMO_FOP, dst, reinterpret_cast<T *>(source), value,
                       pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
                       ro_net_win_id, block_handle, true, get_status_flag(),
-                      ROCSHMEM_SUM, GetROType<T>::Type);
+                      is_default_ctx, ROCSHMEM_SUM, GetROType<T>::Type);
   __threadfence();
+  if (is_default_ctx) {
+    block_handle->default_ctx_atomic_ret->enqueue(source);
+  }
   return *source;
 }
 
@@ -223,8 +232,11 @@ __device__ T ROContext::amo_swap(void *dst, T value, int pe) {
   build_queue_element(RO_NET_AMO_FOP, dst, reinterpret_cast<void *>(source),
                       value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
                       ro_net_win_id, block_handle, true, get_status_flag(),
-                      ROCSHMEM_REPLACE, GetROType<T>::Type);
+                      is_default_ctx, ROCSHMEM_REPLACE, GetROType<T>::Type);
   __threadfence();
+  if (is_default_ctx) {
+    block_handle->default_ctx_atomic_ret->enqueue(source);
+  }
   return *source;
 }
 
@@ -239,8 +251,11 @@ __device__ T ROContext::amo_fetch_and(void *dst, T value, int pe) {
   build_queue_element(RO_NET_AMO_FOP, dst, reinterpret_cast<void *>(source),
                       value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
                       ro_net_win_id, block_handle, true, get_status_flag(),
-                      ROCSHMEM_AND, GetROType<T>::Type);
+                      is_default_ctx, ROCSHMEM_AND, GetROType<T>::Type);
   __threadfence();
+  if (is_default_ctx) {
+    block_handle->default_ctx_atomic_ret->enqueue(source);
+  }
   return *source;
 }
 
@@ -255,8 +270,11 @@ __device__ T ROContext::amo_fetch_or(void *dst, T value, int pe) {
   build_queue_element(RO_NET_AMO_FOP, dst, reinterpret_cast<void *>(source),
                       value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
                       ro_net_win_id, block_handle, true, get_status_flag(),
-                      ROCSHMEM_OR, GetROType<T>::Type);
+                      is_default_ctx, ROCSHMEM_OR, GetROType<T>::Type);
   __threadfence();
+  if (is_default_ctx) {
+    block_handle->default_ctx_atomic_ret->enqueue(source);
+  }
   return *source;
 }
 
@@ -271,8 +289,11 @@ __device__ T ROContext::amo_fetch_xor(void *dst, T value, int pe) {
   build_queue_element(RO_NET_AMO_FOP, dst, reinterpret_cast<void *>(source),
                       value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL,
                       ro_net_win_id, block_handle, true, get_status_flag(),
-                      ROCSHMEM_XOR, GetROType<T>::Type);
+                      is_default_ctx, ROCSHMEM_XOR, GetROType<T>::Type);
   __threadfence();
+  if (is_default_ctx) {
+    block_handle->default_ctx_atomic_ret->enqueue(source);
+  }
   return *source;
 }
 
@@ -294,7 +315,8 @@ __device__ void ROContext::broadcast(rocshmem_team_t team, T *dest,
   build_queue_element(RO_NET_TEAM_BROADCAST, dest, const_cast<T *>(source),
                       nelems, 0, 0, 0, pe_root, nullptr, nullptr,
                       team_obj->mpi_comm, ro_net_win_id, block_handle, true,
-                      get_status_flag(), ROCSHMEM_SUM, GetROType<T>::Type);
+                      get_status_flag(), is_default_ctx, ROCSHMEM_SUM,
+                      GetROType<T>::Type);
 
   __syncthreads();
 }
@@ -312,7 +334,8 @@ __device__ void ROContext::alltoall(rocshmem_team_t team, T *dest,
   build_queue_element(RO_NET_ALLTOALL, dest, const_cast<T *>(source), nelems, 0,
                       0, 0, 0, team_obj->ata_buffer, nullptr,
                       team_obj->mpi_comm, ro_net_win_id, block_handle, true,
-                      get_status_flag(), ROCSHMEM_SUM, GetROType<T>::Type);
+                      get_status_flag(), is_default_ctx, ROCSHMEM_SUM,
+                      GetROType<T>::Type);
 
   __syncthreads();
 }
@@ -330,7 +353,8 @@ __device__ void ROContext::fcollect(rocshmem_team_t team, T *dest,
   build_queue_element(RO_NET_FCOLLECT, dest, const_cast<T *>(source), nelems, 0,
                       0, 0, 0, team_obj->ata_buffer, nullptr,
                       team_obj->mpi_comm, ro_net_win_id, block_handle, true,
-                      get_status_flag(), ROCSHMEM_SUM, GetROType<T>::Type);
+                      get_status_flag(), is_default_ctx, ROCSHMEM_SUM,
+                      GetROType<T>::Type);
 
   __syncthreads();
 }

tests/functional_tests/CMakeLists.txt

@@ -45,6 +45,7 @@ target_sources(
     ping_all_tester.cpp
     primitive_tester.cpp
     primitive_mr_tester.cpp
+    default_ctx_primitive_tester.cpp
     team_ctx_primitive_tester.cpp
     team_ctx_infra_tester.cpp
     amo_bitwise_tester.cpp

tests/functional_tests/default_ctx_primitive_tester.cpp

@@ -0,0 +1,193 @@
+/******************************************************************************
+ * 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 "default_ctx_primitive_tester.hpp"
+
+ #include <rocshmem/rocshmem.hpp>
+ 
+ using namespace rocshmem;
+ 
+ /******************************************************************************
+  * DEVICE TEST KERNEL
+  *****************************************************************************/
+ __global__ void DefaultCTXPrimitiveTest(int loop, int skip,
+                               long long int *start_time,
+                               long long int *end_time, char *source,
+                               char *dest, int size, TestType type,
+                               ShmemContextType ctx_type, int wf_size) {
+   int wg_id = get_flat_grid_id();
+   int t_id  = get_flat_block_id();
+   int wf_id = t_id / wf_size;
+   rocshmem_wg_init();
+ 
+   /**
+    * Shared array to capture the start time for each wavefront
+    * Max threads per block = 1024, wavefront size = 64 (in most GPUs)
+    * Maximum array size required = 1024/64 = 16
+    */
+   __shared__ long long int wf_start_time[16];
+ 
+   /**
+    * Calculate start index for each thread within the grid
+    */
+   uint64_t offset = size * get_flat_id();
+   source += offset;
+   dest += offset;
+ 
+   for (int i = 0; i < loop + skip; i++) {
+     if (i == skip) {
+       __syncthreads();
+       // Ensures all RMA calls from the skip loops are completed
+       if(is_thread_zero_in_block()) {
+         rocshmem_quiet();
+       }
+       __syncthreads();
+       // Capture the start time of each wavefront to identify the earliest one
+       wf_start_time[wf_id] = wall_clock64();
+     }
+ 
+     switch (type) {
+       case DefaultCTXGetTestType:
+         rocshmem_getmem(dest, source, size, 1);
+         break;
+       case DefaultCTXGetNBITestType:
+         rocshmem_getmem_nbi(dest, source, size, 1);
+         break;
+       case DefaultCTXPutTestType:
+         rocshmem_putmem(dest, source, size, 1);
+         break;
+       case DefaultCTXPutNBITestType:
+         rocshmem_putmem_nbi(dest, source, size, 1);
+         break;
+       case DefaultCTXPTestType:
+         for (int s = 0; s < size; s++) {
+           char val = source[s];
+           rocshmem_char_p(&dest[s], val, 1);
+         }
+         break;
+       case DefaultCTXGTestType:
+         for (int s = 0; s < size; s++) {
+           char ret = rocshmem_char_g(&source[s], 1);
+           dest[s] = ret;
+         }
+         break;
+       default:
+         break;
+     }
+   }
+ 
+   __syncthreads();
+   if(is_thread_zero_in_block()) {
+     rocshmem_quiet();
+   }
+ 
+   /**
+    * End time of the last wavefront is recorded by overwriting
+    * the value previously set by earlier wavefronts.
+    */
+   end_time[wg_id] = wall_clock64();
+ 
+   // Find the earliest start time
+   int num_wfs = (get_flat_block_size() - 1 ) / wf_size + 1;
+   for (int i = num_wfs / 2; i > 0; i >>= 1 ) {
+     if(t_id < i) {
+       wf_start_time[t_id] = min(wf_start_time[t_id], wf_start_time[t_id + i]);
+     }
+   }
+   __syncthreads();
+ 
+   if (t_id == 0) {
+     start_time[wg_id] = wf_start_time[0];
+   }
+ 
+   rocshmem_wg_finalize();
+ }
+ 
+ /******************************************************************************
+  * HOST TESTER CLASS METHODS
+  *****************************************************************************/
+ DefaultCTXPrimitiveTester::DefaultCTXPrimitiveTester(TesterArguments args)
+  : Tester(args) {
+   size_t buff_size = args.max_msg_size * args.wg_size * args.num_wgs;
+   source = (char *)rocshmem_malloc(buff_size);
+   dest = (char *)rocshmem_malloc(buff_size);
+ 
+   if (source == nullptr || dest == nullptr) {
+     std::cerr << "Error allocating memory from symmetric heap" << std::endl;
+     std::cerr << "source: " << source << ", dest: " << dest << std::endl;
+     if (source) {
+       rocshmem_free(source);
+     }
+     if (dest) {
+       rocshmem_free(dest);
+     }
+     rocshmem_global_exit(1);
+   }
+ 
+   for(size_t i = 0; i < buff_size; i++) {
+     source[i] = static_cast<char>('a' + i % 26);
+   }
+ }
+ 
+ DefaultCTXPrimitiveTester::~DefaultCTXPrimitiveTester() {
+   rocshmem_free(source);
+   rocshmem_free(dest);
+ }
+ 
+ void DefaultCTXPrimitiveTester::resetBuffers(uint64_t size) {
+   size_t buff_size = size * args.wg_size * args.num_wgs;
+   memset(dest, '1', buff_size);
+ }
+ 
+ void DefaultCTXPrimitiveTester::launchKernel(dim3 gridSize, dim3 blockSize,
+                                              int loop, uint64_t size) {
+   size_t shared_bytes = 0;
+ 
+   hipLaunchKernelGGL(DefaultCTXPrimitiveTest, gridSize, blockSize,
+                      shared_bytes, stream, loop, args.skip, start_time,
+                      end_time, source, dest, size, _type, _shmem_context,
+                      wf_size);
+ 
+   num_msgs = (loop + args.skip) * gridSize.x * blockSize.x;
+   num_timed_msgs = loop * gridSize.x * blockSize.x;
+ }
+ 
+ void DefaultCTXPrimitiveTester::verifyResults(uint64_t size) {
+   int check_id =
+       (_type == DefaultCTXGetTestType ||
+        _type == DefaultCTXGetNBITestType || _type == DefaultCTXGTestType)
+           ? 0
+           : 1;
+ 
+   if (args.myid == check_id) {
+     size_t buff_size = size * args.wg_size * args.num_wgs;
+     for (uint64_t i = 0; i < buff_size; i++) {
+       if (dest[i] != source[i]) {
+         std::cerr << "Data validation error at idx " << i << std::endl;
+         std::cerr << " Got " << dest[i] << ", Expected "
+                   << source[i] << std::endl;
+         exit(-1);
+       }
+     }
+   }
+ }
+ 

tests/functional_tests/default_ctx_primitive_tester.hpp

@@ -0,0 +1,48 @@
+/******************************************************************************
+ * 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 _DEFAULT_CTX_PRIMITIVE_TESTER_HPP_
+#define _DEFAULT_CTX_PRIMITIVE_TESTER_HPP_
+
+#include "tester.hpp"
+
+/******************************************************************************
+ * HOST TESTER CLASS
+ *****************************************************************************/
+class DefaultCTXPrimitiveTester : public Tester {
+ public:
+  explicit DefaultCTXPrimitiveTester(TesterArguments args);
+  virtual ~DefaultCTXPrimitiveTester();
+
+ protected:
+  virtual void resetBuffers(uint64_t size) override;
+
+  virtual void launchKernel(dim3 gridSize, dim3 blockSize, int loop,
+                            uint64_t size) override;
+
+  virtual void verifyResults(uint64_t size) override;
+
+  char *source = nullptr;
+  char *dest = nullptr;
+};
+
+#endif

tests/functional_tests/tester.cpp

@@ -33,6 +33,7 @@
 #include "amo_bitwise_tester.hpp"
 #include "amo_extended_tester.hpp"
 #include "amo_standard_tester.hpp"
+#include "default_ctx_primitive_tester.hpp"
 #include "barrier_all_tester.hpp"
 #include "empty_tester.hpp"
 #include "ping_all_tester.hpp"
@@ -119,6 +120,26 @@ std::vector<Tester*> Tester::create(TesterArguments args) {
       if (rank == 0) std::cout << "Non-Blocking Puts ###" << std::endl;
       testers.push_back(new PrimitiveTester(args));
       return testers;
+    case DefaultCTXGetTestType:
+      if (rank == 0)
+        std::cout << "Default context Blocking Gets ###" << std::endl;
+      testers.push_back(new DefaultCTXPrimitiveTester(args));
+      return testers;
+    case DefaultCTXGetNBITestType:
+      if (rank == 0)
+        std::cout << "Default context Non-Blocking Gets ###" << std::endl;
+      testers.push_back(new DefaultCTXPrimitiveTester(args));
+      return testers;
+    case DefaultCTXPutTestType:
+      if (rank == 0)
+        std::cout << "Default context Blocking Puts ###" << std::endl;
+      testers.push_back(new DefaultCTXPrimitiveTester(args));
+      return testers;
+    case DefaultCTXPutNBITestType:
+      if (rank == 0)
+        std::cout << "Default context Non-Blocking Puts ###" << std::endl;
+      testers.push_back(new DefaultCTXPrimitiveTester(args));
+      return testers;
     case TeamCtxInfraTestType:
       if (rank == 0) std::cout << "Team Ctx Infra test ###" << std::endl;
       testers.push_back(new TeamCtxInfraTester(args));

tests/functional_tests/tester.hpp

@@ -93,6 +93,12 @@ enum TestType {
   WGSignalFetchTestType = 56,
   WAVESignalFetchTestType = 57,
   TeamBarrierTestType = 58,
+  DefaultCTXGetTestType = 59,
+  DefaultCTXGetNBITestType = 60,
+  DefaultCTXPutTestType = 61,
+  DefaultCTXPutNBITestType = 62,
+  DefaultCTXPTestType = 63,
+  DefaultCTXGTestType = 64,
 };
 
 enum OpType { PutType = 0, GetType = 1 };