portelli/Grid
1
0 Fork 0
mirror of https://github.com/paboyle/Grid.git synced 2026-10-10 09:48:06 +01:00
Code Releases Activity
Files
865c338dd2594d1e0bd6b249a13bcbd5e0b77b06
Grid/examples/Example_pvdagm_mrhs.cc
T

384 lines
16 KiB
C++
Raw Blame History

/*************************************************************************************
Grid physics library, www.github.com/paboyle/Grid
Source file: ./examples/Example_pvdagm_mrhs.cc
Copyright (C) 2026
Author: Peter Boyle <paboyle@ph.ed.ac.uk>
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.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
See the full license in the file "LICENSE" in the top level distribution directory
*************************************************************************************/
/* END LEGAL */
// MultiRHS (valence) two-level multigrid for PVdagM.
//
// Philosophy: NO block Krylov, NO per-RHS lockstep coefficients. A single
// GCR polynomial for the enlarged block-diagonal system diag(A,...,A):
// all inner products are summed over the RHS index, giving one alpha/beta
// per step shared by every RHS.
//
// Level structure mirrors Example_pvdagm:
// outer: MrhsPGCRNonHermitian on PVdagM over std::vector<LatticeFermionD>
// precon: V-cycle -- per-RHS fine post-smoother (16-step shifted GCR),
// batched restriction (MultiRHSBlockProject / GEMM),
// ONE coarse PGCR on the 6D mrhs coarse operator
// (MultiGeneralCoarsenedMatrix, GEMM mults -- the ~10x win),
// batched prolongation.
//
// The coarse operator is coarsened once with the standard single-RHS
// machinery (subspace cache reused) and imported via CopyMatrix.
//
// Memory note: outer restart history is 2*mmax*nrhs fine fields
// (49GB/field global at 48^3x96,Ls=24). Default mmax=8, nrhs=12 needs
// ~10TB for history; use OuterMmax / NRHS to fit the partition.
//
// Env vars: MASS, SUBSPACE_FILE, BLOCK (dotted e.g. 4.4.3.4),
// NRHS (default 12, multiple of Nsimd),
// FineSmootherShift, FineSmootherOrder,
// CoarseSolverTol, CoarseSolverOrder,
// OuterMmax, OuterNstep, OuterTol
#include <Grid/Grid.h>
#include <Grid/lattice/PaddedCell.h>
#include <Grid/stencil/GeneralLocalStencil.h>
#include <Grid/algorithms/iterative/PrecGeneralisedConjugateResidualNonHermitian.h>
#include <Grid/algorithms/multigrid/MrhsMultiGrid.h>
using namespace std;
using namespace Grid;
RealD FineSmootherShift = 0.1;
int FineSmootherOrder = 16;
RealD CoarseSolverTol = 0.03;
int CoarseSolverOrder = 200;
RealD OuterTol = 1.0e-8;
int OuterMmax = 8;
int OuterNstep = 8;
int Nrhs = 12;
RealD mass = 0.00078;
void ParseEnvironment(void)
{
if(getenv("MASS")) mass = atof(getenv("MASS"));
if(getenv("FineSmootherShift")) FineSmootherShift = atof(getenv("FineSmootherShift"));
if(getenv("FineSmootherOrder")) FineSmootherOrder = atoi(getenv("FineSmootherOrder"));
if(getenv("CoarseSolverTol")) CoarseSolverTol = atof(getenv("CoarseSolverTol"));
if(getenv("CoarseSolverOrder")) CoarseSolverOrder = atoi(getenv("CoarseSolverOrder"));
if(getenv("OuterTol")) OuterTol = atof(getenv("OuterTol"));
if(getenv("OuterMmax")) OuterMmax = atoi(getenv("OuterMmax"));
if(getenv("OuterNstep")) OuterNstep = atoi(getenv("OuterNstep"));
if(getenv("NRHS")) Nrhs = atoi(getenv("NRHS"));
std::cout << GridLogMessage << "PARAM: MASS " << mass << std::endl;
std::cout << GridLogMessage << "PARAM: NRHS " << Nrhs << std::endl;
std::cout << GridLogMessage << "PARAM: FineSmootherShift " << FineSmootherShift << std::endl;
std::cout << GridLogMessage << "PARAM: FineSmootherOrder " << FineSmootherOrder << std::endl;
std::cout << GridLogMessage << "PARAM: CoarseSolverTol " << CoarseSolverTol << std::endl;
std::cout << GridLogMessage << "PARAM: CoarseSolverOrder " << CoarseSolverOrder << std::endl;
std::cout << GridLogMessage << "PARAM: OuterTol " << OuterTol << std::endl;
std::cout << GridLogMessage << "PARAM: OuterMmax " << OuterMmax << std::endl;
std::cout << GridLogMessage << "PARAM: OuterNstep " << OuterNstep << std::endl;
}
template <class Field>
void saveSubspace(std::vector<Field> &subspace, std::string const fname){
#ifdef HAVE_LIME
std::cout << Grid::GridLogMessage << "Saving subspace (" << subspace.size() << " vectors) to: " << fname << std::endl;
Grid::emptyUserRecord record;
Grid::ScidacWriter SW(subspace[0].Grid()->IsBoss());
SW.open(fname);
for (int k = 0; k < (int)subspace.size(); k++)
SW.writeScidacFieldRecord(subspace[k], record);
SW.close();
#endif
}
template <class Field>
void loadSubspace(std::vector<Field> &subspace, std::string const fname){
#ifdef HAVE_LIME
std::cout << Grid::GridLogMessage << "Loading subspace (" << subspace.size() << " vectors) from: " << fname << std::endl;
Grid::emptyUserRecord record;
Grid::ScidacReader SR;
SR.open(fname);
for (int k = 0; k < (int)subspace.size(); k++)
SR.readScidacFieldRecord(subspace[k], record);
SR.close();
#endif
}
template<class Matrix,class Field>
class PVdagMLinearOperator : public LinearOperatorBase<Field> {
Matrix &_Mat;
Matrix &_PV;
public:
PVdagMLinearOperator(Matrix &Mat,Matrix &PV): _Mat(Mat),_PV(PV) {};
void OpDiag (const Field &in, Field &out) { assert(0); }
void OpDir (const Field &in, Field &out,int dir,int disp) { assert(0); }
void OpDirAll (const Field &in, std::vector<Field> &out){ assert(0); };
void Op (const Field &in, Field &out){
Field tmp(in.Grid());
_Mat.M(in,tmp);
_PV.Mdag(tmp,out);
}
void AdjOp (const Field &in, Field &out){
Field tmp(in.Grid());
_PV.M(in,tmp);
_Mat.Mdag(tmp,out);
}
void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){
HermOp(in,out);
ComplexD dot = innerProduct(in,out);
n1=real(dot);
n2=norm2(out);
}
void HermOp(const Field &in, Field &out){
Field tmp(in.Grid());
Op(in,tmp);
AdjOp(tmp,out);
}
};
template<class Matrix,class Field>
class ShiftedPVdagMLinearOperator : public LinearOperatorBase<Field> {
Matrix &_Mat;
Matrix &_PV;
public:
RealD shift;
ShiftedPVdagMLinearOperator(RealD _shift,Matrix &Mat,Matrix &PV): shift(_shift),_Mat(Mat),_PV(PV){};
void OpDiag (const Field &in, Field &out) { assert(0); }
void OpDir (const Field &in, Field &out,int dir,int disp) { assert(0); }
void OpDirAll (const Field &in, std::vector<Field> &out){ assert(0); };
void Op (const Field &in, Field &out){
Field tmp(in.Grid());
_Mat.M(in,tmp);
_PV.Mdag(tmp,out);
out = out + shift * in;
}
void AdjOp (const Field &in, Field &out){
Field tmp(in.Grid());
_PV.M(tmp,out);
_Mat.Mdag(in,tmp);
out = out + shift * in;
}
void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){ assert(0); }
void HermOp(const Field &in, Field &out){
Field tmp(in.Grid());
Op(in,tmp);
AdjOp(tmp,out);
}
};
//////////////////////////////////////////////////////////////////////
// Trivial multi-RHS preconditioner
//////////////////////////////////////////////////////////////////////
template<class Field>
class TrivialMrhsPrecon : public MrhsLinearFunction<Field> {
public:
virtual void operator()(std::vector<Field> &in, std::vector<Field> &out){
for(int r=0;r<(int)in.size();r++) out[r]=in[r];
}
};
int main (int argc, char ** argv)
{
Grid_init(&argc,&argv);
ParseEnvironment();
const int Ls=24;
RealD M5=1.8;
RealD b=1.5;
RealD c=0.5;
const int nbasis = 60;
const int nrhs = Nrhs;
GRID_ASSERT(nrhs % vComplex::Nsimd() == 0);
std::cout << GridLogMessage << "MultiRHS PVdagM MG: mass=" << mass << " Ls=" << Ls
<< " nbasis=" << nbasis << " nrhs=" << nrhs << std::endl;
std::vector<int> lat_size {48, 48, 48, 96};
GridCartesian * UGrid = SpaceTimeGrid::makeFourDimGrid(lat_size, GridDefaultSimd(Nd,vComplex::Nsimd()),GridDefaultMpi());
GridRedBlackCartesian * UrbGrid = SpaceTimeGrid::makeFourDimRedBlackGrid(UGrid);
GridCartesian * FGrid = SpaceTimeGrid::makeFiveDimGrid(Ls,UGrid);
GridRedBlackCartesian * FrbGrid = SpaceTimeGrid::makeFiveDimRedBlackGrid(Ls,UGrid);
// Blocking, env-overridable: BLOCK=2.2.2.2
Coordinate clatt = lat_size;
Coordinate Block({4,4,3,4});
if ( getenv("BLOCK") ) {
GridCmdOptionIntVector(std::string(getenv("BLOCK")),Block);
GRID_ASSERT(Block.size()==4);
}
for(int d=0;d<clatt.size();d++){
GRID_ASSERT(lat_size[d] % Block[d] == 0);
clatt[d] = lat_size[d]/Block[d];
}
std::cout << GridLogMessage << "Block " << Block << " coarse lattice " << clatt << std::endl;
GridCartesian *Coarse4d = SpaceTimeGrid::makeFourDimGrid(clatt, GridDefaultSimd(Nd,vComplex::Nsimd()),GridDefaultMpi());
GridCartesian *Coarse5d = SpaceTimeGrid::makeFiveDimGrid(1,Coarse4d);
////////////////////////////////////////////////////////////
// 6D multi-RHS coarse grid: rhs is dim 0, SIMD across rhs
// (pattern: tests/debug/Test_general_coarse_hdcg_phys48.cc)
////////////////////////////////////////////////////////////
Coordinate mpi=GridDefaultMpi();
Coordinate rhMpi ({1,1,mpi[0],mpi[1],mpi[2],mpi[3]});
Coordinate rhLatt({nrhs,1,clatt[0],clatt[1],clatt[2],clatt[3]});
Coordinate rhSimd({vComplex::Nsimd(),1, 1,1,1,1});
GridCartesian *CoarseMrhs = new GridCartesian(rhLatt,rhSimd,rhMpi);
GridParallelRNG RNG5(FGrid); RNG5.SeedFixedIntegers({5,6,7,8});
GridParallelRNG RNG4(UGrid); RNG4.SeedFixedIntegers({1,2,3,4});
LatticeGaugeField Umu(UGrid);
std::cout << GridLogMessage << "Reading gauge field" << std::endl;
FieldMetaData header;
std::string file("/ccs/home/poare/ckpoint_lat.1000");
NerscIO::readConfiguration(Umu,header,file);
MobiusFermionD Ddwf(Umu,*FGrid,*FrbGrid,*UGrid,*UrbGrid,mass,M5,b,c);
MobiusFermionD Dpv (Umu,*FGrid,*FrbGrid,*UGrid,*UrbGrid,1.0, M5,b,c);
typedef PVdagMLinearOperator<MobiusFermionD,LatticeFermionD> PVdagM_t;
typedef ShiftedPVdagMLinearOperator<MobiusFermionD,LatticeFermionD> ShiftedPVdagM_t;
typedef GeneralCoarsenedMatrix<vSpinColourVector,vTComplex,nbasis> LittleDiracOperator;
typedef MultiGeneralCoarsenedMatrix<vSpinColourVector,vTComplex,nbasis> MrhsLittleDiracOperator;
typedef LittleDiracOperator::CoarseVector CoarseVector;
typedef Aggregation<vSpinColourVector,vTComplex,nbasis> Subspace;
PVdagM_t PVdagM(Ddwf,Dpv);
ShiftedPVdagM_t ShiftedPVdagM(FineSmootherShift,Ddwf,Dpv);
NextToNearestStencilGeometry5D geom(Coarse5d);
////////////////////////////////////////////////////////////
// Subspace: load from cache or generate
////////////////////////////////////////////////////////////
std::string subspace_file = "/lustre/orion/phy157/proj-shared/phy157_dwf/paboyle/subspace_nb"
+ std::to_string(nbasis) + ".scidac";
if ( getenv("SUBSPACE_FILE") ) subspace_file = std::string(getenv("SUBSPACE_FILE"));
uint64_t file_exists = 0;
if ( UGrid->IsBoss() ) { std::ifstream f(subspace_file); file_exists = f.good() ? 1 : 0; }
UGrid->GlobalSum(file_exists);
const int cb = 0;
Subspace AggregatesGCR(Coarse5d,FGrid,cb);
if ( file_exists ) {
std::cout << GridLogMessage << "*** Loading subspace from disk ***" << std::endl;
loadSubspace(AggregatesGCR.subspace, subspace_file);
} else {
std::cout << GridLogMessage << "*** GCR subspace generation ***" << std::endl;
AggregatesGCR.CreateSubspaceGCR(RNG5,PVdagM,nbasis);
saveSubspace(AggregatesGCR.subspace, subspace_file);
}
////////////////////////////////////////////////////////////
// Coarsen once (single-RHS machinery), import into mrhs op.
// NB CoarsenOperator block-orthogonalises subspace in place;
// ImportBasis AFTER so the projector matches the coarse op.
////////////////////////////////////////////////////////////
MrhsLittleDiracOperator mrhsLittleDiracOpPV(geom,CoarseMrhs);
{
// Scope the single-RHS operator so its padded _A is freed after import.
// At small local volumes the depth-2 padded cell inflates ~8x
// (e.g. 2^4 blocking on 432 ranks: local 8x4x4x12 -> padded 12x8x8x16,
// ~23GB/GCD for _A alone -> OOM if kept alive).
LittleDiracOperator LittleDiracOpPV(geom,FGrid,Coarse5d);
LittleDiracOpPV.CoarsenOperator(PVdagM, AggregatesGCR);
mrhsLittleDiracOpPV.CopyMatrix(LittleDiracOpPV);
}
MultiRHSBlockProject<LatticeFermionD> MrhsProjector;
MrhsProjector.Allocate(nbasis,FGrid,Coarse5d);
MrhsProjector.ImportBasis(AggregatesGCR.subspace);
////////////////////////////////////////////////////////////
// Solvers
////////////////////////////////////////////////////////////
NonHermitianLinearOperator<MrhsLittleDiracOperator,CoarseVector> mrhsLinOpCoarse(mrhsLittleDiracOpPV);
TrivialPrecon<CoarseVector> simpleC;
PrecGeneralisedConjugateResidualNonHermitian<CoarseVector>
L2PGCRmrhs(CoarseSolverTol, CoarseSolverOrder/20, mrhsLinOpCoarse, simpleC, 20, 20);
L2PGCRmrhs.Level(2);
L2PGCRmrhs.Name("Couter");
L2PGCRmrhs.SetZeroGuess(1); // caller zeroes CsolMrhs
TrivialPrecon<LatticeFermionD> simple_fine;
PrecGeneralisedConjugateResidualNonHermitian<LatticeFermionD>
SmootherGCR(0.0, 1, ShiftedPVdagM, simple_fine, FineSmootherOrder, FineSmootherOrder);
SmootherGCR.Level(1);
SmootherGCR.Name("Fsmoother");
SmootherGCR.SetZeroGuess(1); // caller zeroes vec2[r]
typedef PrecGeneralisedConjugateResidualNonHermitian<LatticeFermionD> FineSmoother_t;
MrhsTwoLevelMG<LatticeFermionD,CoarseVector,FineSmoother_t>
TwoLevelPrecon(PVdagM, SmootherGCR, MrhsProjector, L2PGCRmrhs, Coarse5d, CoarseMrhs);
MrhsPGCRNonHermitian<LatticeFermionD>
L1PGCRmrhs(OuterTol, 1000, PVdagM, TwoLevelPrecon, OuterMmax, OuterNstep);
L1PGCRmrhs.Level(1);
L1PGCRmrhs.Name("Fouter");
L1PGCRmrhs.SetZeroGuess(1); // sol[r]=Zero() at source setup
////////////////////////////////////////////////////////////
// Sources and solve
////////////////////////////////////////////////////////////
std::vector<LatticeFermionD> src(nrhs,FGrid);
std::vector<LatticeFermionD> sol(nrhs,FGrid);
for(int r=0;r<nrhs;r++){
gaussian(RNG5,src[r]);
sol[r]=Zero();
}
std::cout << GridLogMessage << "**********************************************" << std::endl;
std::cout << GridLogMessage << " MultiRHS two-level solve: " << nrhs << " RHS " << std::endl;
std::cout << GridLogMessage << "**********************************************" << std::endl;
GridStopWatch w; w.Start();
L1PGCRmrhs(src,sol);
w.Stop();
std::cout << GridLogMessage << "MultiRHS solve total " << w.Elapsed()
<< " (per RHS: " << w.useconds()/1.0e6/nrhs << " s)" << std::endl;
// Independent final verification, per RHS
{
LatticeFermionD Ax(FGrid);
RealD worst=0.0;
for(int r=0;r<nrhs;r++){
PVdagM.Op(sol[r],Ax);
Ax = Ax - src[r];
RealD rn = std::sqrt(norm2(Ax)/norm2(src[r]));
std::cout << GridLogMessage << "FINAL: rhs["<<r<<"] true residual = " << rn << std::endl;
worst = std::max(worst,rn);
}
std::cout << GridLogMessage << "FINAL: worst-case residual = " << worst << std::endl;
}
std::cout << GridLogMessage << "Done" << std::endl;
Grid_finalize();
return 0;
}
View Git Blame Copy Permalink