algebraic multigrid

This work addresses a critical gap in the JAX ecosystem: the absence of a sparse linear solver that supports GPU-accelerated algebraic multigrid (AMG), automatic differentiation, and distributed multi-GPU execution. The authors present the first native integration of NVIDIA AmgX as a JAX primitive, unifying AMG with Krylov subspace methods within JAX’s computational framework. This integration enables just-in-time (JIT) compilation, reverse-mode automatic differentiation, batching, and MPI-based distributed execution across multiple GPUs. To enhance efficiency, the implementation incorporates a solver caching mechanism that mitigates repeated setup overhead. By delivering a high-performance, scalable sparse linear algebra layer, this work bridges a key tooling gap between differentiable simulation and scientific computing, enabling seamless incorporation into scientific machine learning pipelines—particularly for PDE-constrained optimization and inverse problems.

This work addresses the challenge of efficiently scaling dense linear algebra operations—specifically linear systems and eigenvalue problems—to multiple GPUs while maintaining compatibility with just-in-time (JIT) compilation in modern programming frameworks. We present the first implementation of composable, JIT-compatible multi-GPU dense linear algebra primitives in JAX by integrating NVIDIA’s cuSOLVERMg library through the XLA Foreign Function Interface. This integration enables out-of-core Cholesky factorization and symmetric eigendecomposition at unprecedented scales, overcoming single-GPU memory limitations. Crucially, our approach seamlessly interoperates with JAX’s native transformations, including automatic differentiation and vectorization, thereby delivering scalable, high-performance distributed dense linear algebra capabilities for end-to-end scientific computing workflows.

This work addresses the key challenge of enabling efficient, cross-architecture execution of the DD-αAMG multigrid solver for lattice QCD simulations on heterogeneous exascale supercomputers (e.g., ORISE). Method: We present the first full port of DD-αAMG to the HIP platform, unifying acceleration across both NVIDIA and AMD GPUs; introduce an improved Richardson smoother that achieves superior convergence–speed trade-offs—outperforming GCR significantly and lagging SAP by only ~10%; and extend the odd-even preconditioned GMRES/Richardson hybrid iteration, applying advanced SIMD vectorization and computational restructuring to coarse-grid operations. Contribution/Results: Experiments on ORISE demonstrate efficient heterogeneous acceleration, with substantial performance gains in critical coarse-grained operations. The proposed framework establishes a scalable, cross-platform AMG solver paradigm tailored for large-scale lattice QCD computations.

High-order discontinuous Galerkin (HDG) methods face severe iterative inefficiency and memory bottlenecks when solving PDEs on GPUs, primarily due to the unstructured sparsity and large size of the global stiffness matrix. Method: This work introduces the first GPU-native, block-structured storage format tailored for HDG global stiffness matrices, coupled with a full-stack parallel iterative solver and preconditioning framework. Key techniques include: (1) elimination of local degrees of freedom to maximize arithmetic intensity; (2) batched block-Jacobi, additive Schwarz, and polynomial smoothers, co-designed with shared-memory optimization and memory-coalesced access patterns; and (3) native support for Newton–GMRES nonlinear solvers and domain-decomposition preconditioning. Results: The framework demonstrates high efficiency across six PDE classes—from Poisson to Reynolds-Averaged Navier–Stokes—on multiple generations of NVIDIA and AMD GPUs. It supports both structured and unstructured meshes and arbitrary polynomial orders, achieving superior throughput and strong scalability compared to state-of-the-art approaches.

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.

This work addresses the scalability bottleneck of large-scale incompressible fluid simulations on multi-node CPU/GPU clusters by proposing an MPI-based heterogeneous parallel optimization framework integrated with an enhanced geometric multigrid Poisson solver. The core innovations include an adaptive under-relaxed red-black Gauss–Seidel smoother and an anisotropic coarsening operator, both implemented purely in Julia. Implemented within WaterLily.jl, the method achieves near-ideal strong scaling and weak scaling efficiency exceeding 85%. Notably, it attains over 96% weak scaling efficiency across nodes at billion-cell resolution, substantially enhancing solver performance and memory concurrency.

This work addresses the longstanding challenge in algebraic multigrid (AMG) methods of balancing sparsity and convergence quality in coarse-grid operators. The authors propose RAPNet, the first graph neural network framework that directly learns nongalerkin coarse operators from sparse algebraic systems. Employing a hierarchical training strategy, RAPNet generalizes from small subgraphs to problems with millions of unknowns, achieving an effective trade-off among sparsity, convergence, and generalization while preserving computational efficiency during the solve phase. Experimental results demonstrate that RAPNet significantly outperforms classical nongalerkin baselines across a range of discretized partial differential equations and graph Laplacian problems, substantially accelerating multi-query tasks such as eigenvalue computations, time-dependent simulations, inverse problems, and design optimization.

This work addresses the underutilization of CPU resources in modern heterogeneous high-performance computing (HPC) systems when solving large-scale symmetric positive-definite linear systems using GPU-only approaches. Leveraging the SYCL programming model, the authors present the first heterogeneous implementations of the conjugate gradient (CG) method and Cholesky decomposition that operate across multi-vendor CPU-GPU platforms, including NVIDIA, AMD, and Intel architectures. Experimental results demonstrate that the heterogeneous CG solver achieves up to 32% speedup over GPU-only execution on large matrices, while the heterogeneous Cholesky decomposition attains a 29% acceleration. Furthermore, across diverse hardware vendors, the Cholesky solver consistently delivers at least a 12% performance improvement, significantly enhancing both computational efficiency and portability.

This work addresses the high memory and computational complexity typically associated with solving three-dimensional partial differential equations on Cartesian grids. By exploiting tensor-product structure, the proposed method decomposes the 3D operator into one-dimensional banded kernels aligned with coordinate axes, thereby avoiding explicit assembly of the global matrix and enabling a matrix-free solution strategy. Within a unified framework that integrates diverse numerical approaches—including Kronecker product algebra, compact finite differences, isogeometric analysis, and direct diagonalization—the study systematically identifies three key techniques: multi-right-hand-side reshaping, sum factorization, and pencil-style MPI decomposition. These innovations collectively enhance hardware affinity and parallel scalability, reducing algorithmic complexity to O(N) and storage requirements to O(Nₓ + Nᵧ + N_z), thus enabling efficient large-scale 3D PDE simulations.