穴日記

どうだ明るくなったろう

enokiで遊んでみた(Visual Studio 2017編)

これはレイトレアドベントカレンダー2019の記事です。
https://qiita.com/advent-calendar/2019/raytracing


https://enoki.readthedocs.io/en/master/

enokiというのは最近公開されたC++テンプレートライブラリで、SIMDやCUDAによるベクタライズを抽象的にサポートした算術ベクトルライブラリとなっており、自動微分も可能なものになっています。SIGGRAPH ASIA 2019で発表されたMitsuba2(https://rgl.epfl.ch/publications/NimierDavidVicini2019Mitsuba2)はenokiを全面的に採用しており、SIMDによる"wide"なレンダラを容易に記述できる、Differentiable Rendererを容易に記述できる、といった恩恵が得られると主張しています。

以下は上記オフィシャルページからのコピペですが、以下の様なsRGB Gammaのコードを

float srgb_gamma(float x) {
   if (x <= 0.0031308f)
      return x * 12.92f;
   else
      return std::pow(x * 1.055f, 1.f / 2.4f) - 0.055f;
}

以下のようにenokiスタイルに書き換えると

template <typename Value> Value srgb_gamma(Value x) {
   return enoki::select(
      x <= 0.0031308f,
      x * 12.92f,
      enoki::pow(x * 1.055f, 1.f / 2.4f) - 0.055f
   );
}

コンパイルターゲットをSSE4.2、AVX2、AVX512などに変えることで自動的にバックエンドでSIMD化されて高速なコードが生成されるという寸法です。
やってることとしては、各種オペレーターをオーバーロードしてSIMD intrinsicをラップしている感じではあるのですが、そういう典型的なコードを包括的に実装してくれています。
また、各種数学関数(powとか)のSIMD用実装も提供しているので、そのあたりも便利ですね。

で、CPU向けのSIMDラッパライブラリとして使う分には必要なヘッダをインクルードするだけで使えます。
ここから、GPU向け(CUDA向け)ライブラリとして使ったり、自動微分機能を使ったりするには、いくつかのcppファイルをコンパイルする必要がでてきます。
また、CUDAを実行するにはCUDA toolsetなんかもインストールする必要があります。

思想としては、上記のようなスタイルで抽象的にアルゴリズムを記述したら、あとはenokiがよしなにしてくれて、CPU SIMDからCUDA、自動微分まで何もかもサポートされる!というものです。
が、現実には、たとえばCUDAサポートしようと思うとCPU SIMDむけのenoki::Arrayとは別のenoki::CUDAArrayが用意されてたりして、完全にenokiレイヤで抽象化されてるわけでもないです。(まあ自前で抽象化してもいけそうではありますが)
一方で、SoAデータ構造をストレートに記述できたりするあたりは賢い感じします。(例えば、enoki::Array< enoki::Array, 3>でRGB x 16のSoAになる)

導入(SIMDライブラリ部分)

基本的にLinuxターゲットぽいんですけどgithubリポジトリをみるとWinビルドもテスト通っています。
なので、Win上でVisual Studio使ってみます。
最初はVS2019使ってみたんですけど、途中でビルドエラーが解決困難になったので、VS2017を使いました。

まず、ビルドにはcmakeが必要なので最新バージョンをインストールします。
githubからenokiの最新のリポをcloneして、以下コマンドを実行してみます。
https://github.com/mitsuba-renderer/enoki

cmake -G "Visual Studio 15 2017" -A "x64" -DENOKI_TEST=ON  .

こちらで生成されたソリューションファイルからRUN_TESTS選んでビルドして実行すると基本的なテストが一通り通るはずです。

そしたら、自前のプログラムから呼ぶことを考えます。
これも簡単で、単にenoki/includeフォルダにパス通して#include とかすればいいです。
バックエンド(SSEだとかAVXだとか)の指定はプロジェクトファイル側で適当に設定すれば自動で検出されます。
ちなみに、C++17が前提になっているので、その辺りの指定も必要です。
なんか細かい話として、iostreamヘッダをenokiヘッダより先にインクルードするとM_PIとかが死ぬとかあります。

あとはドキュメントを読むといいです。かなり丁寧に書いてあります。以下のような感じでベクタライズされたarrayを作って演算を記述するイメージです。

enoki::Array<float, 20> arr;

ちなみに、上記20幅のarray作ると以下のように適切な幅のSIMDレジスタで埋まるらしいです。すごいですね。
f:id:h013:20191209181124p:plain:w360

GPUと自動微分

CPUのSIMDバックエンド使うだけならヘッダインクルードするだけです。とても簡単ですね。
といっても、テンプレートを駆使しているため、MSVCだと挙動不審になることはあります。たぶんclangとかだともっとちゃんと?動くのかもしれない。

さて、GPUや自動微分もしてみたくなります。
まずCUDA toolsetをインストールします。そして、以下のようにcmakeコマンドを打ってみます。

cmake -G "Visual Studio 15 2017" -A "x64" -DENOKI_TEST=ON -DENOKI_CUDA=ON -DENOKI_AUTODIFF=ON -DENOKI_PYTHON=ON .

今度はCUDAとAUTODIFFをONにしてます。PYTHONというのは各種enoki機能のPythonバインダなのですが、なんか手元だとうまく動きませんでした。誰か動かしてみてください!

で、enoki-cudaとenoki-autodiffプロジェクトをlibビルドするわけですが、そのままだと通りませんでした。(コミット:dfc8e8978813659308ad86a09ffc89b5d431f734 )
そこで以下のような修正をしました。(気が向いたらプルリクエスト出すかもしれない)

  • include/enoki/autodiff.hの1436行目以降、ENOKI_IMPORTの位置をtemplateの直後に移動。
  • src/autodiff/autodiff.cppの1113行目以降、ENOKI_EXPORTの位置をtemplateの直後に移動。
  • src/python/common.hでlambda式の中でif constexpr使っている箇所があるが、VS2017だとconstexprとしてコンパイルされないっぽいのでlambda式をfree functionとして分離。

以上修正するとビルドが通るので、autodiff_nativeプロジェクト側でリンクすると、テストプロジェクトのビルドも成功し、テストも通りました。
このあたりはちょっと怪しくて、dllとかでも動くような感じもしつつ、うまくいかないような感じもしつつという感じです。VisualStudio版はまだちょっと発展途上って感じですね。

しかしまあ、これで自動微分もCUDAでの実行も普通に行えます。(テストプロジェクトのコードを参考にするといいですね)CUDAで実行するとプログラムが終了しきらないことがあるのですが、まあご愛敬ということで。以下のような自動微分コードが普通に動きます(しかもベクタライズされて)

DiffArray<DynamicArray<Packet<float>>> a(1.0f, 2.0f, 3.0f, 5.0f);
set_requires_gradient(a);

auto b = a * a;
forward(a);
std::cout << gradient(b) << std::endl;

上記の出力結果は、a^2の微分が2 * aなので、[2, 4, 6, 10]になります。

所感

SIMDラッパライブラリ的に使うのであればかなり容易に使えます。それだけでも価値が高いのではないかという気がします。
自動微分も、わりと簡単に利用できます。CUDAが絡んでくると、ちょっと面倒かも。でもそこまで大変ではないです。
Pythonバインド周りは諦めちゃいましたが、基本的にはpybind11使ってるだけっぽいし、ということはWin/VS2017でも問題なく動くはずです(理論的には)

今回は触れませんでしたが、SoA化とか、他の様々な機能も豊富です(ドキュメント読むといいでしょう)
C++のいい感じのベクタライズライブラリでちゃんとメンテもされてる最新のもの、という意味では割と未来があるのではないでしょうか。

おわり。