TensorFlow as a DSL for stencil-based computation on the Cerebras Wafer Scale Engine

TL;DR

This paper evaluates the Cerebras Wafer Scale Engine's potential for accelerating stencil-based HPC codes using TensorFlow, demonstrating significant performance advantages over GPUs and CPUs.

Contribution

It introduces TensorFlow as a programming interface for the Cerebras WSE and assesses its performance for stencil computations, highlighting its potential for HPC acceleration.

Findings

WSE outperforms 4 V100 GPUs by 2.5x

WSE is approximately 114x faster than two Intel Xeon CPUs

TensorFlow can effectively interface with WSE for HPC workloads

Abstract

The Cerebras Wafer Scale Engine (WSE) is an accelerator that combines hundreds of thousands of AI-cores onto a single chip. Whilst this technology has been designed for machine learning workloads, the significant amount of available raw compute means that it is also a very interesting potential target for accelerating traditional HPC computational codes. Many of these algorithms are stencil-based, where update operations involve contributions from neighbouring elements, and in this paper we explore the suitability of this technology for such codes from the perspective of an early adopter of the technology, compared to CPUs and GPUs. Using TensorFlow as the interface, we explore the performance and demonstrate that, whilst there is still work to be done around exposing the programming interface to users, performance of the WSE is impressive as it out performs four V100 GPUs by two and a half times and two Intel Xeon Platinum CPUs by around 114 times in our experiments. There is significant potential therefore for this technology to play an important role in accelerating HPC codes on future exascale supercomputers.