Improving Performance by Vectorizing Particle-in-Cell Codes

A Practical Guide

Bei Wang, HPC Software Engineer, Princeton University; Carlos Rosales-Fernandez, Software Technical Consulting Engineer, Intel Corporation; and William Tang, Professor, Princeton Plasma Physics Laboratory

The basic particle method is a well-established approach to simulating the behavior of charged particles interacting with each other through pairwise electromagnetic forces. At each step, the particle properties are updated according to these calculated forces. For applications on powerful modern supercomputers with deep cache hierarchies, a pure particle method is efficient with respect to both locality and arithmetic intensity (compute-bound).

Unfortunately, the O(N²) complexity makes a particle method impractical for plasma simulations using millions of particles per process. Instead of calculating O(N²) forces, the particle-in-cell (PIC) method uses a grid as the medium to calculate long-range electromagnetic forces. This reduces the complexity from O(N²) to O(N + M log M), where M is the number of grid points, generally much smaller than N. However, achieving high parallel and architectural efficiency is a significant challenge for PIC methods due to the gather/scatter nature of the algorithm.

Attaining performance becomes even more complex as HPC technology moves to the era of multi- and many-core architectures with increased thread and vector parallelism on shared memory processors. A deep understanding of how to improve the associated scalability will have a wide-ranging influence on numerous physical applications that use particle-mesh algorithms―including molecular dynamics, cosmology, accelerator physics, and plasma physics.

This article is a practical guide to improving performance by enabling vectorization for PIC codes.

Optimization for PIC Codes

The code example we use for this demonstration is the particle class in Athena++, an astrophysical magnetohydrodynamics (MHD) code written in C++ ¹. The particle class encapsulates basic data structures and functions in PIC methods. Particle properties are represented by phase space position―that is, physical space position (x1,x2,x3) and velocity (v1,v2,v3). The class functions implement three essential particle-based operations in PIC methods:

Data Layout and Alignment

Particle properties can be stored as array-of-structures (AoS) or structure-of-arrays (SoA) data layout (Figure 1).

AoS helps to pack and unpack particles for shift, but doesn’t help to enable vectorization for deposit and move with stride-one memory access. Since deposit and move are the hotspots of the code, we choose SoA data layout for particle representation. Specifically, we allocate the memory with alignment using the posix_memalign() function (Figure 2).

Figure 1 – Particle representation using AoS and SoA layout

Figure 2 – Allocating the memory with alignment using the posix_memalign() function

We start the optimization by checking the vectorization report of the original code for deposit.cpp and move.cpp (Figure 3).

Figure 3 – Original code for deposit.cpp (left) and move.cpp (right)

We compile the code with the Intel® C++ Compiler 19.0 and the following options: -O3 -g -qopt -report5
-xCORE-AVX512 -qopenmp simd. By default, the compiler will try to vectorize the innermost loop for both deposit (L106) and move (L121). From the vectorization report, we see that the compiler fails to do so because of a potential data dependency in deposit.cpp (Figure 4) and efficiency reasons in move.cpp (Figure 5).

Figure 4 – Vectorization report for deposit.cpp

Figure 5 – Vectorization report for move.cpp

Explicit SIMD via OpenMP Directives

For both deposit and move, the loop count for the outermost loop (e.g., nparticle loop) is large, while the loop count for the innermost loop is small. Therefore, the outermost loop is a more suitable candidate for vectorization. Our first attempt is to enable vectorization for the outermost loop in both deposit and move using OpenMP* SIMD directives. For both kernels, there’s no data dependency with respect to pointer aliasing. In considering vectorization for deposit in the outermost loop, a potential data dependency will appear when two particles within a single vector length try to write on the same memory location for the grid-based array (e.g., mcoup). When compiling for AVX2, it’s important to use #pragma omp ordered simd such that the loop can be safely vectorized. Since AVX512* provides conflict detection instructions (vpconflict), this is no longer necessary for AVX512. Also, we use aligned and simdlen clauses in the OpenMP SIMD directives to have the compiler generate more efficient vector instructions, e.g.:

Note the current Intel C++ Compiler only allows a local variable name in the aligned clause list. Using a class data member name directly will cause a compiler error. Our solution here is to define local variables and assign class data member variables to local variables, e.g.:

When we check the vectorization report after adding the OpenMP SIMD directives, we see that the compiler has now successfully vectorized the outermost loop with aligned access to the particle data (Figure 6). Unfortunately, the vectorized version is quite inefficient and doesn’t give a speedup over the scalar version due to the gather/scatter nature of the kernels. For example, the estimated potential speedup is only 0.55 and 0.67 for deposit and move, respectively.

