コンパイラGPUパフォーマンス

FutharkをGPUから遠ざける

Futhark は通常、配列プログラムを GPU カーネルにコンパイルするが、一部のワークロードは多くの CPU コアで回した方が速い。本記事はスパースなヤコビ行列の計算を例に、ターゲットが GPU ではなく CPU になったときコンパイラのどこを変える必要があるかを追う。

FutharkTroels Henriksen約 9 分

日本語
中文日本語
コピー

今年初め、Elias Smedegaard がBSc 論文を仕上げた。自動微分を使って疎な Jacobian 行列を効率よく計算する方法の研究だ。この記事のために、それが何で何の役に立つのかを正確に理解する必要はない——自動微分の話であり、これについてはまだブログを書いていない。本記事で重要なのは、彼の手法の一部に、行列のスパースパターンに対応するグラフの彩色が含まれているという点だ。

最適なグラフ彩色は NP-hard として知られているが、必要なのは最適な彩色ではない——十分に良い彩色であり、しかも効率的なアルゴリズムで実装できるものであればよい。Elias は distance-2 彩色のためのかなり高速な貪欲アルゴリズムを見つけた（興味があれば、この論文の Algorithm 3.1 だ）が、残念ながらこのアルゴリズムは本質的に逐次的だ。Futhark ではこれは問題にならない。Futhark はかなり効率的な in-place 更新をサポートしているからだ。問題は、Elias のプログラム全体が 2 つのステップに分かれていることにある：

  1. グラフを彩色する。
  2. その彩色結果を使って、Jacobian を大規模並列に計算する。

私たちはステップ 2 に並列ハードウェア（GPU など）を使い、ステップ 1 は CPU で実行したいと強く望んでいる。ところが、Futhark のいくつかの実装上の選択により、これが非常に厄介なものになっている。この記事では、この問題がどのように生じたのか、原則に基づいた解決策はどのような形になりうるのか、そしてなぜ結局その場しのぎのパッチで済ませたのかを説明する。

問題はどのように生じたか

Futhark のコンパイルモデルは単純だ。配列はメモリ上に置く必要があり、ある GPU パイプラインを使う場合、すべての配列が GPU メモリに置かれる。これは内部表現が CPU と GPU のメモリの混在を表現できないからではなく、純粋に実装上の便宜だ。安全な設計でもある。CPU コードは GPU メモリにアクセスできる（コストの高い API 経由で）が、GPU コードは一般に CPU メモリにアクセスできない。動くからといって速いとは限らない。次の仮想的な Futhark ループを考えてみよう：

loop sum = 0 for i < n do
  sum + xs[i]

xs が GPU メモリに格納された配列だとすると、ループ本体には次のようなコードが含まれる：1 つの要素を GPU メモリから CPU メモリのどこかへ苦労してコピーし、それから sum に加算する（後者はおそらくレジスタにある）。これは遅い！配列要素 1 つをコピーするの自体はほとんど時間がかからないが、通信を確立してコピー完了を待つオーバーヘッドは極めて高く、配列要素 1 つあたりに均せば軽く数マイクロ秒になる。

より良い方法は、ループの前に xs 配列をまとめてコピーし、通信オーバーヘッドを薄めることだ。実際にこれがより良い方法であり、後で触れる。しかし GPU 配列の添字が CPU コードに現れる限り、極端に遅い GPU メモリ読み出しを生成してしまい、ループの中にあれば破滅的だ。

Elias が実装した distance-2 彩色アルゴリズムは、本質的に何層か入れ子になった逐次ループで、グラフとスタックを表す配列を操作する。このコードを GPU バックエンドでコンパイルすると、逐次的な c バックエンドと比べて性能が壊滅的で、軽く 3 桁以上悪化する。c バックエンドに切り替えれば性能は非常に良い（C で書いたのとほぼ同等）が、そうするとプログラムの後半部分も逐次実行になってしまう。折衷案は Futhark の multicore バックエンドで、並列な CPU コードを生成する——だが Jacobian は GPU で計算したいのだ！

原則に基づいた解決策はどのような形になりうるか

Futhark は「分割コンパイル」をサポートしておらず、プログラムの異なる部分で異なるバックエンドを使うこともできない。そもそもかなり不格好な解決策に思える。本当に最良の方法は上でほのめかしたものだ：配列が実際にどのように（というよりどこで）使われるかを調べ、それに基づいて CPU メモリに置くか GPU メモリに置くかを決める——プログラムが必要とするなら、両方に置いてもよい。

実は Philip Børgesen が実装した最適化がすでに似たことをしている：ある単純な逐次計算（データではない！）の結果が GPU 上でのみ使われる場合、その計算を GPU に移す。狙いは通信オーバーヘッドの削減だ——計算自体が取るに足らないもので、時間が配列要素をいくつか往復コピーするのに費やされるなら、その逐次計算をシングルスレッドの GPU kernel に入れて走らせたほうが、通信を減らせる分だけましだ。残念ながらこの最適化はここでは使えない。第一に、distance-2 彩色アルゴリズムは（ループを伴うため）複雑すぎてこの最適化を発火させられない。仮にチェックを無効にして無理やり GPU で走らせても、（驚くことではないが）このようなループの多い逐次コードを単一の GPU スレッドで走らせるのは途方もなく遅いことがわかる。

