A Portable Framework for Accelerating Stencil Computations on Modern Node Architectures

TL;DR

StencilPy is a high-level framework that simplifies writing portable, high-performance stencil computations across diverse modern architectures, achieving performance comparable to hand-optimized code with significantly improved productivity.

Contribution

It introduces StencilPy, a domain-specific language and backend system that enables efficient, portable stencil computations on a wide range of current and emerging hardware architectures.

Findings

StencilPy achieves performance comparable to hand-written code.

It reduces code length significantly compared to manual implementations.

Provides cross-architecture portability and high productivity.

Abstract

Finite-difference methods based on high-order stencils are widely used in seismic simulations, weather forecasting, computational fluid dynamics, and other scientific applications. Achieving HPC-level stencil computations on one architecture is challenging, porting to other architectures without sacrificing performance requires significant effort, especially in this golden age of many distinctive architectures. To help developers achieve performance, portability, and productivity with stencil computations, we developed StencilPy. With StencilPy, developers write stencil computations in a high-level domain-specific language, which promotes productivity, while its backends generate efficient code for existing and emerging architectures, including modern many-core CPUs (such as AMD Genoa-X, Fujitsu A64FX, and Intel Sapphire Rapids), latest generations of GPUs (including NVIDIA H100 and A100, AMD MI200, and Intel Ponte Vecchio), and accelerators (including Cerebras and STX). StencilPy demonstrates promising performance results on par with hand-written code, maintains cross-architectural performance portability, and enhances productivity. Its modular design enables easy configuration, customization, and extension. A 25-point star-shaped stencil written in StencilPy is one-quarter of the length of a hand-crafted CUDA code and achieves similar performance on an NVIDIA H100 GPU. In addition, the same kernel written using our tool is 7x shorter than hand-optimized code written in Cerebras Software Language (CSL), and it delivers comparable performance that code on a Cerebras CS-2.