Gather based distributed dense inverse

Grid/algorithms/multigrid/RecursiveSchurInverse.h

@@ -107,7 +107,9 @@ public:
   // Phase timers/counters, reported by ReportTelemetry
   double                tMemset;       // device panel zero-fill
   double                tDeposit;      // owner rows -> panel (device kernel)
-  double                tAllreduce;    // GlobalSumVector on device panels
+  double                tAllreduce;    // panel collective (see tBarrier)
+  double                tBarrier;      // arrival skew, when DENSE_BARRIER_PROBE
+  double                tRepack;       // rank-major -> panel (gather path)
   double                tGemm;         // strided gemm + synchronise
   double                tLeaf;         // leaf inversions
   double                tARmin;        // fastest single panel collective
@@ -119,6 +121,28 @@ public:
   // Persistent grow-only device panel; assembly and collectives are
   // device-resident.  Device builds require GPU-aware MPI.
   deviceVector<ComplexD> dPanelBuf;
+  // Rank-major receive staging for the AllGatherV path, plus the per-panel-row
+  // owner maps that drive the repack.  Grow-only, same discipline as dPanelBuf:
+  // a fresh device allocation per collective is catastrophic (measured).
+  deviceVector<ComplexD> dRecvBuf;
+  deviceVector<int64_t>  dOwnerOff;    // panel row -> owner's first panel row
+  deviceVector<int64_t>  dOwnerRows;   // panel row -> owner's row count
+
+  // DENSE_GATHER            : 1 => AllGatherV in place of zero-fill+GlobalSum
+  // DENSE_GATHER_MIN_BYTES  : panels below this stay on the allreduce path.
+  //                           Default 0 (always gather).  The Schur tree puts
+  //                           94% of its bytes in panels with ~3.7 MB per-rank
+  //                           messages, i.e. squarely bandwidth bound; only the
+  //                           two deepest levels (1.2% of bytes, 16-65 KB per
+  //                           rank) are latency bound.  This exists so that
+  //                           tail can be excluded with an env var rather than
+  //                           a rebuild, if it ever proves to matter.
+  // DENSE_BARRIER_PROBE     : 1 => Barrier() before each collective, so that
+  //                           arrival skew lands in tBarrier and tAllreduce
+  //                           measures transfer alone.
+  int                    useGather;
+  int64_t                gatherMinBytes;
+  int                    barrierProbe;
 
   ///////////////////////////////////////////////////////////////////////////
   // Ownership-table validation: a proper partition of [0,N).
@@ -179,6 +203,11 @@ public:
     myNrows = rowStart[me+1] - rowStart[me];
 
     telLeafMaxInv = 0.0;
+
+    useGather      = getenv("DENSE_GATHER") ? atoi(getenv("DENSE_GATHER")) : 0;
+    gatherMinBytes = getenv("DENSE_GATHER_MIN_BYTES")
+                   ? atol(getenv("DENSE_GATHER_MIN_BYTES")) : 0;
+    barrierProbe   = getenv("DENSE_BARRIER_PROBE") ? 1 : 0;
   }
 
   ///////////////////////////////////////////////////////////////////////////
@@ -236,75 +265,145 @@ public:
       GRID_ASSERT( colB + widthB <= B.cols );
     }
 
-    int64_t kc = panelBytes / ( (int64_t)sizeof(ComplexD) * n );
+    // COLUMN chunking, not row chunking.  Each rank's contribution to a
-    if ( kc < 1 ) kc = 1;
+    // column chunk is rows_r x nchunk at ld = B.rows -- i.e. exactly the
-    if ( kc > k ) kc = k;
+    // contiguous window B.ColumnWindow(colB+j0) -- so the AllGatherV send
-    GRID_ASSERT( kc*n < 2147483647L );   // GlobalSumVector count is int
+    // needs no pack.  Row chunking would slice each column and force one.
+    // It also makes every chunk a full-k product, so beta applies directly
+    // and the cross-chunk accumulation disappears.
+    int64_t bufs = useGather ? 2 : 1;         // panel, plus receive staging
+    int64_t nc = panelBytes / ( bufs * (int64_t)sizeof(ComplexD) * k );
+    if ( nc < 1 ) nc = 1;
+    if ( nc > n ) nc = n;
+    GRID_ASSERT( k*nc < 2147483647L );        // collective counts are int
 
     deviceVector<ComplexD> &dPanel = dPanelBuf;
