From a14a445007461f5186eefebe92b8c5fb56765f66 Mon Sep 17 00:00:00 2001
From: Cosmin G Petra <petra1@llnl.gov>
Date: Wed, 28 Aug 2024 15:48:41 -0700
Subject: [PATCH] added user options for checkpointing

---
 src/Optimization/hiopAlgFilterIPM.cpp | 20 ++++++++++++++
 src/Optimization/hiopAlgFilterIPM.hpp |  5 ++++
 src/Utils/hiopOptions.cpp             | 40 +++++++++++++++++++++++++--
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/src/Optimization/hiopAlgFilterIPM.cpp b/src/Optimization/hiopAlgFilterIPM.cpp
index a179341a..1f6b52a7 100644
--- a/src/Optimization/hiopAlgFilterIPM.cpp
+++ b/src/Optimization/hiopAlgFilterIPM.cpp
@@ -1106,6 +1106,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
      ************************************************/
@@ -1553,6 +1558,21 @@ void hiopAlgFilterIPMQuasiNewton::load_state_from_data_store(const sidre::DataSt
   const double* x = it_curr->get_x()->local_data_host();
   //destination = nlp_group->createViewAndAllocate("x", ::axom::sidre::DOUBLE_ID, size);
 }
+
+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
+  ::std::string path = nlp->options->GetString("checkpoint_file");
+
+
+  
+  // replace #
+  
+}
 #endif // HIOP_USE_AXOM
 
 /******************************************************************************************************
diff --git a/src/Optimization/hiopAlgFilterIPM.hpp b/src/Optimization/hiopAlgFilterIPM.hpp
index a44208b6..d155ffed 100644
--- a/src/Optimization/hiopAlgFilterIPM.hpp
+++ b/src/Optimization/hiopAlgFilterIPM.hpp
@@ -379,6 +379,11 @@ class hiopAlgFilterIPMQuasiNewton : public hiopAlgFilterIPMBase
 #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  
 private:
   hiopNlpDenseConstraints* nlpdc;
 private:
diff --git a/src/Utils/hiopOptions.cpp b/src/Utils/hiopOptions.cpp
index f2d1cbb7..0f5d7eef 100644
--- a/src/Utils/hiopOptions.cpp
+++ b/src/Utils/hiopOptions.cpp
@@ -1290,9 +1290,25 @@ void hiopOptionsNLP::register_options()
   }
 
   // checkpointing and restarting
-  
-}
+  //  - currently only for IPM Quasi-Newton solver
+  //  - only available with HIOP_USE_AXOM
+  {
+    vector<string> range = {"yes", "no"};
+    constexpr char msgcs[] = "Save state of NLP solver to file specified by 'checkpoint_file'.";
+    register_str_option("checkpoint_save", range[1], range, msgcs);
+    
+    constexpr char msgcsN[] = "Iteration frequency of saving checkpoints to disk.";
+    register_int_option("checkpoint_save_every_N_iter", 10, 1, 1e+6, msgcsN);
+
+    constexpr char msgcf[] = "Path to checkpoint file. If present character '#' will be replaced "
+      "with iteration number.";
+    register_str_option("checkpoint_file", "hiop_state_#.chk", msgcf);
 
+    constexpr char msgclos[] = "On (re)start the NLP solver will load checkpoint file "
+      "specified by checkpoint_file'.";
+    register_str_option("checkpoint_load_on_start", range[1], range, msgclos);
+  }
+}
 void hiopOptionsNLP::ensure_consistence()
 {
   //check that the values of different options are consistent
@@ -1553,6 +1569,26 @@ void hiopOptionsNLP::ensure_consistence()
     }
     set_val("moving_lim_rel", 0.);
   }
+
+#ifndef HIOP_USE_AXOM
+  const vector<string> chkpnt_opts = {"checkpoint_save",
+                                      "checkpoint_save_every_N_iter",
+                                      "checkpoint_file",
+                                      "checkpoint_load_on_start"};
+  for(string opt : chkpnt_opts) {
+    if(is_user_defined(opt.c_str())) {
+      log_printf(hovWarning,
+                 "Checkpointing not available since HiOp was not built with AXOM. All checkpointing options "
+                 "are ignored.\n");
+      //reset them to as not being user defined to avoid triggering the message. 
+      for(auto opt2 : chkpnt_opts) {
+        mOptions_[opt2]->specifiedInFile = false;
+        mOptions_[opt2]->specifiedAtRuntime = false;
+      }
+      break;
+    }
+  }
+#endif  
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////