Source-linked AI summary
Mixed-Precision Coalesced Matrix-Free Multigrid with Vertex-Patch Smoothing
Michał Wichrowski
TL;DR
GPU multigrid with cell-wise storage faces both data-movement demands and loss of cancellation when unassembled residuals are rounded in FP32. This paper combines vertex-patch smoothing with mixed precision and FP64 rebalancing; the resulting solver matches FP64 cycle convergence and exceeds 400 million degrees of freedom per second on one A100.
Problem
GPU multigrid needs data-efficient patch smoothing, while FP32 rounding of unassembled cell-wise residuals can lose cancellation between opposite-sign fluxes.
Source: §1 · p. 1, §8 · p. 21Method
Results
416 million degrees of freedom per second on one A100 is 2.2 times the throughput reported in; rebalanced cycles converge as in FP64 independently of mesh size.
Source: §1 · p. 3, §8 · p. 22Takeaways & Limitations
The solver preserves continuous iterates without DSS while achieving mesh-independent convergence and optimal error convergence down to the precision of the solve.
Source: §1 · p. 2, §8 · p. 22Fast diagonalization is limited to separable problems; more general local problems may require local p-multigrid, and higher-degree mappings remain to be tested.
Source: §8 · p. 22
Abstract
from arXiv · showhide
I present a geometric multigrid solver for the Poisson equation with continuous cubic elements in three dimensions, smoothed by a multiplicative vertex-patch Schwarz method with exact local solves. All fields are kept in persistent cell-wise storage, where every cell holds its own copy of each shared degree of freedom. The patch smoother yields a continuous solution without any exchange between cells. On the GPU, the patch update is computed in registers, one cell per thread, and the cells of a patch exchange data by warp shuffles, so shared memory is avoided almost entirely. Despite the low occupancy that this design implies, the arithmetic of the patch update runs close to the FP32 peak. Fast diagonalization, merged transforms, and a reflection symmetry of the patch then reduce the arithmetic until a sweep moves its data at $87\%$ of the bandwidth of a plain device copy. The hierarchy therefore runs in single precision inside a double-precision outer iteration. Rounding the unassembled residual destroys the cancellation between opposite-sign copies of shared degrees of freedom; exchanging the copies between cells in double precision performs this cancellation before rounding, without changing the dual vector they represent. With this \emph{rebalancing}, the cycle converges as in double precision and independently of the mesh size, the error converges at the optimal rate down to the precision of the solve, and the solver reaches over $400$ million degrees of freedom per second on a single NVIDIA A100 GPU.
1 Introduction
The paper combines cell-wise geometric multigrid, multiplicative vertex-patch smoothing, and mixed precision to make GPU solves efficient without assembling shared degrees of freedom. Rebalancing the residual in FP64 preserves convergence, while register-based patch updates and arithmetic reductions improve throughput.
1 Introduction
- Cell-wise storage keeps copies of shared degrees of freedom in each cell, while the patch smoother preserves continuity without DSS or other synchronization.The patch writes the same correction into every copy of each patch degree of freedom.Source: §1 · p. 2
- Register-based vertex-patch smoothing runs near the FP32 arithmetic peak, addressing the shared-memory bandwidth limit of the GPU smoother in.The Q3 kernel reaches 96% of FP32 peak, with its pipeline 99.6% busy despite low occupancy.Source: §1 · pp. 1–2
- FP32 rounding of unassembled residuals loses cancellation between opposite-sign cell fluxes; FP64 rebalancing cancels them before rounding without changing the represented dual vector.The stationary iteration rebalances each cycle, while CG can reuse one rebalancing over many iterations.Source: §1 · p. 2
- Fast diagonalization, merged transforms, and patch reflection symmetry reduce arithmetic until the sweep approaches the memory-traffic limit.In a prototype moving only the iterate and correction, the final variant takes 14% longer than a kernel performing only that traffic.Source: §1 · p. 3
- 416 million degrees of freedom per second on one A100 is 2.2 times the throughput reported in.A smoothing sweep moves data at 87% of plain device-copy bandwidth, and rebalancing takes 5% of solve time.Source: §1 · p. 3
2 Problem Setting
The solver treats the three-dimensional Poisson problem on hexahedral meshes with continuous tensor-product elements, using a cell-wise representation rather than a vector of unique degrees of freedom. Its hierarchy applies operators and transfers cell by cell, with a vertex-patch smoother and either CG or a stationary outer iteration.
2 Problem Setting
- The experiments use continuous Q3 elements with Gauss–Lobatto nodal bases on Cartesian meshes and uniform refinement levels.The problem has homogeneous Dirichlet boundary conditions on a bounded domain in three dimensions.Source: §2 · p. 3
- Cartesian meshes are a deliberate scope choice because their uniform interior patches are the kernel of the smoother, including for background meshes used by unfitted discretizations.The GPU patch-smoother reference is also restricted to Cartesian meshes.Source: §2 · p. 3
- Cell-wise storage retains each cell's copies of shared degrees of freedom; the multigrid preconditioner returns continuous corrections by construction, without intercell exchange.The operator maps primal vectors to unassembled dual residuals, and the preconditioner depends on a residual through its assembly.Source: §2 · pp. 3–4
- The V-cycle uses cell-wise transfers and matrix-free level operators, while the outer iteration is CG or a stationary update monitored by a preconditioned residual pairing.For the unsymmetric cycle, the pairing is only a monitor; in the stationary iteration it lags one cycle behind the iterate.Source: §2 · p. 4
3 The Vertex-Patch Smoother
The multiplicative vertex-patch smoother assembles each patch residual from cell copies, solves the local problem exactly, and writes corrections consistently across those copies. Fast diagonalization is adapted to the cell-wise patch representation so the local solve can be distributed across cells without forming unique patch values.
3 The Vertex-Patch Smoother
- Each interior-vertex patch contains eight cells, and its update solves the local problem exactly, annihilating the assembled residual on the patch space.The patch residual is assembled from the raw cell residuals belonging to that patch.Source: §3 · pp. 4–5
- Writing the same correction into every copy keeps the iterate continuous, so the smoother needs no DSS, gather–scatter, valence, or explicit continuity constraints.Rebalancing is the solver's only intercell exchange and serves residual precision rather than continuity.Source: §3 · p. 5
- Eight patch colors update concurrently without conflicts; reversing their order gives the adjoint sweep, enabling a symmetric V-cycle for CG.Forward post-smoothing instead produces the unsymmetric cycle used as a stationary iteration.Source: §3 · p. 5
3.1 Fast Diagonalization of the Cell-Wise Patch
- The assembled patch solve has 125 unique unknowns, but applying it directly would require assembling and redistributing values and would distribute work unevenly across cells.The cell-wise patch holds 216 copies of its values across eight cells.Source: §3.1 · p. 6
- Fast diagonalization instead acts on cell-wise patch values by composing eigenvectors with duplication and padding, incorporating shared-node assembly and scatter into the transforms.The forward transform sums shared-node copies, and the backward transform produces corrections in both copies without forming the 125 unique values.Source: §3.1 · pp. 5–6
- Padding makes the modal dimension six per axis, allowing an even split between cells and mapping equal amounts of work onto eight threads.The resulting transforms have six rather than five rows; later reductions recover more than this added cost.Source: §3.1 · p. 6
4 Mixed Precision
The FP32 hierarchy reduces memory traffic and uses the GPU’s higher FP32 arithmetic peak, but residual cancellation must be preserved for convergence. Rebalancing the cell-wise residual before rounding restores FP64-like convergence, with error rates remaining optimal until solve precision is reached.
4 Mixed Precision
- FP32 halves vector storage and traffic and doubles arithmetic peak relative to FP64, benefits that matter only if each cycle converges as fast.A smoother sweep visits every cell eight times, reading u and the right-hand side and writing u.Source: §4 · p. 7
4.1 Precision Split
- The level operators and transfers use TF32 tensor-core contractions with split operands; their relative error is about 10^-6 against FP64.The omitted low-by-low product has relative size 2^-22, and terms are accumulated in FP32.Source: §4.1 · p. 7
4.2 Test Configuration
- Experiments use Q3 elements on the unit cube, global refinements, a 10^-9 monitor tolerance, and one pre- and post-smoothing sweep per V-cycle.Throughput counts unique DoFs; solve timings exclude hierarchy setup and one warm-up cycle.Source: §4.2 · pp. 7–8
4.3 Stagnation
- The FP32 cycle stalls without rebalancing because rounding the unassembled residual loses the small residual hidden by opposite-sign O(1) fluxes.On the L = 6 mesh, the stationary cycle stalls near 6 · 10^-8 and CG fluctuates near 2 · 10^-8.Source: §4.3 · pp. 8–9
4.4 Rebalancing
- FP64 exchange averages shared DoF copies before rounding, removing flux cancellation while preserving the represented dual vector.The patch smoother needs no modification, and stationary iterations rebalance each cycle while CG can rebalance every few iterations.Source: §4.4 · pp. 9–10
4.5 Convergence with Rebalancing
- With rebalancing in every cycle, four stationary V-cycles reach tolerance 10−9 on every mesh; the L2 error converges optimally through L=6, then reaches solve precision from L=7 onward.The fractional cycle count is virtually identical across meshes and to the same cycle in FP64.Source: §4.5 · p. 10, Table 2 · p. 10
- Mixed-precision GMRES is left to future work because its Arnoldi-based residual estimate is only single-precision accurate with an FP32 V-cycle.The paper notes that a Krylov method gains less than one cycle over the stationary iteration.Source: §4.5 · p. 11
5 In-Register Evaluation
The section maps patch work onto registers and warp shuffles, then organizes memory access to preserve coalescing. The final kernel avoids an on-chip-memory limit and is bound by DRAM bandwidth.
5.1 Thread Mapping
- The cell-wise mapping keeps each contraction within a thread's registers and uses warp shuffles only for patch-cell coupling.The patch is balanced across cells, unlike assembled-patch mapping, which redistributes data through shared memory after each contraction stage.Source: §5.1 · pp. 12–13
5.2 Memory Layout
- Cell-contiguous storage makes one cell's coefficients adjacent, but the one-cell-per-thread access pattern runs 2.9 times slower than blocked storage without coalescing support.Vector loads alone gain 6%; the kernel instead stages row-wise vector loads through shared memory to hand values to their owning threads.Source: §5.2 · p. 14
- 4.06 to 3.06 ms on 223 cells when removing the second read of u and the store; keeping cell data in cache reduces the time to 2.95 ms.The second read and store are hidden behind arithmetic in the final kernel.Source: §5.2 · p. 15
5.3 Relation to Shared-Memory Realizations
- Registers avoid the shared-memory bandwidth bound of staged contractions, whose traffic can match the arithmetic they serve at p = 3.The shared-memory vertex-patch design stages operands between steps and separates them with block-wide synchronizations.Source: §5.3 · p. 15
- The final kernel is DRAM-bound, not on-chip-memory-bound, and leaves 59% of the FP32 peak unused.Its shared-memory traffic is 1,536 bytes per cell visit by count, while measurement gives 1,600 bytes, including 4.2% bank conflicts.Source: §5.3 · p. 15
6 Reducing the Arithmetic
The section progressively reduces arithmetic in an in-register patch update through accumulator chains, eigenbasis residuals, merged transforms, and reflection symmetry. The final variant approaches the traffic bound for the constant-coefficient Laplacian.
6 Reducing the Arithmetic
- Table 8 compares result-equivalent color-step variants against a same-traffic, no-arithmetic stub; its prototype omits the right-hand-side transform and read-modify-write traffic.Source: §6 · p. 16, Table 8 · p. 16
6.1 Accumulator Chains
- Accumulator chaining removes 488 instructions per thread and 11% of run time without changing the arithmetic result.Source: §6.1 · p. 16
6.5 Arithmetic Rate
- For the constant-coefficient Laplacian, the final variant is 1.7 times faster and, at one update per load, within 14% of the traffic stub.Source: Figure 2 · p. 17, §6.5 · p. 19
6.2 Residual in the Cell Eigenbasis
- For cubic elements, the cell eigenbasis replaces seven tensor contractions with six and a pointwise product, without accuracy loss in single precision.The change of basis has cond(S) = 2.94.Source: §6.2 · p. 17
6.3 Merged Transforms
- Merging adjacent transforms removes three contractions because the nodal residual need not be materialized between the back-to-back diagonalizations.This is the largest single gain in Table 8.Source: §6.3 · p. 18
6.4 Reflection Symmetry
- Equal-sized cells let reflection parity make a partner block a signed copy, so lanes fold partner coefficients and contract only their own block.Source: §6.4 · p. 18
7 Performance
The final solver pairs an optimized smoother with an unsymmetric stationary outer iteration, reaching 416 million DoFs per second on the finest mesh. Small-mesh performance remains limited by kernel-launch overhead.
7.1 Smoother Throughput
- 1521 GB/s, or 87% of the 1756 GB/s plain device-copy bandwidth, is reached by a complete eight-color smoothing sweep.The correction agrees with an FP64 evaluation of the same color step to FP32 accuracy.Source: §7.1 · p. 19, Table 9 · p. 20
7.2 Choice of the Outer Iteration
- The best Krylov combination reduces the residual by only a factor of two to five beyond the unsymmetric cycle’s result, less than one cycle’s gain.The fractional iteration counts compare outer iterations using the same mixed-precision V-cycle.Source: §7.2 · p. 20, Table 10 · p. 20
- 3.7 cycles for the unsymmetric stationary iteration beats 4.7 CG iterations and 7.2 cycles for the symmetric stationary cycle.The symmetric stationary cycle has a contraction factor fifteen times larger than the unsymmetric one.Source: §7.2 · p. 19, Table 10 · p. 20
7.3 Solve Time
- The patch smoother takes 60% of solve time, while rebalancing takes 5%.The timed solves apply four cycles; stopping on the lagging measured monitor requires a fifth cycle, while an estimated monitor avoids it but is not a bound.Source: §7.3 · p. 20
- Small-mesh solve time is launch-bound; L=4 and L=5 take nearly the same time despite L=5 having eight times as many cells.The number of smoothing launches per V-cycle is independent of mesh size, and small levels may not contain enough patches to saturate the device.Source: §7.3 · pp. 20–21
- On the finest mesh, throughput reaches 416 million DoFs/s; on L=8, it is 2.2 times Cui et al.’s 188 million DoFs/s on the same GPU model.Both solves apply four V-cycles.Source: §7.3 · pp. 20–21, Table 12 · p. 21
8 Conclusions
FP64 rebalancing prevents residual stagnation in the FP32 hierarchy, restoring mesh-independent cycle convergence and optimal error convergence to solve precision. The register-based patch kernel approaches copy bandwidth, while its arithmetic reductions rely on a restricted patch structure.
8 Conclusions
- Persistent cell-wise storage lets each vertex patch assemble its residual and write identical corrections to its cell copies, preserving continuity without inter-cell exchange.Diagonalizing cell-wise rather than assembled patch matrices absorbs assembly and scatter into transforms and splits the local solve into eight identical cell shares.Source: §8 · p. 21
- FP64 rebalancing cancels opposite-sign inter-element fluxes before rounding, restoring FP64-like cycle convergence and optimal error convergence down to solve precision.Four cycles reduce the residual by nine orders of magnitude on every mesh; this is the solver’s only inter-cell exchange and serves residual precision, not continuity.Source: §8 · pp. 21–22
- The unsymmetric cycle as a stationary iteration is the fastest outer iteration, and a Krylov method gains less than one cycle over it.The stated comparison concerns the outer iterations used with the mixed-precision V-cycle.Source: §8 · p. 22
- Register-resident patch updates with warp-shuffle exchanges run near the FP32 arithmetic peak; fast diagonalization, merged transforms, and patch reflection symmetry reduce a sweep to 87% of plain-copy bandwidth.The solver reaches about 400 million DoFs per second on one A100, 2.2 times the throughput of the mixed-precision GPU patch smoother at the same number of V-cycles.Source: §8 · p. 22
- Fast diagonalization requires Kronecker-sum patch matrices, limiting this arithmetic reduction to separable problems with constant coefficients and equal-sized cells within a patch.For more general problems, the exact local solve can be replaced by local p-multigrid; locally refined hierarchies and higher-degree mappings remain to be tested.Source: §8 · p. 22