Checkpointing (quasi-Newton solver)

CMakeLists.txt

@@ -27,7 +27,8 @@ option(HIOP_USE_EIGEN "Build with Eigen support" ON)
 option(HIOP_USE_MPI "Build with MPI support" ON)
 option(HIOP_USE_GPU "Build with support for GPUs - CUDA or HIP libraries" OFF)
 option(HIOP_TEST_WITH_BSUB "Use `jsrun` instead of `mpirun` commands when running tests" OFF)
-option(HIOP_USE_RAJA   "Build with portability abstraction library RAJA" OFF)
+option(HIOP_USE_RAJA "Build with portability abstraction library RAJA" OFF)
+option(HIOP_USE_AXOM "Build with AXOM to use Sidre for scalable checkpointing" OFF)
 option(HIOP_DEEPCHECKS "Extra checks and asserts in the code with a high penalty on performance" OFF)
 option(HIOP_WITH_KRON_REDUCTION "Build Kron Reduction code (requires UMFPACK)" OFF)
 option(HIOP_DEVELOPER_MODE "Build with extended warnings and options" OFF)
@@ -289,6 +290,15 @@ if(HIOP_USE_RAJA)
   message(STATUS "Found umpire pkg-config: ${umpire_CONFIG}")
 endif()
 
+if(HIOP_USE_AXOM)
+  find_package(AXOM CONFIG
+    PATHS ${AXOM_DIR} ${AXOM_DIR}/lib/cmake/
+    REQUIRED)
+  target_link_libraries(hiop_tpl INTERFACE axom)
+  message(STATUS "Found AXOM pkg-config: ${AXOM_CONFIG}")
+endif()
+
+
 if(HIOP_WITH_KRON_REDUCTION)
   set(HIOP_UMFPACK_DIR CACHE PATH "Path to UMFPACK directory")
   include(FindUMFPACK)

src/Interface/hiopInterface.hpp

@@ -467,8 +467,8 @@ class hiopInterfaceBase
   }
 
   /** 
-   * This method is used to provide an user all the hiop iterate
-   * procedure. @see solution_callback() for an explanation of the parameters.
+   * This method is used to provide user all the internal hiop iterates. @see solution_callback() 
+   * for an explanation of the parameters.
    * 
    * @param[in] x array of (local) entries of the primal variables (managed by Umpire, see note below)
    * @param[in] z_L array of (local) entries of the dual variables for lower bounds (managed by Umpire, see note below)
@@ -496,7 +496,7 @@ class hiopInterfaceBase
   {
     return true;
   }
-
+  
   /**
    * A wildcard function used to change the primal variables.
    *

src/Interface/hiop_defs.hpp.in

@@ -12,6 +12,7 @@
 #cmakedefine HIOP_USE_PARDISO
 #cmakedefine HIOP_USE_RESOLVE
 #cmakedefine HIOP_USE_GINKGO
+#cmakedefine HIOP_USE_AXOM
 #define HIOP_VERSION  "@PROJECT_VERSION@"
 #define HIOP_VERSION_MAJOR @PROJECT_VERSION_MAJOR@
 #define HIOP_VERSION_MINOR @PROJECT_VERSION_MINOR@

src/Optimization/hiopAlgFilterIPM.cpp

@@ -66,6 +66,11 @@
 
 #include "hiopCppStdUtils.hpp"
 
+#ifdef HIOP_USE_AXOM
+#include "SidreHelper.hpp"
+using namespace axom;
+#endif
+
 #include <cmath>
 #include <cstring>
 #include <cassert>
@@ -976,8 +981,38 @@ hiopSolveStatus hiopAlgFilterIPMQuasiNewton::run()
 
   nlp->runStats.tmOptimizTotal.start();
 
-  startingProcedure(*it_curr, _f_nlp, *_c, *_d, *_grad_f, *_Jac_c, *_Jac_d); //this also evaluates the nlp
-  _mu=mu0;
+  //
+  // starting point:
+  // - user provided (with slack adjustments and lsq eq. duals initialization)
+  // or
+  // - loaded checkpoint
+  //
+  if(nlp->options->GetString("checkpoint_load_on_start") != "yes") {
+    //this also evaluates the nlp
+    startingProcedure(*it_curr, _f_nlp, *_c, *_d, *_grad_f, *_Jac_c, *_Jac_d); 
+    _mu=mu0;
+    iter_num = 0;
+  } else {
+    //
+    //checkpoint load
+    //
+    //load from file: will populate it_curr, _Hess_lagr, and algorithmic parameters
+    auto chkpnt_ok = load_state_from_file(nlp->options->GetString("checkpoint_file"));
+    if(chkpnt_ok) {
+      //additionally: need to evaluate the nlp
+      if(!this->evalNlp_noHess(*it_curr, _f_nlp, *_c, *_d, *_grad_f, *_Jac_c, *_Jac_d)) {
+        nlp->log->printf(hovError, "Failure in evaluating user NLP functions at loaded checkpoint.");
+        return Error_In_User_Function;
+      }
+    } else {
+      nlp->log->printf(hovWarning, "Using default starting procedure (no checkpoint load!).\n");
+      //this also evaluates the nlp
+      startingProcedure(*it_curr, _f_nlp, *_c, *_d, *_grad_f, *_Jac_c, *_Jac_d); 
+      _mu=mu0;
+      iter_num = 0;
+    }
+    solver_status_ = NlpSolve_SolveNotCalled;
+  }
 
   //update log bar
   logbar->updateWithNlpInfo(*it_curr, _mu, _f_nlp, *_c, *_d, *_grad_f, *_Jac_c, *_Jac_d);
@@ -987,7 +1022,7 @@ hiopSolveStatus hiopAlgFilterIPMQuasiNewton::run()
 
   nlp->log->write("First residual-------------", *resid, hovIteration);
 
-  iter_num=0; nlp->runStats.nIter=iter_num;
+  nlp->runStats.nIter = iter_num;
   bool disableLS = nlp->options->GetString("accept_every_trial_step")=="yes";
 
   theta_max = theta_max_fact_*fmax(1.0,resid->get_theta());
@@ -1095,6 +1130,11 @@ hiopSolveStatus hiopAlgFilterIPMQuasiNewton::run()
       solver_status_ = User_Stopped; break;
     }
 
+#ifdef HIOP_USE_AXOM    
+    //checkpointing - based on options provided by the user
+    checkpointing_stuff();
+#endif
+
     /*************************************************
      * Termination check
      ************************************************/
