FFT improvement by ~2x

This commit is contained in:
Peter Boyle
2026-09-03 18:13:02 -04:00
parent 357ede3664
commit 59f4a3729a
2 changed files with 70 additions and 30 deletions
+7 -2
View File
@@ -428,8 +428,13 @@ static void FFT_dim_execute(
scalar *rbuf_v = &rbuf[0];
scalar *pgbuf_v = &pgbuf[0];
// deterministic ceil-pad slots (never read back, but keeps padded FFT lines finite)
acceleratorMemSet(sbuf_v, 0, nbuf*sizeof(scalar));
// Pad slots (olin in [Nperp, Oloc*P)) are never packed, so they would carry
// garbage device memory into the FFT. Zero them so the padded lines stay
// finite -- but ONLY when padding is actually present. In the common
// Nperp % P == 0 case pack writes every sbuf entry bijectively, so skip the
// whole-buffer memset and its device sync entirely.
if ( (int64_t)Oloc*P != Nperp )
acceleratorMemSet(sbuf_v, 0, nbuf*sizeof(scalar));
const Coordinate ldims = grid->_ldimensions;
const Coordinate rdims = grid->_rdimensions;
+63 -28
View File
@@ -33,43 +33,77 @@ using namespace Grid;
template<class LatticeObject>
void bench(GridCartesian *grid, std::string name)
{
LatticeComplexD C(grid);
LatticeComplexD coor(grid);
ComplexD ci(0.0,1.0);
Coordinate p({1,2,3,4});
Coordinate latt_size = grid->_fdimensions;
std::cout<<"*************************************************"<<std::endl;
std::cout<<" Benchmarking FFT of "<<name<<" on plane wave "<<std::endl;
std::cout<<" Benchmarking FFT of "<<name<<" on random field "<<std::endl;
std::cout<<"*************************************************"<<std::endl;
C=Zero();
for(int mu=0;mu<4;mu++){
RealD TwoPiL = M_PI * 2.0/ latt_size[mu];
LatticeCoordinate(coor,mu);
C = C + (TwoPiL * p[mu]) * coor;
// Random source: works for every Lattice type (incl. LatticeFermionD, where the
// scalar plane-wave broadcast S=S+C does not compile). Content is irrelevant to
// FFT timing and plan cost; Parseval still holds -- a full forward FFT scales
// norm2 by exactly vol for ANY input.
GridParallelRNG RNG(grid); RNG.SeedFixedIntegers(std::vector<int>({1,2,3,4}));
LatticeObject S(grid); gaussian(RNG,S);
typedef typename LatticeObject::vector_object vobj;
const int nrep = 10;
// Correctness + Parseval pass (unplanned), also the WARMUP that absorbs
// FFTW/allocator/MPI first-touch so the two timed passes below are both warm
// and the comparison is fair (order confound removed).
{
LatticeObject Stilde(grid); Stilde=S;
FFT theFFT(grid);
std::cout << " norm2(s) "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,0,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,1,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,2,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,3,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
}
C = exp(C*ci);
LatticeObject S(grid);
LatticeObject Stilde(grid);
// NOTE: norm2() is a GlobalSum (MPI all-reduce over all ranks) -- it must NOT
// sit inside the timed loops or it dominates the number (that was the 80% gap
// in the Frontier LatticeFermionD run: 4 norm2 reductions + I/O bracketed by
// the outer timer, NOT plan creation). Timed loops below do FFT only.
S=Zero();
S = S+C;
double t_unplanned, t_planned, t_build;
FFT theFFT(grid);
// ---- UNPLANNED: the FFT class builds AND destroys an FFTW plan on EVERY
// FFT_dim call -- averaged over nrep so the recurring plan cost shows.
{
LatticeObject Stilde(grid);
double tt= -usecond();
for(int r=0;r<nrep;r++){
Stilde=S;
FFT theFFT(grid);
for(int mu=0;mu<4;mu++) theFFT.FFT_dim(Stilde,Stilde,mu,FFT::forward);
}
tt+= usecond();
t_unplanned = tt/1.e6/nrep;
}
Stilde=S;
std::cout << " norm2(s) "<<norm2(Stilde)<<std::endl;
double tt= -usecond();
theFFT.FFT_dim(Stilde,Stilde,0,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,1,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,2,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
theFFT.FFT_dim(Stilde,Stilde,3,FFT::forward); std::cout << theFFT.MFlops()<<" mflops "<<norm2(Stilde)<<std::endl;
tt+= usecond();
// ---- PLANNED: PlannedFFT builds all plans once in its ctor, reused every
// call -- the plan-free steady state (what repeated applies pay).
{
LatticeObject Stilde(grid);
double tc= -usecond();
PlannedFFT<vobj> theFFT(grid); // one-time plan build (all dims, fwd+bwd)
tc+= usecond();
t_build = tc/1.e6;
double tt= -usecond();
for(int r=0;r<nrep;r++){
Stilde=S;
for(int mu=0;mu<4;mu++) theFFT.FFT_dim(Stilde,Stilde,mu,FFT::forward);
}
tt+= usecond();
t_planned = tt/1.e6/nrep;
}
std::cout<<"*************************************************"<<std::endl;
std::cout<<" FFT of "<<latt_size <<" "<<name<<" took "<<tt/1.e6<<" s"<<std::endl;
// The decisive line: (unplanned - planned) is the per-call plan create+destroy
// cost, since the execute path is identical. If ~0, PlannedFFT buys nothing.
std::cout<<" PLANCOST "<<name<<" : unplanned "<<t_unplanned<<" s/call planned "<<t_planned
<<" s/call => plan create+destroy "<<(t_unplanned - t_planned)<<" s/call"
<<" (one-time build "<<t_build<<" s)"<<std::endl;
std::cout<<"*************************************************"<<std::endl;
}
@@ -95,6 +129,7 @@ int main (int argc, char ** argv)
bench<LatticeComplexD>(&GRID,std::string("LatticeComplexD"));
bench<LatticeColourMatrixD>(&GRID,std::string("LatticeColourMatrixD"));
bench<LatticeFermionD>(&GRID,std::string("LatticeFermionD")); // Ncomp=12, the FreePropagator/Fourier-precon path
bench<LatticePropagatorD>(&GRID,std::string("LatticePropagatorD"));
Grid_finalize();