Use batching where possible in dirac_quda.h

lattice · Dec 11, 2024 · 5d6e6e8 · 5d6e6e8
1 parent 49a0211
commit 5d6e6e8
Showing 1 changed file with 7 additions and 13 deletions.
diff --git a/include/dirac_quda.h b/include/dirac_quda.h
@@ -2347,8 +2347,7 @@ namespace quda {
     void operator()(cvector_ref<ColorSpinorField> &out, cvector_ref<const ColorSpinorField> &in) const override
     {
       dirac->M(out, in);
-      for (auto i = 0u; i < in.size(); i++)
-        if (shift != 0.0) blas::axpy(shift, in[i], out[i]);
+      if (shift != 0.0) blas::axpy(shift, in, out);
     }
 
     int getStencilSteps() const override { return dirac->getStencilSteps(); }
@@ -2369,8 +2368,7 @@ namespace quda {
     void operator()(cvector_ref<ColorSpinorField> &out, cvector_ref<const ColorSpinorField> &in) const override
     {
       dirac->MdagM(out, in);
-      for (auto i = 0u; i < in.size(); i++)
-        if (shift != 0.0) blas::axpy(shift, in[i], out[i]);
+      if (shift != 0.0) blas::axpy(shift, in, out);
     }
 
     int getStencilSteps() const override
@@ -2421,8 +2419,7 @@ namespace quda {
     void operator()(cvector_ref<ColorSpinorField> &out, cvector_ref<const ColorSpinorField> &in) const override
     {
       dirac->MMdag(out, in);
-      for (auto i = 0u; i < in.size(); i++)
-        if (shift != 0.0) blas::axpy(shift, in[i], out[i]);
+      if (shift != 0.0) blas::axpy(shift, in, out);
     }
 
     int getStencilSteps() const override
@@ -2448,8 +2445,7 @@ namespace quda {
     void operator()(cvector_ref<ColorSpinorField> &out, cvector_ref<const ColorSpinorField> &in) const override
     {
       dirac->Mdag(out, in);
-      for (auto i = 0u; i < in.size(); i++)
-        if (shift != 0.0) blas::axpy(shift, in[i], out[i]);
+      if (shift != 0.0) blas::axpy(shift, in, out);
     }
 
     int getStencilSteps() const override { return dirac->getStencilSteps(); }
@@ -2496,7 +2492,7 @@ namespace quda {
 
       @param vec[in,out] vector to which gamma5 is applied in place
     */
-    void applyGamma5(ColorSpinorField &vec) const
+    void applyGamma5(cvector_ref<ColorSpinorField> &vec) const
     {
       auto dirac_type = dirac->getDiracType();
       auto pc_type = dirac->getMatPCType();
@@ -2573,10 +2569,8 @@ namespace quda {
     void operator()(cvector_ref<ColorSpinorField> &out, cvector_ref<const ColorSpinorField> &in) const override
     {
       dirac->M(out, in);
-      for (auto i = 0u; i < in.size(); i++) {
-        if (shift != 0.0) blas::axpy(shift, in[i], out[i]);
-        applyGamma5(out[i]);
-      }
+      if (shift != 0.0) blas::axpy(shift, in, out);
+      applyGamma5(out);
     }
 
     int getStencilSteps() const override { return dirac->getStencilSteps(); }