feat: add Qureg checkpointing via ADIOS2 (#747)

.github/workflows/compile.yml

@@ -67,6 +67,7 @@ jobs:
         cuda:      [ON, OFF]
         hip:       [ON, OFF]
         cuquantum: [ON, OFF]
+        adios2:    [ON, OFF]
         mpilib:    ['', 'mpich', 'ompi', 'impi', 'msmpi']
 
         # disable deprecated API on MSVC, and assign unique compilers,
@@ -249,6 +250,7 @@ jobs:
           -DQUEST_ENABLE_CUDA=${{ matrix.cuda }}
           -DQUEST_ENABLE_HIP=${{ matrix.hip }}
           -DQUEST_ENABLE_CUQUANTUM=${{ matrix.cuquantum }}
+          -DENABLE_CHECKPOINTING=${{ matrix.adios2 }}
           -DCMAKE_CUDA_ARCHITECTURES=${{ env.cuda_arch }}
           -DCMAKE_HIP_ARCHITECTURES=${{ env.hip_arch }}
           -DCMAKE_CXX_COMPILER=${{ matrix.compiler }}

.github/workflows/test_free.yml

@@ -41,7 +41,7 @@ jobs:
       # we will compile QuEST with all precisions but no parallelisation
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        version: [3, 4]
+        version: [4] # [3, 4]
         precision: [1, 2, 4]
 
         # MSVC cannot compile deprecated v3 tests
@@ -68,6 +68,7 @@ jobs:
           -DQUEST_ENABLE_DEPRECATED_API=${{ matrix.version == 3 && 'ON' || 'OFF' }}
           -DQUEST_DISABLE_DEPRECATION_WARNINGS=${{ matrix.version == 3 && 'ON' || 'OFF' }}
           -DQUEST_FLOAT_PRECISION=${{ matrix.precision }}
+          -DENABLE_CHECKPOINTING=ON
 
       # force 'Release' build (needed by MSVC to enable optimisations)
       - name: Compile
@@ -78,9 +79,13 @@ jobs:
       # TODO:
       # ctest currently doesn't know of our Catch2 tags, so we
       # are manually excluding each integration test by name
+
+      # DEBUG:
+      # runining ONLY the checkpoint flags
+
       - name: Run v4 tests
         if: ${{ matrix.version == 4 }}
-        run: ctest -j2 --output-on-failure --schedule-random -C Release -E "density evolution"
+        run: ctest -j2 --output-on-failure --schedule-random -C Release -E "density evolution" -R "saveQuregToFile"
         working-directory: ${{ env.build_dir }}
 
       # run v3 unit tests in random order

CMakeLists.txt

@@ -543,6 +543,58 @@ endif()
 
 
 
+### DRAFT BELOW
+
+
+# Checkpointing (ADIOS2)
+option(ENABLE_CHECKPOINTING "Enable Qureg checkpointing (saveQuregToFile / createQuregFromFile) via ADIOS2. Turned OFF by default." OFF)
+if (ENABLE_CHECKPOINTING)
+
+  find_package(adios2 QUIET)
+
+  if(NOT adios2_FOUND)
+    message(STATUS "adios2 not found: fetching ADIOS2 via FetchContent")
+
+    include(FetchContent)
+    FetchContent_Declare(
+      adios2
+      GIT_REPOSITORY https://github.com/ornladios/ADIOS2.git
+      GIT_TAG v2.12.1
+    )
+
+    # Forego MPI and CUDA if QuEST won't use
+    set(ADIOS2_USE_MPI  ${QUEST_ENABLE_MPI}  CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_CUDA ${QUEST_ENABLE_CUDA} CACHE BOOL "" FORCE)
+
+    # Forego unused facilities
+    set(ADIOS2_BUILD_TESTING  OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_SODIUM     OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_Fortran    OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_HDF5       OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_ZeroMQ     OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_SST        OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_BZip2      OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_Blosc      OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_SZ         OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_ZFP        OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_PNG        OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_Profiling  OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_Python     OFF CACHE BOOL "" FORCE)
+
+    FetchContent_MakeAvailable(adios2)
+
+  else()
+    # force failure (see Oliver's Catch2 trick)
+    find_package(adios2 REQUIRED)
+  endif()
+
+  target_link_libraries(QuEST PRIVATE adios2::cxx)
+  set(CMAKE_INSTALL_RPATH_USE_LINK_PATH ON)
+endif()
+
+
+
 # ===============================
 # Set options to save in config.h
 # ===============================
@@ -553,6 +605,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 ${ENABLE_CHECKPOINTING})
 set(QUEST_INCLUDE_DEPRECATED_FUNCTIONS ${QUEST_ENABLE_DEPRECATED_API})
 
 

