diff --git a/src/reverse_offload/backend_ro.cpp b/src/reverse_offload/backend_ro.cpp index 9fc5fcd503..bdcd54df98 100644 --- a/src/reverse_offload/backend_ro.cpp +++ b/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)); diff --git a/src/reverse_offload/backend_ro.hpp b/src/reverse_offload/backend_ro.hpp index e4907c1c84..62c6b1c8f0 100644 --- a/src/reverse_offload/backend_ro.hpp +++ b/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 ctx_free_list{}; + /** + * @brief AtomicWFQueue containing status flag buffers for default context + */ + AtomicWFQueueProxy default_ctx_status_{}; + + /** + * @brief AtomicWFQueue containing rocshmem_g return buffers for default + * context + */ + AtomicWFQueueProxy default_ctx_g_ret_buffer_{}; + + /** + * @brief AtomicWFQueue containing rocshmem return buffers for default + * context + */ + AtomicWFQueueProxy 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 diff --git a/src/reverse_offload/block_handle.hpp b/src/reverse_offload/block_handle.hpp index e0d3e61ec3..ad9301094d 100644 --- a/src/reverse_offload/block_handle.hpp +++ b/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; +using AWF_Queue_ret_buffT = AtomicWFQueue; + 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 @@ -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; diff --git a/src/reverse_offload/context_proxy.hpp b/src/reverse_offload/context_proxy.hpp index f0a8fdcb44..2d77c8a125 100644 --- a/src/reverse_offload/context_proxy.hpp +++ b/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), -1); + new (ctx) ROContext(reinterpret_cast(backend), -1, true); rocshmem_ctx_t local{ctx, tinfo}; set_internal_ctx(&local); } diff --git a/src/reverse_offload/context_ro_device.cpp b/src/reverse_offload/context_ro_device.cpp index cef85b5f54..bad1028358 100644 --- a/src/reverse_offload/context_ro_device.cpp +++ b/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(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(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(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(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(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(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(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(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(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(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(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 diff --git a/src/reverse_offload/context_ro_device.hpp b/src/reverse_offload/context_ro_device.hpp index 9ebfa2a30f..dfda274caa 100644 --- a/src/reverse_offload/context_ro_device.hpp +++ b/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 diff --git a/src/reverse_offload/context_ro_tmpl_device.hpp b/src/reverse_offload/context_ro_tmpl_device.hpp index a40eee2cad..d37e2d7b2d 100644 --- a/src/reverse_offload/context_ro_tmpl_device.hpp +++ b/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(source), nreduce, 0, 0, 0, 0, nullptr, nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle, true, get_status_flag(), - Op, GetROType::Type); + is_default_ctx, Op, GetROType::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(reinterpret_cast(dest), source, 1, pe); + if (is_default_ctx) { + block_handle->default_ctx_g_ret->enqueue(dest); + } return *(reinterpret_cast(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(static_cast(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::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(source), value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle, true, get_status_flag(), - ROCSHMEM_SUM, GetROType::Type); + is_default_ctx, ROCSHMEM_SUM, GetROType::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(source), value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle, true, get_status_flag(), - ROCSHMEM_REPLACE, GetROType::Type); + is_default_ctx, ROCSHMEM_REPLACE, GetROType::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(source), value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle, true, get_status_flag(), - ROCSHMEM_AND, GetROType::Type); + is_default_ctx, ROCSHMEM_AND, GetROType::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(source), value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle, true, get_status_flag(), - ROCSHMEM_OR, GetROType::Type); + is_default_ctx, ROCSHMEM_OR, GetROType::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(source), value, pe, 0, 0, 0, nullptr, nullptr, (MPI_Comm)NULL, ro_net_win_id, block_handle, true, get_status_flag(), - ROCSHMEM_XOR, GetROType::Type); + is_default_ctx, ROCSHMEM_XOR, GetROType::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(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::Type); + get_status_flag(), is_default_ctx, ROCSHMEM_SUM, + GetROType::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(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::Type); + get_status_flag(), is_default_ctx, ROCSHMEM_SUM, + GetROType::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(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::Type); + get_status_flag(), is_default_ctx, ROCSHMEM_SUM, + GetROType::Type); __syncthreads(); } diff --git a/tests/functional_tests/CMakeLists.txt b/tests/functional_tests/CMakeLists.txt index 58851c4b74..2070e8fcb5 100644 --- a/tests/functional_tests/CMakeLists.txt +++ b/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 diff --git a/tests/functional_tests/default_ctx_primitive_tester.cpp b/tests/functional_tests/default_ctx_primitive_tester.cpp new file mode 100644 index 0000000000..3ea5e81966 --- /dev/null +++ b/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 + + 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('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); + } + } + } + } + \ No newline at end of file diff --git a/tests/functional_tests/default_ctx_primitive_tester.hpp b/tests/functional_tests/default_ctx_primitive_tester.hpp new file mode 100644 index 0000000000..8a7ff3d615 --- /dev/null +++ b/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 diff --git a/tests/functional_tests/tester.cpp b/tests/functional_tests/tester.cpp index 9bc16a1b73..bd0f58baa2 100644 --- a/tests/functional_tests/tester.cpp +++ b/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::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)); diff --git a/tests/functional_tests/tester.hpp b/tests/functional_tests/tester.hpp index 76871f4b33..fe691bd1db 100644 --- a/tests/functional_tests/tester.hpp +++ b/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 };