LLVM/C++日記　Halideから読み解くLLVM

1年ほど前から，swiftコンパイラ周りの興味からLLVMを勉強し始めていましたが，チュートリアルを一通り読み終えた後から，モチベーションが続きませんでした．やはり，何かしら，作りたい目標がないと私は動けないようです．しかし，ひょうなことからそのモチベーションが復活しました．それは，Halideというソフトウェアの勉強をしていたときでした．

LLVMの勉強メモ・・・https://github.com/sonsongithub/llvm-tutorial

計算グラフの実装・・・https://github.com/sonsongithub/llvm_jit_compile

Halide

Halideは，画像コミュニティではよく知られたソフトウェアなのですが，そこまで一般的というものではないので，概要をお話しします．一言で言うと，画像処理等を高速処理したいときの並列計算を，SIMD，マルチスレッド，データのアクセス方法の工夫などを用いて，自動的に，コードを生成してくれるというものです．もう少し具体的にいうと，プログラマがアルゴリズムを"Define by run"のような形でC++上で実装すると，Halideがそれを環境に合わせて自動的に並列化してくれるというものです．

Maruさんからコメントを頂き追記します．HalideはAOTもできます．つまり，実行前にバイナリとして最適化された関数を生成することも可能です．

SIMDは，いわゆる"single instruction, multiple data"のことで，１回のCPUの命令で複数のデータ処理を実行できるCPUの特殊命令群です．例えば，ワンクロックで足算を8個同時に処理したり，絶対値の計算結果を同時にX個取得したりすることができ，ベクトルや行列の演算の高速化に必ずといっていいほど，利用されるものです．しかし，CPUごとによって使えるSIMD命令が異なるため，SIMDを使う場合，コードの可搬性に注意する必要があります．マルチスレッドは言うまでもなく，CPUのマルチコアを利用して，コードを同時に実行し，計算を高速化することです．

最後にデータのアクセスの工夫について．どんなデータもメモリ上に並んでいます．画像の全ピクセルを処理するとき，画像のピクセルをひとつひとつ読み込んでいくと遅くなるので，大概の場合，まとめて読み込みます．実装によりますが，画像のピクセルがメモリ上で"1行"ずつにならんでいると，当然．列方向にアクセスするとき，メモリのアドレスが大きくことなることになります．こういったジャンプが大きくなると，メモリとキャッシュとCPUのアクセスが頻繁に発生し，処理速度が悪化します．この問題は，アルゴリズムにおいて，データを処理する順番などを工夫すると，改善することができます．

Halideで書くとこんなにお手軽

実際に，[1]で示された例を見ると，その生産性の高さは一目瞭然です．画像にボカシを入れる処理を，SIMD，マルチスレッド，データアクセスの工夫を凝らして，実際にx86環境で実装すると以下のようなコードになるそうです．

void fast_blur(const Image &in, Image &blurred) {
   m128i one_third = _mm_set1_epi16(21846);
   #pragma omp parallel for
   for (int yTile = 0; yTile < in.height(); yTile += 32) {
       m128i a, b, c, sum, avg;
       m128i tmp[(256/8)*(32+2)];
       for (int xTile = 0; xTile < in.width(); xTile += 256) {
           m128i *tmpPtr = tmp;
           for (int y = -1; y < 32+1; y++) {
               const uint16_t *inPtr = &(in(xTile, yTile+y));
               for (int x = 0; x < 256; x += 8) {
                   a = _mm_loadu_si128(( m128i*)(inPtr-1));
                   b = _mm_loadu_si128(( m128i*)(inPtr+1));
                   c = _mm_load_si128(( m128i*)(inPtr));
                   sum = _mm_add_epi16(_mm_add_epi16(a, b), c);
                   avg = _mm_mulhi_epi16(sum, one_third);
                   _mm_store_si128(tmpPtr++, avg);
                   inPtr += 8;
               }
           }
       }
       tmpPtr = tmp;
       for (int y = 0; y < 32; y++) {
           m128i *outPtr = ( m128i *)(&(blurred(xTile, yTile+y)));
           for (int x = 0; x < 256; x += 8) {
               a = _mm_load_si128(tmpPtr+(2*256)/8);
               b = _mm_load_si128(tmpPtr+256/8);
               c = _mm_load_si128(tmpPtr++);
               sum = _mm_add_epi16(_mm_add_epi16(a, b), c);
               avg = _mm_mulhi_epi16(sum, one_third);
               _mm_store_si128(outPtr++, avg);
           }
       }
   }
}

しかし，これをHalideを使って書くと，以下のようにコンパクトになります．さらに，実行/コンパイル時に環境に合わせて，コードを展開してくれるので，ポータブルな実装にもなります．

Func halide_blur(Func in) { Func tmp, blurred;
   Var x, y, xi, yi;
   // The algorithm
   tmp(x, y) = (in(x-1, y) + in(x, y) + in(x+1, y))/3;
   blurred(x, y) = (tmp(x, y-1) + tmp(x, y) + tmp(x, y+1))/3;
   // The schedule
   blurred.tile(x, y, xi, yi, 256, 32).vectorize(xi, 8).parallel(y);
   tmp.chunk(x).vectorize(x, 8);
   return blurred;
}

Halideの実装

Halideの特殊な書き方を見ると，C++上の文法で計算グラフを構築し，おそらく，それをJITコンパイルしているのだろうと推測できます．Halideのサイトや解説文書などを見ていると，やはり，そのようになっているようです．

計算グラフは，Chainerに端を発し，現在，深層学習の実装とは切っても切り離せない概念です．そこで，私は，Halideではそれがどう実装されているのか学ぶため，自分で計算グラフを実装し，JITコンパイルして，動的に実行できるソフトウェアを書いてみることにしました．しかし・・・そこには大きな壁があったのです．

C++がわからない

早速Halideのコードを眺めてみます．全然読めない・・・・．まじで読めない・・・・・・．え・・・・つらい・・・・．さて・・・私は，自分の職務経歴書の技術からC++を消しました．そして，わたしのC++とLLVMをさまよう旅が始まったのです．

続く

参考文献

[1] Jonathan Ragan-Kelley, Andrew Adams, Sylvain Paris, Marc Levoy, Saman Amarasinghe, and Frédo Durand. 2012. Decoupling algorithms from schedules for easy optimization of image processing pipelines. ACM Trans. Graph. 31, 4, Article 32 (July 2012), 12 pages. DOI:https://doi.org/10.1145/2185520.2185528

[2] 「画像処理を簡単に高速化してみませんか！？」 Halideによる画像処理プログラミング入門 https://www.slideshare.net/fixstars/halide-82788728