Search papers, labs, and topics across Lattice.
This paper introduces a novel blocked matrix type for smoothed-aggregation algebraic multigrid (AMG) in PETSc, specifically tailored for GPU implementations that handle rectangular-block structures typical in elasticity problems. By avoiding the expansion to scalar form during the coarsening path and keeping operations resident on the device, the authors achieve significant performance improvements on A100 GPUs, demonstrating up to 2.27x speedup in coarse-operator recompute at scale. This advancement not only enhances memory efficiency but also optimizes computational intensity, making AMG more viable for large-scale problems in finite-element discretizations.
Achieving up to 2.27x speedup in coarse-operator recompute on A100 GPUs reveals the untapped potential of natively blocked formats in GPU-accelerated algebraic multigrid methods.
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.
