arXiv:2606.24748v1 [cs.DC] 23 Jun 2026
A Natively Blocked, Device-Resident Algebraic Multigrid GPU Path in PETSc Mark F. Adams Lawrence Berkeley National Laboratory Berkeley, CA 94720 [email protected]
Abstract Smoothed-aggregation algebraic multigrid (AMG) is widely used for the linear systems that arise from finite-element discretizations of vector partial differential equations such as elasticity, but its GPU implementations have used scalar sparse matrix formats. These problems carry a natural block structure: the matrix nonzeros occur in dense bs × bs blocks that share 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 in NVIDIA cuSPARSE [1] and Kokkos Kernels [2] 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 three dimensions) 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 [3] (MATBAIJKOKKOS) with rectangular-block kernels, and make every step of the smoothed-aggregation setup operate on the block format directly, with no expansion to scalar AIJ on the coarsening path. The two phases that recur when the hierarchy is reused across solves—the Galerkin coarse-operator recompute (Ac = P T AP ) and the V-cycle—are kept resident on the device in blocks, using a native blocked off-process prolongator gather over a dedicated PetscSF and a new blocked COO assembly path (MatCOOUseBlockIndices) that generalizes PETSc’s coordinate-format assembly to dense bsr × bsc blocks. On NVIDIA A100 GPUs for 3D linear elasticity the vendor cuSPARSE Galerkin product runs out of GPU memory on a 1283 grid (6.3M unknowns) packed onto 8 GPUs, where the blocked format fits; the portable native-Kokkos-Kernels scalar path also fits there, 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.24× faster on the V-cycle (KSPSolve), 1.42× faster on SpMV, and 1.80× faster on the coarse-operator recompute (PtAP), with the PtAP advantage reaching 2.27× at 64 GPUs.
1
Introduction
Multigrid methods solve the linear systems arising from elliptic PDEs in operations proportional to the number of unknowns [4]. For unstructured, matrix-defined problems, algebraic multigrid (AMG) builds the multilevel hierarchy from the system matrix itself, and smoothed-aggregation AMG [5] is particularly effective for systems of PDEs such as elasticity, where the near-null space of rigid-body modes must be preserved on every coarse level. Two hardware trends shape any modern multigrid implementation. First, accelerators (GPUs) supply the bulk of floating-point throughput on current leadership machines, so the solver must run on the device. Second, on those devices floating-point throughput vastly exceeds memory bandwidth, so the cost of nearly every sparse kernel—sparse matrix–vector products in the smoother and Krylov
1
accelerator, sparse matrix–matrix products in the Galerkin coarse-operator construction—is set by how many bytes move, not how many flops execute [4]. These two trends point directly at blocked sparse storage. A discretization of a vector PDE associates several unknowns with each mesh point—three displacement components per node for 3D elasticity—so the matrix has a natural fixed block size bs, and its nonzeros occur in dense bs × bs blocks that share one (row-block, col-block) index. Storing those blocks directly amortizes one integer index over bs2 values, shrinks index storage by a factor of about bs2 , and raises arithmetic intensity by reusing each loaded index across an entire block. PETSc has provided such a format (BAIJ) on the CPU for decades. Why a vendor blocked kernel is not enough. The blocked formats in NVIDIA cuSPARSE [1] and Kokkos Kernels [2] (block compressed sparse row, BSR) require equal row and column block sizes. Smoothed-aggregation AMG for elasticity violates this at every level transition: the near-null space—six rigid-body modes in 3D—gives each aggregate six coarse unknowns while fine nodes carry three, so the prolongator P has rectangular 3 × 6 blocks and the Galerkin triple product P T AP composes blocks of differing sizes. Among the vendor libraries only AMD’s rocSPARSE [6] exposes a general block format with independent row and column block sizes. To obtain a portable, vendorneutral blocked AMG we add a new Kokkos-backed blocked matrix type (MATBAIJKOKKOS) whose kernels are templated on independent row and column block sizes. The gap in PETSc’s AMG. Adding the matrix type was necessary but not sufficient. PETSc’s smoothed-aggregation multigrid (GAMG) has no blocked path: every step of its setup phase is written against scalar AIJ machinery—operator inspection, the aggregation graph, the off-process gather of prolongator rows, and the assembly of the Galerkin coarse operator all require the operator in scalar form. A blocked operator therefore has no path through GAMG without expanding to scalar AIJ, which forfeits the blocked memory advantage and, as Section 4.5 shows, can exhaust GPU memory forming the coarse Galerkin product with the vendor cuSPARSE backend. Contributions. This paper makes the blocked path native end to end: 1. A portable, rectangular-block GPU matrix type (MATBAIJKOKKOS) with Kokkos-backed kernels supporting independent row and column block sizes. 2. A conversion-free blocked GAMG setup: operator inspection, the aggregation graph, and the off-process prolongator gather all operate on the block format directly, with no AIJ expansion on the coarsening path. 3. Device-resident reuse of the coarse-operator recompute: the prolongator-side data is cached and state-gated so the hot PtAP numeric phase incurs no host round trip on reuse. 4. A reusable blocked COO assembly primitive (MatCOOUseBlockIndices) generalizing PETSc’s coordinate-format assembly to dense bsr × bsc blocks, of independent value for finite-element assembly. 5. A measured characterization of when blocked wins, with an arithmetic-intensity argument for why the dominant kernels are bandwidth-bound. The remainder is organized as follows. Section 2 reviews multigrid and block structure. Section 3 describes the native blocked design. Section 4 reports performance on A100 GPUs. Section 5 presents the blocked COO primitive. Section 6 concludes.
2
2
Background
2.1
Multigrid methods
Multigrid is motivated by the observation that a low-resolution discretization of an operator can capture error components that are expensive to resolve on a highly resolved discretization. A simple pointwise relaxation (smoother) such as Gauss–Seidel or damped Jacobi efficiently damps the high-frequency, locally determined part of the error but leaves the smooth, globally coupled part nearly untouched; a coarser representation resolves that smooth error cheaply. Applying this idea recursively over a sequence of coarser levels yields a method whose total work converges to a small multiple of the finest-level work. For a linear system Ax = b, a two-level cycle is defined by a smoother S, restriction operator R, prolongation operator P , and coarse operator Ac . One cycle pre-smooths the iterate, forms and restricts the residual, solves the coarse correction, prolongs and adds it, and post-smooths. Recursion gives the V-cycle. The coarse operator is most commonly formed by the Galerkin triple product Ac = P T AP . A consideration growing steadily more important is that modern processors have limited memory bandwidth relative to floating-point throughput, and minimizing memory traffic is therefore a primary design criterion [4].
2.2
Algebraic multigrid by smoothed aggregation
Algebraic multigrid (AMG) constructs the entire hierarchy from the system matrix. We use the smoothed-aggregation (SA) variant [5], which is well suited to systems of PDEs and to the block structure exploited here. SA builds each coarse level in two stages: aggregation and prolongator construction, driven by strong coupling. For threshold ε, the strongly-coupled neighborhood of node i is √ Ni (ε) = {j : |aij | ≥ ε aii ajj }. Aggregates are formed by a greedy disjoint covering. Each aggregate gives rise to one coarse node. A tentative prolongator P̃ is defined so that the near-null space is represented exactly on the coarse level. For 3D elasticity the near-null space is the six rigid-body modes, so each aggregate contributes six coarse degrees of freedom and the coarse operator naturally carries a larger block size than the fine operator (e.g. bs: 3 → 6). Preserving these zero-energy modes on every level is essential for mesh-independent convergence [5]. The tentative prolongator is improved by one step of a damped-Jacobi prolongator smoother, P = (I − ωD−1 A)P̃ , where D is the diagonal of A. The coarse operator is then Ac = P T AP and the process repeats.
2.3
Block structure of systems of PDEs
Discretizations of vector PDEs associate several unknowns with each mesh point—three displacement components per node for 3D elasticity—so the system matrix has a natural fixed block size bs, and its nonzeros occur in dense bs × bs blocks that share a single (row-block, col-block) index. Storing and operating on these blocks directly, rather than as bs2 independent scalar entries each carrying its own column index, is the classical motivation for PETSc’s BAIJ format. Crucially, the block structure is not uniform across levels: the fine elasticity operator has 3 × 3 blocks, but the coarse block size equals the dimension of the preserved near-null space—six 3
rigid-body modes in 3D—so coarse operators carry 6 × 6 blocks and the prolongator has rectangular 3 × 6 blocks. A blocked AMG hierarchy therefore demands kernels that handle independent row and column block sizes.
2.4
The blocked-GPU gap in PETSc AMG
PETSc’s GPU support routes sparse kernels through backend libraries: cuSPARSE [1] for scalar AIJ on NVIDIA hardware, and Kokkos Kernels [2] for the portable *KOKKOS types. Scalar AMG on the GPU is mature. The vendor libraries do provide blocked sparse formats, but both cuSPARSE and Kokkos Kernels implement BSR with a single fixed square bs, which cannot represent the rectangular 3 × 6 prolongator or mixed-size Galerkin products of Section 2. Among the major vendor libraries only AMD’s rocSPARSE [6] offers a general block format (GEBSR) with independent row and column block sizes. To obtain a portable blocked AMG we implement a Kokkos-backed blocked matrix type with rectangular-block kernels. Even with that matrix type in hand, PETSc has no device-resident blocked GAMG setup: the GAMG setup is written against scalar AIJ, so operator inspection, the aggregation graph, and coarse-operator assembly all require the operator in scalar form, making a conversion-free blocked path impossible without this work.
3
Design
The design rests on two invariants. First, the blocked operator is never expanded to scalar AIJ anywhere on the coarsening path. Second, the two recurring phases of a production solve—the Galerkin coarse-operator recompute (PtAP) and the V-cycle—are fully device-resident in blocks. The one-time aggregation graph setup remains a host computation (Section 3.2); because the hierarchy is reused, it is a cold, amortized cost outside the recurring work.
3.1
What is hot, and why it matters
A GAMG solve has two phases. The setup coarsens the operator: it builds the aggregation graph, the tentative and smoothed prolongators, and the Galerkin coarse operators Ac = P T AP on every level. The solve runs the V-cycle. In a production nonlinear or time-dependent simulation the operator A changes on every Newton or time step, but the hierarchy—the aggregates and prolongators—is reused across steps (-pc gamg reuse interpolation). On each step the symbolic setup is computed once and reused, while the numeric coarse operators are recomputed because A is new. The recurring operations are therefore: • hot PtAP: the numeric Galerkin recompute Ac = P T AP with reused sparsity (MAT REUSE MATRIX), on every level; and • hot KSPSolve: the V-cycle, dominated by SpMV in the smoother and Krylov accelerator.
3.2
Native block operator inspection and aggregation graph (host)
Smoothed aggregation needs two pieces of information from the operator before forming any product: a scalar measure of each block row for the strength-of-connection threshold, and a graph whose edges are the strong couplings Ni (ε). GAMG’s existing code requires a scalar AIJ operator to supply them.
4
We compute both directly from the block format. Operator inspection (MatGetInfo/MatNorm) is computed over the bs × bs blocks of the Kokkos BAIJ storage. The aggregation graph is built from the block sparsity: each (row-block, col-block) index is one graph edge, and the strength weight is the block norm. These operations run on the host—graph construction is irregular, serial-leaning work poorly suited to the GPU—but, critically, they no longer pay the bs2 scalar expansion, and because the graph and aggregates are built once and reused across solves they sit in the cold setup outside the recurring per-step cost.
3.3
Native block off-process prolongator gather (Poth )
The Galerkin product P T AP on a distributed matrix needs, on each rank, the rows of P corresponding to the off-process columns of the local A—the off-process prolongator rows Poth . We carry the gather natively in the blocked format. A dedicated PetscSF is constructed whose roots and leaves are block rows of P , and the broadcast moves whole prolongator blocks (PetscSFBcast over a block-typed unit) rather than separate scalars. Because P is rectangular (3 × 6 blocks), the block-typed unit is the full bsr × bsc payload. The communication is GPU-aware: the block payload is staged and exchanged directly from device memory (PetscSFBcastWithMemType with PETSC MEMTYPE KOKKOS), so the gather never round-trips through the host.
3.4
Blocked COO assembly of the coarse operator
The coarse operator is assembled from local triple-product contributions plus off-process contributions routed to their owning ranks. PETSc’s device-assembly path uses coordinate format (COO): the caller supplies, once, the (i, j) coordinates of every contribution (MatSetPreallocationCOO), PETSc builds a cached communication-and-scatter plan, and thereafter each numeric assembly is a single MatSetValuesCOO that scatters a flat value array and sums duplicates, entirely on the device. We generalize this to blocks. A new option, MatCOOUseBlockIndices, tells the COO setup that each coordinate addresses a dense bsr × bsc block. The preallocation plan is built over block coordinates, and the numeric scatter moves block payloads. The block size need not be square—the coarse-operator assembly uses the same mixed sizes as the Galerkin product—so the primitive covers the rectangular case the vendor square-block formats cannot.
3.5
Device-resident reuse of the hot PtAP
The three pieces above make a single blocked PtAP native. The production case repeats it under reused interpolation, where P is fixed across hot solves and only A changes. We exploit this by caching the prolongator-side data—the local prolongator rows Plocal , the off-process rows Poth , the stacked operand, and the transpose R—and gating its rebuild on the prolongator’s object state (PetscObjectState). On a hot recompute, if P ’s state matches the cached value, the path reuses the cached device-resident values directly: in particular the Poth gather is not re-broadcast but served from cache. The hot PtAP numeric phase becomes a local blocked triple product plus the off-process reduction of the new coarse values, with everything on the prolongator side served from device-resident cache. Combined, Sections 3.2–3.5 deliver the two invariants. No step of the coarsening path expands the operator to scalar AIJ, and across the recurring phases the operator stays in blocks on the device with no host round trip on reuse.
5
4
Performance
4.1
Experimental setup
All results are on the Perlmutter system at NERSC. Each node has four NVIDIA A100 GPUs, and we use GPU-aware MPI. Unless noted we run one MPI rank per A100: sharing a GPU between two ranks inflated contention and is not representative of production placement. The build is optimized, with the CUDA backend for Kokkos [7] and Kokkos Kernels [2] configured with NVIDIA’s vendor libraries (cuSPARSE, cuBLAS, cuSOLVER) as third-party libraries (TPLs)—the default for a CUDA --download-kokkos-kernels build. With these TPLs enabled, the scalar AIJKOKKOS matrix type dispatches its sparse kernels to cuSPARSE; built without them the same type runs the portable native Kokkos Kernels (KK) implementations. The blocked kernels are our own in either build: cuSPARSE’s BSR format supports only square, equal-size blocks, so rectangular Galerkin and prolongator products cannot route through it. We use cuSPARSE and Kokkos Kernels (KK) to name the two scalar backends, reserving TPL for the build-configuration discussion (Section 4.3). The model problem is 3D linear elasticity (block size bs = 3), preconditioned by GAMG with a point-block Jacobi smoother (pbjacobi) and a CG Krylov accelerator. Sections 4.2–4.5 and 4.7–4.9 use the hand-assembled trilinear (Q1) hexahedral elasticity example src/ksp/ksp/tutorials/ex56 [3]. We identify each problem by its node grid m3 ; the weak-scaling ladder is m3 ∈ {323 , 643 , 963 , 1283 } for {1, 8, 27, 64} GPUs, holding 98,304 unknowns per GPU. The separate nonzeros-per-row study of Section 4.6 uses a DMPlex/PetscFE harness (src/snes/tutorials/ex56) with selectable element order (Q1/Q2). The two harnesses use different boundary/load configurations and are not directly comparable across tables. We use the unpreconditioned residual norm throughout; with this norm the two formats converge in the same iteration count to the same true residual on every problem we report. We compare three GPU backends: two scalar (AIJKOKKOS) backends—cuSPARSE (the vendortuned default) and Kokkos Kernels (KK) (TPLs off)—and one block (BAIJKOKKOS) backend. cuSPARSE is the primary baseline. We follow the production model of Section 3.1: a hierarchy is built once and reused across solves with a changing operator (-pc gamg reuse interpolation true, -two solves). Hot numbers are taken from the second/third solve and the second setup. A note on timers. PETSc’s -log view mis-attributes asynchronous GPU kernel time unless per-event device synchronization is enabled (-log view gpu time), but that synchronization also serializes MPI waits and inflates end-to-end solve time. We read each metric from the run in which it is accurate: KSPSolve from the run without the GPU timer (true wall time), and SpMV from the run with it (true GPU-event time). For the communication-bound hot PtAP, neither -log view mode is trustworthy, so we report hot PtAP from Nsight Systems traces scoped to the hot MatPtAPNumeric spans.
4.2
Hot KSPSolve and SpMV—block wins at scale
The V-cycle is SpMV-bound: the matrix–vector products in the smoother and CG accelerator account for the large majority of KSPSolve time. Table 1 gives the weak-scaling comparison, and Figure 1 summarizes the speedup ratios across the weak-scaling ladder. At eight GPUs the block format wins both KSPSolve (0.143 vs. 0.149 s, 1.04×) and SpMV (0.089 vs. 0.100 s, 1.12×). The advantage widens at 27 GPUs: KSPSolve 0.272 vs. 0.338 s (1.24×) and SpMV 0.149 vs. 0.211 s (1.42×). At 64 GPUs the block format continues to win: KSPSolve 0.279 vs. 0.325 s (1.16×) and SpMV 0.163 vs. 0.213 s (1.30×). The block SpMV moves one column index per bs × bs block instead of bs2 indexed scalar entries, moving less data per useful flop on a 6
Table 1: Hot timings (seconds), weak scaling at one rank per A100, 3D elasticity. KSPSolve is off-timer wall time; SpMV is gpu-timer; PtAP is Nsight-scoped. Iteration counts match at each scale (1 GPU: 24; 8 GPUs: 30; 27 GPUs: 35; 64 GPUs: 38). Bold = faster. Event (hot)
1 GPU, 323
8 GPUs, 643
27 GPUs, 963
64 GPUs, 1283
Block
Block
Block
Block
Scalar
Scalar
Scalar
Scalar
KSPSolve (s) 0.0510 0.0428 0.1428 0.1487 0.2718 0.3380 0.2793 0.3252 SpMV (s) 0.0361 0.0243 0.0893 0.1003 0.1492 0.2113 0.1632 0.2125 PtAP (ms, scoped) 5.99 5.95 10.21 14.80 10.50 18.89 10.88 24.69
parity (1×) KSPSolve SpMV PtAP
2.50
Block / scalar speedup
2.25
2.27×
2.00
1.80×
1.75 1.45×
1.50
1.24× 1.42×
1.25
1.04× 1.12×
0.99× block wins
1.00 0.84× scalar wins 0.75 0.50 0.67× 1
8
GPU count (1 rank/GPU)
27
1.16× 1.30×
64
Figure 1: Block/scalar speedup ratio (scalar time / block time; > 1 means block faster) at one rank per A100. SpMV advantage peaks at 27 GPUs (1.42×, at the theoretical index-bandwidth ceiling) and narrows slightly at 64 GPUs as cuSPARSE handles the larger communication pattern more efficiently. PtAP advantage grows monotonically (1.45× → 2.27×). bandwidth-bound device. The SpMV advantage peaks at 27 GPUs (1.42×) and narrows slightly at 64 GPUs, where scalar cuSPARSE handles the larger-rank communication pattern more efficiently. On a single GPU the scalar path wins KSPSolve/SpMV by ∼ 19–33%: with no halo communication to amortize the index traffic against, the tuned scalar cuSPARSE SpMV is faster. Bounding the SpMV advantage. Counting bytes a 3 × 3 block (nine scalar nonzeros) moves in each format, with int32 indices and fp64 values: Format Scalar CSR Blocked BAIJ
Values
Indices
Total
Per scalar nz
9 × 8 = 72 B 9 × 8 = 72 B
9 × 4 = 36 B 1 × 4 = 4B
108 B 76 B
12.0 B 8.44 B
The blocked format moves 76/108 = 0.70 of the scalar matrix bytes—a 30% reduction, for a traffic-bound ceiling of 108/76 ≈ 1.42×. The reason is that fp64 values dominate: the index is only 4 of the 12 B per scalar nonzero, and blocking removes 8 of every 9 indices, saving at most (8/9) · 4/12 ≈ 30%. The measured 1.42× at 27 GPUs is at this theoretical ceiling, so the blocked 7
1.04×
0.14
0.1428
8 GPUs, 643 grid, 1 rank/GPU
0.1487
block (BAIJ) scalar (AIJ)
Time (s), hot
0.12
1.12×
0.1003
0.10
0.0893
0.08 0.06 0.04 1.45×
0.02 0.00
0.0102
KSPSolve
SpMV
0.0148
PtAP
Figure 2: Hot timings at 8 GPUs (643 grid, one rank per A100). Grouped bars show block (dark) vs. scalar cuSPARSE (light) for KSPSolve, SpMV, and PtAP . PtAP values are Nsight-scoped. Table 2: Scalar backend with and without the cuSPARSE TPL, Q1 n = 8, 643 grid. SpMV is gpu-timer; PtAP is Nsight-scoped (ms). Block (BAIJ) is the same code in both builds (±1% run-to-run variance). Bold = faster scalar backend. Metric (hot)
Scalar cuSPARSE
Scalar KK
Block (BAIJ)
SpMV (s) KSPSolve (s) PtAP (ms)
0.1015 0.1496 14.60
0.0949 0.1368 111.95
0.0886 0.1408 9.80
SpMV realizes essentially all of the available index-bandwidth saving. The same accounting explains why the largest effect is on capacity (Section 4.5), not SpMV speed: the scalar out-of-memory failure comes from the bs2 -expanded SpGEMM symbolic buffers in the Galerkin product, a far larger factor than the 30% steady-state SpMV saving.
4.3
cuSPARSE vs. Kokkos Kernels (scalar backends)
To verify that the block win is measured against the stronger scalar backend, we rebuilt from the same source with --with-kokkos-kernels-tpl=0, so the scalar AIJKOKKOS type runs native Kokkos Kernels (KK) instead of cuSPARSE. The TPL choice affects only the scalar backend; the block kernels are the same code in both builds. Table 2 shows the result. At this index-bound, mid-size configuration native KK CSR SpMV is faster than cuSPARSE (0.0949 vs. 0.1015 s, −6%). The block backend still wins both scalar backends on SpMV: 1.07× over KK and 1.15× over cuSPARSE. Because KK is the faster scalar SpMV here, the weak-scaling SpMV ratios in Table 1, computed against cuSPARSE, slightly overstate the block advantage relative to
8
Table 3: Hot PtAP at 8 GPUs (643 grid), Nsight-scoped, milliseconds. “Ungated” re-broadcasts Poth every recompute; “state-gated” reuses it from cache when P is unchanged. Inside hot PtAP
Block ungated
Block state-gated
Scalar (cuSPARSE)
Triple-product compute Poth broadcast Off-process reduce
7.83 9.93 13.94
7.4 0 2.2
10.57 3.17 0.45
Wall total
31.85
10.2
14.80
the best scalar backend. The PtAP comparison goes the other way: native KK spgemm is 7.7× slower than cuSPARSE for the hot numeric phase (112 vs. 14.6 ms), so cuSPARSE is the stricter PtAP baseline. We report cuSPARSE throughout as the vendor-default scalar backend.
4.4
Hot PtAP—device-resident reuse
The Galerkin coarse-operator recompute is the second production-critical metric, and it is where the device-resident reuse of Section 3.5 has the largest effect. Table 3 shows the ablation at eight GPUs. Without device-resident reuse, the block hot PtAP at eight GPUs is 31.8 ms (Table 3, ungated), against the scalar path’s 14.8 ms—2.15× slower. Nsight scoping locates the gap in communication, not arithmetic: the block triple-product compute (7.83 ms) is faster than cuSPARSE’s spgemm (10.57 ms), but the ungated path repeats the Poth PetscSF broadcast and the off-process reduction inside every numeric recompute (together ∼ 24 ms), work that the scalar cuSPARSE path folds into its cached symbolic phase. State-gating the prolongator-side rebuild removes this. With P unchanged the cached Poth gather is reused rather than re-broadcast, so its cost falls to zero and the device-to-host transfers fall to zero; removing the redundant per-rank work also tightens the coarse-grid synchronization, shrinking the imbalance-gated wait at the off-process reduction from 13.9 to ∼ 2.2 ms. The block hot PtAP drops from 31.8 to 10.2 ms (3.1×), moving from 2.15× slower than scalar to 1.45× faster. On a single GPU there is no communication and the PtAP comparison is pure compute: block 5.99 ms vs. scalar 5.95 ms, within 1% (parity).
4.5
Memory capacity—a backend-specific limit
A memory effect appears at the largest local problem sizes, and it is backend-specific. The scalar GPU format expands every bs×bs block into bs2 scalar entries, each with its own column index, so its sparse matrix–matrix product allocates much larger temporary buffers than the blocked equivalent. Placing a 1283 grid (6.3M unknowns) on only eight A100 GPUs—eight times the per-GPU load of the weak-scaling ladder—the scalar AIJ path backed by cuSPARSE runs out of GPU memory: the cudaMalloc of the cuSPARSE spgemm symbolic buffer fails forming the first coarse-level Galerkin product. This is a property of the cuSPARSE spgemm, not of the scalar format. Rebuilt with the portable native Kokkos Kernels spgemm (the TPL-off backend of Section 4.3), the same scalar AIJ problem fits and converges in 41 iterations, matching the blocked format. The two paths differ sharply in the Galerkin product: the native-KK scalar coarse-operator recompute costs ≈ 0.33 s and dominates its setup, against ≈ 0.04 s for the blocked format (gpu-timer), the ∼ 8× gap of the slow native-KK spgemm seen in Section 4.3. At this 8×-overloaded, compute-bound load the scalar V-cycle is in fact slightly faster than the blocked one (hot KSPSolve 0.39 vs. 0.50 s), as in the high-nonzeros-per-row 9
Block (BAIJ) vs scalar (AIJ): cuSPARSE runs out of memory at 1283 (native KK fits) 0.05
323, 1 GPU
0.0510
block (BAIJ) scalar (AIJ)
0.0428
Time (s), hot
0.04
0.14
0.143
0.149
block (BAIJ) scalar (AIJ)
0.12
0.0361
0.03
0.0060 0.0060
KSPSolve
SpMV
PtAP
block (BAIJ) scalar (AIJ)
0.367 0.292
0.2 0.1
0.02 0.00
0.387
0.325
0.04
0.01
cuSPARSE spgemm: OOM
0.3
0.06
0.02
1283, 8 GPUs
0.496
0.100 0.0893
0.08
0.0243
0.5 0.4
0.10
0.00
643, 8 GPUs
0.0102
KSPSolve
SpMV
0.0148
PtAP
0.0410
0.0
KSPSolve
SpMV
PtAP
Figure 3: Capacity at 1283 on 8 A100s (8× the weak-scaling per-GPU load). The scalar AIJ path backed by cuSPARSE runs out of memory forming the first coarse Galerkin product; the portable native Kokkos Kernels backend (shown) and the blocked format both fit, in 41 iterations. The 323 and 643 panels (scalar = cuSPARSE) are for reference. Hot timings: KSPSolve off-timer, SpMV/PtAP gpu-timer. At 1283 the blocked Galerkin product (PtAP) is ∼ 8× cheaper than nativeKK scalar, while the compute-bound V-cycle slightly favors scalar.
Figure 4: Block ÷ scalar hot-timing ratios (lower = block faster) for Q1 (∼ 78 nonzeros/row) and Q2 (∼ 180 nonzeros/row) at one and eight GPUs. The block advantage shrinks with nonzeros per row at eight GPUs, as the kernels shift from index-bound to more compute-bound. regime of Section 4.6; the blocked advantage at this size is the Galerkin product, not the V-cycle. Two readings follow (Figure 3): against the vendor-default cuSPARSE backend the blocked format extends the problem size a machine can solve; against the portable backend it removes the Galerkin product as the dominant cost.
4.6
Sensitivity to nonzeros per row (Q1 vs. Q2)
A natural hypothesis is that the block advantage grows with nonzeros per row: more entries per row amortize the block index overhead further and raise arithmetic intensity. We tested it with a finite-element harness comparing Q1 (∼ 78 nonzeros/row) against Q2 (∼ 180 nonzeros/row), at one and eight GPUs, matched in total degrees of freedom. Table 4 and Figure 4 show the result. The hypothesis is refuted. At one GPU the block format loses, and the deficit grows with nonzeros per row (SpMV 1.44× → 1.62× slower). At eight GPUs the block format wins, but its advantage shrinks with nonzeros per row (SpMV 1.67× → 1.23× faster). The reason is that the block format’s edge comes from compressing column indices—a benefit that is proportionally largest in the index-bound, low-nonzeros-per-row regime. As nonzeros per row grow, the kernels become
10
Table 4: Block ÷ scalar ratio (< 1 means block faster), hot solve, finite-element harness. Size
Degree
n=1 n=1 n=8 n=8
Q1 Q2 Q1 Q2
nnz/row
KSPSolve
SpMV
1.19 1.33 0.65 0.81
1.44 1.62 0.60 0.81
76 169 78 180
Table 5: ncu metrics for the dominant fine-level SpGEMM kernel, single A100. Block kernel: RunNumericAB SeqBAIJKokkos<3,3,6>. Scalar: cuSPARSE load balancing kernel. All runs unmodified GPU clocks (no clock lock). Config Block, Q2, fine ⟨3, 3, 6⟩ Scalar, Q2, fine Block, Q1, fine ⟨3, 3, 6⟩
DRAM traffic
%peak HBM
%peak SM
415 MB 4.22 GB 163 MB
8.9% 46.5% 9.4%
19.6% 22.9% 21.0%
more flop-bound, where the scalar cuSPARSE implementation is strongest and the gap closes. The low-nonzeros-per-row Q1 discretization is thus the better case for the block speed advantage.
4.7
Why the kernels are bandwidth-bound
The performance story rests on the dominant kernels being bandwidth-bound. The simplest case is the SpMV accounted above: the block format moves 76 B per 3 × 3 block against 108 B for scalar CSR, and the resulting 1.42× ceiling is reached at 27 GPUs. The Galerkin SpGEMM is bandwidth- bound as well. A block-pair product moves O(bs2 ) bytes of A and P values plus one amortized 4-byte index while performing O(bs3 ) flops, so its arithmetic intensity grows with bs but stays far below the A100 double-precision ridge of ≈ 6.2 flop/byte at both bs = 3 and bs = 6; the scalar spgemm, at ≈ 0.17 flop/byte (a 12-byte value-plus-index per 2-flop entry), is lower still. Floating-point peak is therefore not the limiting factor—all configurations run at < 5% of fp64 peak—and data volume is. ncu measurements on a single A100 (one rank, Q1 and Q2 harnesses) confirm both kernels are well below the ridge. Table 5 shows the dominant fine-level kernel for each backend. The DRAM traffic ratio is 4,220 MB/415 MB ≈ 10.2×, close to the theoretical bs2 = 9×, confirming the index-amortization argument. Both kernels sit far below the A100 ridge (6.2 flop/byte): even the scalar kernel at 46.5% of peak HBM (≈ 930 GB/s) is bandwidth-dominated. The block kernel reaches a lower fraction of peak HBM (8.9%, ≈ 180 GB/s) because the matrices are small at n = 1, but it wins in absolute time because it moves 10× less data: the block format’s advantage is a reduction in data volume, not a higher achieved bandwidth fraction.
4.8
Communication profile and load imbalance
The hot PtAP at scale is gated by an exposed, load-imbalance-driven wait at the off-process reduction. The block format moves fewer messages in every case (one block reduce versus bs2 scalar reduces per entry), so its fastest rank waits the least; the imbalance ratio is correspondingly higher for block (e.g. 75× at Q1, 94× at Q2 versus 15× and 36× for scalar) precisely because the block path’s fast ranks finish their lighter communication sooner and then wait on the slow rank. This is consistent
11
with the scoped PtAP story of Section 4.4: the block path moves less, so what remains is dominated by the coarse-grid imbalance common to both formats. Overlapping that exposed reduction with compute is a remaining lever.
4.9
Cold setup
The one-time cold setup is dominated by symbolic coarsening and prolongator smoothing. The native-block setup steps (operator inspection, graph, Poth ) perform no AIJ expansion on the coarsening path—confirmed by per-stage logging showing zero conversions in the hot second setup. One residual conversion remains in the cold path: the prolongator smoother’s MatAXPY falls back to a scalar conversion when the two operands do not share a sparsity pattern. On a 1283 grid (four A100s, two ranks per GPU) the cold setup is 13.7 s, of which ∼ 6 s is this residual conversion (18 calls, all in the first setup, none in the hot setup). Because it is a one-time cost amortized over every reused solve, it does not affect the production metrics; a native block MatAXPY would remove it and is future work (Section 6).
5
Blocked COO as a Reusable Primitive
The coarse-operator assembly of Section 3.4 is built on a primitive of independent value. PETSc’s COO assembly interface lets a caller declare, once, the (i, j) coordinates of every contribution to a matrix—including duplicates and entries destined for other ranks—after which PETSc constructs a cached communication-and-scatter plan and each subsequent numeric assembly is a single device kernel. This is the standard high-performance device-assembly path in PETSc, used most heavily by finite-element codes, where per-element dense matrices generate exactly such a stream of duplicated, possibly-off-process contributions. Our MatCOOUseBlockIndices option generalizes this from scalar entries to dense bsr × bsc value blocks. Each declared coordinate addresses a block, the value stream is a sequence of dense blocks (square or rectangular), and the scatter writes and sums contiguous bsr · bsc runs. Everything that the scalar COO plan stores—the coordinate arrays, the duplicate-summation map, and the off-process PetscSF—shrinks by a factor of about the block area. Within this paper the primitive serves the Galerkin coarse-operator assembly. Its reach is broader: any caller assembling a blocked matrix on the device—finite-element assembly of vector PDEs being the immediate example—can emit block coordinates and a block value stream and obtain the same block-area reduction in index overhead and assembly traffic. The primitive is tested for sequential and parallel blocked matrices. A standalone characterization of blocked COO against element-by-element assembly for a finite-element workload is left to future work.
6
Conclusion and Future Work
We have made PETSc’s smoothed-aggregation algebraic multigrid run in a blocked format with no scalar AIJ expansion on the coarsening path, and with its two recurring phases—the Galerkin coarse-operator recompute (PtAP) and the V-cycle—resident on the GPU in blocks. Each setup step (operator inspection, the aggregation graph, the off-process prolongator gather, and coarse-operator assembly) operates on the block format directly, and cached, state-gated reuse keeps the repeated Galerkin recompute on the device with no host round trip. On A100 GPUs for 3D elasticity the blocked format extends the problem size solvable with the vendor cuSPARSE backend—which runs out of memory forming the coarse Galerkin product for a 1283 grid on 8 GPUs, where the blocked
12
format and the portable native-Kokkos-Kernels scalar path both fit—and, where the formats run, is faster at scale (1.04 × /1.24 × /1.16× KSPSolve, 1.12 × /1.42 × /1.30× SpMV, 1.45 × /1.80 × /2.27× hot PtAP at 8/27/64 GPUs) and at parity on one GPU. The PtAP advantage grows monotonically with GPU count; the SpMV and KSPSolve advantages peak at 27 GPUs, where the SpMV ratio reaches the 1.42× index-bandwidth ceiling, and narrow at 64 GPUs as the scalar cuSPARSE path scales more efficiently. The advantage is bandwidth-bound in origin: it is largest in the low-nonzeros-per-row regime, where index compression matters most, and narrows as the kernels become flop-bound. The blocked COO assembly primitive (MatCOOUseBlockIndices) is a contribution beyond GAMG: it brings the same bs2 index-and-traffic reduction to any device-resident blocked assembly, with finite-element assembly the natural next application. Several directions remain. Device-resident MIS coarsening (the aggregation graph phase of Section 3.2) has been prototyped: a new Kokkos coarser (MATCOARSENMISKOKKOS) implements parallel Luby-round MIS on the device using deterministic hash weights, is integrated into GAMG’s auto-selection path for device-resident operators, and runs the aggregation without leaving the GPU on single-rank problems. Extending it to the MPI case—where the boundary-vertex neighborhood spans ranks and requires a ghost-layer exchange—and to aggressive coarsening passes are the remaining steps to complete the fully device-resident cold setup. A native block MatAXPY would remove the one residual scalar conversion in the cold prolongator-smoother path (Section 4.9), making the cold setup as conversion-free as the recurring path already is. Overlapping the off-process reduction in the hot PtAP with compute would attack the exposed, imbalance-gated wait that now dominates that phase at scale (Section 4.8). Vendor-specific blocked kernels are a further direction: the portable Kokkos MATBAIJKOKKOS kernels developed here serve as a reference implementation and template from which hand-tuned vendor versions—CUDA for NVIDIA and HIP for AMD—can be derived, and, as the vendor sparse libraries mature, the rectangular-block kernels could eventually migrate onto cuSPARSE and rocSPARSE (the latter’s GEBSR format already exposing the independent row and column block sizes that smoothed-aggregation elasticity requires). Finally, a standalone evaluation of blocked COO for finite-element assembly would quantify the primitive’s benefit in its own right.
Code Availability The implementation described in this paper is part of PETSc [3] and is available in the PETSc git repository at https://gitlab.com/petsc/petsc on the branch adams/seqbaijkokkos-gamg. The model problems are the elasticity examples src/ksp/ksp/tutorials/ex56 and src/snes/ tutorials/ex56.
Use of AI Portions of this work—including code development and debugging, performance-data reduction and analysis, and preparation of this manuscript—were carried out with the assistance of Anthropic’s Claude (Claude Code). All design decisions, experiments, and results were directed, generated, and verified by the author, who takes full responsibility for the content.
13
Acknowledgment This material is based upon work supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) Program through the FASTMath Institute, under contract number DE-AC02-05CH11231 at Lawrence Berkeley National Laboratory.
References [1] NVIDIA Corporation, “cuSPARSE library,” 2024, block compressed sparse row (BSR) format supports equal-size blocks only. [Online]. Available: https://docs.nvidia.com/cuda/cusparse/ [2] S. Rajamanickam, M. Deveci, C. Trott, S. Kim, N. Ellingwood, S. Deveci, M. Perego, and D. Sunderland, “KokkosKernels: Performance portable sparse/dense linear algebra and graph kernels,” in Proc. IEEE High Performance Extreme Computing Conf. (HPEC), 2021. [3] S. Balay, S. Abhyankar, M. F. Adams, S. Benson, J. Brown, P. Brune, K. Buschelman, E. Constantinescu, L. Dalcin, A. Dener, V. Eijkhout, J. Faibussowitsch, W. D. Gropp, V. Hapla, T. Isaac, P. Jolivet, D. Karpeev, D. Kaushik, M. G. Knepley, F. Kong, S. Kruger, D. A. May, L. C. McInnes, R. T. Mills, L. Mitchell, T. Munson, J. E. Roman, K. Rupp, P. Sanan, J. Sarich, B. F. Smith, S. Zampini, H. Zhang, H. Zhang, and J. Zhang, “PETSc/TAO users manual,” Argonne National Laboratory, Tech. Rep. ANL-21/39 - Revision 3.22, 2024. [Online]. Available: https://petsc.org/release/manual/ [4] M. F. Adams, R. Samtaney, and A. Brandt, “Toward textbook multigrid efficiency for fully implicit resistive magnetohydrodynamics,” J. Comput. Phys., vol. 229, no. 16, pp. 6208–6219, 2010. [5] P. Vaněk, J. Mandel, and M. Brezina, “Algebraic multigrid by smoothed aggregation for second and fourth order elliptic problems,” Computing, vol. 56, no. 3, pp. 179–196, 1996. [6] Advanced Micro Devices, Inc., “rocSPARSE library,” 2024, gEBSR (general block sparse row) format supports independent row/column block sizes. [Online]. Available: https://rocm.docs.amd.com/projects/rocSPARSE/ [7] C. R. Trott, D. Lebrun-Grandié, D. Arndt, J. Ciesko, V. Dang, N. Ellingwood, R. Gayatri, E. Harvey, D. S. Hollman, D. Ibanez, N. Liber, J. Madsen, J. Miles, D. Poliakoff, A. Powell, S. Rajamanickam, M. Simberg, D. Sunderland, B. Turcksin, and J. Wilke, “Kokkos 3: Programming model extensions for the exascale era,” IEEE Trans. Parallel Distrib. Syst., vol. 33, no. 4, pp. 805–817, 2022.
14