SMART-PETSc: Smart Middleware for Accelerating PETSc

The efficient parallelization of algebraic solvers for partial differential equations, as exemplified in PETSc, requires nearest-neighbor ghost-point communication and global reductions. The ghost-point communication involves non-contiguous memory access. For scalability, the global reductions in pipelined Krylov methods must be non-blocking. The challenge to utilizing GPUs in an MPI environment for algebraic solvers is using all available hardware resources simultaneously. Achieving this requires MPI support for streams to overlap communication and computation, efficient non-contiguous memory access, nearest neighbor collectives, and GPU-efficient non-blocking global reductions. We propose to develop SMART-PETSc a Smart Middleware for Accelerating PETSc on modern HPC systems with the following capabilities: 1) optimized datatype processing for intra-/inter-node transfers using advanced hardware capabilities, 2) overlapping kernel launch and execution to use all hardware resources simultaneously, 3) co-designing PETSc and MVAPICH to support user-specified CUDA streams, 4) in-network support for non-blocking CPU-/GPUbased collectives, 5) in-network support for neighborhood collectives. We enable full-stack observability and explainability through tight integration with the TAU analysis tools. We have shown positive results in 1) intra-/inter-node transferring using kernel fusion, 2) overlapping kernel launch and execution, 3) communication with compression support, and 4) efficient neighborhood communication. Our prototype product (X-ScalePETSc 2024.6) was tested and evaluated using up to 2,048 processes and representative applications including Ratel. It improves the performance of high-order finite difference Laplacian and non-periodic boundary Possion by up to 10% and 9%, respectively (on TACC Frontera). For matrix-vector kernel operations, we have demonstrated 24% performance gain. In addition, we have developed a new capability to further PETSc stack observability and explainability through tight integration with TAU. This capability has been contributed to and incorporated into the PETSc open-source repository. Justification for Proposed Work in Phase-IIA While we are progressing on the SMART-PETSc designs for existing HPC clusters in Phase II, the market has moved significantly towards new technologies including NVIDIA Grace CPU architecture, Hopper/Blackwell GPU architecture, Bluefield-3 DPU, HPE/Cray Slingshot Interconnect, etc. These technologies are leading to a new generation of HPC clusters and cloud environments, exhibiting new characteristics in compute to memory bandwidth and computation to communication ratio using CPU and GPU. Several SMART-PETSc designs (compression and neighborhood collectives) need to be enhanced for these emerging technologies so that the product can appeal to a large number of customers. The proposed Phase IIA will focus on these efforts while remaining consistent with the development goals in the original Phase II grant. Benefits The transformative impact of the proposed SMARTPETSc product will enable many HPC and AI applications that use PETSc routines to take advantage of the novel features and capabilities of emerging hardware technologies and reap the additive benefits of co-design and tighter integration between PETSc and MVAPICH libraries. We expect that SMART-PETSc can reduce the communication overhead by up to 2x, leading to a significant boost in performance and scalability of many HPC and AI applications that use PETSc APIs. The collaboration with teams at ANL and ParaTools allows DOE Labs to use SMART-PETSc on modern HPC systems and provide a commercial version for supercomputing centers and cloud providers worldwide.