はじめに
前置きはここ．
今回は，Halide::Func::compute_withについて紹介．
Halideのチュートリアルを終えたばかりの方は，"compute_なんとか"というスケジュールメソッドについて，compute_inline・compute_root・compute_atはご存じかと思います（一応チュートリアル終わった人向けに書いているつもりなので，これらの詳細はここでは省略）．
しかし，computeシリーズはもう一つ存在しており，それが本記事のメインであるcompute_withです．
compute_withは有能な機能ですが，チュートリアルでは触れられていないためか，あまり知られていない機能（な気がする）ので，布教の手助けになればと本記事を書きました．
compute_withとは
compute_withは，あるFuncに対して異なるFuncとループを共有させるスケジュールです．
詳しくは，公式ドキュメントにて説明されていますので，そちらを参考にしてください．
下に，そのドキュメントからコードを交えた解説を引用します．
f(x, y) = x + y;
g(x, y) = x - y;
h(x, y) = f(x, y) + g(x, y);
f.compute_root();
g.compute_root();
f.split(x, xo, xi, 8);
g.split(x, xo, xi, 8);
g.compute_with(f, xo);
This is equivalent to:
for y:
for xo:
for xi:
f(8*xo + xi) = (8*xo + xi) + y
for xi:
g(8*xo + xi) = (8*xo + xi) - y
for y:
for x:
h(x, y) = f(x, y) + g(x, y)
このように，ループ構造が一部または全て一致している2つのFuncについて，引数を用いてどのループレベルを共有するかを指定したうえで，ループの冗長性を減らします．
上の解説では，超単純な計算パイプラインでのデモでしたが，実際に扱うHalideコードの計算パイプラインをもっと複雑なはずです...よね？
おそらくHalidユーザの方が息を吐くように使用している，vectorizeやunroll・parallelといったスケジュールメソッドは，compute_withと共存させるのに，少し工夫が求められます．
（それらの対処法についての備忘録としても書いてます）
compute_with使用例
チュートリアルで登場する"compute_なんとか"は純粋定義でしか使用できませんでしたが，compute_withは純粋・更新二つの定義に対して使用できます．
そのため，畳み込みなどで使用するRDomが生成するループもcompute_withはまとめることができます．
下のコード例では，同じループ構造を持つ関数fと関数gに対して，純粋・更新定義のループをまとめます．
（コードについて，処理内容に特に意味はありません）
using namespace Halide;
Buffer<uint8_t> src_f(512, 512), src_g(512, 512);
Func f("f"), g("g"), h("h");
Var x("x"), y("y");
Func clamped_f = BoundaryConditions::repeat_edge(src_f);
Func clamped_g = BoundaryConditions::repeat_edge(src_g);
RDom r(-3, 7, -3, 7, "r");
// algorithm part
f(x, y) = 0;
f(x, y) = f(x, y) + clamped_f(x + r.x, y + r.y) / 49;
g(x, y) = 0;
g(x, y) = g(x, y) + clamped_g(x + r.x, y + r.y) / 49;
h(x, y) = f(x, y) + g(x, y);
// scheduling part
f.compute_root();
g.compute_root().compute_with(f, y).update().compute_with(f.update(), r.x);
h.compute_root();
h.print_loop_nest();
produce f:
produce g:
for fused.y:
for x:
f(...) = ...
for x:
g(...) = ...
for y:
for x:
for r in [-3, 3]:
for r in [-3, 3]:
f(...) = ...
g(...) = ...
consume f:
consume g:
produce h:
for y:
for x:
h(...) = ...
このように，fとgで更新定義を含め，ループを共有していることが分かります．
また，上の例ではfとgは完全に同じループ構造を持つため，gについてcompute_with(f,x)とすれば，ループ全てを共有させることも可能です．
また，compute_withの使用制限として，fとgのように完全に独立（互いに干渉していない）している必要があります．
なのでfとgを消費するhとはcompute_withでループをまとめることはできません（これはcompute_atの仕事）．
．．．一方，おそらく先ほどのコードを見たHalide使いの方は，「vectorizeとかもしたくね...？」と思われたはず．
しかし先述の通り，通常のやり方ではvectorizeとcompute_withは共存が難しいです．
vectorizeとcompute_with
プログラムを高速化する定番の手法であるSIMD命令によるベクトル化，これをHalide上で行うにはvectorizeメソッドを用います．
このvectorizeメソッドは，ちょっと厄介な挙動をするため，compute_withを共存させるのは少しテクニックがいります．
例えば，次のように先ほどのコードにおいてスケジュール部分を変更させて，xを幅8でベクトル化さる場合，（なにも考えなければ）次のようになります．
// scheduling part
f.compute_root().vectorize(x, 8)
.update().vectorize(x, 8);
g.compute_root().vectorize(x, 8).compute_with(f, x)
.update().vectorize(x, 8).compute_with(f.update(), r.x);
Error: Invalid compute_with: names of dim 2 of f.s1(x.v5) and g.s1(x.v7) do not match.
エラーが返ってきました．
違う変数どうしてcompute_withできないと怒られました．
違うとされているx.v5とx.v7ですが，これらはアルゴリズム記述では登場しておらず，スケジュールメソッドで自動的に生成されたものだと推測できます．
この謎の変数x.vはvectorizeメソッド，特にベクトルサイズとしてfactorを引数に持たせたときに出現するものです．
vectorizeは本来，引数に与えられた変数のループをベクトル化するもので，factorが設定された場合はsplitで変数を分離した上でvectorizeを設定します．
これをコードで示すとしたのようになってます．
// vectorize with factor
f(x, y) = hoge;
f.vectorize(x, 8);
// above equal below
f(x, y) = hoge;
Var xv("x.v");
f.split(x, x, xv, 8).vectorize(xv);
このように，factor付きvectorizeはsplitを内包しており，上のコードで使用したxvがvectorize内部で自動的に生成され，変数名に固有名が割り当てられます．
そのため，fのvectorizeとgのvectorizeで生成されるxvの名前が異なります．
Halideコンパイラでは，変数名が異なる場合異なるループと認識するらしく，これが原因でcompute_withがはじかれます．
解決策1 ... hにcompute_atさせる
ベクトル化をしつつcompute_withも適用させるために，今まで空気な存在だったhが活躍します．
上手く動作するコードを下に示します．
// scheduling part
h.vectorize(x, 8);
string var_name = h.function().definition().schedule().dims()[0].var;
f.compute_at(h, Var(var_name));
g.compute_at(h, Var(var_name)).update().compute_with(f.update(), r.x);
produce h:
for y:
for x.x:
vectorized x.v4 in [0, 7]:
produce f:
produce g:
f(...) = ...
g(...) = ...
for r in [-3, 3]:
for r in [-3, 3]:
f(...) = ...
g(...) = ...
consume f:
consume g:
h(...) = ...
この方法では，hに対してxをベクトル化し，fとgをhのベクトル化したxにcompute_atします．
すると，fとgが持つループがRDomによるループだけに限定されるため，fとgでcopute_withが通ります．
一方，この方法では，hのベクトル化した変数を得るところが，一般的なHalideの記述ではできないため，少しトリッキーな形になります．
別記事で少し紹介した通り，Halide::Funcの本体はHalide::Internal::Functionが担っており，スケジュールにおけるループ（の次元）はFunction::definition().schedule().dims()で取得できます．
この戻り値はvector<Halide::Internal::Dim>で，Funcのスケジュールにおける次元を内側から格納されています．
したがって，この0番目（つまり最内側）を取得することで，ベクトル化したxの情報を引っ張ってくることができます．
このように，fとgを消費するhにcompute_atを行うことで，解決できますが，これは常に有効に行えるとは限りません．
また，ベクトル化した変数の取得がめんどくさいですが，これは後述する別の解決策と組み合わせれば何とかなります．
解決策2 ... vectorizeとsplitを分離
vectorizeがsplitを内包し，隠蔽するため，名前が異なるVarが生まれるので，vectorizeに含まれるsplitもユーザ側で行えば解決できます．
Var xv("xv");
f.compute_root().split(x, x, xv, 8).vectorize(xv)
.update().split(x, x, xv, 8).vectorize(xv);
g.compute_root().split(x, x, xv, 8).vectorize(xv).compute_with(f, xv)
.update().split(x, x, xv, 8).vectorize(xv).compute_with(f.update(), r.x);
produce f:
produce g:
for fused.y:
for x.fused.x:
vectorized x.fused.xv in [0, 7]:
f(...) = ...
g(...) = ...
for y:
for x.x:
vectorized x.xv in [0, 7]:
for r in [-3, 3]:
for r in [-3, 3]:
f(...) = ...
g(...) = ...
consume f:
consume g:
produce h:
for y:
for x:
h(...) = ...
このように，ベクトル化を含んだ状態で，compute_withを適用することができました．
（ただし，hへのcompute_atもしたほうがキャッシュ効率的にはいいかも...?）
vectorize以外のスケジュールとcompute_with
vectorizeのように，ループ構造を変形して最適化をかけるスケジュールメソッドとして，unrollやparallelがありますが，これらもvectorizeと同じく，factorを設定すると内部的にsplitも発行されます．
そのため，これらとcompute_withの共存は同じ問題が発生しますが，同じように解決できます．
compute_atとcompute_with
先ほどまでの例において，fとgで共通して消費される関数があった場合，その関数を事前計算して再利用した方が効率が良い場合があります．
そこで，compute_withとcompute_atを組合せて使用します．
下にコード例を示します．
Buffer<float> src_f(512, 512), src_g(512, 512);
Func e("e") ,f("f"), g("g"), h("h");
Var x("x"), y("y");
Func clamped_f = BoundaryConditions::repeat_edge(src_f);
Func clamped_g = BoundaryConditions::repeat_edge(src_g);
RDom r(-3, 7, -3, 7, "r");
// algorithm part
e(x, y) = Halide::exp(-pow(clamped_f(x, y) - clamped_g(x, y), 2));
f(x, y) = 0.f;
f(x, y) = f(x, y) + e(x + r.x, y + r.y) * clamped_f(x + r.x, y + r.y);
g(x, y) = 0.f;
g(x, y) = g(x, y) + e(x + r.x, y + r.y) * clamped_g(x + r.x, y + r.y);
h(x, y) = f(x, y) + g(x, y);
// scheduling part
Var xv("xv");
h.split(x, x, xv, 8).vectorize(xv);
f.compute_at(h, xv);
g.compute_at(h, xv).update().compute_with(f.update(), r.x);
e.compute_at(f, x);
h.print_loop_nest();
produce h:
for y:
for x.x:
vectorized x.xv in [0, 7]:
produce f:
produce g:
g(...) = ...
f(...) = ...
produce e:
for y:
for x:
e(...) = ...
consume e:
for r in [-3, 3]:
for r in [-3, 3]:
f(...) = ...
g(...) = ...
consume f:
consume g:
h(...) = ...
このように，fおよびgで消費されるeをcompute_atさせ，fとgでの再計算を防ぎます．
また，スケジュールでは，eはfに対してcompute_atすることが必要です．
gに対してcompute_atした場合は，下のようなエラーがでます．
Legal locations for this function are:
e.compute_root();
e.compute_at(h, Var::outermost());
e.compute_at(h, y);
e.compute_at(h, x);
e.compute_at(h, xv);
e.compute_at(f, Var::outermost());
e.compute_at(f, y);
e.compute_at(f, x);
e.compute_at(f, r$y);
e.compute_at(f, r$x);
"e" is used in the following places:
for h.s0.y:
for h.s0.x.x:
for h.s0.x.xv:
...
for f.s1.y:
for f.s1.x:
for f.s1.r$y:
for f.s1.r$x:
g uses e
for g.s1.fused.y:
for g.s1.fused.x:
for g.s1.fused.r$y:
for g.s1.fused.r$x:
g uses e
このように，fにcompute_atすることが求められます．
おわりに
compute_withの説明と，記述方法や応用を紹介しました．
compute_withの弱点（？）として，今回のコード例におけるhのような，fとgを両方消費する関数が必要になります．
つまり，fとgを別々にrealizeしたい場合などでは，compute_withできないようになっています（ここで議論されてました）．
これが解消されると，いろいろ応用範囲が広がる気もしますが．．．
あと，vectorizeがsplitを内包する点はチュートリアルにも書いてありましたね．．．（さっき見て気がついた）