Faster implementation of conserved current site contraction. Added 5D vectorised support, but not G-parity.

lib/qcd/action/fermion/FermionOperatorImpl.h

@@ -213,6 +213,13 @@ namespace QCD {
       mult(&phi(), &U(mu), &chi());
     }
     
+    inline void multLinkProp(SitePropagator &phi,
+                             const SiteDoubledGaugeField &U,
+                             const SitePropagator &chi,
+                             int mu) {
+       mult(&phi(), &U(mu), &chi());
+    }
+      
     template <class ref>
     inline void loadLinkElement(Simd &reg, ref &memory) {
       reg = memory;
@@ -341,6 +348,19 @@ class DomainWallVec5dImpl :  public PeriodicGaugeImpl< GaugeImplTypes< S,Nrepres
     mult(&phi(), &UU(), &chi());
   }
 
+  inline void multLinkProp(SitePropagator &phi,
+                           const SiteDoubledGaugeField &U,
+                           const SitePropagator &chi,
+                           int mu) {
+    SiteGaugeLink UU;
+    for (int i = 0; i < Nrepresentation; i++) {
+      for (int j = 0; j < Nrepresentation; j++) {
+        vsplat(UU()()(i, j), U(mu)()(i, j));
+      }
+    }
+    mult(&phi(), &UU(), &chi());
+  }
+
   inline void DoubleStore(GridBase *GaugeGrid, DoubledGaugeField &Uds,const GaugeField &Umu) 
   {
     SiteScalarGaugeField  ScalarUmu;
@@ -538,6 +558,13 @@ class GparityWilsonImpl : public ConjugateGaugeImpl<GaugeImplTypes<S, Nrepresent
    
  }
 
