parallel linear BVH (LBVH) をCUDA+thrustで書いた

久々に少し非自明なことをCUDA+thrustでやった。大体動くようになるまでまる一日（半日+半日なので別のmetricsでは2日）かかった。

BVHとはなんぞやという話は、以下の記事にとてもよくまとまっている。というか、多分この記事の著者は私のような空間分割エンジョイ勢ではなくその道のプロと思われる。

shinjiogaki.github.io

なのでちゃんとした話は上の記事を見てもらって、ここでは細かいことをすっ飛ばしてイメージだけ喋ろう。オブジェクトのそれぞれに、それをすっぽり覆うような長方形を被せることを考える。これは各軸にアラインした長方形、直方体、またはN次元矩形とする。これをAxisAlignedBoundingBox、AABBという。計算が簡単なのでよく使われる。その後空間的に近いオブジェクト同士をまとめていって木を作る。ここで、木のノードにも子ノード全てを覆うような長方形が被せられている。葉ノードだった場合は持っているオブジェクトを全て覆うような長方形が被せられている。

ここで、あるオブジェクトが他の何かに衝突していないか調べたいとする。もしノードを覆う長方形がそのオブジェクトと衝突していないなら、ノードの中にある要素もまた今調べているオブジェクトとは衝突しない。ノードを覆う長方形の中にあるからだ。なので、このような場合にはノードを丸ごと無視できる。これにより、木の品質にもよるが、衝突判定が高速にできるようになる。

さて、で、今回やったことの話に入ろう。LBVHという種類のBVHがある。この手法は、ツリーを構築する際、Z-order curveなどの空間充填曲線を使ってオブジェクトにインデックスを割り振る。空間充填曲線なので、インデックスの値が近いもの同士は近くに並んでいるだろう（近くにあってもインデックスが大きく違ってしまうということはあり得る。まあある程度は仕方がない）。これを使えば、「近いものをまとめる」という操作が簡単になる。インデックスが近いもの同士をくっつけて行けば良いのだ。

このLBVHで、ノードの構築を完全に並列化するという手法が2012年に発表されている。nvidiaの中の人が考えたらしく、nvidiaのdevblogでも記事が出ている。構築方法については、先に紹介したbvhの記事でも詳しく紹介されている。ちなみに日本語の方の記事を読もうと思っている・読んだ人のために一応補足しておくと、δ(i, j) は i か j の少なくとも一つが範囲外になるとき δ(i, j) = -1 になる。

devblogs.nvidia.com

大まかな構築方法、探索方法に関しては、nvidia の devblog と bvh の記事で十分説明があると思うので、ここでは流し見しただけではわかりにくかった実装の詳細をメモしておこうと思う。

morton codeが衝突したとき

このアルゴリズムでは、どのオブジェクトが同じノードにアサインされるかがmorton codeで決まる。これは空間充填曲線の一種であるZ-order curveによって付けられたインデックスで、どういう順序になるかは先に紹介したnvidiaの記事か以下のWikipedia記事が参考になると思う。

ja.wikipedia.org

これによって、明示的に距離を計算したりグリッドを切ることなく、近くにあるものをまとめる操作ができる。しかもこの曲線はフラクタルなので好きなだけ細かくできる、つまり空間をどれだけ細分化するかは比較的簡単に選ぶことができる。とはいえ有限の情報量しか格納できないマシンを使っている以上無限に細かく分割することはできない。今回実装したアルゴリズムでは32bit整数のビットをxyz方向に均等に割り振って、10bitずつ、つまり1024分の1まで割るようになっている。

これは稀に衝突する。空間をXYZのそれぞれの方向で1024分割したグリッドがあると思ってほしい。ほとんど衝突しないだろうが、オブジェクトが十分密な領域があると、一つのグリッドに複数のオブジェクトが入る可能性はゼロではない。で、論文のSection4で触れられているが、このアルゴリズムはmorton codeがユニークであることを前提としているので、morton codeが衝突した場合はそれ用の処理が必要になるそうだ。

その場合にどうしたらよいかだが、nvidiaの記事のコメント欄では「そうならないようにオブジェクトを取り除け」とか「morton codeにインデックスをアペンドしろ」とかが飛び交っていたが（論文ではmorton codeとインデックスのビット表現を結合したもの（多分全部で64bitになっている）を使っていた。今読み直して気づいた）、実際に実装してコードも公開している（CUDAではなくOpenMPのようだが）人は determine_range と find_split 関数で隣とmorton codeが等しいかどうかを見て、それらについては特別な処理をしている。同じmorton codeを持つものは共通のparent nodeを持つようにし、そのparent nodeはleafまで階段状に割られているようだ。二分していくよりも深くなってしまいそうだが、衝突がそもそも起きにくいので実用上問題にならないだろう。

