A GEMM-based direct solver for finite-difference Poisson problems in non-uniform grids

This paper presents a robust, GEMM-based direct solver for finite-difference Poisson problems on non-uniform 3D Cartesian grids that leverages tensor formulations and matrix-matrix multiplications to achieve superior time-to-solution and parallel efficiency compared to traditional multigrid and FFT-based methods on modern heterogeneous hardware.

The Gist

Imagine you are trying to solve a massive, three-dimensional puzzle. This puzzle represents the flow of a fluid (like water or air) around an object. To solve it, you need to figure out the "pressure" at every single point in the fluid. Mathematically, this is called solving the Poisson equation.

In the past, scientists had two main ways to solve this puzzle:

1. The "Uniform Grid" Method: Imagine a perfectly flat, square checkerboard. This is easy to solve quickly using a magic trick called the FFT (Fast Fourier Transform). It's like having a super-fast calculator that knows the answer instantly if the grid is perfectly even.
2. The "Non-Uniform Grid" Method: Real life isn't a perfect checkerboard. Near a wall, you need tiny, detailed squares to see the turbulence. Far away, you can use huge, lazy squares. This is a stretched grid. The problem? The "magic trick" (FFT) breaks when the squares aren't equal. The old way to solve this was to use a slow, step-by-step method (like a multigrid solver) that often got stuck or took forever.

The Breakthrough: The "GEMM" Solver

This paper introduces a new, super-smart way to solve the puzzle on these uneven, stretched grids. The authors, working with NVIDIA, created a method that swaps the old "magic trick" for a different kind of super-power: GEMM (General Matrix-Matrix Multiplication).

Here is the simple analogy:

The Analogy: The Library and the Librarian

Imagine you are a librarian trying to organize a massive library of books (the fluid data).

• The Old Way (FFT): If the library shelves are all the same height and perfectly aligned, you can use a special conveyor belt system (FFT) to sort the books in seconds. But if the shelves are different heights (non-uniform grid), the conveyor belt jams. You have to manually move every single book one by one, which takes forever.
• The New Way (GEMM): Instead of a conveyor belt, you hire a team of incredibly fast robots (the GPU) that are experts at stacking boxes.
  • The authors realized that even if the shelves are uneven, you can mathematically "reshape" the problem so the robots can still stack the books efficiently.
  • Instead of sorting one book at a time, the robots grab huge blocks of books and stack them all at once. This is Matrix Multiplication.
  • Modern computer chips (GPUs) are built specifically to do this "stacking" incredibly fast. In fact, they are so good at it that they can often beat the old conveyor belt, even if the math is slightly more complex.

How It Works (The "Secret Sauce")

1. Symmetrizing the Problem: The uneven grid makes the math messy and asymmetrical (like a wobbly table). The authors found a clever trick to "level the table" using a simple scaling factor. This makes the math look like a perfect, symmetrical problem again, allowing them to use their powerful robot stackers.
2. The Hybrid Approach: The best part? The system is flexible. If a part of your grid is perfectly even, it uses the fast conveyor belt (FFT). If a part is stretched and uneven, it switches to the robot stackers (GEMM). You can mix and match them in the same simulation!
3. The Result:
   • Speed: On a single computer, this new method is up to 100 times faster than the old slow methods for stretched grids.
   • Scale: When they ran this on supercomputers with thousands of GPUs, the "robot stackers" (GEMM) were so efficient that they didn't get slowed down by the time it took to talk to each other (communication overhead). The old methods slowed down significantly as you added more computers, but this new one kept flying.

Why Does This Matter?

This is a game-changer for weather forecasting, airplane design, and climate modeling.

• Realism: To simulate a real airplane wing, you need tiny details near the surface and huge spaces far away. You can't use a uniform grid without wasting millions of hours of computing power.
• Efficiency: This new solver allows scientists to use those detailed, stretched grids without the simulation taking weeks to finish. It saves massive amounts of time and energy.

In a nutshell: The authors took a problem that was hard to solve on uneven grids and figured out how to use the most powerful, optimized math operation available on modern supercomputers (matrix multiplication) to solve it. They turned a "slow, manual process" into a "high-speed robot assembly line," making high-resolution fluid simulations faster and more accessible than ever before.

Technical Summary

Here is a detailed technical summary of the paper "A GEMM-based direct solver for finite-difference Poisson problems in non-uniform grids."

1. Problem Statement

The numerical solution of the 3D Poisson equation is a critical bottleneck in Direct Numerical Simulations (DNS) of incompressible Navier–Stokes equations, where it is used to enforce the divergence-free constraint (pressure projection).

• The Challenge: High-fidelity simulations of turbulent flows often require non-uniform (stretched) grids to resolve boundary layers efficiently. However, standard direct solvers based on Fast Fourier Transforms (FFT) rely on uniform grid spacing.
• Limitations of Alternatives:
  • Geometric Multigrid (MG): While flexible for non-uniform grids, MG performance degrades significantly on strongly stretched meshes due to reduced effectiveness in smoothing low-frequency errors, often requiring many more iterations.
  • GPU Inefficiency: Many existing efficient solvers (like Block Cyclic Reduction) do not map well to modern GPU architectures, which thrive on massive parallelism and high arithmetic intensity, whereas MG and reduction methods often lead to small local problems and under-utilization of GPU resources.
• Goal: Develop a direct Poisson solver that supports non-uniform grids, maintains high efficiency on modern heterogeneous hardware (GPUs), and offers better performance than iterative multigrid methods on stretched meshes.

