Performance Portable Adjoints for Structured Mesh Applications with OPS
TL;DR
This work extends the Oxford Parallel library for Structured mesh solvers (OPS) domain-specific language to compute derivatives using the reverse mode of AD, and focuses on adjoint-mode AD for structured mesh stencil applications.
Abstract
Derivatives are crucial for engineering and scientific applications like optimization and inverse problems. While finite difference approximations can estimate derivatives, they are computationally expensive and inaccurate. Algorithmic Differentiation (AD) provides an efficient and exact method to compute derivatives by treating computer programs as mathematical functions and applying the chain rule of calculus. Our work focuses on adjoint-mode AD for structured mesh stencil applications. We extend the Oxford Parallel library for Structured mesh solvers (OPS) domain-specific language to compute derivatives using the reverse mode of AD. OPS allows developers to express mesh algorithms from a high-level code targeting multiple hardware from the same source. Taking advantage of the domain-specific abstraction, the extension creates a compact adjoint tape at the level of computational loops and generates the platform-specific (OpenMP and CUDA) parallel implementations for the adjoint loops, using the user-provided primal and adjoint stencil-kernels. We differentiate three example applications written in OPS, demonstrating similar performance to the original applications on both CPUs and GPUs. On these applications, computing derivatives only took \(3.7{-}9.7x\) time (including the evaluation of the application) compared to the original applications, which is in line with state-of-the-art tools.