Bit-Efficient Tensor Core Acceleration for High-Order Shallow-Water Dynamical Core
TL;DR
This work addresses the difficulty in HOPE, a PyTorch-based high-order finite-volume dynamical core, to enable efficient mixed-precision execution on Tensor Cores and proposes a modification of the Ozaki II scheme to improve input mantissa bit utilization in compute-intensive kernels, such as convolution-reformulated stencil computations.
Abstract
Modern accelerators offer substantial performance opportunities through high-throughput low-precision arithmetic. However, exploiting this capability in high-order PDE solvers is challenging: the numerical method imposes strict accuracy requirements and prescribes structured computation patterns. This work addresses the difficulty in HOPE, a PyTorch-based high-order finite-volume dynamical core, to enable efficient mixed-precision execution on Tensor Cores. In particular, this work proposes (1) a modification of the Ozaki II scheme to improve input mantissa bit utilization in compute-intensive kernels, such as convolution-reformulated stencil computations; and (2) a low-precision storage strategy for error-tolerant components in memory-bound sections, such as Riemann solvers, without degrading solution fidelity. Across standard shallow-water benchmarks, the proposed optimizations improve end-to-end performance by 1.13 × –1.34 × relative to the existing channel-stacked Ozaki convolution baseline, while preserving the expected convergence behavior across the validated resolutions and orders. These results demonstrate the effectiveness of exploiting low-precision Tensor Core hardware in high-order dynamical cores.