diff --git a/CMakeLists.txt b/CMakeLists.txt index b5a438713..486606d70 100644 --- a/CMakeLists.txt +++ b/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}) diff --git a/quest/include/config.h.in b/quest/include/config.h.in index 1bb8a0470..ff2c7f07f 100644 --- a/quest/include/config.h.in +++ b/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) diff --git a/quest/include/qureg.h b/quest/include/qureg.h index 4ff4c5627..d2b6525df 100644 --- a/quest/include/qureg.h +++ b/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 } diff --git a/quest/src/api/qureg.cpp b/quest/src/api/qureg.cpp index 84bcd2bd0..51184eaa1 100644 --- a/quest/src/api/qureg.cpp +++ b/quest/src/api/qureg.cpp @@ -25,6 +25,12 @@ #include #include +// 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 +#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("numQubits"); + auto vIsDensMatr = io.DefineVariable("isDensityMatrix"); + auto vAmpBytes = io.DefineVariable("ampSizeBytes"); + auto vAmps = io.DefineVariable("amplitudes", + {static_cast(totalBytes)}, // global shape + {static_cast(offsetBytes)}, // this node's start + {static_cast(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(ampBytes)); + } + engine.Put(vAmps, reinterpret_cast(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("numQubits"), numQubits, adios2::Mode::Sync); + engine.Get(io.InquireVariable("isDensityMatrix"), isDensMatr, adios2::Mode::Sync); + engine.Get(io.InquireVariable("ampSizeBytes"), fileAmpBytes, adios2::Mode::Sync); + validate_checkpointFileMatchesPrecision(fileAmpBytes, static_cast(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(sizeof(qcomp)); + qindex offsetBytes = util_getGlobalIndexOfFirstLocalAmp(qureg) * static_cast(sizeof(qcomp)); + adios2::Variable vAmps = io.InquireVariable("amplitudes"); + vAmps.SetSelection({{static_cast(offsetBytes)}, {static_cast(localBytes)}}); + engine.Get(vAmps, reinterpret_cast(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 } diff --git a/quest/src/core/validation.cpp b/quest/src/core/validation.cpp index 62ff93166..1cf52cf81 100644 --- a/quest/src/core/validation.cpp +++ b/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 */ diff --git a/quest/src/core/validation.hpp b/quest/src/core/validation.hpp index 87f81a0d6..32f42c641 100644 --- a/quest/src/core/validation.hpp +++ b/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 */