Performance portable adjoints for structured mesh applications with OPS

Source 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 Structured mesh solver library (OPS) domain-specific language to compute derivatives using the reverse mode of algorithmic differentiation. 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.