Accelerating cosmological simulations on GPUs: a portable approach using OpenMP

TL;DR

The paper tackles the computational bottleneck in cosmological simulations by porting PINOCCHIO's collapse-time kernel to GPUs using OpenMP target directives, complemented by GPU-native cubic spline and bilinear interpolation to replace GSL. The approach delivers portable performance across NVIDIA and AMD platforms, achieving strong single-node speedups (up to 8× on AMD and at least 4× on NVIDIA) and over 80% of FP64 peak efficiency in roofline analyses, while preserving scientific accuracy (collapse statistics and HMF within ~1%). In production-scale runs, the GPU-accelerated kernel reduces time-to-solution by about 6× per run, translating to substantial cumulative savings across thousands of realizations. Overall, the work demonstrates that directive-based OpenMP offloading is a viable, maintainable path to accelerate large-scale cosmological codes on heterogeneous HPC hardware without bespoke GPU code paths.

Abstract

In this work we present the porting to Graphics Processing Units (GPUs, using OpenMP target directives) and optimization of a key module within the cosmological {\pinocchio} code, a Lagrangian Perturbation Theory (LPT)-based framework widely used for generating dark matter (DM) halo catalogs. Our optimization focuses on a specific segment of the code responsible for calculating the collapse time of each particle involved in the simulation. Due to the embarrassingly parallel nature of this computation, it represents an ideal candidate for GPU offloading. As part of the porting process, we developed fully GPU-native implementations of both cubic spline and bilinear interpolation routines, required for evaluating collapse times. Since GNU Scientific Library (GSL) does not support GPU offloading, these custom implementations run entirely on the GPU and achieve residuals of only $\sim0.003\%$ when compared to the CPU-based implementation of GSL. Comparative benchmarking on the LEONARDO (NVIDIA-based) and SETONIX (AMD-based) supercomputers reveals notable portability and performance, with speedups of~\textit{4x} and up to~\textit{8x}, respectively. While collapse time calculation is not a primary bottleneck in the overall workflow, the acceleration reduces full production runs by $\sim 100$ seconds each leading to a cumulative saving of $\sim 160000$ Standard-h ($\sim28$ hours wall time) across thousands of simulations. Roofline analysis confirms that our GPU porting achieves over 80\% of the theoretical FP64 peak performance, confirming efficient compute-bound execution. This work demonstrates that OpenMP directives offer a portable, effective strategy for accelerating large-scale cosmological simulations on heterogeneous hardware.

Figures (11)

• Figure 1: Overview of the GPU offloading strategy used in the PINOCCHIO collapse time kernel. Host-side memory is allocated and selectively transferred to the GPU using OpenMP data mapping directives. The target kernel, including the custom GPU-native interpolation routines and binary-masked control flow, runs entirely on the device. Only minimal, necessary data is transferred back to the host, improving performance by reducing memory traffic.
• Figure 2: Left panel displays the comparison between the GPU-custom cubic spline interpolation routine, the GSL implementation, and the analytical 1D reference function described in Section \ref{['interpolation_validation']}. A total of only 35 evaluation points are shown for clarity. Right panel shows the histogram of the residuals between the GPU-custom and GSL interpolations, aggregated over the 35 evaluation points. Residuals are shown in absolute units, while the dashed line indicates the average residual expressed as a percentage.
• Figure 3: Comparison of interpolation wall-time between the GPU-native cubic spline routine and the GSL implementation, as a function of the number of evaluation points. GPU timings include memory transfers between the host and device. Speedup values relative to the GSL implementation are reported above each GPU bar.
• Figure 4: Same as Figure \ref{['Cubic_spline_vs_GSL']} but for the bilinear interpolation routines. In this case 50 evaluation points are shown in the left panel.
• Figure 5: Same as Figure \ref{['Cubic_spline_vs_GSL_timing']} but for the bilinear interpolation routines.