Fix process divergence issues (could hang when autotuning) in generic…

…PrintMatrix routine. Apply same patch to genericPrintVector for future proofing
lattice · Nov 27, 2024 · 0a3f608 · 0a3f608
1 parent 2b81309
commit 0a3f608
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/lib/color_spinor_util.in.cu b/lib/color_spinor_util.in.cu
@@ -378,8 +378,6 @@ namespace quda {
 
   void genericPrintVector(const ColorSpinorField &a, int parity, unsigned int x_cb, int rank)
   {
-    if (rank != comm_rank()) return;
-
     ColorSpinorParam param(a);
     param.location = QUDA_CPU_FIELD_LOCATION;
     param.create = QUDA_COPY_FIELD_CREATE;
@@ -388,6 +386,8 @@ namespace quda {
     std::unique_ptr<ColorSpinorField> clone_a = !host_clone ? nullptr : std::make_unique<ColorSpinorField>(param);
     const ColorSpinorField &a_ = !host_clone ? a : *clone_a.get();
 
+    if (rank != comm_rank()) return; // rank returns after potential copy to host to prevent tuning hang
+
     switch (a.Precision()) {
     case QUDA_DOUBLE_PRECISION:  genericPrintVector<double>(a_, parity, x_cb); break;
     case QUDA_SINGLE_PRECISION:  genericPrintVector<float>(a_, parity, x_cb); break;

diff --git a/lib/gauge_norm.in.cu b/lib/gauge_norm.in.cu
@@ -159,8 +159,6 @@ namespace quda {
 
   void genericPrintMatrix(const GaugeField &a, int d, int parity, unsigned int x_cb, int rank)
   {
-    if (rank != comm_rank()) return;
-
     GaugeFieldParam param(a);
     param.field = const_cast<GaugeField *>(&a);
     param.location = QUDA_CPU_FIELD_LOCATION;
@@ -172,6 +170,8 @@ namespace quda {
     std::unique_ptr<GaugeField> clone_a = !host_clone ? nullptr : std::make_unique<GaugeField>(param);
     const GaugeField &a_ = !host_clone ? a : *clone_a.get();
 
+    if (rank != comm_rank()) return;  // rank returns after potential copy to host to prevent tuning hang
+
     switch (a.Precision()) {
     case QUDA_DOUBLE_PRECISION: genericPrintMatrix<double>(a_, d, parity, x_cb); break;
     case QUDA_SINGLE_PRECISION: genericPrintMatrix<float>(a_, d, parity, x_cb); break;