Simplify the compressor interface again.

lib/qcd/action/fermion/FermionOperatorImpl.h

@@ -130,7 +130,7 @@ namespace Grid {
 
       typedef WilsonCompressor<SiteHalfSpinor,SiteSpinor> Compressor;
       typedef WilsonImplParams ImplParams;
-      typedef CartesianStencil<SiteSpinor,SiteHalfSpinor,Compressor> StencilImpl;
+      typedef WilsonStencil<SiteSpinor,SiteHalfSpinor> StencilImpl;
 
       ImplParams Params;
 
@@ -205,7 +205,7 @@ PARALLEL_FOR_LOOP
       typedef Lattice<SiteDoubledGaugeField> DoubledGaugeField;
 
       typedef WilsonCompressor<SiteHalfSpinor,SiteSpinor> Compressor;
-      typedef CartesianStencil<SiteSpinor,SiteHalfSpinor,Compressor> StencilImpl;
+      typedef WilsonStencil<SiteSpinor,SiteHalfSpinor> StencilImpl;
 
       typedef GparityWilsonImplParams ImplParams;
 

lib/qcd/action/fermion/WilsonCompressor.h

@@ -48,12 +48,7 @@ namespace QCD {
       mu=p;
     };
 
-    virtual SiteHalfSpinor operator () (const SiteSpinor &in,int dim,int plane,int osite,GridBase *grid) {
-      return spinproject(in);
-    }
-
-    SiteHalfSpinor spinproject(const SiteSpinor &in)
-    {
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
       SiteHalfSpinor ret;
       int mudag=mu;
       if (!dag) {
@@ -92,6 +87,173 @@ namespace QCD {
     }
   };
 
+  /////////////////////////
+  // optimised versions
+  /////////////////////////
+
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonXpCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjXp(ret,in);
+      return ret;
+    }
+  };
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonYpCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjYp(ret,in);
+      return ret;
+    }
+  };
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonZpCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjZp(ret,in);
+      return ret;
+    }
+  };
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonTpCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjTp(ret,in);
+      return ret;
+    }
+  };
+
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonXmCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjXm(ret,in);
+      return ret;
+    }
+  };
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonYmCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjYm(ret,in);
+      return ret;
+    }
+  };
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonZmCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjZm(ret,in);
+      return ret;
+    }
+  };
+  template<class SiteHalfSpinor,class SiteSpinor>
+  class WilsonTmCompressor {
+  public:
+    inline SiteHalfSpinor operator () (const SiteSpinor &in) {
+      SiteHalfSpinor ret;
+      spProjTm(ret,in);
+      return ret;
+    }
+  };
+
+    // Fast comms buffer manipulation which should inline right through (avoid direction
+    // dependent logic that prevents inlining
+  template<class vobj,class cobj>
+  class WilsonStencil : public CartesianStencil<vobj,cobj> {
+  public:
+
+    WilsonStencil(GridBase *grid,
+		int npoints,
+		int checkerboard,
+		const std::vector<int> &directions,
+		const std::vector<int> &distances)  : CartesianStencil<vobj,cobj> (grid,npoints,checkerboard,directions,distances) 
+      {    };
+
+    template < class compressor>
+    std::thread HaloExchangeOptBegin(const Lattice<vobj> &source,compressor &compress) {
+      this->Mergers.resize(0); 
+      this->Packets.resize(0);
+      this->HaloGatherOpt(source,compress);
+      return std::thread([&] { this->Communicate(); });
+    }
+
+    template < class compressor>
+    void HaloExchangeOpt(const Lattice<vobj> &source,compressor &compress) 
+    {
+      auto thr = this->HaloExchangeOptBegin(source,compress);
+      this->HaloExchangeOptComplete(thr);
+    }
+
+    void HaloExchangeOptComplete(std::thread &thr) 
+    {
+	this->CommsMerge(); // spins
+	this->jointime-=usecond();
+	thr.join();
+	this->jointime+=usecond();
+    }
+
+    template < class compressor>
+    void HaloGatherOpt(const Lattice<vobj> &source,compressor &compress)
+    {
+	// conformable(source._grid,_grid);
+	assert(source._grid==this->_grid);
+	this->halogtime-=usecond();
+
+	assert (this->comm_buf.size() == this->_unified_buffer_size );
+	this->u_comm_offset=0;
+
+	int dag = compress.dag;
+	static std::vector<int> dirs(8);
+	for(int mu=0;mu<4;mu++){
+	  if ( dag ) {
+	    dirs[mu]  =mu;
+	    dirs[mu+4]=mu+4;
+	  } else { 
+	    dirs[mu]  =mu+4;
+	    dirs[mu+4]=mu;
+	  }
+	}
+
+
+	WilsonXpCompressor<cobj,vobj> XpCompress;
+	this->HaloGatherDir(source,XpCompress,dirs[0]);
+
+	WilsonYpCompressor<cobj,vobj> YpCompress;
+	this->HaloGatherDir(source,YpCompress,dirs[1]);
+
+	WilsonZpCompressor<cobj,vobj> ZpCompress;
+	this->HaloGatherDir(source,ZpCompress,dirs[2]);
+
+	WilsonTpCompressor<cobj,vobj> TpCompress;
+	this->HaloGatherDir(source,TpCompress,dirs[3]);
+
+	WilsonXmCompressor<cobj,vobj> XmCompress;
+	this->HaloGatherDir(source,XmCompress,dirs[4]);
+
+	WilsonYmCompressor<cobj,vobj> YmCompress;
+	this->HaloGatherDir(source,YmCompress,dirs[5]);
+
+	WilsonZmCompressor<cobj,vobj> ZmCompress;
+	this->HaloGatherDir(source,ZmCompress,dirs[6]);
+
+	WilsonTmCompressor<cobj,vobj> TmCompress;
+	this->HaloGatherDir(source,TmCompress,dirs[7]);
+
+	assert(this->u_comm_offset==this->_unified_buffer_size);
+	this->halogtime+=usecond();
+      }
+
+  };
+
 
 }} // namespace close
 #endif

lib/qcd/action/fermion/WilsonFermion5D.cc

@@ -304,8 +304,8 @@ void WilsonFermion5D<Impl>::DhopInternalCommsThenCompute(StencilImpl & st, Lebes
   int nwork = U._grid->oSites();
   
   commtime -=usecond();
-  auto handle = st.HaloExchangeBegin(in,compressor);
-  st.HaloExchangeComplete(handle);
+  auto handle = st.HaloExchangeOptBegin(in,compressor);
+  st.HaloExchangeOptComplete(handle);
   commtime +=usecond();
 
   jointime -=usecond();
@@ -440,7 +440,7 @@ void WilsonFermion5D<Impl>::DhopInternalCommsOverlapCompute(StencilImpl & st, Le
   int nwork = U._grid->oSites();
   
   commtime -=usecond();
-  auto handle = st.HaloExchangeBegin(in,compressor);
+  auto handle = st.HaloExchangeOptBegin(in,compressor);
   commtime +=usecond();
   
   // Dhop takes the 4d grid from U, and makes a 5d index for fermion
@@ -498,7 +498,7 @@ PARALLEL_FOR_LOOP
   dslashtime +=usecond();
 
   jointime -=usecond();
-  st.HaloExchangeComplete(handle);
+  st.HaloExchangeOptComplete(handle);
   jointime +=usecond();
 
   local    = false;