A performance portable fast Ewald summation for Stokes flow

2026-06-17

AI Summary

This work proposes a high-performance, portable algorithm for Ewald summation in the N-body problem of Stokes flow within periodic domains, featuring a synergistic computation of near-field (P2P) and far-field (P2G/G2P+FFT) interactions. Built upon the PyKokkos-wrapped Kokkos parallel programming model, the method supports diverse hardware architectures including AMD and NVIDIA GPUs as well as ARM and x86 CPUs. A novel P2G algorithm enables processing up to 8 million particles per second with 9-digit accuracy on an NVIDIA H200, achieving up to 16× speedup over baseline GPU implementations. Experimental results demonstrate 84% computational efficiency for P2P interactions and show that communication overhead remains manageable under weak scaling to 256 million particles, effectively alleviating the performance bottlenecks between near- and far-field computations across multiple hardware platforms.

Abstract

We present GPU algorithms for Ewald summation methods for accelerating N-body Stokes flow problems in periodic domains. Like most N-body codes, Ewald sums use a near-field/far-field decomposition. The near field involves particle-to-particle (P2P) interactions. The far field primarily involves particle-to-grid (P2G) and grid-to-particle (G2P) interactions, as well as Fast Fourier Transforms. For each interaction, we investigate several algorithmic variants. Our implementation uses PyKokkos, a Python interface for the Kokkos C++ parallel programming framework, which supports portability to AMD/NVIDIA GPU and ARM/x86 CPU architectures. Double and single-precision numerical results, alongside analytical performance models, confirm the efficiency of our algorithms on AMD and NVIDIA GPU and on ARM and AMD CPU architectures. The P2P interaction achieves around 73% compute efficiency on NVIDIA H200, 84% on NVIDIA A100, 60% on AMD MI300, 52% on Grace CPU, and 68% on AMD Epyc CPU. A straightforward implementation of the P2G kernel can become a computational bottleneck. We introduce a novel P2G algorithm that achieves up to 16$\times$ speedup compared to a baseline GPU implementation. The overall Ewald sum code processes approximately 8 million particles per second on a H200 GPU, and about a half-million particles per second on a Grace CPU, for nine digits of accuracy. We also perform a multi-GPU weak scaling test on up to 256 million particles (64 GPUs) that shows bounded communication cost for all stages except the all-to-all particle sorting, which can be reduced to neighbor communication in the relevant time-stepping regime.

Problem

• Ewald summation
• Stokes flow
• performance portability
• N-body problem
• GPU acceleration

Innovation

• Ewald summation
• performance portability
• GPU acceleration
• Stokes flow
• particle-in-cell