From dbbd5dba35b4c4c4da61f3167ecdc5fba751d758 Mon Sep 17 00:00:00 2001 From: Peter Boyle Date: Tue, 25 Aug 2026 07:29:02 -0400 Subject: [PATCH] Potentially faster GCR linalg --- ...GeneralisedConjugateResidualNonHermitian.h | 62 ++++---- Grid/lattice/Lattice_reduction.h | 145 ++++++++++++++++++ ...mple_pvdagm_v2_3level_DenseCoarseMatrix.cc | 68 +++++--- 3 files changed, 224 insertions(+), 51 deletions(-) diff --git a/Grid/algorithms/iterative/PrecGeneralisedConjugateResidualNonHermitian.h b/Grid/algorithms/iterative/PrecGeneralisedConjugateResidualNonHermitian.h index 5a389dff1..af043ff4c 100644 --- a/Grid/algorithms/iterative/PrecGeneralisedConjugateResidualNonHermitian.h +++ b/Grid/algorithms/iterative/PrecGeneralisedConjugateResidualNonHermitian.h @@ -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) <(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 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 qwin(northog), pwin(northog); for(int back=0;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["< bcoef; + innerProductMulti(bcoef,qwin,q[peri_kp]); + for(int back=0;back delete SliceGrid; } +///////////////////////////////////////////////////////////////////////////// +// Batched linear algebra for Krylov orthogonalisation (GCR history windows). +// +// innerProductMulti(out,left,right): out[j] = 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: fine fermion fields and coarse +// multi-RHS fields (nrhs folded into the grid) alike. +///////////////////////////////////////////////////////////////////////////// +template +void rankInnerProductMultiChunk(ComplexD *out,int m, + const std::vector*> &left, + const Lattice &right) +{ + typedef decltype(innerProductD(vobj(),vobj())) inner_t; + typedef iVector 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 h_left_v(m); + deviceVector d_left_v(m); + for(int j=0;jView(AcceleratorRead); + } + acceleratorCopyToDevice(&h_left_v[0],&d_left_v[0],m*sizeof(View)); + View *left_vp = &d_left_v[0]; + + deviceVector 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 +void rankInnerProductMulti(std::vector &out, + const std::vector*> &left, + const Lattice &right) +{ + int m = left.size(); + out.resize(m); + for(int j0=0;j0*> 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 +void innerProductMulti(std::vector &out, + const std::vector*> &left, + const Lattice &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 +RealD axpyMultiNormImpl(Lattice &z,const std::vector &b, + const std::vector*> &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 h_x_v(std::max(m,1)); + deviceVector d_x_v(std::max(m,1)); + hostVector h_b(std::max(m,1)); + deviceVector d_b(std::max(m,1)); + for(int j=0;jView(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_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;jGlobalSum(nrm); + } + return nrm; +} +template +void axpyMulti(Lattice &z,const std::vector &b,const std::vector*> &x) +{ + axpyMultiNormImpl(z,b,x,0); +} +template +RealD axpyMultiNorm(Lattice &z,const std::vector &b,const std::vector*> &x) +{ + return axpyMultiNormImpl(z,b,x,1); +} + NAMESPACE_END(Grid); diff --git a/examples/Example_pvdagm_v2_3level_DenseCoarseMatrix.cc b/examples/Example_pvdagm_v2_3level_DenseCoarseMatrix.cc index bf75768fa..fea175bfb 100644 --- a/examples/Example_pvdagm_v2_3level_DenseCoarseMatrix.cc +++ b/examples/Example_pvdagm_v2_3level_DenseCoarseMatrix.cc @@ -79,18 +79,26 @@ int Ls = 24; int CoarsenBatch = 9; std::vector 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< &src,std::vector &psi,RealD rsq){ - RealD cp; ComplexD a,b,rq; RealD zAAz; int nrhs=src.size(); GridBase *grid=src[0].Grid(); - std::vector 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 r(nrhs,grid),Az(nrhs,grid); // Az: restart residual scratch only std::vector< std::vector > q(mmax,std::vector(nrhs,grid)); std::vector< std::vector > p(mmax,std::vector(nrhs,grid)); std::vector qq(mmax); if (ZeroGuess && FirstCycle) { for(int rr=0;rr(mmax-1))?(mmax-1):(kp); { GRID_TRACE("MrhsPGCR orthog"); - for(int back=0;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() 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 bcoef(northog,ComplexD(0.0)), part; + for(int rr=0;rr qwin(northog); + for(int back=0;back=0); qwin[back]=&q[peri_back][rr]; } + rankInnerProductMulti(part,qwin,q[peri_kp][rr]); + for(int back=0;backGlobalSumVector(&bcoef[0],northog); + for(int back=0;back qwin(northog), pwin(northog); + for(int back=0;backL2 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;