Adding ring all reduce for large vectors

2026-08-28 21:39:35 +01:00 · 2026-08-25 12:07:11 -04:00
parent dbbd5dba35
commit 167753aa3f
7 changed files with 607 additions and 6 deletions
@@ -0,0 +1,117 @@
+/*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid
+
+    Source file: ./tests/debug/Test_multi_reduction.cc
+
+    Copyright (C) 2026
+
+Author: Peter Boyle <pboyle@bnl.gov>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    See the full license in the file "LICENSE" in the top level distribution
+    directory
+*************************************************************************************/
+/*  END LEGAL */
+
+//////////////////////////////////////////////////////////////////////////////
+// Gate for the batched Krylov linear algebra in Lattice_reduction.h:
+//   innerProductMulti  vs  innerProduct, one j at a time
+//   axpyMulti          vs  sequential axpy
+//   axpyMultiNorm      vs  sequential axpy + norm2
+// on TWO field types through the SAME code path:
+//   fine  : LatticeFermionD on the 4D grid           (vSpinColourVector)
+//   coarse: Lattice<iVector<CComplex,nbasis>> on a 5D grid with the extra
+//           dimension standing in for nrhs, as the coarse mRHS solver does.
+// Batch sizes cover every template width and the >16 chunking path.
+//
+//   mpirun -n 1 ./Test_multi_reduction --grid 8.8.8.8 --mpi 1.1.1.1
+//////////////////////////////////////////////////////////////////////////////
+
+#include <Grid/Grid.h>
+
+using namespace Grid;
+
+static int failures = 0;
+static void Report(const std::string &name, bool pass, const std::string &detail="")
+{
+  std::cout << GridLogMessage << "  " << name << (pass ? "  PASS" : "  ** FAIL **");
+  if ( detail.size() ) std::cout << "   " << detail;
+  std::cout << std::endl;
+  if ( !pass ) failures++;
+}
+static double Cabs(const ComplexD &z){ double re=z.real(), im=z.imag(); return std::sqrt(re*re+im*im); }
+
+template<class Field>
+void Check(const std::string &tag, GridBase *grid, GridParallelRNG &RNG, const std::vector<int> &batches)
+{
+  const int Mmax = 20;
+  std::vector<Field> x; for(int j=0;j<Mmax;j++){ x.emplace_back(grid); gaussian(RNG,x[j]); }
+  Field r(grid); gaussian(RNG,r);
+  Field z0(grid); gaussian(RNG,z0);
+  Field za(grid), zb(grid), d(grid);
+
+  for(int m : batches){
+    std::vector<const Field*> win(m);
+    std::vector<ComplexD> b(m);
+    for(int j=0;j<m;j++){ win[j]=&x[j]; b[j]=ComplexD(0.3*j-1.0, 0.7-0.2*j); }
+
+    // dots
+    std::vector<ComplexD> ip;
+    innerProductMulti(ip,win,r);
+    double worst=0.0;
+    for(int j=0;j<m;j++){
+      ComplexD ref = innerProduct(x[j],r);
+      worst = std::max(worst, Cabs(ip[j]-ref)/std::max(Cabs(ref),1.0e-300));
+    }
+    Report(tag+" innerProductMulti m="+std::to_string(m), worst<1.0e-12, "worst rel "+std::to_string(worst));
+
+    // update
+    za = z0; axpyMulti(za,b,win);
+    zb = z0; for(int j=0;j<m;j++) axpy(zb,b[j],x[j],zb);
+    d = za - zb;
+    double rel = std::sqrt(norm2(d)/norm2(zb));
+    Report(tag+" axpyMulti m="+std::to_string(m), rel<1.0e-12, "rel "+std::to_string(rel));
+
+    // update + norm
+    za = z0; RealD nn = axpyMultiNorm(za,b,win);
+    RealD nref = norm2(zb);
+    d = za - zb;
+    rel = std::sqrt(norm2(d)/norm2(zb));
+    double nrel = std::fabs(nn-nref)/nref;
+    Report(tag+" axpyMultiNorm m="+std::to_string(m), rel<1.0e-12 && nrel<1.0e-12,
+           "rel "+std::to_string(rel)+" norm rel "+std::to_string(nrel));
+  }
+}
+
+int main(int argc, char **argv)
+{
+  Grid_init(&argc, &argv);
+
+  GridCartesian *UGrid = SpaceTimeGrid::makeFourDimGrid(GridDefaultLatt(),
+                             GridDefaultSimd(Nd, vComplexD::Nsimd()), GridDefaultMpi());
+  const int nrhs = 3;   // stands in for the coarse mRHS dimension
+  GridCartesian *CGrid = SpaceTimeGrid::makeFiveDimGrid(nrhs, UGrid);
+
+  std::vector<int> seeds({1,2,3,4});
+  GridParallelRNG RNG4(UGrid); RNG4.SeedFixedIntegers(seeds);
+  GridParallelRNG RNG5(CGrid); RNG5.SeedFixedIntegers(seeds);
+
+  std::vector<int> batches({1,2,3,4,5,8,9,16,17,20});
+
+  Check<LatticeFermionD>("fine  ", UGrid, RNG4, batches);
+
+  const int nbasis = 8;
+  typedef iVector<iScalar<iScalar<iScalar<vComplexD>>>,nbasis> CoarseSiteVector;
+  typedef Lattice<CoarseSiteVector> CoarseField;
+  Check<CoarseField>("coarse", CGrid, RNG5, batches);
+
+  std::cout << GridLogMessage << (failures ? "Test_multi_reduction: FAILURES"
+                                           : "Test_multi_reduction: ALL PASS") << std::endl;
+  Grid_finalize();
+  return failures ? 1 : 0;
+}