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.

• The paper introduces NektarIR, a domain-specific MLIR dialect that abstracts high-order finite element operations for improved performance portability.
• It employs a multi-stage IR transformation pipeline, including decomposition, loop restructuring, and vectorization, to convert high-level formulations into hardware-optimized kernels.
• Benchmarking shows NektarIR surpasses hand-tuned implementations on AVX512 CPUs and offers competitive GPU performance, highlighting its potential in hp-adaptive CFD simulations.

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

Motivation and Context

The increasing complexity and heterogeneity of modern HPC hardware necessitates solutions that enable performance portability while delivering near-peak efficiency in critical scientific workloads. High-order and spectral/hp element methods, widely adopted in computational fluid dynamics (CFD), achieve superior numerical accuracy and computational intensity by leveraging dense operator kernels optimized for memory bandwidth limitations. However, conventional approaches—either hand-tuned for specific architectures or relying on general-purpose compilers—struggle to maintain both portability and performance as simulation parameters and hardware configurations evolve dynamically, especially in $hp$-adaptive contexts involving mixed element types and polynomial orders.

NektarIR Dialect and Abstraction

NektarIR introduces a domain-specific MLIR dialect designed to express finite element operator evaluations at a high abstraction level closely matched to established mathematical formulations. It features a custom block type encapsulating mesh element shape, polynomial expansion order, quadrature configuration, and memory layout metadata. The dialect supports both batched and elemental operations on expansion modes and physical-space quantities, facilitating optimizations by exposing domain-relevant information directly in the IR.

Figure 1: Examples of the block type demonstrating block-level grouping for mesh elements with consistent shape, order, and quadrature.

The block type abstraction, together with attributes encoding data layout and sparsity (particularly relevant for simplicial elements such as tetrahedra), provides the foundational interface for geometry-specific transformations and vectorization. This enables efficient memory interleaving and loop restructuring for SIMD-aware code generation.

Multi-stage IR Transformations

NektarIR executes a multi-stage lowering pipeline, leveraging both custom and upstream MLIR passes to progressively transform domain-specific operators into hardware-optimized kernels. The transformation sequence includes:

• Element-wise decomposition of batch operations (e.g., block backward transforms to single-element transforms within loops)
• Bufferization, mapping block types to memory regions compatible with vector loads and device-specific allocation semantics
• Conversion to explicit loop nests using the affine dialect, optimizing sum-factorization for tensor-product expansions

Figure 2: Conversion from block backward transform (nir.bwd) to loops over elements and elemental backward transforms (nir.elemental_bwd).

Figure 3: Transformation from NektarIR elemental operator representation to affine loop nests, matching sum-factorized matrix operations.

Through these stages, NektarIR exploits compile-time knowledge of element topology and expansion structure, automating loop unrolling, fusion, and memory pool allocation for temporary storage. Custom loop coalescing transformations address non-rectangular iteration domains, essential for collapsed-coordinate integration in tetrahedral elements.

Figure 4: Loop coalescing of triangular loop nests, enabling efficient execution of sum-factorization for non-tensorial shapes.

Compilation Pipeline and Hardware Targeting

The compilation pipeline extends from high-level spectral/hp element abstractions in C++ or Python frontends, traversing NektarIR and multiple MLIR dialects (vector, affine, gpu, nvvm/rocdl), before emitting LLVM IR for final code generation. For CPU targets, vectorization is tailored to the element batch dimension, with layout optimizations for contiguous SIMD loads and stores.

GPU-targeting transformations introduce device-grid parallelism via gpu.launch, automatically mapping loop domains to grid/block/thread hierarchies. Two distinct parallelization strategies are supported: threading across elements (each thread processes a full element) and threading across expansion coefficients (threads process individual mode indices, with loop coalescing for reduction kernels). Temporary storage allocations are managed with awareness of device register and shared/global memory constraints, including padding and expansion for efficient access.

Figure 5: NektarIR lowering pipeline from frontend abstraction to hardware-specific kernel, spanning MLIR dialects and LLVM IR.

Performance and JIT Overhead

Comprehensive benchmarking quantifies both compilation overhead and runtime efficiency. Kernel lowering and JIT compilation times are consistently below 1 second—negligible in typical CFD workloads—validating suitability for adaptive and dynamic simulation environments.

Figure 6: Time to lower the Helmholtz operator for hexahedral and tetrahedral elements from NektarIR to LLVM IR across host and device targets.

Figure 7: Time to compile Helmholtz operator for hexahedral and tetrahedral elements from NektarIR to LLVM IR for host and device targets.

Numerical Throughput Results

NektarIR achieves strong throughput for matrix-free high-order kernels. For AVX512-equipped CPUs, code-generated kernels consistently surpass hand-tuned Nektar++ implementations, maintaining high throughput up to large polynomial orders and across both hexahedral and tetrahedral element shapes.

Figure 8: Throughput comparison of AVX512 Helmholtz kernels in NektarIR and Nektar++ on AMD EPYC 9554 CPU for both element types and polynomial orders.

On GPU (NVIDIA H100), throughput is competitive with specialized template-based Nektar++ kernels for moderate problem sizes and element types, though register pressure in NektarIR kernels induces performance degradation for large-scale problems and certain threading strategies. The threading-over-coefficient strategy exhibits especially pronounced performance gaps for tetrahedral elements, warranting further low-level IR optimization.

Figure 9: Throughput comparison of Helmholtz kernels in NektarIR and Nektar++ on NVIDIA H100 GPU for both threading strategies and element types.

Implications and Future Directions

The NektarIR framework demonstrates the efficacy of domain-specific compilation for high-order finite element operations, closing the gap between performance portability and hardware efficiency for spectral/hp element solvers. Its dialect-level abstraction and multi-stage IR rewriting pipeline generalize to other structured scientific domains, with architectural extensibility inherent to MLIR and LLVM ecosystems.

Practical implications include:

• Deployment in $hp$-adaptive CFD simulations where runtime performance tuning is essential
• Integration with other high-order frameworks (Neko, MFEM, deal.II) leveraging the IR abstraction for diverse basis types and element configurations
• Extension to broader PDE solver pipelines, supporting a wider range of test functions, quadrature rules, and mesh shapes

Theoretically, NektarIR's approach offers a model for leveraging structural similarities between high-order FEM kernels and tensor contractions, unlocking further compiler-assisted optimizations via MLIR's linalg and sparse dialects.

Conclusion

NektarIR provides a robust domain-specific compiler for high-order finite element kernels, advancing state-of-the-art performance portability in spectral/hp element CFD solvers. Its MLIR-based approach achieves fast, adaptive code generation and strong runtime throughput—surpassing hand-tuned kernels on CPUs and approaching them on GPUs. Remaining GPU optimization challenges highlight the importance of low-level IR tuning and motivate continued research on compiler strategies for complex non-tensorial domains. The abstraction, pipeline, and transforms introduced in NektarIR are directly relevant to scientific computing and compiler research and have broad applicability as high-order methods are increasingly adopted across simulation disciplines (2606.20917).