面积级机器学习中间表示 (MLIR) 下降管道用于晶圆级阵列
分布式、并行与集群计算
2026-01-27 v1 新兴技术
编程语言
摘要
Cerebras 面积级发动机 (Wafer-Scale Engine, WSE) 交付了超过 90 万个计算单元的性能，这些计算单元通过单片晶圆芯片互连全部连接。虽然最初为 AI 设计，WSE 架构也非常适合高性能计算 (HPC)。然而，其分布式异步编程模型与通常从数学程序描述派生出的简单顺序或批量同步程序有显著差异。针对 WSE 需要针对性地重新实现以移植现有代码。由于 MLIR 等编译器缺乏对 WSE 的支持，自动化这一过程几乎不可能。由于 Stencil 在 HPC 中十分常见，本文探讨了利用 Stencil 的领域特定信息由编译器自动针对 WSE 而无需应用程序级别代码更改的假设。我们提出了将基于 Stencil 的内核转换为为 WSE 优化的 CSL 代码的编译管道，弥合了问题数学表示与 WSE 异步执行模型之间的语义差距。基于 5 个基准测试，分别在 Cerebras WSE2 和 WSE3 上运行，我们的方法在性能上与手动优化的代码相当，甚至略好。此外，由于不需要任何应用程序级别的代码更改，WSE3 的性能比 128 个 Nvidia A100 GPU 快约 14 倍，比 128 节点的基于 CPU 的 Cray-EX 超级计算机快约 20 倍。
引用
@article{arxiv.2601.17754,
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},
journal= {arXiv preprint arXiv:2601.17754},
year = {2026}
}
备注
Paper in ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS '26)