-    if ( dPanel.size() < (uint64_t)kc*n ) dPanel.resize((uint64_t)kc*n);
+    if ( dPanel.size() < (uint64_t)k*nc ) dPanel.resize((uint64_t)k*nc);
     deviceVector<ComplexD*> ap(1);
     deviceVector<ComplexD*> bp(1);
     deviceVector<ComplexD*> cp(1);
     std::vector<ComplexD*>  ptr(1);
 
-    for(int64_t k0=0; k0<k; k0+=kc)
+    // Panel-row -> owner maps for the repack.  Depend only on [rB0,rB1),
+    // so they are built once per call rather than once per chunk.
+    if ( useGather )
     {
-      int64_t kchunk = std::min(kc, k-k0);
+      if ( dRecvBuf.size()   < (uint64_t)k*nc ) dRecvBuf.resize((uint64_t)k*nc);
-
+      if ( dOwnerOff.size()  < (uint64_t)k )    dOwnerOff.resize((uint64_t)k);
-      // Zeroing must complete before MPI reads the panel; it is the sole
+      if ( dOwnerRows.size() < (uint64_t)k )    dOwnerRows.resize((uint64_t)k);
-      // producer on non-owner ranks.
+      std::vector<int64_t> hoff(k), hrows(k);
-      tMemset -= usecond();
+      for(int r=rB0; r<rB1; r++)
-      acceleratorMemSet(&dPanel[0], 0, (uint64_t)kchunk*n*sizeof(ComplexD));
-      tMemset += usecond();
-      if ( owner )
       {
-        int64_t i0 = std::max(k0,        myOff);
+        int64_t off_r  = rowStart[r]   - rowStart[rB0];
-        int64_t i1 = std::min(k0+kchunk, myOff+B.rows);
+        int64_t rows_r = rowStart[r+1] - rowStart[r];
-        if ( i1 > i0 )
+        for(int64_t i=0;i<rows_r;i++){ hoff[off_r+i]=off_r; hrows[off_r+i]=rows_r; }
+      }
+      acceleratorCopyToDevice(&hoff[0], &dOwnerOff[0], (uint64_t)k*sizeof(int64_t));
+      acceleratorCopyToDevice(&hrows[0],&dOwnerRows[0],(uint64_t)k*sizeof(int64_t));
+      if ( owner ) GRID_ASSERT( B.rows == rowStart[me+1]-rowStart[me] );
+    }
+
+    for(int64_t j0=0; j0<n; j0+=nc)
+    {
+      int64_t nchunk = std::min(nc, n-j0);
+      uint64_t panelWords = (uint64_t)k*nchunk;
+      uint64_t panelBytesThis = panelWords*sizeof(ComplexD);
+
+      int gatherThis = useGather && ( (int64_t)panelBytesThis >= gatherMinBytes );
+
+      // Arrival skew is charged to the barrier, so that tAllreduce measures
+      // transfer alone.  Diagnostic only; off by default.
+      if ( barrierProbe ) { double tb = -usecond(); grid->Barrier(); tBarrier += tb+usecond(); }
+
+      double tar = -usecond();
+      if ( gatherThis )
+      {
+        //////////////////////////////////////////////////////////////////
+        // AllGatherV: move the payload once, no arithmetic, no zero fill.
+        // Receive is rank major -- a concatenation of rows_r x nchunk
+        // column-major blocks -- which is the correct ROW order but the
+        // wrong LAYOUT, so one device repack follows.
+        //////////////////////////////////////////////////////////////////
+        std::vector<int> counts(P,0), displs(P,0);
+        for(int r=rB0; r<rB1; r++)
+        {
+          counts[r] = (int)((rowStart[r+1]-rowStart[r])*nchunk);
+          displs[r] = (int)((rowStart[r]  -rowStart[rB0])*nchunk);
+        }
+        void *send = owner ? (void *)B.ColumnWindow(colB+j0) : (void *)&dRecvBuf[0];
+        grid->AllGatherV(send, counts[me],
+                         (void *)&dRecvBuf[0], counts, displs, sizeof(ComplexD));
+        tar += usecond();
+
+        tRepack -= usecond();
+        {
+          ComplexD *pan = &dPanel[0];
+          ComplexD *rcv = &dRecvBuf[0];
+          int64_t  *ooff  = &dOwnerOff[0];
+          int64_t  *orows = &dOwnerRows[0];
+          int64_t   kk = k, ncw = nchunk;
+          accelerator_for(idx, panelWords, 1, {
+            int64_t j = idx / kk;
+            int64_t p = idx - j*kk;
+            int64_t o = ooff[p];
+            int64_t rr= orows[p];
+            pan[idx] = rcv[(uint64_t)(o*ncw + (p-o) + j*rr)];
+          });
+        }
+        tRepack += usecond();
+      }
+      else
+      {
+        //////////////////////////////////////////////////////////////////
+        // Zero fill + deposit + GlobalSumVector.  Exact because the zero
+        // fill leaves exactly one contributing rank per element, but it
+        // moves the payload twice and reduces over zeros.
+        //////////////////////////////////////////////////////////////////
+        tar += usecond();
+        tMemset -= usecond();
+        acceleratorMemSet(&dPanel[0], 0, panelBytesThis);
+        tMemset += usecond();
+        if ( owner )
         {
-          // Deposit my rows: strided block copy, len rows x n columns
-          int64_t   len   = i1-i0;
           int64_t   brows = B.rows;
-          int64_t   dof   = i0-k0;
+          int64_t   kk    = k;
-          int64_t   sof   = i0-myOff;
+          int64_t   dof   = myOff;
-          ComplexD *src   = B.ColumnWindow(colB);
+          ComplexD *src   = B.ColumnWindow(colB+j0);
           ComplexD *dst   = &dPanel[0];
           tDeposit -= usecond();
-          accelerator_for(idx, (uint64_t)(len*n), 1, {
+          accelerator_for(idx, (uint64_t)(brows*nchunk), 1, {
-            int64_t j = idx / len;
+            int64_t j = idx / brows;
-            int64_t i = idx % len;
+            int64_t i = idx - j*brows;
-            dst[(uint64_t)(dof + i + j*kchunk)] = src[(uint64_t)(sof + i + j*brows)];
+            dst[(uint64_t)(dof + i + j*kk)] = src[(uint64_t)(i + j*brows)];
           });
           tDeposit += usecond();
         }
+        tar = -usecond();
+        grid->GlobalSumVector(&dPanel[0], (int)panelWords);
+        tar += usecond();
       }
-      double tar = -usecond();
-      grid->GlobalSumVector(&dPanel[0], (int)(kchunk*n));
-      tar += usecond();
       tAllreduce += tar;
       tARmin = std::min(tARmin, tar);
       tARmax = std::max(tARmax, tar);
-      bytesAllreduce += (uint64_t)kchunk*n*sizeof(ComplexD);
+      bytesAllreduce += panelBytesThis;
       nAllreduce++;
 
       if ( m > 0 )
       {
-        ComplexD beta_use = ( k0==0 ) ? beta : ComplexD(1.0,0.0);
+        // Full-k product into this column chunk of C: beta applies directly.
-
+        ptr[0] = A.ColumnWindow(colA);
-        ptr[0] = A.ColumnWindow(colA + k0);
         acceleratorCopyToDevice(&ptr[0], &ap[0], sizeof(ComplexD*));
         ptr[0] = &dPanel[0];
         acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*));
-        ptr[0] = C.ColumnWindow(colC);
+        ptr[0] = C.ColumnWindow(colC+j0);
         acceleratorCopyToDevice(&ptr[0], &cp[0], sizeof(ComplexD*));
 
         tGemm -= usecond();
         BLAS.gemmBatched(GridBLAS_OP_N, GridBLAS_OP_N,
-                         (int)m, (int)n, (int)kchunk,
+                         (int)m, (int)nchunk, (int)k,
                          alpha,    ap, (int)A.ld,
-                                   bp, (int)kchunk,
+                                   bp, (int)k,
-                         beta_use, cp, (int)C.ld);
+                         beta,     cp, (int)C.ld);
         BLAS.synchronise();
         tGemm += usecond();
       }
