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();
 }