[論文レビュー] tcFFT: Accelerating Half-Precision FFT through Tensor Cores

Binrui Li, Shenggan Cheng+1 more|arXiv (Cornell University)|Apr 23, 2021

Parallel Computing and Optimization Techniques参考文献 28被引用数 6

ひとこと要約

tcFFT は、Tensor Core フラグメントにおける単一要素操作を可能にし、メモリアクセスパターンを最適化することで、NVIDIA Tensor Cores を活用して半精度 1D および 2D FFT の処理を高速化する高性能 FFT ライブラリである。さまざまな FFT サイズと次元において、V100 および A100 GPU で cuFFT に対して 1.29x ～ 3.24x の高速化を達成している。

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.

研究の動機と目的

• NVIDIA Tensor Cores を活用して、低精度 GEMM 専用に最適化されているが、FFT の特殊な演算には natively 向いていない現代の GPU で半精度 FFT を高速化すること。
• 大規模または多次元 FFT における不規則なメモリアクセスパターンが引き起こす性能ボトルネックを克服すること。
• 幅広いサイズ範囲でバッチ処理可能な 1D および 2D 変換をサポートする、ポータブルで高性能な FFT ライブラリの設計。
• FFT タスクにおける Tensor Cores の効率的利用を可能にし、cuFFT のような既存のライブラリが十分にサポートしていない分野に対応すること。
• Tensor Cores の高速化が、従来の密行列線形代数を凌駕する科学的計算ワークロードにも効果的に適用可能であることを示すこと。

提案手法

• FFT 特有の演算（複素行列アクセス、要素単位の乗算など）をサポートするため、Tensor Core フラグメントにおける単一要素操作技術を提案した。
• 可変な FFT サイズに適応可能な細粒度のデータ配置戦略を設計し、連続的かつコalesced（凝集）されたメモリアクセスパターンを実現した。
• 大規模または 2D FFT のマージ段階における、非凝集的なストライドアクセスを最小限に抑えるために、グローバルメモリ上のデータレイアウトを再編成した。
• 共有メモリを活用してグローバルメモリ帯域幅の圧力を軽減し、演算強度を向上させた。
• FFT の計算ステージを Tensor Core 操作にマッピングする最小限のデータ移動で済む、カーネル融合戦略を実装した。
• 独自のメモリレイアウトとスレッドマッピングを用いた Tensor Core WMMA インターフェースを活用し、スレッドの占有率と利用効率を最大化した。

実験結果

リサーチクエスチョン

• RQ1 固定された GEMM 向け設計を持つ Tensor Cores を、その設計に反するが半精度 FFT の高速化に効果的に活用できるか？
• RQ2 大規模または多次元 FFT における不規則なメモリアクセスパターンは、どのように緩和可能か？ これにより GPU メモリボトルネックを回避できるか？
• RQ3 複素数演算や要素単位のスケーリングといった FFT の特殊な演算を Tensor Cores で実行するには、どのような最適化が必要か？
• RQ4 1 つの FFT ライブラリが、Tensor Cores を用いて幅広い 1D および 2D FFT サイズで一貫した性能向上を達成できるか？
• RQ5 実際の科学的計算ワークロードにおいて、Tensor Cores 最適化された FFT ライブラリの性能は、広く使われている cuFFT と比べてどの程度優れているか？

主な発見

• tcFFT は V100 GPU における 1D FFT で、cuFFT に対して平均 1.90x の高速化を達成した。
• 2D FFT では、V100 で 1.29x ～ 3.24x、A100 で 1.10x ～ 3.03x の高速化を達成した。
• tcFFT の性能優位性は、FFT サイズが大きく、バッチ数が多い場合に顕著であり、この場合 cuFFT のメモリアクセスパターンが著しく劣化する。
• tcFFT は小さなバッチサイズでも高い性能を維持し、1D FFT ではバッチサイズが 4 を超えると、2D FFT ではバッチサイズが 2 を超えると、cuFFT を上回った。
• 提案された単一要素フラグメント操作技術により、Tensor Cores 上での FFT 特有の演算が効率的に実行可能となり、ネイティブハードウェアサポートの欠如という主要な制限要因を克服した。
• 最適化されたメモリレイアウトにより、グローバルメモリへの連続的かつ凝集されたアクセスが確保され、大規模 FFT におけるメモリ帯域幅ボトルネックが顕著に低減された。