@@ -552,6 +651,8 @@ public:
     tMemset        = 0.0;
     tDeposit       = 0.0;
     tAllreduce     = 0.0;
+    tBarrier       = 0.0;
+    tRepack        = 0.0;
     tGemm          = 0.0;
     tLeaf          = 0.0;
     tARmin         = 1.0e30;
@@ -596,10 +697,14 @@ public:
               << "  memset "    << tMemset/1.0e6
               << "  deposit "   << tDeposit/1.0e6
               << "  allreduce " << tAllreduce/1.0e6
+              << "  barrier "   << tBarrier/1.0e6
+              << "  repack "    << tRepack/1.0e6
               << "  gemm "      << tGemm/1.0e6
               << "  leaf "      << tLeaf/1.0e6
               << std::endl;
     std::cout << GridLogMessage << "Schur comms:"
+              << ( useGather ? "  [AllGatherV]" : "  [zero-fill+GlobalSum]" )
+              << ( barrierProbe ? " [barrier probe on]" : "" )
               << "  GatherGemm calls " << nGatherGemm
               << "  panel allreduces " << nAllreduce
               << "  allreduce GB "     << bytesAllreduce/1024./1024./1024.
@@ -610,6 +715,8 @@ public:
               << "  min " << (nAllreduce ? tARmin/1.0e3 : 0.0) << " ms"
               << "  avg " << (nAllreduce ? tAllreduce/nAllreduce/1.0e3 : 0.0) << " ms"
               << "  max " << tARmax/1.0e3 << " ms"
