An MLIR Lowering Pipeline for Stencils at Wafer-Scale

Nicolai Stawinoga, David Katz, Anton Lydike, Justs Zarins, Nick Brown, George Bisbas and 1 other

The Cerebras Wafer-Scale Engine is a single silicon wafer holding over 900,000 tiny independent processors, and it is superb for scientific computing — but you program it in an alien, fully asynchronous message-passing style so far from ordinary math code that porting an existing simulation means rewriting it by hand, and no compiler could automate that; this paper (TU Berlin, Edinburgh, Cambridge — Nicolai Stawinoga, Tobias Grosser, and colleagues) builds an MLIR-based compiler pipeline that takes a high-level stencil kernel and, through a chain of intermediate dialects down to one that mirrors Cerebras's own language, lowers it automatically into optimized wafer code — running about 14 times faster than 128 Nvidia A100 GPUs with no application code changes.

The argument in one glance

Problem

A wafer-sized supercomputer that almost nobody can program

The Cerebras Wafer-Scale Engine (WSE) is a single chip the size of an entire silicon wafer, holding over 900,000 tiny processors, called processing elements (PEs), each with its own small memory and its own router, all wired together on one piece of silicon. Built for artificial intelligence, it turns out to be excellent for high-performance scientific computing too — the kind of physics simulations that model weather, earthquakes, or heat flow. The problem is how you program it. The WSE runs a distributed, fully asynchronous model: each PE acts independently, like an actor that reacts to messages arriving from its neighbors, with no global clock or shared control. That is nothing like the simple sequential or bulk-synchronous programs a scientist naturally writes from a mathematical description. To run existing code on the WSE, you have to re-implement it by hand in Cerebras's low-level language, painstakingly wiring up asynchronous message exchanges and callbacks. And mainstream compilers, including MLIR, had no support for the WSE at all, so there was little hope of automating this port. A remarkable machine sat largely out of reach.
What people did

Hand-port each application into the wafer's asynchronous language

The established path is manual. For each application you want on the WSE, an expert rewrites the computation from scratch in Cerebras's low-level language, explicitly programming how the processing elements exchange data asynchronously, react to incoming messages with callbacks, and route information across the wafer. It is a bespoke re-implementation for every code.
Why it fell short

Hand-porting is slow, expert-only, and does not scale to many codes

Rewriting each program by hand is enormously laborious and demands rare expertise in the wafer's asynchronous, actor-based model — a model with no top-level loop construct, where a time-step loop must be expressed through message callbacks and routing configuration rather than ordinary control flow. Every new simulation is a fresh, months-long effort, and the result is brittle and hard to change. Because compilers like MLIR did not understand the WSE, none of this could be automated: there was no way to take a high-level mathematical program and have a tool generate correct, efficient wafer code. So the WSE's extraordinary hardware was accessible only to the few teams willing to hand-craft each kernel, and the vast body of existing scientific code could not benefit. What was missing was a compiler that could bridge the semantic gap between the clean mathematical form of a problem and the wafer's asynchronous execution, automatically.
This paper's idea

Lower high-level stencils to wafer code through a chain of MLIR dialects

This paper presents a compiler pipeline, built on MLIR and the compatible xDSL framework, that automatically turns high-level stencil kernels into optimized code for the WSE, with no application-level changes. Its guiding hypothesis is that domain-specific knowledge about stencils — computations where each grid point is updated from a fixed pattern of its neighbors, the backbone of finite-difference physics solvers — gives the compiler enough structure to target the wafer on its own. MLIR is a compiler framework built from reusable dialects, each a set of operations at some level of abstraction, and compilation proceeds by lowering: progressively rewriting a program from high-level dialects toward low-level ones. The pipeline starts from a high-level stencil dialect that expresses the computation mathematically, then applies a series of intermediate domain-specific dialects, and finally a low-level csl-ir dialect that re-implements a large subset of Cerebras's own language (CSL) inside MLIR. A strategy of lowering transformations carries the program down this chain, generating the asynchronous message exchanges, callbacks, and routing the wafer needs — bridging the gap from the mathematical representation to the WSE's actor-based execution. The result matches or beats hand-optimized code: on the WSE3, with no application-level changes, it runs around 14 times faster than 128 Nvidia A100 GPUs and 20 times faster than 128 nodes of a CPU-based Cray-EX supercomputer.

