PFVM は、Preferred Networks が開発した高性能・マルチプラットフォーム対応の深層学習アクセラレータライブラリです。
ONNX 形式の深層学習モデル (計算グラフ) に対し独自のアルゴリズムに基づく最適化を適用することで、通常よりも省メモリかつ高速な推論を実現します。

このページでは、PFVM で採用されている最適化技術の一部をご紹介します。

自動カーネル融合#

自動カーネル融合による高速化#

PFVM では、複数の命令を自動的にひとつの命令にまとめることで、GPU カーネル呼び出し回数を減らす最適化を行っています。
例えば、以下のような計算列をまとめる場合について考えてみましょう。

fusion-sample

左側の図は、Gather, Mul, Sqrt の 3 つの命令を順に実行する計算列を示しています。
それぞれの命令は独立した GPU カーネルを用いて計算しており、命令のたびに計算結果を GPU のグローバルメモリ領域に書き込んでいます。

右側の図は PFVM が左側の計算列をひとつの命令にまとめて得られる計算列の図です。
このとき、Gather, Mul, Sqrt の 3 つの命令において、計算途中の変数 tmp1, tmp2 の読み書きを同一スレッドで行うような CUDA カーネルを実装することができます。
したがって、計算の実行中にスレッド間で同期を行う必要がないため、 tmp1tmp2 をグローバルメモリ領域に読み書きする必要がなくなり、カーネル実行にかかる実行時間を短縮することができます。
以上のような、複数の CUDA カーネルを用いて実行していた命令をひとつの CUDA カーネルにまとめる高速化を 「カーネル融合」 と呼びます。[1]

PFVM では、中間変数をスレッド間で同期する必要がないような命令の組み合わせを適切に探索することで、高性能なプログラムへと自動的に変換されます。
最適な命令の組み合わせの探索が完了した後、PFVM はカーネル融合後の CUDA プログラムを文字列で生成してコンパイルしています。
以下は、PFVM が実際に 6 つの命令をカーネル融合して生成した CUDA プログラムに対して、読みやすく若干の修正を加えたものです。

extern "C" __global__ void kernel(size_t n, float* x1, float* x2, float* x3, float* x4, float* x5, double* y1) {
    size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    const float v1 = x1[tid];  // input
    const float v2 = x2[tid];  // input
    const float v3 = x3[tid];  // input
    const float v4 = x4[tid];  // input
    const float v5 = x5[tid];  // input
    const float v6 = v4 * v3;  // Mul
    const float v7 = v6 + v6;  // Add
    const float v8 = v5 + v7;  // Add
    const float v9 = v8 + v1;  // Add
    const float v10 = v9 + v2;  // Add
    const double v11 = static_cast<double>(v10);  // Cast
    y1[tid] = v11;  // output
}

カーネル融合がメモリ消費量に与える影響#

カーネル融合が可能な条件を満たすような全ての命令に対して闇雲にカーネル融合を処理すると、消費メモリ量が極端に増えてしまう場合があります。
PFVM では、独自で開発した高速なメモリシミュレータを用いることで、カーネル融合を行っても消費メモリ量が増えないような命令の組み合わせを適切に探索しています。
メモリシミュレーターが高速に消費メモリ量のシミュレーションを行うメカニズムについては特許出願中であるためここでは説明できませんが、従来のシミュレーターと比べて数百倍程度の命令の組み合わせについて探索を行うことが可能です。

Reference#

[1] https://arxiv.org/abs/1305.1183