Acceleration of Tensor-Product Operations with Tensor Cores

TL;DR

This paper demonstrates how leveraging NVIDIA A100 GPU Tensor Cores significantly accelerates tensor product operations in finite element methods, achieving over twofold speedups and maintaining accuracy.

Contribution

It introduces a novel implementation strategy using inline PTX instructions and shared memory, resulting in substantial performance improvements over traditional CUDA Cores.

Findings

2.3-fold increase in double precision performance

Fourfold speedup in half-precision Poisson solver

Maintains discretization accuracy comparable to double precision

Abstract

In this paper, we explore the acceleration of tensor product operations in finite element methods, leveraging the computational power of the NVIDIA A100 GPU Tensor Cores. We provide an accessible overview of the necessary mathematical background and discuss our implementation strategies. Our study focuses on two common programming approaches for NVIDIA Tensor Cores: the C++ Warp Matrix Functions in nvcuda::wmma and the inline Parallel Thread Execution (PTX) instructions mma.sync.aligned. A significant focus is placed on the adoption of the versatile inline PTX instructions combined with a conflict-free shared memory access pattern, a key to unlocking superior performance. When benchmarked against traditional CUDA Cores, our approach yields a remarkable 2.3-fold increase in double precision performance, achieving 8 TFLOPS/s-45% of the theoretical maximum. Furthermore, in half-precision computations, numerical experiments demonstrate a fourfold enhancement in solving the Poisson equation using the flexible GMRES (FGMRES) method, preconditioned by a multigrid method in 3D. This is achieved while maintaining the same discretization error as observed in double precision computations. These results highlight the considerable benefits of using Tensor Cores for finite element operators with tensor products, achieving an optimal balance between computational speed and precision.