From defd814750bb248944f266bccffbd54f205ca24e Mon Sep 17 00:00:00 2001
From: Peter Boyle <paboyle@ph.ed.ac.uk>
Date: Fri, 22 Dec 2023 18:07:03 -0500
Subject: [PATCH] Speed up the coarsened matrix matrix evaluation. It is block
 project limited. Could be sped up with calls to Batched GEMM and a data
 layout change.

---
 .../multigrid/GeneralCoarsenedMatrix.h        | 50 ++++++++++++-------
 1 file changed, 33 insertions(+), 17 deletions(-)

diff --git a/Grid/algorithms/multigrid/GeneralCoarsenedMatrix.h b/Grid/algorithms/multigrid/GeneralCoarsenedMatrix.h
index bedcf199..31172fa4 100644
--- a/Grid/algorithms/multigrid/GeneralCoarsenedMatrix.h
+++ b/Grid/algorithms/multigrid/GeneralCoarsenedMatrix.h
@@ -50,6 +50,7 @@ public:
   typedef iVector<CComplex,nbasis >  Cvec;
   typedef Lattice< CComplex >   CoarseScalar; // used for inner products on fine field
   typedef Lattice<Fobj >        FineField;
+  typedef Lattice<CComplex >    FineComplexField;
   typedef CoarseVector Field;
   ////////////////////
   // Data members
@@ -308,6 +309,7 @@ public:
     RealD teigen=0.0;
     RealD tmat=0.0;
     RealD tphase=0.0;
+    RealD tphaseBZ=0.0;
     RealD tinv=0.0;
 
     /////////////////////////////////////////////////////////////
@@ -362,28 +364,41 @@ public:
     ///////////////////////////////////////////////////////////////////////
     FineField phaV(grid); // Phased block basis vector
     FineField MphaV(grid);// Matrix applied
+    std::vector<FineComplexField> phaF(npoint,grid);
+    std::vector<CoarseComplexField> pha(npoint,CoarseGrid());
+    
     CoarseVector coarseInner(CoarseGrid());
+    
+    typedef typename CComplex::scalar_type SComplex;
+    FineComplexField one(grid); one=SComplex(1.0);
+    FineComplexField zz(grid); zz = Zero();
+    tphase=-usecond();
+    for(int p=0;p<npoint;p++){ // Loop over momenta in npoint
+      /////////////////////////////////////////////////////
+      // Stick a phase on every block
+      /////////////////////////////////////////////////////
+      CoarseComplexField coor(CoarseGrid());
+      pha[p]=Zero();
+      for(int mu=0;mu<Nd;mu++){
+	LatticeCoordinate(coor,mu);
+	RealD TwoPiL =  M_PI * 2.0/ clatt[mu];
+	pha[p] = pha[p] + (TwoPiL * geom.shifts[p][mu]) * coor;
+      }
+      pha[p]  =exp(pha[p]*ci);
 
+      blockZAXPY(phaF[p],pha[p],one,zz);
+      
+    }
+    tphase+=usecond();
+    
     std::vector<CoarseVector> ComputeProj(npoint,CoarseGrid());
     std::vector<CoarseVector>          FT(npoint,CoarseGrid());
     for(int i=0;i<nbasis;i++){// Loop over basis vectors
       std::cout << GridLogMessage<< "CoarsenMatrixColoured vec "<<i<<"/"<<nbasis<< std::endl;
       for(int p=0;p<npoint;p++){ // Loop over momenta in npoint
-	/////////////////////////////////////////////////////
-	// Stick a phase on every block
-	/////////////////////////////////////////////////////
-	tphase-=usecond();
-	CoarseComplexField coor(CoarseGrid());
-	CoarseComplexField pha(CoarseGrid());	pha=Zero();
-	for(int mu=0;mu<Nd;mu++){
-	  LatticeCoordinate(coor,mu);
-	  RealD TwoPiL =  M_PI * 2.0/ clatt[mu];
-	  pha = pha + (TwoPiL * geom.shifts[p][mu]) * coor;
-	}
-	pha  =exp(pha*ci);
-	phaV=Zero();
-	blockZAXPY(phaV,pha,Subspace.subspace[i],phaV);
-	tphase+=usecond();
+	tphaseBZ-=usecond();
+	phaV = phaF[p]*Subspace.subspace[i];
+	tphaseBZ+=usecond();
 
 	/////////////////////////////////////////////////////////////////////
 	// Multiple phased subspace vector by matrix and project to subspace
@@ -394,8 +409,8 @@ public:
 	tmat+=usecond();
 
 	tproj-=usecond();
-	blockProject(coarseInner,MphaV,Subspace.subspace);
-	coarseInner = conjugate(pha) * coarseInner;
+	blockProjectFast(coarseInner,MphaV,Subspace.subspace);
+	coarseInner = conjugate(pha[p]) * coarseInner;
 
 	ComputeProj[p] = coarseInner;
 	tproj+=usecond();
@@ -431,6 +446,7 @@ public:
     ExchangeCoarseLinks();
     std::cout << GridLogMessage<<"CoarsenOperator eigen  "<<teigen<<" us"<<std::endl;
     std::cout << GridLogMessage<<"CoarsenOperator phase  "<<tphase<<" us"<<std::endl;
+    std::cout << GridLogMessage<<"CoarsenOperator phaseBZ "<<tphaseBZ<<" us"<<std::endl;
     std::cout << GridLogMessage<<"CoarsenOperator mat    "<<tmat <<" us"<<std::endl;
     std::cout << GridLogMessage<<"CoarsenOperator proj   "<<tproj<<" us"<<std::endl;
     std::cout << GridLogMessage<<"CoarsenOperator inv    "<<tinv<<" us"<<std::endl;