From 2bf3b4d5765b9db6d73fc4b3795971f851cb7d0f Mon Sep 17 00:00:00 2001
From: Peter Boyle <paboyle@ph.ed.ac.uk>
Date: Tue, 7 Dec 2021 09:02:02 -0800
Subject: [PATCH] Update to reduce memory footpring in benchmark test

---
 Grid/communicator/Communicator_mpi3.cc |  4 +-
 Grid/threads/Accelerator.h             |  2 +-
 benchmarks/Benchmark_dwf_fp32.cc       | 61 ++++++++++++++++++--------
 3 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/Grid/communicator/Communicator_mpi3.cc b/Grid/communicator/Communicator_mpi3.cc
index 305a3a9b..162180bc 100644
--- a/Grid/communicator/Communicator_mpi3.cc
+++ b/Grid/communicator/Communicator_mpi3.cc
@@ -388,6 +388,7 @@ double CartesianCommunicator::StencilSendToRecvFromBegin(std::vector<CommsReques
     // TODO : make a OMP loop on CPU, call threaded bcopy
     void *shm = (void *) this->ShmBufferTranslate(dest,recv);
     assert(shm!=NULL);
+    std::cout <<"acceleratorCopyDeviceToDeviceAsynch"<< std::endl;
     acceleratorCopyDeviceToDeviceAsynch(xmit,shm,bytes);
   }
 
