An MLIR Lowering Pipeline for Stencils at Wafer-Scale

TL;DR

This paper addresses the challenge of running stencil HPC kernels on the Cerebras WSE without changing application code. It introduces an MLIR-based stencil lowering pipeline that preserves stencil semantics while targeting the WSE’s asynchronous, actor-like execution using three domain-specific dialects (csl-stencil, csl-wrapper, csl-ir) and a runtime library. The approach is validated on five benchmarks across three frontends and two hardware generations, delivering competitive or superior performance to hand-tuned CSL and to large GPU/CPU clusters (e.g., up to 14x and 20x faster on WSE3 for certain workloads). The work demonstrates substantial programmer productivity gains via automatic lowering, and it releases the tooling as open source to enable broader adoption and extension to other CGRA-like architectures.

Abstract

The Cerebras Wafer-Scale Engine (WSE) delivers performance at an unprecedented scale of over 900,000 compute units, all connected via a single-wafer on-chip interconnect. Initially designed for AI, the WSE architecture is also well-suited for High Performance Computing (HPC). However, its distributed asynchronous programming model diverges significantly from the simple sequential or bulk-synchronous programs that one would typically derive for a given mathematical program description. Targeting the WSE requires a bespoke re-implementation when porting existing code. The absence of WSE support in compilers such as MLIR, meant that there was little hope for automating this process. Stencils are ubiquitous in HPC, and in this paper we explore the hypothesis that domain specific information about stencils can be leveraged by the compiler to automatically target the WSE without requiring application-level code changes. We present a compiler pipeline that transforms stencil-based kernels into highly optimized CSL code for the WSE, bridging the semantic gap between the mathematical representation of the problem and the WSE's asynchronous execution model. Based upon five benchmarks across three HPC programming technologies, running on both the Cerebras WSE2 and WSE3, our approach delivers comparable, if not slightly better, performance than manually optimized code. Furthermore, without requiring any application level code changes, performance on the WSE3 is around 14 times faster than 128 Nvidia A100 GPUs and 20 times faster than 128 nodes of a CPU-based Cray-EX supercomputer when using our approach.

Figures (11)

• Figure 1: A time-step loop surrounding asynchronous data exchanges in CSL.
• Figure 2: Through the use of three WSE MLIR dialects, we incrementally move from the stencil mathematical representation to the actor-based WSE execution model.
• Figure 3: Our staged compilation flow extends existing infrastructure to target novel hardware through multiple stages, allowing efficient code generation for the Cerebras WSE.
• Figure 4: Example stencil code used as a running example in this section.
• Figure 5: After distributing $x$, $y$ dimensions across PEs and tensorizing $z$ dimension.