From 82d91433c9cb609a124062f8b798f88df36fd476 Mon Sep 17 00:00:00 2001 From: Aurelien Bouteiller Date: Tue, 6 Jan 2026 20:41:51 -0500 Subject: [PATCH] Enable new a2a (pr 334) on ionic as well (#366) * Enable new a2a (pr 334) on ionic as well * Apply suggestions from AI code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/gda/context_gda_tmpl_device.hpp | 3 ++- src/gda/queue_pair.cpp | 14 ++++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/gda/context_gda_tmpl_device.hpp b/src/gda/context_gda_tmpl_device.hpp index bf6436ad2d..5e223c5121 100644 --- a/src/gda/context_gda_tmpl_device.hpp +++ b/src/gda/context_gda_tmpl_device.hpp @@ -605,7 +605,8 @@ __device__ void GDAContext::internal_broadcast(T *dst, const T *src, int nelems, template __device__ void GDAContext::alltoall(rocshmem_team_t team, T *dst, const T *src, int nelems) { - if (gda_provider_ == GDAProvider::BNXT) { + if (gda_provider_ == GDAProvider::BNXT || + gda_provider_ == GDAProvider::IONIC) { alltoall_linear_thread_puts(team, dst, src, nelems); } else { alltoall_linear(team, dst, src, nelems); diff --git a/src/gda/queue_pair.cpp b/src/gda/queue_pair.cpp index 7adf1ff41c..c3a3d4d612 100644 --- a/src/gda/queue_pair.cpp +++ b/src/gda/queue_pair.cpp @@ -176,7 +176,10 @@ __device__ void QueuePair::post_wqe_rma_single(int32_t size, uintptr_t laddr, ui case GDAProvider::BNXT: return bnxt_post_wqe_rma_single(size, laddr, raddr, opcode, ring_db); #endif +#if defined(GDA_IONIC) case GDAProvider::IONIC: + return ionic_post_wqe_rma(0 /*pe (unused)*/, size, laddr, raddr, opcode, Collectivity::THREAD); +#endif case GDAProvider::MLX5: default: assert(false /* invalid nic provider */); @@ -212,8 +215,11 @@ __device__ uint64_t QueuePair::post_wqe_amo_single(uintptr_t raddr, uint8_t opco case GDAProvider::BNXT: return bnxt_post_wqe_amo_single(raddr, opcode, atomic_data, atomic_cmp, fetching); #endif - case GDAProvider::MLX5: +#if defined(GDA_IONIC) case GDAProvider::IONIC: + return ionic_post_wqe_amo(0 /*pe (unused)*/, 8 /*size_bytes (only 8-byte atomics implemented)*/, raddr, opcode, atomic_data, atomic_cmp, fetching); +#endif + case GDAProvider::MLX5: default: assert(false /* invalid nic provider */); return 0; @@ -253,8 +259,12 @@ __device__ void QueuePair::quiet_single() { bnxt_quiet_single(); return; #endif - case GDAProvider::MLX5: +#if defined(GDA_IONIC) case GDAProvider::IONIC: + ionic_quiet(); + return; +#endif + case GDAProvider::MLX5: default: assert(false /* invalid nic provider */); }