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,11 +66,21 @@
 
 #include "hiopCppStdUtils.hpp"
 
+#ifdef HIOP_USE_AXOM
+#include <axom/sidre/core/DataStore.hpp>
+#include <axom/sidre/core/Group.hpp>
+#include <axom/sidre/core/View.hpp>
+#include <axom/sidre/spio/IOManager.hpp>
+using namespace axom;
+#endif
+
 #include <cmath>
 #include <cstring>
 #include <cassert>
 #include <stdio.h>
 #include <ctype.h>
+#include <exception>
+#include <sstream>
 
 namespace hiop
 {
@@ -976,8 +986,37 @@ 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;
+  } 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;     
+    }
+
+    solver_status_ = NlpSolve_SolveNotCalled;
+  }
 
   //update log bar
   logbar->updateWithNlpInfo(*it_curr, _mu, _f_nlp, *_c, *_d, *_grad_f, *_Jac_c, *_Jac_d);
@@ -1095,6 +1134,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 +1529,170 @@ void hiopAlgFilterIPMQuasiNewton::outputIteration(int lsStatus, int lsNum, int u
   }
 }
 
+#ifdef HIOP_USE_AXOM
+
+bool hiopAlgFilterIPMQuasiNewton::save_state_to_file(const ::std::string& path) noexcept
+{
+  auto success = true;
+  sidre::DataStore* ds = nullptr;
+  try {
+    ds = new sidre::DataStore();
+    this->save_state_to_data_store(ds);
+
+    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.c_str(), sidre::Group::getDefaultIOProtocol());
+  } catch(const std::exception& exp) {
+    nlp->log->printf(hovError, "Error in saving checkpoint to file '%s'\n", path.c_str());
+    nlp->log->printf(hovError,  "  Addtl info: %s\n", exp.what());
+    success = false;    
+  }  
+  delete ds;
+  return success;
+}
+
+bool hiopAlgFilterIPMQuasiNewton::load_state_from_file(const ::std::string& path) noexcept
+{
+  sidre::DataStore* ds = nullptr;
+  auto success = true;
+  try {
+    ds = new sidre::DataStore();
+    sidre::IOManager reader(this->get_nlp()->get_comm());
+    reader.read(ds->getRoot(), path, false);
+
+    this->load_state_from_data_store(ds);
+
+  } 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());
+    success = false;
+  }
+  delete ds;
+  return success;
+}
+
+void hiopAlgFilterIPMQuasiNewton::
+copy_vec_to_new_view(const ::std::string& name, const hiopVector* vec, sidre::Group* nlp_group)
+{
+  using IndType = sidre::IndexType;
+  const IndType size = vec->get_local_size();
+  sidre::View* dest = nlp_group->createViewAndAllocate(name, sidre::DOUBLE_ID, size);
+
+  const auto stride(dest->getStride());
+  double *const dest_ptr(dest->getArray());
+  const double* arr = vec->local_data_host_const();
+  if(1==stride) {
+    ::std::copy(arr, arr+size, dest_ptr);
+  } else {
+    for(IndType i=0; i<size; ++i) {
+      dest_ptr[i*stride] = arr[i];
+    }
+  }  
+}
+void hiopAlgFilterIPMQuasiNewton::
+copy_vec_from_view(const ::std::string& name, hiopVector* vec, const sidre::Group* nlp_group)
+{
+  const sidre::View* view_const = nlp_group->getView(name);
+  if(!view_const) {
+    ::std::stringstream ss;
+    ss << "Could not find view '" << name << "in sidre::Group.\n";
+    throw ::std::runtime_error(ss.str());
+  }
+  // const_cast becase View does not have a const getArray()
+  sidre::View* view = const_cast<sidre::View*>(view_const);  
+  if(view) {
+    const hiop::size_type size = vec->get_local_size();
+    if(view->getNumElements() != size) {
+      ::std::stringstream ss;
+      ss << "Size mismatch for state/view '" << name << "' between hiop and sidre::View: " <<
+        " HiOp state is " << size << " doubles, while the view has " << view->getNumElements() <<
+        " double elements.\n";
+      throw ::std::runtime_error(ss.str());
+    }
+    double* arr_dest = vec->local_data_host();
+    const double *const arr_src = view->getArray();
+    const auto stride(view->getStride());
+    if(1==stride) {
+      ::std::copy(arr_src, arr_src+size, arr_dest);
+    } else {
+      for(hiop::index_type i=0; i<size; ++i) {
+        arr_dest[i] = arr_src[i*stride];
+      }
+    }    
+  } else {
+    assert(false && "const cast failed for Sidre::View");
+  }
+}
+
+void hiopAlgFilterIPMQuasiNewton::save_state_to_data_store(::axom::sidre::DataStore* ds)
+{
+  using IndType = sidre::IndexType;
+  sidre::Group* nlp_group = ds->getRoot()->createGroup("hiop solver");
+
+  //create views for each member that needs to be saved
+  copy_vec_to_new_view("x", it_curr->get_x(), nlp_group);
+  copy_vec_to_new_view("d", it_curr->get_d(), nlp_group);
+  copy_vec_to_new_view("sxl", it_curr->get_sxl(), nlp_group);
+  copy_vec_to_new_view("sxu", it_curr->get_sxu(), nlp_group);
+  copy_vec_to_new_view("sdl", it_curr->get_sdl(), nlp_group);
+  copy_vec_to_new_view("sdu", it_curr->get_sdu(), nlp_group);
+  copy_vec_to_new_view("yc", it_curr->get_yc(), nlp_group);
+  copy_vec_to_new_view("zl", it_curr->get_zl(), nlp_group);
+  copy_vec_to_new_view("zu", it_curr->get_zu(), nlp_group);
+  copy_vec_to_new_view("vl", it_curr->get_vl(), nlp_group);
+  copy_vec_to_new_view("vu", it_curr->get_vu(), nlp_group);
+
+  //quasi-Newton Hessian approximation 
+
+  //algorithmic parameters for this state
+  //mu, iteration number
+  const double alg_params[] = {_mu};
+
+}
+
+void hiopAlgFilterIPMQuasiNewton::load_state_from_data_store(const sidre::DataStore* ds)
+{
+  const sidre::Group* nlp_group = ds->getRoot()->getGroup("hiop solver");
+
+  copy_vec_from_view("x", it_curr->get_x(), nlp_group);
+  copy_vec_from_view("d", it_curr->get_d(), nlp_group);
+  copy_vec_from_view("sxl", it_curr->get_sxl(), nlp_group);
+  copy_vec_from_view("sxu", it_curr->get_sxu(), nlp_group);
+  copy_vec_from_view("sdl", it_curr->get_sdl(), nlp_group);
+  copy_vec_from_view("sdu", it_curr->get_sdu(), nlp_group);
+  copy_vec_from_view("yc", it_curr->get_yc(), nlp_group);
+  copy_vec_from_view("zl", it_curr->get_zl(), nlp_group);
+  copy_vec_from_view("zu", it_curr->get_zu(), nlp_group);
+  copy_vec_from_view("vl", it_curr->get_vl(), nlp_group);
+  copy_vec_from_view("vu", it_curr->get_vu(), nlp_group);
+
+
+}
+
+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,15 @@
 #include "hiopPDPerturbation.hpp"
 #include "hiopFactAcceptor.hpp"
 
