Skip to content

GPU implementation

Three-layer parallelisation

gevolution 2.0 combines MPI, OpenMP, and CUDA:

flowchart TD
    A[Simulation volume] --> B[MPI rod decomposition]
    B --> C[Rank-local host memory]
    B --> D[Rank-local GPU memory]
    C --> E[OpenMP Fourier loops]
    C --> F[File I/O]
    D --> G[Particle projection, kick, drift]
    D --> H[Real-space mesh updates]
    D --> I[cuFFT stages]
    D <--> J[CUDA-aware MPI or host-staged MPI]

The preferred production layout is one MPI rank per GPU. OpenMP allows each rank to use its associated CPU cores without increasing the number of MPI endpoints.

Domain decomposition

LATfield2 cuts the cubic volume along two dimensions, creating rod-shaped domains. The \(x\) coordinate is the fast, contiguous, rank-local direction. Halo cells hold values needed for operations near neighbouring domains.

The process grid is two-dimensional and requires at least 2 x 2 ranks. Neighbour communication assumes that a neighbour is a distinct MPI rank.

Memory model

Particle-mesh operations and real-space mesh updates run on the GPU. Particle positions, momenta, IDs, row metadata, and many real-space fields therefore reside in device high-bandwidth memory.

Fourier-space fields and some operations can reside in host memory. Bulk device/host transfers occur at controlled stages of FFTs and I/O. GPU memory is usually the first capacity constraint for large runs.

Particle handler

The GPU particle handler replaces the linked-list organisation used by older CPU versions. Particles are densely packed in separate arrays for positions, momenta, and IDs.

Particles are binned by mesh row and sorted by \(x\) within each row. Each CUDA thread block processes one row. This improves data locality for cloud-in-cell projection and metric interpolation.

After a drift:

  1. particles leaving the local domain are sorted into special exchange rows;
  2. those rows are communicated to neighbouring ranks;
  3. incoming particles are merged and sorted into their new rows.

The implementation uses cub::DeviceRadixSort and cub::DeviceSegmentedSort.

Projection, kick, and drift

Projection deposits particle quantities onto nearby mesh points. Halo values carry spill-over to adjacent ranks and are exchanged after projection.

Kick updates each particle momentum independently from interpolated metric fields. It requires no particle exchange and is usually very fast.

Drift updates particle positions. It is followed by sorting and MPI exchange for particles crossing domain boundaries.

Callable structs replace CPU function-pointer callbacks so the compiler can instantiate GPU kernels for each operation.

Mesh updates and reductions

The GPU-enabled LATfield2 backend provides lattice_for_each, a generic CUDA kernel templated on a callable struct. It replaces serial Site::next() loops for real-space mesh work and supports device reductions.

Fourier-space host loops use OpenMP where appropriate.

Distributed FFT pipeline

LATfield2 decomposes a 3D FFT into local 1D transforms separated by distributed transpositions. Most 1D transforms use cuFFT. Where the first real-to-complex stage cannot meet cuFFT alignment requirements, the implementation falls back to FFTW rather than copying to a fresh aligned array.

For multi-component fields, transforms and transpositions are overlapped where possible. CUDA-aware MPI passes device buffers directly when explicitly enabled on a supported MPI stack. Otherwise device buffers are copied to host staging buffers before MPI and copied back afterwards. At scale, MPI transpositions commonly dominate FFT cost.

MPI communication modes

The runtime communication mode is selected from environment-variable hints; it is not selected by a gevolution compiler definition.

Mode Selection Behaviour
CUDA-aware MPI A recognised support hint is enabled and LATFIELD2_DISABLE_CUDA_AWARE_MPI is not enabled Device pointers are passed directly to MPI
Host-staged MPI No support hint is enabled, or LATFIELD2_DISABLE_CUDA_AWARE_MPI=1 Device data is copied into host buffers around MPI calls

Host-staged fallbacks are implemented for:

  • LATfield2 FFT transpositions, halo updates, and particle exchanges;
  • particle light-cone ID-log routing;
  • balanced particle light-cone output;
  • HEALPix neighbour exchange and write-buffer aggregation.

Blocking light-cone neighbour and ID exchanges use temporary host buffers. Nonblocking balanced particle and HEALPix output keeps host send buffers alive until MPI_Waitall completes.

The support hints are trusted declarations rather than active MPI capability tests. See CUDA-aware MPI selection before enabling direct device-pointer MPI.

I/O staging

Particle and HEALPix pixel output is packed on the device. Buffers are then moved in bulk to host memory for file output. When CUDA-aware MPI is disabled, inter-rank output communication is also staged through host memory. Large output schedules therefore increase device and host memory pressure as well as filesystem traffic.

Performance bottlenecks

  • Distributed FFT transpositions and network topology.
  • Parallel filesystem bandwidth for large snapshots and light cones.
  • GPU memory capacity and particle imbalance.
  • Non-CUDA-aware MPI paths that stage communication through host memory.
  • Poor CPU, GPU, or NUMA binding.