Stencil Kernels on Wafer-Scale Architectures

• The paper demonstrates efficient integration of stencil kernels with wafer-scale architectures by leveraging on-chip SRAM and deterministic mesh routing for near-linear scalability.
• It details innovative algorithmic mapping and multi-phase halo exchange protocols that align data partitioning with local processing constraints.
• Compiler pipelines using high-level DSLs and auto-tuning techniques substantially reduce code complexity while matching or exceeding traditional performance benchmarks.

Stencil kernels are iterative numerical algorithms in which each element of a multidimensional data grid is updated as a function of a weighted pattern of neighboring points. This computational motif arises pervasively in the discretization of partial differential equations for scientific simulation, including seismic imaging, fluid dynamics, and climate modeling. Wafer-scale architectures such as the Cerebras Wafer-Scale Engine (WSE) integrate hundreds of thousands to millions of processing elements (PEs) with a flat, high-bandwidth, on-chip memory system and specialized mesh communication fabric. These properties make them uniquely suited to high-performance stencil computations, by transforming the memory wall that constrains traditional cache-based platforms into a communication-bound or compute-bound regime, permitting nearly ideal scaling for regular stencils. Recent research details both architectural properties and algorithmic transformations required to exploit these machines, the resulting performance characteristics, domain-specific compilation strategies, and productivity implications for scientific computation (Jacquelin et al., 2022, Stawinoga et al., 25 Jan 2026, Brown et al., 2022, Rocki et al., 2020, Belli et al., 8 May 2026, Sai et al., 2023).

1. Wafer-Scale Architecture for Stencil Kernels

The Cerebras WSE series (CS-1, WSE-2, WSE-3) comprises an entire silicon wafer with 350,000–2.6 million PEs, each equipped with local SRAM (e.g., 48 KB/PE) and interconnected by a 2D mesh network with deterministic, low-latency routing. Critically, no off-chip DRAM or cache hierarchy is present; instead, all working data reside in on-chip SRAM, eliminating DRAM latency and bandwidth bottlenecks inherent to conventional architectures (Jacquelin et al., 2022, Belli et al., 8 May 2026). Each PE communicates directly with its nearest neighbors via single-cycle "wavelets" or router messages, supporting both synchronous and asynchronous communication models (Stawinoga et al., 25 Jan 2026, Rocki et al., 2020). This mesh fabric sustains aggregate bandwidths up to 215 PB/s (WSE-3), with scalable channel multiplexing (24 "colors") for minimal congestion (Belli et al., 8 May 2026). The distributed memory design, with explicit PE-local state and deterministic routing, aligns naturally with the spatial locality of stencil algorithms.

2. Algorithmic Mapping: Stencil Formulation and Data Partitioning

Stencil kernels are discretized as regular patterns—for example, a 25-point 3D wave-propagation stencil:

$$u_{i,j,k}^{n+1} = \alpha\, u_{i,j,k}^n + \beta \sum_{(d_x, d_y, d_z)\in N} u_{i+d_x, j+d_y, k+d_z}^n$$

where the neighbor set $N$ encodes all offsets relevant to the stencil (axes and diagonals for 25-point). The global computational grid is partitioned across the PE array, with each PE assigned a contiguous subdomain and its associated halo regions—typically the width equals the maximum stencil radius. For the 25-point wave stencil, halos of width 2 on all faces are required (Jacquelin et al., 2022, Belli et al., 8 May 2026, Sai et al., 2023). Subdomains are often assigned such that the major axes (X, Y) tile the 2D wafer grid, and the depth (Z) dimension is stretched within each PE, matching architectural SIMD width or buffer constraints (Stawinoga et al., 25 Jan 2026, Sai et al., 2023). Efficient mapping demands careful alignment of subdomain and halo layout with PE-local SRAM capacity and the hardware’s routing granularity.

3. Communication Strategies: Halo Exchange and Overlap

On cache-based systems, temporal and spatial tiling attempt to maximize cache reuse; on wafer-scale architectures, traditional inter-cache bandwidth is replaced by explicit, asynchronous halo exchange across the mesh (Jacquelin et al., 2022, Belli et al., 8 May 2026). For each time step, PEs exchange boundary data ("halos") with their six (3D, axis-aligned) or four (2D) neighbors. For higher-order stencils (e.g., 25-point 3D or 9-point 2D), a multi-phase protocol, sometimes including edge or corner forwarding via store-and-forward, is employed (Jacquelin et al., 2022, Belli et al., 8 May 2026). Halo data are packed into router-aligned buffers (DSD/DSR based streaming), and transfer is initiated by fabric commands; arrival of all required halo points triggers local stencil computation (Sai et al., 2023, Belli et al., 8 May 2026). Hardware support for pipelining (multiple microthreads per PE) and asynchronous event callbacks allows compute to overlap with communication, hiding most transfer latency. Double-buffering and chunk-based streaming further amortize synchronization delays (Stawinoga et al., 25 Jan 2026).

4. Programming Models and Compiler Pipelines

