From 44f4f5c8e2f5b64329cbd8bf4648476cf7b18d8b Mon Sep 17 00:00:00 2001
From: paboyle <paboyle@ph.ed.ac.uk>
Date: Fri, 27 Jul 2018 23:00:16 +0100
Subject: [PATCH] Momentum loop

---
 benchmarks/Benchmark_meson_field.cc | 183 ++++++++++++++++++++++++++--
 1 file changed, 176 insertions(+), 7 deletions(-)

diff --git a/benchmarks/Benchmark_meson_field.cc b/benchmarks/Benchmark_meson_field.cc
index aa2dbd0a..10a69525 100644
--- a/benchmarks/Benchmark_meson_field.cc
+++ b/benchmarks/Benchmark_meson_field.cc
@@ -180,7 +180,6 @@ void sliceInnerProductMesonFieldGamma(std::vector< std::vector<ComplexD> > &mat,
   const int Nsimd = grid->Nsimd();
   int Nt     = grid->GlobalDimensions()[orthogdim];
   int Ngamma = gammas.size();
-  //  int Nmom   = mom.size();
 
   assert(mat.size()==Lblock*Rblock*Ngamma);
   for(int t=0;t<mat.size();t++){
@@ -196,7 +195,6 @@ void sliceInnerProductMesonFieldGamma(std::vector< std::vector<ComplexD> > &mat,
   // splitting the SIMD
   int MFrvol = rd*Lblock*Rblock*Ngamma;
   int MFlvol = ld*Lblock*Rblock*Ngamma;
-  int MFfvol = fd*Lblock*Rblock*Ngamma;
 
   std::vector<vector_type,alignedAllocator<vector_type> > lvSum(MFrvol);
   parallel_for (int r = 0; r < MFrvol; r++){
@@ -330,8 +328,6 @@ void sliceInnerProductMesonFieldGamma1(std::vector< std::vector<ComplexD> > &mat
   int Nt     = grid->GlobalDimensions()[orthogdim];
   int Ngamma = gammas.size();
 
-  //  int Nmom   = mom.size();
-
   assert(mat.size()==Lblock*Rblock*Ngamma);
   for(int t=0;t<mat.size();t++){
     assert(mat[t].size()==Nt);
@@ -347,8 +343,6 @@ void sliceInnerProductMesonFieldGamma1(std::vector< std::vector<ComplexD> > &mat
   int MFrvol = rd*Lblock*Rblock;
   int MFlvol = ld*Lblock*Rblock;
 
-  int MFfvol = fd*Lblock*Rblock*Ngamma; // Do to dirac matrices here
-
   Vector<SpinMatrix_v > lvSum(MFrvol);
   parallel_for (int r = 0; r < MFrvol; r++){
     lvSum[r] = zero;
@@ -450,6 +444,161 @@ void sliceInnerProductMesonFieldGamma1(std::vector< std::vector<ComplexD> > &mat
   return;
 }
 
+template<class vobj>
+void sliceInnerProductMesonFieldGammaMom(std::vector< std::vector<ComplexD> > &mat, 
+					 const std::vector<Lattice<vobj> > &lhs,
+					 const std::vector<Lattice<vobj> > &rhs,
+					 int orthogdim,
+					 std::vector<Gamma::Algebra> gammas,
+					 const std::vector<LatticeComplex > &mom) 
+{
+  typedef typename vobj::scalar_object sobj;
+  typedef typename vobj::scalar_type scalar_type;
+  typedef typename vobj::vector_type vector_type;
+  typedef iSpinMatrix<vector_type> SpinMatrix_v;
+  typedef iSpinMatrix<scalar_type> SpinMatrix_s;
+  
+  int Lblock = lhs.size();
+  int Rblock = rhs.size();
+
+  GridBase *grid = lhs[0]._grid;
+  
+  const int    Nd = grid->_ndimension;
+  const int Nsimd = grid->Nsimd();
+  int Nt     = grid->GlobalDimensions()[orthogdim];
+  int Ngamma = gammas.size();
+  int Nmom   = mom.size();
+
+  assert(mat.size()==Lblock*Rblock*Ngamma*Nmom);
+  for(int t=0;t<mat.size();t++){
+    assert(mat[t].size()==Nt);
+  }
+
+  int fd=grid->_fdimensions[orthogdim];
+  int ld=grid->_ldimensions[orthogdim];
+  int rd=grid->_rdimensions[orthogdim];
+
+  // will locally sum vectors first
+  // sum across these down to scalars
+  // splitting the SIMD
+  int MFrvol = rd*Lblock*Rblock*Nmom;
+  int MFlvol = ld*Lblock*Rblock*Nmom;
+
+  Vector<SpinMatrix_v > lvSum(MFrvol);
+  parallel_for (int r = 0; r < MFrvol; r++){
+    lvSum[r] = zero;
+  }
+
+  Vector<SpinMatrix_s > lsSum(MFlvol);             
+  parallel_for (int r = 0; r < MFlvol; r++){
+    lsSum[r]=scalar_type(0.0);
+  }
+
+  int e1=    grid->_slice_nblock[orthogdim];
+  int e2=    grid->_slice_block [orthogdim];
+  int stride=grid->_slice_stride[orthogdim];
+  
+  std::cout << GridLogMessage << " Entering first parallel loop "<<std::endl;
+
+  // Parallelise over t-direction doesn't expose as much parallelism as needed for KNL
+  parallel_for(int r=0;r<rd;r++){
+
+    int so=r*grid->_ostride[orthogdim]; // base offset for start of plane 
+
+    for(int n=0;n<e1;n++){
+      for(int b=0;b<e2;b++){
+	int ss= so+n*stride+b;
+	Vector<iSinglet<vector_type> > phase(Nmom);
+	for(int m=0;m<Nmom;m++) phase[m] = mom[m]._odata[ss];
+
+	for(int i=0;i<Lblock;i++){
+
+	  auto left = conjugate(lhs[i]._odata[ss]);
+	  for(int j=0;j<Rblock;j++){
+
+	    SpinMatrix_v vv;
+	    auto right = rhs[j]._odata[ss];
+	    for(int s1=0;s1<Ns;s1++){
+	    for(int s2=0;s2<Ns;s2++){
+	      vv()(s1,s2)() = left()(s1)(0) * right()(s2)(0)
+		+             left()(s1)(1) * right()(s2)(1)
+		+             left()(s1)(2) * right()(s2)(2);
+	    }}
+	    
+	    // After getting the sitewise product do the mom phase loop
+	    for ( int m=0;m<Nmom;m++){
+	      int idx = m+Nmom*i+Nmom*Lblock*j+Nmom*Lblock*Rblock*r;
+	      lvSum[idx]=lvSum[idx]+vv*phase[m];
+	    }
+	  
+	  }
+	}
+      }
+    }
+  }
+
+  std::cout << GridLogMessage << " Entering second parallel loop "<<std::endl;
+  // Sum across simd lanes in the plane, breaking out orthog dir.
+  parallel_for(int rt=0;rt<rd;rt++){
+
+    std::vector<int> icoor(Nd);
+    std::vector<SpinMatrix_s> extracted(Nsimd);               
+
+
+    for(int i=0;i<Lblock;i++){
+    for(int j=0;j<Rblock;j++){
+    for(int m=0;m<Nmom;m++){
+
+      int ij_rdx = m+Nmom*i+Nmom*Lblock*j+Nmom*Lblock*Rblock*rt;
+
+      extract(lvSum[ij_rdx],extracted);
+
+      for(int idx=0;idx<Nsimd;idx++){
+
+	grid->iCoorFromIindex(icoor,idx);
+
+	int ldx    = rt+icoor[orthogdim]*rd;
+
+	int ij_ldx = m+Nmom*i+Nmom*Lblock*j+Nmom*Lblock*Rblock*ldx;
+
+	lsSum[ij_ldx]=lsSum[ij_ldx]+extracted[idx];
+
+      }
+    }}}
+  }
+
+  std::cout << GridLogMessage << " Entering third parallel loop "<<std::endl;
+  parallel_for(int t=0;t<fd;t++)
+  {
+    int pt = t / ld; // processor plane
+    int lt = t % ld;
+    for(int i=0;i<Lblock;i++){
+    for(int j=0;j<Rblock;j++){
+      if (pt == grid->_processor_coor[orthogdim]){
+	for(int m=0;m<Nmom;m++){
+	  int ij_dx = m+Nmom*i + Nmom*Lblock * j + Nmom*Lblock * Rblock * lt;
+	  for(int mu=0;mu<Ngamma;mu++){
+	    mat[ mu
+		+m*Ngamma
+		+i*Nmom*Ngamma
+		+j*Nmom*Ngamma*Lblock][t] = trace(lsSum[ij_dx]*Gamma(gammas[mu]));
+	  }
+	}
+      }
+      else{
+	for(int mu=0;mu<Ngamma;mu++){
+	for(int m=0;m<Nmom;m++){
+	  mat[mu+m*Ngamma+i*Nmom*Ngamma+j*Nmom*Lblock*Ngamma][t] = scalar_type(0.0);
+	}}
+      }
+    }}
+  }
+  std::cout << GridLogMessage << " Done "<<std::endl;
+  // defer sum over nodes.
+  return;
+}
+
+
 
 /*
 template void sliceInnerProductMesonField<SpinColourVector>(std::vector< std::vector<ComplexD> > &mat, 
@@ -491,7 +640,8 @@ int main (int argc, char ** argv)
   std::vector<int> simd_layout = GridDefaultSimd(Nd,vComplex::Nsimd());
   std::vector<int> mpi_layout  = GridDefaultMpi();
   GridCartesian               Grid(latt_size,simd_layout,mpi_layout);
-
+  
+  int Nmom=9;
   int nt = latt_size[Tp];
   uint64_t vol = 1;
   for(int d=0;d<Nd;d++){
@@ -507,12 +657,17 @@ int main (int argc, char ** argv)
 
   std::vector<LatticeFermion> v(Nm,&Grid);
   std::vector<LatticeFermion> w(Nm,&Grid);
+  std::vector<LatticeComplex> phases(Nmom,&Grid);
 
   for(int i=0;i<Nm;i++) { 
     random(pRNG,v[i]);
     random(pRNG,w[i]);
   }
 
+  for(int i=0;i<Nmom;i++) { 
+    phases[i] = Complex(1.0);
+  }
+
   double flops = vol * (11.0 * 8.0 + 6.0) * Nm*Nm;
   double byte  = vol * (12.0 * sizeof(Complex) ) * Nm*Nm;
 
@@ -520,6 +675,7 @@ int main (int argc, char ** argv)
   std::vector<std::vector<ComplexD> > MesonFields   (Nm*Nm);
   std::vector<std::vector<ComplexD> > MesonFields4  (Nm*Nm*4);
   std::vector<std::vector<ComplexD> > MesonFields16 (Nm*Nm*16);
+  std::vector<std::vector<ComplexD> > MesonFields16mom (Nm*Nm*16*Nmom);
   std::vector<std::vector<ComplexD> > MesonFieldsRef(Nm*Nm);
 
   for(int i=0;i<MesonFields.size();i++   ) MesonFields   [i].resize(nt);
@@ -527,6 +683,8 @@ int main (int argc, char ** argv)
   for(int i=0;i<MesonFields4.size();i++  ) MesonFields4  [i].resize(nt);
   for(int i=0;i<MesonFields16.size();i++ ) MesonFields16 [i].resize(nt);
 
+  for(int i=0;i<MesonFields16mom.size();i++ ) MesonFields16mom [i].resize(nt);
+
   GridLogMessage.TimingMode(1);
 
   std::cout<<GridLogMessage << "Running loop with sliceInnerProductVector"<<std::endl;
@@ -586,6 +744,17 @@ int main (int argc, char ** argv)
   std::cout<<GridLogMessage << "Done "<< flops/(t1-t0) <<" mflops " <<std::endl;
   std::cout<<GridLogMessage << "Done "<< byte /(t1-t0) <<" MB/s " <<std::endl;
 
+  std::cout<<GridLogMessage << "Running loop with Sixteen gammas "<<Nmom<<" momenta "<<std::endl;
+  flops = vol * ( 2 * 8.0 + 6.0 + 8.0*Nmom) * Nm*Nm*16;
+  byte  = vol * (12.0 * sizeof(Complex) ) * Nm*Nm
+        + vol * ( 2.0 * sizeof(Complex) *Nmom ) * Nm*Nm* 16;
+  t0 = usecond();
+  sliceInnerProductMesonFieldGammaMom(MesonFields16mom,w,v,Tp,Gmu16,phases);
+  t1 = usecond();
+  std::cout<<GridLogMessage << "Done "<< (t1-t0) <<" usecond " <<std::endl;
+  std::cout<<GridLogMessage << "Done "<< flops/(t1-t0) <<" mflops " <<std::endl;
+  std::cout<<GridLogMessage << "Done "<< byte /(t1-t0) <<" MB/s " <<std::endl;
+
 
 
   RealD err = 0;