arXiv CSopen access

A performance portable fast Ewald summation for Stokes flow

A PERFORMANCE PORTABLE FAST EWALD SUMMATION FOR STOKES FLOW

arXiv:2606.19059v1 [math.NA] 17 Jun 2026

GABRIEL KOSMACHER∗ , ZIYU DU† , JOAR BAGGE† , AND GEORGE BIROS† 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/farfield 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× 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. Key words. N-body problem, Ewald summation, fast Fourier transform, Kokkos, GPU, performance portability, parallel algorithms MSC codes. 31-04, 33-04, 35-04, 45-04, 65R20, 65Y20, 76D07

1. Introduction. Stokes suspensions involve the interaction of viscous flows with rigid or deformable particles [51]. They appear in microfluidic devices [17], swimming microorganisms [13, 20], capillary blood flow [42], antibodies [21], and industrial emulsions [53]. Boundary integral equations are a commonly used numerical method for simulating Stokes suspensions [32]. Upon discretization, the velocity field (henceforth referred to as potential per N-body literature conventions) is written as a sum over Ns source points y j , of the form (1.1)

u(xi ) =

Ns X j=1

G(xi − y j )f (y j ),

i = 1, . . . , Nt ,

where G is a Stokes Green’s function and f (y j ) is a known force density (or simply density) associated with y j . The sum in (1.1) needs to be evaluated at Nt target points xi ; assuming N = Nt = Ns , direct evaluation costs O(N 2 ). N-body methods reduce this cost to O(N ) for the fast multipole method (FMM) [50, 12, 22] and to O(N log N ) for an Ewald method [23]. Although the FMM has favorable asymptotic complexity, the prefactor constants can be quite high compared to the Ewald method when the spatial particle distribution is nearly uniform [11]. As a result, Ewald sum methods are widely used for dense suspensions with periodic boundary conditions [4, 49, 2, 52]. Yet, as far as we are aware, there is no existing open source GPU implementation for a Stokes Ewald method. ∗ Oden Institute for Computational Engineering and Sciences, The University of Texas, Austin,

TX ({gkosmacher, ziyu du}@utexas.edu, [email protected], [email protected]). 1

2

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Contributions: Following the Stokes Ewald algorithm introduced in [5] (see section 2), we present the following methodological and experimental results: • The design and implementation of GPU algorithms for different steps in Ewald sums; in particular, our novel particle-to-grid (P2G) method achieves a speedup of more than 16× compared to a baseline implementation (see section 3 and subsection 3.2). • A novel rational function approximation for erf(x)/x that appears in Ewald sums; it is faster than system libraries in lower precision (see section 4). • A roofline analysis for our kernels with various assumptions that aligns with the measured times, enabling automatic selection of algorithm parameters (see section 5). • Evaluation on multiple platforms, including NVIDIA H200 and A100 GPUs, AMD MI300A GPU, AMD Epyc CPU, and NVIDIA Grace ARM CPU (see section 6). • An open source Python library ParkiPy1 (supplement A), developed using PyKokkos, supporting single node APIs for Stokes and Poisson potentials in single and double precision and a distributed API for the Stokes potential in double precision. Using standard techniques, we also present results using the message-passing interface (MPI) on up to 64 GPUs. Related work. Ewald methods, also known as smooth particle-mesh Ewald sums, have been extensively studied for the Poisson problem for molecular dynamics simulations. The Ewald decomposition for Stokes was introduced by Hasimoto [16]. As mentioned, our summation algorithm is based on the fast spectrally accurate Ewald summation method in [5], which is a refinement of the scheme introduced in [23]. Spectral Ewald summations are similar to smoothed particle mesh methods but more accurate [23, 52, 37]. The development of parallel Ewald sums has been largely focused on Poisson kernels for molecular dynamics [15, 31, 39, 38, 36]. Examples of work on GPUs include GROMACS [33], NAMD [30], AMBER [35] and many others [18]. The primary objective in most studies is to assess the overall performance of an atomistic simulation in terms of nanoseconds per day. Surprisingly, detailed discussions and performance evaluations for the Ewald sum and its components are rare compared to the richer FMM literature. For Stokes flows, implementations of GPU-accelerated Ewald methods are mentioned in Wang and Brady [49] and Fiore et al. [10]. These are restricted to smaller systems for Brownian dynamics (fewer than 10K particles), with no assessment of the Ewald summation component. The work on GPU nonuniform fast Fourier transforms [41] is one of the few exceptions that examines different compute steps that closely resemble parts of Ewald sum. In terms of portability, frameworks like OpenCL [43], Kokkos [9, 46], and RAJA [7] have been developed to support code execution on CPUs and GPUs. Kokkos has been used to develop Poisson Ewald codes in LAMMPS [44, 14]. Another portable implementation is from Mayani et al. [25], who implemented a free-space spectral Poisson solver via Kokkos. PyKokkos [3, 45] enables programming Kokkos in Python with minimal overheads using just-in-time compiling.

1 https://github.com/ut-padas/parki

3

FAST EWALD SUMMATION FOR STOKES FLOW

Fig. 1. Two examples of Stokesian flows. Left: A simulation from ongoing work in our group showing deformable particles (red blood cells) flowing in a micrometer-wide pipe. Right: A suspension of a large number of rigid particles (slender rods) in a pipe, taken from [4]. Both problems are periodic in the direction parallel to the axis of the pipe. Each time step involves particle-related calculations followed by a Stokes Ewald sum.

2. Sequential Ewald Summation. We summarize the Ewald sum scheme from [5]. The periodized version of (1.1) reads (2.1)

u(xi ) =

Ns X X j=1 p∈P

G(xi − y j + p)f (y j ),

i = 1, . . . , Nt ,

where P denotes a set of periodic images that tile Rd to create a periodic field. Depending on the problem, we can have periodicity in one, two, or all three directions. Here for simplicity we only consider the 1-periodic case, P = {(L1 α, 0, 0) : α ∈ Z}, with the other directions being free; Ewald sums of this case can be used to accelerate the solution of confined Stokes solvers in pipe geometries, such as those in Figure 1. The extension to the 2- and 3-periodic cases is straightforward and does not require any changes to the kernels discussed in section 3. We assume {xi }, {y j } ⊂ B := [0, L1 ) × [0, L2 ) × [0, L3 ), and that B is periodically replicated along the x1 axis. The r⊗r I + ∥r∥ Green’s function G is the Stokeslet given by G(r) = ∥r∥ 3 , where I is the 3-by-3 identity matrix. Note that if ∥r∥ ≈ 0, G is singular; if y j ≈ xi , the term is skipped. Ewald summation splits u = uN + uF by setting G = GN + GF . The near-field component GN rapidly decays with ∥r∥; the far-field component GF is smooth and slowly decays with ∥r∥. A split parameter ξ > 0 controls the decay of GN . 2.1. GN : Near-Field (P2P). The near-field sum (or P2P interaction) is given

by N

(2.2)

u (xi ) =

Ns X X j=1 p∈P

GN (xi − y j + p) · f (y j ),

and the expression for the near-field Stokeslet GN is 2

(2.3)

∥r∥2ξe−ξ ∥r∥ √ G (r) = G(r) erfc(ξ∥r∥) + π N

2

!

2

−I

2

4ξe−ξ ∥r∥ √ , π

where erfc(x) = 1 − erf(x) is the complementary error function. Let the cutoff radius rc ∼ ξ −1 be a function of the split parameter ξ whose form is derived via (2.9). GN can then be efficiently computed by dropping any interaction with ∥r∥ > rc . To exploit this locality in (2.3), we spatially partition B into uniform cubic cells of width rc and then assign each source and target point to a cell. When evaluating (2.2) on a target xi , it is enough to determine which cell xi belongs to and then restrict its P2P interactions to the 27-cell neighborhood of that cell. To control the near-field error, we need to choose ξ appropriately (see subsection 2.4).

4

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

2.2. GF : Far-Field Computation (P2G, G2P, and FFTs). The far field is smooth and is defined in Fourier space: uF (xi ) =

(2.4a)

Ns X j=1

(2.4b)