Figure 6 – Vectorization report for loops in deposit (after adding OpenMP SIMD directives)

Strip Mining

Strip mining is a common technique to help the compiler automatically vectorize code by exposing data parallelism². The idea is to transform a single loop into a nested loop where the outer loop strides through a strip and the inner loop strides all iterations within the strip (Figure 7).

Usually, the strip size is a multiple of the vector length. In some computationally-intensive loops, the strip mining technique alone can lead to a significant performance boost. For deposit and move, replacing the particle loop with two nested loops doesn’t help with performance due to the gather/scatter. After applying strip mining, for example, the estimated potential speedup is still only 0.52 and 0.67 for deposit and move, respectively (Figure 8).

Figure 7 – Strip mining for particle loop

Figure 8 – Vectorization report for deposit and move after applying strip mining

However, strip mining can be a powerful technique when you use it as a basis for other techniques, such as loop fission.

Loop Fission for move

The move loop includes two operations:

While the former won’t benefit from vector instructions, the latter will. The idea here is to use loop fission to separate the non-vectorizable operation from the vectorizable one to improve performance. At the same time, by applying strip mining with loop fission, we can significantly reduce the storage requirement for passing local variables from the first loop to the second. Figure 9 shows the new move implementation.

The vectorization report shows that although there is no potential speedup for the first loop, the second achieves a 6.59 estimated potential speedup (Figure 10).

Vectorizable Charge Deposition Algorithm for deposit

Compared with move, deposit is more challenging to optimize. For the scatter operation, the performance challenges involve not only random memory accesses and potential data conflicts, but also increased pressure on the memory subsystem. Also, unlike move, which also includes many independent computationally-intensive operations (i.e., updating particles properties), deposit is mostly memory operations. A key optimization here is to reduce (or even avoid) data conflicts and regularize memory accesses. Motivated by the portable SIMD charge deposition algorithm in the PICSAR* code³, we implemented a vectorizable charge deposition algorithm for deposit (Figure 11).

Figure 9 – New move implementation with strip mining and loop fission 

Figure 10 – Vectorization report for the revised move code

First, the outermost particle loop is changed to two nested loops using strip mining. In the first nested loop (L69-L133), the particle properties are deposited on a set of local arrays with stride-one memory access, so the compiler should have no trouble generating highly vectorized machine code.

In the second nested loop (L135-L159), the information saved in the local arrays is transferred to an extended global grid array (L163-L205), where each grid point is associated with particle deposition on surrounding grid points (i.e., a 27-point stencil). The price we pay here is the extra storage required for the extended global array. The storage can be reduced if we consider the 27-point stencil as three 9-point stencils.

The grid-based values in the extended global array need to be merged at the end. Though it will be difficult to vectorize the merging operation, the cost isn’t significant, since the loop is over grid indices at least an order of magnitude smaller than the particle one.

Figure 11 – Vectorizable charge deposition algorithm for deposit

Figure 12 shows the vectorization report for the new implementation. The first and second loops now achieve 5.72 and 5.69 estimated potential speedup, respectively.

Performance Evaluation on Intel Architectures

Performance was evaluated on a 10x10x10 test grid with 100 particles per cell. We ran the simulations for 1,001 timesteps and measured the total wall clock time. We measured performance on single-core Intel® Xeon® and Intel® Xeon Phi™ processors: “BDW” (Intel® Xeon® E5-2680 v4 processor, 2.4GHz, 2 sockets, 14 cores), “SKX” (Intel® Xeon® Gold 6148 processor, 2.40GHz, 2 sockets, 20 cores), and “KNL” (Intel® Xeon Phi™ 7250 processor, 1.4GHz, 1 sockets, 68 cores) (Figure 13).

Figure 12 – Vectorization report for deposit

Figure 13 – Performance comparison of various optimization techniques on Intel “BDW,” “SKX,” and “KNL” processors. The optimization techniques include SIMD, strip mining (SM), loop fission for move (LF), and vectorizable charge deposition (VD). For each case, the reported number is the averaged time of multiple runs on a dedicated system.

We can see that the actual performance is consistent with the estimated potential speedup from the compiler report. On BDW, we see significant performance overhead in using SIMD for deposit. This is due to the #pragma omp ordered simd clause ensuring the right order of writing to the grid array when compiling with AVX2 instructions. Overall, we see that strip mining, combined with other techniques, has led to 1.3x to 1.9x performance boost on Intel® processors.

Acknowledgements

This work was supported by the Intel® Parallel Computing Center (IPCC) program. We would like to thank Jason Sewall at Intel for providing strong technical support for this work. We also would like to thank Lev Arzamasskiy and Matthew Kunz in the Princeton University Department of Astrophysical Sciences for sharing their source code.