では私はどうしたかというと、先にmorton codeでオブジェクトをソートしたあと、同じmorton codeについてreduceし、各ノードがオブジェクトプール内のレンジを持つようにした。これは分子動力学シミュレーションや粒子法による流体シミュレーションなどでセルリストから近傍リストを作るときにも使われる手法なので、私にとっては馴染みのあるものだ。以下のように衝突しまくりのシーンがあるとしよう。

object list | 9| 5| 8| 2| 3| 6| 1| 7| 4| 0| (sorted by morton code)
morton code | 1| 2| 2| 3| 3| 3| 4| 4| 5| 6|

この morton code をkeyに使って、 thrust::constant_iterator<unsigned int>(1) を reduce_by_key する。 reduce_by_key は key が同じであるような要素についてreduceする（典型的には総和を取る）。要素として constant_iterator(1) を渡しておけば、実質的にcountとuniqueが同時に達成できる。

morton code | 1| 2| 2| 3| 3| 3| 4| 4| 5| 6|
constant    | 1| 1| 1| 1| 1| 1| 1| 1| 1| 1|
| reduce_by_key
v
reduced key | 1| 2| 3| 4| 5| 6|
reduced val | 1| 2| 3| 2| 1| 1|

続いてこれの inclusive_scan を取る。これは、その要素までの和を計算するものだ。

reduced val | 1| 2| 3| 2| 1| 1|
| inclusive_scan
v
scanned val | 1| 3| 6| 8| 9|10|

このスキャン済みの配列を後ろに一つずらし、先頭に0を追加すると、以下のようになる。実際には後で pop_front すると大変なので、最初に一つ分大きいリストを作っておき、書き込み先を begin() の一つ先にしておくとよい。

object list | 9| 5| 8| 2| 3| 6| 1| 7| 4| 0|
morton code | 1| 2| 2| 3| 3| 3| 4| 4| 5| 6|
              ^  ^     ^        ^
              |  |  +--+--------+
              |  |  |  | ...
scanned val | 0| 1| 3| 6| 8| 9|10|

この scanned val の隣り合う要素が、レンジになっていることにお気づきだろうか。最初のノードは、 object list の[0, 1) の範囲に対応しているし、その次のノードは [1, 3) に、その次は [3, 6) に対応している。これで、葉ノードがオブジェクトリストのどの範囲に対応しているかを高速にルックアップできるテーブルが作れた。これにはもう一つついでの利点があって、葉ノードは常に自分のインデックス+1を指すようにしておくことができる（ scanned val は先頭に足された 0 の分だけ、葉ノードの数よりひとつだけ大きくなる）。よって 0 を指す葉ノードはいない。なので、このインデックスをそのまま内部ノードかどうかの判定に使うことができる。 range_idx が 0 なら内部ノード、そうでなければ object list の [range[range_idx-1], range[range_idx])の範囲を見れば対応するオブジェクトのインデックスがわかる。

というわけで、各葉ノードは厳密に一つの morton code に対応し、衝突があった場合は単にそれらすべてのオブジェクトが入ったレンジが葉ノードに対応することになる。これによってmorton code衝突問題は解決される。

悪い点としては、常に葉ノードはレンジに対応するので、めったに起きないと思われるmorton code衝突のためにすべての葉ノードが長さ1のレンジを持っており、ルックアップが一回余分に発生してしまうということだ。なので、衝突を心配しなくていいという安定を取った結果少し遅くなっていることになる。これがどう効いてくるかはアプリケーションに依るだろう。キーが64bitになるのを気にしていたが、これに比べると単純にインデックスを足したほうが速かったかもしれない。たいていの場合は衝突しないのだし、ソートが終わったタイミングでuniqueして被っていなかったらそのまま、被っていたらmorton codeをindexで拡張する、というのでよかった気もしてきた。あとで変えようかな？

追記(7/31 1:13): 衝突を検出して、もしあればmorton code(32bit)にオブジェクトのインデックス(32bit)を結合して64bitに拡張し、それを使うように変えた。

