A shared compilation stack for distributed-memory parallelism in stencil DSLs

TL;DR

This paper presents a shared compilation framework based on MLIR for distributed-memory parallelism in stencil DSLs, enabling shared infrastructure and high-performance code generation across multiple HPC stencil compilers.

Contribution

It adapts MLIR for HPC stencil DSLs, introducing new abstractions for message passing and demonstrating shared compiler components across three different stencil DSLs.

Findings

Shared compiler components improve development efficiency.

Framework generates high-performance distributed stencil code.

Shared infrastructure enhances scalability and maintainability.

Abstract

Domain Specific Languages (DSLs) increase programmer productivity and provide high performance. Their targeted abstractions allow scientists to express problems at a high level, providing rich details that optimizing compilers can exploit to target current- and next-generation supercomputers. The convenience and performance of DSLs come with significant development and maintenance costs. The siloed design of DSL compilers and the resulting inability to benefit from shared infrastructure cause uncertainties around longevity and the adoption of DSLs at scale. By tailoring the broadly-adopted MLIR compiler framework to HPC, we bring the same synergies that the machine learning community already exploits across their DSLs (e.g. Tensorflow, PyTorch) to the finite-difference stencil HPC community. We introduce new HPC-specific abstractions for message passing targeting distributed stencil computations. We demonstrate the sharing of common components across three distinct HPC stencil-DSL compilers: Devito, PSyclone, and the Open Earth Compiler, showing that our framework generates high-performance executables based upon a shared compiler ecosystem.