Optimal checkpointing for adjoint multistage time-stepping schemes

• Published in: Journal of computational science Vol. 66; p. 101913
• Main Authors: Zhang, Hong, Constantinescu, Emil M.
• Format: Journal Article
• Language: English
• Published: Elsevier B.V 01.01.2023
• Subjects: Adjoint method; Checkpointing; Dynamic programming; Multistage time-stepping schemes; PETSc; Checkpointing; PETSc; Adjoint method; Dynamic programming; Multistage time-stepping schemes
• ISSN: 1877-7503; 1877-7511; 1877-7511
• DOI: 10.1016/j.jocs.2022.101913

We consider checkpointing strategies that minimize the number of recomputations needed when performing discrete adjoint computations using multistage time-stepping schemes that require computing several substeps within one complete time step. Specifically, we propose two algorithms that can generate optimal checkpointing schedules under weak assumptions. The first is an extension of the seminal Revolve algorithm adapted to multistage schemes. The second algorithm, named CAMS, is developed based on dynamic programming, and it requires the least number of recomputations when compared with other algorithms. The CAMS algorithm is made publicly available in a library with bindings to C and Python. Numerical results show that the proposed algorithms can deliver up to two times the speedup compared with that of classical Revolve. Moreover, we discuss the utilization of the CAMS library in mature scientific computing libraries and demonstrate the ease of using it in an adjoint workflow. The proposed algorithms have been adopted by the PETScTSAdjoint library. Their performance has been demonstrated with a large-scale PDE-constrained optimization problem on a leadership-class supercomputer. This work is a significant extension of the authors’ conference paper (Zhang and Constantinescu, 2021). •Optimal checkpointing algorithms for multistage time-stepping schemes.•A library for efficient large-scale adjoint calculation with limited storage budget.•The library can be easily incorporated into scientific libraries such as PETSc.•The performance advantages are demonstrated on a leadership-class supercomputer.