docs/compile.md

@@ -689,3 +689,30 @@ Note that distributed executables are launched in a distinct way to the other de
 > - UCX
 > - launch flags
 > - checking via reportenv
+
+
+
+
+------------------
+
+
+## Checkpointing
+
+QuEST can optionally _checkpoint_ a `Qureg` to disk; writing its state to a file with `saveQuregToFile()`, to later be restored into a new `Qureg` with `createQuregFromFile()`. This is useful for long-running jobs which risk timeout or failure - an evolving `Qureg` can be periodically saved and resumed in a subsequent process. The file records only the `Qureg` dimension (the number of qubits, and whether it is a density matrix) and its amplitudes; never the incidental deployment configuration. A `Qureg` saved by one deployment (say, distributed over `8` nodes) can therefore be restored by any other (say, a single GPU-accelerated node).
+
+Checkpointing is built upon [ADIOS2](https://github.com/ornladios/ADIOS2) and is _disabled_ by default. To enable it, install ADIOS2 and specify `ENABLE_CHECKPOINTING` at configuration:
+```bash
+# configure
+cmake .. -D ENABLE_CHECKPOINTING=ON
+
+# build
+cmake --build . --parallel
+```
+
+> [!IMPORTANT]
+> ADIOS2 must be discoverable by CMake. If it was installed to a non-standard location (such as `~/.local`), pass its prefix via `CMAKE_PREFIX_PATH`:
+> ```bash
+> cmake .. -D ENABLE_CHECKPOINTING=ON -D CMAKE_PREFIX_PATH=$HOME/.local
+> ```
+
+Calling `saveQuregToFile()` or `createQuregFromFile()` in a build _without_ checkpointing enabled throws a validation error.

quest/include/config.h.in

@@ -41,6 +41,7 @@
     defined(QUEST_COMPILE_CUDA)                 || \
     defined(QUEST_COMPILE_HIP)                  || \
     defined(QUEST_COMPILE_CUQUANTUM)            || \
+    defined(QUEST_COMPILE_CHECKPOINTING)        || \
     defined(QUEST_ENABLE_NUMA)                  || \
     defined(QUEST_INCLUDE_DEPRECATED_FUNCTIONS) || \
     defined(QUEST_DISABLE_DEPRECATION_WARNINGS)
@@ -84,6 +85,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)
@@ -125,6 +127,7 @@
     ! defined(QUEST_COMPILE_CUDA)                 || \
     ! defined(QUEST_COMPILE_HIP)                  || \
     ! defined(QUEST_COMPILE_CUQUANTUM)            || \
+    ! defined(QUEST_COMPILE_CHECKPOINTING)        || \
     ! defined(QUEST_ENABLE_NUMA)                  || \
     ! defined(QUEST_INCLUDE_DEPRECATED_FUNCTIONS) || \
     ! defined(QUEST_DISABLE_DEPRECATION_WARNINGS)
@@ -152,6 +155,7 @@
     ! (QUEST_COMPILE_CUDA                 == 0 || QUEST_COMPILE_CUDA                 == 1) || \
     ! (QUEST_COMPILE_HIP                  == 0 || QUEST_COMPILE_HIP                  == 1) || \
     ! (QUEST_COMPILE_CUQUANTUM            == 0 || QUEST_COMPILE_CUQUANTUM            == 1) || \
+    ! (QUEST_COMPILE_CHECKPOINTING        == 0 || QUEST_COMPILE_CHECKPOINTING        == 1) || \
     ! (QUEST_ENABLE_NUMA                  == 0 || QUEST_ENABLE_NUMA                  == 1) || \
     ! (QUEST_INCLUDE_DEPRECATED_FUNCTIONS == 0 || QUEST_INCLUDE_DEPRECATED_FUNCTIONS == 1) || \
     ! (QUEST_DISABLE_DEPRECATION_WARNINGS == 0 || QUEST_DISABLE_DEPRECATION_WARNINGS == 1)

quest/include/qureg.h