Why this problem is worth solving

Wafer-scale hardware offers a step change in performance for science, but only if the enormous gap between how scientists write code and how the wafer runs can be bridged.

Ordinary chips are cut from a silicon wafer into many small dies; the Cerebras Wafer-Scale Engine keeps the whole wafer as one chip, giving over 900,000 processing elements on a single piece of silicon, connected by an on-chip interconnect with enormous bandwidth. Each processing element has a compute core, its own router, and a small local memory (48 kilobytes of fast SRAM), and the elements operate fully independently. This design is a dataflow architecture, and its execution model is asynchronous and actor-based: a processing element sits idle until data arrives from a neighbor, then reacts — like an actor responding to a message — running a small piece of work and forwarding results onward, with routing across the wafer configured through virtual channels. There is no global sequential program and, in the wafer’s language, not even a normal top-level loop. Scientists, by contrast, describe a simulation mathematically and implement it as straightforward sequential or bulk-synchronous code. The distance between those two worlds — the clean math on one side and the swarm of asynchronous actors on the other — is the semantic gap, and it is why the wafer, despite being ideal for high-performance computing, has been usable only by experts willing to hand-write everything in its native style.

The cost of that gap is that a whole field of existing software cannot reach the hardware.

To put a simulation on the WSE today, an expert must re-implement it from scratch in Cerebras’s low-level language, explicitly orchestrating how each processing element exchanges data with its neighbors, handling completion callbacks when asynchronous transfers finish, and setting up the routes that carry messages across the wafer. Because the language has no ordinary loop construct, even something as basic as iterating over time steps must be encoded through the actor-and-callback machinery. This is slow, error-prone work that only a handful of specialists can do well, and it must be redone for every new application. Scientific computing, however, runs on a large accumulated base of code and on high-level tools that let domain scientists express problems in familiar terms without worrying about low-level parallelism. None of that can target the wafer if each program needs a months-long bespoke rewrite. And since mainstream compiler infrastructure such as MLIR had no notion of the WSE, there was no foundation on which to build automation. The value of closing this gap is therefore large and direct: it is the difference between a wafer-scale machine that only its makers can fully exploit and one that the broad scientific-computing community can compile to from the high-level code it already writes.

What people did before, and where it broke

Hand-write each kernel in Cerebras’s low-level language. Have an expert re-implement the computation directly in the wafer’s native language, orchestrating asynchronous exchanges and routing. Where it broke. It is enormously laborious, needs rare expertise, must be repeated for every application, and produces brittle code — so only a few teams and a few kernels ever reach the wafer.
Rely on the wafer’s asynchronous model with no top-level loop. Encode ordinary control flow, like a time-step loop, through message callbacks and routing configuration by hand. Where it broke. Expressing normal program structure through the actor-and-callback machinery is unintuitive and hard to get right, widening the gap between the math and the implementation.
Use existing compilers to target the wafer. Hope that mainstream infrastructure such as MLIR could generate WSE code. Where it broke. There was simply no WSE support in these compilers, so there was no path to automate the port at all.
The gap. No compiler could take a high-level mathematical program, exploit its structure, and automatically generate correct, efficient, asynchronous wafer code — leaving the WSE’s hardware unreachable for the broad base of existing scientific software.

How the pipeline works