2. Methodology

The authors propose a tensor-based direct solver that generalizes the classical eigenfunction expansion method to non-uniform grids.

Core Algorithm

1. Separation of Variables: The 3D Poisson operator is decomposed using Kronecker sums. The solver diagonalizes the operator along two directions (e.g., $x$ and $y$) and solves the remaining direction ($z$) directly via a tridiagonal solver.
2. Symmetrization for Non-Uniform Grids:
   • On non-uniform grids, the 1D finite-difference operator is non-symmetric, making eigendecomposition difficult.
   • The authors apply a diagonal similarity transformation ($D^{1/2} T D^{-1/2}$) to symmetrize the tridiagonal matrix. This preserves eigenvalues while allowing the use of efficient symmetric eigensolvers.
   • The original eigenvectors are recovered via scaling: $Q = D^{-1/2} \tilde{Q}$.
3. GEMM-Based Transforms:
   • Instead of using FFTs (which require uniform grids), the solver performs dense eigenbasis transforms along the non-uniform directions.
   • These transforms are implemented as General Matrix-Matrix Multiplications (GEMM). By aggregating many independent 1D transforms into a single large matrix operation, the method achieves high arithmetic intensity, mapping efficiently to optimized BLAS kernels and GPU Tensor Cores.
4. Hybrid Capability: The implementation (an extension of the CaNS solver) allows a hybrid approach:
   • Uniform directions: Use FFTs (fast, $O(N \log N)$).
   • Non-uniform directions: Use GEMM-based dense transforms (flexible, $O(N^2)$).
   • The third direction is always solved via a tridiagonal solver (Thomas algorithm or PCR-TDMA).

Implementation Details

• Parallelization: Uses a 2D "pencil" domain decomposition with MPI for distributed memory.
• Communication: Relies on collective transpose operations (handled by 2DECOMP&FFT on CPUs and cuDecomp on GPUs) to align data for transforms in different directions.
• Hardware: Optimized for both CPUs (OpenBLAS) and GPUs (cuBLAS, cuFFT).

3. Key Contributions

• First Direct Solver for Non-Uniform Grids on GPUs: Extends the efficiency of direct eigenfunction methods to arbitrary stretched grids by replacing FFTs with GEMM-based dense transforms.
• Symmetrization Technique: Demonstrates a robust method to symmetrize non-uniform finite-difference operators via diagonal scaling, enabling the use of standard symmetric eigensolvers (LAPACK xSTEDC).
• Hybrid FFT/GEMM Framework: Provides a flexible solver where users can mix FFTs and GEMMs depending on grid uniformity in specific directions, optimizing the trade-off between computational cost and grid flexibility.
• GPU Optimization: Leverages the high arithmetic intensity of GEMM to amortize communication and transpose overheads, making it superior to FFT-heavy methods on massively parallel GPU systems.

4. Results and Performance

The solver was validated against reference DNS data (lid-driven cavity and turbulent square duct) and benchmarked against state-of-the-art alternatives (Geometric Multigrid, FFT+BLKTRI).

Accuracy

• The solver produces solutions with machine-precision residuals.
• Validation cases (Re=1000 cavity, Re$\approx$150 turbulent duct) showed excellent agreement with reference DNS data.

CPU Performance

• Single Core: The direct GEMM-based method is 1.5–2 orders of magnitude faster than Geometric Multigrid on strongly stretched grids (where MG struggles with convergence).
• Strong Scaling: GEMM-based variants (GG) scale better than FFT-based ones (FF) on many-core CPUs. The compute-intensive nature of GEMM better amortizes communication overheads.
• Weak Scaling: FFT-based synthesis shows mild growth as the grid size increases ($O(N \log N)$), while GEMM-based synthesis grows quadratically ($O(N^2)$). However, the absolute time-to-solution for GEMM remains competitive due to the massive reduction in total grid points enabled by stretching.

GPU Performance (NVIDIA GB200)

• Single GPU: The fully non-uniform (GG) case incurred a 2.8x slowdown in Poisson solve time compared to the uniform (FF) case, but only a 1.8x slowdown in the total Navier–Stokes time step.
• Strong Scaling: The solver scales efficiently up to 64 GPUs. GEMM-heavy variants maintain higher parallel efficiency (45–66%) than FFT-heavy variants because they are less limited by memory bandwidth and communication latency.
• Break-even Point: The paper estimates that using non-uniform grids to reduce the total grid count by a factor of 2–3 is sufficient to offset the higher computational cost of GEMM transforms compared to FFTs on uniform grids.

5. Significance

• Enables High-Resolution Stretched Meshes: This method allows researchers to perform DNS on highly stretched grids (essential for wall-bounded turbulence) without resorting to slow iterative multigrid solvers or sacrificing the accuracy of direct solvers.
• Future-Proof for Heterogeneous Hardware: By prioritizing GEMM operations, the solver is naturally aligned with the trajectory of modern supercomputers, which are increasingly dominated by GPU accelerators and AI-driven workloads that favor dense linear algebra.
• Versatility: The hybrid approach allows the solver to adapt to mixed grid types (e.g., uniform in streamwise, stretched in wall-normal), maximizing efficiency across different problem configurations.

In conclusion, the paper presents a robust, high-performance direct solver that bridges the gap between the efficiency of FFT-based methods and the geometric flexibility required for non-uniform grids, specifically optimized for the next generation of GPU-accelerated supercomputers.