@@ -1485,6 +1525,149 @@ void hiopAlgFilterIPMQuasiNewton::outputIteration(int lsStatus, int lsNum, int u
   }
 }
 
+#ifdef HIOP_USE_AXOM
+
+bool hiopAlgFilterIPMQuasiNewton::save_state_to_file(const ::std::string& path) noexcept
+{
+  try {
+    sidre::DataStore ds;
+    sidre::Group* group = ds.getRoot()->createGroup("HiOp quasi-Newton alg state");
+    this->save_state_to_sidre_group(*group);
+
+    sidre::IOManager writer(this->get_nlp()->get_comm());
+    int n_files;
+    MPI_Comm_size(this->get_nlp()->get_comm(), &n_files);
+    writer.write(ds.getRoot(), n_files, path, sidre::Group::getDefaultIOProtocol());
+    return true;
+  } catch(const std::exception& exp) {
+    nlp->log->printf(hovError, "Error when saving checkpoint to file '%s'\n", path.c_str());
+    nlp->log->printf(hovError,  "  Addtl info: %s\n", exp.what());
+    return false;    
+  }  
+}
+
+bool hiopAlgFilterIPMQuasiNewton::load_state_from_file(const ::std::string& path) noexcept
+{
+  try {
+    sidre::DataStore ds;
+
+    sidre::IOManager reader(this->get_nlp()->get_comm());
+    auto path2 = SidreHelper::check_path(path);
+    reader.read(ds.getRoot(), path2, false);
+    nlp->log->printf(hovScalars, "Loaded checkpoint file [%s].\n", path2.c_str());
+
+    const sidre::Group* group = ds.getRoot()->getGroup("HiOp quasi-Newton alg state");
+    this->load_state_from_sidre_group(*group);
+    return true;
+  } catch(const std::exception& exp) {
+    nlp->log->printf(hovError, "Error in loading checkpoint from file '%s'\n", path.c_str());
+    nlp->log->printf(hovError,  "  Addtl info: %s\n", exp.what());
+    return false;
+  }
+}
+
+void hiopAlgFilterIPMQuasiNewton::save_state_to_sidre_group(::axom::sidre::Group& group)
+{
+  using IndType = sidre::IndexType;
+
+  //metadata
+
+  //iterate states
+  //create views for each member that needs to be saved
+  SidreHelper::copy_vec_to_view(group, "x", *it_curr->get_x());
+  SidreHelper::copy_vec_to_view(group, "d", *it_curr->get_d());
+  SidreHelper::copy_vec_to_view(group, "sxl", *it_curr->get_sxl());
+  SidreHelper::copy_vec_to_view(group, "sxu", *it_curr->get_sxu());
+  SidreHelper::copy_vec_to_view(group, "sdl", *it_curr->get_sdl());
+  SidreHelper::copy_vec_to_view(group, "sdu", *it_curr->get_sdu());
+  SidreHelper::copy_vec_to_view(group, "yc", *it_curr->get_yc());
+  SidreHelper::copy_vec_to_view(group, "yd", *it_curr->get_yd());
+  SidreHelper::copy_vec_to_view(group, "zl", *it_curr->get_zl());
+  SidreHelper::copy_vec_to_view(group, "zu", *it_curr->get_zu());
+  SidreHelper::copy_vec_to_view(group, "vl", *it_curr->get_vl());
+  SidreHelper::copy_vec_to_view(group, "vu", *it_curr->get_vu());
+
+  //state of quasi-Newton Hessian approximation 
+
+  //algorithmic parameters for this state
+  //mu, iteration number, num MPI ranks
+  int nranks=1;
+#ifdef HIOP_USE_MPI  
+  MPI_Comm_size(get_nlp()->get_comm(), &nranks);
+#endif
+
+  const double alg_params[] = {_mu, (double)iter_num, (double)nranks};
+  const size_type nparams = sizeof(alg_params) / sizeof(double);
+
+  SidreHelper::copy_array_to_view(group, "alg_params", alg_params, nparams);
+}
+
+void hiopAlgFilterIPMQuasiNewton::load_state_from_sidre_group(const sidre::Group& group)
+{
+  //metadata
+
+  //algorithmic parameters
+  //!!! dev note: nparams needs to match the nparams from save_state_to_data_store
+  const int nparams = 3; 
+  double alg_params[nparams];
+  SidreHelper::copy_array_from_view(group, "alg_params", alg_params, nparams);
+  //!!! dev note: match order in save_state_to_data_store
+  _mu = alg_params[0];
+  iter_num = alg_params[1];
+
+  int nranks=1;
+#ifdef HIOP_USE_MPI  
+  MPI_Comm_size(get_nlp()->get_comm(), &nranks);
+#endif
+  if( (int)alg_params[2] != nranks ) {
+    ::std::stringstream ss;
+    ss << "Mismatch in the number of MPI ranks used to checkpoint. Checkpointing was " <<
+      "done on " << (int)alg_params[2] << " ranks while HiOp currently runs on " <<
+      nranks << " ranks.\n";
+    //throw std::runtime_error(ss.str());
+  }
+
+  //iterate states
+  SidreHelper::copy_vec_from_view(group, "x", *it_curr->get_x());
+  SidreHelper::copy_vec_from_view(group, "d", *it_curr->get_d());
+  SidreHelper::copy_vec_from_view(group, "sxl", *it_curr->get_sxl());
+  SidreHelper::copy_vec_from_view(group, "sxu", *it_curr->get_sxu());
+  SidreHelper::copy_vec_from_view(group, "sdl", *it_curr->get_sdl());
+  SidreHelper::copy_vec_from_view(group, "sdu", *it_curr->get_sdu());
+  SidreHelper::copy_vec_from_view(group, "yc", *it_curr->get_yc());
+  SidreHelper::copy_vec_from_view(group, "yd", *it_curr->get_yd());
+  SidreHelper::copy_vec_from_view(group, "zl", *it_curr->get_zl());
+  SidreHelper::copy_vec_from_view(group, "zu", *it_curr->get_zu());
+  SidreHelper::copy_vec_from_view(group, "vl", *it_curr->get_vl());
+  SidreHelper::copy_vec_from_view(group, "vu", *it_curr->get_vu());
+
+  //state of quasi-Newton Hessian approximation
+
+}
+
+void hiopAlgFilterIPMQuasiNewton::checkpointing_stuff()
+{
+  if(nlp->options->GetString("checkpoint_save")=="no") {
+    return;
+  }
+  int chk_every_N = nlp->options->GetInteger("checkpoint_save_every_N_iter");
+  //check iteration
+  if(iter_num>0 && iter_num % chk_every_N==0) {
+    using ::std::string;
+    // replace "#" in checkpointing file with iteration number
+    string path = nlp->options->GetString("checkpoint_file");
+    auto pos = path.find("#");
+    if(string::npos != pos) {
+      auto s_it_num = ::std::to_string(iter_num);
+      path.replace(pos, 1, s_it_num);
+    }
+
+    nlp->log->printf(hovSummary, "Saving checkpoint at iter %d in '%s'.\n", iter_num, path.c_str());
+    //actual checkpointing via axom::sidre
+    save_state_to_file(path);
+  }
+}
+#endif // HIOP_USE_AXOM
 
 /******************************************************************************************************
  * FULL NEWTON IPM

src/Optimization/hiopAlgFilterIPM.hpp

@@ -67,11 +67,19 @@
 #include "hiopPDPerturbation.hpp"
 #include "hiopFactAcceptor.hpp"
 
+#ifdef HIOP_USE_AXOM
+namespace axom {
+namespace sidre {
+class Group; // forward declaration
+}
+}
+#endif
+
 #include "hiopTimer.hpp"
 
 namespace hiop
 {
-
+   
 class hiopAlgFilterIPMBase {
 public:
   hiopAlgFilterIPMBase(hiopNlpFormulation* nlp_, const bool within_FR = false);
@@ -117,6 +125,8 @@ class hiopAlgFilterIPMBase {
   { 
     return filter.contains(theta, logbar_obj); 
   }
+
+  /// Setter for the primal steplength.
   inline void set_alpha_primal(const double alpha_primal) { _alpha_primal = alpha_primal; }
 
 protected:
@@ -339,8 +349,81 @@ class hiopAlgFilterIPMQuasiNewton : public hiopAlgFilterIPMBase
   virtual ~hiopAlgFilterIPMQuasiNewton();
 
   virtual hiopSolveStatus run();
+
+  // note that checkpointing is only available with a axom-enabled build
+#ifdef HIOP_USE_AXOM
+  /**
+   * @brief Save state of HiOp algorithm to a sidre::Group as a checkpoint.
+   *
+   * @param group a reference to the group where state will be saved to
+   *
+   * @exception std::runtime indicates the group contains a view whose size does not match
+   * the size of the corresponding HiOp algorithm state variable of parameter. 
+   *
+   * @details 
+   * Each state variable of each parameter of HiOp algorithm will be saved in a named 
+   * view within the group. A new view will be created within the group if it does not 
+   * already exist. If it exists, the view must have same number of elements as the 
+   * as the size of the corresponding state variable. This means that this method will
+   * throw an exception if an existing group is reused to save a problem that changed
+   * sizes since the group was created.
+   */
+  virtual void save_state_to_sidre_group(::axom::sidre::Group& group);
+
+  /**
+   * @brief Load state of HiOp algorithm from a sidre::Group checkpoint.
+   *
+   * @param group a pointer to group containing the a prevously saved HiOp algorithm state.
+   *
+   * @exception std::runtime indicates the group does not contain a view expected by this 
+   * method or the view's number of elements mismatches the size of the corresponding HiOp
+   * state. The latter can occur if the file was saved with a different number of MPI ranks.
+   *
+   * @details 
+   * Copies views from the sidre::Group passed as argument to HiOp algorithm's state variables
+   * and parameters. The group should be created by first calling save_state_to_sidre_group 
+   * for a problem/NLP of the same sizes as the problem for which this method is called. 
+   * The method expects views within the group with certain names. If one such view is not 
+   * found or has a number of elements different than the size of the corresponding HiOp state, 
+   * then a std::runtime_error exception is thrown. The latter can occur when the loading 
+   * occurs for a instance of HiOp that is not ran on the same number of MPI ranks used to
+   * save the file.
+   */ 
+  virtual void load_state_from_sidre_group(const ::axom::sidre::Group& group);
+
+  /**
+   * @brief Save the state of the algorithm to the file for checkpointing.
+   *
+   * @param path   the name of the file
+   * @return true if successful, false otherwise
+   * 
+   * @details
+   * Internally, HiOp uses axom::sidre::DataStore and sidre's scalable IO. A detailed
+   * error description is sent to the log if an error or exception is caught.
+   */
+  bool save_state_to_file(const ::std::string& path) noexcept;
+
+  /**
+   * @brief Load the state of the algorithm from checkpoint file.  
+   *
+   * @param path   the name of the file to load from
+   * @return true if successful, false otherwise
+   * 
+   * @details 
+   * The file should contains a axom::sidre::DataStore that was previously saved using 
+   * save_state_to_file(). A detailed error description is sent to the log if an error 
+   * or exception is caught.
+   */
+  bool load_state_from_file(const ::std::string& path) noexcept;
+#endif // HIOP_USE_AXOM
 private:
   virtual void outputIteration(int lsStatus, int lsNum, int use_soc = 0, int use_fr = 0);
+
+#ifdef HIOP_USE_AXOM  
+  ///@brief The options-based logic for saving checkpoint and the call to save_state().
+  void checkpointing_stuff();
+#endif // HIOP_USE_AXOM
+
 private:
   hiopNlpDenseConstraints* nlpdc;
 private: