CUDAでカーネル関数がスキップされる(ように見えた)

1ヶ月くらい、土日や夜の自由時間を主に使ってGPUを使ったレイトレを書いていた。この話もまとめたいな〜と思いながらダラダラしていたら1週間が経過してしまった。みんなどこから時間をひねり出しているんだ? 私以外の人々はその人生で一切ダラダラしないに違いない。この調子で1ヶ月くらい経過して忘れる前に、その道中で一番理解に苦しんだバグを取り上げておきたい。

レポジトリ自体はここにある。そのバグは直っている。まだ機能は少ないが。

github.com

何をしようとしていたかというと、分子モデルを画面に表示しようとしていた。画面に表示するといえばカメラだ。最近のカメラはレンズの後ろにCCDやCMOSといった光センサが並んでおり、そこに当たっている光を検出して画像データを取得し、保存している。プログラムでもこれを真似ることになる。とはいえ光学系全てをシミュレーションするほどの正確性は必要ないので、必要のない部分はばっさり切り落とすことになる。この手の手法は本当に無限に研究されており深入りすると戻ってこれなくなるので、基本的にこの記事ではほぼ全てをばっさりカットすることをご了承いただきたい。

とにかく、GUIでリアルタイムに描画するとなるとレイトレ部分以外をまともに動かすだけで死ぬほど手間なので、レイトレーシングのコードを書く前にまずはモックアップ的なものを作ろうと思った。そのモックアップは以下のように動作する。それぞれのピクセルについて、そのピクセルに入射して来る光線を逆に辿って、最初にぶつかったオブジェクトを探す。そのオブジェクトから光が来たと考えられるので、そのピクセルはそのオブジェクトの色にする。リアルさは欠けらもないが、単純で間違えようがない。これは動いた。陰影はなくのっぺりしているし輪郭もギザギザだが、どこに何があるのかはわかる。

// コードはイメージです
color camera::render_pixel(std::size_t w, std::size_t h,
                           const world& wld)
{
    const auto ray = this->cast_ray(w, h);
    const auto obj = wld.first_hit(ray);
    return obj.albedo();
}

ではレイトレーシングはどのようにするかというと、最初にぶつかった物体を見つけたら、そこから来る光の強度を計算する。光の強度は、物体に差し込んで来る光の波長毎の強さに波長毎の物体の反射率をかけたものになる。物体自身が光っている場合はそこに物体から出る光の強度が加算される。なので、最初にぶつかった物体を見つけたら、その物体のその点に入射して来る光をサンプリングして、その強さをまた計算することになる。物体のない方向から来た光は、環境を満たしている光(空の色など)がそこに入って来ていたということになる。実際には反射するたびに全天を走査することは不可能なので、ランダムに方向を選んで毎回の反射毎に1本の光線を選び、色を計算することになる。そうなるとたまたま光源にヒットしなかったピクセルが真っ暗になったり、その逆が起きたりするので、一つのピクセル毎に100本も1000本も光線を辿って平均することになる。モンテカルロ積分だ。

// コードはイメージです
color camera::render_pixel(std::size_t w, std::size_t h,
                           const world& wld)
{
    auto ray = this->cast_ray(w, h);
    return this->trace_ray(ray, wld);
}

color camera::trace_ray(const ray_t& ray, const world& wld)
{
    const auto [obj, collision_point] = wld.first_hit(ray);
    const auto next_ray = obj.scatter_at(collision_point, ray);
    return obj.albedo() * trace_ray(next_ray, wld) +
           obj.emission();
}

この過程で、光が反射するたびに、その光がどこから来たかを辿って、物体から来ていたらまたその物体の前はどこから来たかを……という計算をする。ある物体からきた光の強度を計算するために、その前の物体からきた光の強度を使う。これは再帰だ。だが末尾再帰なので、人間でもループに置き換えることができた。GPUでは特に再帰よりもループの方が早いと思われるので、CUDAカーネルに単純なforループを書いた。すると画面が真っ暗になった。

よくわからなくなってループを消すと元ののっぺりした画像に戻る。完全に脳内を?マークが巡っている中、手元の本でCUDAカーネルでループが書けるか確認したり(当たり前すぎる)、forと等価なwhileループを書いてそれでも動かないことを確認したりした。いよいよもってわからず、カーネルからprintfを呼んでみた。何も表示されない。forを消して動くコードに戻すと、ターミナルがprintfからの出力で埋まる。ここから、コードを間違えたから画面が黒くなっているのではなく、CUDAカーネル自体が実行されていないことがわかる。

