Step-based checkpointing with high-level algorithmic differentiation

TL;DR

This paper introduces a novel approach combining high-level algorithmic differentiation with step-based checkpointing to efficiently handle adjoint calculations in time-dependent PDE solvers, improving performance by leveraging data structure insights.

Contribution

It presents a new method that integrates high-level AD with dynamic checkpointing schedules, enhancing efficiency for adjoint computations without prior knowledge of the forward calculation structure.

Findings

Improved checkpointing strategies reduce memory usage and computation time.

High-level AD simplifies data management for adjoint calculations.

Method is particularly effective for time-dependent PDE solvers.

Abstract

Automated code generation allows for a separation between the development of a model, expressed via a domain specific language, and lower level implementation details. Algorithmic differentiation can be applied symbolically at the level of the domain specific language, and the code generator reused to implement code required for an adjoint calculation. However the adjoint calculations are complicated by the well-known problem of storing or recomputing the forward data required by the adjoint, and different checkpointing strategies have been developed to tackle this problem. This article considers the combination of high-level algorithmic differentiation with step-based checkpointing schedules, with the primary application being for solvers of time-dependent partial differential equations. The focus is on algorithmic differentiation using a dynamically constructed record of forward operations, where the precise structure of the original forward calculation is unknown ahead-of-time. In addition, high-level approaches provide a simplified view of the model itself. This allows data required to restart and advance the forward, and data required to advance the adjoint, to be identified. The difference between the two types of data is here leveraged to implement checkpointing strategies with improved performance.