cupyで2次元ガウス窓を生成する

8000メートルなどを超えるような超深海では、超高圧のためタンパク質の構造自体が保てなくなる(水分子が侵入し破壊される)ことから、通常の構造の延長としての生物は生息できないのだそうです。そこにいる生物は、TMAOという物質がもつタンパク質に侵入する水分子をブロックする性質を利用して生きているのだとか。したたかです。

さて、Pythonを使って2次元のガウス窓を生成したくなりました。
1次元でのscipy.signal.gaussianの多次元版にあたるものです。

具体的には、画像として描画するとこんなかんじになるような配列の生成が目的です。

Numpy芸

まずNumpy芸であればいくつか方法はありますが、こんなかんじでできます。
今回は、簡単のため2次元配列の中央が最大値1をとるようなガウス関数を考えます。

1次元のガウスの配列の直積をとって2次元にするという雰囲気です。sizeは作りたい配列のサイズ、sigmaはガウスの標準偏差。

※このNumpy芸の一部はomeometoさんのアドバイスによるものです!多謝!

Numpy芸の応用としてのCupy芸

自分はデータはGPU上に欲しかったので、Numpy互換のGPUテンソルライブラリであるCupyを使います。
この場合も基本的な考え方は同じでさくっとできます。
scipy.signal.gaussianにあたるものが無いため、これを手作業で作る必要がありますが大したことはありません。

ElementwiseKernelでのベタ実装

上2つはnumpy芸とそのcupyへの応用ですが、愚直にループを書いて各座標ごとにガウス関数の値を計算するアプローチの応用版として、CupyのElementwiseKernelの仕組みを使う方法もあります。

ElementwiseKernelはPythonで言うところのmapと同じで添字はカーネル本体に渡されないので、arrayとrangeをzipしてmapするのと同じ考え方で、座標を表す配列をmeshgridで生成しています。その分GPU上にメモリを無駄(+2倍)に食うことになります。

また、消費するFLOPの回数も、expやsqrtをsize*sizeの全要素に対して計算することからかなり多めになります。outerを使うアプローチだとこのような重たい計算は愚直な実装の平方根程度しか呼ばれずにすみます。が、それでもGPUの並列性とシンプルな処理により、実時間としてはそれなりに速いことが期待できます。

マイクロベンチマーク

このようなコードで、生成する2次元ガウス関数の配列のサイズを8×8から16384×16384まで変化させながら時間を測ってみました。横軸は配列の一辺のサイズ、縦軸は1回あたりの生成時間(ミリ秒)です。

1〜3項目目は上で説明したとおり、4項目目は(データはGPU上に欲しいので)numpyで生成したあとGPUに送る処理まで含んだ時間です。
環境はi7-7700K/GeForce GTX 1080Ti/Ubuntu16.04/CUDA8.0/Cupy1.0.3/Numpy1.11.3/Python3.6.0です。

GPUが絡むとどうしてもある一定以上速くはならないので、生成したい配列のサイズによって戦略を変える必要がありそう。

自分の場合はせいぜい数十というサイズが必要だったので、numpyで作ってGPUに送ってしまうほうが速い可能性が高いようです。
一方で、CupyのElementWise実装(gen_2d_gaussian_cupy_elementwise)におけるrng, x, yは、出力サイズが同じであれば毎回作る必要はなく完全に使いまわせるので、その方法でベンチマークをとるとコンスタントに(生成する配列のサイズによらず)3倍ほど速くなりました。これは常にnumpyで作ってGPUに送るアプローチよりも速いです。ガウス生成処理を複数回呼んで混合ガウスみたいなものを作りたい、などの場合は、このアプローチが最も効率的です。

同じことはcupy芸の実装(gen_2d_gaussian_cupy)でも言えて、arangeで生成しているxは同じサイズの出力に対して常に流用できるので速くなるのですが、arangeに比べてexpさらにはouterは遥かに重たいので、それほどは速くなりません(1割以下)。

柔軟性としては、上記の混合ガウス生成などに加えて、例えば斜めに分布するようなgaussianを描きたい場合はElementwiseKernelの方法だと何も考えずにシンプルに拡張すればよさそう。Numpy芸の方法だと、また少し頭をひねらないといけなそうです。

たぶんこの記事に書いた方法よりずっと速くてエレガントな方法がnumpy/cupyとも存在することだろうと思います。ぜひマサカリをください。

続報:CUDAカーネルを手で書いてみた

そもそも最大速度はどのぐらいになるんだろうと思って、Pythonを経由せず直接CUDAカーネルを書いたC++コードで同じことをしてみました。

これを合わせたベンチマークの結果がこちら。
生CUDA版を書く過程で、cudaFreeがアホほど遅いことが判明したので、これを時間計測の外に吐き出した版も比較しました。

Cupy経由などより遙かに速くなること、Pythonオーバヘッド分がないので当然なのですが、それにしても1桁近くの差となりました。
今回使ったGTX1080Tiは3584個のCUDAコアですので、16384×16384の場合は単純計算で1コアあたり75000画素を担当することになります。
このコアが1.5GHzで駆動し、またptxで見ると1画素の計算は概ね60ステップ程度だったので、メモリ待ちなどは無視して強引に1クロック1ステップで計算するとだいたい2.5msで全部終わることになる計算です。実際の数字が4ms程度ということなので、妥当な結果と言えそうです。
(生成結果は別途画像に描画して確認しています。)

もしこのサイズの出力のときもっと高速化しようとしたら、いくつか方法がありそうです。

まずカーネルの中身でやってる処理を減らす基本的な方法で、例えば中心μからの距離によって決まるガウス値を予め計算したLUTのようなものを作っておいて(この作業は辺のサイズnに対して高々O(n)ですむ)、カーネルでは距離計算をしてこのLUTをひくだけ…とすれば、LUTのメモリアクセスでキャッシュが効くなら、計算量的には数倍程度速くなる可能性があります。

あるいは、今回の問題設定では上下左右に対称であるため、カーネルを発行するのは左上にあたる領域だけにして、それらが対称になる他の領域にも書き込むようにする、とすれば、計算量自体は4倍を上限とする高速化が見込めます。汎用性はあまりないですが。LUTとの合わせ技で、さらに速くなり得ます。

一方で、16384×16384のfloat配列は1.07GBなので、メモリ帯域が高々500GB/s弱しかないGTX1080Tiでは、理論値が発揮できてもせいぜい1/500sつまり約2msが限界ということになります。
すでにこのカーネルでメモリ帯域の50%を引き出していたのですね。

GPGPUは奥が深く楽しいです。

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です