Ring allgather too. Let's nail the Dense CoarseCoarse

Grid/algorithms/multigrid/DenseCoarseMatrix.h

@@ -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)",
+      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 allreduce (P2P)","DEVICE flat ring allreduce (P2P)",
-      GRID_ASSERT(devSum>=0 && devSum<=3);
+                                "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);

Grid/communicator/RingAllReduce.h

@@ -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);

systems/Frontier/smoother_modes.job

@@ -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
+# The (order x shift) table at Csn=2, one axis at a time from the overshoot,
-run_cell   L0_overshoot 12  1.0   6    replay  gcr     # the converged overshoot (M3), now with last-call selection + refresh 5
+# so a failing cell identifies WHICH knob it needed.  Reference: 28.57 s.
-run_cell   L1_csn2      12  1.0   2    replay  gcr     # coarse smoother back to the banked 2 steps
+#          name          Fso  Fss   Csn  fine    coarse
-run_cell   L2_fso8      8   1.0   2    replay  gcr
+run_cell   L0_overshoot   12   1.0   6    replay  gcr     # M3's point with last-call selection + record 8..16 + refresh 5
-run_cell   L3_fss05     8   0.5   2    replay  gcr
+run_cell   L1_12_10       12   1.0   2    replay  gcr     # coarse smoother back to 2 steps; the row/column anchor
-run_cell   L4_banked    6   0.5   2    replay  gcr     # nearest to the banked adaptive point
+run_cell   L2_08_10       8    1.0   2    replay  gcr     # order axis
-# Coarse replay: same lesson as the fine level -- record a DECENT polynomial
+run_cell   L3_06_10       6    1.0   2    replay  gcr
-# first (overshoot: 6-step coarse smoother, shift 2.0) and back off from there.
+run_cell   L4_12_05       12   0.5   2    replay  gcr     # shift axis
-run_cell   L5_coarse6    8   0.5   6    replay  replay  # coarse frozen at a 6-step polynomial (overshoot)
+run_cell   L5_08_05       8    0.5   2    replay  gcr
-run_cell   L6_coarse4    8   0.5   4    replay  replay  # back off
+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"

tests/debug/Test_ring_allreduce.cc

@@ -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;