From 594fee5791d25c53b09f125451aaa90a32fb24d5 Mon Sep 17 00:00:00 2001 From: Avinash Kethineedi Date: Mon, 29 Sep 2025 12:38:49 -0400 Subject: [PATCH] GDA 64 bit atomics APIs (#254) feat(gda): add support for * Standard atomics - atomic_CAS - atomic_fetch_CAS * Extended atomics - atomic set - atomic swap * Bitwise atomics - atomic_fetch_and - atomic_and - atomic_fetch_or - atomic_or - atomic_fetch_xor - atomic_xor [ROCm/rocshmem commit: 98323a608651fc4d2b91d9344d0771b6cbd582e4] --- .../src/gda/context_gda_tmpl_device.hpp | 143 ++++++++++++++---- 1 file changed, 111 insertions(+), 32 deletions(-) diff --git a/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp b/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp index 5a2d91175b..9c3790506c 100644 --- a/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp +++ b/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp @@ -108,70 +108,141 @@ __device__ void GDAContext::amo_add(void *dst, T value, int pe) { template __device__ void GDAContext::amo_set(void *dst, T value, int pe) { - if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_set not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t - uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; - T ret_val; - T cond = 0; - for (int i = 0; i < WF_SIZE; i++) { //TODO: this looks wrong - while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, cond, pe))) { - if (ret_val == cond) { break; } - cond = ret_val; - } - } + amo_swap(dst, value, pe); } template __device__ T GDAContext::amo_swap(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_swap not implemented\n"); - abort(); - return 0; + if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_set not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t + uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; + bool need_turn {true}; + uint64_t turns = __ballot(need_turn); + T ret_val; + T cond = 0; + while (turns) { + uint8_t lane = __ffsll((unsigned long long)turns) - 1; + int pe_turn = __shfl(pe, lane); + if (pe_turn == pe) { + /** + * Guess that the remote memory is zero by setting condition to zero. + * The compare-and-swap loop will execute at least twice if wrong. + * It may run additional times if contention on memory location. + */ + while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, + cond, pe)) != cond) { + cond = ret_val; + } + need_turn = false; + } + turns = __ballot(need_turn); + } + return ret_val; } template __device__ T GDAContext::amo_fetch_and(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_fetch_and not implemented\n"); - abort(); - return 0; + if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fetch_and not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t + uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; + bool need_turn {true}; + uint64_t turns = __ballot(need_turn); + T ret_val; + T cond = 0; + T desired_val = cond & value; + while (turns) { + uint8_t lane = __ffsll((unsigned long long)turns) - 1; + int pe_turn = __shfl(pe, lane); + if (pe_turn == pe) { + while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, + desired_val, cond, pe)) != cond) { + cond = ret_val; + desired_val = ret_val & value; + } + need_turn = false; + } + turns = __ballot(need_turn); + } + return ret_val; } template __device__ void GDAContext::amo_and(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_and not implemented\n"); - abort(); + amo_fetch_and(dst, value, pe); } template __device__ T GDAContext::amo_fetch_or(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_fetch_or not implemented\n"); - abort(); - return 0; + if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fetch_or not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t + uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; + bool need_turn {true}; + uint64_t turns = __ballot(need_turn); + T ret_val; + T cond = 0; + T desired_val = cond | value; + while (turns) { + uint8_t lane = __ffsll((unsigned long long)turns) - 1; + int pe_turn = __shfl(pe, lane); + if (pe_turn == pe) { + while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, + desired_val, cond, pe)) != cond) { + cond = ret_val; + desired_val = ret_val | value; + } + need_turn = false; + } + turns = __ballot(need_turn); + } + return ret_val; } template __device__ void GDAContext::amo_or(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_or not implemented\n"); - abort(); + amo_fetch_or(dst, value, pe); } template __device__ T GDAContext::amo_fetch_xor(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_fetch_xor not implemented\n"); - abort(); - return 0; + if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fetch_xor not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t + uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; + bool need_turn {true}; + uint64_t turns = __ballot(need_turn); + T ret_val; + T cond = 0; + T desired_val = cond ^ value; + while (turns) { + uint8_t lane = __ffsll((unsigned long long)turns) - 1; + int pe_turn = __shfl(pe, lane); + if (pe_turn == pe) { + while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, + desired_val, cond, pe)) != cond) { + cond = ret_val; + desired_val = ret_val ^ value; + } + need_turn = false; + } + turns = __ballot(need_turn); + } + return ret_val; } template __device__ void GDAContext::amo_xor(void *dst, T value, int pe) { - printf("rocshmem::gda:amo_xor not implemented\n"); - abort(); + amo_fetch_xor(dst, value, pe); } template __device__ void GDAContext::amo_cas(void *dst, T value, T cond, int pe) { if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_cas not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; - for (int i = 0; i < WF_SIZE; i++) { //TODO: this looks wrong - qps[pe].atomic_cas_nofetch(base_heap[pe] + L_offset, value, cond, pe); + bool need_turn {true}; + uint64_t turns = __ballot(need_turn); + while (turns) { + uint8_t lane = __ffsll((unsigned long long)turns) - 1; + int pe_turn = __shfl(pe, lane); + if (pe_turn == pe) { + qps[pe].atomic_cas_nofetch(base_heap[pe] + L_offset, value, cond, pe); + need_turn = false; + } + turns = __ballot(need_turn); } } @@ -198,9 +269,17 @@ template __device__ T GDAContext::amo_fetch_cas(void *dst, T value, T cond, int pe) { if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fcas not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t uint64_t L_offset = reinterpret_cast(dst) - base_heap[my_pe]; + bool need_turn {true}; + uint64_t turns = __ballot(need_turn); T ret_val; - for (int i = 0; i < WF_SIZE; i++) { - ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, cond, pe); + while (turns) { + uint8_t lane = __ffsll((unsigned long long)turns) - 1; + int pe_turn = __shfl(pe, lane); + if (pe_turn == pe) { + ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, cond, pe); + need_turn = false; + } + turns = __ballot(need_turn); } return ret_val; }