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>

[ROCm/rocshmem commit: 82d91433c9]
This commit is contained in:
Aurelien Bouteiller
2026-01-06 20:41:51 -05:00
committed by GitHub
parent e38f98fad5
commit bcdf60def6
2 changed files with 14 additions and 3 deletions
@@ -605,7 +605,8 @@ __device__ void GDAContext::internal_broadcast(T *dst, const T *src, int nelems,
template <typename T>
__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);
+12 -2
View File
@@ -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 */);
}