Lowering a stencil, dialect by dialect, down to wafer code

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Use a chain of intermediate representations to cross a huge gap in steps. The distance from a mathematical stencil to asynchronous wafer code is too large to jump at once, so the pipeline builds a ladder of MLIR dialects — high-level stencil, intermediate domain-specific dialects, and a low-level csl-ir dialect — and lowers the program one manageable rewrite at a time. Each dialect adds a bit more wafer-specific detail while keeping the program analyzable. Bridging the distance between a high-level program and an unusual hardware target by defining intermediate representations that step down toward the hardware is the same approach the work on an MLIR-based intermediate representation for hardware accelerators covered elsewhere on this site is built on, giving the compiler a structured path to the device. This pipeline applies that layered-dialect strategy to reach an entire wafer.

Let the mathematical structure of stencils drive the mapping to a dataflow machine. The pipeline’s central bet is that knowing the computation is a stencil — a regular, local neighbor pattern over a grid — tells the compiler how to spread the grid across the processing elements and which boundary data neighbors must exchange, so it can generate the wafer’s dataflow automatically. Turning a clean mathematical description into an efficient mapping onto a dataflow architecture by exploiting the structure of the math is the same principle the work on compiling Einstein summations to locality-aware dataflow covered elsewhere on this site relies on, using the form of the computation to place data and work well. Here the structure being exploited is the locality of stencils, and the target is the wafer’s grid of independent elements.

Automate the descent into an asynchronous, distributed execution model. The hardest part of the wafer is its actor-based, message-passing execution, and the pipeline’s lowering transformations generate exactly that — the asynchronous exchanges, callbacks, and routing — from a program that never mentioned them. Compiling a high-level program down to a distributed, dataflow hardware target by having the compiler synthesize the low-level communication and control is the same undertaking the work on a hierarchical dataflow compiler for high-level synthesis covered elsewhere on this site takes on, mechanically producing the dataflow implementation a person would otherwise hand-write. This pipeline does it for the WSE’s uniquely asynchronous fabric, including expressing loops the wafer’s language cannot state directly.

Following a stencil through the pipeline

Compiling one finite-difference solver onto the wafer

A trace, drawn from the paper, of how a stencil kernel is lowered through MLIR dialects into asynchronous wafer code.

Step 1 of 5

A scientist writes a stencil in high-level form

The input is a finite-difference kernel — for example, updating each grid cell from its neighbors to simulate a wave — written in a high-level way, or coming from an existing scientific tool, with no wafer-specific code.

The scientist works in familiar mathematical terms, not the wafer's language.

Why it works

A ladder of dialects makes an impossible jump tractable. The gap from math to asynchronous wafer code is too wide to cross directly → but lowering through high-level, intermediate, and csl-ir dialects crosses it in small verified steps → so each rewrite stays correct and analyzable.
Stencil structure tells the compiler how to distribute and communicate. A stencil is a regular, local neighbor pattern → so the compiler can split the grid across processing elements and know exactly which boundary values neighbors must exchange → so it generates the wafer’s dataflow without human guidance.
A CSL-mirroring dialect keeps MLIR’s power at the lowest level. The wafer’s language is far from MLIR → but a csl-ir dialect that re-implements a large subset of CSL brings it inside MLIR → so the compiler can transform and optimize wafer code with standard infrastructure.
Automating the async model removes the expert bottleneck. Hand-writing actors, callbacks, and routing needs rare skill → but lowering transformations synthesize all of it from the high-level program → so any stencil code can target the wafer with no application-level changes.
Matching the math to the hardware yields top performance. Stencils are local and the wafer has huge on-chip bandwidth and many elements → so a well-mapped stencil keeps the elements busy with cheap neighbor exchanges → so the generated code beats 128 A100 GPUs by about 14 times.

What it costs (trade-offs and scope)

