QuEST #747 — Qureg checkpointing (save/load to file via ADIOS2) [unitaryHACK]

CMakeLists.txt

@@ -223,6 +223,15 @@ if (QUEST_INSTALL_BINARIES)
 endif()
 
 
+# Checkpointing (issue #747): save/load a Qureg to file via ADIOS2
+option(
+  QUEST_ENABLE_CHECKPOINTING
+  "Whether QuEST will be built with Qureg checkpointing (saveQuregToFile / createQuregFromFile) via ADIOS2. Turned OFF by default."
+  OFF
+)
+message(STATUS "Checkpointing is turned ${QUEST_ENABLE_CHECKPOINTING}. Set QUEST_ENABLE_CHECKPOINTING to modify.")
+
+
 
 # ============================
 # Validate options
@@ -486,6 +495,18 @@ if (QUEST_ENABLE_MPI)
 endif()
 
 
+# Checkpointing (issue #747): link ADIOS2 (its MPI-enabled component when distributed)
+if (QUEST_ENABLE_CHECKPOINTING)
+  find_package(ADIOS2 REQUIRED)
+
+  if (QUEST_ENABLE_MPI)
+    target_link_libraries(QuEST PRIVATE adios2::cxx_mpi)
+  else()
+    target_link_libraries(QuEST PRIVATE adios2::cxx)
+  endif()
+endif()
+
+
 # CUDA
 if (QUEST_ENABLE_CUDA)
 
@@ -553,6 +574,7 @@ set(QUEST_COMPILE_OMP ${QUEST_ENABLE_OMP})
 set(QUEST_COMPILE_MPI ${QUEST_ENABLE_MPI})
 set(QUEST_COMPILE_SUBCOMM ${QUEST_ENABLE_SUBCOMM})
 set(QUEST_COMPILE_CUQUANTUM ${QUEST_ENABLE_CUQUANTUM})
+set(QUEST_COMPILE_CHECKPOINTING ${QUEST_ENABLE_CHECKPOINTING})
 set(QUEST_INCLUDE_DEPRECATED_FUNCTIONS ${QUEST_ENABLE_DEPRECATED_API})
 
 

quest/include/config.h.in

@@ -84,6 +84,7 @@
 #cmakedefine01 QUEST_COMPILE_CUDA
 #cmakedefine01 QUEST_COMPILE_CUQUANTUM
 #cmakedefine01 QUEST_COMPILE_HIP
+#cmakedefine01 QUEST_COMPILE_CHECKPOINTING
 
 
 // crucial to QuEST source (informs optional NUMA usage)

quest/include/qureg.h

@@ -488,6 +488,29 @@ void getDensityQuregAmps(qcomp** outAmps, Qureg qureg, qindex startRow, qindex s
 /** @} */
 
 
+
+/**
+ * @defgroup qureg_checkpoint Checkpointing
+ * @brief Functions for saving a Qureg to file and restoring it (issue #747).
+ * @details Available only when QuEST is compiled with -DQUEST_ENABLE_CHECKPOINTING=ON,
+ *          which links ADIOS2. The saved file records the qubit count, statevector/
+ *          density-matrix type and the full amplitudes, but no deployment details, so
+ *          a Qureg can be resumed under a different GPU/distribution configuration.
+ * @{
+ */
+
+
+/// @notyetdoced
+void saveQuregToFile(Qureg qureg, const char* fn);
+
+
+/// @notyetdoced
+Qureg createQuregFromFile(const char* fn);
+
+
+/** @} */
+
+
 // end de-mangler
 #ifdef __cplusplus
 }

quest/src/api/qureg.cpp

@@ -25,6 +25,12 @@
 #include <string>
 #include <vector>
 
+// issue #747: optional Qureg checkpointing via ADIOS2, enabled at compile time
+// with -DQUEST_ENABLE_CHECKPOINTING=ON (which sets QUEST_COMPILE_CHECKPOINTING)
+#if QUEST_COMPILE_CHECKPOINTING
+    #include <adios2.h>
+#endif
+
 using std::string;
 using std::vector;
 