+#ifdef HIOP_USE_AXOM
+#include <axom/sidre/core/DataStore.hpp>
+#endif
+
 #include "hiopTimer.hpp"
 
 namespace hiop
 {
-
+   
 class hiopAlgFilterIPMBase {
 public:
   hiopAlgFilterIPMBase(hiopNlpFormulation* nlp_, const bool within_FR = false);
@@ -117,6 +121,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 +345,74 @@ 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 data store.
+   * @param data_store a pointer to DataStore
+   *
+   * @details 
+   * A new sidre::group "hiop solver" is created within data_store. Then a sidre::View 
+   * is created within the group for each of the algorithm's states and the states are
+   * copied in the corresponding views. 
+   */
+  virtual void save_state_to_data_store(::axom::sidre::DataStore* data_store);
+
+  /**
+   * @brief Load state of HiOp algorithm from a sidre data store.
+   * @param data_store a pointer to DataStore
+   *
+   * @details 
+   * Copies views from the data store sidre::group (named "hiop solver") to HiOp algorithm's
+   * state variables. The group should be created by first calling save_state_to_data_store 
+   * for a problem/NLP of the same sizes as the problem for which load_state_from_data_store
+   * is called. This ensures the views have the names and sizes expected by this method. 
+   * Otherwise a std::runtime_error exception is thrown.
+   */ 
+  virtual void load_state_from_data_store(const ::axom::sidre::DataStore* data_store);
+
+  /**
+   * @brief Save the state of the algorithm to the file
+   * @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 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();
+
+  /**
+   * @brief Copy HiOp vector to a (new) axom::sidre::View.
+   *
+   * @details A new view is created/allocated within the sidre group.
+   */
+  void copy_vec_to_new_view(const ::std::string& name, const hiopVector* vec, ::axom::sidre::Group* nlp_group);
+
+  /// Copy content of the named sidre view from nlp_group into HiOp Vector.
+  void copy_vec_from_view(const ::std::string& name, hiopVector* vec, const axom::sidre::Group* nlp_group);
+#endif // HIOP_USE_AXOM
+
 private:
   hiopNlpDenseConstraints* nlpdc;
 private: