From cfe3cd76d1543d1f13b3b53c6b029f154347cc1a Mon Sep 17 00:00:00 2001
From: Azusa Yamaguchi <ayamaguc@staffmail.ed.ac.uk>
Date: Mon, 19 Jun 2017 14:04:21 +0100
Subject: [PATCH] Block solver improvements

---
 .../iterative/BlockConjugateGradient.h        |  19 +-
 lib/lattice/Lattice_reduction.h               | 189 +++++++++++++++++-
 .../solver/Test_staggered_block_cg_unprec.cc  |   7 +-
 3 files changed, 192 insertions(+), 23 deletions(-)

diff --git a/lib/algorithms/iterative/BlockConjugateGradient.h b/lib/algorithms/iterative/BlockConjugateGradient.h
index d90194ae..53e11fa7 100644
--- a/lib/algorithms/iterative/BlockConjugateGradient.h
+++ b/lib/algorithms/iterative/BlockConjugateGradient.h
@@ -42,7 +42,7 @@ class BlockConjugateGradient : public OperatorFunction<Field> {
 
   typedef typename Field::scalar_type scomplex;
 
-  const int blockDim = 0;
+  int blockDim ;
 
   int Nblock;
   bool ErrorOnNoConverge;  // throw an assert when the CG fails to converge.
@@ -51,14 +51,15 @@ class BlockConjugateGradient : public OperatorFunction<Field> {
   Integer MaxIterations;
   Integer IterationsToComplete; //Number of iterations the CG took to finish. Filled in upon completion
   
-  BlockConjugateGradient(RealD tol, Integer maxit, bool err_on_no_conv = true)
+  BlockConjugateGradient(int _Orthog,RealD tol, Integer maxit, bool err_on_no_conv = true)
     : Tolerance(tol),
+    blockDim(_Orthog),
     MaxIterations(maxit),
     ErrorOnNoConverge(err_on_no_conv){};
 
 void operator()(LinearOperatorBase<Field> &Linop, const Field &Src, Field &Psi) 
 {
-  int Orthog = 0; // First dimension is block dim
+  int Orthog = blockDim; // First dimension is block dim; this is an assumption
   Nblock = Src._grid->_fdimensions[Orthog];
 
   std::cout<<GridLogMessage<<" Block Conjugate Gradient : Orthog "<<Orthog<<" Nblock "<<Nblock<<std::endl;
@@ -179,7 +180,7 @@ void operator()(LinearOperatorBase<Field> &Linop, const Field &Src, Field &Psi)
 
       Linop.HermOp(Psi, AP);
       AP = AP-Src;
-      std::cout << GridLogMessage <<"\tTrue residual is " << std::sqrt(norm2(AP)/norm2(Src)) <<std::endl;
+      std::cout << GridLogMessage <<"\t A__ True residual is " << std::sqrt(norm2(AP)/norm2(Src)) <<std::endl;
 
       std::cout << GridLogMessage << "Time Breakdown "<<std::endl;
       std::cout << GridLogMessage << "\tElapsed    " << SolverTimer.Elapsed()     <<std::endl;
@@ -209,8 +210,7 @@ class MultiRHSConjugateGradient : public OperatorFunction<Field> {
 
   typedef typename Field::scalar_type scomplex;
 
-  const int blockDim = 0;
-
+  int blockDim;
   int Nblock;
   bool ErrorOnNoConverge;  // throw an assert when the CG fails to converge.
                            // Defaults true.
@@ -218,14 +218,15 @@ class MultiRHSConjugateGradient : public OperatorFunction<Field> {
   Integer MaxIterations;
   Integer IterationsToComplete; //Number of iterations the CG took to finish. Filled in upon completion
   
-   MultiRHSConjugateGradient(RealD tol, Integer maxit, bool err_on_no_conv = true)
+  MultiRHSConjugateGradient(int Orthog,RealD tol, Integer maxit, bool err_on_no_conv = true)
     : Tolerance(tol),
+    blockDim(Orthog),
     MaxIterations(maxit),
     ErrorOnNoConverge(err_on_no_conv){};
 
 void operator()(LinearOperatorBase<Field> &Linop, const Field &Src, Field &Psi) 
 {
-  int Orthog = 0; // First dimension is block dim
+  int Orthog = blockDim; // First dimension is block dim
   Nblock = Src._grid->_fdimensions[Orthog];
 
   std::cout<<GridLogMessage<<"MultiRHS Conjugate Gradient : Orthog "<<Orthog<<" Nblock "<<Nblock<<std::endl;
@@ -285,12 +286,10 @@ void operator()(LinearOperatorBase<Field> &Linop, const Field &Src, Field &Psi)
     MatrixTimer.Stop();
 
     // Alpha
-    //    sliceInnerProductVectorTest(v_pAp_test,P,AP,Orthog);
     sliceInnerTimer.Start();
     sliceInnerProductVector(v_pAp,P,AP,Orthog);
     sliceInnerTimer.Stop();
     for(int b=0;b<Nblock;b++){
-      //      std::cout << " "<< v_pAp[b]<<" "<< v_pAp_test[b]<<std::endl;
       v_alpha[b] = v_rr[b]/real(v_pAp[b]);
     }
 
diff --git a/lib/lattice/Lattice_reduction.h b/lib/lattice/Lattice_reduction.h
index 14234fe0..78f88ce3 100644
--- a/lib/lattice/Lattice_reduction.h
+++ b/lib/lattice/Lattice_reduction.h
@@ -1,4 +1,4 @@
- /*************************************************************************************
+/*************************************************************************************
     Grid physics library, www.github.com/paboyle/Grid 
     Source file: ./lib/lattice/Lattice_reduction.h
     Copyright (C) 2015
@@ -462,13 +462,15 @@ static void sliceMaddMatrix (Lattice<vobj> &R,Eigen::MatrixXcd &aa,const Lattice
   typedef typename vobj::vector_type vector_type;
 
   int Nblock = X._grid->GlobalDimensions()[Orthog];
-  
+
   GridBase *FullGrid  = X._grid;
   GridBase *SliceGrid = makeSubSliceGrid(FullGrid,Orthog);
-  
+
   Lattice<vobj> Xslice(SliceGrid);
   Lattice<vobj> Rslice(SliceGrid);
-  
+
+#if 0
+  // R[i] = Y[i] + X[j] a(j,i) 
   for(int i=0;i<Nblock;i++){
     ExtractSlice(Rslice,Y,i,Orthog);
     for(int j=0;j<Nblock;j++){
@@ -477,8 +479,90 @@ static void sliceMaddMatrix (Lattice<vobj> &R,Eigen::MatrixXcd &aa,const Lattice
     }
     InsertSlice(Rslice,R,i,Orthog);
   }
+#endif
+#if 0
+  int nh =  FullGrid->_ndimension;
+  int nl = SliceGrid->_ndimension;
+
+#pragma omp parallel 
+{ 
+
+  std::vector<int> lcoor(nl); // sliced coor
+  std::vector<int> hcoor(nh); // unsliced coor
+  std::vector<sobj> s_x(Nblock);
+
+#pragma omp for
+  for(int idx=0;idx<SliceGrid->lSites();idx++){
+
+    SliceGrid->LocalIndexToLocalCoor(idx,lcoor); 
+
+    int ddl=0;
+    for(int d=0;d<nh;d++){
+      if ( d!=Orthog ) { 
+	hcoor[d]=lcoor[ddl++];
+      }
+    }
+
+    sobj dot;
+    for(int i=0;i<Nblock;i++){
+      hcoor[Orthog] = i;
+      peekLocalSite(s_x[i],X,hcoor);
+    }
+
+    for(int i=0;i<Nblock;i++){
+      hcoor[Orthog] = i;
+      peekLocalSite(dot,Y,hcoor);
+      for(int j=0;j<Nblock;j++){
+	dot = dot + s_x[j]*(scale*aa(j,i));
+      }
+      pokeLocalSite(dot,R,hcoor);
+    }
+  }
+}
+#endif
+
+#if 1
+  assert( FullGrid->_simd_layout[Orthog]==1);
+  int nh =  FullGrid->_ndimension;
+  int nl = SliceGrid->_ndimension;
+
+
+  //FIXME package in a convenient iterator
+  //Should loop over a plane orthogonal to direction "Orthog"
+  int stride=FullGrid->_slice_stride[Orthog];
+  int block =FullGrid->_slice_block [Orthog];
+  int nblock=FullGrid->_slice_nblock[Orthog];
+  int ostride=FullGrid->_ostride[Orthog];
+#pragma omp parallel 
+  {
+
+    std::vector<vobj> s_x(Nblock);
+
+#pragma omp for collapse(2)
+    for(int n=0;n<nblock;n++){
+    for(int b=0;b<block;b++){
+      int o  = n*stride + b;
+
+
+      for(int i=0;i<Nblock;i++){
+	s_x[i] = X[o+i*ostride];
+      }
+
+      vobj dot;
+
+      for(int i=0;i<Nblock;i++){
+	dot = Y[o+i*ostride];
+	for(int j=0;j<Nblock;j++){
+	  dot = dot + s_x[j]*(scale*aa(j,i));
+	}
+	R[o+i*ostride]=dot;
+      }
+    }}
+  }
+#endif
 };
 
+
 template<class vobj>
 static void sliceInnerProductMatrix(  Eigen::MatrixXcd &mat, const Lattice<vobj> &lhs,const Lattice<vobj> &rhs,int Orthog) 
 {
@@ -497,7 +581,8 @@ static void sliceInnerProductMatrix(  Eigen::MatrixXcd &mat, const Lattice<vobj>
   Lattice<vobj> Rslice(SliceGrid);
   
   mat = Eigen::MatrixXcd::Zero(Nblock,Nblock);
-  
+
+#if 0  
   for(int i=0;i<Nblock;i++){
     ExtractSlice(Lslice,lhs,i,Orthog);
     for(int j=0;j<Nblock;j++){
@@ -505,12 +590,96 @@ static void sliceInnerProductMatrix(  Eigen::MatrixXcd &mat, const Lattice<vobj>
       mat(i,j) = innerProduct(Lslice,Rslice);
     }
   }
-#undef FORCE_DIAG
-#ifdef FORCE_DIAG
-  for(int i=0;i<Nblock;i++){
-    for(int j=0;j<Nblock;j++){
-      if ( i != j ) mat(i,j)=0.0;
+#endif
+
+#if 0
+  int nh =  FullGrid->_ndimension;
+  int nl = SliceGrid->_ndimension;
+
+#pragma omp parallel 
+{ 
+  std::vector<int> lcoor(nl); // sliced coor
+  std::vector<int> hcoor(nh); // unsliced coor
+  std::vector<sobj> Left(Nblock);
+  std::vector<sobj> Right(Nblock);
+  Eigen::MatrixXcd  mat_thread = Eigen::MatrixXcd::Zero(Nblock,Nblock);
+
+#pragma omp for
+  for(int idx=0;idx<SliceGrid->lSites();idx++){
+
+    SliceGrid->LocalIndexToLocalCoor(idx,lcoor); 
+
+    int ddl=0;
+    for(int d=0;d<nh;d++){
+      if ( d!=Orthog ) { 
+	hcoor[d]=lcoor[ddl++];
+      }
     }
+
+    // Get the scalar objects
+    for(int i=0;i<Nblock;i++){
+      hcoor[Orthog] = i;
+      peekLocalSite(Left[i] ,lhs,hcoor);
+      peekLocalSite(Right[i],rhs,hcoor);
+    }
+
+    for(int i=0;i<Nblock;i++){
+    for(int j=0;j<Nblock;j++){
+      std::complex<double> ip = innerProduct(Left[i],Right[j]);
+      mat_thread(i,j) += ip;
+    }}
+  }
+
+#pragma omp critical
+  {
+    mat += mat_thread;
+  }  
+
+}
+#endif
+
+#if 1
+  assert( FullGrid->_simd_layout[Orthog]==1);
+  int nh =  FullGrid->_ndimension;
+  int nl = SliceGrid->_ndimension;
+
+  //FIXME package in a convenient iterator
+  //Should loop over a plane orthogonal to direction "Orthog"
+  int stride=FullGrid->_slice_stride[Orthog];
+  int block =FullGrid->_slice_block [Orthog];
+  int nblock=FullGrid->_slice_nblock[Orthog];
+  int ostride=FullGrid->_ostride[Orthog];
+
+  typedef typename vobj::vector_typeD vector_typeD;
+
+#pragma omp parallel 
+  {
+    std::vector<vobj> Left(Nblock);
+    std::vector<vobj> Right(Nblock);
+    Eigen::MatrixXcd  mat_thread = Eigen::MatrixXcd::Zero(Nblock,Nblock);
+
+#pragma omp for collapse(2)
+    for(int n=0;n<nblock;n++){
+    for(int b=0;b<block;b++){
+
+      int o  = n*stride + b;
+
+      for(int i=0;i<Nblock;i++){
+	Left [i] = lhs[o+i*ostride];
+	Right[i] = rhs[o+i*ostride];
+      }
+
+      for(int i=0;i<Nblock;i++){
+      for(int j=0;j<Nblock;j++){
+	auto tmp = innerProduct(Left[i],Right[j]);
+	vector_typeD rtmp = TensorRemove(tmp);
+	mat_thread(i,j) += Reduce(rtmp);
+      }}
+    }}
+#pragma omp critical
+    {
+      mat += mat_thread;
+    }  
   }
 #endif
   return;
diff --git a/tests/solver/Test_staggered_block_cg_unprec.cc b/tests/solver/Test_staggered_block_cg_unprec.cc
index 44e8fb52..8da93195 100644
--- a/tests/solver/Test_staggered_block_cg_unprec.cc
+++ b/tests/solver/Test_staggered_block_cg_unprec.cc
@@ -74,13 +74,14 @@ int main (int argc, char ** argv)
 
   LatticeGaugeField Umu(UGrid); SU3::HotConfiguration(pRNG,Umu);
 
-  RealD mass=0.01;
+  RealD mass=0.003;
   ImprovedStaggeredFermion5DR Ds(Umu,Umu,*FGrid,*FrbGrid,*UGrid,*UrbGrid,mass);
   MdagMLinearOperator<ImprovedStaggeredFermion5DR,FermionField> HermOp(Ds);
 
   ConjugateGradient<FermionField> CG(1.0e-8,10000);
-  BlockConjugateGradient<FermionField> BCG(1.0e-8,10000);
-  MultiRHSConjugateGradient<FermionField> mCG(1.0e-8,10000);
+  int blockDim = 0;
+  BlockConjugateGradient<FermionField>    BCG(blockDim,1.0e-8,10000);
+  MultiRHSConjugateGradient<FermionField> mCG(blockDim,1.0e-8,10000);
 
   std::cout << GridLogMessage << "************************************************************************ "<<std::endl;
   std::cout << GridLogMessage << " Calling 4d CG "<<std::endl;