Halide小ネタ(malloc編)

はじめに

前回の続きです．前置きは前回と同じなので省略．
今回のテーマはHalide::Funcにおけるmallocについて．

compute_root VS realize

早速ですが，下のコードはセパラブル実装したガウシアンフィルタです．

using namespace Halide;

void Gaussian_compute_root(Halide::Buffer<float>& src, int rad, float sigma)
{
    Var x("x"), y("y");
    Func clamped = BoundaryConditions::repeat_edge(src);

    RDom r(-rad, 2 * rad + 1);
    Func blur_x("blur_x"), blur_y("blur_y"), total("total"), kernel("kernel");

    Expr d = -1.f / (2.f * sigma * sigma);
    total() = sum(fast_exp((r * r) * d));

    kernel(x) = fast_exp((x * x) * d);
    blur_x(x, y) = sum(kernel(r) * clamped(x + r, y)) / total();
    blur_y(x, y) = sum(kernel(r) * blur_x(x, y + r)) / total();

    Var xi("xi"), yi("yi"), tile_index("tile_index");
    clamped.compute_inline();
    total.compute_root();
    kernel.compute_root();
    blur_x.compute_root().tile(x, y, xi, yi, 64, 64).fuse(x, y, tile_index).parallel(tile_index)
        .vectorize(xi, 16);
    blur_y.compute_root().tile(x, y, xi, yi, 64, 64).fuse(x, y, tile_index).parallel(tile_index)
        .vectorize(xi, 16);

    Halide::Buffer<float> dst(src.width(), src.height());
    // 実行時間計測
    double result = Halide::Tools::benchmark(20, 20, [&]()
        {
            // 実行部分
            blur_y.realize(dst);
        }
    );
    std::cout << "compute_root : " << result * 1e3 << "ms\n";
}

2次以上の畳み込みを複数の1次の畳み込みに分割するのがセパラブル実装です．上のコードでは2次元畳み込みをx方向(blur_x)とy方向(blur_y)に分割して行います．blur_xが入力画像に対してx方向のフィルタリングを行い，blur_yがblur_xに対してy方向のフィルタリングを行うことで，出力を得ています．

ここで，blur_xのスケジュールについて，compute_rootを指定していることに注目してください（通常なら，blur_yでタイリング処理しているので，blur_xはcompute_at(blur_y,tile_index)とかにした方が速くなりますが，あえてcompute_rootに）．つまり，blur_yが参照するblur_x全体を計算してから，blur_yを計算しています．

次に，下のコードも同じくセパラブル実装のガウシアンフィルタです．

void Gaussian_realize(Halide::Buffer<float>& src, int rad, float sigma)
{
    bool auto_schedule = false;
    using namespace Halide;
    Var x("x"), y("y");
    Func clamped = BoundaryConditions::repeat_edge(src);

    Buffer<float> inter_out(src.width(), src.height() + rad * 2);
    inter_out.set_min(0, -rad);

    RDom r(-rad, 2 * rad + 1);
    Func blur_x("blur_x"), blur_y("blur_y"), total("total"), kernel("kernel");

    Expr d = -1.f / (2.f * sigma * sigma);
    total() = sum(fast_exp((r * r) * d));

    kernel(x) = fast_exp((x * x) * d);
    blur_x(x, y) = sum(kernel(r) * clamped(x + r, y)) / total();
    blur_y(x, y) = sum(kernel(r) * inter_out(x, y + r)) / total(); // <- ここ

    Var xi("xi"), yi("yi"), tile_index("tile_index");
    clamped.compute_inline();
    total.compute_root();
    kernel.compute_root();
    blur_x.compute_root().tile(x, y, xi, yi, 64, 64).fuse(x, y, tile_index).parallel(tile_index)
        .vectorize(xi, 16);
    blur_y.compute_root().tile(x, y, xi, yi, 64, 64).fuse(x, y, tile_index).parallel(tile_index)
        .vectorize(xi, 16);

    Halide::Buffer<float> dst(src.width(), src.height());
    // 実行時間計測
    double result = Halide::Tools::benchmark(20, 20, [&]()
        {
            // 実行部分
            blur_x.realize(inter_out); // <- ここ
            blur_y.realize(dst);
        }
    );
    std::cout << "realize : " << result * 1e3 << "ms\n";
}

上のコードは先ほどのガウシアンフィルタと殆ど一緒ですが，blur_yはinter_outというBufferに対してフィルタリングを行っていることに注目してください．inter_outは中間出力で，blur_xの出力が格納されるBufferです．そのため，実行部分でblur_xはinter_outに対してrealizeを行い，続けてblur_yがrealizeを行っています（FuncのBufferへの参照はポインタで行っているようで，このような実装でもちゃんと動く）．つまり，realizeを2回実行することで，blur_x全体を計算してから，blur_yの計算を行っています．