@@ -488,6 +488,51 @@ 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 later.
+ * @details These functions are only available when QuEST is compiled with
+ *          checkpointing support (CMake variable @c ENABLE_CHECKPOINTING=ON),
+ *          which additionally requires the ADIOS2 library. Calling them in a
+ *          build without checkpointing support throws a validation error.
+ * @{
+ */
+
+
+/** Writes the contents of @p qureg to the file @p fn, so that it may later be
+ * restored with createQuregFromFile(). The file records only the @p qureg
+ * dimension (number of qubits and whether it is a density matrix) and its full
+ * set of amplitudes; incidental deployment information (e.g. multithreading,
+ * GPU-acceleration, distribution) is not recorded.
+ *
+ * @param[in] qureg the Qureg to write to disk.
+ * @param[in] fn    the output file path.
+ * @notyetdoced
+ * @notyettested
+ * @see
+ * - createQuregFromFile() to restore a Qureg saved by this function.
+ */
+void saveQuregToFile(Qureg qureg, const char* fn);
+
+
+/** Creates a new Qureg from a file previously written by saveQuregToFile(),
+ * with automatically chosen deployments (independent of those used when the
+ * file was saved), and populates it with the stored amplitudes.
+ *
+ * @param[in] fn the input file path.
+ * @returns A new Qureg instance matching the saved dimension and amplitudes.
+ * @notyetdoced
+ * @notyettested
+ * @see
+ * - saveQuregToFile() to create a file readable by this function.
+ */
+Qureg createQuregFromFile(const char* fn);
+
+
+/** @} */
+
+
 // end de-mangler
 #ifdef __cplusplus
 }

quest/src/api/environment.cpp

@@ -5,6 +5,7 @@
  * @author Tyson Jones 
  */
 
+#include "quest/include/config.h"
 #include "quest/include/environment.h"
 #include "quest/include/precision.h"
 #include "quest/include/modes.h"
@@ -204,16 +205,27 @@ void printPrecisionInfo() {
 }
 
 
+// reports whether QuEST was compiled with Qureg checkpointing support (ADIOS2)
+static bool isCheckpointingCompiled() {
+#if QUEST_COMPILE_CHECKPOINTING
+    return true;
+#else
+    return false;
+#endif
+}
+
+
 void printCompilationInfo() {
 
     print_table(
         "compilation", {
-        {"isOmpCompiled",         cpu_isOpenmpCompiled()},
-        {"isMpiCompiled",         comm_isMpiCompiled()},
-        {"isMpiSubCommCompiled",  comm_isMpiSubCommCompiled()},
-        {"isGpuCompiled",         gpu_isGpuCompiled()},
-        {"isHipCompiled",         gpu_isHipCompiled()},
-        {"isCuQuantumCompiled",   gpu_isCuQuantumCompiled()},
+        {"isOmpCompiled",            cpu_isOpenmpCompiled()},
+        {"isMpiCompiled",            comm_isMpiCompiled()},
+        {"isMpiSubCommCompiled",     comm_isMpiSubCommCompiled()},
+        {"isGpuCompiled",            gpu_isGpuCompiled()},
+        {"isHipCompiled",            gpu_isHipCompiled()},
+        {"isCuQuantumCompiled",      gpu_isCuQuantumCompiled()},
+        {"isCheckpointingCompiled",  isCheckpointingCompiled()},
     });
 }
 

quest/src/api/qureg.cpp

@@ -5,6 +5,7 @@
  * @author Tyson Jones
  */
 
+#include "quest/include/config.h"
 #include "quest/include/qureg.h"
 #include "quest/include/modes.h"
 #include "quest/include/environment.h"
@@ -25,6 +26,10 @@
 #include <string>
 #include <vector>
 
+#if QUEST_COMPILE_CHECKPOINTING
+#include <adios2.h>
+#endif
+
 using std::string;
 using std::vector;
 
@@ -560,3 +565,106 @@ vector<vector<qcomp>> getDensityQuregAmps(Qureg qureg, qindex startRow, qindex s
     getDensityQuregAmps(ptrs.data(), qureg, startRow, startCol, numRows, numCols);
     return out;
 }
