A GPU-accelerated implementation of the semi-implicit barotropic mode solver for the MPAS-Ocean

A semi-implicit barotropic mode solver for the Model for Prediction Across Scales Ocean (MPAS-O), an ocean component of the Energy Exascale Earth System Model (E3SM), has been ported on GPU using OpenACC directives. Since the semi-implicit solver in MPAS-O consists of a linear iterative solver and a preconditioner that requires linear algebra operations, we incorporated the Matrix Algebra on GPU and Multicore Architecture (MAGMA) and CUBLAS which are collections of linear algebra libraries for heterogeneous architectures. We applied several techniques such as algorithmic changes of the iterative solver, refactorization of loops, and the GPU-aware Message Passing Interface for the global all-to-all node communications to obtain optimized GPU performance. Several numerical tests are performed to evaluate computational performance on Summit supercomputer that equips NVIDIA Tesla V100 GPUs and IBM Power9 CPUs. For the overall runtime of the main solver including data staging, we achieved a 5.4x speedup on 20 nodes (high workload regime per GPU) and a 1.4x speedup on 100 nodes (low workload per GPU) of Summit. We will also show the GPU-accelerated solver performance using hipBLAS supporting AMD MI250X GPU on Frontier.