@@ -399,12 +400,13 @@ double CartesianCommunicator::StencilSendToRecvFromBegin(std::vector<CommsReques
 }
 void CartesianCommunicator::StencilSendToRecvFromComplete(std::vector<CommsRequest_t> &list,int dir)
 {
+  acceleratorCopySynchronise(); std::cout << "Copy Synchronised\n"<<std::endl;
+
   int nreq=list.size();
 
   if (nreq==0) return;
 
   std::vector<MPI_Status> status(nreq);
-  acceleratorCopySynchronise(); 
   int ierr = MPI_Waitall(nreq,&list[0],&status[0]);
   assert(ierr==0);
   list.resize(0);
diff --git a/Grid/threads/Accelerator.h b/Grid/threads/Accelerator.h
index cec0600f..8be712ba 100644
--- a/Grid/threads/Accelerator.h
+++ b/Grid/threads/Accelerator.h
@@ -306,7 +306,7 @@ inline void acceleratorFreeDevice(void *ptr){free(ptr,*theGridAccelerator);};
 inline void acceleratorCopyDeviceToDeviceAsynch(void *from,void *to,size_t bytes)  {
   theGridAccelerator->memcpy(to,from,bytes);
 }
-inline void acceleratorCopySynchronise(void) {  theGridAccelerator->wait(); }
+inline void acceleratorCopySynchronise(void) {  theGridAccelerator->wait(); std::cout<<"acceleratorCopySynchronise() wait "<<std::endl; }
 inline void acceleratorCopyToDevice(void *from,void *to,size_t bytes)  { theGridAccelerator->memcpy(to,from,bytes); theGridAccelerator->wait();}
 inline void acceleratorCopyFromDevice(void *from,void *to,size_t bytes){ theGridAccelerator->memcpy(to,from,bytes); theGridAccelerator->wait();}
 inline void acceleratorMemSet(void *base,int value,size_t bytes) { theGridAccelerator->memset(base,value,bytes); theGridAccelerator->wait();}
diff --git a/benchmarks/Benchmark_dwf_fp32.cc b/benchmarks/Benchmark_dwf_fp32.cc
index d48486c0..4edf7c16 100644
--- a/benchmarks/Benchmark_dwf_fp32.cc
+++ b/benchmarks/Benchmark_dwf_fp32.cc
@@ -126,19 +126,10 @@ int main (int argc, char ** argv)
   // Naive wilson implementation
   ////////////////////////////////////
   // replicate across fifth dimension
-  LatticeGaugeFieldF Umu5d(FGrid);
-  std::vector<LatticeColourMatrixF> U(4,FGrid);
-  {
-    autoView( Umu5d_v, Umu5d, CpuWrite);
-    autoView( Umu_v  , Umu  , CpuRead);
-    for(int ss=0;ss<Umu.Grid()->oSites();ss++){
-      for(int s=0;s<Ls;s++){
-	Umu5d_v[Ls*ss+s] = Umu_v[ss];
-      }
-    }
-  }
+  //  LatticeGaugeFieldF Umu5d(FGrid);
+  std::vector<LatticeColourMatrixF> U(4,UGrid);
   for(int mu=0;mu<Nd;mu++){
-    U[mu] = PeekIndex<LorentzIndex>(Umu5d,mu);
+    U[mu] = PeekIndex<LorentzIndex>(Umu,mu);
   }
   std::cout << GridLogMessage << "Setting up Cshift based reference " << std::endl;
 
@@ -147,10 +138,28 @@ int main (int argc, char ** argv)
     ref = Zero();
     for(int mu=0;mu<Nd;mu++){
 
-      tmp = U[mu]*Cshift(src,mu+1,1);
+      tmp = Cshift(src,mu+1,1);
+      {
+	autoView( tmp_v  , tmp  , CpuWrite);
+	autoView( U_v  , U[mu]  , CpuRead);
+	for(int ss=0;ss<U[mu].Grid()->oSites();ss++){
+	  for(int s=0;s<Ls;s++){
+	    tmp_v[Ls*ss+s] = U_v[ss]*tmp_v[Ls*ss+s];
+	  }
+	}
+      }
       ref=ref + tmp - Gamma(Gmu[mu])*tmp;
 
-      tmp =adj(U[mu])*src;
+      {
+	autoView( tmp_v  , tmp  , CpuWrite);
+	autoView( U_v  , U[mu]  , CpuRead);
+	autoView( src_v, src    , CpuRead);
+	for(int ss=0;ss<U[mu].Grid()->oSites();ss++){
+	  for(int s=0;s<Ls;s++){
+	    tmp_v[Ls*ss+s] = adj(U_v[ss])*src_v[Ls*ss+s];
+	  }
+	}
+      }
       tmp =Cshift(tmp,mu+1,-1);
       ref=ref + tmp + Gamma(Gmu[mu])*tmp;
     }
@@ -242,16 +251,30 @@ int main (int argc, char ** argv)
     for(int mu=0;mu<Nd;mu++){
 
       //    ref =  src - Gamma(Gamma::Algebra::GammaX)* src ; // 1+gamma_x
-      tmp = U[mu]*Cshift(src,mu+1,1);
+      tmp = Cshift(src,mu+1,1);
       {
 	autoView( ref_v, ref, CpuWrite);
 	autoView( tmp_v, tmp, CpuRead);
-	for(int i=0;i<ref_v.size();i++){
-	  ref_v[i]+= tmp_v[i] + Gamma(Gmu[mu])*tmp_v[i]; ;
+	autoView( U_v  , U[mu]  , CpuRead);
+	for(int ss=0;ss<U[mu].Grid()->oSites();ss++){
+	  for(int s=0;s<Ls;s++){
+	    int i=s+Ls*ss;
+	    ref_v[i]+= U_v[ss]*(tmp_v[i] + Gamma(Gmu[mu])*tmp_v[i]); ;
+	  }
 	}
       }
-
-      tmp =adj(U[mu])*src;
+      
+      {
+	autoView( tmp_v  , tmp  , CpuWrite);
+	autoView( U_v  , U[mu]  , CpuRead);
+	autoView( src_v, src    , CpuRead);
+	for(int ss=0;ss<U[mu].Grid()->oSites();ss++){
+	  for(int s=0;s<Ls;s++){
+	    tmp_v[Ls*ss+s] = adj(U_v[ss])*src_v[Ls*ss+s];
+	  }
+	}
+      }
+      //      tmp =adj(U[mu])*src;
       tmp =Cshift(tmp,mu+1,-1);
       {
 	autoView( ref_v, ref, CpuWrite);