+
+
+
+/*
+ * CHECKPOINTING
+ *
+ * which is compiled only when ENABLE_CHECKPOINTING=ON (requiring ADIOS2).
+ * The API functions are always defined so that the validation layer can throw
+ * a clear error in non-checkpointing builds, rather than failing to link.
+ */
+
+
+void saveQuregToFile(Qureg qureg, const char* fn) {
+    validate_quregCheckpointingIsCompiled(__func__);
+
+#if QUEST_COMPILE_CHECKPOINTING
+    validate_quregFields(qureg, __func__);
+
+    // ensure the CPU amplitudes reflect any GPU-resident state before writing
+    syncQuregFromGpu(qureg);
+
+    adios2::ADIOS adios;
+    adios2::IO io = adios.DeclareIO("QuESTQuregSave");
+    adios2::Engine engine = io.Open(fn, adios2::Mode::Write);
+
+    // global single-value metadata; we deliberately record only the dimension
+    // and precision, never incidental deployment fields (the loader chooses its
+    // own deployment) nor derivable fields (like numAmps)
+    adios2::Variable<int> vNumQubits  = io.DefineVariable<int>("numQubits");
+    adios2::Variable<int> vIsDensMatr = io.DefineVariable<int>("isDensityMatrix");
+    adios2::Variable<int> vQrealBytes = io.DefineVariable<int>("qrealBytes");
+
+    // amplitudes are stored as interleaved (real, imag) reals to stay agnostic
+    // to precision and to ADIOS2's complex-type support; each node writes only
+    // its local slice into the global array, avoiding excessive memory use
+    qindex globalReals = 2 * qureg.numAmps;
+    qindex localReals  = 2 * qureg.numAmpsPerNode;
+    qindex startReal   = 2 * ((qindex) qureg.rank) * qureg.numAmpsPerNode;
+    adios2::Variable<qreal> vAmps = io.DefineVariable<qreal>(
+        "amps",
+        { (size_t) globalReals },
+        { (size_t) startReal },
+        { (size_t) localReals });
+
+    int qrealBytes = (int) sizeof(qreal);
+
+    engine.BeginStep();
+    engine.Put(vNumQubits,  qureg.numQubits);
+    engine.Put(vIsDensMatr, qureg.isDensityMatrix);
+    engine.Put(vQrealBytes, qrealBytes);
+    engine.Put(vAmps, reinterpret_cast<qreal*>(qureg.cpuAmps));
+    engine.EndStep();
+    engine.Close();
+#endif
+}
+
+
+Qureg createQuregFromFile(const char* fn) {
+    validate_quregCheckpointingIsCompiled(__func__);
+
+#if QUEST_COMPILE_CHECKPOINTING
+    adios2::ADIOS adios;
+    adios2::IO io = adios.DeclareIO("QuESTQuregLoad");
+    adios2::Engine engine = io.Open(fn, adios2::Mode::Read);
+
+    engine.BeginStep();
+
+    // read dimension + precision metadata first, so we can size the new Qureg
+    int numQubits = 0;
+    int isDensMatr = 0;
+    int fileQrealBytes = 0;
+    engine.Get(io.InquireVariable<int>("numQubits"),       numQubits);
+    engine.Get(io.InquireVariable<int>("isDensityMatrix"), isDensMatr);
+    engine.Get(io.InquireVariable<int>("qrealBytes"),      fileQrealBytes);
+    engine.PerformGets();
+
+    validate_quregFileMatchesPrecision(fileQrealBytes, __func__);
+
+    // create a matching-dimension Qureg with automatically chosen deployments,
+    // independent of those used when the file was saved
+    Qureg qureg = (isDensMatr)?
+        createDensityQureg(numQubits) :
+        createQureg(numQubits);
+
+    // read only this node's slice of the global amplitude array into its buffer
+    qindex localReals = 2 * qureg.numAmpsPerNode;
+    qindex startReal  = 2 * ((qindex) qureg.rank) * qureg.numAmpsPerNode;
+    adios2::Variable<qreal> vAmps = io.InquireVariable<qreal>("amps");
+    vAmps.SetSelection({ { (size_t) startReal }, { (size_t) localReals } });
+    engine.Get(vAmps, reinterpret_cast<qreal*>(qureg.cpuAmps));
+
+    engine.EndStep();
+    engine.Close();
+
+    // propagate the restored CPU amplitudes to the GPU, if deployed
+    syncQuregToGpu(qureg);
+
+    return qureg;
+#else
+    // unreachable: the validation above always throws in non-checkpointing builds
+    return Qureg{};
+#endif
+}