+              << "   effective " << (tAllreduce>0 ? bytesAllreduce/tAllreduce*1.0e6/1.0e9 : 0.0)
+              << " GB/s"
               << std::endl;
   }
 };

Grid/communicator/Communicator_base.h

@@ -249,6 +249,23 @@ public:
                  void *out,const std::vector<int> &recvcounts,const std::vector<int> &recvdispls,
                  uint64_t bytes);
 
+  ////////////////////////////////////////////////////////////////////////////
+  // Gather to all. Every rank contributes sendcount words from "in"; every
+  // rank receives the concatenation of all contributions in rank order.
+  // Counts and displacements are in units of "bytes" sized words and are
+  // indexed by rank within this communicator; recvcounts[ThisRank()] must
+  // equal sendcount.  AllGather below is the uniform count special case.
+  //
+  // This is the primitive form of the zero-fill + GlobalSumVector idiom used
+  // wherever each element of the result has exactly one contributing rank.
+  // That idiom moves the payload twice and reduces over zeros; this moves it
+  // once and performs no arithmetic.
+  ////////////////////////////////////////////////////////////////////////////
+  void AllGatherV(void *in ,int sendcount,
+                  void *out,const std::vector<int> &recvcounts,const std::vector<int> &recvdispls,
+                  uint64_t bytes);
+  void AllGather (void *in ,void *out,uint64_t words,uint64_t bytes);
+  
   template<class obj> void Broadcast(int root,obj &data)
   {
     Broadcast(root,(void *)&data,sizeof(data));

Grid/communicator/Communicator_mpi3.cc

@@ -966,4 +966,39 @@ void CartesianCommunicator::AllToAllV(void *in ,const std::vector<int> &sendcoun
   MPI_Type_free(&object);
 }
 
+void CartesianCommunicator::AllGatherV(void *in ,int sendcount,
+                                       void *out,const std::vector<int> &recvcounts,const std::vector<int> &recvdispls,
+                                       uint64_t bytes)
+{
+  FlightRecorder::StepLog("AllGatherV");
+  GRID_ASSERT(recvcounts.size()==(size_t)_Nprocessors);
+  GRID_ASSERT(recvdispls.size()==(size_t)_Nprocessors);
+  GRID_ASSERT(recvcounts[_processor]==sendcount);
+  // MPI counts are "int"; the caller sizes the word to keep them in range
+  int ibytes = bytes;
+  GRID_ASSERT(bytes == (uint64_t)ibytes);
+  MPI_Datatype object;
+  MPI_Type_contiguous(ibytes,MPI_BYTE,&object);
+  MPI_Type_commit(&object);
+  int ierr = MPI_Allgatherv(in ,sendcount,object,
+                            out,(int *)&recvcounts[0],(int *)&recvdispls[0],object,communicator);
+  GRID_ASSERT(ierr==0);
+  MPI_Type_free(&object);
+}
+
+void CartesianCommunicator::AllGather(void *in,void *out,uint64_t words,uint64_t bytes)
+{
+  FlightRecorder::StepLog("AllGather");
+  int iwords = words;
+  int ibytes = bytes;
+  GRID_ASSERT(words == (uint64_t)iwords);
+  GRID_ASSERT(bytes == (uint64_t)ibytes);
+  MPI_Datatype object;
+  MPI_Type_contiguous(ibytes,MPI_BYTE,&object);
+  MPI_Type_commit(&object);
+  int ierr = MPI_Allgather(in,iwords,object,out,iwords,object,communicator);
+  GRID_ASSERT(ierr==0);
+  MPI_Type_free(&object);
+}
+
 NAMESPACE_END(Grid);

Grid/communicator/Communicator_none.cc

@@ -125,6 +125,23 @@ void CartesianCommunicator::AllToAllV(void *in ,const std::vector<int> &sendcoun
         (char *)out+(uint64_t)recvdispls[0]*bytes,bytes*(uint64_t)sendcounts[0]);
 }
 