The programming abstraction for stencil kernels on wafer-scale systems has evolved from low-level domain-specific languages (CSL) and explicit actor/task orchestration to high-productivity DSLs and automated compilation. In bare CSL, developers write explicit control-state machines for each PE, orchestrating state transitions, receive/send, and vectorized compute (Jacquelin et al., 2022, Sai et al., 2023). Recent work introduces high-level domain-specific languages (e.g., StencilPy, TensorFlow as DSL, MLIR-based toolchains) that parse symbolic stencil definitions and automatically generate in-place, wafer-optimized code (Brown et al., 2022, Sai et al., 2023, Stawinoga et al., 25 Jan 2026). These pipelines perform IR lowering, tiling/blocking analysis, halo pattern inference, and backend-specific code generation (e.g., actor tasks, router configuration, DSD scheduling). Automated fusion of communication and compute, minimization of task and buffer footprint, and cross-platform backend support (e.g., CUDA, HIP, CSL) have resulted in tool-generated kernels matching or exceeding the performance of prior hand-optimized codes for nontrivial stencils (Sai et al., 2023, Stawinoga et al., 25 Jan 2026).

5. Performance Results and Comparative Analysis

Empirical measurements demonstrate wafer-scale architectures' capability to sustain performance unreachable by comparable traditional platforms. For a 25-point 3D seismic kernel, the WSE-2 achieves 503 TFLOPs, more than twice the throughput of an 8-node (64×V100) GPU cluster, and only matched by full-scale supercomputing installations (Jacquelin et al., 2022). In WSE-3, the CStencil framework exhibits >90% utilization of compute peak and delivers speedups up to 342× over a single-precision port of ConvStencil on NVIDIA A100, for weakly scaled 2D heat-diffusion benchmarks (Belli et al., 8 May 2026). Roofline analysis shows that, as arithmetic intensity increases, the computational roof is reached rather than memory bound. Communication overheads are marginal (≪1% per iteration on WSE-3, <5% on WSE-2 until fabric saturation), and scaling is near-linear for domain sizes utilizing over 250,000 active PEs (Belli et al., 8 May 2026, Jacquelin et al., 2022). High-productivity pipelines (MLIR, StencilPy) enable 14× and 20× speedup over 128×A100 GPU and Cray-EX CPU supercomputer partitions, respectively, with no change to scientific source code (Stawinoga et al., 25 Jan 2026).

Platform	Peak Performance	Delivered Stencil Perf.	Relative Speedup (WSE)
Cerebras WSE-2/3	11.7 PFLOPS (theor)	503–678 TFLOPS	2–14× 128×A100, 20× 128×CPU
NVIDIA A100 (128 GPUs)	156 TF32 TFLOPS	~19.5 TFLOPS (ConvStencil)	1/14×–1/342× (vs WSE-3)
Cray-EX (128 nodes)	(varies)	(see above)	1/20× (vs WSE-3, acoustic)
CS-1/WSE-2 (StencilPy)	11.7 PFLOPS	232 TFLOPS (25-point)	≈2% peak (auto-gen code)

6. Portability, Productivity, and Tuning Considerations

StencilPy and MLIR-based flows yield substantial reductions in code size (from 1600 + 360 LOC for hand-written CSL+Python to ~70 LOC in DSL) and provide platform-agnostic output (CUDA, HIP, CSL, STX backends) (Sai et al., 2023). Compiler transformations expose auto-tuning opportunities (e.g., spatial block sizes, streaming/3D tiling, register/SRAM buffer allocation), and domain-specific fusion passes optimize communication/computation overlap. However, such architectures impose limits tied to local SRAM (e.g., 48 KB/PE), which restricts maximum tile size or achievable stencil radius; maximizing PE utilization and fabric bandwidth demands problem decomposition matched to hardware geometry. While high utilization is achievable for regular, memory-bound stencils, codes requiring double precision or complex irregular communication (e.g., adaptive meshes) may see reduced efficiency or require further architectural features (Belli et al., 8 May 2026, Rocki et al., 2020).

7. Challenges, Generalization, and Future Directions

Critical ingredients for wafer-scale stencil efficiency include: flat fully on-chip SRAM, low-latency/ high-radix mesh fabrics, and programming models that provide explicit, orchestrated halo exchange rather than relying on caches. Porting to future (or custom) wafer-scale engines such as Anton 3 for molecular dynamics is plausible where these features persist (Jacquelin et al., 2022). Anticipated obstacles include SRAM-limited problem decomposition, challenging network contention at extreme scales (mitigated via communication scheduling), and the need for software stacks capable of managing fine-grained actor/task semantics. The field continues to develop higher-order compiler frameworks, flexible map/reduce abstractions for non-canonical stencils, and integration of dynamic adaptation strategies. A plausible implication is that widespread adoption of high-productivity, cross-architecture DSLs for stencil computation will be catalyzed by continued improvements in automated communication orchestration and backend portability for emerging wafer-scale and dataflow architectures (Sai et al., 2023, Stawinoga et al., 25 Jan 2026).