From 148fc052bd57696d210e061790fee80940bd2668 Mon Sep 17 00:00:00 2001
From: Antonin Portelli <antonin.portelli@me.com>
Date: Fri, 5 Oct 2018 21:04:10 +0100
Subject: [PATCH] Hadrons: Aslash field, tested

---
 Grid/qcd/utils/A2Autils.h                     | 191 +++++++++++++++
 Hadrons/Modules.hpp                           |   1 +
 .../Modules/MContraction/A2AAslashField.cc    |   7 +
 .../Modules/MContraction/A2AAslashField.hpp   | 223 ++++++++++++++++++
 Hadrons/modules.inc                           |   2 +
 5 files changed, 424 insertions(+)
 create mode 100644 Hadrons/Modules/MContraction/A2AAslashField.cc
 create mode 100644 Hadrons/Modules/MContraction/A2AAslashField.hpp

diff --git a/Grid/qcd/utils/A2Autils.h b/Grid/qcd/utils/A2Autils.h
index ccdbc7b9..89b4d4bd 100644
--- a/Grid/qcd/utils/A2Autils.h
+++ b/Grid/qcd/utils/A2Autils.h
@@ -60,6 +60,14 @@ public:
 			  const FermionField *vj,
 			  int orthogdim);
 
+  template <typename TensorType> // output: rank 5 tensor, e.g. Eigen::Tensor<ComplexD, 5>
+  static void AslashField(TensorType &mat, 
+        const FermionField *lhs_wi,
+        const FermionField *rhs_vj,
+        const std::vector<ComplexField> &emB0,
+        const std::vector<ComplexField> &emB1,
+        int orthogdim, double *t_kernel = nullptr, double *t_gsum = nullptr);
+
   static void ContractWWVV(std::vector<PropagatorField> &WWVV,
 			   const Eigen::Tensor<ComplexD,3> &WW_sd,
 			   const FermionField *vs,
@@ -617,6 +625,189 @@ void A2Autils<FImpl>::PionFieldVV(Eigen::Tensor<ComplexD,3> &mat,
   PionFieldXX(mat,vi,vj,orthogdim,nog5);
 }
 
+// "A-slash" field w_i(x)^dag * i * A_mu * gamma_mu * v_j(x)
+//
+// With:
+//
+// B_0 = A_0 + i A_1
+// B_1 = A_2 + i A_3
+// 
+// then in spin space
+// 
+//                 ( 0          0          -conj(B_1) -B_0 )
+// i * A_mu g_mu = ( 0          0          -conj(B_0)  B_1 )
+//                 ( B_1        B_0        0          0    )
+//                 ( conj(B_0)  -conj(B_1) 0          0    )
+template <class FImpl>
+template <typename TensorType>
+void A2Autils<FImpl>::AslashField(TensorType &mat, 
+          const FermionField *lhs_wi,
+          const FermionField *rhs_vj,
+          const std::vector<ComplexField> &emB0,
+          const std::vector<ComplexField> &emB1,
+          int orthogdim, double *t_kernel, double *t_gsum) 
+{
+    typedef typename FermionField::vector_object vobj;
+    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;
+    typedef iSinglet<vector_type>    Singlet_v;
+    typedef iSinglet<scalar_type>    Singlet_s;
+    
+    int Lblock = mat.dimension(3); 
+    int Rblock = mat.dimension(4);
+
+    GridBase *grid = lhs_wi[0]._grid;
+    
+    const int    Nd = grid->_ndimension;
+    const int Nsimd = grid->Nsimd();
+
+    int Nt  = grid->GlobalDimensions()[orthogdim];
+    int Nem = emB0.size();
+    assert(emB1.size() == Nem);
+
+    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*Nem;
+    int MFlvol = ld*Lblock*Rblock*Nem;
+
+    Vector<vector_type> lvSum(MFrvol);
+    parallel_for (int r = 0; r < MFrvol; r++)
+    {
+        lvSum[r] = zero;
+    }
+
+    Vector<scalar_type> 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];
+
+    // Nested parallelism would be ok
+    // Wasting cores here. Test case r
+    if (t_kernel) *t_kernel = -usecond();
+    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;
+
+            for(int i=0;i<Lblock;i++)
+            {
+                auto left = conjugate(lhs_wi[i]._odata[ss]);
+
+                for(int j=0;j<Rblock;j++)
+                {
+                    SpinMatrix_v vv;
+                    auto right = rhs_vj[j]._odata[ss];
+
+                    for(int s1=0;s1<Ns;s1++)
+                    for(int s2=0;s2<Ns;s2++)
+                    {
+                        vv()(s1,s2)() = left()(s2)(0) * right()(s1)(0)
+                                        + left()(s2)(1) * right()(s1)(1)
+                                        + left()(s2)(2) * right()(s1)(2);
+                    }
+                    
+                    // After getting the sitewise product do the mom phase loop
+                    int base = Nem*i+Nem*Lblock*j+Nem*Lblock*Rblock*r;
+
+                    for ( int m=0;m<Nem;m++)
+                    {
+                        int idx  = m+base;
+                        auto b0  = emB0[m]._odata[ss];
+                        auto b1  = emB1[m]._odata[ss];
+                        auto cb0 = conjugate(b0);
+                        auto cb1 = conjugate(b1);
+
+                        lvSum[idx] += - vv()(3,0)()*b0()()()  - vv()(2,0)()*cb1()()()
+                                      + vv()(3,1)()*b1()()()  - vv()(2,1)()*cb0()()()
+                                      + vv()(0,2)()*b1()()()  + vv()(1,2)()*b0()()()
+                                      + vv()(0,3)()*cb0()()() - vv()(1,3)()*cb1()()();
+                    }
+                }
+            }
+        }
+    }
+
+    // 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<scalar_type> extracted(Nsimd);               
+
+        for(int i=0;i<Lblock;i++)
+        for(int j=0;j<Rblock;j++)
+        for(int m=0;m<Nem;m++)
+        {
+
+            int ij_rdx = m+Nem*i+Nem*Lblock*j+Nem*Lblock*Rblock*rt;
+
+            extract<vector_type,scalar_type>(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+Nem*i+Nem*Lblock*j+Nem*Lblock*Rblock*ldx;
+
+                lsSum[ij_ldx]=lsSum[ij_ldx]+extracted[idx];
+            }
+        }
+    }
+    if (t_kernel) *t_kernel += usecond();
+
+    // ld loop and local only??
+    int pd = grid->_processors[orthogdim];
+    int pc = grid->_processor_coor[orthogdim];
+    parallel_for_nest2(int lt=0;lt<ld;lt++)
+    {
+        for(int pt=0;pt<pd;pt++)
+        {
+            int t = lt + pt*ld;
+            if (pt == pc)
+            {
+                for(int i=0;i<Lblock;i++)
+                for(int j=0;j<Rblock;j++)
+                for(int m=0;m<Nem;m++)
+                {
+                    int ij_dx = m+Nem*i + Nem*Lblock * j + Nem*Lblock * Rblock * lt;
+
+                    mat(m,0,t,i,j) = lsSum[ij_dx];
+                }
+            } 
+            else 
+            { 
+                const scalar_type zz(0.0);
+
+                for(int i=0;i<Lblock;i++)
+                for(int j=0;j<Rblock;j++)
+                for(int m=0;m<Nem;m++)
+                {
+                    mat(m,0,t,i,j) = zz;
+                }
+            }
+        }
+    }
+    if (t_gsum) *t_gsum = -usecond();
+    grid->GlobalSumVector(&mat(0,0,0,0,0),Nem*Nt*Lblock*Rblock);
+    if (t_gsum) *t_gsum += usecond();
+}
 
 ////////////////////////////////////////////
 // Schematic thoughts about more generalised four quark insertion