+void CartesianCommunicator::AllGatherV(void *in ,int sendcount,
+                                       void *out,const std::vector<int> &recvcounts,const std::vector<int> &recvdispls,
+                                       uint64_t bytes)
+{
+  // Single rank: the gather degenerates to a copy of our own contribution
+  GRID_ASSERT(recvcounts.size()==1);
+  GRID_ASSERT(recvdispls.size()==1);
+  GRID_ASSERT(recvcounts[0]==sendcount);
+  bcopy((char *)in,
+        (char *)out+(uint64_t)recvdispls[0]*bytes,bytes*(uint64_t)sendcount);
+}
+
+void CartesianCommunicator::AllGather(void *in,void *out,uint64_t words,uint64_t bytes)
+{
+  bcopy((char *)in,(char *)out,bytes*words);
+}
+
 int  CartesianCommunicator::RankWorld(void){return 0;}
 void CartesianCommunicator::Barrier(void){}
 void CartesianCommunicator::Broadcast(int root,void* data, uint64_t bytes) {}

skills/multigrid-design-notes.md

@@ -153,6 +153,94 @@ MPI latency: 1188 μs = 37% of 3.2 ms per call. Irreducible for single-RHS.
 Multi-RHS is the fundamental solution for throughput, but cannot be used in HMC
 (each trajectory has a new gauge field → new coarse operator).
 