さて，これら2つのガウシアンフィルタについて，どちらもblur_x全体を計算してから，blur_yの計算を行っていますが，どちらの方が実行速度が速いでしょうか?

実験結果

環境はIntel Core i7-8550U CPU @ 1.80GHzです．
同じ入力，パラメータで実行した結果がこちら．

compute_root : 1.086505ms
realize : 0.508755ms

実は，realizeを2回呼び出す方が速いです．

mallocをトレースしてみる

Halide::Toolsには様々な便利機能があり，上のコードで使用している実行時間を計測するHalide::Tools::benchmarkもその1つです．
Halide::Tools::halide_malloc_traceとHalide::Tools::halide_free_traceは，Func::set_custom_allocatorに渡すことでFuncが実行されるときに呼び出されるmallocやfreeをトレースすることができます．
これを使って先程の2つのガウシアンフィルタを比較してみます．

compute_rootで実装したガウシアンフィルタについて，// 実行時間計測以下を下のように変更．

    blur_y.set_custom_allocator(Halide::Tools::halide_malloc_trace, Halide::Tools::halide_free_trace);

    std::cout << "compute_root\n1st try\n";
    blur_y.realize(dst);
    std::cout << "\n\n2nd try\n";
    blur_y.realize(dst);
    std::cout << "\n\n3rd try\n";
    blur_y.realize(dst);

set_custom_allocatorはrealizeを呼ぶFuncに対して行います．

同様にrealize2回で実装したガウシアンフィルタについて，下のように変更．

    blur_y.set_custom_allocator(Halide::Tools::halide_malloc_trace, Halide::Tools::halide_free_trace);
    blur_x.set_custom_allocator(Halide::Tools::halide_malloc_trace, Halide::Tools::halide_free_trace);

    std::cout << "realize\n1st try\n";
    blur_x.realize(inter_out);
    blur_y.realize(dst);
    std::cout << "\n\n2nd try\n";
    blur_x.realize(inter_out);
    blur_y.realize(dst);
    std::cout << "\n\n3rd try\n";
    blur_x.realize(inter_out);
    blur_y.realize(dst);

実験結果

compute_root
1st try
halide_malloc => [0x1ec34c13080, 0x1ec34d22083], # size:1110020, align:128
halide-header => [0x1ec34c13040, 0x1ec34c1307f], # size:64, align:64
halide_free => [0x1ec34c13040, 0x1ec34c1307f], # size:64, align:64


2nd try
halide_malloc => [0x1ec34c19080, 0x1ec34d28083], # size:1110020, align:128
halide-header => [0x1ec34c19040, 0x1ec34c1907f], # size:64, align:64
halide_free => [0x1ec34c19040, 0x1ec34c1907f], # size:64, align:64


3rd try
halide_malloc => [0x1ec34c12080, 0x1ec34d21083], # size:1110020, align:128
halide-header => [0x1ec34c12040, 0x1ec34c1207f], # size:64, align:64
halide_free => [0x1ec34c12040, 0x1ec34c1207f], # size:64, align:64

realize
1st try


2nd try


3rd try

compute_root実装ではrealize毎でmallocが呼ばれいる（blur_xの結果を格納するため）のに対し，realize2回実装ではmallocが呼ばれいません（中間結果はinter_outに保存されるため，mallocする必要がない）．
この違いが実行速度の差を生んでいると考えられます．

最後に

今回はHalideにおけるmallocの仕様について紹介しました．
Funcのrealize毎でmallocとfreeが呼び出されているので，同じFuncを複数回realizeする場合だと冗長な動作になる場合があります．

ちなみに，mallocが呼ばれる条件として，Funcの保存先がHeapになっている必要があるそうです．
デフォルトで保存先はFuncが必要な領域のサイズ等で自動的に判別されるそう（例:小さい領域ならStack）ですが，明示的に指定したい場合はFunc::store_in(Halide::MemoryType::Heap)で指定できます．

個人的にmallocの呼び出し有無だけで，果たしてここまで速度が変わるのか？と疑問に思います．もっと他に要因があるのでしょうか？(毎回メモリ領域が変わるからキャッシュ効率がめちゃ悪い？)

おまけ

realizeを複数回呼ぶのが嫌な人向け

Pipeline p({ blur_x,blur_y });
p.realize({ inter_out,dst });

詳しい説明は割愛．