Optimal checkpointing for adjoint multistage time-stepping schemes

Abstract

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).