Merge pull request #1504 from lattice/feature/full-qdpjit-fields

Mods to allow copy in/out of full QDP-JIT spinors
lattice · Oct 29, 2024 · aebd9b2 · aebd9b2
2 parents 9521952 + 76947a3
commit aebd9b2
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 11 deletions.
diff --git a/include/color_spinor_field_order.h b/include/color_spinor_field_order.h
@@ -1799,8 +1799,8 @@ namespace quda
       {
         for (int s = 0; s < Ns; s++) {
           for (int c = 0; c < Nc; c++) {
-            v[s * Nc + c] = complex(field[(((0 * Nc + c) * Ns + s) * 2 + (1 - parity)) * volumeCB + x],
-                                    field[(((1 * Nc + c) * Ns + s) * 2 + (1 - parity)) * volumeCB + x]);
+            v[s * Nc + c] = complex(field[(((0 * Nc + c) * Ns + s) * 2 + parity) * volumeCB + x],
+                                    field[(((1 * Nc + c) * Ns + s) * 2 + parity) * volumeCB + x]);
           }
         }
       }
@@ -1809,8 +1809,8 @@ namespace quda
       {
         for (int s = 0; s < Ns; s++) {
           for (int c = 0; c < Nc; c++) {
-            field[(((0 * Nc + c) * Ns + s) * 2 + (1 - parity)) * volumeCB + x] = v[s * Nc + c].real();
-            field[(((1 * Nc + c) * Ns + s) * 2 + (1 - parity)) * volumeCB + x] = v[s * Nc + c].imag();
+            field[(((0 * Nc + c) * Ns + s) * 2 + parity) * volumeCB + x] = v[s * Nc + c].real();
+            field[(((1 * Nc + c) * Ns + s) * 2 + parity) * volumeCB + x] = v[s * Nc + c].imag();
           }
         }
       }

diff --git a/include/kernels/copy_color_spinor.cuh b/include/kernels/copy_color_spinor.cuh
@@ -16,6 +16,30 @@ namespace quda
 
   using namespace colorspinor;
 
+  /**
+   * @brief A helper function to figure out what parity to use for input and output.
+   * @details Pick parity from input field site order. Addditionally QDPJIT fields
+   *  may need a relative parity flip compared to what is expected when dealing with
+   *  only the odd parity since the pointer is always to the top of the full spinort.
+   * @param[in] f Reference to the field for parity computation
+   * @return the computed parity
+   */
+  inline int computeParity(const ColorSpinorField &f)
+  {
+
+    // Account for odd-even vs. even-odd site orders
+    int ret_val = f.SiteOrder() == QUDA_ODD_EVEN_SITE_ORDER ? 1 : 0;
+
+    // Account for potential parity flip to access single parity subset QDP-JIT fields
+    // The Flip is only needed fir offsetting into Odd Parity Fields
+    if (f.FieldOrder() == QUDA_QDPJIT_FIELD_ORDER && f.SiteSubset() == QUDA_PARITY_SITE_SUBSET
+        && f.SuggestedParity() == QUDA_ODD_PARITY) {
+      ret_val = 1 - ret_val;
+    }
+
+    return ret_val;
+  }
+
   template <typename FloatOut, typename FloatIn, int nSpin_, int nColor_, typename Out, typename In,
             template <int, int> class Basis_>
   struct CopyColorSpinorArg : kernel_param<> {
@@ -32,8 +56,8 @@ namespace quda
       kernel_param(dim3(in.VolumeCB(), in.SiteSubset(), 1)),
       out(out, 1, Out_),
       in(in, 1, const_cast<FloatIn *>(In_)),
-      outParity(out.SiteOrder() == QUDA_ODD_EVEN_SITE_ORDER ? 1 : 0),
-      inParity(in.SiteOrder() == QUDA_ODD_EVEN_SITE_ORDER ? 1 : 0)
+      outParity(computeParity(out)),
+      inParity(computeParity(in))
     {
     }
   };

diff --git a/lib/copy_color_spinor.cuh b/lib/copy_color_spinor.cuh
@@ -181,11 +181,6 @@ namespace quda
       errorQuda("Copying to full fields with lexicographical ordering is not currently supported");
     }
 
-    if (dst.SiteSubset() == QUDA_FULL_SITE_SUBSET
-        && (src.FieldOrder() == QUDA_QDPJIT_FIELD_ORDER || dst.FieldOrder() == QUDA_QDPJIT_FIELD_ORDER)) {
-      errorQuda("QDPJIT field ordering not supported for full site fields");
-    }
-
     genericCopyColorSpinor<Ns, Nc>(param);
   }