Added Qureg checkpointing via ADIOS2 (#780)

Added new functions saveQuregToFile() and createQuregFromFile(), which are optionally enabled through ADIOS2 with new CMake option QUEST_ENABLE_ADIOS2.

Ashmit added the core functionality, doc and prototyped tests. Tyson added input validation, bug fixes, and auto ADIOS2 downloading.

---------

Co-authored-by: Tyson Jones <tyson.jones.input@gmail.com>

Author: ashmitjsg

.github/workflows/compile.yml

@@ -47,6 +47,7 @@ jobs:
       ${{ matrix.cuda      == 'ON' && 'CUDA' || '' }}
       ${{ matrix.hip       == 'ON' && 'HIP'  || '' }}
       ${{ matrix.cuquantum == 'ON' && 'CUQ'  || '' }}
+      ${{ matrix.adios2    == 'ON' && 'CKPT' || '' }}
 
     runs-on: ${{ matrix.os }}
 
@@ -67,6 +68,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,14 +251,16 @@ jobs:
           -DQUEST_ENABLE_CUDA=${{ matrix.cuda }}
           -DQUEST_ENABLE_HIP=${{ matrix.hip }}
           -DQUEST_ENABLE_CUQUANTUM=${{ matrix.cuquantum }}
+          -DQUEST_ENABLE_ADIOS2=${{ matrix.adios2 }}
           -DCMAKE_CUDA_ARCHITECTURES=${{ env.cuda_arch }}
           -DCMAKE_HIP_ARCHITECTURES=${{ env.hip_arch }}
           -DCMAKE_CXX_COMPILER=${{ matrix.compiler }}
           -DCMAKE_CXX_FLAGS=${{ matrix.mpi == 'ON' && matrix.cuda == 'ON' && '-fno-lto' || '' }}
 
-      # force 'Release' build (needed by MSVC to enable optimisations)
+      # force 'Release' build (needed by MSVC to enable optimisations),
+      # and force serial compilation to avoid ADIOS2 OOM error
       - name: Compile
-        run: cmake --build ${{ env.build_dir }} --config Release --parallel
+        run: cmake --build ${{ env.build_dir }} --config Release --parallel 1
 
       # run all compiled isolated examples to test for link-time errors,
       # continuing if any fail (since some deliberately fail)

.github/workflows/test_free.yml

@@ -59,7 +59,7 @@ jobs:
       - name: Get QuEST
         uses: actions/checkout@main
 
-      # compile serial unit tests, optionally include deprecated test
+      # compile serial unit tests, optionally include deprecated test, always including ADIOS2
       - name: Configure CMake
         run: >
           cmake -B ${{ env.build_dir }}
@@ -68,16 +68,18 @@ jobs:
           -DQUEST_ENABLE_DEPRECATED_API=${{ matrix.version == 3 && 'ON' || 'OFF' }}
           -DQUEST_DISABLE_DEPRECATION_WARNINGS=${{ matrix.version == 3 && 'ON' || 'OFF' }}
           -DQUEST_FLOAT_PRECISION=${{ matrix.precision }}
+          -DQUEST_ENABLE_ADIOS2=ON
 
-      # force 'Release' build (needed by MSVC to enable optimisations)
+      # force 'Release' build (needed by MSVC to enable optimisations), and force serial (to avoid ADIOS2 OOM)
       - name: Compile
-        run: cmake --build ${{ env.build_dir }} --config Release --parallel
+        run: cmake --build ${{ env.build_dir }} --config Release --parallel 1
 
       # run v4 unit tests in random order, excluding the integration tests,
       # using the default environment variables (e.g. test all permutations)
       # TODO:
       # ctest currently doesn't know of our Catch2 tags, so we
       # are manually excluding each integration test by name
+
       - name: Run v4 tests
         if: ${{ matrix.version == 4 }}
         run: ctest -j2 --output-on-failure --schedule-random -C Release -E "density evolution"

CMakeLists.txt

@@ -187,7 +187,6 @@ message(STATUS "AMD GPU acceleration is turned ${QUEST_ENABLE_HIP}. Set QUEST_EN
 
 # GPU Performance Tuning
 # (We do not print this value when configuring CMake as it is for advanced users only)
-
 set(quest_tpb_description # (the games we play for multi-line set() strings!)
   "The default number of threads per block QuEST will use when offloading to a GPU. Set to 128 by default. "
   "Must be a multiple of 32 (on NVIDIA GPUs) or 64 (on AMD GPUs). Can be overridden at executable launch "
@@ -199,6 +198,21 @@ set(QUEST_DEFAULT_NUM_GPU_THREADS_PER_BLOCK 128
 mark_as_advanced(QUEST_DEFAULT_NUM_GPU_THREADS_PER_BLOCK)
 
 
+# Checkpointing with ADIOS2
+option(
+  QUEST_ENABLE_ADIOS2
+  "Whether QuEST will be built with ADIOS2, enabling checkpointing (via saveQuregToFile / createQuregFromFile). Turned OFF by default."
+  OFF
+)
+message(STATUS "ADIOS2 integration is turned ${QUEST_ENABLE_ADIOS2}. Set QUEST_ENABLE_ADIOS2 to modify.")
+
+option(
+  QUEST_DOWNLOAD_ADIOS2
+  "Whether ADIOS2 will be downloaded if it is enabled but not found. Turned ON by default."
+  ON
+)
+
+
 # Deprecated API
 option(
   QUEST_ENABLE_DEPRECATED_API
@@ -543,6 +557,81 @@ endif()
 
 
 
+# Checkpointing (ADIOS2)
+if (QUEST_ENABLE_ADIOS2)
+
+  find_package(adios2 QUIET)
+
+  # A distributed QuEST needs an MPI-enabled ADIOS2 (which provides the
+  # adios2::cxx_mpi target). A serial system install lacks it, so in that case we
+  # ignore the found package and fetch an MPI-enabled build instead of failing.
+  set(quest_use_found_adios2 ${adios2_FOUND})
+  if (adios2_FOUND AND QUEST_ENABLE_MPI AND NOT TARGET adios2::cxx_mpi)
+    message(STATUS "Found ADIOS2 lacks MPI support (no adios2::cxx_mpi target); fetching an MPI-enabled build instead")
+    set(quest_use_found_adios2 FALSE)
+  endif()
+
+  if(NOT quest_use_found_adios2 AND QUEST_DOWNLOAD_ADIOS2)
+    message(STATUS "fetching ADIOS2 via FetchContent")
+
+    include(FetchContent)
+    FetchContent_Declare(
+      adios2
+      GIT_REPOSITORY https://github.com/ornladios/ADIOS2.git
+      GIT_TAG v2.12.1
+    )
+
+    # Match ADIOS2's MPI to QuEST's so distributed runs write per-rank slices
+    # into one shared file. ADIOS2's CUDA support is deliberately left OFF:
+    # checkpointing copies amps to host memory (syncQuregFromGpu/syncQuregToGpu)
+    # before any I/O, so ADIOS2 never touches device pointers. Building it with
+    # CUDA is unnecessary and stalls the Windows CUDA CI job.
+    set(ADIOS2_USE_MPI  ${QUEST_ENABLE_MPI} CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_CUDA OFF 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_DataMan    OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_SSC        OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_MHS        OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_DAOS       OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_MGARD      OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_BZip2      OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_Blosc      OFF CACHE BOOL "" FORCE)
+    set(ADIOS2_USE_Blosc2     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()
+    # re-run non-QUIET so configuration fails with a clear error if the package
+    # somehow became unavailable between the two calls
+    find_package(adios2 REQUIRED)
+  endif()
+
+  # In distributed builds link ADIOS2's MPI-enabled C++ interface: it defines
+  # ADIOS2_USE_MPI, which exposes the adios2::ADIOS(MPI_Comm) constructor used in
+  # qureg.cpp for collective per-rank I/O. The serial target lacks it.
+  if (QUEST_ENABLE_MPI)
+    target_link_libraries(QuEST PRIVATE adios2::cxx_mpi)
+  else()
+    target_link_libraries(QuEST PRIVATE adios2::cxx)
+  endif()
+  set(CMAKE_INSTALL_RPATH_USE_LINK_PATH ON)
+endif()
+
+
+
 # ===============================
 # Set options to save in config.h
 # ===============================
@@ -553,6 +642,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_ADIOS2 ${QUEST_ENABLE_ADIOS2})
 set(QUEST_INCLUDE_DEPRECATED_FUNCTIONS ${QUEST_ENABLE_DEPRECATED_API})
 
 

docs/cmake.md

@@ -44,6 +44,8 @@ make
 | `QUEST_ENABLE_CUDA` | (`OFF`), `ON` | Determines whether QuEST will be built with support for NVIDIA GPU acceleration. If turned on, `CMAKE_CUDA_ARCHITECTURES` should probably also be set. |
 | `QUEST_ENABLE_CUQUANTUM` | (`OFF`), `ON` | Determines whether QuEST will make use of the NVIDIA CuQuantum library. Cannot be turned on if `QUEST_ENABLE_CUDA` is off. |
 | `QUEST_ENABLE_HIP` | (`OFF`), `ON` | Determines whether QuEST will be built with support for AMD GPU acceleration. If turned on, `CMAKE_HIP_ARCHITECTURES` should probably also be set. |
+| `QUEST_ENABLE_ADIOS2` | (`OFF`), `ON` | Determines whether QuEST will be built with ADIOS2 to enable checkpointing, via functions `saveQuregToFile()` and `createQuregFromFile()`. |
+| `QUEST_DOWNLOAD_ADIOS2` | (`ON`), `OFF` | Determines whether to download ADIOS2 from Github, when ADIOS2 is enabled but not found. |
 | `QUEST_ENABLE_DEPRECATED_API` | (`OFF`), `ON` | Determines whether QuEST will be built with support for the deprecated (v3) API. ***Note**: this will generate compiler warnings and is not supported by MSVC.* |
 | `QUEST_DISABLE_DEPRECATION_WARNINGS` | (`OFF`), `ON` | Whether to disable the compile-time deprecation warnings when using the deprecated (v3) API. |
 | `USER_SOURCE_NAMES` | (Undefined), String | The source file for a user program which will be compiled alongside QuEST. `USER_OUTPUT_EXE_NAME` *must* also be defined. |

docs/compile.md

@@ -45,6 +45,7 @@ Compiling is configured with variables supplied by the [`-D` flag](https://cmake
 > - <a href="#compile_cuquantum">cuQuantum</a>
 > - <a href="#compile_distribution">Distribution</a>
 > - <a href="#compile_multi-gpu">Multi-GPU</a>
+> - <a href="#compile_checkpointing">Checkpointing</a>
 
 > **See also**:
 > - [`cmake.md`](cmake.md) for the full list of passable compiler variables.
@@ -689,3 +690,32 @@ Note that distributed executables are launched in a distinct way to the other de
 > - UCX
 > - launch flags
 > - checking via reportenv
+
+
+
+
+------------------
+
+
+<!-- permit doxygen to reference section -->
+<a id="compile_checkpointing"></a>
+
+## Checkpointing
+
+QuEST has optional facilities for _checkpointing_ a `Qureg`; writing its state to a file with [`saveQuregToFile()`](https://quest-kit.github.io/QuEST/group__experimental.html#gaf9a1aec34fdfdb3c650dc60e5a8ac9d9), to be later restored into a new `Qureg` with [`createQuregFromFile()`](https://quest-kit.github.io/QuEST/group__experimental.html#gab1ebe89e2ff15470fa340d4c2ced5703). 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. 
+
+Checkpointing is built upon [ADIOS2](https://github.com/ornladios/ADIOS2) and is _disabled_ by default. To enable it, simply specify `QUEST_ENABLE_ADIOS2` at configuration:
+```bash
+# configure
+cmake .. -D QUEST_ENABLE_ADIOS2=ON
+
+# build
+cmake --build . --parallel
+```
+
+If a compatible ADIOS2 is not found, it will be automatically downloaded and installed from the ADIOS2 [Github](https://github.com/ornladios/ADIOS2), unless `QUEST_DOWNLOAD_ADIOS2` is overridden to be `OFF`. If an existing ADIOS2 is installed in a non-standard location (such as `~/.local`), pass its prefix via [`CMAKE_PREFIX_PATH`](https://cmake.org/cmake/help/latest/variable/CMAKE_PREFIX_PATH.html):
+> ```bash
+> cmake .. -D QUEST_ENABLE_ADIOS2=ON -D CMAKE_PREFIX_PATH=$HOME/.local
+> ```
+
+Calling `saveQuregToFile()` or `createQuregFromFile()` in a build _without_ checkpointing enabled will trigger a runtime 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_ADIOS2)               || \
     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_ADIOS2
 
 
 // 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_ADIOS2)               || \
     ! 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_ADIOS2               == 0 || QUEST_COMPILE_ADIOS2               == 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)