mirror of
https://github.com/paboyle/Grid.git
synced 2026-10-08 16:58:06 +01:00
Clean up and unification of PVdagM and HDCG, mixed precision support
This commit is contained in:
Peter Boyle committed
1 parent
2043072d9c
commit
eeb56e6b39
46 files changed
+1404
-3815
No files matched your search
@@ -50,6 +50,7 @@ Author: Peter Boyle <paboyle@ph.ed.ac.uk>
|
||||
#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;
|
||||
@@ -167,206 +168,6 @@ public:
|
||||
void HermOp (const Field &in, Field &out) { Field tmp(in.Grid()); Op(in,tmp); AdjOp(tmp,out); }
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// mrhs interfaces + single-polynomial mrhs PGCR (verbatim from Example_pvdagm_mrhs.cc):
|
||||
// reductions summed over rhs -> one alpha/beta per step for the enlarged system.
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template<class Field>
|
||||
class MrhsLinearFunction {
|
||||
public:
|
||||
virtual void operator()(std::vector<Field> &in, std::vector<Field> &out) = 0;
|
||||
};
|
||||
|
||||
template<class Field>
|
||||
class MrhsPGCRNonHermitian {
|
||||
public:
|
||||
RealD Tolerance; Integer MaxIterations; int mmax,nstep,steps,level;
|
||||
int ZeroGuess = 0; int FirstCycle = 0; // caller contract: zero guess => first-cycle r0 = src
|
||||
std::string name = "Level 1";
|
||||
LinearOperatorBase<Field> &Linop;
|
||||
MrhsLinearFunction<Field> &Preconditioner;
|
||||
void Level(int lv){ name = "Level " + std::to_string(lv); level=lv; }
|
||||
void Name(std::string n){ name = n; }
|
||||
void SetZeroGuess(int z){ ZeroGuess=z; }
|
||||
MrhsPGCRNonHermitian(RealD tol,Integer maxit,LinearOperatorBase<Field> &_Linop,MrhsLinearFunction<Field> &Prec,int _mmax,int _nstep)
|
||||
: Tolerance(tol),MaxIterations(maxit),Linop(_Linop),Preconditioner(Prec),mmax(_mmax),nstep(_nstep){ level=1; }
|
||||
static RealD vnorm2(std::vector<Field> &x){ RealD s=0; for(auto &f:x) s+=norm2(f); return s; }
|
||||
static ComplexD vinnerProduct(std::vector<Field> &x,std::vector<Field> &y){ ComplexD s(0); for(int r=0;r<(int)x.size();r++) s+=innerProduct(x[r],y[r]); return s; }
|
||||
static void vaxpy(std::vector<Field> &z,ComplexD a,std::vector<Field> &x,std::vector<Field> &y){ for(int r=0;r<(int)z.size();r++) axpy(z[r],a,x[r],y[r]); }
|
||||
void vOp(std::vector<Field> &in,std::vector<Field> &out){ for(int r=0;r<(int)in.size();r++) Linop.Op(in[r],out[r]); }
|
||||
void operator()(std::vector<Field> &src,std::vector<Field> &psi){
|
||||
RealD cp,ssq,rsq; int nrhs=src.size(); GridBase *grid=src[0].Grid();
|
||||
ssq=vnorm2(src); rsq=Tolerance*Tolerance*ssq;
|
||||
std::vector<Field> r(nrhs,grid);
|
||||
GridStopWatch T; T.Start(); steps=0; FirstCycle=1;
|
||||
for(int k=0;k<MaxIterations;k++){
|
||||
cp=GCRnStep(src,psi,rsq);
|
||||
std::cout<<GridLogMessage<<std::string(level,'\t')<<" "<<name<<" MrhsPGCR("<<mmax<<","<<nstep<<") "<<steps<<" steps cp = "<<cp<<" target "<<rsq<<std::endl;
|
||||
if(cp<rsq){
|
||||
T.Stop(); vOp(psi,r); for(int rr=0;rr<nrhs;rr++) axpy(r[rr],-1.0,src[rr],r[rr]);
|
||||
RealD tr=vnorm2(r);
|
||||
std::cout<<GridLogMessage<<std::string(level,'\t')<<" "<<name<<" MrhsPGCR: Converged on iteration "<<steps
|
||||
<<" computed residual "<<std::sqrt(cp/ssq)<<" true residual "<<std::sqrt(tr/ssq)<<" target "<<Tolerance<<std::endl;
|
||||
std::cout<<GridLogMessage<<std::string(level,'\t')<<" "<<name<<" MrhsPGCR Time elapsed: Total "<<T.Elapsed()<<std::endl;
|
||||
for(int rr=0;rr<nrhs;rr++){ RealD rn=std::sqrt(norm2(r[rr])/norm2(src[rr])); std::cout<<GridLogMessage<<"MrhsPGCR per-rhs true residual["<<rr<<"] = "<<rn<<std::endl; }
|
||||
return;
|
||||
}
|
||||
}
|
||||
std::cout<<GridLogMessage<<"MrhsPGCR: did not converge"<<std::endl;
|
||||
}
|
||||
RealD GCRnStep(std::vector<Field> &src,std::vector<Field> &psi,RealD rsq){
|
||||
RealD cp; ComplexD a,b,rq; RealD zAAz; int nrhs=src.size(); GridBase *grid=src[0].Grid();
|
||||
std::vector<Field> r(nrhs,grid),z(nrhs,grid),Az(nrhs,grid);
|
||||
std::vector< std::vector<Field> > q(mmax,std::vector<Field>(nrhs,grid));
|
||||
std::vector< std::vector<Field> > p(mmax,std::vector<Field>(nrhs,grid));
|
||||
std::vector<RealD> qq(mmax);
|
||||
std::cout<<GridLogMessage<<std::string(level,'\t')<<" "<<name<<" MrhsPGCR nStep("<<nstep<<")"<<std::endl;
|
||||
if (ZeroGuess && FirstCycle) { for(int rr=0;rr<nrhs;rr++){ psi[rr]=Zero(); r[rr]=src[rr]; } }
|
||||
else { vOp(psi,Az); for(int rr=0;rr<nrhs;rr++) r[rr]=src[rr]-Az[rr]; }
|
||||
FirstCycle=0;
|
||||
Preconditioner(r,z); vOp(z,Az); zAAz=vnorm2(Az);
|
||||
p[0]=z; q[0]=Az; qq[0]=zAAz; cp=vnorm2(r);
|
||||
for(int k=0;k<nstep;k++){
|
||||
steps++; int kp=k+1, peri_k=k%mmax, peri_kp=kp%mmax;
|
||||
rq=vinnerProduct(q[peri_k],r); a=rq/qq[peri_k];
|
||||
vaxpy(psi,a,p[peri_k],psi); vaxpy(r,-a,q[peri_k],r); cp=vnorm2(r);
|
||||
std::cout<<GridLogMessage<<std::string(level,'\t')<<" "<<name<<" MrhsPGCR step["<<steps<<"] resid "<<cp<<" target "<<rsq<<std::endl;
|
||||
if((k==nstep-1)||(cp<rsq)) return cp;
|
||||
Preconditioner(r,z); vOp(z,Az); zAAz=vnorm2(Az);
|
||||
q[peri_kp]=Az; p[peri_kp]=z;
|
||||
int northog=((kp)>(mmax-1))?(mmax-1):(kp);
|
||||
for(int back=0;back<northog;back++){ int peri_back=(k-back)%mmax; GRID_ASSERT((k-back)>=0);
|
||||
b=-real(vinnerProduct(q[peri_back],Az))/qq[peri_back];
|
||||
vaxpy(p[peri_kp],b,p[peri_back],p[peri_kp]); vaxpy(q[peri_kp],b,q[peri_back],q[peri_kp]); }
|
||||
qq[peri_kp]=vnorm2(q[peri_kp]);
|
||||
}
|
||||
GRID_ASSERT(0); return cp;
|
||||
}
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// L2->L3 mrhs V-cycle: a LinearFunction on the 6D mrhs COARSE field.
|
||||
// Mirrors Example_pvdagm_mrhs.cc's MrhsTwoLevelMG one level down, and the
|
||||
// single-RHS MGPreconditioner of Example_pvdagm_3level_SVDdefl.cc:
|
||||
// out = in (trivial pre)
|
||||
// r = in - A_coarse out
|
||||
// restrict (unpack 6D coarse -> blockProject -> pack 6D coarse-coarse)
|
||||
// ONE coarse-coarse solve (L3, GEMM)
|
||||
// prolong (unpack -> blockPromote -> pack); out += correction
|
||||
// r = in - A_coarse out
|
||||
// coarse smoother (shifted 6D coarse op); out += smooth(r)
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template<class CoarseField, class CoarseCoarseField>
|
||||
class MrhsCoarseThreeLevelPrec : public LinearFunction<CoarseField> {
|
||||
public:
|
||||
LinearOperatorBase<CoarseField> &_CoarseOp; // mrhs coarse op (6D)
|
||||
LinearFunction<CoarseField> &_CoarseSmoother; // shifted 6D coarse smoother
|
||||
MultiRHSBlockProject<CoarseField> &_Projector; // L2->L3 (vector-based)
|
||||
LinearFunction<CoarseCoarseField> &_CoarseCoarseSolve; // L3 solve (6D cc)
|
||||
GridBase *_Coarse5d, *_CoarseCoarse5d, *_CoarseCoarseMrhs;
|
||||
int _nrhs;
|
||||
|
||||
MrhsCoarseThreeLevelPrec(LinearOperatorBase<CoarseField> &CoarseOp,
|
||||
LinearFunction<CoarseField> &CoarseSmoother,
|
||||
MultiRHSBlockProject<CoarseField> &Projector,
|
||||
LinearFunction<CoarseCoarseField> &CoarseCoarseSolve,
|
||||
GridBase *Coarse5d, GridBase *CoarseCoarse5d, GridBase *CoarseCoarseMrhs, int nrhs)
|
||||
: _CoarseOp(CoarseOp), _CoarseSmoother(CoarseSmoother), _Projector(Projector),
|
||||
_CoarseCoarseSolve(CoarseCoarseSolve),
|
||||
_Coarse5d(Coarse5d), _CoarseCoarse5d(CoarseCoarse5d), _CoarseCoarseMrhs(CoarseCoarseMrhs), _nrhs(nrhs) {}
|
||||
|
||||
using LinearFunction<CoarseField>::operator();
|
||||
virtual void operator()(const CoarseField &in, CoarseField &out) {
|
||||
int nrhs=_nrhs; double t;
|
||||
CoarseField vec1(in.Grid());
|
||||
CoarseField vec2(in.Grid());
|
||||
|
||||
// trivial pre-smoother
|
||||
out = in;
|
||||
|
||||
// residual (6D coarse)
|
||||
_CoarseOp.Op(out,vec1); sub(vec1,in,vec1);
|
||||
|
||||
// restrict: unpack 6D coarse -> vector<Coarse> -> blockProject -> vector<CoarseCoarse> -> pack 6D cc
|
||||
std::vector<CoarseField> csplit(nrhs,_Coarse5d);
|
||||
std::vector<CoarseCoarseField> ccsplit(nrhs,_CoarseCoarse5d);
|
||||
CoarseCoarseField CCsrc(_CoarseCoarseMrhs);
|
||||
CoarseCoarseField CCsol(_CoarseCoarseMrhs);
|
||||
|
||||
t=-usecond();
|
||||
for(int r=0;r<nrhs;r++) ExtractSliceFast(csplit[r],vec1,r,0);
|
||||
_Projector.blockProject(csplit,ccsplit);
|
||||
for(int r=0;r<nrhs;r++) InsertSliceFast(ccsplit[r],CCsrc,r,0);
|
||||
t+=usecond();
|
||||
std::cout<<GridLogMessage<<"L2->L3 restrict took "<<t/1000.0<<"ms"<<std::endl;
|
||||
|
||||
// L3 solve (6D coarse-coarse, GEMM)
|
||||
t=-usecond();
|
||||
CCsol=Zero();
|
||||
_CoarseCoarseSolve(CCsrc,CCsol);
|
||||
t+=usecond();
|
||||
std::cout<<GridLogMessage<<"L3 coarse-coarse solve took "<<t/1000.0<<"ms ("<<t/1000.0/nrhs<<"ms/rhs)"<<std::endl;
|
||||
|
||||
// prolong: unpack 6D cc -> blockPromote -> pack 6D coarse; add correction
|
||||
t=-usecond();
|
||||
for(int r=0;r<nrhs;r++) ExtractSliceFast(ccsplit[r],CCsol,r,0);
|
||||
_Projector.blockPromote(csplit,ccsplit);
|
||||
for(int r=0;r<nrhs;r++) InsertSliceFast(csplit[r],vec1,r,0);
|
||||
add(out,out,vec1);
|
||||
t+=usecond();
|
||||
std::cout<<GridLogMessage<<"L2->L3 prolong took "<<t/1000.0<<"ms"<<std::endl;
|
||||
|
||||
// residual + coarse smoother (6D coarse)
|
||||
_CoarseOp.Op(out,vec1); sub(vec1,in,vec1);
|
||||
vec2=Zero();
|
||||
_CoarseSmoother(vec1,vec2);
|
||||
add(out,out,vec2);
|
||||
}
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// L1->L2 mrhs V-cycle (verbatim from Example_pvdagm_mrhs.cc):
|
||||
// per-rhs fine smoother + batched restriction + ONE coarse solve + batched prolong.
|
||||
// The coarse solve passed in is now itself three-level (preconditioned by L2->L3).
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template<class FineField, class MrhsCoarseVector, class FineSmoother>
|
||||
class MrhsTwoLevelMG : public MrhsLinearFunction<FineField> {
|
||||
public:
|
||||
typedef MrhsCoarseVector CoarseVector;
|
||||
LinearOperatorBase<FineField> &_FineOperator;
|
||||
FineSmoother &_PostSmoother;
|
||||
MultiRHSBlockProject<FineField> &_Projector;
|
||||
LinearFunction<CoarseVector> &_CoarseSolve;
|
||||
GridBase *_CoarseGrid, *_CoarseGridMrhs;
|
||||
MrhsTwoLevelMG(LinearOperatorBase<FineField> &FineOp, FineSmoother &Post,
|
||||
MultiRHSBlockProject<FineField> &Projector, LinearFunction<CoarseVector> &CoarseSolve,
|
||||
GridBase *CoarseGrid, GridBase *CoarseGridMrhs)
|
||||
: _FineOperator(FineOp),_PostSmoother(Post),_Projector(Projector),_CoarseSolve(CoarseSolve),
|
||||
_CoarseGrid(CoarseGrid),_CoarseGridMrhs(CoarseGridMrhs){}
|
||||
virtual void operator()(std::vector<FineField> &in, std::vector<FineField> &out){
|
||||
int nrhs=in.size(); GridBase *fgrid=in[0].Grid(); double t;
|
||||
std::vector<FineField> vec1(nrhs,fgrid),vec2(nrhs,fgrid);
|
||||
for(int r=0;r<nrhs;r++) out[r]=in[r];
|
||||
for(int r=0;r<nrhs;r++){ _FineOperator.Op(out[r],vec1[r]); sub(vec1[r],in[r],vec1[r]); }
|
||||
std::vector<CoarseVector> Csrc_split(nrhs,_CoarseGrid), Csol_split(nrhs,_CoarseGrid);
|
||||
CoarseVector CsrcMrhs(_CoarseGridMrhs), CsolMrhs(_CoarseGridMrhs);
|
||||
t=-usecond();
|
||||
_Projector.blockProject(vec1,Csrc_split);
|
||||
for(int r=0;r<nrhs;r++) InsertSliceFast(Csrc_split[r],CsrcMrhs,r,0);
|
||||
t+=usecond(); std::cout<<GridLogMessage<<"Mrhs project+pack took "<<t/1000.0<<"ms"<<std::endl;
|
||||
t=-usecond(); CsolMrhs=Zero(); _CoarseSolve(CsrcMrhs,CsolMrhs); t+=usecond();
|
||||
std::cout<<GridLogMessage<<"Mrhs coarse solve took "<<t/1000.0<<"ms ("<<t/1000.0/nrhs<<"ms/rhs)"<<std::endl;
|
||||
t=-usecond();
|
||||
for(int r=0;r<nrhs;r++) ExtractSliceFast(Csol_split[r],CsolMrhs,r,0);
|
||||
_Projector.blockPromote(vec1,Csol_split);
|
||||
for(int r=0;r<nrhs;r++) add(out[r],out[r],vec1[r]);
|
||||
t+=usecond(); std::cout<<GridLogMessage<<"Mrhs unpack+promote took "<<t/1000.0<<"ms"<<std::endl;
|
||||
for(int r=0;r<nrhs;r++){ _FineOperator.Op(out[r],vec1[r]); sub(vec1[r],in[r],vec1[r]); }
|
||||
t=-usecond();
|
||||
for(int r=0;r<nrhs;r++){ vec2[r]=Zero(); _PostSmoother(vec1[r],vec2[r]); add(out[r],out[r],vec2[r]); }
|
||||
t+=usecond(); std::cout<<GridLogMessage<<"Mrhs post-smooth took "<<t/1000.0<<"ms ("<<t/1000.0/nrhs<<"ms/rhs)"<<std::endl;
|
||||
}
|
||||
};
|
||||
|
||||
int main (int argc, char ** argv)
|
||||
{
|
||||
|
||||
Reference in new issue
Block a user