diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000..567609b123 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +build/ diff --git a/CMakeLists.txt b/CMakeLists.txt index f9c3d83d10..5cf74fb6cb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -68,6 +68,7 @@ option(USE_FUNC_CALL "Force compiler to use function calls on library API" OFF) option(USE_SHARED_CTX "Request support for shared ctx between WG" OFF) option(USE_SINGLE_NODE "Enable single node support only." OFF) option(USE_HOST_SIDE_HDP_FLUSH "Use a polling thread to flush the HDP cache on the host." OFF) +option(USE_COOPERATIVE_GROUPS "Use cooperative groups for internal syncronization" OFF) option(BUILD_FUNCTIONAL_TESTS "Build the functional tests" ON) option(BUILD_SOS_TESTS "Build the host-facing tests" OFF) option(BUILD_UNIT_TESTS "Build the unit tests" ON) diff --git a/cmake/config.h.in b/cmake/config.h.in index 8067651a01..21619a6c3e 100644 --- a/cmake/config.h.in +++ b/cmake/config.h.in @@ -14,3 +14,4 @@ #cmakedefine USE_FUNC_CALL #cmakedefine USE_SINGLE_NODE #cmakedefine USE_HOST_SIDE_HDP_FLUSH +#cmakedefine USE_COOPERATIVE_GROUPS diff --git a/scripts/build_configs/ipc_single_cg b/scripts/build_configs/ipc_single_cg new file mode 100755 index 0000000000..c0904bad93 --- /dev/null +++ b/scripts/build_configs/ipc_single_cg @@ -0,0 +1,31 @@ +#!/bin/bash +# Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. + +if [ -z $1 ] +then + install_path=~/rocshmem +else + install_path=$1 +fi + +src_path=$(dirname "$(realpath $0)")/../../ + +cmake \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=$install_path \ + -DCMAKE_VERBOSE_MAKEFILE=OFF \ + -DDEBUG=OFF \ + -DPROFILE=OFF \ + -DUSE_GPU_IB=OFF \ + -DUSE_RO=OFF \ + -DUSE_DC=OFF \ + -DUSE_IPC=ON \ + -DUSE_COHERENT_HEAP=ON \ + -DUSE_THREADS=OFF \ + -DUSE_WF_COAL=OFF \ + -DUSE_SINGLE_NODE=ON \ + -DUSE_HOST_SIDE_HDP_FLUSH=OFF \ + -DUSE_COOPERATIVE_GROUPS=ON \ + $src_path +cmake --build . --parallel 8 +cmake --install . diff --git a/scripts/unit_tests/driver.sh b/scripts/unit_tests/driver.sh new file mode 100755 index 0000000000..8eea7daef7 --- /dev/null +++ b/scripts/unit_tests/driver.sh @@ -0,0 +1,62 @@ +#!/bin/bash + +# Function to display help information +function display_help { + echo "Usage:" + echo " $0 binary_name all # Runs all standard tests" + echo " $0 binary_name custom # Runs custom test configuration" + echo + echo "Arguments:" + echo " binary_name: Name of the binary to run." + echo " all: Executes predefined test configurations." + echo " custom: Executes a test with custom MPI ranks and GTest filter." + echo " ranks: Number of MPI ranks (required for custom mode)." + echo " filter: GTest filter string (required for custom mode)." + echo +} + +# Validate number of arguments for each mode +if [[ "$#" -lt 2 ]] || + { [[ "$2" == "all" ]] && [[ "$#" -ne 2 ]]; } || + { [[ "$2" == "custom" ]] && [[ "$#" -ne 4 ]]; }; then + display_help + exit 1 +fi + +binary_name=$1 +mode=$2 +timestamp=$(date "+%Y-%m-%d-%H:%M:%S") +log_file="unit_tests_${timestamp}.log" + +# Function to execute mpirun command +function run_mpirun { + local np=$1 + local gtest_filter=$2 + echo "mpirun -np $np $binary_name --gtest_filter='$gtest_filter'" | tee -a "$log_file" + mpirun -np "$np" "$binary_name" --gtest_filter="$gtest_filter" >> "$log_file" 2>&1 +} + +# Processing modes +case $mode in + all) + run_mpirun 4 "-IPCImplSimpleCoarseTestFixture.*:IPCImplSimpleFineTestFixture.*" + run_mpirun 2 "IPCImplSimpleCoarseTestFixture.*:IPCImplSimpleFineTestFixture.*" + ;; + custom) + # Check if ranks is a positive integer + if [[ "$3" -le 1 ]]; then + echo "Error: 'ranks' must be a positive integer." + display_help + exit 1 + fi + run_mpirun $3 $4 + ;; + *) + echo "Error: Invalid mode '$mode'." | tee -a "$log_file" + display_help + exit 1 + ;; +esac + +echo "Tests Completed" +echo "log file: '$log_file'" diff --git a/src/assembly.hpp b/src/assembly.hpp index 076600067e..a14907a0bd 100644 --- a/src/assembly.hpp +++ b/src/assembly.hpp @@ -135,12 +135,12 @@ NOWARN(-Wdeprecated-volatile, : "v"(src)); #endif break; + case 8: #if defined(__gfx906__) #endif #if defined(__gfx908__) #endif #if defined(__gfx90a__) - case 8: asm volatile( "global_load_dwordx2 %0 %1 off glc slc \n" "s_waitcnt vmcnt(0)" diff --git a/src/gpu_ib/backend_ib.cpp b/src/gpu_ib/backend_ib.cpp index 9ae3acd482..4692235738 100644 --- a/src/gpu_ib/backend_ib.cpp +++ b/src/gpu_ib/backend_ib.cpp @@ -133,8 +133,6 @@ __device__ bool GPUIBBackend::create_ctx(int64_t options, return false; } ctx_ = pop_result.value; - - ctx_->dev_mtx_.shareable_ = static_cast(options >> 3); ctx->ctx_opaque = ctx_; return true; } @@ -296,7 +294,8 @@ void GPUIBBackend::setup_default_ctx() { CHECK_HIP(hipGetSymbolAddress(reinterpret_cast(&symbol_address), HIP_SYMBOL(ROC_SHMEM_CTX_DEFAULT))); - roc_shmem_ctx_t ctx_default_host{default_ctx_, nullptr}; + TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world; + roc_shmem_ctx_t ctx_default_host{default_ctx_, tinfo}; hipStream_t stream; CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking)); diff --git a/src/gpu_ib/segment_builder.cpp b/src/gpu_ib/segment_builder.cpp index 205c890861..6f2a77a23e 100644 --- a/src/gpu_ib/segment_builder.cpp +++ b/src/gpu_ib/segment_builder.cpp @@ -22,6 +22,7 @@ #include "segment_builder.hpp" +#include "../util.hpp" #include "endian.hpp" namespace rocshmem { @@ -35,9 +36,9 @@ __device__ SegmentBuilder::SegmentBuilder(uint64_t wqe_idx, void *base) { __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; + mlx5_wqe_ctrl_seg ctrl_seg; - ctrl_seg->opmod_idx_opcode = (opcode << 24) | (wqe_idx << 8); + ctrl_seg.opmod_idx_opcode = (opcode << 24) | (wqe_idx << 8); uint32_t DS = 2; if (zero_byte_rd == false) { @@ -48,39 +49,43 @@ __device__ void SegmentBuilder::update_cntrl_seg( DS += connection_policy->wqeCntrlOffset(); - ctrl_seg->qpn_ds = (DS << 24) | ctrl_qp_sq; + ctrl_seg.qpn_ds = (DS << 24) | ctrl_qp_sq; - ctrl_seg->signature = ctrl_sig; + ctrl_seg.signature = ctrl_sig; - ctrl_seg->fm_ce_se = ctrl_sig >> 24; + ctrl_seg.fm_ce_se = ctrl_sig >> 24; - ctrl_seg->imm = ctrl_sig >> 32; + ctrl_seg.imm = ctrl_sig >> 32; + + memcpy(&seg_ptr->ctrl_seg, &ctrl_seg, sizeof(mlx5_wqe_ctrl_seg)); 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; + mlx5_wqe_atomic_seg atomic_seg; - swap_endian_store(reinterpret_cast(&atomic_seg->swap_add), + swap_endian_store(reinterpret_cast(&atomic_seg.swap_add), atomic_data); - swap_endian_store(reinterpret_cast(&atomic_seg->compare), + swap_endian_store(reinterpret_cast(&atomic_seg.compare), atomic_cmp); + memcpy(&seg_ptr->atomic_seg, &atomic_seg, sizeof(mlx5_wqe_atomic_seg)); seg_ptr++; } __device__ void SegmentBuilder::update_rdma_seg(uintptr_t *raddr, uint32_t rkey) { - mlx5_wqe_raddr_seg *raddr_seg = &seg_ptr->raddr_seg; + mlx5_wqe_raddr_seg raddr_seg; - raddr_seg->rkey = rkey; + raddr_seg.rkey = rkey; - swap_endian_store(reinterpret_cast(&raddr_seg->raddr), + swap_endian_store(reinterpret_cast(&raddr_seg.raddr), reinterpret_cast(raddr)); + memcpy(&seg_ptr->raddr_seg, &raddr_seg, sizeof(mlx5_wqe_raddr_seg)); seg_ptr++; } @@ -90,33 +95,36 @@ __device__ void SegmentBuilder::update_data_seg(uintptr_t *laddr, int32_t size, return; } - mlx5_wqe_data_seg *data_seg = &seg_ptr->data_seg; + mlx5_wqe_data_seg data_seg; + data_seg.lkey = lkey; - data_seg->lkey = lkey; - - swap_endian_store(&data_seg->byte_count, size & 0x7FFFFFFFU); - - swap_endian_store(reinterpret_cast(&data_seg->addr), + swap_endian_store(&data_seg.byte_count, size & 0x7FFFFFFFU); + swap_endian_store(reinterpret_cast(&data_seg.addr), reinterpret_cast(laddr)); + memcpy(&seg_ptr->data_seg, &data_seg, sizeof(mlx5_wqe_data_seg)); 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; + mlx5_wqe_inl_data_seg inl_data_seg; - swap_endian_store(&inl_data_seg->byte_count, (size & 0x3FF) | 0x80000000); + swap_endian_store(&inl_data_seg.byte_count, (size & 0x3FF) | 0x80000000); // Assume fence HDP flush // TODO(khamidou): Rework fence interface to avoid this + size_t field_size{sizeof(mlx5_wqe_inl_data_seg)}; if (!laddr) { uint8_t flush_val = 1; - memcpy(inl_data_seg + 1, &flush_val, sizeof(flush_val)); + memcpy(&inl_data_seg + 1, &flush_val, sizeof(flush_val)); + field_size += sizeof(flush_val); } else { - memcpy(inl_data_seg + 1, laddr, size); + memcpy(&inl_data_seg + 1, laddr, size); + field_size += size; } + memcpy(&seg_ptr->inl_data_seg, &inl_data_seg, field_size); seg_ptr++; } diff --git a/src/ipc/backend_ipc.cpp b/src/ipc/backend_ipc.cpp index 58caaf53f7..7475fb2c8b 100644 --- a/src/ipc/backend_ipc.cpp +++ b/src/ipc/backend_ipc.cpp @@ -56,18 +56,18 @@ int get_ls_non_zero_bit(char *bitmask, int mask_length) { IPCBackend::IPCBackend(MPI_Comm comm) : Backend() { type = BackendType::IPC_BACKEND; - + if (auto maximum_num_contexts_str = getenv("ROC_SHMEM_MAX_NUM_CONTEXTS")) { std::stringstream sstream(maximum_num_contexts_str); sstream >> maximum_num_contexts_; } init_mpi_once(comm); - + initIPC(); - + auto *bp{ipc_backend_proxy.get()}; - + bp->heap_ptr = &heap; /* Initialize the host interface */ @@ -82,16 +82,19 @@ IPCBackend::IPCBackend(MPI_Comm comm) allocate_atomic_region(&bp->atomic_ret, MAX_NUM_BLOCKS); - default_context_proxy_ = IPCDefaultContextProxyT(this); - setup_team_world(); + TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world; + roc_shmem_collective_init(); + setup_fence_buffer(); + teams_init(); + default_context_proxy_ = IPCDefaultContextProxyT(this, tinfo); + setup_ctxs(); - } IPCBackend::~IPCBackend() { @@ -105,7 +108,7 @@ IPCBackend::~IPCBackend() { * Free the atomic_ret array. */ CHECK_HIP(hipFree(bp->atomic_ret->atomic_base_ptr)); - + // TODO(Avinash) Free g_ret // delete host_interface; @@ -141,6 +144,8 @@ __device__ bool IPCBackend::create_ctx(int64_t options, roc_shmem_ctx_t *ctx) { ctx_ = pop_result.value; ctx->ctx_opaque = ctx_; + + ctx_->tinfo = reinterpret_cast(ctx->team_opaque); return true; } @@ -287,6 +292,14 @@ void IPCBackend::teams_destroy() { free(reduced_bitmask_); } +void IPCBackend::setup_fence_buffer() { + /* + * Allocate heap space for fence + */ + fence_pool = reinterpret_cast(roc_shmem_malloc( + sizeof(int) * num_pes)); +} + void IPCBackend::roc_shmem_collective_init() { /* * Allocate heap space for barrier_sync @@ -389,4 +402,4 @@ void IPCBackend::teams_init() { NET_CHECK(MPI_Barrier(thread_comm)); } -} // namespace rocshmem \ No newline at end of file +} // namespace rocshmem diff --git a/src/ipc/backend_ipc.hpp b/src/ipc/backend_ipc.hpp index cc385f16ee..de8cb9c617 100644 --- a/src/ipc/backend_ipc.hpp +++ b/src/ipc/backend_ipc.hpp @@ -32,6 +32,10 @@ #include "ipc_context_proxy.hpp" #include "../ipc_policy.hpp" +#ifndef USE_COOPERATIVE_GROUPS +#include "../memory/notifier.hpp" +#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */ + namespace rocshmem { class IPCBackend : public Backend { @@ -104,7 +108,7 @@ class IPCBackend : public Backend { 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) */ @@ -156,6 +160,15 @@ class IPCBackend : public Backend { */ void *pAta_pool{nullptr}; + /** + * @brief Handle for raw memory for fence/quiet + */ + int *fence_pool{nullptr}; + +#ifndef USE_COOPERATIVE_GROUPS + NotifierProxy notifier_{}; +#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */ + protected: /** * @copydoc Backend::dump_backend_stats() @@ -203,6 +216,11 @@ class IPCBackend : public Backend { */ void roc_shmem_collective_init(); + /** + * @brief Allocate buffer for fence/quiet operation + */ + void setup_fence_buffer(); + private: /** * @brief Proxy for the default context @@ -249,4 +267,4 @@ class IPCBackend : public Backend { } // namespace rocshmem -#endif // LIBRARY_SRC_IPC_BACKEND_HPP_ \ No newline at end of file +#endif // LIBRARY_SRC_IPC_BACKEND_HPP_ diff --git a/src/ipc/context_ipc_device.cpp b/src/ipc/context_ipc_device.cpp index 6fbb362d73..d208a070fb 100644 --- a/src/ipc/context_ipc_device.cpp +++ b/src/ipc/context_ipc_device.cpp @@ -47,6 +47,13 @@ __host__ IPCContext::IPCContext(Backend *b) barrier_sync = backend->barrier_sync; g_ret = bp->g_ret; atomic_base_ptr = bp->atomic_ret->atomic_base_ptr; + fence_pool = backend->fence_pool; + + orders_.store = detail::atomic::rocshmem_memory_order::memory_order_seq_cst; + +#ifndef USE_COOPERATIVE_GROUPS + notifier_ = backend->notifier_.get(); +#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */ } __device__ void IPCContext::threadfence_system() { @@ -85,12 +92,17 @@ __device__ void IPCContext::getmem_nbi(void *dest, const void *source, } __device__ void IPCContext::fence() { + for (int i{0}, j{tinfo->pe_start}; i < tinfo->size; i++, j += tinfo->stride) { + detail::atomic::store(&fence_pool[j], 1, orders_); + } } __device__ void IPCContext::fence(int pe) { + detail::atomic::store(&fence_pool[pe], 1, orders_); } __device__ void IPCContext::quiet() { + fence(); } __device__ void *IPCContext::shmem_ptr(const void *dest, int pe) { diff --git a/src/ipc/context_ipc_device.hpp b/src/ipc/context_ipc_device.hpp index 4cc6d08d56..065d37cec4 100644 --- a/src/ipc/context_ipc_device.hpp +++ b/src/ipc/context_ipc_device.hpp @@ -24,6 +24,15 @@ #define LIBRARY_SRC_IPC_CONTEXT_DEVICE_HPP_ #include "../context.hpp" +#include "../atomic.hpp" +#include "../team.hpp" + +#ifdef USE_COOPERATIVE_GROUPS +#include +namespace cg = cooperative_groups; +#else +#include "../memory/notifier.hpp" +#endif /* USE_COOPERATIVE_GROUPS */ namespace rocshmem { @@ -232,6 +241,24 @@ class IPCContext : public Context { //Temporary scratchpad memory used by internal barrier algorithms. int64_t *barrier_sync{nullptr}; + +#ifndef USE_COOPERATIVE_GROUPS + Notifier *notifier_{nullptr}; +#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */ + + //Struct defining memory ordering for atomic operations. + detail::atomic::rocshmem_memory_orders orders_{}; + + //Buffer to perform Atomic store to enforce memory ordering + int *fence_pool{nullptr}; + + public: + //TODO(Avinash): + //Make tinfo private variable, it requires changes to the context + //creation API in backend + + //Team information for the team associated with the context + TeamInfo *tinfo{nullptr}; }; } // namespace rocshmem diff --git a/src/ipc/context_ipc_device_coll.cpp b/src/ipc/context_ipc_device_coll.cpp index 804be2efbf..f128d76d3f 100644 --- a/src/ipc/context_ipc_device_coll.cpp +++ b/src/ipc/context_ipc_device_coll.cpp @@ -35,6 +35,9 @@ __device__ void IPCContext::internal_direct_barrier(int pe, int PE_start, if (pe == PE_start) { // Go through all PE offsets (except current offset = 0) // and wait until they all reach +#if defined(__gfx90a__) + __threadfence_system(); +#endif /* __gfx90a__ */ 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; @@ -49,6 +52,9 @@ __device__ void IPCContext::internal_direct_barrier(int pe, int PE_start, // Mark current PE offset as reached size_t pe_offset = (pe - PE_start) / stride; put_nbi(&pSync[pe_offset], &flag_val, 1, PE_start); +#if defined(__gfx90a__) + __threadfence_system(); +#endif /* __gfx90a__ */ wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val); pSync[0] = ROC_SHMEM_SYNC_VALUE; threadfence_system(); @@ -78,8 +84,15 @@ __device__ void IPCContext::internal_atomic_barrier(int pe, int PE_start, // Uses PE values that are relative to world __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride, int PE_size, int64_t *pSync) { - __syncthreads(); - if (is_thread_zero_in_block()) { +#ifdef USE_COOPERATIVE_GROUPS + cg::grid_group grid = cg::this_grid(); + grid.sync(); + if (0 == grid.thread_rank()) +#else + notifier_->sync(); + if (0 == get_flat_id()) +#endif /* USE_COOPERATIVE_GROUPS */ + { if (PE_size < 64) { internal_direct_barrier(pe, PE_start, stride, PE_size, pSync); } else { @@ -87,7 +100,11 @@ __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride, } } __threadfence(); - __syncthreads(); +#ifdef USE_COOPERATIVE_GROUPS + grid.sync(); +#else + notifier_->sync(); +#endif /* USE_COOPERATIVE_GROUPS */ } __device__ void IPCContext::sync(roc_shmem_team_t team) { diff --git a/src/ipc/ipc_context_proxy.hpp b/src/ipc/ipc_context_proxy.hpp index 867f199094..87ca4371e0 100644 --- a/src/ipc/ipc_context_proxy.hpp +++ b/src/ipc/ipc_context_proxy.hpp @@ -41,10 +41,11 @@ class IPCDefaultContextProxy { /* * Placement new the memory which is allocated by proxy_ */ - explicit IPCDefaultContextProxy(IPCBackend* backend) : constructed_{true} { + explicit IPCDefaultContextProxy(IPCBackend* backend, TeamInfo *tinfo) + : constructed_{true} { auto ctx{proxy_.get()}; new (ctx) IPCContext(reinterpret_cast(backend)); - roc_shmem_ctx_t local{ctx, nullptr}; + roc_shmem_ctx_t local{ctx, tinfo}; set_internal_ctx(&local); } diff --git a/src/reverse_offload/backend_ro.cpp b/src/reverse_offload/backend_ro.cpp index 1c1bf4645c..96471181e3 100644 --- a/src/reverse_offload/backend_ro.cpp +++ b/src/reverse_offload/backend_ro.cpp @@ -94,7 +94,9 @@ ROBackend::ROBackend(MPI_Comm comm) default_block_handle_proxy_ = DefaultBlockHandleProxyT( bp->g_ret, bp->atomic_ret, &queue_, &ipcImpl, hdp_proxy_.get()); - default_context_proxy_ = DefaultContextProxyT(this); + + TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world; + default_context_proxy_ = DefaultContextProxyT(this, tinfo); block_handle_proxy_ = BlockHandleProxyT(bp->g_ret, bp->atomic_ret, &queue_, &ipcImpl, hdp_proxy_.get()); diff --git a/src/reverse_offload/context_proxy.hpp b/src/reverse_offload/context_proxy.hpp index d89281d921..4ae94d3851 100644 --- a/src/reverse_offload/context_proxy.hpp +++ b/src/reverse_offload/context_proxy.hpp @@ -42,10 +42,11 @@ class DefaultContextProxy { /* * Placement new the memory which is allocated by proxy_ */ - explicit DefaultContextProxy(ROBackend* backend) : constructed_{true} { + explicit DefaultContextProxy(ROBackend* backend, TeamInfo *tinfo) + : constructed_{true} { auto ctx{proxy_.get()}; new (ctx) ROContext(reinterpret_cast(backend), -1); - roc_shmem_ctx_t local{ctx, nullptr}; + roc_shmem_ctx_t local{ctx, tinfo}; set_internal_ctx(&local); } diff --git a/src/roc_shmem_gpu.cpp b/src/roc_shmem_gpu.cpp index 086a6deed1..1afd5b10ff 100644 --- a/src/roc_shmem_gpu.cpp +++ b/src/roc_shmem_gpu.cpp @@ -262,9 +262,9 @@ __device__ int roc_shmem_wg_ctx_create(long option, roc_shmem_ctx_t *ctx) { GPU_DPRINTF("Function: roc_shmem_ctx_create\n"); bool result{true}; if (get_flat_block_id() == 0) { + ctx->team_opaque = reinterpret_cast(ROC_SHMEM_CTX_DEFAULT.team_opaque); device_backend_proxy->create_ctx(option, ctx); reinterpret_cast(ctx->ctx_opaque)->setFence(option); - ctx->team_opaque = nullptr; } __syncthreads(); return result == true ? 0 : -1; @@ -279,11 +279,11 @@ __device__ int roc_shmem_wg_team_create_ctx(roc_shmem_team_t team, long options, bool result{true}; if (get_flat_block_id() == 0) { - result = device_backend_proxy->create_ctx(options, ctx); - reinterpret_cast(ctx->ctx_opaque)->setFence(options); Team *team_obj{get_internal_team(team)}; TeamInfo *info_wrt_world = team_obj->tinfo_wrt_world; ctx->team_opaque = info_wrt_world; + result = device_backend_proxy->create_ctx(options, ctx); + reinterpret_cast(ctx->ctx_opaque)->setFence(options); } __syncthreads(); @@ -412,7 +412,9 @@ __device__ void roc_shmem_ctx_fence(roc_shmem_ctx_t ctx) { __device__ void roc_shmem_ctx_fence(roc_shmem_ctx_t ctx, int pe) { GPU_DPRINTF("Function: roc_shmem_ctx_fence\n"); - get_internal_ctx(ctx)->fence(pe); + int pe_in_world = translate_pe(ctx, pe); + + get_internal_ctx(ctx)->fence(pe_in_world); } __device__ void roc_shmem_ctx_quiet(roc_shmem_ctx_t ctx) { diff --git a/tests/functional_tests/sync_tester.cpp b/tests/functional_tests/sync_tester.cpp index 54d4a99f99..1afecaa45c 100644 --- a/tests/functional_tests/sync_tester.cpp +++ b/tests/functional_tests/sync_tester.cpp @@ -83,9 +83,21 @@ void SyncTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, roc_shmem_team_split_strided(ROC_SHMEM_TEAM_WORLD, 0, 1, n_pes, nullptr, 0, &team_sync_world_dup); +#ifdef USE_COOPERATIVE_GROUPS + void* kernelParams[] = {(void*)&loop, + (void*)&args.skip, + (void*)&timer, + (void*)&_type, + (void*)&_shmem_context, + (void*)&team_sync_world_dup}; + + CHECK_HIP(hipLaunchCooperativeKernel(SyncTest, gridSize, blockSize, + kernelParams, shared_bytes, stream)); +#else hipLaunchKernelGGL(SyncTest, gridSize, blockSize, shared_bytes, stream, loop, args.skip, timer, _type, _shmem_context, team_sync_world_dup); +#endif /* USE_COOPERATIVE_GROUPS */ num_msgs = (loop + args.skip) * gridSize.x; num_timed_msgs = loop; diff --git a/tests/unit_tests/shmem_gtest.cpp b/tests/unit_tests/shmem_gtest.cpp index 3115ef5e52..8536ca638a 100644 --- a/tests/unit_tests/shmem_gtest.cpp +++ b/tests/unit_tests/shmem_gtest.cpp @@ -21,8 +21,27 @@ *****************************************************************************/ #include "gtest/gtest.h" +#include int main(int argc, char **argv) { ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); + + int initialized; + MPI_Initialized(&initialized); + if (!initialized) { + int provided; + MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided); + if (provided != MPI_THREAD_MULTIPLE) { + std::cerr << "MPI_THREAD_MULTIPLE support disabled.\n"; + } + } + + int ret_val = RUN_ALL_TESTS(); + + int finalized{0}; + MPI_Finalized(&finalized); + if (!finalized) { + MPI_Finalize(); + } + return ret_val; }