Fast Stencil-Code Computation on a Wafer-Scale Processor

TL;DR

This paper demonstrates that a wafer-scale processor can efficiently solve large PDE linear systems, achieving significant performance improvements over traditional CPU and GPU systems.

Contribution

It introduces the use of a wafer-scale processor for PDE solutions, showcasing high performance and discussing system architecture, programming, and potential for full applications.

Findings

Achieved 0.86 PFLOPS on a single wafer-scale system.

Solved a large linear system with about one third of the machine's peak performance.

Discussed memory, precision, and future application extensions.

Abstract

The performance of CPU-based and GPU-based systems is often low for PDE codes, where large, sparse, and often structured systems of linear equations must be solved. Iterative solvers are limited by data movement, both between caches and memory and between nodes. Here we describe the solution of such systems of equations on the Cerebras Systems CS-1, a wafer-scale processor that has the memory bandwidth and communication latency to perform well. We achieve 0.86 PFLOPS on a single wafer-scale system for the solution by BiCGStab of a linear system arising from a 7-point finite difference stencil on a 600 X 595 X 1536 mesh, achieving about one third of the machine's peak performance. We explain the system, its architecture and programming, and its performance on this problem and related problems. We discuss issues of memory capacity and floating point precision. We outline plans to extend this work towards full applications.