@@ -474,6 +480,114 @@ void getDensityQuregAmps(qcomp** outAmps, Qureg qureg, qindex startRow, qindex s
 }
 
 
+/*
+ * CHECKPOINTING (issue #747)
+ *
+ * saveQuregToFile() and createQuregFromFile() persist a Qureg to disk and
+ * restore it, via ADIOS2. Only the essential, non-derivable, deployment-
+ * independent state is stored: numQubits, isDensityMatrix, and the global
+ * amplitude array. Deployment (GPU / distribution / threads) is NOT stored -
+ * the restored Qureg auto-deploys, and ADIOS2's global-array model lets a file
+ * written under one distribution be read under another.
+ *
+ * Amplitudes are stored as raw bytes (with sizeof(qcomp) recorded for a
+ * load-time compatibility check), so all three qcomp precisions - including
+ * long double, which ADIOS2 cannot represent as a native type - are supported
+ * uniformly. The only extra memory is one GPU->host copy of the already-
+ * resident local amplitudes.
+ *
+ * Compiled only when QUEST_COMPILE_CHECKPOINTING; otherwise the functions
+ * report a user error (and are no-ops thereafter).
+ */
+
+void saveQuregToFile(Qureg qureg, const char* fn) {
+    validate_quregFields(qureg, __func__);
+    validate_quregCheckpointingIsCompiled(__func__);
+
+#if QUEST_COMPILE_CHECKPOINTING
+    // ensure host amps reflect the (possibly GPU-resident) state
+    syncQuregFromGpu(qureg);
+
+    // this node's contiguous slice of the global amplitude array, in bytes
+    qindex ampBytes    = sizeof(qcomp);
+    qindex localBytes  = qureg.numAmpsPerNode * ampBytes;
+    qindex totalBytes  = qureg.numAmps        * ampBytes;
+    qindex offsetBytes = util_getGlobalIndexOfFirstLocalAmp(qureg) * ampBytes;
+
+    #if QUEST_COMPILE_MPI
+        adios2::ADIOS adios(comm_getMpiComm());
+    #else
+        adios2::ADIOS adios;
+    #endif
+    adios2::IO io = adios.DeclareIO("QuESTCheckpointWrite");
+
+    auto vNumQubits  = io.DefineVariable<int>("numQubits");
+    auto vIsDensMatr = io.DefineVariable<int>("isDensityMatrix");
+    auto vAmpBytes   = io.DefineVariable<int>("ampSizeBytes");
+    auto vAmps = io.DefineVariable<int8_t>("amplitudes",
+        {static_cast<size_t>(totalBytes)},   // global shape
+        {static_cast<size_t>(offsetBytes)},  // this node's start
+        {static_cast<size_t>(localBytes)});  // this node's count
+
+    adios2::Engine engine = io.Open(fn, adios2::Mode::Write);
+    engine.BeginStep();
+
+    // global scalars are written once, by the root node
+    if (qureg.rank == 0) {
+        engine.Put(vNumQubits,  qureg.numQubits);
+        engine.Put(vIsDensMatr, qureg.isDensityMatrix);
+        engine.Put(vAmpBytes,   static_cast<int>(ampBytes));
+    }
+    engine.Put(vAmps, reinterpret_cast<int8_t*>(qureg.cpuAmps));
+
+    engine.EndStep();
+    engine.Close();
+#endif
+}
+
+
+Qureg createQuregFromFile(const char* fn) {
+    validate_envIsInit(__func__);
+    validate_quregCheckpointingIsCompiled(__func__);
+
+#if QUEST_COMPILE_CHECKPOINTING
+    #if QUEST_COMPILE_MPI
+        adios2::ADIOS adios(comm_getMpiComm());
+    #else
+        adios2::ADIOS adios;
+    #endif
+    adios2::IO io = adios.DeclareIO("QuESTCheckpointRead");
+    adios2::Engine engine = io.Open(fn, adios2::Mode::ReadRandomAccess);
+
+    // read the global scalars describing the saved Qureg
+    int numQubits = 0, isDensMatr = 0, fileAmpBytes = 0;
+    engine.Get(io.InquireVariable<int>("numQubits"),       numQubits,    adios2::Mode::Sync);
+    engine.Get(io.InquireVariable<int>("isDensityMatrix"), isDensMatr,   adios2::Mode::Sync);
+    engine.Get(io.InquireVariable<int>("ampSizeBytes"),    fileAmpBytes, adios2::Mode::Sync);
+    validate_checkpointFileMatchesPrecision(fileAmpBytes, static_cast<int>(sizeof(qcomp)), __func__);
+
+    // create a Qureg of the saved dimension, with auto-chosen deployment
+    Qureg qureg = (isDensMatr)?
+        createDensityQureg(numQubits) : createQureg(numQubits);
+
+    // read this node's slice of the global amplitude array into its host buffer
+    qindex localBytes  = qureg.numAmpsPerNode * static_cast<qindex>(sizeof(qcomp));
+    qindex offsetBytes = util_getGlobalIndexOfFirstLocalAmp(qureg) * static_cast<qindex>(sizeof(qcomp));
+    adios2::Variable<int8_t> vAmps = io.InquireVariable<int8_t>("amplitudes");
+    vAmps.SetSelection({{static_cast<size_t>(offsetBytes)}, {static_cast<size_t>(localBytes)}});
+    engine.Get(vAmps, reinterpret_cast<int8_t*>(qureg.cpuAmps), adios2::Mode::Sync);
+    engine.Close();
+
+    // propagate the loaded host amps to the GPU, if accelerated
+    syncQuregToGpu(qureg);
+    return qureg;
+#else
+    // unreachable - the validation above aborts - but required for compilation
+    return Qureg{};
+#endif
+}
+
+
 
 // end de-mangler
 }