It targets stencil computations specifically. The whole approach leans on the regular, local structure of stencils to drive distribution and communication; computations without that structure — irregular or global-dependency problems — would not map through this pipeline the same way and are outside its stated scope.
It re-implements a subset of the wafer’s language, not all of it. The csl-ir dialect covers a large subset of CSL, not the entire language; capabilities of the wafer that fall outside that subset are not expressed, so some hand-written CSL can still reach features the pipeline cannot.
It is tied to the Cerebras wafer and its toolchain. The lowering, the routing model, and the csl-ir dialect are built for the WSE’s specific architecture; the effort does not transfer directly to other accelerators without building analogous dialects and lowerings.
The headline comparison is against non-hand-tuned baselines. The roughly 14-times-over-128-A100 and 20-times-over-128-Cray-EX-nodes figures are strong, but the GPU baseline used OpenACC rather than hand-tuned CUDA (the tuned version being proprietary), so the exact multiples depend on the baselines chosen; the paper notes the advantage would remain significant regardless.
Peak performance still depends on the problem fitting the wafer. Results were measured across five benchmarks and three problem sizes on the WSE2 and WSE3; how well any given stencil performs depends on how its grid and memory footprint fit the wafer’s many small local memories, so gains vary with the problem.

The evaluation

Getting a scientific stencil onto the Wafer-Scale Engine

Property	Hand-port to CSL	Existing compilers (MLIR)	This MLIR pipeline
Automates targeting the wafer
Application-level code changes needed
Handles the async, actor-based model
Performance vs manual code
WSE3 performance (no code changes)

A comparison of ways to run a scientific stencil on the Cerebras Wafer-Scale Engine. The MLIR pipeline lowers high-level stencils through a chain of dialects to asynchronous wafer code automatically, matching or beating hand-optimized implementations and, on the WSE3 with no application-level changes, running about 14 times faster than 128 Nvidia A100 GPUs and 20 times faster than 128 nodes of a CPU-based Cray-EX supercomputer.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s full text (arXiv:2601.17754). The mechanism described is the paper’s: a compiler pipeline, built on MLIR and the compatible xDSL framework, that automatically transforms high-level stencil kernels into optimized code for the Cerebras Wafer-Scale Engine (WSE) without application-level changes, based on the hypothesis that domain-specific information about stencils lets the compiler target the wafer. The WSE is a dataflow architecture with a systolic array of up to about 900,000 processing elements, each with a compute core, router, program counter, clock, and 48 kilobytes of SRAM, operating under an asynchronous, actor-based execution model. The pipeline uses a high-level stencil dialect, a set of intermediate domain-specific dialects, and a low-level csl-ir dialect that re-implements a large subset of Cerebras’s CSL, with a strategy of lowering transformations generating the asynchronous exchanges, callbacks, and routing that bridge the semantic gap to the wafer’s execution model. It is evaluated on five benchmarks across three HPC programming technologies, at three problem sizes, on the WSE2 and WSE3, delivering comparable or higher performance than manually optimized code and, on the WSE3 with no application-level changes, around 14 times the performance of 128 Nvidia A100 GPUs and 20 times that of 128 nodes of a CPU-based Cray-EX supercomputer. Exact dialect definitions, the benchmarks and problem sizes, and full results should be read from the paper.

Background

Why exotic accelerators now live or die by their compilers

It helps to see this work as part of a broad shift in computing: as general-purpose processors stopped getting much faster, the industry produced a wave of radically specialized accelerators, and the hardest problem with each one is not building it but programming it. A machine like the Wafer-Scale Engine can deliver extraordinary performance, but its execution model is so far from how people write programs that, without a compiler to bridge the gap, only its creators can use it well. The recurring answer, across many such architectures, is to build layered compiler infrastructure — reusable intermediate representations and lowering passes — that lets high-level programs be translated, in structured steps, down to the peculiar hardware, so the accelerator becomes reachable from the software people already have. MLIR has become a common substrate for exactly this, precisely because its dialects let a compiler meet each new device partway. The broader effort to systematically compile and realize computations onto novel spatial and reconfigurable hardware, rather than hand-mapping each one, is the same undertaking the work on a systematic framework for heterogeneous coarse-grained reconfigurable array realization covered elsewhere on this site pursues, making an unusual fabric a compiler target rather than an artisanal one. This pipeline brings that philosophy to the largest chip of all, arguing that even a whole-wafer, fully asynchronous machine can be reached automatically from clean, high-level scientific code.

