Halideによる高速演算のすすめ - 数値要約編

9 min read読了の目安(約8400字

Halideによる高速演算のすすめ - 数値要約編

はじめに

Halideで数値要約を行う場合、次元削減処理を実行する必要があります。
具体例を挙げると、画像の平均値の場合、2次元のデータを0次元の数値への変換と見ることができます。

そのような処理を記述するには Halide::RDom を使用します。

RDomのコンストラクタは(開始値, 要素数)をペアとした引数を取ります。

Halide::RDom rd(start, count);

この処理は以下のforループと同等です。

for (int rd=start; rd < (start+count); ++rd) {
   // ...
}

画像横幅全ての要素にアクセスするループを実装する場合、次のように定義します。

Halide::RDom rd(0, width);

また、RDomのコンストラクタの引数にもう一つペアを追加することで、次元を増やすことができます。

Halide::RDom rd(start_x, count_x, start_y, count_y);

この処理は以下の二重forループと同等です。

for (int rd_y=start_y; rd_y < (start_y+count_y); ++rd_y) {
    for (int rd_x=start_x; rd_x < (start_x+count_x); ++rd_x) {
        // ...
    }
}

Halide::RDomを使用することで、データの数値要約処理を記述することが可能になります。

ここでは一般的な統計処理をHalideで記述する方法について紹介します。

総和

総和は統計処理で頻出します。

\sum^{N-1}_{i=0} x_i

総和処理はHalide::RDomを用いて以下のように記述します。

class Sum : public Halide::Generator<Sum> {
public:
    Halide::GeneratorInput<Halide::Buffer<std::uint8_t>> input{"input", 2};
    Halide::GeneratorOutput<Halide::Func> output{"output"};

    Halide::RDom rd;

    void generate() {
        // 初期値定義
        output() = Halide::Expr(0.0);
        auto old = output();
        
        // 値更新
        rd = Halide::RDom(0, input.width(), 0, input.height());
        auto value = input(rd.x, rd.y);
        output() = old + value;
    }
};

初期値の定義と、データ更新の定義の2部に分かれています。
これは通常の総和処理を見れば一目瞭然です。

double output = 0.0;
for (int y=0; y < height; ++y) {
    for (int x=0; x < width; ++x) {
        output += input[y][x];
    }
}

さて、Halide::RDomを使用した場合、処理にHalide::Varが関与しません。
そのため、parallel(y)といった次元を指定した高速化が行えません。

そこで、rfactor()を使用してHalide::RDomのひとつの次元をHalide::Varに割り当てます。

void schedule() {
    output.compute_root();
    Halide::Var y{"y"};
    Halide::Func intm = output.update().rfactor(rd.y, y);
    intm.compute_root().update().parallel(y);
}

このスケジュールによって、総和処理をyの次元で並列に実行することが可能になります。
ここで紹介する数値要約処理は基本的に上記スケジュールで並列実行することができます。

さて、総和処理をHalideに生成させると次のような関数が定義されます。

int sum(struct halide_buffer_t *_input_buffer, struct halide_buffer_t *_output_buffer);

0次元のデータ、すなわちスカラー値を返す場合にもhalide_buffer_t *を要求します。
この処理は次のように呼び出します。

Halide::Runtime::Buffer<std::uint8_t> input(image, width, height);
auto output = Halide::Runtime::Buffer<double>::make_scalar();
sum(input, output);
double result = output();

平均値

平均値はデータの総和をデータ数で除算することで求められます。

\mu = \frac{1}{N} \sum^{N-1}_{i=0} x_i

ところで、上式を変形した以下の式でも同じ値が求められます。

\mu = \sum^{N-1}_{i=0} {\frac{x_i}{N}}

前者は総和の結果が型の最大値を超える場合に正しい結果が得られない問題があります。
後者は除算の総和のため、誤差が生じる可能性がありますが、基本的にdoubleの演算であれば問題のない精度で結果が得られます。

後者をHalideで記述すると次のようになります。

class Mean : public Halide::Generator<Mean> {
public:
    Halide::GeneratorInput<Halide::Buffer<std::uint8_t>> input{"input", 2};
    Halide::GeneratorOutput<Halide::Func> output{"output"};

    Halide::RDom rd;

    void generate() {
        auto count = Halide::cast<double>(input.width() * input.height());

        // 初期値定義
        output() = Halide::Expr(0.0);
        auto old = output();

        // 値更新
        rd = Halide::RDom(0, input.width(), 0, input.height());
        auto value = input(rd.x, rd.y) / count;
        output() = old + value;
    }
};

最小・最大値

最小値を求めるHalideコードはこれまでの処理を応用して次のように書けます。

class Min : public Halide::Generator<Min> {
public:
    Halide::GeneratorInput<Halide::Buffer<std::uint8_t>> input{"input", 2};
    Halide::GeneratorOutput<Halide::Func> output{"output"};

    Halide::RDom rd;

    void generate() {
        // 初期値定義
        output() = input.type().max();
        auto old = output();
        
        // 値更新
        rd = Halide::RDom(0, input.width(), 0, input.height());
        auto value = input(rd.x, rd.y);
        output() = Halide::min(old, value);
    }
};

この処理は以下のコードと同等になります。

std::uint8_t min = std::numeric_limits<std::uint8_t>::max();
for (int y=0; y < height; ++y) {
    for (int x=0; x < width; ++x) {
        min = std::min(min, input[y][x]);
    }
}

さて、最小値を単独で求めることができましたが、次のような場合はどうしたらよいでしょうか。

std::uint8_t min = std::numeric_limits<std::uint8_t>::max();
std::uint8_t max = std::numeric_limits<std::uint8_t>::min();
for (int y=0; y < height; ++y) {
    for (int x=0; x < width; ++x) {
        min = std::min(min, input[y][x]);
	max = std::max(max, input[y][x]);
    }
}

ひとつのループの中で、最小値と最大値の両方を求めています。
もし最小値を求める処理とは別に最大値を求める処理を作成し、minを求めた後でmaxを求める処理を実行した場合、単純計算で2倍の処理時間がかかることになり非効率です。

ひとつのループの中で複数の値を求める方法としてHalide::Tupleがあります。

class MinMax : public Halide::Generator<MinMax> {
public:
    Halide::GeneratorInput<Halide::Buffer<std::uint8_t>> input{"input", 2};
    Halide::GeneratorOutput<Halide::Func> output{"output"};

    Halide::RDom rd;

    void generate() {
        // 初期値定義
        output() = Halide::Tuple(
            input.type().max(),
            input.type().min()
	);
        auto old = output();
        auto old_min = old[0];
        auto old_max = old[1];
        
        // 値更新
        rd = Halide::RDom(0, input.width(), 0, input.height());
        auto value = input(rd.x, rd.y);
        output() = Halide::Tuple(
            Halide::min(old_min, value),
           Halide::max(old_max, value)
        );
    }
};

Halideに処理を生成させると次のような3引数をとる関数が定義されます。

int minmax(struct halide_buffer_t *_input_buffer, struct halide_buffer_t *_output_0_buffer, struct halide_buffer_t *_output_1_buffer);

第二引数と第三引数が_output_0_buffer, _output_1_bufferとなっています。
この01Halide::Tupleのインデックスに対応します。
すなわち、関数の第二引数_output_0_bufferに最小値、関数の第三引数_output_1_bufferに最大値の結果が格納されます。

呼び出し例:

Halide::Runtime::Buffer<std::uint8_t> input(image, width, height);
auto min = Halide::Runtime::Buffer<std::uint8_t>::make_scalar();
auto max = Halide::Runtime::Buffer<std::uint8_t>::make_scalar();
minmax(input, min, max);
auto result_min = min();
auto result_max = max();

ヒストグラム

これまでは2次元のデータを0次元に要約する処理でした。
さて、ヒストグラムのような1次元の要約はどのように記述したら良いでしょうか。

ヒストグラム算出を通常のC++で記述すると次のようになると思います。

std::vector<int> histogram(256, 0);
for (int y=0; y < height; ++y) {
    for (int x=0; x < width; ++x) {
        histogram[image[y][x]] += 1;
    }
}

これをHalideに移植してみましょう。

class Histogram : public Halide::Generator<Histogram> {
public:
    Halide::GeneratorInput<Halide::Buffer<std::uint8_t>> input{"input", 2};
    Halide::GeneratorOutput<Halide::Buffer<int>> output{"output", 1};

    Halide::Var i{"i"};

    Halide::RDom rd;

    void generate() {
        // 初期値定義
        output(i) = 0;
        
        // 値更新
        rd = Halide::RDom(0, input.width(), 0, input.height());
        auto value = input(rd.x, rd.y);
        output(value) += 1;
    }
};

上記コードは冒頭に書いたC++のコードと同等の処理を生成します。

ところで、この処理を並列処理することはできるでしょうか?

もし、次のようなコードでyを並列に実行した場合、histogramの結果は正しいものになるでしょうか?

std::vector<int> histogram(256, 0);
#pragma omp parallel for
for (int y=0; y < height; ++y) {
    for (int x=0; x < width; ++x) {
        histogram[image[y][x]] += 1;
    }
}

残念ながらこの処理ではhistogramの中身はでたらめな結果になります。
histogram[i] += 1;の処理が同一のiで同時に生じた場合、そこに記録される値はどちらか一方のhistogram[i] += 1;の値となり、一方の+1は無かったことにされてしまうためです。

並列で求めるにはちょっとした工夫を施す必要があります。

class Histogram : public Halide::Generator<Histogram> {
public:
    Halide::GeneratorInput<Halide::Buffer<std::uint8_t>> input{"input", 2};
    Halide::GeneratorOutput<Halide::Buffer<int>> output{"output", 1};

    Halide::Var i{"i"}, y{"y"};

    Halide::Func hist_rows{"hist_rows"};

    Halide::RDom rx, ry;

    void generate() {
        // 初期値定義
        output(i) = 0;
        hist_rows(i, y) = 0;
        
        // 値更新
        rx = Halide::RDom(0, input.width());
        auto value = input(rx, y);
        hist_rows(value, y) += 1;
        ry = Halide::RDom(0, input.height());
        output(i) += hist_rows(i, ry);
    }

    void schedule() {
        output.compute_root().parallel(i);
        hist_rows.compute_root().parallel(y);
    }
};

上記Halideコードは行ごとに並列でヒストグラムを求め、最後にインデックスごとに並列で合算を行います。

参考