Optimized thread-block arrangement in a GPU implementation of a linear solver for atmospheric chemistry mechanisms

TL;DR

This work tackles the high computational cost of chemical solvers in atmospheric chemistry models and introduces Block-cells, a GPU-centric strategy that distributes the solver workload across CUDA thread blocks within the CAMP framework. By coupling a GPU-accelerated BCG linear solver and reorganizing work so that each block handles multiple cells, the approach reduces costly CPU-GPU reductions and improves memory and compute utilization. Empirical results show substantial speedups: up to 35x over a single-core CPU baseline and 50% faster than a node-wide 40-core CPU, with Block-cells outperforming traditional MPI CPU deployments. The findings indicate Block-cells as a powerful path toward scalable, GPU-based chemical solvers in Earth system models, with clear avenues for extending to other solver components and hybrid architectures.

Abstract

Earth system models (ESM) demand significant hardware resources and energy consumption to solve atmospheric chemistry processes. Recent studies have shown improved performance from running these models on GPU accelerators. Nonetheless, there is room for improvement in exploiting even more GPU resources. This study proposes an optimized distribution of the chemical solver's computational load on the GPU, named Block-cells. Additionally, we evaluate different configurations for distributing the computational load in an NVIDIA GPU. We use the linear solver from the Chemistry Across Multiple Phases (CAMP) framework as our test bed. An intermediate-complexity chemical mechanism under typical atmospheric conditions is used. Results demonstrate a 35x speedup compared to the single-CPU thread reference case. Even using the full resources of the node (40 physical cores) on the reference case, the Block-cells version outperforms them by 50%. The Block-cells approach shows promise in alleviating the computational burden of chemical solvers on GPU architectures.

Figures (7)

• Figure 1: Workflow diagram of the One-cell (left) and Multi-cells (right) solving strategies.
• Figure 2: Diagram of BCG Multi-cells interactions between CPU and GPU (excluding data transfers). $GPU_1$ contains all the BCG operations in the GPU except for $Reduce$ (e.g., functions like dot vector, SPMV, etc.). The $Reduce$ computation is divided between CPU and GPU. Each thread block performs a reduction operation on the GPU, resulting in a convergence error for each block. Then, another reduction is performed over these errors on the CPU side. The final value of the reduction is used to check for convergence, i.e., if the algorithm needs to iterate again or finish.
• Figure 3: Diagram of Multi-cells and Block-cells interactions between CPU and GPU (excluding data transfers).
• Figure 4: Reduction in the number of solving iterations of Block-cells (1) against Block-cells (N) using ideal (orange dashed line) or realistic (blue line) initial conditions. This reduction is calculated by dividing the solving iterations of Block-cells (N) by Block-cells (1) and corresponds to the iterations of the last thread block to finish the algorithm. The reduction is averaged over 720 time steps. The blue shade indicates the standard deviation of the reduced iterations for all time steps using realistic conditions.
• Figure 5: Speedups and iterations of GPU linear solver configurations. The speedup and iterations are normalized over the worst case, which corresponds to Block-cells (N) and Block-cells (3), respectively. The Block-cells (1), (2),(3), and (N) implementations correspond to block sizes 1,2,3 and 6. The results are averaged over 720 time steps and are configured with 10,000 cells and realistic conditions.