AMSS-NCKU Numerical Relativity Code Optimization
Topic Background
General Relativity (GR), formulated by Einstein, describes gravity as the geometric curvature of spacetime and serves as the theoretical cornerstone of modern astrophysics and cosmology. Its core governing equations---the Einstein Field Equations (EFEs)---elucidate the nonlinear coupling between matter distribution and spacetime structure. In the era of gravitational-wave astronomy, the research value of GR has become increasingly prominent.
In the highly dynamic, strong-field regime of BBH mergers, extreme spacetime curvature and relativistic velocities cause the nonlinear terms of the EFEs to dominate, rendering analytical approximation methods (such as post-Newtonian approximations and perturbation theory) invalid. Consequently, the discrete iterative schemes of Numerical Relativity (NR) emerge as the only viable approach for accurately solving the EFEs in these scenarios. Currently, mainstream technical frameworks typically rely on the BSSN formulation or the CCZ4 constraint-damping system under the ADM decomposition to ensure numerical stability. These frameworks utilize high-order finite difference methods coupled with Runge-Kutta time integration for discretization. Furthermore, to bridge the vast scale disparity between the black hole horizon singularity and gravitational wave wavelengths, core algorithms widely integrate Berger-Oliger-style adaptive mesh refinement (AMR) and employ MPI parallel communication strategies based on domain decomposition.
However, these mainstream techniques still face severe challenges when addressing the demands of ultra-high precision and large-scale simulations. The low arithmetic intensity of finite difference operators is heavily bound by memory bandwidth, making it difficult to fully exploit the peak floating-point performance of modern hardware. Within the Berger-Oliger AMR architecture, frequent inter-level interpolations during mesh regridding induce massive MPI communication overhead, making the code highly susceptible to the "communication wall" bottleneck during large-scale parallelization.
This work focuses on the in-depth performance optimization of the open-source numerical relativity code AMSS-NCKU, aiming to resolve computational bottlenecks in high-precision spacetime evolution. By fine-tuning key modules---including the finite difference computational kernels, AMR, and MPI parallel communication---we employ a comprehensive suite of optimization strategies. These include the decoupling of evolution-analysis loops, GPU heterogeneous acceleration, process affinity binding, communication aggregation optimization, 3D interpolation algorithm improvements, recursion elimination, and compiler optimization flag tuning.
Clusters Configuration in AMSS-NCKU Problem
Hardware Configuration:
| Component | Specification |
|---|---|
| CPU | Intel Xeon Gold 6348 Processor 2.60 GHz 28 cores × 2 |
| Memory | DDR4 ECC 3200MT/s 16GB × 16 |
| Disk | 480GB SSD SATA × 1 |
| HCA | Mellanox ConnectX-5 EDR 100 Gb/s InfiniBand HCA |
| Switch | Mellanox SB7800 EDR 100Gb/s InfiniBand Smart Edge Switches |
| Cable | Mellanox MCP1600-E0xxEyy 100Gb/s QSFP28 DAC Cable, 3m |
Software Configuration:
| Component | Specification |
|---|---|
| Compiler | GCC-14.2.0 |
| Operating System | Rocky Linux 8.10 (Green Obsidian) |
| Software | Python-3.13 |
| Performance Analysis Tool | Intel Vtune Profiler |
We adopted the Intel oneAPI 2025.3 suite to leverage its advanced vectorization capabilities and high-performance MPI libraries.
Analysis of the Overall Structure
The overall architecture of the AMSS-NCKU simulation code can be abstracted into three primary stages: data initialization, spacetime evolution-analysis iteration, and result visualization. Data initialization is responsible for generating the initial spacetime configuration and mesh topology, while result visualization handles gravitational waveform extraction and post-processing of spacetime evolution features. Since both stages involve fixed logic frameworks invoked only once or infrequently, the vast majority of computational cost during program execution is concentrated in the spacetime evolution-analysis iteration stage.
During the evolution phase, the main evolution loop is driven by the Evolve function, which invokes the RecursiveStep function at each time step to update the adaptive mesh. Specifically, RecursiveStep employs a top-down hierarchical traversal strategy: it first calls the Step function on the coarsest grid level (Level 0) to perform finite-difference evolution, then recursively proceeds to finer refinement levels until all AMR levels have been traversed.
According to the conservation requirements of the Berger-Oliger AMR algorithm, refined grid levels must obtain interpolated boundary conditions from coarser levels via ghost zones at their boundaries. Conversely, after completing time integration on a refined level, high-accuracy data must be projected back onto the coarser level using a restriction operator. This bidirectional data coupling between hierarchy levels introduces strict temporal dependencies, which constitute the primary bottleneck hindering task parallelization of the evolution process.
Recursion Elimination Strategy Based on Iteration
The deep recursive invocation mechanism of RecursiveStep within the core evolution step is constrained not only by bidirectional data dependencies but also introduces significant system-level overhead. On one hand, in scenarios involving extremely high refinement levels, frequent recursive calls continuously consume substantial stack space, posing a potential risk of stack overflow. On the other hand, the deeply nested recursive structure greatly increases the complexity of performance profiling, making it difficult for hotspot tracing tools to accurately isolate the true computational cost of the underlying finite-difference operators.
Although directly decoupling the spatiotemporal dependencies of this algorithm at the mathematical logic level presents significant challenges, we have adopted a refactoring strategy that replaces traditional recursive calls with explicit iteration (loops). This approach aims to eliminate the aforementioned performance profiling obstacles and reduce runtime system overhead.
By flattening the function call stack, the true computational costs of grid evolution at each underlying level are clearly isolated and revealed. This structural adjustment effectively eliminates the masking effect of the recursive framework on the execution time of core underlying operators.
Task-Based Process Splitting
Further analysis of the execution flow reveals that the original evolution-analysis pipeline employed a strongly serially coupled architecture: upon completing the evolution calculation for each time step, the main loop was forced to block and wait for the numerical analysis module to finish before proceeding to the next step. However, from a physical logic perspective, numerical analysis depends unidirectionally only on historical field data; its post-processing results provide no feedback to subsequent spacetime dynamical evolution. This unnecessary temporal dependency constitutes an artificial synchronization barrier, preventing the full utilization of computational resources.
To address this performance bottleneck, we propose and implement an asynchronous parallelization strategy based on the decoupling of evolution and analysis tasks. The core concept involves domain decomposition via process group mapping: the global MPI process pool is strictly partitioned into an evolution process group and an analysis process group. In terms of the specific mapping strategy, available processes are evenly split by rank, with an offset equal to half the total number of processes used to establish point-to-point communication mappings between the two groups.
At the communication level, a message-passing-based pipeline parallelism mode is constructed between the evolution and analysis processes. Specifically, after completing the RecursiveStep integration for the current time step, the evolution process immediately packs the grid state data designated for analysis and invokes MPI_Send to transmit it to the corresponding bound analysis process. Upon completion of the send operation, the evolution process seamlessly transitions to the computation of the next time step without waiting for the analysis task to return.
In the specific code implementation, a GROUP_ROLE enumeration identifier (divided into EVOLVE and ANALYSIS) was introduced within the Step function. This identifier is injected as a context parameter into the BSSN solver to enable fine-grained branching scheduling of the instruction flow. Furthermore, to ensure topological robustness for heterogeneous task parallelism, the underlying communication domain was refactored to derive a dedicated intra-group communicator, COMM_SOLVER.
Odd-Even Communication Scheme
In the initial phase of architectural refactoring, to minimize code intrusiveness and maintain the logical coherence of the original BSSN solver, we adopted a coarse-grained "half-split block mapping" strategy: global MPI processes were split by rank, with the first half dedicated exclusively to evolution and the second half responsible for analysis.
However, as the simulation scale expands, this topology-agnostic mapping suffers from severe bottlenecks. When the total number of processes exceeds the physical core count of a single node, the evolution and corresponding analysis processes are highly likely to be completely isolated on different physical nodes. In this case, massive intermediate field data must be transmitted entirely across nodes via the InfiniBand network.
To break through the network bandwidth barrier, we adopt a topology-aware communication strategy named "Interleaved Pairing." This strategy abandons block-based partitioning and binds processes at a fine-grained level according to even and odd ranks: adjacent even-ranked processes and odd-ranked processes undertake evolution and analysis tasks respectively.
Given that modern supercomputing nodes are generally equipped with an even number of cores and adopt a NUMA architecture, this interleaved mapping ensures that any data-dependent pair is precisely placed within the same compute node or NUMA domain at full load. As a result, the originally expensive point-to-point message passing is automatically downgraded by the underlying MPI to highly efficient intra-node shared-memory copies.
Distributed File I/O
The code is designed to simulate gravitational wave radiation during binary black hole mergers. It requires real-time tracking of the puncture coordinates of the two black holes throughout the evolution and the calculation of gravitational wave strain amplitudes based on the Weyl scalar Monitor objects are instantiated during the initialization phase of the BSSN solver, responsible for: tracking the trajectories of the binary black hole punctures; extracting
However, after adopting the process group partitioning strategy, the original global output logic faced a fragmentation issue regarding the communication domain: Rank 0 within MPI_COMM_WORLD belongs to only one specific process group and thus cannot collect and unify data output across both groups.
To address this, we refactored the Monitor instantiation logic and the data writing protocol. Specifically, Rank 0 within the Evolution Process Group is now exclusively responsible for writing the binary black hole puncture trajectory file and the ADM constraint residual log. Conversely, Rank 0 within the Analysis Process Group independently handles the output of gravitational wave waveform data and apparent horizon physical quantities.
Optimization of 3D Spatial Interpolation
After establishing the pipeline-parallel architecture for the evolution and analysis tasks, further performance benchmarking revealed severe load imbalance between the two pipeline stages: the evolution process group requires approximately 4 seconds per computational step, whereas the analysis process group incurs a post-processing latency as high as 8 seconds. This 1:2 timing disparity forces the evolution processes to remain blocked for extended periods after completing their current step, awaiting the analysis processes.
Microarchitectural hotspot profiling using Intel VTune pinpointed the core computational bottleneck in the analysis phase to the Fortran-based 3D interpolation routine polin3.
In the original implementation, 3D interpolation was performed by repeatedly invoking the 1D interpolation function polint within nested loops inside polin3. Although mathematically correct, this pointwise invocation approach suffers from several critical performance deficiencies at the microarchitectural level: deep control-flow nesting hinders the compiler's ability to identify vectorization opportunities; frequent creation and destruction of temporary arrays within the inner loops severely degrade cache locality; and function-call overhead scales linearly with grid size, exacerbating inefficiency.
To address the aforementioned dual bottlenecks in memory access and computation, the optimized implementation refactors the scalar pointwise invocations into array-level batch operations, enabling the 1D interpolation operator to process entire contiguous data blocks in a single pass. Leveraging the Fortran compiler's robust automatic vectorization capabilities for array expressions, this restructuring significantly enhances memory access continuity and improves the pipeline utilization of floating-point execution units.
Benchmark results demonstrate that this optimization reduces the analysis phase latency from approximately 8 seconds to 3 seconds, successfully restoring temporal masking with the evolution stage and bringing the total per-step execution time back to roughly 4 seconds.
Exploration of FCFS Communication Scheduling
During program execution, multiple global synchronizations occur, all facilitated by the transfer function to handle inter-process data exchange. The original implementation adopted a "bulk initiation and collective wait" pattern: all non-blocking MPI_Isend and MPI_Irecv requests were initiated simultaneously, and only after all communication operations completed were the received data unpacked and processed in a unified step.
Based on these observations, we introduced a "first-come, first-served" strategy. This approach leverages MPI_Waitany to poll for ready requests; upon detecting the completion of any message reception, it immediately triggers the unpacking and processing workflow.
However, in full-scale testing, this strategy failed to deliver the anticipated performance gains. On one hand, the frequent invocations of MPI_Waitany introduced additional runtime overhead that effectively negated the potential speedup from overlapping computation and communication. On the other hand, under complex communication patterns, this interface proved prone to race conditions or state misjudgments.
Considering stability, maintainability, and actual performance outcomes, we ultimately decided to revert to the original "bulk communication + unified processing" scheme.
Compiler-Level Optimization
In the early stages of the competition, we attempted to enable aggressive optimization flags for the entire program. While this reduced the total runtime by approximately 1 second, it introduced significant numerical errors---with the RMS error reaching as high as 5%, far exceeding the tolerance limits specified in the problem statement. Consequently, this approach was abandoned.
Subsequent stability tests identified that the errors primarily originated from the C/C++ modules. Further analysis revealed that these modules extensively utilize complex C++ language features, which can cause the compiler to generate vectorized code in inappropriate locations when advanced optimizations are enabled, thereby compromising numerical stability.
In contrast, the Fortran code demonstrated high sensitivity to compiler optimizations while maintaining excellent compatibility. Under the premise of ensuring result correctness, we enabled maximum optimization levels for the Fortran modules:
| Category | Flag | Description |
|---|---|---|
| C++ | -O3 | Highest optimization level |
| -Wno-deprecated | Disable deprecated feature warnings | |
| -Dfortran3 | Define preprocessor macro fortran3 | |
| -Dnewc | Define preprocessor macro newc | |
| -DENABLE_GROUP_ROLE | Enable group role functionality | |
| Fortran | -O3 | Highest optimization level |
| -mtune=icelake-server | Tune for Intel Ice Lake server CPUs | |
| -ffast-math | Enable fast floating-point optimizations | |
| -funroll-loops | Enable loop unrolling | |
| -march=native | Optimize for host CPU architecture | |
| -fpp | Enable Fortran preprocessor |
GPU Implementation
Concurrent with the routine optimization of the CPU implementation, we also conducted in-depth explorations and trials regarding GPU acceleration. Specifically, our team deployed 8x NVIDIA A100 80GB SXM GPUs and carried out the development within a CUDA 12.4 environment.
Although this heterogeneous mixed computing paradigm yielded significant isolated gains in compute-intensive segments, it also introduced an undeniable bottleneck---the communication overhead induced by the frequent transmission of massive data between the host and the device. Performance testing revealed that the frequent PCIe bus synchronization negated the dividends of computational power, resulting in a final speedup ratio for the GPU version that was merely half that of the CPU version.
GPU-Based BSSN Solver Optimizations
In the process of solving the BSSN equations, to ensure that the computational results of the CUDA version are strictly consistent with the baseline Fortran code, we performed systematic logical corrections and standard updates on the core operators:
Synchronization primitive APIs: We replaced cudaThreadSynchronize with cudaDeviceSynchronize, enhancing the forward compatibility and execution stability of the code in modern CUDA environments.
Finite difference operator logic: We refactored the boundary condition handling and finite difference schemes of the derivative computation kernels. In the original CUDA implementation, boundary detection in x, y, and z dimensions was evaluated independently. Referring to the holistic evaluation logic of the Fortran code, we corrected this mechanism to achieve strict bit-for-bit equivalence.
Memory management optimization: We introduced a memory pool pre-allocation mechanism. During the program initialization phase, a sufficient global video memory workspace is allocated at once based on the grid topology scale. This strategy completely eliminated the unnecessary blocking caused by the frequent allocation and deallocation of video memory at each time step.
GPU Analysis Module Acceleration
Performance profiling indicates that the AnalysisStuff module incurs significant computational overhead. Among them, the time consumed by the internal interpolation routine Interp_Points is nearly equivalent to that of solving the right-hand side of the BSSN equations.
To break through this bottleneck, we developed the Interp_Points_gpu module, migrating the compute-intensive interpolation logic entirely to the GPU side. We refactored the three-dimensional interpolation logic within the kernel function, fully utilizing the massive concurrent computational resources of the GPU to allocate independent threads for tens of thousands of detection points.
We introduced a constant memory optimization mechanism: parameters that remain globally read-only during the evolution are mapped to the __constant__ memory space. This strategy not only effectively alleviates register pressure but also fully exploits the low-latency broadcast characteristics of the constant cache.
Ultimately, through the comprehensive parallelization of the interpolation calculations for the
GPU prolong3 Routine Acceleration
The performance profiling reveals that the prolong3 function (a sixth-order three-dimensional prolongation operator) within the data_packer module also constitutes a significant computational bottleneck. We designed and implemented a complex sixth-order three-dimensional interpolation kernel function based on the CUDA framework.
At the level of algorithmic refactoring, the empirical results indicate that although the GPU core computation speed vastly exceeds that of the CPU serial execution, the overall performance enhancement yielded by this module is relatively marginal. The underlying reason is that prolongation operations typically occur during the data exchange phase at the mesh refinement boundaries, resulting in a low ratio of computation to data transfer volume.
Performance Comparison
Determining the Baseline
The baseline performance testing was conducted in two main phases. The exploratory phase adopted a single-node testing mode, configuring 56 cores as the baseline. Subsequently entering the certainty assurance phase, the architecture was upgraded to a dual-node deployment, totaling 64 cores.
| Process Number | Time/(per iter) | RMS |
|---|---|---|
| 4 | 115.0041 s | 0.0203% |
| 8 | 83.0201 s | 0.0203% |
| 16 | 66.1602 s | 0.0203% |
| 32 | 49.9385 s | 0.0203% |
| 56 | 26.9602 s | 0.0203% |
| 64 | 19.9866 s | 0.0000% |
| 128 | 17.1985 s | 0.0203% |
Final Results
For the AMSS-NCKU project, following multiple rounds of iteration and performance profiling, the CPU-optimized version was ultimately selected as the final delivery solution. Experimental data demonstrates that this version significantly outperforms the GPU version in overall execution efficiency, reducing the execution time by over 53% (equivalent to a performance improvement of approximately 2.1 times).
The total execution time of the program upon completing all computational tasks was 9,531.03 seconds. Compared to the initial baseline version (which took 34,740.93 seconds), the optimized solution achieved a significant acceleration effect, reaching an overall speedup of 5.02x.
| Optimizing Strategy | Total Time | Rate |
|---|---|---|
| Baseline | 47894.33 s | 1.0x |
| Final optimized version | 9531.03 s | 5.02x |
While pursuing performance enhancements, we strictly ensured the numerical precision of the computational results. Through a step-by-step comparison with the baseline results (Total Steps: 2400), the root mean square error of the final solution was determined to be 0.0000%, fully satisfying the precision requirements.