理想的には、この最適化に対応するものを思い描ける：計算ではなくデータを移動するもので、根底にある原理の多くはおそらく同じだ。特に注意すべきは、過剰なコピーを避けつつ、同時にそれらのコピーが無駄に長く存在し続けないようにすることだ。長く存在すればメモリ使用量を押し上げる。この最適化がおおよそどんな形になるかは想像できるが、今の自分が割ける以上の労力と注意を注ぐ価値があることも分かっている。

というわけで……

実際に採用したその場しのぎの方法

コンパイラに賢い仕組みを仕込む代わりに、#[cpu_function] という新しい属性を追加した。当初の設計では、これを非インライン関数（通常は #[noinline] で強制する）に付けると、#[cpu_function] は関数本体を CPU コードとしてコンパイルし、すべての入力・中間・出力配列を CPU メモリに置く。どのコンパイラバックエンドを使っていてもそうなる。配列引数が正しいメモリ空間にあることを保証するのは関数の呼び出し側の責任で、コンパイラがそれを確実にするコードを挿入する。使用例は次のとおり：

#[noinline] #[cpu_function]
def seqsum [n] (xs: [n]i32) =
  loop sum = 0 for i < n do
    sum + xs[i]

こうしたのは、簡単にできるはずだと思ったからだ。実際ほぼそのとおりだった。メモリを表すための中間表現は、以前から（意図的にそう設計されている）複数のメモリを同時に扱えるようになっていた。具体的には、mem 型にメモリ空間のパラメータがあり、GPU 共有メモリの表現にも使っている。問題は、この機能を実際に大規模に使ったことがこれまで一度もなかったことで、そのせいで片づけるべき問題がいくつか表面化した。同じ精神で、選んだ解決策はどれもその場で一番楽なものにしてある。正しいことをやるのは、時間ができてからだ。

ぶつかった問題のひとつは、GPU kernel のコードを生成するとき、関数本体で参照される配列はすべて GPU メモリ上にあると決め打ちしていたことだ。今はそうとは限らない。たとえば次のようなプログラムを考える。

#[noinline] #[cpu_function]
def frob (n: i64) =
  iota n

entry main n = let xs = frob n
               in map (\x -> x + 1) xs

main の map は GPU kernel にコンパイルされるが、配列 xs は #[cpu_function] のせいで CPU メモリ上にある！ これは不正だ。正しいやり方は、GPU kernel の中で実際にどの配列がアクセスされるかを調べ、必要なら GPU にコピーすることだが、これを正しくやるのは実は簡単ではない。

  1. GPU kernel がループの中で呼ばれるなら、ループの各反復でコピーすべきではない。
  2. ……が同時に、配列のコピーの寿命はできるだけ短くあるべきだ。

そこで #[cpu_function] の意味を少し変えることにした。入力と中間結果はすべて CPU メモリに置き、結果は GPU メモリに置いて、関数の最後に一度コピーして終わりにする。こうすればコンパイラは、#[cpu_function] が付いた関数を除いて、すべての配列が GPU メモリにあると引き続きみなせる。実際、それ以外の場合はそうなっている。

……もちろん、これらの関数は後で厄介者になった。GPU バックエンドは、iota と replicate の演算（どちらも IR のプリミティブだ）の結果が GPU メモリにあるとみなすが、#[cpu_function] の内部ではそうならない。これを直すのはそれほど難しくなかったが、似たバグはまだどこかに潜んでいそうだ。

最後の問題は、#[cpu_function] を map したとき何が起きるべきか。その場の判断で、このときは #[cpu_function] を無視して、普通の並列化された関数を生成することにした。ここで何らかの形で逐次関数を生成すべきだと主張する人もいるだろう。

現状

#[cpu_function] はかなり厄介なハックだが、目の前の問題は解決した。これで、グラフ彩色を非常に効率のよい逐次コードで行いつつ、同じ Futhark プログラムの中で並列の処理もできる。実際、グラフ彩色が全体の実行時間に占める割合は今ではごくわずかで、大半はその後の数値計算だ。#[cpu_function] の手法は粗粒度で、これらの関数を呼ぶと高価なコピーが発生する。ただし要素ごとではなくまとめてコピーするので、オーバーヘッドはそれほど大きくない。

コンパイラに必要な変更は多くなく、そのほとんどはバグ修正の類だ。Futhark コンパイラが配列の置き場所を自動で判断できるようになった理想の未来でも、この属性は有用なヒントとして残すだろう。

DIKU による開発（著作権、謝辞、ハリネズミ）

出典: Futhark← ホームへ戻る