Stencil Computations on Cerebras Wafer-Scale Engine

TL;DR

This paper demonstrates that the Cerebras WSE-3 can efficiently run 2D stencil computations, achieving significant speedups over GPU implementations by leveraging its on-chip memory and parallel architecture.

Contribution

It introduces CStencil, a framework for stencil computations on WSE-3, and shows it outperforms GPU-based solvers, bridging AI hardware and scientific computing.

Findings

CStencil achieves up to 342x speedup over GPU-based solver.

WSE-3's on-chip memory eliminates off-chip bottlenecks.

The architecture effectively saturates compute and memory resources.

Abstract

Stencil computations are a fundamental kernel in scientific computing, critical for simulations in domains such as fluid dynamics and climate modeling. However, these computations are often memory-bound on traditional High-Performance Computing architectures like GPUs, struggling against the "Memory Wall". Simultaneously, the rise of AI-oriented hardware, such as the Cerebras Wafer-Scale Engine, offers massive core parallelism and high-bandwidth on-chip memory, though typically optimized for lower-precision workloads. This work investigates the viability of bridging this divergence by mapping stencil algorithms onto the Cerebras WSE-3. The study introduces CStencil, a novel framework designed to implement two-dimensional stencil computations on the WSE-3. To ensure a rigorous and fair performance evaluation, the research also adapts ConvStencil, a state-of-the-art GPU stencil solver, porting it from its original double-precision design to single-precision for execution on an NVIDIA A100 GPU. Experimental results show that the WSE-3's distributed SRAM and mesh interconnect effectively eliminate the off-chip memory bottlenecks common in GPU implementations. CStencil achieves speedups of up to 342x over the adapted ConvStencil version. A roofline model analysis further confirms that CStencil saturates the available compute and memory resources, demonstrating that the WSE dataflow architecture can be successfully repurposed for traditional scientific algorithms. These findings highlight the potential of the WSE-3 to deliver hardware utilization levels unattainable on conventional systems, offering a promising path toward overcoming the memory limitations of current HPC architectures.