Ring allgather too. Let's nail the Dense CoarseCoarse

This commit is contained in:
Peter Boyle
2026-08-26 21:05:36 -04:00
parent 8e4a304626
commit dc1ae3185a
4 changed files with 144 additions and 16 deletions
+35 -4
View File
@@ -119,6 +119,7 @@ public:
deviceVector<ComplexF> dSlab;
deviceVector<ComplexF> dX; // N x MRHS_MAX
deviceVector<ComplexF> dY; // nrows x MRHS_MAX
deviceVector<ComplexF> dG; // N x MRHS_MAX rank-major staging for the allgather (devSum==4)
deviceVector<ComplexF> dPartial; // NK x (nrows x MRHS_MAX)
deviceVector<ComplexF*> aptrs; // slab K-chunk pointers (lda = N)
deviceVector<ComplexF*> xptrs; // X K-chunk pointers (ldb = N)
@@ -252,9 +253,11 @@ public:
acceleratorCopyToDevice(&h[0],&cptrs[0],NK*sizeof(ComplexF*));
devSum = getenv("DENSE_DEVICE_SUM") ? atoi(getenv("DENSE_DEVICE_SUM")) : 0;
const char *sumName[4] = {"host allreduce","DEVICE-buffer allreduce (GPU-aware MPI)",
"DEVICE cartesian ring allreduce (P2P)","DEVICE flat ring allreduce (P2P)"};
GRID_ASSERT(devSum>=0 && devSum<=3);
const char *sumName[5] = {"host allreduce","DEVICE-buffer allreduce (GPU-aware MPI)",
"DEVICE cartesian ring allreduce (P2P)","DEVICE flat ring allreduce (P2P)",
"DEVICE cartesian ring ALLGATHER (P2P, ~8x fewer bytes than the padded allreduce)"};
GRID_ASSERT(devSum>=0 && devSum<=4);
if ( devSum==4 ) dG.resize((uint64_t)N*MRHS_MAX);
std::cout << GridLogMessage << "DenseCoarseMatrix: slab resident on device ("
<< sbytes/1024./1024. << " MB/rank), split-K NK=" << NK << " (Kc=" << Kc << "); "
<< sumName[devSum] << std::endl;
@@ -957,7 +960,34 @@ public:
int64_t Kc = N / NK;
double t1 = usecond();
double t2, t3;
if (devSum) {
if (devSum==4) {
// ALLGATHER: x is not a reduction -- every rank owns rows
// [me*nrows,(me+1)*nrows) of x and needs all of it. Only MY rows go
// host->device (nrows x nr, ~15 KB at nr=1), rank-major staged, gathered
// along the process grid, then scattered into the column-major dX
// (ld = N) the split-K GEMM reads.
const int me = grid->ThisRank();
const uint64_t chunk = (uint64_t)nrows*nr; // my block: [r][i]
{ GRID_TRACE("DenseH2D");
std::vector<ComplexF> hG(chunk);
for(int r=0;r<nr;r++)
memcpy(&hG[(uint64_t)r*nrows], &hX[(uint64_t)r*N + (uint64_t)me*nrows], nrows*sizeof(ComplexF));
acceleratorCopyToDevice(&hG[0], &dG[(uint64_t)me*chunk], chunk*sizeof(ComplexF));
}
t2 = usecond();
{ GRID_TRACE("DenseAllgather");
CartesianRingAllGather(grid, (ComplexF *)&dG[0], chunk);
// scatter [q][r][i] -> dX[r*N + q*nrows + i]
ComplexF *g = &dG[0]; ComplexF *x = &dX[0];
const int64_t nrw = nrows; const int64_t NN = N; const int nrr = nr;
accelerator_for(idx, (uint64_t)N*nr, 1, {
int64_t r = idx / NN; int64_t gi = idx - r*NN;
int64_t q = gi / nrw; int64_t i = gi - q*nrw;
x[idx] = g[q*(nrw*nrr) + r*nrw + i];
});
}
t3 = usecond();
} else if (devSum) {
{ GRID_TRACE("DenseH2D");
acceleratorCopyToDevice(&hX[0],&dX[0],nX*sizeof(ComplexF));
}
@@ -967,6 +997,7 @@ public:
// above ~8 MB: 12 RHS at N=138240 is 13.3 MB)
// DENSE_DEVICE_SUM=2 : CartesianRingAllReduce, P2P only, no size cliff
// DENSE_DEVICE_SUM=3 : flat RingAllReduce, P2P only
// DENSE_DEVICE_SUM=4 : CartesianRingAllGather (branch above)
if (devSum==2) CartesianRingAllReduce(grid,(ComplexF *)&dX[0],nX);
else if (devSum==3) RingAllReduce(grid,(ComplexF *)&dX[0],nX);
else grid->GlobalSumVector((ComplexF *)&dX[0], (int)nX);
+57
View File
@@ -116,4 +116,61 @@ void CartesianRingAllReduce(CartesianCommunicator *comm, T *buf, uint64_t n)
}
}
/////////////////////////////////////////////////////////////////////////////
// Cartesian ring ALLGATHER, point-to-point only.
//
// CartesianRingAllGather(comm, buf, chunk)
// buf holds P*chunk elements of T. On entry rank r's chunk is at
// buf[r*chunk]; on exit every rank holds all P chunks in RANK order.
//
// Dimension by dimension from the fastest-varying process coordinate
// (dim Nd-1) to the slowest: each stage is a ring over the P_d ranks of that
// line, after which the held block is the concatenation over that
// coordinate; because MPI Cartesian ranks are lexicographic with the last
// coordinate fastest, the final concatenation IS rank order -- no
// permutation. Bytes sent per rank ~ chunk*(P-1) ... dominated by the last
// stage, i.e. ~N = P*chunk total: 8x less than a zero-padded
// CartesianRingAllReduce of the same vector (which reduce-scatters AND
// gathers along every dimension). Steps: sum_d (P_d-1). Exact (no
// arithmetic): the result is bitwise the same as the padded allreduce.
//
// Written for the dense coarse-coarse apply (every rank owns rows of A^{-1}
// and needs the whole x), measured 1.86 ms for 4.4 MB at 288 ranks with the
// allreduce ring -- at wire speed, but moving 35 MB per rank to deliver 4.4.
/////////////////////////////////////////////////////////////////////////////
template<class T>
void CartesianRingAllGather(CartesianCommunicator *comm, T *buf, uint64_t chunk)
{
int P = comm->ProcessorCount();
int me = comm->ThisRank();
if ( P==1 || chunk==0 ) return;
int Nd = comm->_ndimension;
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)me*chunk;
T *oth = &work[0];
uint64_t blk = chunk; // elements in the held block
for(int d=Nd-1; d>=0; d--){
int Pd = comm->_processors[d];
if ( Pd==1 ) continue;
int med = comm->_processor_coor[d];
int next, prev;
comm->ShiftedRanks(d, 1, prev, next); // (dim, shift, source, dest)
GRID_ASSERT( (blk*sizeof(T))%4 == 0 );
// place my block in slot med of the staging area (oth[0 .. Pd*blk))
acceleratorCopyDeviceToDevice((void *)(cur+off), (void *)(oth+(uint64_t)med*blk), blk*sizeof(T));
for(int t=1;t<Pd;t++){
int sendslot = (med - t + 1 + Pd) % Pd;
int recvslot = (med - t + Pd) % Pd;
comm->SendToRecvFrom((void *)(oth+(uint64_t)sendslot*blk), next,
(void *)(oth+(uint64_t)recvslot*blk), prev, blk*sizeof(T));
}
// the staging area is the new held block
T *tmp = cur; cur = oth; oth = tmp; off = 0;
blk *= Pd;
}
GRID_ASSERT( blk == (uint64_t)P*chunk );
if ( cur != buf ) acceleratorCopyDeviceToDevice((void *)cur, (void *)buf, blk*sizeof(T));
}
NAMESPACE_END(Grid);
+15 -12
View File
@@ -4,7 +4,7 @@
#SBATCH --ntasks-per-node=8
#SBATCH --cpus-per-task=7
#SBATCH --gpus-per-node=8
#SBATCH --time=1:00:00
#SBATCH --time=1:45:00
#SBATCH --account=phy157_dwf
#SBATCH --gpu-bind=none
#SBATCH --exclusive
@@ -89,7 +89,7 @@ export DENSE_CC=1
export DENSE_APPLY_PROFILE=1
unset DENSE_CC_CHECK
export DENSE_SPLITK=128
export DENSE_DEVICE_SUM=2 # cartesian P2P ring: no 8 MB device-allreduce cliff (NRHS=12 abort)
export DENSE_DEVICE_SUM=4 # cartesian P2P ring ALLGATHER: ~8x fewer bytes than the padded allreduce (=2); no collectives, no size cliff
export GRID_ALLOC_NCACHE_LARGE=64
export NRHS=4
export PowerIterations=0
@@ -129,16 +129,19 @@ run_cell () {
grep -h "SCHUR fp64 distributed invert took\|GB/s/rank" $fname | sed 's/^Grid : Message : [0-9.]* s : //' | cut -c1-120 | head -2
}
# name Fso Fss Csn fine coarse
run_cell L0_overshoot 12 1.0 6 replay gcr # the converged overshoot (M3), now with last-call selection + refresh 5
run_cell L1_csn2 12 1.0 2 replay gcr # coarse smoother back to the banked 2 steps
run_cell L2_fso8 8 1.0 2 replay gcr
run_cell L3_fss05 8 0.5 2 replay gcr
run_cell L4_banked 6 0.5 2 replay gcr # nearest to the banked adaptive point
# Coarse replay: same lesson as the fine level -- record a DECENT polynomial
# first (overshoot: 6-step coarse smoother, shift 2.0) and back off from there.
run_cell L5_coarse6 8 0.5 6 replay replay # coarse frozen at a 6-step polynomial (overshoot)
run_cell L6_coarse4 8 0.5 4 replay replay # back off
# The (order x shift) table at Csn=2, one axis at a time from the overshoot,
# so a failing cell identifies WHICH knob it needed. Reference: 28.57 s.
# name Fso Fss Csn fine coarse
run_cell L0_overshoot 12 1.0 6 replay gcr # M3's point with last-call selection + record 8..16 + refresh 5
run_cell L1_12_10 12 1.0 2 replay gcr # coarse smoother back to 2 steps; the row/column anchor
run_cell L2_08_10 8 1.0 2 replay gcr # order axis
run_cell L3_06_10 6 1.0 2 replay gcr
run_cell L4_12_05 12 0.5 2 replay gcr # shift axis
run_cell L5_08_05 8 0.5 2 replay gcr
run_cell L6_06_05 6 0.5 2 replay gcr # nearest to the banked adaptive point (Fso6/Fss0.1)
# Coarse replay: record a DECENT polynomial first (6-step, shift 2.0) and back off.
run_cell L7_coarse6 8 1.0 6 replay replay
run_cell L8_coarse4 8 1.0 4 replay replay
echo "========================================================="
echo "summary"
+37
View File
@@ -26,6 +26,7 @@ Author: Peter Boyle <pboyle@bnl.gov>
// T2 cartesian ring == GlobalSumVector, same sweep
// T3 bitwise repeatable (deterministic order)
// T4 timing at 16 MB, both rings vs GlobalSumVector
// T5 CartesianRingAllGather bitwise == padded GlobalSumVector; timing at the dense-apply shape
//
// mpirun -n 4 ./Test_ring_allreduce --grid 16.16.16.32 --mpi 1.1.2.2
// (2D process grid so the cartesian variant exercises more than one ring)
@@ -101,6 +102,42 @@ int main(int argc, char **argv)
Check<RealF> ("RealF ", grid, 1.0e-5);
Check<ComplexF>("ComplexF", grid, 1.0e-5);
// T5: CartesianRingAllGather == zero-padded GlobalSumVector, BITWISE
// (no arithmetic in either path for disjoint chunks), all types, chunk
// sizes including 1 element and non-multiples of anything.
{
int P=grid->ProcessorCount(), me=grid->ThisRank();
for(uint64_t chunk : std::vector<uint64_t>({1,3,64,1000,65537})){
uint64_t n=chunk*P;
std::vector<ComplexD> h(n,ComplexD(0.0,0.0)), ref;
for(uint64_t i=0;i<chunk;i++) h[me*chunk+i]=Fill<ComplexD>(me*chunk+i,me);
ref=h; grid->GlobalSumVector(&ref[0],(int)n);
deviceVector<ComplexD> d(n); acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexD));
CartesianRingAllGather(grid,&d[0],chunk);
std::vector<ComplexD> out(n); acceleratorCopyFromDevice(&d[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("T5 CartesianRingAllGather bitwise == padded GlobalSumVector, ComplexD chunk="+std::to_string(chunk), diff==0.0);
}
{ uint64_t chunk=1001, n=chunk*P;
std::vector<ComplexF> h(n,ComplexF(0.0,0.0)), ref;
for(uint64_t i=0;i<chunk;i++) h[me*chunk+i]=Fill<ComplexF>(me*chunk+i,me);
ref=h; grid->GlobalSumVector(&ref[0],(int)n);
deviceVector<ComplexF> d(n); acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
CartesianRingAllGather(grid,&d[0],chunk);
std::vector<ComplexF> out(n); acceleratorCopyFromDevice(&d[0],&out[0],n*sizeof(ComplexF));
RealD diff=(memcmp(&out[0],&ref[0],n*sizeof(ComplexF))!=0)?1.0:0.0; grid->GlobalSum(diff);
Report("T5 CartesianRingAllGather bitwise, ComplexF chunk=1001", diff==0.0);
}
// timing: the dense-apply shape, N=138240 x 4 rhs of ComplexF, chunk = N*4/P
{ uint64_t chunk=(uint64_t)138240*4/P, n=chunk*P;
deviceVector<ComplexF> d(n); std::vector<ComplexF> h(n,ComplexF(1.0,0.0)); acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
double t0=usecond(); CartesianRingAllGather(grid,&d[0],chunk); double t1=usecond();
acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
double t2=usecond(); CartesianRingAllReduce(grid,&d[0],n); double t3=usecond();
std::cout << GridLogMessage << "T5 timing N=138240 x 4 ComplexF (" << n*8/1.0e6 << " MB): allgather " << (t1-t0)/1000. << " ms, cartesian allreduce " << (t3-t2)/1000. << " ms" << std::endl;
}
}
// T4 timing at 16 MB of ComplexF (the dense-apply size at 12 RHS is 13.3 MB)
{
uint64_t n = 2*1024*1024;