Analyzing GPU Mixed-Precision and Tensor Cores for High-Order FEM for Seismic Wave Propagation

Abstract

Spektrale elementers metode (SEM) er en høyere ordens endelige elementers metode (FEM) som brukes til å finne tilnærmede løsninger av partielle differensialligninger. Metoden kan blant annet brukes når man simulerer seismisk bølgforplantning, noe som er sentralt innenfor geofysikk for å undersøke jordas egenskaper. SEM egner seg godt til parallelisering ved hjelp av GPUer. Derimot kreves det store mengder langringsplass og båndbredde for å kunne gjennomføre seismiske simuleringer, og man ønsker stadig å prosessere mer data, med høyere oppløsning og med mer komplekse modeller. Dette understreker behovet for implementasjoner med høy ytelse som kan utnytte egenskapene til nye GPU-arkitekturer.

Det er en økende interesse for algoritmer som bruker blandet presisjon og for å bruke flere nivå av presisjon i samme applikasjon, med mål om å redusere mengden lagring og båndbredde som kreves i en applikasjon. Videre finnes det nå maskinvare som har innebygd støtte for å gjøre beregninger med blandet presisjon. Et eksempel er tensorkjernene (Tensor Cores) fra NVIDIA som er spesielt utviklet for å utføre generell matrise-multiplikasjon (GEMM).

I dette arbeidet undersøker vi bruken av NVIDIAs tensorkjerner og blandet presisjon til å utføre tensorsammentrekningene i SEM. Fokuset er på seismisk bølgeforplantning i tre dimensjoner. Vi evaluerer ytelsen til og sammenligner bibliotekene cuBLAS og CUTLASS for å utføre GEMM-sekvensen i SEM, og inkluderer biblioteksrutiner i en grunnimplementasjon av SEM med enkel presisjon. Undersøkelsene viser at blandet presisjon er tilstrekkelig når man simulerer seismisk bølgeforplantning i et homogent medium, både ved bruk av presisjonsformatet TensorFloat32 fra NVIDIA, og ved bruk av halv presisjon når man inkluderer en skaleringskonstant for å unngå underflyt i beregningene. Med tanke på ytelse har både cuBLAS og CUTLASS høyt minnebruk og CPU-overheng når de brukes i vert-side CUDA-kode, og våre resultater viser at de ikke kan konkurrere med implementasjoner som utfører beregningene i farten uten blandet presisjon og tensorkjerner. Derimot viser vi at man kan lagre data i halv presisjon og bruke blandet presisjon uten tensorkjerner til å utføre tensorsammentrekningene i grunnimplementasjonen av SEM, noe som gir en økt hastighet på ~5%. Dette demonstrerer at å bruke lavere presisjon i implementasjoner av SEM burde tas i betraktning både for økt ytelse og reduserte lagringskrav. Vi presenterer også mulige retninger for fremtidig arbeid.

The Spectral Element Method (SEM) is a high-order Finite Element Method (FEM) used to find approximate solutions to partial differential equations (PDEs). Solving PDEs is a central activity in forward simulation of seismic wave propagation, which is important in geophysical research to understand the earth's structures and properties. SEM is suitable for parallelization using GPUs. However, a large amount of computing power and memory is needed to realize a simulation, and there is a desire to process more seismic data at higher resolutions and use more complex models. This motivates the search for high-performance implementations that can fully utilize emerging GPU architectures.

There is a growing interest in mixed-precision algorithms and multi-precision implementations to reduce storage and bandwidth requirements of applications. This entails using low precision computations where appropriate or using low precision calculations to emulate the accuracy that can be achieved with higher precision calculations. Furthermore, recent years have seen the appearance of hardware support for mixed-precision computations, one example being NVIDIA's Tensor Cores specialized at performing general matrix multiplication (GEMM).

In this work, we implement a basic SEM applications which we then use to investigate the use of NVIDIA Tensor Cores and mixed-precision GEMM for performing tensor contractions in SEM. The focus is on the simulation of seismic wave propagation in three dimensions. We evaluate the two libraries cuBLAS and CUTLASS for performing the GEMM sequence in SEM and include batched GEMM library calls in the single-precision baseline SEM implementation. The experiments show that the accuracy obtained when using mixed-precision for the tensor contractions in SEM is sufficient when applied to seismic wave propagation in a homogeneous medium when using NVIDIA's TensorFloat32 precision, and when using half-precision with scaling to avoid underflow. In terms of performance, our results show that using cuBLAS and CUTLASS calls in host-side CUDA code has a large memory and CPU overhead and cannot compete with on the fly calculations without mixed-precision and Tensor Cores. However, we show that using mixed-precision computations without Tensor Cores for the tensor contractions in the baseline SEM solution and storing data in half-precision leads to a speedup of ~5%, which demonstrates that lower-precision should be considered both for storage and computations in SEM implementations. Directions for future work are also presented.