NektarIR: A Domain-Specific Compiler for High-Order Finite Element Operations on Heterogeneous Hardware

Abstract

Modern high performance computing (HPC) applications must target heterogeneous hardware. This requires significant work to ensure domain specific implementations translate to highly performant kernels across a range hardware types and vendors, each requiring bespoke optimization to make use of the specific target architecture. Through the development of a domain specific compiler built with the multi-level intermediate representations (MLIR) project, one can express a high-level, close to the specific domain, abstraction that is progressively lowered to a low, close to metal, abstraction. At each intermediate representation (IR), appropriate optimizations can be applied without costly analysis due to the knowledge embedded in the domain specific IRs. We apply this method to the construction of discrete differential operators for use in spectral/hp element method solvers for computational fluid dynamics (CFD). Here, the performance is driven by a small set of common finite element operators that are composed to create kernels for the discrete differential operators used to solve weak partial differential equations. We create our own MLIR dialect to represent these operators and implement a bespoke lowering pipeline to facilitate the just-in-time compilation of these kernels for both CPU and GPU architecture and illustrate performance comparisons with the Nektar++ spectral/hp element framework.