+## Fine operator performance model (measured 2026-08-20, Frontier, 288 ranks)
+
+48³×96, Ls=24, mpi 3.6.4.4 (36 nodes), `--accelerator-threads 8 --shm-mpi 1 --comms-overlap`.
+
+### Measured decomposition of one Möbius `M` (roctx trace, MG smoother)
+
+| | ms |
+|---|---|
+| `DhopInternalOverlappedComms` (wall) | 5.549 |
+| ... of which self | 4.746 |
+| `hipStreamSynchronize` x7 | 2.158 |
+| whole `L1L2-Vcycle - Smoothers` step = one `M` | 7.128 |
+
+Per `M`: one `Meooe`/DW (~4.7 ms, comms dominated), two `M5D` (~0.5 ms each),
+axpy (~0.4 ms). Dslash arithmetic is only 1.4 (interior) + 0.4 (exterior) ms.
+**~63% of the sequence has MPI in flight; comms and compute are essentially
+serialised despite `--comms-overlap`.**
+
+`PVdagM.Op` = `M` + `Mdag` = 2x7.128 = 14.3 ms. `CoarsenOperator` `tmat`
+measured 30.695 s / 1980 Ops = 15.5 ms/Op, i.e. **within 9% of the fine
+operator's own in-production cost** -- the coarsening is not the problem.
+
+### Benchmark_dwf, fp64 (mflop/s per node, avg of 4 runs)
+
+| accelerator-threads | comms fp64 | comms fp32 (`SloppyComms`) |
+|---|---|---|
+| 16 | 1,769,808 | 2,842,248 |
+| **8** | **1,803,904** | **2,862,619** |
+| 4 | 1,802,937 | 2,731,611 |
+
+- **Comms precision is the dominant knob: +59%.** `--shm-mpi 1` a further 2-4%.
+- Thread count is second-order (2-5% spread). A wavefront-fill model
+  (`block = Nsimd x nt`, fp64 halves Nsimd, so nt=16 refills 64) predicts 16 and
+  is WRONG -- the Dslash is bandwidth bound, not occupancy bound. The flatness is
+  itself the disproof. 8 remains best, as tuned under fp32.
+- **Trap:** `Benchmark_dwf` defaults to `Ls=16` and takes `-Ls` (single dash).
+  Check the reported mflop/s against `1320*V5*ncall/t` before comparing to Ls=24
+  work, or you will infer a 2x deficit that is not there.
+
+### Optimisation model (per `M`, from 7.1 ms)
+
+Overlappable work = DslashInterior + DslashExterior + addQmu(noop) + axpby + M5D
++ gather = 1.4+0.4+0+0.4+0.5+0.4 ~ **3.1 ms**, against **4.7 ms** irreducible comms.
+
+| | ms | speedup |
+|---|---|---|
+| now | 7.5 | - |
+| overlap only | max(3.1,4.7) = 4.7 | 1.6x |
+| overlap + sloppy comms | 3.1 (compute bound) | 2.3x |
+
+**The two are independent.** Overlap is 1.6x at *zero* precision cost -- do it
+unconditionally. Sloppy is the further 1.4x and is the only piece carrying a
+judgement.
+
+### Where sloppy comms may be applied
+
+- **Whole V-cycle (fine smoother + coarse solves): safe.** Only the outer PGCR
+  operator application must be exact; the outer Krylov corrects the rest every
+  iteration.
+- **`CoarsenOperator`: different risk category.** Error goes into `A_c`
+  permanently and propagates to its inverse; it is not corrected by the outer
+  solver.
+  **The dense bottom's certificate does NOT detect this.** The import
+  certificate compares `Dense x` against `Op.M x`; if `A_c` was built sloppily
+  both sides carry the same error and it still reads ~1e-8. VERIFY is likewise
+  the exact inverse of whatever matrix it was handed. Both certify the import
+  and inversion, not the accuracy of `A_c`.
+  The instruments that DO see it: coarsen twice and compare `BLAS_A` element by
+  element against an exact run (the comparison `Test_coarse_v2_coarsen` already
+  performs between V1 and V2), or the outer iteration count and time to
+  solution.
+- **Does not reach the coarse levels.** `SloppyComms` sets a flag on
+  `Stencil`/`StencilEven`/`StencilOdd`; V2's `M` uses `PaddedCell::Face_exchange`,
+  which sends raw `vobj` bytes with no precision option. Coarse levels are less
+  bandwidth sensitive, so this matters less than it sounds.
+- **Object sharing caveat:** the flag is per fermion operator. A single
+  `MobiusFermionD` shared between `PVdagM` and `ShiftedPVdagM` cannot be sloppy
+  for one and exact for the other -- the preconditioner needs its own pair on the
+  same gauge field.
+
+### Pipelined mrhs (idea iii)
+
+Ripple independent vectors so `Dwbegin(n)` overlaps `Dwcomplete(n-1)`; needs
+double-buffered comms (two DWF operators alternating, or the stencil
+`preserve_shm` flag) and a larger SHM segment. **`CoarsenOperator` is the natural
+first customer**: its 1980 applications are already fully independent and the
+single-RHS driver is already the loop shape, so no algorithmic change is needed.
+
 ## Three-level perspective
 
 The chi deflation of the coarse solve can be viewed as a third level: coarsening all