Note: 分子動力学などでセルリストを使って同様のことをする場合、空のセルに気をつけないといけない。空のセルのCellIDは reduce 操作によって姿を消してしまうので、空のセルがあるとそこからあとのセルのインデックスとレンジのインデックスがずれてしまう。そのため、私が以前練習で書いたCUDA MDプログラムは、スキャッタを一度挟んで空のセルに入っている粒子数を明示的に0にしている。粒子密度が極端に薄い領域が出てくるようなシミュレーションは標準的ではないせいか、詳細すぎるので大抵の解説ではあまり触れられていない落とし穴なので注意。今回は、空のグリッドは単にどのノードからも参照されない（なので対応するノードがそもそも作られない）ので問題にはならない。

各ノードに対応するAABBの構築

さて、分割位置や子ノードが決まっても、AABBはすぐには決まらない。葉ノードのAABBはオブジェクトしか持っていないので最初からAABBが決まるが、葉ノードより上のノードのAABBは両方の子ノードのAABBが決まらないと決まらない。なので、完全に並列にしようとすると、すべてのノードで自身の子ノードが構築を終えるまで待つか、各ノードがそれぞれ葉ノードにたどり着くまで探索して自力ですべてのオブジェクトをマージしていくかのどちらかしかない。

後者は苦痛だ。明らかに、根ノードは全体を覆い尽くしているので、根ノードはすべてのオブジェクトのAABBをマージしなければならない。これは時間がかかるだろう。そしてそれが律速になることは明らかだ。なんとなく、少し待つとしても2つのAABBをマージするだけでいいほうが短時間で済みそうな気がする。

というわけでnvidiaのブログで（自然言語で）説明されていたやり方を実装した。まず、葉ではないノードのそれぞれに対応するアトミックに更新可能なフラグを用意しておく。その後、葉ノードすべてに対応するGPUスレッドを立ち上げる。葉ノードなので、オブジェクトしかもっていない。なのでAABBが即決まる。続いて、自身の親ノードを見に行く。BVHでは親ノードは2つの子ノードを持っているので、後に来た方がAABBを計算すればいい。両方の子ノードのAABBが必要なので、先に来た方はもう片方のAABBを知ることができず、処理ができない。なので先に来た方は片側のAABBが計算済みだというフラグを立てて休憩し、後で来た方が両方をマージしてそのノードの代表としてそのまた親に向かえばよい。

この計算はアトミックなCompare And Swapを使うことで達成できる。この命令は、あるメモリ領域に値を書き込むのだが、書き込む前に今の値を読む。そして期待している値と比べる。値が期待したもののままだったなら、そのまま書き込む。そうでなければそのままにしておく。CUDAでは、呼び出した側には読み込んだ値が帰ってくる（失敗したかどうかを返す環境もある）。その値が期待していたものと同じだったら書き込めたはずだし、そうでなければ今の値は帰ってきた値と同じになっている。

今回は、フラグをすべて例えば 0 で埋めておいて、 0 を期待して 1 に変更するような atomicCAS を発行すればよい。もし 0 が帰ってきたら、フラグは期待通り 0 だったことになり、それまで誰もフラグを変えていなかったことになる。安心して次にやってくるスレッドに処理を任せていい。もし 1 が帰ってきたら、既に誰かがフラグを立てていたことがわかる。この場合、そのスレッドが仕事をする必要がある。

なので以下のような形になるはずだ。

unsigned int parent = self.nodes[idx].parent_idx;
while(parent != 0xFFFFFFFF) // means idx == 0
{
    const int old = atomicCAS(/* flag address = */ flags + parent,
                              /* compare = */ 0, /* value = */ 1);
    if(old == 0)
    {
        // もう片側の子ノードの処理がまだ終わっていない
        return;
    }
    assert(old == 1);
    // 両側の子ノードの処理が終わった。AABBを計算
    //（既に反対側の子ノードがフラグを立てている）

    const auto lidx = self.nodes[parent].left_idx;
    const auto ridx = self.nodes[parent].right_idx;
    const auto lbox = self.aabbs[lidx];
    const auto rbox = self.aabbs[ridx];
    self.aabbs[parent] = merge(lbox, rbox);

    // 親ノードへ……
    parent = self.nodes[parent].parent_idx;
}
return;

任意形状のサポート

これは少し毛色が変わって、アルゴリズムではなく実装・ライブラリ設計の話になる。

BVHはAABBさえ作ることができればどんな形状のオブジェクトも格納できる。なので、格納するオブジェクトを例えば点に絞ってしまうのはもったいない。

このような場合、複数のデザインチョイスがありえる。