そんなことがあるのか? と思って普段は使わないcuda-gdbを取り出し、CUDAカーネルを呼び出す関数にブレークポイントを設置して、カーネルを呼んでから帰って来るまでをステップ実行してみた。すると、CUDAカーネルの呼び出し直後に何も起きずに戻って来てしまった。驚いたが、同時に「cudaLaunch returned (0x7)」という表示が出たので、呼び出しそのものが失敗していることがわかった。これは、「Launch out of resources」を意味するらしい。指定したCUDAカーネルを指定したブロック・スレッド数で実行するにはGPU上のなんらかの資源が足りないというわけだ。

そこで少しGPUアーキテクチャについて調べてみた。CUDAカーネルを起動する際はブロックの数とスレッドの数を指定することはよく知られているが、どのようにしてその数を決めればいいかの具体的な指針はあまり見つからない。基本的にはスレッドを遊ばせておくよりも多く起動した方が並列性が上がり、GPUがよく燃えることはわかる。CUDAを書いている際はGPUに負荷をかけて燃やしたいものなので、スレッドはタスクがある限り最大限起動した方がいいだろう(タスクがある限り、というのは、1000要素同士の足し算にスレッドを2000個立ち上げてもしょうがないからだ)。

// CUDAカーネル呼び出しのイメージ
const dim3 blocks(...);
const dim3 threads(...);
cuda_kernel<<<blocks, threads>>>(...);

ハード側の上限は割とリスト化されている。このGPUではブロック毎に立ち上げられるスレッド数はいくつで−−、のような話だ。あと、スレッドはワープ単位(32スレッド)で管理されるため、ワープの倍数で指定した方が良いということも知られている(50スレッドだけ立ち上げようとしても64スレッド立ち上がってしまうので、最初から64スレッド立ち上げると良い)。このあたりの情報からすると、常にスレッドをハード的な上限まで立ち上げ、データの数をそれで割った分のブロックを立ち上げれば良いと思ってしまう。実際、私はそうしていた。

ところがGPU上の資源はコアだけではなく、メモリやレジスタも含まれる。レジスタやシェアードメモリなどの一部のメモリはブロック毎に管理されており、ブロックがスポーンさせるスレッドの全てに均等に分割される(シェアードメモリは共有されるので分割されないが)。なので、スレッドを上限までスポーンさせた場合、意外と簡単にこの上限を超えてしまうようだ。最終的にブロックあたりのスレッド数を制限する(代わりにたくさんのブロックを作った)ことで解決したので、この部分でやらかしていたのは確実だと思う。なんとなく計算してみた限りレジスタが尽きたのかな、と思っていたが(スレッドを上限まで作ると1スレッドあたりint64個までしか置けなくなる)、レジスタの内容は退避させることもあるとかなんとか(Register Spilling)、よくわからなくなってきたので明言は避けておこう。よくわかりません。みなさん頑張ってください。

よくわからないままどうやって解決したんだ、ということだが、CUDAには便利なAPIがあって、デバイスのハード情報だけではなく、実行するカーネルが消費する資源の情報まで取得できる。これを使うことで回避が可能だ。以下の関数に呼びたいカーネルの関数ポインタを渡して呼び出すことで、カーネルの情報を得ることができる。

CUDA Runtime API

帰ってくる情報は以下のようなクラスだ。

class cudaFuncAttributes
{
public:
int     binaryVersion;
int     cacheModeCA;
size_t  constSizeBytes;
size_t  localSizeBytes;
int     maxDynamicSharedSizeBytes;
int     maxThreadsPerBlock;
int     numRegs;
int     preferredShmemCarveout;
int     ptxVersion;
size_t  sharedSizeBytes;
};

ここにmaxThreadsPerBlockというものがある。普通に考えるとデバイスのハード情報に思えるが、これは関数の属性だ。関数が消費する資源とデバイス情報から逆算して、ブロックあたりいくつのスレッドを立ち上げられるかの値が入っている。なので、何も考えずにこれを超えない最大の数のスレッドを立ち上げれば良い。ブロックの数はデータの数をスレッドの数で割って計算すると良いだろう。試しに見てみると、動かなかったカーネルでは確かにハード的な上限をわずかに下回る数が入っていた。なるほどね。

というわけで、スレッドの数をこの数までに制限すると動いた。まだまだ30FPS出るかどうかくらいの速度だが、綺麗な画像がインタラクティブに表示されるのは最高だ。みなさんもレイトレやりましょうレイトレ