Technology

中間データのメモリ往復を省く線形代数・FFTのカーネル内融合

この記事のポイント

  1. 実現したこと

    行列積、三角行列求解、FFTを、アプリケーション固有の処理と同じCUDAカーネル内で実行できる。

  2. 実現の仕組み

    共有メモリやレジスタ上で中間結果を引き継ぎ、数値計算ルーチンと前後の演算を融合する。

  3. 得られた結果

    融合した演算間では、中間データをグローバルメモリへ書き戻して読み直す往復を省ける。

  4. 従来との違い

    cuBLASDx 0.5.0以降では、非同期ロードと非同期計算を複数段重ねるパイプライン拡張が利用可能になった。

行列タイルとスペクトル波が単一の計算領域内で結合し、後続演算へ流れる概念図
AI生成画像

線形代数とFFTをアプリケーション固有の処理と同じCUDAカーネル内で実行し、中間データを共有メモリやレジスタから次の演算へ渡せる。グローバルメモリへの書き戻しと再読み込みを省く実行構成だ。

CUDAが備える数値計算の実行経路

NVIDIAのGPU開発環境であるCUDA Toolkitは、GPU対応アプリケーションの開発、最適化、配備に必要なライブラリ、デバッグ・最適化ツール、C/C++コンパイラ、ランタイムライブラリをまとめて提供する。組み込みシステムからワークステーション、データセンター、HPCスーパーコンピューターまでを対象とする開発基盤である。

数値計算向けには、CUDAランタイム上でBLASを実装する線形代数ライブラリのcuBLASと、NVIDIA GPU上で高速フーリエ変換を実行するcuFFTがある。Device Extensionsに当たるcuBLASDxとcuFFTDxは、こうした計算の一部をユーザーのCUDAカーネル内部へ置く別の実行経路を提供する。

GEMMとTRSMをカーネル内部へ組み込むcuBLASDx

線形代数向けライブラリのcuBLASDxは、行列積のGEMMと三角行列求解のTRSMをユーザーのCUDAカーネルへ埋め込める。線形代数ルーチンをアプリケーション固有の処理と同じカーネル内に配置できるため、計算結果を後続演算へ直接つなぐ構成を選べる。

埋め込むルーチンは、行列サイズ、精度、データ型、対象CUDAアーキテクチャなどに合わせて調整できる。cuBLASDxは指定された条件に応じてFMAまたはMMAの命令構成を振り分け、LDSMやSTSMを利用する読み込み命令も自動的に選択する。

共有メモリと非同期パイプラインで演算を接続

cuBLASDxは、線形代数の累積と別の演算との融合を共有メモリまたはレジスタ上で実行できる。中間結果をカーネル内に保持して次の演算へ渡すことで、グローバルメモリへ保存した後に別の処理が読み直す往復を省く。

バージョン0.5.0以降のパイプライン拡張は、非同期ロードと非同期計算をそれぞれ複数ステージ進行させ、両者を重ねて実行する。K方向の累積が完了するとエピローグ処理を1段実行し、データ読み込み、行列演算、後処理を段階的につなぐ。

FFTの出力を後続演算へ直接渡すcuFFTDx

FFT向けライブラリのcuFFTDxは、高速フーリエ変換の関数をユーザーのCUDAカーネルへ埋め込める。変換サイズ、精度、バッチ数などに応じてFFTルーチンを調整し、アプリケーション固有の処理と同じカーネル内で実行できる。

カーネル内で融合した後続演算は、FFTの出力をグローバルメモリへいったん戻さずに受け取れる。NVIDIAは、不要なデータ移動を省くこの構成によって、遅延を減らし、アプリケーション全体の性能を改善できるとしている。