基底クラスの（スマート）ポインタを持つことにし、その基底クラスが get_aabb() のようなメソッドを持っておく。ユーザーは適切にget_aabbを実装した派生クラスをそこに差し込む。BVHはget_aabb()を使ってオブジェクトのAABBを取得し、なんやかんやする。
template を使って任意の型のオブジェクトを詰め込めるようにする。ただし、その型からAABBを取得する方法をなにか指定する必要がある。

1. はクラスベースオブジェクト指向言語で書かれたライブラリでよく見られるやり方で、形状を追加するのが容易だ（継承したクラスを作って、後から差し込めば木自体をコンパイルし直さなくても動くので）。ただし、既存の型（float4など）を使おうと思ったとき、基底クラスを追加できないので、継承しただけのラッパークラスを作る必要がある。また、オブジェクトにアクセスする際に必ずポインタの間接参照をするために遅くなることがあり得る。もちろんこれも場合に依って、オブジェクトを動かすときにポインタしか動かさずに済む分速いこともある。

2. は、Boostなどテンプレートをバリバリ使うライブラリが選ぶやり方で、間接参照がないこととコンパイル時にどの関数を呼ぶかが静的に決まるのでパフォーマンスが出やすい。代わりにコンパイル時間とコンパイルエラーメッセージの短さを犠牲にする。また、あとで違う型を入れることになった際の作業量が少し増える。少なくともコンパイルはやり直す必要がある。

私はGPU上で継承するのがつらそうだと思ったので2.を選んだ。

どちらの場合でもAABBを計算するメソッドを実装しなければならないが、今回はユーザー定義のAABBを取得する関数オブジェクトを追加で渡してもらうことにした。以下のようなノリになる。実際のコードではないが。

template<typename Object, typename AABBGetter>
struct bvh
{
    void f()
    {
        Object obj = /* ... */;
        AABBGetter get_aabb;
        const auto box = get_aabb(obj);
        return;
    }
};

やり残していること

intersection queryで見つかるオブジェクトの数を数える

今、衝突判定のために、ある矩形領域とAABBが被っているオブジェクトをすべて取ってくる、というような関数が用意されている。この関数はユーザーにバッファを作ってもらって、そのバッファの上限サイズと先頭イテレータを受け取っている。被っているオブジェクトが見つかったらイテレータに書き込んだあとイテレータを進める。こうすれば、staticなストレージに書き込んでほしいユーザーは上限と静的配列のbegin()を渡せばいいし、動的メモリ確保をして構わないから全部返してほしいユーザーはback_insert_iteratorを渡せば良い。その場合は上限はsize_tの上限でいいだろう。実質無限だ。

だが、CUDAで使うとき、最初に必要な量を教えてもらって、その後全体で必要なメモリを一気に確保して、めいめいそこに書き込んでいくという戦略もありな気がしてきた。なので、書き込まずにカウントだけして返すという関数を実装するのは決して無意味ではないように思う。

拡張morton codeのサポート

morton codeを使うと空間がおよそ一様に分割される。これは、非常に粗いポリゴンがある世界の中で、例えばゲームの主人公がいる周りだけがとても精密なポリゴンがある、みたいな状況だと不便だ。なぜなら、非常に粗いポリゴンによって広い世界が作られているなかで、主人公の周りだけ非常に密にオブジェクトがあることになる。すると、主人公の周りの、特に重要なはずのポリゴンでmorton codeの衝突が起きまくる。すると重要な箇所なのに木構造の品質が悪くなり、分割の恩恵が受けられない。

これは例えば粗いポリゴンと細かいポリゴンのそれぞれのために木を2つ作ってしまうというやり方で回避することもできるだろう。だが他に、morton codeにAABBの大きさを含めてしまって、小さなAABBと大きなAABBを区別できるようにするというアイデアがあるらしい。これに関しても、MortonCodeCalculatorのようなオブジェクトを受け取ることにすればユーザーが選べるようになる。これは比較的簡単な変更で入れられるので、サポートしておきたい。

User-defined predicatorのサポート

今の所BVHに投げられるクエリが、矩形とのオーバーラップと最近傍しかない。だが、多くのユースケースでは、ユーザーしか知らないオブジェクトの性質によってクエリの結果を分けたいこともあるだろう。例えば、りんごとみかんが入っているが、今は自分の周り1mにあるりんごの位置だけが知りたくて、みかんは必要ないとか。そういうときのために、Predicatorを受け取ってそれでフィルタするという機能は普通にありだ。ただ一般的な形で、使いやすく、高速に、となるとちょっと手間はかかるだろうが。

まとめ

実装しました。これでシミュレーションとかレイトレとか速くできるといいな

in neuro