From a822c48565ccd3352b2d7f4682eabf1365249b6e Mon Sep 17 00:00:00 2001
From: "Henrique B.R" <h.b.rocha@ed.ac.uk>
Date: Fri, 24 Sep 2021 17:13:25 +0100
Subject: [PATCH] Added accelerated pick-set checkerboard functions

---
 Grid/lattice/Lattice_transfer.h | 71 +++++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/Grid/lattice/Lattice_transfer.h b/Grid/lattice/Lattice_transfer.h
index 2292088c..4b1459c3 100644
--- a/Grid/lattice/Lattice_transfer.h
+++ b/Grid/lattice/Lattice_transfer.h
@@ -85,6 +85,77 @@ template<class vobj> inline void setCheckerboard(Lattice<vobj> &full,const Latti
   });
 }
 
+template<class vobj> inline void acceleratorPickCheckerboard(int cb,Lattice<vobj> &half,const Lattice<vobj> &full, int checker_dim_half=0)
+{
+  half.Checkerboard() = cb;
+  autoView(half_v, half, AcceleratorWrite);
+  autoView(full_v, full, AcceleratorRead);
+  Coordinate rdim_full             = full.Grid()->_rdimensions;
+  Coordinate rdim_half             = half.Grid()->_rdimensions;
+  unsigned long ndim_half          = half.Grid()->_ndimension;
+  Coordinate checker_dim_mask_half = half.Grid()->_checker_dim_mask;
+  Coordinate ostride_half          = half.Grid()->_ostride;
+  accelerator_for(ss, full.Grid()->oSites(),full.Grid()->Nsimd(),{
+    
+    Coordinate coor;
+    int cbos;
+    int linear=0;
+
+    Lexicographic::CoorFromIndex(coor,ss,rdim_full);
+    assert(coor.size()==ndim_half);
+
+    for(int d=0;d<ndim_half;d++){ 
+      if(checker_dim_mask_half[d]) linear += coor[d];
+    }
+    cbos = (linear&0x1);
+
+    if (cbos==cb) {
+      int ssh=0;
+      for(int d=0;d<ndim_half;d++) {
+        if (d == checker_dim_half) ssh += ostride_half[d] * ((coor[d] / 2) % rdim_half[d]);
+        else ssh += ostride_half[d] * (coor[d] % rdim_half[d]);
+      }
+      coalescedWrite(half_v[ssh],full_v(ss));
+    }
+  });
+}
+template<class vobj> inline void acceleratorSetCheckerboard(Lattice<vobj> &full,const Lattice<vobj> &half, int checker_dim_half=0)
+{
+  int cb = half.Checkerboard();
+  autoView(half_v , half, AcceleratorRead);
+  autoView(full_v , full, AcceleratorWrite);
+  nvtxRangePop();
+  Coordinate rdim_full             = full.Grid()->_rdimensions;
+  Coordinate rdim_half             = half.Grid()->_rdimensions;
+  unsigned long ndim_half          = half.Grid()->_ndimension;
+  Coordinate checker_dim_mask_half = half.Grid()->_checker_dim_mask;
+  Coordinate ostride_half          = half.Grid()->_ostride;
+  accelerator_for(ss,full.Grid()->oSites(),full.Grid()->Nsimd(),{
+
+    Coordinate coor;
+    int cbos;
+    int linear=0;
+  
+    Lexicographic::CoorFromIndex(coor,ss,rdim_full);
+    assert(coor.size()==ndim_half);
+
+    for(int d=0;d<ndim_half;d++){ 
+      if(checker_dim_mask_half[d]) linear += coor[d];
+    }
+    cbos = (linear&0x1);
+
+    if (cbos==cb) {
+      int ssh=0;
+      for(int d=0;d<ndim_half;d++){
+        if (d == checker_dim_half) ssh += ostride_half[d] * ((coor[d] / 2) % rdim_half[d]);
+        else ssh += ostride_half[d] * (coor[d] % rdim_half[d]);
+      }
+      coalescedWrite(full_v[ss],half_v(ssh));
+    }
+
+  });
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////
 // Flexible Type Conversion for internal promotion to double as well as graceful
 // treatment of scalar-compatible types