diff --git a/Grid/communicator/RingAllReduce.h b/Grid/communicator/RingAllReduce.h index ca7f999a5..b78a28039 100644 --- a/Grid/communicator/RingAllReduce.h +++ b/Grid/communicator/RingAllReduce.h @@ -156,19 +156,31 @@ inline int CartesianLexIndex(CartesianCommunicator *comm) return idx; } +// +// CartesianRingAllGather(comm, buf, chunk, dim) dim in 0..Nd-1: gather along +// ONE processor dimension only. buf holds P_dim*chunk elements; on entry my +// chunk is at buf[coor[dim]*chunk], on exit block c is the chunk of the rank +// at coordinate c along dim with all other coordinates equal to mine. Every +// rank of the line holds the same result (e.g. dim=3 after a +// CartesianRingAllReduce(orthogDim=3): the boss of each spatial line can then +// write the P_t-times-longer vector, not the P-times-longer one). +///////////////////////////////////////////////////////////////////////////// template -void CartesianRingAllGather(CartesianCommunicator *comm, T *buf, uint64_t chunk) +void CartesianRingAllGather(CartesianCommunicator *comm, T *buf, uint64_t chunk, int dim=-1) { int P = comm->ProcessorCount(); - if ( P==1 || chunk==0 ) return; int Nd = comm->_ndimension; - int mylex = CartesianLexIndex(comm); + GRID_ASSERT( dim >= -1 && dim < Nd ); + if ( dim >= 0 ) P = comm->_processors[dim]; // ranks in my line along dim + if ( P==1 || chunk==0 ) return; + int mylex = (dim<0) ? CartesianLexIndex(comm) : comm->_processor_coor[dim]; deviceVector work((uint64_t)P*chunk); // ping-pong between buf and work; the held block lives at offset `off` in `cur` T *cur = buf; uint64_t off = (uint64_t)mylex*chunk; T *oth = &work[0]; uint64_t blk = chunk; // elements in the held block for(int d=0; d=0 && d!=dim ) continue; // single-dimension gather int Pd = comm->_processors[d]; if ( Pd==1 ) continue; int med = comm->_processor_coor[d]; diff --git a/tests/debug/Test_ring_allreduce.cc b/tests/debug/Test_ring_allreduce.cc index 9158d3bfa..c6f110fb0 100644 --- a/tests/debug/Test_ring_allreduce.cc +++ b/tests/debug/Test_ring_allreduce.cc @@ -150,6 +150,60 @@ int main(int argc, char **argv) } } + // T6: orthogDim -- CartesianRingAllReduce(.., orthogDim=d) must equal the sum + // over ranks sharing my coordinate in d. Reference: for each value c of that + // coordinate, GlobalSumVector of (my data if my coord==c else 0); keep c=mine. + { + int me=grid->ThisRank(); int Nd=grid->_ndimension; + uint64_t n=4099; + for(int d=-1; d h(n); for(uint64_t i=0;i(i,me); + std::vector ref(n,ComplexD(0.0,0.0)); + int Pd = (d<0) ? 1 : grid->_processors[d]; + int myc = (d<0) ? 0 : grid->_processor_coor[d]; + for(int c=0;c m(n); for(uint64_t i=0;iGlobalSumVector(&m[0],(int)n); + if ( myc==c ) ref=m; + } + deviceVector dv(n); acceleratorCopyToDevice(&h[0],&dv[0],n*sizeof(ComplexD)); + CartesianRingAllReduce(grid,&dv[0],n,d); + std::vector out(n); acceleratorCopyFromDevice(&dv[0],&out[0],n*sizeof(ComplexD)); + double worst=0.0; for(uint64_t i=0;i(out[i]-ref[i])); + RealD w=worst; grid->GlobalMax(w); + std::ostringstream os; os<<"orthogDim="<ThisRank(); int Nd=grid->_ndimension; + for(int d=0; d_processors[d]; if ( Pd==1 ) continue; + uint64_t chunk=1013, n=chunk*Pd; + int myc=grid->_processor_coor[d]; + std::vector pad(n,ComplexD(0.0,0.0)); + for(uint64_t i=0;i(i,me); + // line index = lex index of my coordinates with dim d removed + auto lineIndex=[&](void){ int idx=0,stride=1; for(int e=0;e_processor_coor[e]*stride; stride*=grid->_processors[e]; } return idx; }; + int nlines=grid->ProcessorCount()/Pd, myline=lineIndex(); + std::vector ref(n); + for(int L=0;L m(n); for(uint64_t i=0;iGlobalSumVector(&m[0],(int)n); + if ( myline==L ) ref=m; + } + deviceVector dv(n); acceleratorCopyToDevice(&pad[0],&dv[0],n*sizeof(ComplexD)); + CartesianRingAllGather(grid,&dv[0],chunk,d); + std::vector out(n); acceleratorCopyFromDevice(&dv[0],&out[0],n*sizeof(ComplexD)); + RealD diff=(memcmp(&out[0],&ref[0],n*sizeof(ComplexD))!=0)?1.0:0.0; grid->GlobalSum(diff); + Report("T7 CartesianRingAllGather(dim="+std::to_string(d)+") bitwise == masked reference, P_d="+std::to_string(Pd), diff==0.0); + } + } + // T4 timing at 16 MB of ComplexF (the dense-apply size at 12 RHS is 13.3 MB) { uint64_t n = 2*1024*1024;