[논문 리뷰] tcFFT: Accelerating Half-Precision FFT through Tensor Cores

Binrui Li, Shenggan Cheng외 1명|arXiv (Cornell University)|2021. 04. 23.

Parallel Computing and Optimization Techniques참고 문헌 28인용 수 6

한 줄 요약

tcFFT는 NVIDIA Tensor Cores를 활용하여 반정밀도 1차원 및 2차원 FFT를 가속화하는 고성능 FFT 라이브러리로, Tensor Core 조각에 대한 단일 요소 조작을 가능하게 하고 메모리 액세스 패턴을 최적화함으로써 성능을 향상시킨다. 다양한 FFT 크기와 차원에서 V100 및 A100 GPU에서 cuFFT 대비 1.29배에서 3.24배의 성능 향상을 달성한다.

ABSTRACT

Fast Fourier Transform (FFT) is an essential tool in scientific and engineering computation. The increasing demand for mixed-precision FFT has made it possible to utilize half-precision floating-point (FP16) arithmetic for faster speed and energy saving. Specializing in lower precision, NVIDIA Tensor Cores can deliver extremely high computation performance. However, the fixed computation pattern makes it hard to utilize the computing power of Tensor Cores in FFT. Therefore, we developed tcFFT to accelerate FFT with Tensor Cores. Our tcFFT supports batched 1D and 2D FFT of various sizes and it exploits a set of optimizations to achieve high performance: 1) single-element manipulation on Tensor Core fragments to support special operations needed by FFT; 2) fine-grained data arrangement design to coordinate with the GPU memory access pattern. We evaluated our tcFFT and the NVIDIA cuFFT in various sizes and dimensions on NVIDIA V100 and A100 GPUs. The results show that our tcFFT can outperform cuFFT 1.29x-3.24x and 1.10x-3.03x on the two GPUs, respectively. Our tcFFT has a great potential for mixed-precision scientific applications.

연구 동기 및 목표

• 저밀도 GEMM에 최적화된 NVIDIA Tensor Cores를 활용하여 현대 GPU에서 반정밀도 FFT를 가속화하고자 하며, 이는 FFT의 특수 연산에 대해 직접 지원되지 않는 점을 보완한다.
• 크거나 다차원 FFT에서 발생하는 비정규적인 메모리 액세스 패턴으로 인한 성능 저하 문제를 해결하고자 한다.
• 다양한 크기의 1D 및 2D 변환을 지원하는 이식성 있고 고성능의 FFT 라이브러리를 설계하고자 한다.
• FFT 워크로드에 대해 Tensor Cores를 효율적으로 활용할 수 있도록 기존 라이브러리인 cuFFT가 잘 지원하지 않는 분야에 대비하고자 한다.
• 기존에 밀도 높은 선형 대수 계산에만 국한되어 있던 Tensor Core 가속화가 과학 계산 워크로드에 효과적으로 적용될 수 있음을 입증하고자 한다.

제안 방법

• FFT 전용 연산(예: 복소수 행렬 액세스, 요소별 곱셈 등)을 지원하기 위해 Tensor Core 조각에 대한 단일 요소 조작 기법을 제안하였다.
• 다양한 FFT 크기에 대응하는 연속적이고 코ales스된 메모리 액세스 패턴을 가능하게 하는 세밀한 데이터 배치 전략을 설계하였다.
• 크거나 2차원 FFT에서 병합 단계 동안의 비연속적인 스트라이드 액세스를 최소화하기 위해 글로벌 메모리 내 데이터 레이아웃을 재구성하였다.
• 공유 메모리를 활용하여 글로벌 메모리 대역폭 부담을 줄이고 연산 밀도를 향상시켰다.
• FFT 계산 단계를 Tensor Core 연산에 매핑하는 커널 융합 전략을 구현하여 데이터 이동을 최소화하였다.
• 고유한 메모리 레이아웃과 스레드 매핑을 사용한 Tensor Core WMMA 인터페이스를 활용하여 최적의 사용도와 할당률을 달성하였다.

실험 결과

연구 질문

• RQ1 고정된 GEMM 중심 설계를 가진 Tensor Cores가 반정밀도 FFT를 효과적으로 가속화하는 데에 활용될 수 있는가?
• RQ2 크거나 다차원 FFT에서 발생하는 비정규적인 메모리 액세스 패턴은 어떻게 완화할 수 있으며, GPU 메모리 병목 현상을 방지할 수 있는가?
• RQ3 복소수 산술 및 요소별 스케일링과 같은 FFT의 특수 연산을 Tensor Cores에서 실행하기 위해 필요한 최적화는 무엇인가?
• RQ4 단일 FFT 라이브러리가 다양한 크기의 1D 및 2D FFT에서 Tensor Cores를 활용해 일관된 성능 향상을 달성할 수 있는가?
• RQ5 실제 과학 계산 워크로드에서 널리 사용되는 cuFFT에 비해 Tensor Core 최적화된 FFT 라이브러리의 성능은 어떠한가?

주요 결과

• tcFFT는 V100 GPU에서 1D FFT에서 cuFFT 대비 평균 1.90배의 성능 향상을 달성하였다.
• 2D FFT에서는 V100 GPU에서 1.29배에서 3.24배, A100 GPU에서 1.10배에서 3.03배의 성능 향상을 기록하였다.
• tcFFT의 성능 우월성은 특히 큰 FFT 크기와 높은 배치 수에서 두드러지며, 이 경우 cuFFT의 메모리 액세스 패턴이 심각하게 저하되기 때문이다.
• tcFFT는 작은 배치 수에서도 높은 성능을 유지하며, 1D FFT에서는 배치 수가 4를 초월할 때, 2D FFT에서는 2를 초월할 때 cuFFT를 능가한다.
• 제안된 단일 요소 조각 조작 기법은 Tensor Cores에서 FFT 전용 연산을 효율적으로 실행할 수 있도록 하여, 하드웨어의 기본 지원에 대한 주요 제약을 극복하였다.
• 최적화된 메모리 레이아웃은 연속적이고 코ales스된 글로벌 메모리 액세스를 보장하여 대규모 FFT에서 메모리 대역폭 병목 현상을 크게 감소시켰다.