From 7abc19dc03700ddf13443a66289cbd0b5496f384 Mon Sep 17 00:00:00 2001 From: Peter Boyle Date: Thu, 27 Aug 2026 17:03:14 -0400 Subject: [PATCH] Allow subsets of dimensions for spatial only sums --- Grid/communicator/RingAllReduce.h | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/Grid/communicator/RingAllReduce.h b/Grid/communicator/RingAllReduce.h index 310f28567..ca7f999a5 100644 --- a/Grid/communicator/RingAllReduce.h +++ b/Grid/communicator/RingAllReduce.h @@ -27,8 +27,15 @@ NAMESPACE_BEGIN(Grid); // communicator-level primitive: needs only CartesianCommunicator. // // RingAllReduce(comm, buf, n) flat ring over all P ranks -// CartesianRingAllReduce(comm, buf, n) ring along each processor dimension -// in turn (P_d ranks per ring) +// CartesianRingAllReduce(comm, buf, n, orthogDim=-1) +// ring along each processor dimension +// in turn (P_d ranks per ring). +// orthogDim in 0..Nd-1: that dimension is +// SKIPPED, so the result is the sum over +// all ranks sharing my coordinate in it +// (e.g. orthogDim=3: sum over each +// time-slice of processors separately). +// -1: all dimensions (full allreduce). // // Why: Cray MPICH device-buffer MPI_Allreduce aborts above ~8 MB (MPI_FLOAT, // measured 4.4 MB pass / 13.3 MB fail) and delivers 5.8 GB/s where P2P rings @@ -95,11 +102,13 @@ void RingAllReduce(CartesianCommunicator *comm, T *buf, uint64_t n) } template -void CartesianRingAllReduce(CartesianCommunicator *comm, T *buf, uint64_t n) +void CartesianRingAllReduce(CartesianCommunicator *comm, T *buf, uint64_t n, int orthogDim=-1) { if ( comm->ProcessorCount()==1 || n==0 ) return; int Nd = comm->_ndimension; + GRID_ASSERT( orthogDim >= -1 && orthogDim < Nd ); for(int d=0;d_processors[d]; if ( P==1 ) continue; int me = comm->_processor_coor[d];