b F · f (y )}(xi − y ), F1−1 {G j j

b F (κ) = 8π I∥κ∥2 − κ ⊗ κ e−∥κ∥2 /(2ξ)2 G ∥κ∥4 F1−1 {g}(r) =

(2.4c)

X 1 2 L1 (2π)

Z

κ1 ∈K



∥κ∥2 1+ (2ξ)2

 ,

g(κ)ei(κ·r) dκ2 dκ3 ,

R2

where κ = (κ1 , κ2 , κ3 ), K = {2πα/L1 : α ∈ Z}, F1−1 is an inverse Fourier transform b F the (series in the periodic direction and continuous in the free directions) and G b F grows with ξ, hence GF has Fourier transform of GF = G − GN . The support of G slower spectral decay as ξ increases. For κ1 = 0, the integral in (2.4c) is singular and is computed using the method by Vico et al. [47]. We omit details here and refer to [5, Section 3.1]. A direct solve of Equation (2.4) has O(N 2 ) complexity as the sources and targets are non-uniformly distributed in B, hence fast Fourier transforms cannot be used. This can be accelerated with a nonuniform FFT [6]: spread points on a regular grid, use FFTs to convolve with GF in spectral space, and interpolate back. Let us assume that the regular grid has grid spacing h. The far-field potential uF is then computed using the following five computational kernels: (1) Particle-to-grid (P2G): Spread the source strengths f (y j ) to the grid points Gh := {g ℓ } of the uniform grid by convolving with a compactly supported window function w(r): (2.5)

ϕh (g ℓ ) :=

Ns X X j=1 p∈P

w(g ℓ − y j + p)f (y j ),

ℓ = 1, . . . , Ng ,

where Ng is the number of grid points and P is as in (2.1). b (κℓ )} = FFT3D {ϕ (g )}, where {κℓ } are the grid points of the (2) FFT: {ϕ h h ℓ uniform grid in the frequency domain. (3) Convolution with GF (CNV): Scale the values in Fourier space by the b F and the Fourier transform w function G b of the window function: (2.6)

bh (κℓ ) := v

b F (κℓ ) G b (κℓ ). ϕ [w(κ b ℓ )]2 h

(4) Inverse FFT: {v h (g ℓ )} = IFFT3D {b v h (κℓ )}. (5) Grid-to-particle (G2P): Interpolate back to target points: (2.7)

3 uF h (xi ) := h

Ng X X ℓ=1 p∈P

w(xi − g ℓ + p)v h (g ℓ ).

FAST EWALD SUMMATION FOR STOKES FLOW

5

The full potential is then approximated by uh (xi ) := uN (xi ) + uF h (xi ). In our discussion, we refer to the combination of the FFT, CNV, and IFFT steps as the Fourier grid convolution (FGC) step. The window function in the P2G and G2P steps is w(r) = w0 (r1 )w0 (r2 )w0 (r3 ), where w0 is the truncated Kaiser–Bessel (KB) function (  √ 2  I0 β 1−r /a2w /I0 (β) if r < a , w (2.8) w0 (r) = 0 otherwise, and I0 is the modified Bessel function. Since I0 is expensive to evaluate (see Table 1), w0 approximated by a polynomial interpolation of degree ν. If P denotes the number of grid points in the support of w0 , then we set P = 2ν + 1, β = 2.5P , aw = P h/2. The truncated KB function was first used for the spectral Ewald method in [34] as an alternative to the classic Gaussian window function [24] with a smaller support needed for the same accuracy. It has recently been shown that window functions based of direct polynomial approximation of prolates [19, Appendix A.6] has slightly smaller support than the truncated KB window; we leave such an implementation to future work. Computing (2.4a) also requires extending the box B and the uniform grid in the free directions to approximate the free space conditions. This is technical and not essential to our discussion; see [5, Section 4.5] for details and a justification for the choices of P , β and aw . Since β and aw are tied to P , the latter is the main window parameter, and it is set according to the accuracy requirements. We summarize the entire spectral Ewald algorithm in Algorithm 2.1. In blue is the complexity of the different steps, with s being the average number of points per cell. Ng is the total number of grid points in the uniform grid, Ng ∼ h−3 ; P is the window parameter related to the P2G and G2P interactions. Algorithm 2.1 Spectral Ewald method and its complexity Input: B = [0, L1 ) × [0, L2 ) × [0, L3 ); Ns source positions y j ∈ B and densities f (y j ); Nt target positions xi ∈ B; user-speified error tolerance ε. Output: Potential approximation uh (xi ) at target positions 1: (P2P) Evaluate uN (xi ) using (2.2) O(3d sNt ) 2: (P2G) f (y j ) → ϕh (g ℓ ) using (2.5) O(Ns P d ) b (κℓ )} 3: (FFT) {ϕh (g ℓ )} → {ϕ O(Ng log Ng ) h b bh (κℓ ) using (2.6) 4: (CNV) ϕh (κℓ ) → v O(Ng ) 5: (IFFT) {b v h (κℓ )} → {v h (g ℓ )} O(Ng log Ng ) 6: (G2P) {v h (g ℓ ) → uF O(Nt P d ) h (xi )} using (2.7) N F 7: (SUM) Set uh (xi ) = u (xi ) + uh (xi ) O(Nt ) In the P2P kernel, 27 is the number of colleagues of a cell — the 3d cells that share a face, edge, or point — so each target interacts with 27s sources in its near field. The P 3 factor in P2G is the cost of spreading one source density to the grid points in the support of the window function. 2.3. Combined Single and Double Layers. In many suspension problems with deformable particles, the Stokes potential is generated not only by Stokeslet sources (single layer densities), but also by stresslet sources (double layer densities). That is, every source point then has an additional two vectors q(y j ) and n(y j ), and
PNs P an additional term j=1 p∈P T (xi − y j + p) q(y j ) ⊗ n(y j ) gets added to the

6

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

right-hand side of (2.1). The stresslet T is a symmetric 3-by-3-by-3 tensor and is split again as T = T N + T F , where the precise form of T N and T F is in [5]. If we combine the T and G kernels in one single call, the amount of data throughout the algorithm changes as follows: • The P2G step has a total of 12 components as input (n, q, y, and f in R3 ), and 12 components as output (f , q ⊗ n). • The FFT has 12 components as input and output. • The CNV steps has 12 components as input, and 3 components as output b). (the GF + T F combined v • The remaining steps (IFFT, G2P) are identical. In the remaining discussion, all analysis is for this simultaneous treatment of Stokeslet and stresslet sources. 2.4. Accuracy and Parameters. Here and in the remainder of the paper, without loss of generality,P let B to be the unit cube so that L1 = L2 = L3 = 1 and, for notational simplicity, let j ∥f (y j )∥2 = 1. We assume a nearly uniform distribution of points so that the number of points per cell s does not vary much across cells. Before proceeding further, we summarize the key algorithmic parameters: • rc is the near-field cutoff radius, which we use as P2P cell spacing; • s is the average number of points per cell; • ξ controls near–far field decomposition scale and is related to rc ; b F; • h is the FFT uniform grid spacing, which controls accuracy and cost of G • P is the number of FFT grid points supported by the window function w0 and controls accuracy and cost of the P2G and G2P operations; • ε is an user-specified absolute error tolerance for the total approximation uh . q PNt 1 /Nt i=1 |uh (xi ) − u(xi )|2 is shown In [5] the root mean square error Erms := to be estimated by r

(2.9)

112 4 3 −ξ2 rc2 4 Erms ≈ ξ rc e + 9 3h {z } | | EN

r

7 −π2 /(2hξ)2 e + C(ξ)e−2.5P . {z } | 2 {z } Ew EF

The first term E N is the near-field error due to the cutoff radius rc . The second term E F is the FFT error due to quadrature and truncation errors in the infinite series and integral in (2.4). The third term E w is the error due to P2G and G2P interpolations. We want to select all parameters to optimize runtime and ensure Erms < ε.2 For a given rc , we select ξ := ξ(rc ) so that E N < ε/3. Given ξ, we select h := h(rc ) and P := P (rc ) so that both E F and E w < ε/3, respectively. Let C N be the cell list set. Once we have these parameters, we set |C N | = 1/rc3 , i.e., the volume of B over the volume of a cell. Thus, the average number of points per cell s = Ns /|C N | = Ns rc3 . Similarly, the number of FFT grid points Ng = 1/h3 . Let T (Ng , Ns , P, s) be the sum of the time complexity of all terms in Algorithm 2.1. Since ξ, P, h, Ng , Ns , P, s are implicit functions of rc , both T and Erms are also implicit functions of rc . Thus, we can find rc by solving minrc T (rc ) subject to Erms (rc ) ≤ ε. Instead of the T from Algorithm 2.1 we can also use the GPU performance models with precise constants from section 5. In conclusion, the algorithm 2 Besides the parameters we discuss here, there are two additional parameters related to the truncation of the integral in (2.4c), but they can be also related to rc . The overall methodology doesn’t change.

FAST EWALD SUMMATION FOR STOKES FLOW

Near-field

P2P

7

P2P Cell list rc

Far-field

P2G

P2G Cell list aw

Ph h Fig. 2. Local structure for near ( top row) and far ( bottom row) interactions. Left column: source particles y j in orange ( top/bottom), target particles xi in blue ( top) and the regular grid points g ℓ in green ( bottom), with grid spacing h. Middle column: interaction radius for a source particle y j . The P2P interaction radius is rc . The P2G interaction radius is aw = P h/2, where P is the number of grid points needed to support w0 (r). Right column: cell list partition; the interaction radius requires each source cell to interact with each 3d of its colleagues. Each source cell is mapped to a GPU block. The G2P kernel is similar to P2G; the sources y j are replaced with targets xi .

has only one free parameter: the error tolerance ε, which is application dependent and set by the user. 3. GPU Acceleration. We now discuss the design and implementations of GPU algorithms for the Ewald sum. Subsections 3.1 to 3.3 present GPU implementations of the P2P, P2G, and G2P steps. Subsection 3.4 briefly addresses the FFT, CNV and the IFFT whose parallelization is more standard. Subsection 3.5 discusses Kokkos-based performance portability to CPUs, section 4 investigates the performance of different operations and special function evaluations on GPUs, and section 5 presents performance models for each stage of the Spectral Ewald method. As mentioned in section 2, P2P requires a cell list C N for efficient computation. To promote coalesced memory access on the GPUs, the cell list is padded so that each cell contains the same number of points. In addition, padded cell lists can accelerate each of the P2P, P2G, and G2P steps by assigning a GPU block B to a cell β, ensuring that each thread (tx , ty ) ∈ B processes spatially close data. Figure 2 summarizes the cell lists for P2P, P2G, and G2P. The construction of a padded cell list can be implemented in two passes through an unsorted particle array. The first pass, carried out in parallel, counts the maximum number of particles per cell s and the number of nonempty cells Nc . Then the cell-list array of shape (s Nc , d) is allocated — where d is the dimension of the data to be stored — as well as a cell counter of size Nc . The second pass reads the particle data from the unsorted array to the cell-list array by assigning one thread to one particle and logging the cell-list positions with AtomicFetchAdd updates to the cell counter. Cells are ordered via a linear index of a 3D partition which enables efficient indexing and memory access on GPU. In the following, we refer to the 1D GPU block size as bx and the 2D GPU block size as (bx , by ). 3.1. Particle-to-Particle Interaction (P2P). Near-field interactions are a standard component of parallel N -body libraries. Typically, a target cell β ∈ C N is assigned to a GPU block; each target point xm ∈ β is assigned to a thread, and a serial loop iterates over the sources points y j in the colleague cells α of β. We

8

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

call this approach P2P-GM-1D: a 1D parallel loop over targets, with sources read from global memory (GM). We can also store sources in shared memory (SM) as each source y j ∈ α is shared by every target xi ∈ β. To express this operation, we let LoadShmem() be a function that loads its arguments from the GPU memory to the shared memory. Another algorithmic variant uses a 2D thread grid and parallelizes on both target and sources. Assuming that GPU threads are indexed by (tx , ty ), tx is assigned to a target xi and each ty handles chunks of source points in the cell α. We refer to these chunks as αc . Thread (tx , ty ) accumulates all the interactions between xi and y j ∈ αc to a chunk potential uc (xi ). Finally, the threads in the row tx perform P a reduction to compute uN (xi ) = c uc (xi ). As in the 1D case, sources can be read from global or shared memory. The 2D shared-memory strategy offers increased parallelism over the sources at the cost of synchronization and diminished target-wise parallelism, as the total number of threads is fixed but more threads are allocated per target. All four variants follow the template in Algorithm 3.1, which is the 2D shared memory variant. P2P-GM-1D doesn’t use lines 4–5; αc = α so there is no loop in line 3; and there is no reduction. P2P-SM-1D uses lines 4–5; αc = α so there is no loop on line 3; and there is no reduction. P2P-GM-2D doesn’t uses lines 4–5. P2P-SM-2D is shown below. Algorithm 3.1 P2P-SM-2D Input: C N Output: uN h 0: Assign target cell β to a GPU block. 1: parfor xi ∈ β do 2: parfor c ← 1, . . . , ⌈s/|αc |⌉ do 3: uc ← 0 4: ∀y j ∈ αc , LoadShmem(y j , f (y j ), q(y j ), n(y j )) 5: synchthreads( ) 6: for y j ∈ αc do 7: uc ← uc + GN (xi − y j + p)f (y j ) 8: uc ← uc + T N (xi − y j + p)(q(y j )n(y j )) 9: end for 10: synchthreads( ) 11: end parfor 12: uN h (xi ) ← RowReduce(uc ) 13: end parfor

// Assign xi to thread tx . // Assign c to thread ty

// row-wise reduction

3.2. Particle-to-Grid Interaction (P2G). The P2G interaction interpolates the density from the source points to the grid. Its implementation involves several design decisions: Do we loop over grid points gathering contributions or do we loop over source points scattering contributions? Do we incur the costs associated with reordering source points in order to optimize streaming memory access or not? Can we avoid conditionals to find the correct interpolant? To explore these questions, we consider four algorithmic variants. Before describing them, it suffices to consider the P2G steps in 1D as the steps are identical across dimensions. Let y be a source point. We use w e0 to denote the polynomial approximation of w0 . To compute w e0 (r) we need to find the r-axis interval y belongs to; we use the term bin to refer to such r-intervals. Figure 3 summarizes the calculation that comprises the following steps: (1) find a bin anchor point a ← ⌊y/h⌋ − (P/2 − 1) on the grid; (2) compute a local distance δ = a − y; (3) compute w e0 (a + lh − y) for each bin l ← 0, . . . , P − 1. Two of the variants require reordering the source points. To do so, we create a new cell list C F . It is constructed identically to C N , but with a cell width equal to

9

FAST EWALD SUMMATION FOR STOKES FLOW

bin 0 bin 1 bin 2 bin 3 bin 4 bin 5 g8 g9 g 10 g 11 g 12 g 13 y w0 (r) 1

h h w̃0

0

r −3h −2h

−h

0

h

2h

3h

aw

Fig. 3. Evaluation of the one-dimensional piecewise approximation w e0 (solid curve) to the KB window function w0 (dashed curve) at a source y on the Fourier grid. An anchor point a = g8 is found and w̃0 is evaluated at bin 2 (i.e., g8+2 ) w e0 (g 8+2 − y) = w e0 (2h − (y − a)). (Here, P = 6.)

aw = P h/2. There is no need to explicitly reorder the grid points as their positions are computed analytically. Writing the grid values for the density does require uncoalesced writes; but the writes enjoy some locality since source points share grid points. Also, let C denote the user-defined input order, which in general can be considered random; and let G indicate the Fourier grid parameters and array. Then, the four variants, all using 1D GPU blocks, are as follows: • Baseline (P2G-base, Algorithm 3.2): Loop over source points in C order; compute contribution to grid points and use atomic add to update ϕ. • Source-ordered (P2G-source, Algorithm 3.2): Same as above, but source points are read in C F order. • Grid-ordered (P2G-grid, Algorithm 3.3): Loop over grid points to avoid atomics; requires repeated reads of source points, hence repeated evaluations of the window function for a source. • Grid-hybrid (P2G-hybrid, Algorithm 3.4): Use shared memory for the source points and traverse all neighbors of a source point to spread the density. P2G-base is easy to implement, but it incurs O(P 3 Ns ) atomic adds and the source point reads are completely random, although it is possible that this randomness reduces the overhead from atomics. P2G-source attempts to improve the source reads, but suffers from similar problems. Both versions avoid thread divergence, since threads in a warp evaluate w e0 in the same bin (lines 69). In contrast, P2G-grid, which parallelizes over the grid points g ℓ and serially loops over the sources y j , completely avoids atomics. However, due to the piecewisepolynomial approximation of w e0 , thread divergence is introduced into the serial loop; for a fixed g ℓ , the bin [k, l, m] ← g ℓ − a, where a is an anchor point dependent on y j , will vary randomly for given sources y j , conflicting with the GPU execution model. Additionally, parallelizing over the grid points increases the work of evaluating w at a source y j from O(2P 3 ) to O(3f (w e0 )P 3 ), where f (w e0 ) is the cost of computing w e0 . Finally, we load the same source repeatedly — in the worst case O(3d ) times — ignoring caching. The drawbacks of the first three variants motivate a two-stage hybrid parallelization scheme. Like in the P2P-base and P2G-source methods, each source cell α is assigned to a GPU block (or more if it has too many source points). Then for each source, we load the densities to the shared memory and save them for the loop beginning at line 15. Lines 4–11 also store the interpolation coefficients per point in shared memory. Once we synchronize on line 14, we assign each thread an index tx

10

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Algorithm 3.2 P2G-base / P2G-source Input: C/C F (source ordering), G (grid) Output: ϕ(G) 0: Assign α in C/C F to a GPU block. 1: parfor y j ∈ α do 2: LoadGlb2Reg(y j , f (y j ), q(y j ), n(y j )) 3: σ ← [f (y j ), vec(q(y j ) ⊗ n(y j )]T 4: a ← ⌊y j /h⌋ − (P/2) + 1 5: δ ← y j /h − a − p 6: for l ← 0, . . . , P − 1 do 7: w1 [l] ← w e0 (lh − δ1 ) 8: w2 [l] ← w e0 (lh − δ2 ) 9: w3 [l] ← w e0 (lh − δ3 ) 10: end for 11: for k ← 0, . . . , P − 1 do 12: for l ← 0, . . . , P − 1 do 13: for m ← 0, . . . , P − 1 do 14: g m ← a + [k, l, m]T 15: w ← w1 [k] w2 [l] w3 [m] 16: AtomicAdd(ϕ(gi ), w σ) 17: end for 18: end for 19: end for 20: end parfor

// Base / Source // Assign y j to a thread

// Compute grid anchor // Compute distance offset // Loop over bin l

// O(P 3 Ns ) atomics

Algorithm 3.3 P2G-grid Input: C F , G Output: ϕ(G) 0: Assign target box α to a GPU block. 1: parfor g ℓ do 2: ϕloc ← 0 3: for β ∈ Colleagues(α) do 4: for y j ∈ α do 5: LoadGlobal(y j , f (y j ), q(y j ), n(y j )) 6: σ ← [f (y j ), vec(q(y j ) ⊗ n(y j )]T 7: a ← ⌊y j /h⌋ − (P/2) + 1 8: δ ← a − y j /h − p 9: [k, l, m]T ← g ℓ − a 10: w←w e0 (kh − δ1 ) w e0 (lh − δ2 ) w e0 (mh − δ3 ) 11: ϕloc ← ϕloc + wσ 12: end for 13: end for 14: ϕ(g ℓ ) ← ϕloc 15: end parfor

// Assign g ℓ to a thread

// Thread divergence!

between 0 and (P/2)3 using the thread id. (If (P/2)3 < bx , we have some idle threads; if (P/2)3 > bx , we stride by bx to assign multiple indices to a thread.) In line 16, we begin a sequential loop over the colleague cells. Regular grid ordering guarantees that every colleague cell β has (P/2)3 grid points. As such, we define a local lexicographical ordering on grid points g ℓ ∈ β, and use integer arithmetic to determine the global grid coordinates of g ℓ from the thread index tx in a function we call Index2Grid to assign each thread a grid point g ℓ . Lines 18–26 mirror the computation in lines 4–12 of P2G-grid, but now aj , f (y j ), q(y j ), n(y j ) are read from shared memory as well as the 1D window values w1 [k], w2 [l], w3 [m], avoiding the thread divergence in P2G-grid line 10. On line 27, we atomically write ϕ(g ℓ ), as O(3d ) source cells α share a given colleague cell. In summary, P2G-hybrid minimizes atomics, reuses

FAST EWALD SUMMATION FOR STOKES FLOW

11

source points, and avoids repeated computations of the window function. In fact, we experimented replacing line 27 with a different variant that first writes at different locations and then calls an additional kernel for reduction, but it was slower. Algorithm 3.4 P2G-hybrid Input: C F , G Output: ϕ(G) 0: Assign source box α to a GPU block 1: parfor y j ∈ α do 2: LoadGlobal(y j ) 3: StoreGlobal2Shmem(f (y j ), q(y j ), n(y j )) 4: aj ← ⌊y j /h⌋ − (P/2) + 1 5: StoreShmem(aj ) 6: δ ← y j /h − aj − p 7: for l ← 0, . . . , P − 1 do 8: w1 [l] ← w e0 (lh − δ1 ) 9: w2 [l] ← w e0 (lh − δ2 ) 10: w3 [l] ← w e0 (lh − δ3 ) 11: StoreShmem(w1 [l], w2 [l], w3 [l]) 12: end for 13: end parfor 14: synchthreads( ) 15: parfor tx ← 0, . . . , (P/2)3 − 1 do 16: for β ∈ Colleagues(α) do 17: g ℓ ← Index2Grid(tx , α) 18: ϕloc ← 0 19: for y j ∈ α do 20: LoadShmem(aj , f (y j ), q(y j ), n(y j )) 21: σ ← [f (y j ), vec(q(y j ) ⊗ n(y j )]T 22: [k, l, m]T ← g ℓ − aj 23: LoadShmem(w1 [k], w2 [l], w3 [m]) 24: w ← w1 [k] w2 [l] w3 [m] 25: ϕloc ← ϕloc + w σ 26: end for 27: AtomicAdd(ϕ(g ℓ ), ϕloc ) 28: end for 29: end parfor

// Assign y j to a thread

// Use (P/2)3 threads // Assign g ℓ to a thread

// O(27Ns ) atomics

3.3. Grid-to-Particle Interaction (G2P). This operation can be thought of as the “transpose” of the P2G interaction. Here we interpolate the potential from the grid points to the target points. For brevity, we only outline this interaction. We consider two variants G2P-base and G2P-target. The first variant is essentially identical to P2G-base, the difference being that in line 1 in Algorithm 3.2 we parallel loop over target points. Then each thread sequentially loops over the grid points in the support of KB window centered at the target, and interpolates the potential at the target point. Unlike the P2G kernel, no atomics are needed as each thread adds to each target point potential. As such, in G2P-target the only optimization we make is to reorder the targets using C F to improve data locality. 3.4. Fourier Grid Convolution (FGC). We now discuss the computations in Fourier space. Single GPU FFT and IFFT are done using vendor libraries. The b F is diagonal; we assign a thread to each frequency κℓ and compute convolution with G the multiplication in (2.6). Moreover, since each of ϕ, GF , and w is real-valued, their Fourier transforms exhibit conjugate symmetry. As a result, the convolution only needs to be computed for the nonnegative wavenumbers κℓ , effectively reducing the computational workload by half compared to the general nonsymmetric case.

12

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Table 1 Observed millicycles per operation for a selection of FP64 CUDA functions on the NVIDIA H200 GPU, measured with NVIDIA Nsight Compute. Flop-equivalents are computed by dividing the millicycles of each operation by those of a DFMA and multiplying by 2. Actual FP64 flops are obtained from Nsight Compute; additional FP32 flops in parenthesis if present. Shaded rows are machine instructions. Operation

DFMA DADD DMUL DIV SQRT RSQRT EXP ERF ERFC SINH I0 J0 J1

Millicycles Flop-equivalents Actual flops 0.121 2 2 0.127 2.10 1 0.127 2.10 1 2.11 35.0 15 (2) 2.50 41.5 13 1.82 30.2 8 3.07 50.8 30 9.33 155 78 (1) 16.1 267 112 21.8 361 49 (2) 5.00 82.9 36 15.8 262 36 39.4 654 36

3.5. Portable Implementation Using Kokkos Abstractions. All interactions and their variants have been implemented using Kokkos abstractions via the Python interface PyKokkos. We map each GPU block to a Kokkos team and utilize Kokkos scratchpad memory to emulate managed shared memory. For 2D kernels, thread rows are mapped to Kokkos vector ranges, which are guaranteed to execute within the same warp. Kokkos also provides built-in support for vector reductions, which we leverage to implement RowReduce in Algorithm 3.1. This level of GPU abstraction enables seamless portability of our kernels across both CUDA, HIP, and CPU backends. A shortcoming for CPUs is that we do not use explicit vectorization but rather rely on the compiler, although Kokkos does offer SIMD data types for abstracting vector intrinsics. When targeting the CPU, Kokkos teams (originally corresponding to GPU blocks) are mapped to OpenMP threads. Accordingly, we set the block size to one, effectively restructuring the cell list so that each team operates on a single particle per cell. 4. Special Function Evaluation. The scheme involves a number of special mathematical functions such as erfc(x), e−x and sinh(x). To estimate the cost of these functions, we provide Table 1, evaluating the double-precision CUDA functions with random input numbers in the range [10−3 , 7], which is typical for our method, on an NVIDIA H200 GPU using NVIDIA Nsight Compute. Performance is measured in the average number of millicycles per operation and normalized by the half the average number of millicycles per fused multiply-add (DFMA) to obtain a FLOPequivalent metric. For example, a division (DIV) requires on average 16.6 as many cycles as a multiplication (DMUL), and is computed using 15 double precision (and 2 single precision) flops. For some functions (e.g., J1 ), the cycle count is much larger than expected from the flop count, likely due to thread divergence. As mentioned in subsection 2.2, the window function w0 (r) is approximated by a piecewise polynomial of degree ν to avoid the cost of evaluating the Bessel function I0 . A separate polynomial is used in each bin (r-axis interval) of width h shown in Figure 3, and these polynomials are precomputed, as described in [34]. The poly-

10 8 6 4 2 0

Rational approximation of erf(x)/x

Baseline erf(x)/x millicycles Rational approx. millicycles Rational approx. max. error

4

6 8 10 12 Approximation order m

13

10 2 10 4 10 6 10 8 10 10 10 12 10 14

Maximum error

Millicycles

FAST EWALD SUMMATION FOR STOKES FLOW

14

Fig. 4. Millicycles per evaluation and maximum error for rational approximation of erf(x)/x on the interval [0, 6].

nomial approximation of w0 can be evaluated using Horner’s rule, requiring only ν DFMA operations per evaluation point. The erfc(x)/x function that appears in the P2P step, cf. (2.3), can also be approximated. First note that erfc(x)/x = 1/x − erf(x)/x, and erf(x)/x is a smooth and slowly decaying function that is suitable for rational function approximation. We apply the AAA algorithm [26, 27] to construct a barycentric rational approximant Pm wj fj Pm wj r(x) = j=1 x−z / j=1 x−zj of erf(x)/x for x ∈ [0, 6]. However, the barycentric j form has m + 1 divisions, which is far too expensive. Instead, we convert the output of the AAA algorithm to the polynomial form r(x) = p(x)/q(x) where both p and q are polynomials of degree m − 1, allowing us to use 2(m − 1) DFMAs and a single division to evaluate r(x). The millicycles observed and the maximum error as a function of m are shown in Figure 4. At m = 8, the error is 10−8 and the speedup of using rational approximation instead of the CUDA function is 1.46×. Empirical measurements suggest that around 22.5% of the P2P runtime is taken up by erf(x)/x evaluation, when using the CUDA erf function. The speedup for the entire P2P runtime when using the rational approximation with m = 8 would then be 1.08×. 5. Performance Models. We construct GPU performance models — to evaluate the algorithms described in section 3 — that account for global memory, L2 cache, registers, and shared (or L1) memory. The models are constructed with increasing levels of cache reuse and are denoted by • T0 (used to model P2G/G2P): A 0-L2 model with no L2 or L1 cache reuse but with managed shared memory and an assumption of coalesced highbandwidth memory (HBM) access for bandwidth values. • T∞ : Same memory characteristics as T0 with an additional infinite L2 cache; data is loaded from HBM once and reused from L2 thereafter. • T∞,l1 : Same memory characteristics as T∞ with an additional infinite L1 cache; data is reused from L1 once prefetched from L2. • Tf (used to model P2P): A compute-only model that accounts solely for floating-point operations, ignoring all data movement costs. The models are defined by the following variables: • τf , the machine time/FLOP; • τm , the global memory time/byte (load or store); 1 • τm , the L1 memory time/byte (load or store);

14

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Performance (Tflop/s)

102

Roofline Model for H200

101

Bandwidth Limit Peak FLOP(64)/s P2P-GM-1D P2P-SM-1D P2G-GRID P2G-HYBRID G2P-TARGET

100

10−1 10−1

100

101

102

103

Arithmetic Intensity (flop/byte) Fig. 5. Analytical roofline and flops of Algorithm 2.1 on the NVIDIA H200 GPU for N = 4 × 106 , s = 224, ϵ = 10−9 . T∞ model is used for P2P kernels, T0 for P2G and G2P kernels. 2 • τm , the L2 memory time/byte (load or store); 1 2 • l1 = τm /τm and l2 = τm /τm ; • f , the algorithm total flop count; • m, sum of bytes for all the input data; • mG and mS , the memory load/stores from the global and shared memory, respectively. Model specific memory trafic estimates are given by m0 , m∞ , and m∞,l1 , defined as m0 = mG + l1 mS , m∞ = m + l2 mG + l1 mS and m∞,l1 = m + l1 mG + l1 mS . The models are explicitly expressed as T0 = τf f + τm m0 , T∞ = τf f + τm m∞ , T∞,l1 = τf f + τm m∞,l1 and Tf = τf f . Arithmetic intensity is given by qk = f/mk for k ∈ {0, ∞, ∞, l1 }. If qk ≥ τm/τf , the kernel is compute-bound under model k; otherwise, it is memory-bound. Subsections 5.1 to 5.3 provide the analytical value of f , m0 , m∞ and m∞,l1 for each kernel (P2P, P2G, and G2P) and compare model predictions with empirical GPU runtimes. Sources and targets are independently generated by a nearly uniform distribution in the unit cube and all numerical quantities and arithmetic are in double precision, unless otherwise noted. We provide a brief summary of the accuracy of our models: • P2P (Table 2): The SM-1D kernel has measured times slightly lower than T∞,l1 , which we attribute to partial overlap between memory transfers and computation. • P2G, G2P (Table 5): Measured times exceed the T0 model as the model assumes coalesced accesses to HBM, but both P2G and G2P have irregular global memory accesses. Efficiency for the P2P kernel is calculated using Tf as a baseline, while efficiency for the P2G and G2P kernels are calculated using T0 as a baseline. The models are optimistic; achieving high efficiency compared to the baseline provides strong evidence that our kernels are highly optimized. Figure 5 presents a roofline analysis of the different kernels. More accurate models such as Volkov [48] can be used, but our models are designed to provide a reasonable approximation to observed performance behavior by capturing the dominant data movement and compute costs in each kernel.

5.1. P2P. The P2P kernel’s flop count is given by fP2P = 27Nt s (Cp2p π/6 + 14), where s is the number of sources in one cell and Cp2p varies across machines due

FAST EWALD SUMMATION FOR STOKES FLOW

15

Table 2 Analytical memory operation (MOP) counts for P2P algorithmic variants under the T∞ model. GPU block size is denoted bx for a 1D block and (bx , by ) for a 2D block. Here, m = 48Nt + 96Ns . The MOP value for the T0 model (m0 ) is obtained by setting m = 0 and l2 = 1, while the MOP value for the T∞,l1 model (m∞,l1 ) is obtained by setting l2 = l1 .

GM-1D m + l2 Nt (48 + 96 · 27s) GM-2D m + l2 Nt (24by + 24 + 96 · 27s) m∞ SM-1D m + l2 Nt (48 + 96·27s/bx ) + l1 Nt (24 + 96 · 27s) SM-2D m + l2 Nt (48 + 96·27s/bx ) + l1 Nt (24by + 96 · 27s) to differing costs for special function evaluations (see Table 1). For instance, on the NVIDIA A100 GPU, Cp2p = 318, while on the NVIDIA H200 GPU, Cp2p = 414. The load operations for the P2P kernel include the target points coordinates (3Nt doubles) and the source points coordinates, normal vectors, and single-layer and double-layer densities (12Ns doubles). The outputs are the potentials at the target points (3Nt doubles). Thus, m = 48Nt + 96Ns for all P2P variants, while mG and mS depend on the variant. In GM-1D, each thread loads one target and 27s sources and globally stores one target. Thus, mGM−1D = Nt (48 + 96 · 27s), and mGM−1D = 0. In G S SM-1D, given a GPU block of size bx ≤ s, each block moves bx targets and 27s sources (in chunks of size bx ) into shared memory. Thus, mSM−1D = Nt (48 + 96·27s/bx ) G SM−1D and mS = Nt (24 + 96 · 27s). For the 2D variants, let (bx , by ) denote the size of the 2D block. Then mGM−2D = Nt (24by + 24 + 96 · 27s) and mGM−2D = 0; G S SM−2D SM−2D 96·27s mG = Nt (48 + /bx ) and mS = Nt (24by + 96 · 27s). Table 2 summarizes the memory operations. Table 3 shows the empirical and analytical runtime under different performance models for different P2P algorithmic variants. The results suggest that the T∞ model is somewhat pessimistic, likely because it does not account for prefetching from L2 to the system managed L1. This observation motivates the inclusion of the T∞,l1 and Tf models at the beginning of this section. Notably, the SM-1D variant achieves measured times even lower than T∞,l1 , indicating that the kernel is highly optimized. Table 3 P2P empirical and analytical time (in milliseconds) on the NVIDIA H200 GPU. T denotes the empirical time. T0 , T∞ , T∞,l1 and Tf denote the analytical time under the 0-L2 cache, ∞L2 cache, ∞-L2 cache with prefetching and compute-only model, respectively. l2 = 1/4, l1 = 1/10. Nt = 4 × 106 , s = 256.

method (bx , by ) T T0 T∞ T∞,l1 Tf GM-1D (64, 1) 261 854 357 257 190 GM-2D (8, 32) 336 855 357 257 190 P2P SM-1D (32, 1) 258 278 262 259 190 SM-2D (64, 2) 323 267 260 258 190 5.2. P2G. In the P2G kernel, each source interacts with P 3 grid points, except for the P2G-grid method, which interacts with each grid point in colleague cells. In the P2G-base, P2G-source, and P2G-hybrid methods, w e0 is calculated per source, and the cost can be neglected; in the P2G-grid method, w e0 needs to be calculated at 3 each interaction. The flop count for P2G-grid is fP2G = 6νNs P 3 + 50 · 27Ns (P/2) , whereas for the remaining variants is fP2G = 46Ns P 3 , P2G load operations involve source points coordinates, normal vectors, and both single layer and double layer densities (12Ns doubles). P2G stores in memory the

16

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Table 4 Analytical MOP values (in bytes) for P2G variants for the T∞ model; bx is the GPU block size and m = 96Ns + 96Ng . The MOP values for the T0 model (m0 ) are obtained by setting m = 0 and l2 = 1, while the MOP values for the T∞,l1 model (m∞,l1 ) are obtained by setting l2 = l1 . 3 base & source m + l2 (24N  s + 168Ns P )

m∞ grid hybrid

 3 m + l2 96 · 27Ns (P/2) + 96Ng
3 m +l2 96Ns + 96·27Ns (P/2) /bx  3 + l1 96Ns P 3 + 12 · 27Ns (P/2)

vector and tensor densities (12Ng doubles) at the grid points. Therefore, m = 96Ns + 96Ng . Let sf be the number of sources in one far-field cell. In the P2G-hybrid method, each block stores bx ≤ sf sources in shared memory, and atomically writes 3 to all grid points in colleague cells. Thus, we have mhybrid = 96Ns + (96·27Ns (P/2) /bx ). G The block loads each source’s anchors from shared memory for each grid point in colleague cells, and if the grid point is in the P 3 cube of the source, it loads the source’s normals, two densities, and its approximated window function values. Therefore, 3 = 96Ns P 3 + 12 · 27Ns (P/2) . Shared memory is not used for other methods. mhybrid S = = 24Ns + 168Ns P 3 ; for grid, mgrid = msource For base and source, mbase G G G 3 base source P 96 · 27Ns ( /2) + 96Ng , and mS = mS = mgrid = 0. Table 4 summarizes the S memory operations. Table 5 shows the empirical and analytical time under different performance models for the P2G kernel. Table 5 P2G and G2P empirical and analytical runtime (in milliseconds) on the NVIDIA H200 GPU. T denotes the empirical time. T0 , T∞ , T∞,l1 and Tf denote the analytical time under the 0-L2, ∞L2, ∞-L2 with prefetching and compute-only model, respectively. l2 = 1/4, l1 = 1/10. Ns = 4 × 106 , P = 8, s = 224, ϵ = 10−9 .

method (bx , by ) T T0 T∞ T∞,l1 Tf base (256, 1) 477 88.9 24.9 12.0 2.81 source (512, 1) 354 88.9 24.9 12.0 2.81 P2G grid (128, 1) 246 98.2 34.2 21.3 12.2 hybrid (64, 1) 39.5 21.7 13.4 11.6 2.81 base (32, 1) 62.8 12.8 3.73 1.88 0.489 G2P target (32, 1) 28.3 12.8 3.73 1.88 0.489 5.3. G2P. In both the G2P-base and G2P-target methods of the G2P kernel, each thread globally reads one target and P 3 grid points, computes the interactions and summations, and then globally writes to one target. Therefore, model parameters are given by f = 8Nt P 3 , m0 = 24Nt P 3 and m∞ = 48Nt + 24Ng + l2 (24Nt P 3 ). Table 5 compares predicted runtimes of the P2G and G2P models with empirical results. The measured times are lower than those predicted by the T0 model. We speculate that the most likely cause of weak performance is due to uncoalesced accesses of loads and stores. The cache provides little benefit due to the irregular memory accesses. 6. Numerical Results and Performance Analysis. The accuracy of the spectral Ewald method is well studied (see [5] for Stokes, [34] for Poisson). Here, we study the HPC performance of our implementation. We summarize the key questions of our investigation:

FAST EWALD SUMMATION FOR STOKES FLOW

17

• Performance comparison of the Ewald P2P with simpler P2P kernels: How does the Ewald Stokes P2P kernel compare to that of free-space Poisson and Stokeslets P2P kernels? How does differing arithmetic intensity affect performance? (Figure 6) • Algorithmic variants: How do the different variants of P2P, P2G, G2P perform on a single device? (Table 6, Table 7 and Table 8, respectively) • Performance portability: What is the performance portability of our implementation? (Figure 7 for P2P, Figure 8 for P2G, Figure 9 for the overall Ewald sum) • Overall Ewald sum on a single device: What is the overall performance for different particle distributions and at different precisions? How does the sum perform and scale on different devices? (Table 8, Table 9 and Figure 9, respectively) • Overall Ewald sum on multiple devices using MPI: What are the MPI overheads for our solver with multi-device computations? (Table 11) Experimental setup: Our experiments were carried out at the Texas Advanced Computing Center (TACC) and the OACISS Center at the University of Oregon. On TACC, we used Lonestar 6 for the A100s and AMD Epyc 7763, and Vista for the H200 and Grace systems; On OACISS we used the Odyssey M1300A system. The ParkiPy library is implemented in Python, combining CuPy [29] for array operations with PyKokkos [3] to interface with the Kokkos kernels described in section 3. The CPU backend runs via the Kokkos::OpenMP execution space using gcc with the -O3, -march=native, and -mtune=native compiler flags, but without explicit vectorization or tuning. GPU code is compiled via the Kokkos::CUDA execution space for NVIDIA devices and the Kokkos::HIP execution space for AMD devices. We profile performance using Tau2 [40] for system-level data and NVIDIA Nsight Compute [28] for kernel-level analysis on NVIDIA GPUs. Unless otherwise noted, source and target y j , xi , as well as f (y j ), q(y j ) and n(y j ), are sampled from a uniform probability distribution in the computational domain with Nt = Ns and all code is compiled in double precision. Flops are counted using the flop-equivalences in Table 1. 6.1. P2P. Table 6 compares execution times and efficiencies for the P2P GM1D, GM-2D, SM-1D, and SM-2D methods on a single NVIDIA H200 GPU for different s (number of points per cell) with N = 106 . We use an empirically computed optimal block size for each method. For 1D methods, bx is selected from B = {32, 64, 128, 256, 512}; for 2D methods, (bx , by ) is selected from {(bx , by ) : bx by ∈ B, by ∈ {2, 8, 32}}. The source chunk size bx by is constrained by the shared memory size on H200. For all methods, we achieve 63%–78% flop efficiency with s = 512. The GM-1D and SM-1D methods exhibit similar execution time and efficiency, and significantly outperform their 2D counterparts. Interestingly, the GM-2D method performs best when the source dimension is highly parallelized (by = 32), while the SM-2D method performs best with near-1D layouts (by = 2). In SM-2D, a large by results in a shared memory allocation of 96Nt by , which may reduce the number of blocks launched concurrently due to shared memory size constraints. Figure 6 shows the performance of the P2P-GM-1D kernel on a single NVIDIA H200 GPU under varying workloads, including the Poisson kernel (lightest), the Stokes single-layer kernel, and the full Ewald kernel (heaviest), with N = 106 . For each workload, performance improves with increasing s due to better latency hiding, as each thread processes more interactions. The full Ewald kernel achieves the highest efficiency (76% flops with s = 512) as the evaluation of special functions makes it more

18

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Table 6 Particle-to-particle (P2P) performance results on a single NVIDIA H200 GPU. Flop efficiency is based on the Tf model. s = 256. (bx , by ) is the optimal.

s = 256 s = 512 Nt t (bx , by ) N Flops (b , b x y ) µs Flops µs 1 GM-1D (64, 1) 14.8 71% (64, 1) 8.0 76% GM-2D (1, 32) 11.8 56% (2, 32) 6.7 63% SM-1D (32, 1) 15.2 72% (32, 1) 8.2 78% SM-2D (64, 2) 12.4 59% (128, 2) 6.6 63%

106 Nt method

P2P-H200-GM-1D 250

37% flops

Workload

N/µs

200 150

57% flops

50 0

46% flops

49% flops

100

64

Poisson Stokes Single Layer Ewald

52% flops

58% flops 66% flops 128

s

65% flops 256

60% 71% flops 71% 76% flops flops flops 512

Fig. 6. Particle-to-particle performance on a single NVIDIA H200 GPU for different workloads (higher is better), N = 106 .

computationally intensive. Even the lightest workload, the Poisson kernel, achieves 60% flops when s = 512, showing that the strong performance of the P2P implementation does not rely on the particular workload and generalizes to other problems. Figure 7 presents a performance comparison of the P2P-GM-1D kernel for the full Ewald sum on various GPUs and CPUs. For GPUs, we achieve flop efficiencies of 84% (A100), 73% (H200), and 60% (MI300A) with N = 4 × 106 , s = 256. NVIDIA devices achieve higher efficiency, but the AMD MI300A attains the lowest execution time due to its highest peak flop performance. On CPUs, the NVIDIA Grace achieves over 50% efficiency, while AMD EPYC exceeds 60%. Efficacy is achieved with the ported Kokkos kernels as is, without hardware-specific tuning. 6.2. P2G. Table 7 compares different P2G methods with varying numbers of source particles and target accuracies E. Each row reports results using an empirically computed optimal 1D thread configuration, bx ∈ {32, 64, 128, 256, 512}, on a single NVIDIA H200 GPU. While P2G-source, P2G-grid, and P2G-hybrid maintain consistent efficiency for fixed accuracy and increasing Ns , the baseline P2G-base method loses efficiency as Ns increases, resulting in larger speedups for P2G-hybrid. This trend probably results from increasing atomic conflicts in P2G-base with larger problem sizes. Interestingly, P2G-source underperforms relative to P2G-base at low accuracies, which we attribute to increased atomic conflicts due to increased spatial locality within thread blocks. At high accuracy tolerance, the advantage of P2G-hybrid becomes more pronounced. For example, with Ns = 4×106 , spectral Ewald tolerance 10−14 and optimal cell size s = 960, P2G-hybrid achieves a 16.18× speedup over P2G-base. In this regime, the large near-field s leads to a smaller Fourier grid, which again increases atomic contention in P2G-base while conversely enhancing spatial locality in P2G-

19

FAST EWALD SUMMATION FOR STOKES FLOW

P2P-GM-1D

30 Device

N/µs

25 20 15

H200 MI300A A100

58% flops

40% flops

10

0

73% flops

71% flops

Device

1.5

Grace Epyc

1.25

51% flops

1

63% flops

52% flops

68% flops

52% flops

65% flops

0.75 72% flops

5

1.75

60% flops

58% flops

84% flops

84% flops

0.5 0.25

2.5 × 105

1 × 106

N = Ns = Nt

4 × 106

0

2.5 × 105

1 × 106

N = Ns = Nt

4 × 106

Fig. 7. Particle-to-particle performance portability (higher is better), s = 256. Note the different y-axes for GPU and CPU plots. Table 7 Particle-to-grid (P2G) performance results on a single NVIDIA H200 GPU. Efficiency is based on the T0 model. Near-field cell size s is 160 for P = 4, E = 10−7 and 224 for P = 8, E = 10−9 . Each row is sub-selected as the fastest run with bx ∈ {32, 64, 128, 256, 512}. Efficiency (Eff ) is reported in mops for the BASE, SOURCE, and HYBRID methods, and flops for the GRID method.

P = 4, E = 10−7 P = 8, E = 10−9 P = 14, E = 10−14 Ns s Speedup Eff Speedup Eff N Speedup Eff µs µs 0.25 base 101.65 — 6% 12.52 — 5% 1.91 — 3% source 48.2 0.47× 3% 11.83 0.95× 5% 2.5 1.31× 3% grid 83.41 0.82× 10% 16.42 1.31× 26% 3.14 1.64× 49% hybrid 147.22 1.45× 9% 90.73 7.25× 37% 21.93 11.48× 31% 1 base 103.02 — 6% 9.97 — 4% 1.61 — 2% source 62.36 0.61× 4% 11.94 1.20× 5% 2.4 1.49× 3% grid 93.58 0.91× 11% 16.87 1.69× 27% 3.1 1.92× 48% hybrid 177.26 1.72× 11% 100.65 10.09× 42% 22.1 13.72× 27% 4 base 80.57 — 5% 8.38 — 3% 1.49 — 2% source 54.17 0.67× 3% 11.3 1.35× 4% 2.52 1.69× 3% grid 82.52 1.02× 10% 16.28 1.94× 26% 3.39 2.28× 52% hybrid 178.23 2.21× 11% 101.28 12.09× 39% 24.05 16.18× 25%

106 Ns method

Ns µs

hybrid. Figure 8 explores the performance portability of P2G-hybrid and P2G-grid. The fastest overall performance is achieved by P2G-hybrid, which processes approximately 6 × 107 particles per second with 30% efficiency. Surprisingly, the AMD MI300A underperforms significantly for the P2G-hybrid method, processing nearly 6× fewer particles per second than NVIDIA GPUs. This drop appears to be related to shared memory limitations in AMD devices, consistent with the trends observed s for P2P-SM-1D. P2G-hybrid has a shared memory workload of (3P + 9) bxf . On the other hand, P2G-grid exhibits consistent performance and efficiency on both NVIDIA and AMD GPUs, and in fact outperforms P2G-hybrid on the MI300A. On CPUs, P2G-hybrid remains the optimal method for P2G computations, although its performance advantage and achieved efficiencies are less pronounced compared to the GPUs. 6.3. G2P. Table 8 compares G2P-base and G2P-target for varying numbers of target particles Nt and spectral Ewald tolerances (controlled by P = 4 and P = 8) on a single NVIDIA H200 GPU. In general, low-accuracy computations (P = 4) are less efficient. For smaller problem sizes, G2P-base outperforms G2P-target due to

20

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

P2G-HYBRID

80 70

N/µs

60

P2G-GRID Device

30% mops

31% mops

Device

H200 MI300A A100

28% mops

H200 MI300A A100

50 40

37% mops

30 20

35% mops

31% mops

3% mops

3% mops

3% mops

10

35% 35% mops mops 26% 26% 25% mops 24% 38% mops mops 36% mops 32% mops mops mops

0 Device

3.5

11% mops

3

N/µs

2.5 2 1.5

9% mops

18% mops

15% mops

Device

Grace Epyc

11% mops

Grace Epyc

21% mops

20% mops

20% mops 28% mops

1

30% mops

18% mops

28% mops

0.5 0

2.5 × 105

1 × 106

N = Ns = Nt

4 × 106

2.5 × 105

1 × 106

N = Ns = Nt

4 × 106

Fig. 8. Particle-to-grid performance portability (higher is better). E = 10−9 , s = 224. Note the different y-axes for GPU and CPU plots. Table 8 Grid-to-particle (G2P) performance results on a single NVIDIA H200 GPU. Mop efficiency is calculated using the T0 model. Near-field cell size s is 160 for P = 4, E = 10−7 and 224 for P = 8, E = 10−9 . Each row is sub-selected as the fastest run with bx ∈ {32, 64, 128, 256, 512}.

P = 4, E = 10−7 P = 8, E = 10−9 Nt Nt µs Speedup Mops µs Speedup Mops 0.25 base 447.73 — 21% 85.56 — 31% target 430.85 0.96× 20% 119.41 1.40× 44% 1 base 661.48 — 28% 66.53 — 24% target 668.02 1.09× 31% 136.97 2.06× 50% 4 base 523.13 — 24% 63.68 — 23% target 760.22 1.45× 35% 141.38 2.22× 52%

106 Nt method

lower GPU overhead and simpler memory access patterns. At higher accuracy (P = 8), G2P-target becomes more efficient, as expected. This improvement is primarily due to better memory coalescing over the target particles, enabled by the introduction of the cell list C F . Figure 9 shows the performance portability of G2P-target in the context of a full Ewald sum with s = 224 and target error E = 10−9 . G2P efficiency is modeled in MOP/s using the 0-cache model. As has become a pattern with our memorybound Kokkos kernels, we observe similar efficiencies on NVIDIA GPUs, while the AMD MI300A exhibits significantly lower efficiency. For the same run on CPUs, i.e., s = 224, Ns = 2.5 × 105 , 1 × 106 , 4 × 106 , we observe efficiencies of 50%, 60%, and 58% on NVIDIA Grace and 65%, 86%, and 93% on AMD Epyc. 6.4. Overall Ewald Sum. Table 8 compares the run-times for each stage in the Ewald sum for different particle distributions with target error E = 10−9 , nearfield cell size 224, and N = 4 × 106 source and target points. The P2P and P2G

FAST EWALD SUMMATION FOR STOKES FLOW

21

Table 9 Ewald sum performance in milliseconds on a single NVIDIA H200 GPU for different particle distributions; N = 4 × 106 , s = 224, P = 8 and E = 10−9 . X/Y indicates the source and target distribution, respectively, where U(0,1) is a uniform-distribution, N(0,0.3) a normal-distribution shifted and truncated to the unit cube, and S is uniform on the surface of the unit sphere. We use the P2P-GM-1D, P2G-HYBRID, and G2P-TARGET methods.

U /U N /N N /U S /S S /U P2P 239.0 351.4 249.4 901.8 268.9 P2G 47.6 77.4 76.3 59.7 59.2 FFT 48.7 48.7 48.7 48.8 48.8 CNV 20.0 20.1 20.2 20.6 20.4 10.8 10.8 IFFT 10.7 10.7 10.7 G2P 36.7 37.1 36.6 21.3 36.8 Total 402.8 545.4 441.9 1063.0 444.8 Table 10 Spectral Ewald performance in milliseconds with N = 4×106 , results on a single NVIDIA H200 GPU for different dtypes. Near-field cell size s is 160 and the target error E = 10−7 . We use the P2P-GM-1D, P2G-HYBRID, and G2P-TARGET methods.

P2P P2G FFT CNV IFFT G2P Total

single-precision double-precision 92.67 191.34 28.24 36.69 4.8 8.17 6.72 11.72 1.27 1.87 7.95 17.51 141.65 267.3

stages are most sensitive to non-uniform particle distribution as both algorithms are reliant on non-adaptive cell lists that grow in size are particles become clustered. Detailed analysis of Ewald sum runtimes for particle distributions with varying levels of non-uniformity are provided in supplement B. Table 9 compares the runtimes of the Ewald sum in single and double precision arithmetic with target error E = 10−7 , near-field cell size 160, and N = 4 × 106 source and target points. All Ewald stages besides P2G and IFFT show about a 2× speedup when performed in single precision arithmetic. The P2G algorithm involves the same costly integer arithmetic and block synchronization in both single and double floating point precision, yielding only a 1.3× speedup. We hypothesize that the IFFT launch is dominated by overhead for such small a size in both precisions. Overall, we see a 1.9× speedup between the single precision and double precision Ewald sums. Figure 9 presents a full Ewald sum with fixed target error E = 10−9 and near-field cell size s = 224. The largest run, with Ns = 4 × 106 source particles, requires an upsampled Fourier grid too large to fit in memory on the A100 GPU. On the MI300A, the absence of a highly efficient P2G kernel causes the total runtime to lag behind that of the H200, despite the MI300As higher theoretical throughput. Both CPUs have similar runtimes and efficiencies for the P2P, P2G, and G2P steps. The Fourier Grid Convolution step — dominated by memory-bound FFT and IFFT operations — is most efficiently executed on the H200 for GPUs, reflecting their superior memory bandwidth and optimized FFT performance. Grace dominates the FCG step on the CPU, which we attribute to optimized library implementations for the FFT and IFFT steps.

22

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Stokes Spectral Ewald Performance Portability A100

Epyc

Stage

FGC G2P

0.30

P2G P2P

Stage

FGC G2P

3 56%

P2G P2P

Epyc

A100

0.25

MI300

70%

Epyc

µs/N

61% 0.20

H200

32%

2 80%

0.15

0.10

13%

H200 67%

53%

47%

H200

15%

1

16%

Grace 12%

50%

51%

8%

7%

1 × 106

4 × 106

13%

5%

73% 0.05

63%

Grace 48%

57%

34% 13%

9%

MI300

MI300 50%

Grace

25% 28%

0.00

2.5 × 105

1 × 106

N = Ns = Nt

4 × 106

0

2.5 × 105

N = Ns = Nt

Fig. 9. Spectral Ewald run on different machines (lower is better). E = 10−9 , s = 224. P2P efficiencies are reported with the L2 prefetching model (i.e., flop dominant) and the P2G and G2P efficiencies are computed with the 0-cache memory model (i.e., mop dominant). Note the different y-axes for GPU and CPU plots.

Overall, as the number of particles increases, we observe clear O(N log N ) scaling in the overall runtime, consistent with the expected complexity of the Ewald sum algorithm. 6.5. Multi-GPU Performance. We consider a standard multi-GPU strategy based on message passing (MPI). Since a number of the kernels (P2P, P2G, CNV, G2P) are spatially local, we divide the domain B into equisized slabs along the periodic spatial dimension, delegating each slab to one GPU. The slab partition is in fact compatible to problems of any periodicity; the only periodicity specific detail being the edge-slab communication pattern. Communication between GPUs can then be reduced to: (1) initial radix sorting of sources and targets so that each particle belongs to the correct GPU (MPI-SORT); (2) exchanging ghost source points with neighbors; the points within distance rc to the slab interface are exchanged before P2P (MPI-GHOST-SOURCE); (3) similar exchange of ghost grid values between IFFT and G2P steps (MPI-GHOST-GRID); (4) FFT and IFFT steps, which we implement using NVIDIA’s cuFFTMp library. Since Kokkos does not provide a performance portable distributed FFT, we rely on vendor libraries; the aforementioned cuFFTMp for NVIDIA GPUs, rocFFT for AMD GPUs, and FFTW MPI for CPUs. A weak scaling test with N = 4 × 106 NGPU source and target particles on NGPU H200 GPUs is shown in Table 11. The GPUs are connected via NVIDIA InfiniBand with 400 Gbit/s bandwidth. Inter-GPU communication takes place in the FFT, IFFT and MPI steps. We observe bounded O(1) communication costs as we increase NGPU in all steps except MPI-SORT, which is all-to-all and grows approximately linearly. We choose to implement MPI-SORT as an all-to-all operation to show worst-case runtime when we have no a-priori information about what particles are on what GPUs; in many settings, e.g., timestepping regimes where GPUs are assigned spatially close particles, MPI-SORT can be implemented as a local neighbor communication, reducing communication complexity from O(NGPU ) to O(1). It should be noted that the method parameters (e.g., rc , ξ, h) were not tuned for each NGPU case, but rather fixed to their single-GPU values. The FFT runtime begins to dominate as the number of GPUs grows. On the GH200, inter-GPU communication is the bottleneck, limited by the RDMA GPUDi-

23

FAST EWALD SUMMATION FOR STOKES FLOW

Table 11 Runtime in milliseconds for multi-GPU weak scaling test with NGPU NVIDIA H200 GPUs, MPI communication, and NVIDIA’s cuFFTMp library. Parameters are as in Figure 9 with N = 4 × 106 NGPU . Note that the FFT stage comprises 12 real-to-complex distributed FFTs while the IFFT stage comprises only 3 complex-to-real distributed inverse FFTs.

Step

NGPU = 1

2

4

8

16

32

64

P2P P2G FFT CNV IFFT G2P

227.4 45.9 48.8 18.1 10.7 33.3

232.2 53.6 190.4 25.4 47.4 35.2

232.2 51.3 245.9 24.8 61.5 34.0

234.2 51.6 278.3 24.6 69.6 32.7

235.1 49.9 299.0 24.2 74.6 32.0

231.5 50.5 307.0 24.3 76.6 31.4

229.9 45.0 330.2 23.4 83.2 30.5

— — —

16.5 7.6 6.9

23.7 8.1 6.2

37.1 8.0 6.1

61.3 8.0 5.7

110.4 8.0 5.5

222.0 7.2 2.7

MPI-SORT MPI-GHOST-SOURCE MPI-GHOST-GRID

rect off-chip over NDR InfiniBand (400 Gbit/s). We estimate the communication cost for the FFT stage, which comprises 12 real-to-complex distributed FFTs, as Tcomm ≈ 12 × (1 − 1/NGPU ) × Ng/2NGPU × 16 B/bcomm , where bcomm is the bandwidth. The infiniband setup on TACC’s Vista computer has bcomm = 50 GB/s and our test problem parameters yield Ng = 256NGPU × 692 × 692. As such, the FFTs achieve 70% of this peak communication model, which is in line with NVIDIA’s results [8]. Note that there are multiple strategies to improve FFT performance, such as tuning the near-field cell size s to shift some cost into the P2P computations; however, we kept this parameter fixed to highlight communication overheads. It should be noted that the P2P computations are independent of the FFTs and can be overlapped with them, but we report each step independently to report isolated costs per step. 7. Conclusions. We introduced algorithms designed for the spectral Stokes Ewald sum, accompanied by a comprehensive analysis and evaluation of their performance. In particular, we compared different algorithmic variants as well as the performance of the method on different hardware and for different particle distributions and precisions. The MPI results served primarily as a preliminary glimpse, highlighting potential overheads associated with off-node operations. In future research, we will study the optimization of CPU kernels and thorough performance characterization and optimization of the MPI implementation, particularly in multiGPU NUMA architectures. We also note that novel Ewald decompositions based on prolate spheroidal wave functions (PSWFs) have the potential to further reduce the computational cost [1] compared to the Hasimoto decomposition that was used here. We plan to introduce such a PSWF decomposition into our algorithm in the future. Acknowledgments. The authors acknowledge the Texas Advanced Computing Center (TACC)3 at The University of Texas at Austin and the Oregon Advanced Computing Institute for Science and Society (OACISS)4 at The University of Oregon for providing HPC resources that have contributed to the research results reported within this paper. JB was supported in part by the Peter O’Donnell Jr. Postdoctoral Fellowship at the Oden Institute. 3 URL: http://www.tacc.utexas.edu 4 URL: https://blogs.uoregon.edu/oaciss/

24

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS REFERENCES

[1] L. af Klinteberg, L. Greengard, S. Jiang, and A.-K. Tornberg, Fast summation of stokes potentials using a new kernel-splitting in the dmk framework, 2025, https://arxiv.org/abs/ 2509.21471, https://arxiv.org/abs/2509.21471. [2] L. af Klinteberg and A.-K. Tornberg, A fast integral equation method for solid particles in viscous flow using quadrature by expansion, Journal of Computational Physics, 326 (2016), pp. 420–445, https://doi.org/10.1016/j.jcp.2016.09.006. [3] N. Al Awar, S. Zhu, G. Biros, and M. Gligoric, A performance portability framework for python, in Proceedings of the 35th ACM International Conference on Supercomputing, ICS 21, New York, NY, USA, June 2021, Association for Computing Machinery, p. 467478, https://doi.org/10.1145/3447818.3460376, https://doi.org/10.1145/3447818.3460376. [4] J. Bagge and A.-K. Tornberg, Highly accurate special quadrature methods for stokesian particle suspensions in confined geometries, International Journal for Numerical Methods in Fluids, 93 (2021), pp. 2175–2224, https://doi.org/https://doi.org/10.1002/fld.4970, https://onlinelibrary.wiley.com/doi/abs/10.1002/fld.4970, https://arxiv.org/abs/https:// onlinelibrary.wiley.com/doi/pdf/10.1002/fld.4970. [5] J. Bagge and A.-K. Tornberg, Fast ewald summation for stokes flow with arbitrary periodicity, Journal of Computational Physics, 493 (2023), p. 112473, https://doi.org/10.1016/j. jcp.2023.112473. [6] A. H. Barnett, J. Magland, and L. af Klinteberg, A parallel nonuniform fast fourier transform library based on an exponential of semicircle” kernel, SIAM Journal on Scientific Computing, 41 (2019), pp. C479–C504, https://doi.org/10.1137/18M120885X. [7] D. A. Beckingsale, J. Burmark, R. Hornung, H. Jones, W. Killian, A. J. Kunen, O. Pearce, P. Robinson, B. S. Ryujin, and T. R. W. Scogland, Raja: Portable performance for large-scale scientific applications, in IEEE/ACM International Workshop on Performance, Portability and Productivity in HPC (P3HPC), 2019. [8] L. Cambier, D. Pan, and L. Ligowski, Multinode Multi-GPU: Using NVIDIA cuFFTMp FFTs at Scale, 2022, https://developer.nvidia.com/blog/ multinode-multi-gpu-using-nvidia-cufftmp-ffts-at-scale/ (accessed 2025-10-23). Blog post. [9] H. C. Edwards, C. R. Trott, and D. Sunderland, Kokkos: Enabling manycore performance portability through polymorphic memory access patterns, Journal of Parallel and Distributed Computing, 74 (2014), pp. 3202 – 3216, https://doi.org/https://doi.org/10.1016/ j.jpdc.2014.07.003, http://www.sciencedirect.com/science/article/pii/S0743731514001257. Domain-Specific Languages and High-Level Frameworks for High-Performance Computing. [10] A. M. Fiore, F. B. Usabiaga, A. Donev, and J. W. Swan, Rapid sampling of stochastic displacements in Brownian dynamics simulations, Journal of Chemical Physics, 146 (2017), p. 124116, https://doi.org/10.1063/1.4978242. [11] A. Gholami, D. Malhotra, H. Sundar, and G. Biros, Fft, fmm, or multigrid? a comparative study of state-of-the-art poisson solvers for uniform and nonuniform grids in the unit cube, SIAM Journal on Scientific Computing, 38 (2016), pp. C280–C306, https://doi.org/ 10.1137/15M1010798, http://dx.doi.org/10.1137/15M1010798. [12] L. Greengard and V. Rokhlin, A fast algorithm for particle simulations, Journal of Computational Physics, 73 (1987), pp. 325–348, https://doi.org/https: //doi.org/10.1016/0021-9991(87)90140-9, https://www.sciencedirect.com/science/article/ pii/0021999187901409. [13] J. S. Guasto, R. Rusconi, and R. Stocker, Fluid Mechanics of Planktonic Microorganisms, Annual Review of Fluid Mechanics, 44 (2012), pp. 373–400, https://doi.org/10.1146/ annurev-fluid-120710-101156. [14] R. Halver, J. H. Meinke, and G. Sutmann, Examining performance portability with kokkos for an ewald sum coulomb solver, in Parallel Processing and Applied Mathematics, R. Wyrzykowski, E. Deelman, J. Dongarra, and K. Karczewski, eds., Cham, 2020, Springer International Publishing, pp. 35–45. [15] M. Harvey and G. De Fabritiis, An implementation of the smooth particle mesh Ewald method on GPU hardware, Journal of chemical theory and computation, 5 (2009), pp. 2371– 2377. [16] H. Hasimoto, On the periodic fundamental solutions of the stokes equations and their application to viscous flow past a cubic array of spheres, Journal of Fluid Mechanics, 5 (1959), pp. 317–328. [17] S. H. Holm, Z. Zhang, J. P. Beech, G. Gompper, D. A. Fedosov, and J. O. Tegenfeldt, Microfluidic particle sorting in concentrated erythrocyte suspensions, Physical Review Ap-

FAST EWALD SUMMATION FOR STOKES FLOW

25

plied, 12 (2019), p. 014051. [18] Á. Jász, Á. Rák, I. Ladjánszki, and G. Cserey, Classical molecular dynamics on graphics processing unit architectures, Wiley Interdisciplinary Reviews: Computational Molecular Science, 10 (2020), p. e1444. [19] S. Jiang and L. Greengard, A dual-space multilevel kernel-splitting framework for discrete and continuous convolution, Communications on Pure and Applied Mathematics, 78 (2025), p. 10861143, https://doi.org/10.1002/cpa.22240. [20] D. R. Ladiges, A. Nonaka, K. Klymko, G. C. Moore, J. B. Bell, S. P. Carney, A. L. Garcia, S. R. Natesh, and A. Donev, Discrete ion stochastic continuum overdamped solvent algorithm for modeling electrolytes, Physical Review Fluids, 6 (2021), p. 044309, https://doi.org/10.1103/PhysRevFluids.6.044309. [21] P.-K. Lai, J. W. Swan, and B. L. Trout, Calculation of therapeutic antibody viscosity with coarse-grained models, hydrodynamic calculations and machine learning-based parameters, mAbs, 13 (2021), p. 1907882, https://doi.org/10.1080/19420862.2021.1907882. [22] I. Lashuk, A. Chandramowlishwaran, T.-A. N. H. Langston, R. Sampath, A. Shringarpure, R. Vuduc, D. Z. L. Ying, and G. Biros, A massively parallel adaptive fast-multipole method on heterogeneous architectures, in SC ’09: Proceedings of the 2009 ACM/IEEE conference on Supercomputing, Piscataway, NJ, USA, 2009, IEEE Press, pp. 1–12. [23] D. Lindbo and A.-K. Tornberg, Spectrally accurate fast summation for periodic Stokes potentials, Journal of Computational Physics, 229 (2010), pp. 8994–9010, https://doi.org/10. 1016/j.jcp.2010.08.026. [24] D. Lindbo and A.-K. Tornberg, Spectral accuracy in fast ewald-based methods for particle simulations, Journal of Computational Physics, 230 (2011), p. 87448761, https://doi.org/ 10.1016/j.jcp.2011.08.022. [25] S. Mayani, V. Montanaro, A. Cerfon, M. Frey, S. Muralikrishnan, and A. Adelmann, A massively parallel performance portable free-space spectral poisson solver, ACM Trans. Math. Softw., 51 (2025), pp. 15:1–15:23, https://doi.org/10.1145/3748815. [26] Y. Nakatsukasa, O. Sète, and L. N. Trefethen, The aaa algorithm for rational approximation, SIAM Journal on Scientific Computing, 40 (2018), pp. A1494–A1522, https://doi.org/ 10.1137/16M1106122, https://doi.org/10.1137/16M1106122, https://arxiv.org/abs/https: //doi.org/10.1137/16M1106122. [27] Y. Nakatsukasa and L. N. Trefethen, An algorithm for real and complex rational minimax approximation, SIAM Journal on Scientific Computing, 42 (2020), pp. A3157–A3179, https: //doi.org/10.1137/19M1281897, https://doi.org/10.1137/19M1281897, https://arxiv.org/ abs/https://doi.org/10.1137/19M1281897. [28] Nsight compute documentation nsightcompute 12.8 documentation, https://docs.nvidia.com/ nsight-compute/. [29] R. Okuta, Y. Unno, D. Nishino, S. Hido, and C. Loomis, Cupy: A numpy-compatible library for nvidia gpu calculations, in Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Thirty-first Annual Conference on Neural Information Processing Systems (NIPS), 2017, http://learningsys.org/nips17/assets/papers/paper 16.pdf. [30] J. C. Phillips, D. J. Hardy, J. D. Maia, J. E. Stone, J. V. Ribeiro, R. C. Bernardi, R. Buch, G. Fiorin, J. Hénin, W. Jiang, et al., Scalable molecular dynamics on CPU and GPU architectures with NAMD, The Journal of chemical physics, 153 (2020). [31] J. C. Phillips, Y. Sun, N. Jain, E. J. Bohm, and L. V. Kalé, Mapping to irregular torus topologies and other techniques for petascale biomolecular simulation, in SC’14: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, IEEE, 2014, pp. 81–91. [32] C. Pozrikidis, Interfacial dynamics for Stokes flow, Journal of Computational Physics, 169 (2001), pp. 250–301. [33] S. Pll, A. Zhmurov, P. Bauer, M. Abraham, M. Lundborg, A. Gray, B. Hess, and E. Lindahl, Heterogeneous parallelization and acceleration of molecular dynamics simulations in gromacs, The Journal of Chemical Physics, 153 (2020), p. 134110, https://doi.org/10. 1063/5.0018516, https://doi.org/10.1063/5.0018516, https://arxiv.org/abs/https://pubs. aip.org/aip/jcp/article-pdf/doi/10.1063/5.0018516/16736127/134110 1 online.pdf. [34] D. Saffar Shamshirgar, J. Bagge, and A.-K. Tornberg, Fast Ewald summation for electrostatic potentials with arbitrary periodicity, Journal of Chemical Physics, 154 (2021), p. 164109, https://doi.org/10.1063/5.0044895. [35] R. Salomon-Ferrer, A. W. Gotz, D. Poole, S. Le Grand, and R. C. Walker, Routine microsecond molecular dynamics simulations with amber on gpus. 2. explicit solvent particle mesh ewald, Journal of chemical theory and computation, 9 (2013), pp. 3878–3888.

26

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

[36] I. Sbalzarini, J. Walther, M. Bergdorf, S. Hieber, E. Kotsalis, and P. Koumoutsakos, Ppm a highly efficient parallel particlemesh library for the simulation of continuum systems, Journal of Computational Physics, 215 (2006), pp. 566–588, https://doi.org/https: //doi.org/10.1016/j.jcp.2005.11.017, https://www.sciencedirect.com/science/article/pii/ S002199910500505X. [37] D. S. Shamshirgar, J. Bagge, and A.-K. Tornberg, Fast Ewald summation for electrostatic potentials with arbitrary periodicity, The Journal of Chemical Physics, 154 (2021). [38] D. E. Shaw, P. J. Adams, A. Azaria, J. A. Bank, B. Batson, A. Bell, M. Bergdorf, J. Bhatt, J. A. Butts, T. Correia, et al., Anton 3: twenty microseconds of molecular dynamics simulation before lunch, in Proceedings of the international conference for high performance computing, networking, storage and analysis, 2021, pp. 1–11. [39] D. E. Shaw, M. M. Deneroff, R. O. Dror, J. S. Kuskin, R. H. Larson, J. K. Salmon, C. Young, B. Batson, K. J. Bowers, J. C. Chao, et al., Anton, a special-purpose machine for molecular dynamics simulation, Communications of the ACM, 51 (2008), pp. 91–97. [40] S. S. Shende and A. D. Malony, The tau parallel performance system, The International Journal of High Performance Computing Applications, 20 (2006), pp. 287–311, https: //doi.org/10.1177/1094342006064482, https://doi.org/10.1177/1094342006064482, https: //arxiv.org/abs/https://doi.org/10.1177/1094342006064482. [41] Y.-h. Shih, G. Wright, J. Andn, J. Blaschke, and A. H. Barnett, cufinufft: a load-balanced gpu library for general-purpose nonuniform ffts, in 2021 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW), June 2021, p. 688697, https://doi.org/10.1109/IPDPSW52791.2021.00105, https://ieeexplore. ieee.org/document/9460591/?arnumber=9460591. [42] T. M. Squires and S. R. Quake, Microfluidics: Fluid physics at the nanoliter scale, Reviews of Modern Physics, 77 (2005), pp. 977–1026, https://doi.org/10.1103/RevModPhys.77.977. [43] J. E. Stone, D. Gohara, and G. Shi, Opencl: A parallel programming standard for heterogeneous computing systems, Computing in Science & Engineering, 12 (2010), pp. 66–73, https://doi.org/10.1109/MCSE.2010.69. [44] A. P. Thompson, H. M. Aktulga, R. Berger, D. S. Bolintineanu, W. M. Brown, P. S. Crozier, P. J. in ’t Veld, A. Kohlmeyer, S. G. Moore, T. D. Nguyen, R. Shan, M. J. Stevens, J. Tranchida, C. Trott, and S. J. Plimpton, LAMMPS - a flexible simulation tool for particle-based materials modeling at the atomic, meso, and continuum scales, Comp. Phys. Comm., 271 (2022), p. 108171, https://doi.org/10.1016/j.cpc.2021.108171. [45] C. Trott, L. Berger-Vergiat, D. Poliakoff, S. Rajamanickam, D. Lebrun-Grandie, J. Madsen, N. Al Awar, M. Gligoric, G. Shipman, and G. Womeldorff, The kokkos ecosystem: Comprehensive performance portability for high performance computing, Computing in Science Engineering, 23 (2021), pp. 10–18, https://doi.org/10.1109/MCSE.2021. 3098509. [46] C. R. Trott, D. Lebrun-Grandi, D. Arndt, J. Ciesko, V. Dang, N. Ellingwood, R. Gayatri, E. Harvey, D. S. Hollman, D. Ibanez, N. Liber, J. Madsen, J. Miles, D. Poliakoff, A. Powell, S. Rajamanickam, M. Simberg, D. Sunderland, B. Turcksin, and J. Wilke, Kokkos 3: Programming model extensions for the exascale era, IEEE Transactions on Parallel and Distributed Systems, 33 (2022), pp. 805–817, https: //doi.org/10.1109/TPDS.2021.3097283. [47] F. Vico, L. Greengard, and M. Ferrando, Fast convolution with free-space Green’s functions, Journal of Computational Physics, 323 (2016), pp. 191–203, https://doi.org/10.1016/ j.jcp.2016.07.028. [48] V. Volkov, Understanding latency hiding on GPUs, University of California, Berkeley, 2016. [49] M. Wang and J. F. Brady, Spectral Ewald Acceleration of Stokesian Dynamics for polydisperse suspensions, Journal of Computational Physics, 306 (2016), pp. 443–477, https: //doi.org/10.1016/j.jcp.2015.11.042. [50] T. Wang, R. Yokota, and L. A. Barba, Exafmm: a high-performance fast multipole method library with c++ and python interfaces, Journal of Open Source Software, 6 (2021), p. 3145. [51] W. Yan, E. Corona, D. Malhotra, S. Veerapaneni, and M. Shelley, A scalable computational platform for particulate stokes suspensions, Journal of Computational Physics, (2020), p. 109524. [52] H. Zhao, A. Isfahani, L. Olson, and J. Freund, A Spectral Boundary Integral Method for Flowing Blood Cells, Journal of Computational Physics, (2010). [53] A. Zinchenko and R. H. Davis, Large-scale simulations of concentrated emulsion flows, Philosophical Transactions Of The Royal Society Of London Series A-Mathematical Physical And Engineering Sciences, 361 (2003), pp. 813–845.

FAST EWALD SUMMATION FOR STOKES FLOW

27

import parkipy ex = parkipy . utils . g e t _ e x e c u t i o n _ s p a c e ( " CUDA " ) # Kokkos backend am = parkipy . utils . g e t _ a r r a y_ m o d u l e ( ex ) # ‘ cupy ’ or ‘ numpy ’, cond . on ex # set s p e c t r a l Ewald box and t o l e r a n c e box = [ 1 .0 , 1 .0 , 1 . 0 ] tol = 1e - 4 # g e n e r a t e sources , targets , densities , and normals nt = 4000000 ns = 4000000 trg = am . random . rand (3 , nt ) * am . array ( box ) . reshape (3 , 1 ) src = am . random . rand (3 , ns ) * am . array ( box ) . reshape (3 , 1 ) dens_sl = am . random . randn (3 , ns ) dens_dl = am . random . randn (3 , ns ) norms = am . random . randn (3 , ns ) dens = am . vstack (( dens_sl , dens_dl ) ) # stack d e n s i t i e s for ewald call # get the p o t e n t i a l for the Stokes single and double layer p o t e n t i a l options = parkipy . ewald . EwaldOptions ( periodicity =1 , box = box , tolerance = tol , cell_size = 224 , ex ec ut i on _s pa c e = ex ) pot = parkipy . ewald . stokes_comb ( trg , src , dens , norms , options ) Fig. 10. Code using the ParkiPy library to compute the combined Stokes single and double layer potential. This example runs on NVIDIA GPUs and uses the cupy array module as a backend.

Appendix A. The ParkiPy Package. The algorithms presented in this paper are implemented in the (Par)ticle (k)ernel (i)nteractions for (Py)thon package, ParkiPy. The package offers easy-to-use and architecture-agnostic APIs to evaluate the Stokes and Poisson kernels in arbitrary periodicity using Ewald summation. ParkiPy’s Ewald summation submodule, parkipy.ewald, currently supports the Ewald summation for the Stokes single layer potential, the combined Stokes single and double layer potential (Figure 10) and the Poisson kernel. There is also a distributed submodule, parkipy.distributed.ewald, which computes the singly periodic combined Stokes potential for multi-node multi-gpu settings in a CUDA execution space. The main namespace has a class parkipy.CellList for efficient cell list construction. Appendix B. Non-uniform Particle Distributions. The performance of the Ewald sum becomes suboptimal for highly non-uniform distributions of the sources and targets. We compare cases where particles are uniformly distributed in the unit cube (U ), normally distributed with variance 0.3 shifted and truncated to the unit cube (N ), and uniformly distributed on the surface of the unit sphere (S ) (see Figure 11). When both sources and targets are distributed as N , our P2P algorithm achieves 68% of the performance relative to a source and target distribution as U , while our P2G algorithm achieves 61% of the performance. When both sources and targets are distributed as S , our P2P algorithm achieves 27% performance and our P2G algorithm 80% performance, again relative to the performance with sources and targets distributed as U . The P2P algorithm performs close to the ideal distribution of U when the sources are distributed as N or S but the targets are distributed as U , as P2P is parallel over target cell lists and assumes a consistent number of points per target cell. The G2P algorithm performs best when the targets are distributed as S , about a 1.7× speedup compared to a distribution of U , as more

28

G. KOSMACHER, Z. DU, AND J. BAGGE, G. BIROS

Particle Distributions Projected to the x-z Plane Uniform (U)

1.0

1.0

Truncated Gaussian (N)

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

0.0

0.0

0.0 0.0

0.2

0.4

0.6

0.8

1.0

400

600

800

1000

1200

Density

Sphere (S)

1.0

0.0

0.2

500

0.4

0.6

1000

Density

0.8

1500

1.0

0.0

0.2

250

0.4

500

0.6

750

0.8

1000

1.0

1250

Density

Fig. 11. N = 4×106 particles projected onto the x-z plane. Densities estimated with a Gaussian KDE. Left: particles distributed uniformly in the unit cube. Middle: particles distributed normally with variance 0.3. Right: particles distributed uniformly on the surface of the unit sphere.

targets read the same grid points, increasing cache locality. In our worst case, where both sources and targets are distributed as S , the total Ewald sum achieves 38% of the performance of the best case where sources and targets are distributed as U .