Merge pull request #1398 from lattice/feature/partfile-colorspinor

`PARTFILE` support for near-null vectors (and eigenvectors)

include/multigrid.h

@@ -157,6 +157,9 @@ namespace quda {
     /** Filename for where to load/store the null space */
     char filename[100];
 
+    /** Whether to save in partfile format (true) or singlefile (false) */
+    bool mg_vec_partfile;
+
     /** Whether or not this is a staggered solve or not */
     QudaTransferType transfer_type;
 
@@ -193,6 +196,7 @@ namespace quda {
       smoother_solve_type(param.smoother_solve_type[level]),
       location(param.location[level]),
       setup_location(param.setup_location[level]),
+      mg_vec_partfile(param.mg_vec_partfile[level]),
       transfer_type(param.transfer_type[level]),
       setup_use_mma(param.setup_use_mma[level] == QUDA_BOOLEAN_TRUE),
       dslash_use_mma(param.dslash_use_mma[level] == QUDA_BOOLEAN_TRUE)
@@ -230,6 +234,7 @@ namespace quda {
       smoother_solve_type(param.mg_global.smoother_solve_type[level]),
       location(param.mg_global.location[level]),
       setup_location(param.mg_global.setup_location[level]),
+      mg_vec_partfile(param.mg_global.mg_vec_partfile[level]),
       transfer_type(param.mg_global.transfer_type[level]),
       setup_use_mma(param.mg_global.setup_use_mma[level] == QUDA_BOOLEAN_TRUE),
       dslash_use_mma(param.mg_global.dslash_use_mma[level] == QUDA_BOOLEAN_TRUE)

include/qio_field.h

@@ -6,8 +6,9 @@ void read_gauge_field(const char *filename, void *gauge[], QudaPrecision prec, c
 void write_gauge_field(const char *filename, void *gauge[], QudaPrecision prec, const int *X, int argc, char *argv[]);
 void read_spinor_field(const char *filename, void *V[], QudaPrecision precision, const int *X, QudaSiteSubset subset,
                        QudaParity parity, int nColor, int nSpin, int Nvec, int argc, char *argv[]);
-void write_spinor_field(const char *filename, const void *V[], QudaPrecision precision, const int *X, QudaSiteSubset subset,
-                        QudaParity parity, int nColor, int nSpin, int Nvec, int argc, char *argv[]);
+void write_spinor_field(const char *filename, const void *V[], QudaPrecision precision, const int *X,
+                        QudaSiteSubset subset, QudaParity parity, int nColor, int nSpin, int Nvec, int argc,
+                        char *argv[], bool partfile = false);
 #else
 inline void read_gauge_field(const char *, void *[], QudaPrecision, const int *, int, char *[])
 {
@@ -25,8 +26,8 @@ inline void read_spinor_field(const char *, void *[], QudaPrecision, const int *
   printf("QIO support has not been enabled\n");
   exit(-1);
 }
-inline void write_spinor_field(const char *, const void *[], QudaPrecision, const int *, QudaSiteSubset, QudaParity, int, int,
-                               int, int, char *[])
+inline void write_spinor_field(const char *, const void *[], QudaPrecision, const int *, QudaSiteSubset, QudaParity,
+                               int, int, int, int, char *[], bool)
 {
   printf("QIO support has not been enabled\n");
   exit(-1);

include/quda.h

@@ -585,6 +585,9 @@ extern "C" {
         MILC I/O) */
     QudaBoolean io_parity_inflate;
 
+    /** Whether to save eigenvectors in QIO singlefile or partfile format */
+    QudaBoolean partfile;
+
     /** The Gflops rate of the eigensolver setup */
     double gflops;
 
@@ -779,6 +782,9 @@ extern "C" {
     /** Filename prefix for where to save the null-space vectors */
     char vec_outfile[QUDA_MAX_MG_LEVEL][256];
 
+    /** Whether to store the null-space vectors in singlefile or partfile format */
+    QudaBoolean mg_vec_partfile[QUDA_MAX_MG_LEVEL];
+
     /** Whether to use and initial guess during coarse grid deflation */
     QudaBoolean coarse_guess;
 

include/vector_io.h

@@ -15,15 +15,17 @@ namespace quda
   {
     const std::string filename;
     bool parity_inflate;
+    bool partfile;
 
   public:
     /**
        Constructor for VectorIO class
        @param[in] filename The filename associated with this IO object
        @param[in] parity_inflate Whether to inflate single_parity
        field to dual parity fields for I/O
+       @param[in] partfile Whether or not to save in partfiles (ignored on load)
     */
-    VectorIO(const std::string &filename, bool parity_inflate = false);
+    VectorIO(const std::string &filename, bool parity_inflate = false, bool partfile = false);
 
     /**
        @brief Load vectors from filename

lib/check_params.h

@@ -197,6 +197,7 @@ void printQudaEigParam(QudaEigParam *param) {
   P(extlib_type, QUDA_EIGEN_EXTLIB);
   P(mem_type_ritz, QUDA_MEMORY_DEVICE);
   P(ortho_block_size, 0);
+  P(partfile, QUDA_BOOLEAN_FALSE);
 #else
   P(use_eigen_qr, QUDA_BOOLEAN_INVALID);
   P(use_poly_acc, QUDA_BOOLEAN_INVALID);
@@ -226,6 +227,7 @@ void printQudaEigParam(QudaEigParam *param) {
   P(extlib_type, QUDA_EXTLIB_INVALID);
   P(mem_type_ritz, QUDA_MEMORY_INVALID);
   P(ortho_block_size, INVALID_INT);
+  P(partfile, QUDA_BOOLEAN_INVALID);
 #endif
 
   // only need to enfore block size checking if doing a block eigen solve
@@ -931,6 +933,12 @@ void printQudaMultigridParam(QudaMultigridParam *param) {
 #else
     P(setup_location[i], QUDA_INVALID_FIELD_LOCATION);
 #endif
+
+#ifdef INIT_PARAM
+    P(mg_vec_partfile[i], QUDA_BOOLEAN_FALSE);
+#else
+    P(mg_vec_partfile[i], QUDA_BOOLEAN_INVALID);
+#endif
   }
 
 #ifdef INIT_PARAM

lib/eigensolve_quda.cpp

@@ -255,7 +255,7 @@ namespace quda
       for (auto &k : kSpace) k.setSuggestedParity(mat_parity);
 
       // save the vectors
-      VectorIO io(eig_param->vec_outfile, eig_param->io_parity_inflate == QUDA_BOOLEAN_TRUE);
+      VectorIO io(eig_param->vec_outfile, eig_param->io_parity_inflate == QUDA_BOOLEAN_TRUE, eig_param->partfile);
       io.save(kSpace, save_prec, n_eig);
     }
 

lib/milc_interface.cpp

@@ -1464,6 +1464,7 @@ struct mgInputStruct {
   int setup_ca_basis_size[QUDA_MAX_MG_LEVEL];    // ignored on first and last level
   char mg_vec_infile[QUDA_MAX_MG_LEVEL][256];    // ignored on first and last level
   char mg_vec_outfile[QUDA_MAX_MG_LEVEL][256];   // ignored on first and last level
+  bool mg_vec_partfile[QUDA_MAX_MG_LEVEL];       // ignored on first and last level
   int geo_block_size[QUDA_MAX_MG_LEVEL][4]; // ignored on first and last level (values on first level are prescribed)
 
   /**
@@ -1496,6 +1497,7 @@ struct mgInputStruct {
   bool deflate_use_poly_acc;
   double deflate_a_min; // ignored if no polynomial acceleration
   int deflate_poly_deg; // ignored if no polynomial acceleration
+  bool deflate_vec_partfile;
 
   void setArrayDefaults()
   {
@@ -1511,6 +1513,7 @@ struct mgInputStruct {
       setup_ca_basis_size[i] = 4;
       mg_vec_infile[i][0] = 0;
       mg_vec_outfile[i][0] = 0;
+      mg_vec_partfile[i] = false;
       for (int d = 0; d < 4; d++) { geo_block_size[i][d] = 2; }
 
       setup_use_mma[i] = true;
@@ -1543,7 +1546,8 @@ struct mgInputStruct {
     deflate_tol(1e-5),
     deflate_use_poly_acc(false),
     deflate_a_min(1e-2),
-    deflate_poly_deg(50)
+    deflate_poly_deg(50),
+    deflate_vec_partfile(false)
   {
     /* initialize internal arrays */
     setArrayDefaults();
@@ -1840,6 +1844,12 @@ struct mgInputStruct {
         strcpy(mg_vec_outfile[atoi(input_line[1].c_str())], input_line[2].c_str());
       }
 
+    } else if (strcmp(input_line[0].c_str(), "mg_vec_partfile") == 0) {
+      if (input_line.size() < 3) {
+        error_code = 1;
+      } else {
+        mg_vec_partfile[atoi(input_line[1].c_str())] = input_line[2][0] == 't' ? true : false;
+      }
     } else /* Begin Solvers */
       if (strcmp(input_line[0].c_str(), "coarse_solve_type") == 0) {
       if (input_line.size() < 3) {
@@ -1947,6 +1957,12 @@ struct mgInputStruct {
         deflate_poly_deg = atoi(input_line[1].c_str());
       }
 
+    } else if (strcmp(input_line[0].c_str(), "deflate_vec_partfile") == 0) {
+      if (input_line.size() < 2) {
+        error_code = 1;
+      } else {
+        deflate_vec_partfile = input_line[1][0] == 't' ? true : false;
+      }
     } else {
       printf("Invalid option %s\n", input_line[0].c_str());
       return false;
@@ -2011,6 +2027,7 @@ void milcSetMultigridEigParam(QudaEigParam &mg_eig_param, mgInputStruct &input_s
   strcpy(mg_eig_param.vec_outfile, "");
   mg_eig_param.io_parity_inflate = QUDA_BOOLEAN_FALSE; // do not inflate coarse vectors
   mg_eig_param.save_prec = QUDA_SINGLE_PRECISION;      // cannot save in fixed point
+  mg_eig_param.partfile = QUDA_BOOLEAN_FALSE;          // ignored, multigrid parameters take precedence
 
   strcpy(mg_eig_param.QUDA_logfile, "" /*eig_QUDA_logfile*/);
 }
@@ -2357,6 +2374,10 @@ void milcSetMultigridParam(milcMultigridPack *mg_pack, QudaPrecision host_precis
     strcpy(mg_param.vec_outfile[i], input_struct.mg_vec_outfile[i]);
     if (strcmp(mg_param.vec_infile[i], "") != 0) mg_param.vec_load[i] = QUDA_BOOLEAN_TRUE;
     if (strcmp(mg_param.vec_outfile[i], "") != 0) mg_param.vec_store[i] = QUDA_BOOLEAN_TRUE;
+    if (i != mg_param.n_level - 1)
+      mg_param.mg_vec_partfile[i] = input_struct.mg_vec_partfile[i] ? QUDA_BOOLEAN_TRUE : QUDA_BOOLEAN_FALSE;
+    else
+      mg_param.mg_vec_partfile[i] = input_struct.deflate_vec_partfile ? QUDA_BOOLEAN_TRUE : QUDA_BOOLEAN_FALSE;
   }
 
   mg_param.coarse_guess = QUDA_BOOLEAN_FALSE; // mg_eig_coarse_guess ? QUDA_BOOLEAN_TRUE : QUDA_BOOLEAN_FALSE;

lib/multigrid.cpp

@@ -672,6 +672,7 @@ namespace quda
           vec_outfile += "_defl_";
           vec_outfile += std::to_string(param.mg_global.n_vec[param.level + 1]);
           strcpy(param_coarse_solver->eig_param.vec_outfile, vec_outfile.c_str());
+          param_coarse_solver->eig_param.partfile = param.mg_global.mg_vec_partfile[param.level + 1];
         }
       }
 
@@ -1382,7 +1383,7 @@ namespace quda
       vec_outfile += std::to_string(param.level);
       vec_outfile += "_nvec_";
       vec_outfile += std::to_string(param.mg_global.n_vec[param.level]);
-      VectorIO io(vec_outfile);
+      VectorIO io(vec_outfile, false, param.mg_global.mg_vec_partfile[param.level]);
       vector_ref<const ColorSpinorField> B_ref;
       for (auto i = 0u; i < B.size(); i++) B_ref.push_back(*B[i]);
       io.save(std::move(B_ref));

lib/qio_field.cpp

@@ -413,8 +413,9 @@ void write_gauge_field(const char *filename, void *gauge[], QudaPrecision precis
   printfQuda("%s: Closed file for writing\n", __func__);
 }
 
-void write_spinor_field(const char *filename, const void *V[], QudaPrecision precision, const int *X, QudaSiteSubset subset,
-                        QudaParity parity, int nColor, int nSpin, int Nvec, int, char *[])
+void write_spinor_field(const char *filename, const void *V[], QudaPrecision precision, const int *X,
+                        QudaSiteSubset subset, QudaParity parity, int nColor, int nSpin, int Nvec, int, char *[],
+                        bool partfile)
 {
   quda_this_node = QMP_get_node_number();
 
@@ -426,7 +427,7 @@ void write_spinor_field(const char *filename, const void *V[], QudaPrecision pre
   sprintf(type, "QUDA_%sNs%dNc%d_ColorSpinorField", (file_prec == QUDA_DOUBLE_PRECISION) ? "D" : "F", nSpin, nColor);
 
   /* Open the test file for reading */
-  QIO_Writer *outfile = open_test_output(filename, QIO_SINGLEFILE, QIO_PARALLEL, QIO_ILDGNO);
+  QIO_Writer *outfile = open_test_output(filename, (partfile ? QIO_PARTFILE : QIO_SINGLEFILE), QIO_PARALLEL, QIO_ILDGNO);
   if (outfile == NULL) { errorQuda("Open file failed\n"); }
 
   /* Read the spinor field record */

lib/vector_io.cpp

@@ -2,13 +2,13 @@
 #include <qio_field.h>
 #include <vector_io.h>
 #include <blas_quda.h>
+#include <timer.h>
 
 namespace quda
 {
 
-  VectorIO::VectorIO(const std::string &filename, bool parity_inflate) :
-    filename(filename),
-    parity_inflate(parity_inflate)
+  VectorIO::VectorIO(const std::string &filename, bool parity_inflate, bool partfile) :
+    filename(filename), parity_inflate(parity_inflate), partfile(partfile)
   {
     if (strcmp(filename.c_str(), "") == 0)
       errorQuda("No eigenspace input file defined (filename = %s, parity_inflate = %d", filename.c_str(), parity_inflate);
@@ -55,8 +55,15 @@ namespace quda
         for (int j = 0; j < Ls; j++) { V[i * Ls + j] = static_cast<char *>(v.V()) + j * stride; }
       }
 
+      // time loading
+      quda::host_timer_t host_timer;
+      host_timer.start(); // start the timer
+
       read_spinor_field(filename.c_str(), V.data(), v0.Precision(), v0.X(), v0.SiteSubset(),
                         spinor_parity, v0.Ncolor(), v0.Nspin(), Nvec * Ls, 0, nullptr);
+
+      host_timer.stop(); // stop the timer
+      logQuda(QUDA_SUMMARIZE, "Time spent loading vectors from %s = %g secs\n", filename.c_str(), host_timer.last());
     } else {
       errorQuda("Unexpected field dimension %d", v0.Ndim());
     }
@@ -114,7 +121,12 @@ namespace quda
       }
     }
 
-    if (getVerbosity() >= QUDA_SUMMARIZE) printfQuda("Start saving %d vectors to %s\n", Nvec, filename.c_str());
+    if (getVerbosity() >= QUDA_SUMMARIZE) {
+      if (partfile)
+        printfQuda("Start saving %d vectors to %s in PARTFILE format\n", Nvec, filename.c_str());
+      else
+        printfQuda("Start saving %d vectors to %s in SINGLEFILE format\n", Nvec, filename.c_str());
+    }
 
     if (v0.Ndim() == 4 || v0.Ndim() == 5) {
       // since QIO routines presently assume we have 4-d fields, we need to convert to array of 4-d fields
@@ -128,8 +140,15 @@ namespace quda
         for (int j = 0; j < Ls; j++) { V[i * Ls + j] = static_cast<const char *>(v.V()) + j * stride; }
       }
 
-      write_spinor_field(filename.c_str(), V.data(), save_prec, v0.X(), v0.SiteSubset(),
-                         spinor_parity, v0.Ncolor(), v0.Nspin(), Nvec * Ls, 0, nullptr);
+      // time saving
+      quda::host_timer_t host_timer;
+      host_timer.start(); // start the timer
+
+      write_spinor_field(filename.c_str(), V.data(), save_prec, v0.X(), v0.SiteSubset(), spinor_parity, v0.Ncolor(),
+                         v0.Nspin(), Nvec * Ls, 0, nullptr, partfile);
+
+      host_timer.stop(); // stop the timer
+      logQuda(QUDA_SUMMARIZE, "Time spent saving vectors to %s = %g secs\n", filename.c_str(), host_timer.last());
     } else {
       errorQuda("Unexpected field dimension %d", v0.Ndim());
     }

tests/io_test.cpp

@@ -71,7 +71,7 @@ TEST_P(GaugeIOTest, verify)
   for (int dir = 0; dir < 4; dir++) { host_free(gauge[dir]); }
 }
 
-using cs_test_t = ::testing::tuple<QudaSiteSubset, bool, QudaPrecision, QudaPrecision, int, QudaFieldLocation>;
+using cs_test_t = ::testing::tuple<QudaSiteSubset, bool, QudaPrecision, QudaPrecision, int, bool, QudaFieldLocation>;
 
 class ColorSpinorIOTest : public ::testing::TestWithParam<cs_test_t>
 {
@@ -81,6 +81,7 @@ class ColorSpinorIOTest : public ::testing::TestWithParam<cs_test_t>
   QudaPrecision prec;
   QudaPrecision prec_io;
   int nSpin;
+  bool partfile;
   QudaFieldLocation location;
 
 public:
@@ -90,7 +91,8 @@ class ColorSpinorIOTest : public ::testing::TestWithParam<cs_test_t>
     prec(::testing::get<2>(GetParam())),
     prec_io(::testing::get<3>(GetParam())),
     nSpin(::testing::get<4>(GetParam())),
-    location(::testing::get<5>(GetParam()))
+    partfile(::testing::get<5>(GetParam())),
+    location(::testing::get<6>(GetParam()))
   {
   }
 };
@@ -139,7 +141,7 @@ TEST_P(ColorSpinorIOTest, verify)
 
   auto file = "dummy.cs";
 
-  VectorIO io(file, inflate);
+  VectorIO io(file, inflate, partfile);
 
   io.save({v.begin(), v.end()}, prec_io, n_vector);
   io.load(u);
@@ -153,7 +155,16 @@ TEST_P(ColorSpinorIOTest, verify)
   }
 
   // cleanup after ourselves and delete the dummy lattice
-  if (::quda::comm_rank() == 0 && remove(file) != 0) errorQuda("Error deleting file");
+  if (partfile && ::quda::comm_size() > 1) {
+    // each rank created its own file, we need to generate the custom filename
+    // an exception is single-rank runs where QIO skips appending the volume string
+    char volstr[9];
+    sprintf(volstr, ".vol%04d", ::quda::comm_rank());
+    std::string part_filename = std::string(file) + volstr;
+    if (remove(part_filename.c_str()) != 0) errorQuda("Error deleting file");
+  } else {
+    if (::quda::comm_rank() == 0 && remove(file) != 0) errorQuda("Error deleting file");
+  }
 }
 
 int main(int argc, char **argv)
@@ -177,13 +188,14 @@ INSTANTIATE_TEST_SUITE_P(Full, ColorSpinorIOTest,
                          Combine(Values(QUDA_FULL_SITE_SUBSET), Values(false),
                                  Values(QUDA_DOUBLE_PRECISION, QUDA_SINGLE_PRECISION, QUDA_HALF_PRECISION),
                                  Values(QUDA_DOUBLE_PRECISION, QUDA_SINGLE_PRECISION), Values(1, 2, 4),
-                                 Values(QUDA_CUDA_FIELD_LOCATION, QUDA_CPU_FIELD_LOCATION)),
+                                 Values(false, true), Values(QUDA_CUDA_FIELD_LOCATION, QUDA_CPU_FIELD_LOCATION)),
                          [](testing::TestParamInfo<cs_test_t> param) {
                            std::string name;
                            name += get_prec_str(::testing::get<2>(param.param)) + std::string("_");
                            name += get_prec_str(::testing::get<3>(param.param)) + std::string("_");
                            name += std::string("spin") + std::to_string(::testing::get<4>(param.param));
-                           name += ::testing::get<5>(param.param) == QUDA_CUDA_FIELD_LOCATION ? "_device" : "_host";
+                           name += ::testing::get<5>(param.param) ? "_singlefile" : "_partfile";
+                           name += ::testing::get<6>(param.param) == QUDA_CUDA_FIELD_LOCATION ? "_device" : "_host";
                            return name;
                          });
 
@@ -192,13 +204,14 @@ INSTANTIATE_TEST_SUITE_P(Parity, ColorSpinorIOTest,
                          Combine(Values(QUDA_PARITY_SITE_SUBSET), Values(false, true),
                                  Values(QUDA_DOUBLE_PRECISION, QUDA_SINGLE_PRECISION, QUDA_HALF_PRECISION),
                                  Values(QUDA_DOUBLE_PRECISION, QUDA_SINGLE_PRECISION), Values(1, 2, 4),
-                                 Values(QUDA_CUDA_FIELD_LOCATION, QUDA_CPU_FIELD_LOCATION)),
+                                 Values(false, true), Values(QUDA_CUDA_FIELD_LOCATION, QUDA_CPU_FIELD_LOCATION)),
                          [](testing::TestParamInfo<cs_test_t> param) {
                            std::string name;
                            if (::testing::get<1>(param.param)) name += std::string("inflate_");
                            name += get_prec_str(::testing::get<2>(param.param)) + std::string("_");
                            name += get_prec_str(::testing::get<3>(param.param)) + std::string("_");
                            name += std::string("spin") + std::to_string(::testing::get<4>(param.param));
-                           name += ::testing::get<5>(param.param) == QUDA_CUDA_FIELD_LOCATION ? "_device" : "_host";
+                           name += ::testing::get<5>(param.param) ? "_singlefile" : "_partfile";
+                           name += ::testing::get<6>(param.param) == QUDA_CUDA_FIELD_LOCATION ? "_device" : "_host";
                            return name;
                          });