mirror of
https://github.com/paboyle/Grid.git
synced 2026-08-26 12:29:36 +01:00
Potentially faster GCR linalg
This commit is contained in:
Peter Boyle
@@ -142,16 +142,13 @@ public:
|
||||
|
||||
RealD cp;
|
||||
ComplexD a, b;
|
||||
// ComplexD zAz;
|
||||
RealD zAAz;
|
||||
ComplexD rq;
|
||||
|
||||
GridBase *grid = src.Grid();
|
||||
|
||||
// Only r and one scratch for the restart residual; the new p/q directions
|
||||
// are produced directly in their persistent history slots.
|
||||
Field r(grid);
|
||||
Field z(grid);
|
||||
Field tmp(grid);
|
||||
Field ttmp(grid);
|
||||
Field Az(grid);
|
||||
|
||||
////////////////////////////////
|
||||
@@ -198,24 +195,19 @@ public:
|
||||
// p = Prec(r)
|
||||
/////////////////////
|
||||
|
||||
// p[0] = Prec(r), q[0] = A p[0], written straight into the history slots
|
||||
PrecTimer.Start();
|
||||
Preconditioner(r,z);
|
||||
Preconditioner(r,p[0]);
|
||||
PrecTimer.Stop();
|
||||
|
||||
MatTimer.Start();
|
||||
Linop.Op(z,Az);
|
||||
Linop.Op(p[0],q[0]);
|
||||
MatTimer.Stop();
|
||||
|
||||
LinalgTimer.Start();
|
||||
|
||||
// zAz = innerProduct(Az,psi);
|
||||
zAAz= norm2(Az);
|
||||
qq[0]= norm2(q[0]);
|
||||
|
||||
//p[0],q[0],qq[0]
|
||||
p[0]= z;
|
||||
q[0]= Az;
|
||||
qq[0]= zAAz;
|
||||
|
||||
cp =norm2(r);
|
||||
LinalgTimer.Stop();
|
||||
GCRLogLevel<< "PGCR true residual "<< sqrt(cp/SSQ) <<std::endl;
|
||||
@@ -247,32 +239,46 @@ public:
|
||||
return cp;
|
||||
}
|
||||
|
||||
// New direction written straight into its history slot: p = Prec(r), q = A p.
|
||||
PrecTimer.Start();
|
||||
Preconditioner(r,z);// solve Az = r
|
||||
Preconditioner(r,p[peri_kp]);
|
||||
PrecTimer.Stop();
|
||||
|
||||
MatTimer.Start();
|
||||
Linop.Op(z,Az);
|
||||
Linop.Op(p[peri_kp],q[peri_kp]);
|
||||
MatTimer.Stop();
|
||||
// zAz = innerProduct(Az,psi);
|
||||
zAAz= norm2(Az);
|
||||
|
||||
LinalgTimer.Start();
|
||||
|
||||
q[peri_kp]=Az;
|
||||
p[peri_kp]=z;
|
||||
|
||||
int northog = ((kp)>(mmax-1))?(mmax-1):(kp); // if more than mmax done, we orthog all mmax history.
|
||||
std::ostringstream bs;
|
||||
|
||||
// Classical Gram-Schmidt: every coefficient is taken against the
|
||||
// UN-updated new q (so all northog inner products are independent and
|
||||
// batchable), then the window is applied. The coefficient is complex:
|
||||
// for a non-Hermitian operator <q_j,Aq> is complex and keeping only the
|
||||
// real part left the q's non-orthogonal.
|
||||
// Batched: one fused kernel + one reduction for all coefficients,
|
||||
// one fused pass per update (independent of mmax).
|
||||
std::vector<const Field*> qwin(northog), pwin(northog);
|
||||
for(int back=0;back<northog;back++){
|
||||
|
||||
int peri_back=(k-back)%mmax; GRID_ASSERT((k-back)>=0);
|
||||
|
||||
b=-real(innerProduct(q[peri_back],Az))/qq[peri_back];
|
||||
p[peri_kp]=p[peri_kp]+b*p[peri_back];
|
||||
q[peri_kp]=q[peri_kp]+b*q[peri_back];
|
||||
if ( LogCoeffs ) bs<<" b["<<back<<"]="<<b;
|
||||
|
||||
GRID_ASSERT(peri_back!=peri_kp);
|
||||
qwin[back] = &q[peri_back];
|
||||
pwin[back] = &p[peri_back];
|
||||
}
|
||||
std::vector<ComplexD> bcoef;
|
||||
innerProductMulti(bcoef,qwin,q[peri_kp]);
|
||||
for(int back=0;back<northog;back++){
|
||||
int peri_back=(k-back)%mmax;
|
||||
bcoef[back] = -bcoef[back]/qq[peri_back];
|
||||
if ( LogCoeffs ) bs<<" b["<<back<<"]="<<bcoef[back];
|
||||
}
|
||||
if ( northog ) {
|
||||
axpyMulti(p[peri_kp],bcoef,pwin);
|
||||
qq[peri_kp]=axpyMultiNorm(q[peri_kp],bcoef,qwin);
|
||||
} else {
|
||||
qq[peri_kp]=norm2(q[peri_kp]);
|
||||
}
|
||||
if ( LogCoeffs && northog ) {
|
||||
GCRLogLevel<<"coeff["<<k<<"]"<<bs.str()<<std::endl;
|
||||
|
||||
@@ -793,6 +793,151 @@ static void sliceInnerProductMatrix( Eigen::MatrixXcd &mat, const Lattice<vobj>
|
||||
delete SliceGrid;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////
|
||||
// Batched linear algebra for Krylov orthogonalisation (GCR history windows).
|
||||
//
|
||||
// innerProductMulti(out,left,right): out[j] = <left[j],right> for all j in
|
||||
// ONE kernel (right read once), ONE device reduction / sync (the per-site
|
||||
// partials are an iVector over the batch) and ONE GlobalSumVector.
|
||||
// axpyMulti(z,b,x): z = z + sum_j b[j] x[j] in ONE pass.
|
||||
// axpyMultiNorm(z,b,x): same, returning global |z|^2 from the
|
||||
// same pass (one reduction, one GlobalSum).
|
||||
//
|
||||
// The batch width is a template parameter chosen at runtime from {2,4,8,16}
|
||||
// so a short window (mmax=2 smoother) does not pay for 16 partial lanes;
|
||||
// windows longer than 16 are processed in chunks of 16, one reduction each.
|
||||
// Same code path for every Lattice<vobj>: fine fermion fields and coarse
|
||||
// multi-RHS fields (nrhs folded into the grid) alike.
|
||||
/////////////////////////////////////////////////////////////////////////////
|
||||
template<int B,class vobj>
|
||||
void rankInnerProductMultiChunk(ComplexD *out,int m,
|
||||
const std::vector<const Lattice<vobj>*> &left,
|
||||
const Lattice<vobj> &right)
|
||||
{
|
||||
typedef decltype(innerProductD(vobj(),vobj())) inner_t;
|
||||
typedef iVector<inner_t,B> batch_t;
|
||||
typedef decltype(right.View(AcceleratorRead)) View;
|
||||
|
||||
GRID_ASSERT(m>=1 && m<=B);
|
||||
GridBase *grid = right.Grid();
|
||||
const uint64_t sites = grid->oSites();
|
||||
|
||||
hostVector<View> h_left_v(m);
|
||||
deviceVector<View> d_left_v(m);
|
||||
for(int j=0;j<m;j++){
|
||||
conformable(*left[j],right);
|
||||
h_left_v[j] = left[j]->View(AcceleratorRead);
|
||||
}
|
||||
acceleratorCopyToDevice(&h_left_v[0],&d_left_v[0],m*sizeof(View));
|
||||
View *left_vp = &d_left_v[0];
|
||||
|
||||
deviceVector<batch_t> partial(sites);
|
||||
batch_t *partial_v = &partial[0];
|
||||
{
|
||||
autoView(right_v,right,AcceleratorRead);
|
||||
accelerator_for(ss,sites,1,{
|
||||
auto r = right_v[ss];
|
||||
batch_t acc;
|
||||
for(int j=0;j<B;j++){
|
||||
if ( j<m ) acc._internal[j] = innerProductD(left_vp[j][ss],r);
|
||||
else zeroit(acc._internal[j]);
|
||||
}
|
||||
partial_v[ss] = acc;
|
||||
});
|
||||
}
|
||||
for(int j=0;j<m;j++) h_left_v[j].ViewClose();
|
||||
|
||||
auto res = sum(partial_v,sites); // one reduction for the whole batch
|
||||
for(int j=0;j<m;j++) out[j] = TensorRemove(res._internal[j]);
|
||||
}
|
||||
|
||||
template<class vobj>
|
||||
void rankInnerProductMulti(std::vector<ComplexD> &out,
|
||||
const std::vector<const Lattice<vobj>*> &left,
|
||||
const Lattice<vobj> &right)
|
||||
{
|
||||
int m = left.size();
|
||||
out.resize(m);
|
||||
for(int j0=0;j0<m;j0+=16){
|
||||
int mm = std::min(16,m-j0);
|
||||
std::vector<const Lattice<vobj>*> sub(left.begin()+j0,left.begin()+j0+mm);
|
||||
if ( mm<=2 ) rankInnerProductMultiChunk<2> (&out[j0],mm,sub,right);
|
||||
else if ( mm<=4 ) rankInnerProductMultiChunk<4> (&out[j0],mm,sub,right);
|
||||
else if ( mm<=8 ) rankInnerProductMultiChunk<8> (&out[j0],mm,sub,right);
|
||||
else rankInnerProductMultiChunk<16>(&out[j0],mm,sub,right);
|
||||
}
|
||||
}
|
||||
|
||||
template<class vobj>
|
||||
void innerProductMulti(std::vector<ComplexD> &out,
|
||||
const std::vector<const Lattice<vobj>*> &left,
|
||||
const Lattice<vobj> &right)
|
||||
{
|
||||
rankInnerProductMulti(out,left,right);
|
||||
if ( out.size() ) right.Grid()->GlobalSumVector(&out[0],(int)out.size());
|
||||
}
|
||||
|
||||
// z = z + sum_j b[j] x[j]; if do_norm, returns global |z|^2 from the same pass.
|
||||
template<class vobj>
|
||||
RealD axpyMultiNormImpl(Lattice<vobj> &z,const std::vector<ComplexD> &b,
|
||||
const std::vector<const Lattice<vobj>*> &x,int do_norm)
|
||||
{
|
||||
typedef decltype(z.View(AcceleratorRead)) View;
|
||||
|
||||
int m = x.size();
|
||||
GRID_ASSERT((int)b.size()>=m);
|
||||
GridBase *grid = z.Grid();
|
||||
const uint64_t nsimd = grid->Nsimd();
|
||||
const uint64_t sites = grid->oSites();
|
||||
|
||||
hostVector<View> h_x_v(std::max(m,1));
|
||||
deviceVector<View> d_x_v(std::max(m,1));
|
||||
hostVector<ComplexD> h_b(std::max(m,1));
|
||||
deviceVector<ComplexD> d_b(std::max(m,1));
|
||||
for(int j=0;j<m;j++){
|
||||
conformable(*x[j],z);
|
||||
GRID_ASSERT(x[j]!=&z); // window must not alias the accumulator
|
||||
h_x_v[j] = x[j]->View(AcceleratorRead);
|
||||
h_b[j] = b[j];
|
||||
}
|
||||
if ( m ) {
|
||||
acceleratorCopyToDevice(&h_x_v[0],&d_x_v[0],m*sizeof(View));
|
||||
acceleratorCopyToDevice(&h_b[0],&d_b[0],m*sizeof(ComplexD));
|
||||
}
|
||||
View *x_vp = &d_x_v[0];
|
||||
ComplexD *b_p = &d_b[0];
|
||||
|
||||
autoView(z_v,z,AcceleratorWrite);
|
||||
typedef decltype(innerProduct(z_v[0],z_v[0])) inner_t;
|
||||
deviceVector<inner_t> inner_tmp(do_norm ? sites : 1);
|
||||
inner_t *inner_tmp_v = &inner_tmp[0];
|
||||
|
||||
accelerator_for(ss,sites,nsimd,{
|
||||
auto acc = coalescedRead(z_v[ss]);
|
||||
for(int j=0;j<m;j++) acc = acc + b_p[j]*coalescedRead(x_vp[j][ss]);
|
||||
coalescedWrite(z_v[ss],acc);
|
||||
if ( do_norm ) coalescedWrite(inner_tmp_v[ss],innerProduct(acc,acc));
|
||||
});
|
||||
for(int j=0;j<m;j++) h_x_v[j].ViewClose();
|
||||
|
||||
RealD nrm = 0.0;
|
||||
if ( do_norm ) {
|
||||
nrm = real(TensorRemove(sumD(inner_tmp_v,sites)));
|
||||
grid->GlobalSum(nrm);
|
||||
}
|
||||
return nrm;
|
||||
}
|
||||
template<class vobj>
|
||||
void axpyMulti(Lattice<vobj> &z,const std::vector<ComplexD> &b,const std::vector<const Lattice<vobj>*> &x)
|
||||
{
|
||||
axpyMultiNormImpl(z,b,x,0);
|
||||
}
|
||||
template<class vobj>
|
||||
RealD axpyMultiNorm(Lattice<vobj> &z,const std::vector<ComplexD> &b,const std::vector<const Lattice<vobj>*> &x)
|
||||
{
|
||||
return axpyMultiNormImpl(z,b,x,1);
|
||||
}
|
||||
|
||||
NAMESPACE_END(Grid);
|
||||
|
||||
|
||||
|
||||
@@ -79,18 +79,26 @@ int Ls = 24;
|
||||
int CoarsenBatch = 9;
|
||||
std::vector<int> lat_size({48,48,48,96});
|
||||
|
||||
// Solver tuning, values as in the V1 example
|
||||
// Solver tuning. PRINCIPLE (PB, 2026-08-24): the defaults ARE the current
|
||||
// optimum, so an unset environment reproduces the best banked result; they
|
||||
// are updated as and when a better point is found, and every change is
|
||||
// dated here. Environment variables of the same names override for sweeps.
|
||||
//
|
||||
// Current optimum: 2026-08-24, slurm-5335492 F4, 48^3x96 Ls=24 on 288 GCDs,
|
||||
// 17.2 s/RHS at Nrhs=4, 32.2 s at Nrhs=1 (exact-halo FINAL ~1e-8 pending
|
||||
// the exact-outer rerun). Smoother mmax == order (full GCR history);
|
||||
// PB's mmax=1 trial gave 72 vs ~60 outer iterations and was slower.
|
||||
RealD FineSmootherShift = 0.1;
|
||||
int FineSmootherOrder = 16;
|
||||
int FineSmootherMmax = 1;
|
||||
int FineSmootherOrder = 6;
|
||||
int FineSmootherMmax = 6;
|
||||
RealD CoarseSmootherShift = 0.1;
|
||||
int CoarseSmootherNstep = 4;
|
||||
int CoarseSmootherMmax = 1;
|
||||
RealD CoarseSolverTol = 0.03;
|
||||
int CoarseSmootherNstep = 2;
|
||||
int CoarseSmootherMmax = 2;
|
||||
RealD CoarseSolverTol = 0.05;
|
||||
int CoarseSolverOrder = 200;
|
||||
int CoarseSolverMmax = 4;
|
||||
int CoarseSolverMmax = 16;
|
||||
RealD OuterTol = 1.0e-8;
|
||||
int OuterMmax = 8;
|
||||
int OuterMmax = 4;
|
||||
int OuterNstep = 8;
|
||||
|
||||
// "It's legal to get the same answer faster, not to get a less correct
|
||||
@@ -319,34 +327,48 @@ public:
|
||||
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);
|
||||
RealD cp; ComplexD a,b,rq; int nrhs=src.size(); GridBase *grid=src[0].Grid();
|
||||
std::vector<Field> r(nrhs,grid),Az(nrhs,grid); // Az: restart residual scratch only
|
||||
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);
|
||||
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);
|
||||
// p[0]=Prec(r), q[0]=A p[0], produced directly in the history slots (no copies)
|
||||
Preconditioner(r,p[0]); vOp(p[0],q[0]); qq[0]=vnorm2(q[0]); 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;
|
||||
// New direction straight into its history slot: p=Prec(r), q=A p.
|
||||
Preconditioner(r,p[peri_kp]);
|
||||
vOp(p[peri_kp],q[peri_kp]);
|
||||
int northog=((kp)>(mmax-1))?(mmax-1):(kp);
|
||||
{
|
||||
GRID_TRACE("MrhsPGCR orthog");
|
||||
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]);
|
||||
}
|
||||
// Classical Gram-Schmidt: all coefficients against the UN-updated new q
|
||||
// (independent, batchable), then apply. Complex coefficient: the
|
||||
// operator is non-Hermitian, real(<q_j,Aq>) alone left q's non-orthogonal.
|
||||
// Batched per rhs (one fused kernel + one reduction each), the shared
|
||||
// coefficient summed over rhs on the host, ONE GlobalSumVector.
|
||||
std::vector<ComplexD> bcoef(northog,ComplexD(0.0)), part;
|
||||
for(int rr=0;rr<nrhs;rr++){
|
||||
std::vector<const Field*> qwin(northog);
|
||||
for(int back=0;back<northog;back++){ int peri_back=(k-back)%mmax; GRID_ASSERT((k-back)>=0); qwin[back]=&q[peri_back][rr]; }
|
||||
rankInnerProductMulti(part,qwin,q[peri_kp][rr]);
|
||||
for(int back=0;back<northog;back++) bcoef[back]+=part[back];
|
||||
}
|
||||
if(northog) grid->GlobalSumVector(&bcoef[0],northog);
|
||||
for(int back=0;back<northog;back++){ int peri_back=(k-back)%mmax; bcoef[back]=-bcoef[back]/qq[peri_back]; }
|
||||
for(int rr=0;rr<nrhs;rr++){
|
||||
std::vector<const Field*> qwin(northog), pwin(northog);
|
||||
for(int back=0;back<northog;back++){ int peri_back=(k-back)%mmax; qwin[back]=&q[peri_back][rr]; pwin[back]=&p[peri_back][rr]; }
|
||||
axpyMulti(p[peri_kp][rr],bcoef,pwin);
|
||||
axpyMulti(q[peri_kp][rr],bcoef,qwin);
|
||||
}
|
||||
}
|
||||
qq[peri_kp]=vnorm2(q[peri_kp]);
|
||||
}
|
||||
@@ -475,7 +497,7 @@ int main (int argc, char ** argv)
|
||||
|
||||
// Level 1 blocking (default 2^4)
|
||||
Coordinate clatt = lat_size;
|
||||
Coordinate Block({2,2,2,2});
|
||||
Coordinate Block({2,2,3,3}); // L1->L2 blocking; banked optimum 2026-08-24 (env BLOCK overrides)
|
||||
if ( getenv("BLOCK") ){ GridCmdOptionIntVector(std::string(getenv("BLOCK")),Block); GRID_ASSERT(Block.size()==4); }
|
||||
for(int d=0;d<4;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;
|
||||
@@ -695,7 +717,7 @@ int main (int argc, char ** argv)
|
||||
// batch grid the L1 operator is currently set to.
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
Coordinate cclatt = clatt;
|
||||
Coordinate Block2({8,4,3,6});
|
||||
Coordinate Block2({4,4,2,4}); // L2->L3 blocking; banked optimum 2026-08-24 (env BLOCK2 overrides)
|
||||
if ( getenv("BLOCK2") ){ GridCmdOptionIntVector(std::string(getenv("BLOCK2")),Block2); GRID_ASSERT(Block2.size()==4); }
|
||||
for(int d=0;d<4;d++){ GRID_ASSERT(clatt[d]%Block2[d]==0); cclatt[d]=clatt[d]/Block2[d]; }
|
||||
std::cout << GridLogMessage << "Block2 " << Block2 << " coarse-coarse lattice " << cclatt << std::endl;
|
||||
|
||||
Reference in New Issue
Block a user