AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gda ionic: use all threads in wave operations (#295)

Browse Source 
Use all available threads for polling the cq to increase the maximum
message rate.  Even when posting a single wqe in the wave, use all
available theads for polling the cq to reserve space in the sq.

Changes were needed in the rocshmem abstraction to avoid disabling gpu
threads, like taking turns or using only the first thread in a wave or
wavefront.  To avoid breaking other gda implementations, reimplement
turn-based or single thread strategy in post_wqe_rma_turn and
post_wqe_rma_single.

Signed-off-by: Allen Hubbe <allen.hubbe@amd.com>

[ROCm/rocshmem commit: 6de67d5d7c]
This commit is contained in:
Allen Hubbe
2025-11-05 12:01:14 -05:00
committed by GitHub
parent 51cf7c6c05
commit e8a7371007
8 changed files with 114 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rocshmem/src/gda/queue_pair.cpp
+44 -13
View File
@@ -118,7 +118,39 @@ QueuePair::~QueuePair() {
/******************************************************************************
************************ PROVIDER-SPECIFIC HELPERS ***************************
*****************************************************************************/
__device__ void QueuePair::post_wqe_rma(int pe, int32_t size, uintptr_t *laddr, uintptr_t *raddr, uint8_t opcode) {
__device__ void QueuePair::post_wqe_rma(int pe, int32_t size, uintptr_t *laddr, uintptr_t *raddr, uint8_t opcode, Collectivity cy) {
switch (gda_provider_) {
#if defined(GDA_IONIC)
case GDAProvider::IONIC:
ionic_post_wqe_rma(pe, size, laddr, raddr, opcode, cy);
return;
#endif
default:
post_wqe_rma_turn(pe, size, laddr, raddr, opcode, cy);
}
}
__device__ void QueuePair::post_wqe_rma_turn(int pe, int32_t size, uintptr_t *laddr, uintptr_t *raddr, uint8_t opcode, Collectivity cy) {
if (cy == THREAD) {
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) {
post_wqe_rma_single(pe, size, laddr, raddr, opcode);
need_turn = false;
}
turns = __ballot(need_turn);
}
} else {
if (is_thread_zero_in_wave()) {
post_wqe_rma_single(pe, size, laddr, raddr, opcode);
}
}
}
__device__ void QueuePair::post_wqe_rma_single(int pe, int32_t size, uintptr_t *laddr, uintptr_t *raddr, uint8_t opcode) {
switch (gda_provider_) {
#if defined(GDA_MLX5)
case GDAProvider::MLX5:
@@ -129,11 +161,6 @@ __device__ void QueuePair::post_wqe_rma(int pe, int32_t size, uintptr_t *laddr,
case GDAProvider::BNXT:
bnxt_post_wqe_rma(pe, size, laddr, raddr, opcode);
return;
#endif
#if defined(GDA_IONIC)
case GDAProvider::IONIC:
ionic_post_wqe_rma(pe, size, laddr, raddr, opcode);
return;
#endif
default:
assert(false /* invalid nic provider */);
@@ -161,16 +188,20 @@ __device__ uint64_t QueuePair::post_wqe_amo(int pe, int32_t size, uintptr_t *rad
}
}
__device__ void QueuePair::quiet() {
__device__ void QueuePair::quiet(Collectivity cy) {
switch (gda_provider_) {
#if defined(GDA_MLX5)
case GDAProvider::MLX5:
mlx5_quiet();
if (cy == THREAD || is_thread_zero_in_wave()) {
mlx5_quiet();
}
return;
#endif
#if defined(GDA_BNXT)
case GDAProvider::BNXT:
bnxt_quiet();
if (cy == THREAD || is_thread_zero_in_wave()) {
bnxt_quiet();
}
return;
#endif
#if defined(GDA_IONIC)
@@ -186,16 +217,16 @@ __device__ void QueuePair::quiet() {
/******************************************************************************
****************************** SHMEM INTERFACE *******************************
*****************************************************************************/
__device__ void QueuePair::put_nbi(void *dest, const void *source, size_t nelems, int pe) {
__device__ void QueuePair::put_nbi(void *dest, const void *source, size_t nelems, int pe, Collectivity cy) {
uintptr_t *src = reinterpret_cast<uintptr_t*>(const_cast<void*>(source));
uintptr_t *dst = reinterpret_cast<uintptr_t*>(dest);
post_wqe_rma(pe, nelems, src, dst, gda_op_rdma_write);
post_wqe_rma(pe, nelems, src, dst, gda_op_rdma_write, cy);
}
__device__ void QueuePair::get_nbi(void *dest, const void *source, size_t nelems, int pe) {
__device__ void QueuePair::get_nbi(void *dest, const void *source, size_t nelems, int pe, Collectivity cy) {
uintptr_t *src = reinterpret_cast<uintptr_t*>(const_cast<void*>(source));
uintptr_t *dst = reinterpret_cast<uintptr_t*>(dest);
post_wqe_rma(pe, nelems, dst, src, gda_op_rdma_read);
post_wqe_rma(pe, nelems, dst, src, gda_op_rdma_read, cy);
}
__device__ int64_t QueuePair::atomic_cas(void *dest, int64_t atomic_data, int64_t atomic_cmp, int pe) {