Masked/orthog dim CartesianRingAllReduce and CartesianRingAllGather

This commit is contained in:
Peter Boyle
2026-08-27 17:10:49 -04:00
parent 7abc19dc03
commit 792d2a8258
2 changed files with 69 additions and 3 deletions
+15 -3
View File
@@ -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<class T>
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<T> 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<Nd; d++){ // dimension 0 first: it is the fastest lex index
if ( dim>=0 && d!=dim ) continue; // single-dimension gather
int Pd = comm->_processors[d];
if ( Pd==1 ) continue;
int med = comm->_processor_coor[d];
+54
View File
@@ -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<Nd; d++){
std::vector<ComplexD> h(n); for(uint64_t i=0;i<n;i++) h[i]=Fill<ComplexD>(i,me);
std::vector<ComplexD> 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<Pd;c++){
std::vector<ComplexD> m(n); for(uint64_t i=0;i<n;i++) m[i] = (myc==c) ? h[i] : ComplexD(0.0,0.0);
grid->GlobalSumVector(&m[0],(int)n);
if ( myc==c ) ref=m;
}
deviceVector<ComplexD> dv(n); acceleratorCopyToDevice(&h[0],&dv[0],n*sizeof(ComplexD));
CartesianRingAllReduce(grid,&dv[0],n,d);
std::vector<ComplexD> out(n); acceleratorCopyFromDevice(&dv[0],&out[0],n*sizeof(ComplexD));
double worst=0.0; for(uint64_t i=0;i<n;i++) worst=std::max(worst,Mag<ComplexD>(out[i]-ref[i]));
RealD w=worst; grid->GlobalMax(w);
std::ostringstream os; os<<"orthogDim="<<d<<" (P_d="<<Pd<<") worst abs "<<w;
Report("T6 CartesianRingAllReduce orthogDim == masked GlobalSumVector", w<1.0e-12, os.str());
}
}
// T7: single-dimension gather. Reference for my line along d: pad my chunk
// at slot coor[d] of a P_d*chunk vector, mask by "other coordinates == this
// line", GlobalSumVector; loop over all lines so every rank gets its own.
{
int me=grid->ThisRank(); int Nd=grid->_ndimension;
for(int d=0; d<Nd; d++){
int Pd=grid->_processors[d]; if ( Pd==1 ) continue;
uint64_t chunk=1013, n=chunk*Pd;
int myc=grid->_processor_coor[d];
std::vector<ComplexD> pad(n,ComplexD(0.0,0.0));
for(uint64_t i=0;i<chunk;i++) pad[myc*chunk+i]=Fill<ComplexD>(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<Nd;e++){ if(e==d) continue; idx+=grid->_processor_coor[e]*stride; stride*=grid->_processors[e]; } return idx; };
int nlines=grid->ProcessorCount()/Pd, myline=lineIndex();
std::vector<ComplexD> ref(n);
for(int L=0;L<nlines;L++){
std::vector<ComplexD> m(n); for(uint64_t i=0;i<n;i++) m[i]=(myline==L)?pad[i]:ComplexD(0.0,0.0);
grid->GlobalSumVector(&m[0],(int)n);
if ( myline==L ) ref=m;
}
deviceVector<ComplexD> dv(n); acceleratorCopyToDevice(&pad[0],&dv[0],n*sizeof(ComplexD));
CartesianRingAllGather(grid,&dv[0],chunk,d);
std::vector<ComplexD> 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;