quest/src/core/validation.cpp

@@ -131,6 +131,17 @@ namespace report {
         "The QuEST environment is not initialised. Please first call initQuESTEnv() or initCustomQuESTEnv().";
 
 
+    /*
+     * CHECKPOINTING (issue #747)
+     */
+
+    string CHECKPOINTING_NOT_COMPILED =
+        "This function requires QuEST to be compiled with checkpointing enabled. Please recompile with the CMake flag -DQUEST_ENABLE_CHECKPOINTING=ON, which links the ADIOS2 library.";
+
+    string CHECKPOINT_FILE_PRECISION_MISMATCH =
+        "The checkpoint file stores ${FILEBYTES}-byte amplitudes, incompatible with this build's ${BUILDBYTES}-byte amplitudes. A checkpoint must be loaded by a QuEST build of the same numerical precision (FLOAT_PRECISION).";
+
+
     /*
      * DEBUG UTILITIES
      */
@@ -1600,6 +1611,32 @@ void validate_envIsInit(const char* caller) {
 
 
 
+/*
+ * CHECKPOINTING (issue #747)
+ */
+
+void validate_quregCheckpointingIsCompiled(const char* caller) {
+
+    if (!global_isValidationEnabled)
+        return;
+
+    assertThat((bool) QUEST_COMPILE_CHECKPOINTING, report::CHECKPOINTING_NOT_COMPILED, caller);
+}
+
+void validate_checkpointFileMatchesPrecision(int fileAmpBytes, int buildAmpBytes, const char* caller) {
+
+    if (!global_isValidationEnabled)
+        return;
+
+    tokenSubs vars = {
+        {"${FILEBYTES}",  fileAmpBytes},
+        {"${BUILDBYTES}", buildAmpBytes}};
+
+    assertThat(fileAmpBytes == buildAmpBytes, report::CHECKPOINT_FILE_PRECISION_MISMATCH, vars, caller);
+}
+
+
+
 /*
  * DEBUG UTILITIES
  */

quest/src/core/validation.hpp

@@ -93,6 +93,16 @@ void validate_envIsInit(const char* caller);
 
 
 
+/*
+ * CHECKPOINTING (issue #747)
+ */
+
+void validate_quregCheckpointingIsCompiled(const char* caller);
+
+void validate_checkpointFileMatchesPrecision(int fileAmpBytes, int buildAmpBytes, const char* caller);
+
+
+
 /*
  * DEBUG UTILITIES
  */