一句话:提出 MLIR lowering 流水线,将 stencil 内核自动编译为 Cerebras WSE 的 CSL 代码,性能媲美手工优化。
论文针对 Cerebras 晶圆级引擎(WSE)缺乏 MLIR 等编译器支持、其分布式异步编程模型与常规顺序或同步 HPC 程序语义差距大的问题,探索利用 stencil 领域信息自动面向 WSE 编译,而无需修改应用代码。作者设计了一条基于 MLIR 的 lowering 流水线,将 stencil 内核变换为高度优化的 CSL 代码,弥合数学表示与 WSE 异步执行模型之间的语义鸿沟。基于三类 HPC 编程技术、五个基准,在 WSE2 与 WSE3 上评估;性能与手工优化代码相当或略优,且在 WSE3 上分别比 128 张 Nvidia A100 GPU 和 128 节点 Cray-EX CPU 超级计算机快约 14 倍和 20 倍。
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.
DOI 原文 · arXiv · PDF(开放获取) · DBLP
@inproceedings{StawinogaKLZ0BG26,
title = {An MLIR Lowering Pipeline for Stencils at Wafer-Scale},
author = {Nicolai Stawinoga and David Katz and Anton Lydike and Justs Zarins and Nick Brown and George Bisbas and Tobias Grosser},
booktitle = {ASPLOS},
year = {2026},
doi = {10.1145/3779212.3790124}
}