If you remember only three things

1. The Cerebras Wafer-Scale Engine is a single-wafer chip with over 900,000 independent processing elements, superb for scientific computing, but it runs a distributed, asynchronous, actor-based model — processing elements reacting to messages from neighbors, with no global control and, in its language, not even a normal top-level loop. That is nothing like the sequential or bulk-synchronous code scientists write, so porting an existing simulation meant a laborious bespoke rewrite by hand, and mainstream compilers like MLIR had no wafer support to automate it.
2. This paper builds an MLIR-based compiler pipeline (using the compatible xDSL framework) that automatically lowers high-level stencil kernels — where each grid point updates from its neighbors — into optimized wafer code, with no application-level changes. It crosses the huge semantic gap in steps: a high-level stencil dialect, intermediate domain-specific dialects, and a low-level csl-ir dialect that re-implements a large subset of Cerebras's CSL, with lowering transformations generating the asynchronous exchanges, callbacks, and routing the wafer needs. It exploits the regular locality of stencils to distribute the grid and place the neighbor communication.
3. The generated code matches or beats hand-optimized implementations, and on the WSE3 with no application-level changes it runs about 14 times faster than 128 Nvidia A100 GPUs and 20 times faster than 128 nodes of a CPU-based Cray-EX supercomputer, evaluated across five benchmarks and three problem sizes on the WSE2 and WSE3. It is specific to stencils, re-implements a subset (not all) of CSL, and is tied to the Cerebras wafer, but it turns a machine only experts could program into an automatic compiler target for existing scientific code.

Check yourself

Check your understanding

1. Why is the Cerebras Wafer-Scale Engine so hard to program despite being ideal for scientific computing?
2. What is the core approach of the compiler pipeline in this paper?
3. Why does the pipeline rely on the computation being a stencil?
4. Predict: a scientist has an existing finite-difference simulation and wants wafer performance. Using this pipeline, what do they change, and what performance can they expect?

Sources

The paper is “An MLIR Lowering Pipeline for Stencils at Wafer-Scale” by Nicolai Stawinoga, David Katz, Anton Lydike, Justs Zarins, Nick Brown, George Bisbas, and Tobias Grosser (Technische Universitat Berlin, EPCC and the University of Edinburgh, and the University of Cambridge), published at ASPLOS 2026, DOI 10.1145/3779212.3790124. It is licensed CC BY.

This explainer is written from the paper’s full text (arXiv:2601.17754). The Cerebras Wafer-Scale Engine (WSE) delivers performance at an unprecedented scale of over 900,000 compute units on a single wafer, and although designed for AI is also well suited to high-performance computing; but its distributed, asynchronous, actor-based programming model diverges sharply from the simple sequential or bulk-synchronous programs derived from a mathematical description, so targeting it required a bespoke re-implementation, and the absence of WSE support in compilers such as MLIR left no path to automate the process. The paper presents a compiler pipeline, based on MLIR and the compatible xDSL framework, that transforms high-level stencil kernels into optimized wafer code without application-level changes, on the hypothesis that domain-specific information about stencils lets the compiler target the WSE. It uses a high-level stencil dialect, intermediate domain-specific dialects, and a low-level csl-ir dialect that re-implements a large subset of Cerebras’s CSL, with a strategy of lowering transformations that generate the asynchronous exchanges, callbacks, and routing bridging the semantic gap to the wafer’s execution model. Evaluated on five benchmarks across three HPC programming technologies at three problem sizes, on the WSE2 and WSE3, it delivers comparable or higher performance than manually optimized code, and on the WSE3 with no application-level changes runs around 14 times faster than 128 Nvidia A100 GPUs and 20 times faster than 128 nodes of a CPU-based Cray-EX supercomputer. No numbers beyond the paper’s own have been introduced; exact dialect definitions, benchmarks, and full results should be confirmed in the paper.