🤖 AI Summary
This work addresses the inefficiency of existing GPU-based algebraic multigrid (AMG) implementations that rely on scalar sparse matrix formats, which fail to exploit the natural rectangular block structure arising from discretizations of vector-valued PDEs such as elasticity, leading to high memory overhead and low arithmetic intensity. The paper presents the first fully GPU-resident, native block AMG pathway in PETSc, built upon Kokkos to support a portable block-sparse matrix type with heterogeneous block sizes across coarse and fine grids while preserving block structure throughout the smoothed aggregation AMG process. It integrates PetscSF for block prolongation communication across processes and introduces a dense rectangular block COO assembly mechanism. Experiments on 3D elasticity problems using A100 GPUs demonstrate significantly reduced GPU memory usage compared to cuSPARSE and scalar Kokkos approaches, achieving a 2.27× speedup in Galerkin construction at 64-GPU scale, along with 1.24× and 1.42× improvements in V-cycle and SpMV performance, respectively.
📝 Abstract
Smoothed-aggregation algebraic multigrid (AMG) is widely used for the linear systems arising from finite-element discretizations of vector PDEs such as elasticity, but its GPU implementations have used scalar sparse matrix formats. These problems carry a natural block structure: matrix nonzeros occur in dense bs x bs blocks sharing one column index, so storing the blocks directly removes most of the index data and raises the arithmetic intensity of the bandwidth-bound kernels that dominate AMG on the GPU. Existing blocked GPU kernels (cuSPARSE, Kokkos Kernels) require equal row and column block sizes, but AMG for elasticity is rectangular-blocked: the near-null space of rigid-body modes makes the coarse block size (6 in 3D) differ from the fine (3), so the prolongator and the Galerkin triple product mix block sizes. We add a portable, Kokkos-backed blocked matrix type to PETSc with rectangular-block kernels, and make every step of the smoothed-aggregation setup operate on the block format directly, with no expansion to scalar form on the coarsening path. The two phases that recur when the hierarchy is reused across solves -- the Galerkin coarse-operator recompute (A_c = P^T A P) and the V-cycle -- are kept resident on the device in blocks, via a native blocked off-process prolongator gather over a PetscSF and a new blocked COO assembly path for dense rectangular blocks. On A100 GPUs for 3D elasticity, the cuSPARSE Galerkin product runs out of GPU memory on a 128^3 grid (6.3M unknowns) packed onto 8 GPUs, where the blocked format fits; the native Kokkos Kernels scalar path also fits, but with a much heavier Galerkin product. Where the formats run, the blocked format is at parity on one GPU and faster at scale: at 27 GPUs it is 1.24x faster on the V-cycle, 1.42x on SpMV, and 1.80x on the coarse-operator recompute, reaching 2.27x on the latter at 64 GPUs.