Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Acceleration of Tensor-Product Operations with Tensor Cores

Cu Cui

TL;DR

The paper addresses accelerating tensor‑product operations in high‑order finite element methods on GPUs by implementing a matrix‑free DG operator with Tensor Cores using both WMMA and inline PTX MMA approaches. It demonstrates that MMA‑based kernels with conflict‑free shared memory layouts can reach up to about 8 TFLOPS/s (roughly 45% of the FP64 Tensor Core peak) and yield 2.3× speedups over optimized CUDA Cores, closely matching roofline predictions. In half precision, the authors achieve around a 3.5× speedup while maintaining accuracy through an error‑correction strategy, and they validate these benefits in a practical Poisson solve using a multigrid preconditioner within FGMRES. The work also provides a roofline and shared‑memory bandwidth analysis to identify bottlenecks and demonstrates a fourfold improvement in solving Poisson problems with mixed‑precision multigrid, highlighting the practical impact of Tensor Cores for tensor‑product FEM operators.

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.

Acceleration of Tensor-Product Operations with Tensor Cores

TL;DR

The paper addresses accelerating tensor‑product operations in high‑order finite element methods on GPUs by implementing a matrix‑free DG operator with Tensor Cores using both WMMA and inline PTX MMA approaches. It demonstrates that MMA‑based kernels with conflict‑free shared memory layouts can reach up to about 8 TFLOPS/s (roughly 45% of the FP64 Tensor Core peak) and yield 2.3× speedups over optimized CUDA Cores, closely matching roofline predictions. In half precision, the authors achieve around a 3.5× speedup while maintaining accuracy through an error‑correction strategy, and they validate these benefits in a practical Poisson solve using a multigrid preconditioner within FGMRES. The work also provides a roofline and shared‑memory bandwidth analysis to identify bottlenecks and demonstrates a fourfold improvement in solving Poisson problems with mixed‑precision multigrid, highlighting the practical impact of Tensor Cores for tensor‑product FEM operators.

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.
Paper Structure (21 sections, 17 equations, 14 figures, 4 tables, 2 algorithms)

This paper contains 21 sections, 17 equations, 14 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Problem description
  3. Matrix-free implementation of discontinuous Galerkin finite element operator
  4. Programming Tensor Cores
  5. C++ Warp Matrix Functions
  6. Inline PTX instructions
  7. Kernel design
  8. Double Precision benchmark
  9. Kernel optimizations
  10. Performance analysis
  11. Half Precision benchmark
  12. Performance analysis
  13. Precision loss
  14. Arbitrary matrix size
  15. Roofline Analysis
  16. ...and 6 more sections

Figures (14)

  • Figure 1: Compute pattern for patch-wise integrals in 2D. Orange indicates cell integrals, while gary indicates face integrals. From left to right: center patch, top boundary patch, right boundary patch and corner patch.
  • Figure 2: Differences of three data movement instructions.
  • Figure 3: Shared memory access pattern for matrix-matrix multiplications in double precision. The number in the grid indicates the bank number corresponding to the shared memory address. $T_i$ represents thread with index $i$.
  • Figure 4: Permuted shared memory data layout in double precision. Within each 4x4 sub-matrix, every thread accesses distinct bank numbers, adhering to a conflict-free access pattern.
  • Figure 5: Arithmetic performance of implementation variants for finite element operator $Au$ with double precision in 3D for $N=8$.
  • ...and 9 more figures