diff --git a/Hadrons/Modules.hpp b/Hadrons/Modules.hpp
index 67ade797..2859f0fc 100644
--- a/Hadrons/Modules.hpp
+++ b/Hadrons/Modules.hpp
@@ -1,4 +1,5 @@
 #include <Hadrons/Modules/MContraction/Baryon.hpp>
+#include <Hadrons/Modules/MContraction/A2AAslashField.hpp>
 #include <Hadrons/Modules/MContraction/A2AMesonField.hpp>
 #include <Hadrons/Modules/MContraction/Meson.hpp>
 #include <Hadrons/Modules/MContraction/WeakHamiltonian.hpp>
diff --git a/Hadrons/Modules/MContraction/A2AAslashField.cc b/Hadrons/Modules/MContraction/A2AAslashField.cc
new file mode 100644
index 00000000..eb76932e
--- /dev/null
+++ b/Hadrons/Modules/MContraction/A2AAslashField.cc
@@ -0,0 +1,7 @@
+#include <Hadrons/Modules/MContraction/A2AAslashField.hpp>
+
+using namespace Grid;
+using namespace Hadrons;
+using namespace MContraction;
+
+template class Grid::Hadrons::MContraction::TA2AAslashField<FIMPL, PhotonR>;
diff --git a/Hadrons/Modules/MContraction/A2AAslashField.hpp b/Hadrons/Modules/MContraction/A2AAslashField.hpp
new file mode 100644
index 00000000..792ff6ee
--- /dev/null
+++ b/Hadrons/Modules/MContraction/A2AAslashField.hpp
@@ -0,0 +1,223 @@
+#ifndef Hadrons_MContraction_A2AAslashField_hpp_
+#define Hadrons_MContraction_A2AAslashField_hpp_
+
+#include <Hadrons/Global.hpp>
+#include <Hadrons/Module.hpp>
+#include <Hadrons/ModuleFactory.hpp>
+#include <Hadrons/A2AMatrix.hpp>
+
+#ifndef ASF_IO_TYPE
+#define ASF_IO_TYPE ComplexF
+#endif
+
+BEGIN_HADRONS_NAMESPACE
+
+/******************************************************************************
+ *                         A2AAslashField                                 *
+ ******************************************************************************/
+BEGIN_MODULE_NAMESPACE(MContraction)
+
+class A2AAslashFieldPar: Serializable
+{
+public:
+    GRID_SERIALIZABLE_CLASS_MEMBERS(A2AAslashFieldPar,
+                                    int, cacheBlock,
+                                    int, block,
+                                    std::string, left,
+                                    std::string, right,
+                                    std::string, output,
+                                    std::vector<std::string>, emField);
+};
+
+class A2AAslashFieldMetadata: Serializable
+{
+public:
+    GRID_SERIALIZABLE_CLASS_MEMBERS(A2AAslashFieldMetadata,
+                                    std::string, emFieldName);
+};
+
+template <typename T, typename FImpl>
+class AslashFieldKernel: public A2AKernel<T, typename FImpl::FermionField>
+{
+public:
+    typedef typename FImpl::FermionField FermionField;
+public:
+    AslashFieldKernel(const std::vector<LatticeComplex> &emB0,
+                      const std::vector<LatticeComplex> &emB1,
+                      GridBase *grid)
+    : emB0_(emB0), emB1_(emB1), grid_(grid)
+    {
+        vol_ = 1.;
+        for (auto &d: grid_->GlobalDimensions())
+        {
+            vol_ *= d;
+        }
+    }
+
+    virtual ~AslashFieldKernel(void) = default;
+    virtual void operator()(A2AMatrixSet<T> &m, const FermionField *left, 
+                            const FermionField *right,
+                            const unsigned int orthogDim, double &t)
+    {
+        A2Autils<FImpl>::AslashField(m, left, right, emB0_, emB1_, orthogDim, &t);
+    }
+
+    virtual double flops(const unsigned int blockSizei, const unsigned int blockSizej)
+    {
+        return 0.;
+    }
+
+    virtual double bytes(const unsigned int blockSizei, const unsigned int blockSizej)
+    {
+        return 0.;
+    }
+private:
+    const std::vector<LatticeComplex> &emB0_, &emB1_;
+    GridBase                          *grid_;
+    double                            vol_;
+};
+
+template <typename FImpl, typename PhotonImpl>
+class TA2AAslashField: public Module<A2AAslashFieldPar>
+{
+public:
+    FERM_TYPE_ALIASES(FImpl,);
+    typedef typename PhotonImpl::GaugeField EmField;
+    typedef A2AMatrixBlockComputation<Complex, 
+                                      FermionField, 
+                                      A2AAslashFieldMetadata, 
+                                      ASF_IO_TYPE> Computation;
+    typedef AslashFieldKernel<Complex, FImpl> Kernel;
+public:
+    // constructor
+    TA2AAslashField(const std::string name);
+    // destructor
+    virtual ~TA2AAslashField(void) {};
+    // dependency relation
+    virtual std::vector<std::string> getInput(void);
+    virtual std::vector<std::string> getOutput(void);
+    // setup
+    virtual void setup(void);
+    // execution
+    virtual void execute(void);
+};
+
+MODULE_REGISTER_TMP(A2AAslashField, ARG(TA2AAslashField<FIMPL, PhotonR>), MContraction);
+
+/******************************************************************************
+ *                 TA2AAslashField implementation                             *
+ ******************************************************************************/
+// constructor /////////////////////////////////////////////////////////////////
+template <typename FImpl, typename PhotonImpl>
+TA2AAslashField<FImpl, PhotonImpl>::TA2AAslashField(const std::string name)
+: Module<A2AAslashFieldPar>(name)
+{}
+
+// dependencies/products ///////////////////////////////////////////////////////
+template <typename FImpl, typename PhotonImpl>
+std::vector<std::string> TA2AAslashField<FImpl, PhotonImpl>::getInput(void)
+{
+    std::vector<std::string> in = par().emField;
+    
+    in.push_back(par().left);
+    in.push_back(par().right);
+
+    return in;
+}
+
+template <typename FImpl, typename PhotonImpl>
+std::vector<std::string> TA2AAslashField<FImpl, PhotonImpl>::getOutput(void)
+{
+    std::vector<std::string> out = {};
+    
+    return out;
+}
+
+// setup ///////////////////////////////////////////////////////////////////////
+template <typename FImpl, typename PhotonImpl>
+void TA2AAslashField<FImpl, PhotonImpl>::setup(void)
+{
+    envTmp(Computation, "computation", 1, envGetGrid(FermionField), 
+           env().getNd() - 1, par().emField.size(), 1, par().block, 
+           par().cacheBlock, this);
+    envTmp(std::vector<ComplexField>, "B0", 1, 
+           par().emField.size(), envGetGrid(ComplexField));
+    envTmp(std::vector<ComplexField>, "B1", 1, 
+           par().emField.size(), envGetGrid(ComplexField));
+    envTmpLat(ComplexField, "Amu");
+}
+
+// execution ///////////////////////////////////////////////////////////////////
+template <typename FImpl, typename PhotonImpl>
+void TA2AAslashField<FImpl, PhotonImpl>::execute(void)
+{
+    auto &left  = envGet(std::vector<FermionField>, par().left);
+    auto &right = envGet(std::vector<FermionField>, par().right);
+
+    int nt         = env().getDim().back();
+    int N_i        = left.size();
+    int N_j        = right.size();
+    int nem        = par().emField.size();
+    int block      = par().block;
+    int cacheBlock = par().cacheBlock;
+
+    LOG(Message) << "Computing all-to-all A-slash fields" << std::endl;
+    LOG(Message) << "Left: '" << par().left << "' Right: '" << par().right << "'" << std::endl;
+    LOG(Message) << "EM fields:" << std::endl;
+    for (auto &name: par().emField)
+    {
+        LOG(Message) << "  " << name << std::endl;
+    }
+    LOG(Message) << "A-slash field size: " << nt << "*" << N_i << "*" << N_j 
+                 << " (filesize " << sizeString(nt*N_i*N_j*sizeof(ASF_IO_TYPE)) 
+                 << "/EM field)" << std::endl;
+    
+    // preparing "B" complexified fields
+    startTimer("Complexify EM fields");
+    envGetTmp(std::vector<ComplexField>, B0);
+    envGetTmp(std::vector<ComplexField>, B1);
+    for (unsigned int i = 0; i < par().emField.size(); ++i)
+    {
+        auto &A = envGet(EmField, par().emField[i]);
+        envGetTmp(ComplexField, Amu);
+
+        B0[i]  = peekLorentz(A, 0);
+        B0[i] += timesI(peekLorentz(A, 1));
+        B1[i]  = peekLorentz(A, 2);
+        B1[i] += timesI(peekLorentz(A, 3));
+    }
+    stopTimer("Complexify EM fields");
+
+    // I/O name & metadata lambdas
+    auto ionameFn = [this](const unsigned int em, const unsigned int dummy)
+    {
+        return par().emField[em];
+    };
+
+    auto filenameFn = [this, &ionameFn](const unsigned int em, const unsigned int dummy)
+    {
+        return par().output + "." + std::to_string(vm().getTrajectory()) 
+               + "/" + ionameFn(em, dummy) + ".h5";
+    };
+
+    auto metadataFn = [this](const unsigned int em, const unsigned int dummy)
+    {
+        A2AAslashFieldMetadata md;
+
+        md.emFieldName = par().emField[em];
+        
+        return md;
+    };
+
+    // executing computation
+    Kernel kernel(B0, B1, envGetGrid(FermionField));
+
+    envGetTmp(Computation, computation);
+    computation.execute(left, right, kernel, ionameFn, filenameFn, metadataFn);
+}
+
+END_MODULE_NAMESPACE
+
+END_HADRONS_NAMESPACE
+
+#endif // Hadrons_MContraction_A2AAslashField_hpp_
diff --git a/Hadrons/modules.inc b/Hadrons/modules.inc
index 0e6ef893..866dfd91 100644
--- a/Hadrons/modules.inc
+++ b/Hadrons/modules.inc
@@ -4,6 +4,7 @@ modules_cc =\
   Modules/MContraction/Meson.cc \
   Modules/MContraction/WeakNeutral4ptDisc.cc \
   Modules/MContraction/WeakHamiltonianNonEye.cc \
+  Modules/MContraction/A2AAslashField.cc \
   Modules/MContraction/WardIdentity.cc \
   Modules/MContraction/A2AMesonField.cc \
   Modules/MContraction/DiscLoop.cc \
@@ -63,6 +64,7 @@ modules_cc =\
 
 modules_hpp =\
   Modules/MContraction/Baryon.hpp \
+  Modules/MContraction/A2AAslashField.hpp \
   Modules/MContraction/A2AMesonField.hpp \
   Modules/MContraction/Meson.hpp \
   Modules/MContraction/WeakHamiltonian.hpp \