High-performance backpropagation for structured-grid solvers

Abstract

This doctoral project is about the solution of inverse problems on hyperbolic PDEs. It includes work on automation of code generation, trading off memory consumption with extra computation, as well as new insights on the use of lossy compression in this setting. Domain-specific languages (DSLs) have emerged as a way to express programs at a high level of abstraction while enabling effective utilisation of modern hardware through specialised compilers. Partial Differential Equations (PDEs) are a standard language for expressing mathematical models in science and engineering, and Devito is a DSL designed to express programs that solve PDEs using the Finite Differ- ence (FD) method on structured meshes. Full Waveform Inversion (FWI) aims to estimate physical parame- ters of the earth by minimising the misfit between an observed signal and one simulated through an FD solution of the Wave PDE. This is solved using gradient-based methods, but the computation of derivatives introduces a high memory requirement. This is because the intermediate states that can be discarded for the forward problem are critical to calculating the derivatives, and storing these intermediate states requires large amounts of memory. A technique used here is checkpointing - where a subset of the states from the forward problem is stored, while the rest are discarded to save memory. The discarded states are recovered later by rerunning the forward computation from the last checkpointed state. A checkpointing library called pyRevolve was developed as part of this project. This library enables the use of checkpointing strategies with minimal changes in application code. We present a strategy to use a combination of compression and recomputation to extend the memory-recomputation tradeoff into a three-way memory-recomputation-compression tradeoff. We first compare lossy and lossless compression for this purpose. Next, we propose a performance model that pyRevolve can use to inform its use of compression for a given problem. The thesis also includes an analysis of the effects of errors introduced during lossy checkpoint compression on the rest of the FWI problem. The last piece of the puzzle is the automatic generation of fast, parallel derivative code corresponding to the objective function specified by the application programmer. We focus on the specific case of explicit FD solvers, which have a data-access pattern that is easy to run in parallel (shared or distributed memory). However, the derivatives generated by reverse mode AD do not preserve this parallelism. We show how derivatives of a finite-difference solver can be generated while preserving the parallelism of the original FD code.