+ // Fixme: Gparity prop * link
+ inline void multLinkProp(SitePropagator &phi, const SiteDoubledGaugeField &U,
+                          const SitePropagator &chi, int mu)
+ {
+   assert(0);
+ }
+
  inline void DoubleStore(GridBase *GaugeGrid,DoubledGaugeField &Uds,const GaugeField &Umu)
  {
    conformable(Uds._grid,GaugeGrid);

lib/qcd/action/fermion/WilsonFermion.cc

@@ -361,8 +361,24 @@ void WilsonFermion<Impl>::ContractConservedCurrent(PropagatorField &q_in_1,
     conformable(_grid, q_in_1._grid);
     conformable(_grid, q_in_2._grid);
     conformable(_grid, q_out._grid);
-    Kernels::ContractConservedCurrentInternal(q_in_1, q_in_2, q_out, 
+    PropagatorField tmp(_grid);
-                                              Umu, curr_type, mu);
+    q_out = zero;
+
+    // Forward, need q1(x + mu), q2(x)
+    tmp = Cshift(q_in_1, mu, 1);
+    parallel_for (unsigned int sU = 0; sU < Umu._grid->oSites(); ++sU)
+    {
+        Kernels::ContractConservedCurrentSiteFwd(tmp, q_in_2, q_out, Umu,
+                                                 mu, sU, sU, sU, sU);
+    }
+
+    // Backward, need q1(x), q2(x + mu)
+    tmp = Cshift(q_in_2, mu, 1);
+    parallel_for (unsigned int sU = 0; sU < Umu._grid->oSites(); ++sU)
+    {
+        Kernels::ContractConservedCurrentSiteBwd(q_in_1, tmp, q_out, Umu,
+                                                 mu, sU, sU, sU, sU);
+    }
 }
 
 template <class Impl>

lib/qcd/action/fermion/WilsonFermion5D.cc

@@ -687,26 +687,44 @@ void WilsonFermion5D<Impl>::ContractConservedCurrent(PropagatorField &q_in_1,
     conformable(q_in_1._grid, q_in_2._grid);
     conformable(_FourDimGrid, q_out._grid);
 
-    PropagatorField q1_s(_FourDimGrid);
+    PropagatorField tmp(FermionGrid());
-    PropagatorField q2_s(_FourDimGrid);
-    PropagatorField tmp(_FourDimGrid);
-
-    // Contract across 5th dimension.
     q_out = zero;
+
+    // Forward, need q1(x + mu, s), q2(x, Ls - 1 - s). 5D lattice so shift
+    // 4D coordinate mu by one.
+    tmp = Cshift(q_in_1, mu + 1, 1);
+    parallel_for (unsigned int sU = 0; sU < Umu._grid->oSites(); ++sU)
+    {
+        unsigned int sF1 = sU * Ls;
+        unsigned int sF2 = (sU + 1) * Ls - 1;
         for (int s = 0; s < Ls; ++s)
         {
-        ExtractSlice(q1_s, q_in_1, s, 0);
+            bool axial_sign = ((curr_type == Current::Axial) && (s < (Ls / 2))) ? \
-        ExtractSlice(q2_s, q_in_2, Ls - s - 1, 0);
+                              true : false;
-        Kernels::ContractConservedCurrentInternal(q1_s, q2_s, tmp, Umu, curr_type, mu);
+            Kernels::ContractConservedCurrentSiteFwd(tmp, q_in_2, q_out, Umu,
-
+                                                     mu, sF1, sF2, sU, sU,
-        // Axial current sign
+                                                     axial_sign);
-        if ((curr_type == Current::Axial) && (s < (Ls / 2)))
+            sF1++;
-        {
+            sF2--;
-            q_out -= tmp;
         }
-        else
+    }
+
+    // Backward, need q1(x, s), q2(x + mu, Ls - 1 - s). 5D lattice so shift
+    // 4D coordinate mu by one.
+    tmp = Cshift(q_in_2, mu + 1, 1);
+    parallel_for (unsigned int sU = 0; sU < Umu._grid->oSites(); ++sU)
     {
-            q_out += tmp;
+        unsigned int sF1 = sU * Ls;
+        unsigned int sF2 = (sU + 1) * Ls - 1;
+        for (int s = 0; s < Ls; ++s)
+        {
+            bool axial_sign = ((curr_type == Current::Axial) && (s < (Ls / 2))) ? \
+                              true : false;
+            Kernels::ContractConservedCurrentSiteBwd(q_in_1, tmp, q_out, Umu,
+                                                     mu, sF1, sF2, sU, sU,
+                                                     axial_sign);
+            sF1++;
+            sF2--;
         }
     }
 }

lib/qcd/action/fermion/WilsonKernels.cc

@@ -290,25 +290,109 @@ void WilsonKernels<Impl>::DhopDir( StencilImpl &st, DoubledGaugeField &U,SiteHal
 #define WilsonCurrentFwd(expr, mu) ((expr - Gamma::gmu[mu]*expr))
 #define WilsonCurrentBwd(expr, mu) ((expr + Gamma::gmu[mu]*expr))
 
+/*******************************************************************************
+ * Name: ContractConservedCurrentSiteFwd
+ * Operation: (1/2) * q2[x] * U(x) * (g[mu] - 1) * q1[x + mu]
+ * Notes: - DoubledGaugeField U assumed to contain -1/2 factor.
+ *        - Pass in q_in_1 shifted in +ve mu direction.
+ ******************************************************************************/
 template<class Impl>
-void WilsonKernels<Impl>::ContractConservedCurrentInternal(const PropagatorField &q_in_1,
+void WilsonKernels<Impl>::ContractConservedCurrentSiteFwd(
+                                                  const PropagatorField &q_in_1,
                                                   const PropagatorField &q_in_2,
                                                   PropagatorField &q_out,
                                                   DoubledGaugeField &U,
-                                                           Current curr_type,
+                                                  unsigned int mu,
-                                                           unsigned int mu)
+                                                  unsigned int sF_in_1,
+                                                  unsigned int sF_in_2,
+                                                  unsigned int sF_out,
+                                                  unsigned int sU,
+                                                  bool switch_sign)
 {
+    SitePropagator result, tmp;
     Gamma g5(Gamma::Algebra::Gamma5);
-    PropagatorField tmp(q_out._grid);
+    multLinkProp(tmp, U._odata[sU], q_in_1._odata[sF_in_1], mu);
-    GaugeLinkField Umu(U._grid);
+    result = g5 * adj(q_in_2._odata[sF_in_2]) * g5 * WilsonCurrentFwd(tmp, mu);
-    Umu  = PeekIndex<LorentzIndex>(U, mu);
+    if (switch_sign)
-
+    {
-    tmp    = this->CovShiftForward(Umu, mu, q_in_1);
+        q_out._odata[sF_out] -= result;
-    q_out  = (g5*adj(q_in_2)*g5)*WilsonCurrentFwd(tmp, mu);
-
-    tmp    = this->CovShiftForward(Umu, mu, q_in_2);
-    q_out -= (g5*adj(tmp)*g5)*WilsonCurrentBwd(q_in_1, mu);
     }
+    else
+    {
+        q_out._odata[sF_out] += result;
+    }
+}
+
+/*******************************************************************************
+ * Name: ContractConservedCurrentSiteBwd
+ * Operation: (1/2) * q2[x + mu] * U^dag(x) * (g[mu] + 1) * q1[x]
+ * Notes: - DoubledGaugeField U assumed to contain -1/2 factor.
+ *        - Pass in q_in_2 shifted in +ve mu direction.
+ ******************************************************************************/
+template<class Impl>
+void WilsonKernels<Impl>::ContractConservedCurrentSiteBwd(
+                                                  const PropagatorField &q_in_1,
+                                                  const PropagatorField &q_in_2,
+                                                  PropagatorField &q_out,
+                                                  DoubledGaugeField &U,
+                                                  unsigned int mu,
+                                                  unsigned int sF_in_1,
+                                                  unsigned int sF_in_2,
+                                                  unsigned int sF_out,
+                                                  unsigned int sU,
+                                                  bool switch_sign)
+{
+    SitePropagator result, tmp;
+    Gamma g5(Gamma::Algebra::Gamma5);
+    multLinkProp(tmp, U._odata[sU], q_in_1._odata[sF_in_1], mu + Nd);
+    result = g5 * adj(q_in_2._odata[sF_in_2]) * g5 * WilsonCurrentBwd(tmp, mu);
+    if (switch_sign)
+    {
+        q_out._odata[sF_out] += result;
+    }
+    else
+    {
+        q_out._odata[sF_out] -= result;
+    }
+}
+
+// G-parity requires more specialised implementation.
+#define NO_CURR_SITE(Impl) \
+template <> \
+void WilsonKernels<Impl>::ContractConservedCurrentSiteFwd( \
+                                                  const PropagatorField &q_in_1, \
+                                                  const PropagatorField &q_in_2, \
+                                                  PropagatorField &q_out,        \
+                                                  DoubledGaugeField &U,          \
+                                                  unsigned int mu,               \
+                                                  unsigned int sF_in_1,          \
+                                                  unsigned int sF_in_2,          \
+                                                  unsigned int sF_out,           \
+                                                  unsigned int sU,               \
+                                                  bool switch_sign)              \
+{ \
+    assert(0); \
+} \
+template <> \
+void WilsonKernels<Impl>::ContractConservedCurrentSiteBwd( \
+                                                  const PropagatorField &q_in_1, \
+                                                  const PropagatorField &q_in_2, \
+                                                  PropagatorField &q_out,        \
+                                                  DoubledGaugeField &U,          \
+                                                  unsigned int mu,               \
+                                                  unsigned int sF_in_1,          \
+                                                  unsigned int sF_in_2,          \
+                                                  unsigned int sF_out,           \
+                                                  unsigned int sU,               \
+                                                  bool switch_sign)              \
+{ \
+    assert(0); \
+}
+
+NO_CURR_SITE(GparityWilsonImplF);
+NO_CURR_SITE(GparityWilsonImplD);
+NO_CURR_SITE(GparityWilsonImplFH);
+NO_CURR_SITE(GparityWilsonImplDF);
 
 
 template <class Impl>
@@ -343,16 +427,6 @@ void WilsonKernels<Impl>::SeqConservedCurrentInternal(const PropagatorField &q_i
 // GParity, (Z)DomainWallVec5D -> require special implementation
 #define NO_CURR(Impl) \
 template <> void  \
-WilsonKernels<Impl>::ContractConservedCurrentInternal(const PropagatorField &q_in_1, \
-                                                      const PropagatorField &q_in_2, \
-                                                      PropagatorField &q_out,        \
-                                                      DoubledGaugeField &U,          \
-                                                      Current curr_type,             \
-                                                      unsigned int mu)               \
-{ \
-    assert(0); \
-} \
-template <> void  \
 WilsonKernels<Impl>::SeqConservedCurrentInternal(const PropagatorField &q_in,       \
                                                  PropagatorField &q_out,            \
                                                  DoubledGaugeField &U,              \

lib/qcd/action/fermion/WilsonKernels.h

@@ -183,12 +183,26 @@ public:
   //////////////////////////////////////////////////////////////////////////////
   // Utilities for inserting Wilson conserved current.
   //////////////////////////////////////////////////////////////////////////////
-  void ContractConservedCurrentInternal(const PropagatorField &q_in_1,
+  void ContractConservedCurrentSiteFwd(const PropagatorField &q_in_1,
                                        const PropagatorField &q_in_2,
                                        PropagatorField &q_out,
                                        DoubledGaugeField &U,
-                                        Current curr_type,
+                                       unsigned int mu,
-                                        unsigned int mu);
+                                       unsigned int sF_in_1,
+                                       unsigned int sF_in_2,
+                                       unsigned int sF_out,
+                                       unsigned int sU,
+                                       bool switch_sign = false);
+  void ContractConservedCurrentSiteBwd(const PropagatorField &q_in_1,
+                                       const PropagatorField &q_in_2,
+                                       PropagatorField &q_out,
+                                       DoubledGaugeField &U,
+                                       unsigned int mu,
+                                       unsigned int sF_in_1,
+                                       unsigned int sF_in_2,
+                                       unsigned int sF_out,
+                                       unsigned int sU,
+                                       bool switch_sign = false);
   void SeqConservedCurrentInternal(const PropagatorField &q_in, 
                                    PropagatorField &q_out,
                                    DoubledGaugeField &U,