mirror of
https://github.com/paboyle/Grid.git
synced 2026-10-14 11:43:18 +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
@@ -58,6 +58,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;
|
||||
@@ -183,194 +184,6 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// Minimal multi-RHS function interface (preconditioner slot)
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template<class Field>
|
||||
class MrhsLinearFunction {
|
||||
public:
|
||||
virtual void operator()(std::vector<Field> &in, std::vector<Field> &out) = 0;
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// Single-polynomial multi-RHS PGCR (non-Hermitian).
|
||||
//
|
||||
// Verbatim adaptation of PrecGeneralisedConjugateResidualNonHermitian
|
||||
// to std::vector<Field>: every innerProduct / norm2 is SUMMED over the
|
||||
// RHS index, so one alpha/beta per step is shared by all RHS -- the
|
||||
// single GCR on the enlarged block-diagonal system.
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template<class Field>
|
||||
class MrhsPGCRNonHermitian {
|
||||
public:
|
||||
RealD Tolerance;
|
||||
Integer MaxIterations;
|
||||
int mmax;
|
||||
int nstep;
|
||||
int steps;
|
||||
int 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; }
|
||||
|
||||
///////////////////////////////////////////////////////////////
|
||||
// vector-of-fields linear algebra, reductions summed over rhs
|
||||
///////////////////////////////////////////////////////////////
|
||||
static RealD vnorm2(std::vector<Field> &x){
|
||||
RealD s=0.0; for(auto &f : x) s+=norm2(f); return s;
|
||||
}
|
||||
static ComplexD vinnerProduct(std::vector<Field> &x, std::vector<Field> &y){
|
||||
ComplexD s(0.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 SolverTimer;
|
||||
SolverTimer.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){
|
||||
SolverTimer.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 "<<SolverTimer.Elapsed()<<std::endl;
|
||||
// Per-RHS true residuals: the honest metric under the summed norm
|
||||
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);
|
||||
std::vector<Field> z (nrhs,grid);
|
||||
std::vector<Field> Az(nrhs,grid);
|
||||
|
||||
////////////////////////////////
|
||||
// history for flexible orthog: [mmax][nrhs]
|
||||
////////////////////////////////
|
||||
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];
|
||||
std::cout<<GridLogMessage<<std::string(level,'\t')<<" "<<name
|
||||
<<" MrhsPGCR true residual r = src - A psi "<<vnorm2(r)<<std::endl;
|
||||
}
|
||||
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;
|
||||
int peri_k = k %mmax;
|
||||
int 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); // never reached
|
||||
return cp;
|
||||
}
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// Trivial multi-RHS preconditioner
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
@@ -382,98 +195,6 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
// MultiRHS two-level V-cycle.
|
||||
//
|
||||
// Mirrors MGPreconditioner in Example_pvdagm:
|
||||
// out = in (trivial pre) [per rhs]
|
||||
// r1 = in - A out [per rhs]
|
||||
// batched blockProject -> pack -> ONE mrhs coarse PGCR -> unpack
|
||||
// -> batched blockPromote; out += correction
|
||||
// r2 = in - A out [per rhs]
|
||||
// per-RHS fine post-smoother; out += smooth(r2)
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template<class FineField, class MrhsCoarseVector, class FineSmoother>
|
||||
class MrhsTwoLevelMG : public MrhsLinearFunction<FineField> {
|
||||
public:
|
||||
typedef MrhsCoarseVector CoarseVector; // same lattice type on Coarse5d and CoarseMrhs
|
||||
|
||||
LinearOperatorBase<FineField> &_FineOperator;
|
||||
FineSmoother &_PostSmoother; // single-RHS smoother, looped
|
||||
MultiRHSBlockProject<FineField> &_Projector;
|
||||
LinearFunction<CoarseVector> &_CoarseSolve; // PGCR on the 6D mrhs field
|
||||
GridBase *_CoarseGrid; // Coarse5d (single rhs)
|
||||
GridBase *_CoarseGridMrhs; // 6D
|
||||
|
||||
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);
|
||||
std::vector<FineField> vec2(nrhs,fgrid);
|
||||
|
||||
// Trivial pre-smoother: out = in (as in Example_pvdagm with simple_fine)
|
||||
for(int r=0;r<nrhs;r++) out[r]=in[r];
|
||||
|
||||
// Residual
|
||||
for(int r=0;r<nrhs;r++){
|
||||
_FineOperator.Op(out[r],vec1[r]);
|
||||
sub(vec1[r],in[r],vec1[r]);
|
||||
}
|
||||
|
||||
// Batched fine->coarse, pack rhs into 6D field
|
||||
std::vector<CoarseVector> Csrc_split(nrhs,_CoarseGrid);
|
||||
std::vector<CoarseVector> Csol_split(nrhs,_CoarseGrid);
|
||||
CoarseVector CsrcMrhs(_CoarseGridMrhs);
|
||||
CoarseVector 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;
|
||||
|
||||
// ONE coarse solve for all rhs (GEMM coarse mults)
|
||||
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;
|
||||
|
||||
// Unpack, batched coarse->fine, add correction
|
||||
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;
|
||||
|
||||
// Residual
|
||||
for(int r=0;r<nrhs;r++){
|
||||
_FineOperator.Op(out[r],vec1[r]);
|
||||
sub(vec1[r],in[r],vec1[r]);
|
||||
}
|
||||
|
||||
// Per-RHS post-smoother (fine level has no batching win; memory-light)
|
||||
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