Zenn
⚗️

Halideによる高速演算のすすめ - 基礎編

2020/11/08に公開
C++
#
Halide
tech

Halideとは

Halideは主に画像処理に用いられるドメイン固有言語です。
並列・ベクトル演算といった高速演算手法を、C++をベースに抽象化されたアルゴリズムを組み、実行先の環境を指定するだけで処理を自動生成することができます。
一つのHalideコードから異なるアーキテクチャ向けの処理を生成できるため、低コストで多様なデバイスに展開することが可能になります。
また、Halideは異なるOS向けのバイナリを生成することが可能です。Linux上でWindowsやMac OSX向けのライブラリを生成できるため、CI等で生成したライブラリをアプリケーション側でリンクして呼び出すといった運用が可能です。

ジェネレーターパターン

Halideの優れている点はアルゴリズムの記述と高速化手法の設定を分離しているところにあります。
従来、高速なコードを記述するには、使用する手法に合わせてコードやデータのメモリ配置を適切に組み替える必要がありましたが、Halideではアルゴリズムと計算手法をそれぞれ与えるだけで処理が適切に自動生成されるため、複雑なコードの組み換えが不要になります。

このアルゴリズムとスケジューリングの大枠を理解するのに、Halideのジェネレーターパターンというのが便利なため、ここでは基本的にジェネレーターパターンで解説を行います。
（LLVMを使用して実行時に動的に処理を生成するパターンもありますが、ここでは扱いません。）

#include <Halide.h>


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

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

    void generate() {
        output(x, y) = input(x, y) + 2;
    }
};


HALIDE_REGISTER_GENERATOR(Sample, sample);

上記ソースを以下の設定でコンパイルすると、01_sampleという実行バイナリが得られます。

g++ -o 01_sample 01_sample.cpp ${HALIDE_PATH}/share/Halide/tools/GenGen.cpp -I ${HALIDE_PATH}/include/ -lHalide -L ${HALIDE_PATH}/lib/ -std=c++17 -Wl,-rpath ${HALIDE_PATH}/lib/

この実行バイナリを次のように呼び出すと、処理が自動生成されます。

./01_sample -g sample target=host -o .
# sample.a  sample.h  sample.registration.cpp

デフォルトでは、静的ライブラリ（sample.a）とそれに対応するヘッダー（sample.h）が生成されます。
他にもオブジェクトファイル、アセンブリ形式でも出力することができるため、既存のシステムに高速化した処理を組み込むことが容易です。

また、C++コードとして出力することもできるので生成された処理が思うように動かなかった場合は、デバッガでステップ実行して確認することも可能です。

生成された処理の呼び出し

さて、自動生成された処理は基本的に次のようなCスタイルの関数として呼び出すことになります。

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

#include <cstdint>
#include <iostream>
#include <numeric>
#include <vector>

// halide_runtime.a のみで動くバッファ
#include <HalideBuffer.h>

#include "sample.h"

int main() {
    const int width = 16;
    const int height = 8;
    const int total = width * height;
    std::vector<std::uint8_t> source(total, 0);
    // [0, 1, 2, ..., 126, 127]
    std::iota(source.begin(), source.end(), 0);
    
    std::vector<std::uint8_t> result(total, 0);
    
    Halide::Runtime::Buffer<std::uint8_t> input(&source[0], width, height);
    Halide::Runtime::Buffer<std::uint8_t> output(&result[0], width, height);
    
    // Halideで作成した sample関数を呼ぶ
    sample(input, output);
    
    bool ok = true;
    for (int i=0; i < total; ++i) {
        if ( result[i] != (i + 2) % 0x100 ) {
            ok = false;
            break;
        }
    }
    
    std::cout << (ok ? "correct!" : "incorrect...") << std::endl;
    
    return 0;
}

上記コードを以下のオプションでコンパイルして実行すると、Halideで生成した処理が正しく動いていることを確認できます。

g++ -o 01_run_sample 01_run_sample.cpp -std=c++17 -I ${HALIDE_PATH}/include/ -l:sample.a -L. -pthread -ldl

スケジューリング

さて、先述のsample関数はアルゴリズムのみを記述しただけの状態で、高速化は行っていません。

HalideがSample::generate()から生成した処理は以下の処理と同等のものになります。

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

並列演算

一番簡単な高速化の方法として並列演算が考えられます。

OpenMPを使用した場合、次のようなコードになるでしょう。

#pragma omp parallel for
for (int y=0; y < height; ++y) {
    for (int x=0; x < width; ++x) {
        output[y][x] = input[y][x] + 2;
    }
}

これをHalideで行う場合、ジェネレーターを次のように書き換えます。

#include <Halide.h>


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

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

    void generate() {
        output(x, y) = input(x, y) + 2;
    }

    // スケジュール追加
    void schedule() {
        output.compute_root();
        output.parallel(y);
    }
};


HALIDE_REGISTER_GENERATOR(Sample, sample);

output.parallel(y)という一文で、yのループについて並列で演算する処理を生成します。

ここまではOpenMPとHalideで大した手間の差はありません。

ベクトル演算

さて、今度はベクトル演算を組み込んでみましょう。

SIMD命令を使って、一回に複数個処理を適用します。

#pragma omp parallel for
for (int y=0; y < height; ++y) {
    __m128i a = _mm_set1_epi8(2);
    for (int x=0; x < width; x+=16) {
        //output[y][x] = input[y][x] + 2;
        __m128i b = _mm_load_si128(reinterpret_cast<__m128i *>(&input[y*width+x]));
        __m128i c = _mm_add_epi8(a, b);
        _mm_store_si128(reinterpret_cast<__m128i *>(&output[y*width+x]), c);
    }
}

もはや output = input + 2 という基本的なアルゴリズムの見る陰もありません。
一見動きそうなこのコード、動きはするものの、実のところ使い物になりません。

SIMD命令を使用する場合、メモリのアライメントを考慮する必要があります。
ロードされるメモリのアドレスは16の倍数でなければならず、newで確保したメモリ（std::vector等デフォルトのアロケータが確保したメモリも同）をそのまま渡すと十中八九クラッシュします。
また、widthを16の倍数でなくした途端、y > 0の条件で、input[y*width+x]のアドレスは16の倍数でなくなるため、クラッシュするようになります。

実用的に動かすためには、メモリ配置のレベルから適切な知識を持って記述する必要があります。
少なくとも上記コード以上の規模の修正が必要となるでしょう。
これが既存のコードの高速化を考えていた場合、非常に頭の痛くなる要素のひとつになります。

それではHalideを使った場合はどうなるでしょうか。

#include <Halide.h>


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

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

    void generate() {
        output(x, y) = input(x, y) + 2;
    }

    void schedule() {
        output.compute_root();
        // ベクトル演算追加
        const int vector_length = 16;
        output.vectorize(x, vector_length).parallel(y);
    }
};


HALIDE_REGISTER_GENERATOR(Sample, sample);

vectorize(x, vector_length)と記述するだけで、ベクトル演算を適用することができます。
処理を呼び出す場合もHalide::Runtime::Bufferに渡すアドレスのアライメントを意識する必要はありません。アライメント周りはHalideが適切に行います（ベクトル化による副作用の問題は頭の片隅に置いておくと生成された処理の予期せぬ挙動の原因に気づくことができます。しかし、複雑なスケジュールでもない限りは適切に処理が生成されるため、ほとんど困ることはないでしょう）。
そのため、既存のコードのメモリ配置等を書き換えることなく、処理だけ高速なものに差し替えることが可能になります。

アルゴリズムとスケジュール

さて、Halideのスケジューリングを変えても、generate()の内容に一切手を付けていないことにお気づきでしょうか。

アルゴリズムとスケジュールが分離されていることで、スケジュールを組み替えながら処理の最適化を容易に行えるようになります。

例えば、画像処理をタイル化する場合、次のように記述します。

void schedule() {
    output.compute_root();
    const int tile_x = 32;
    const int tile_y = 32;
    Halide::Var x_outer, x_inner, y_outer, y_inner;
    output.tile(x, y, x_outer, y_outer, x_inner, y_inner, tile_x, tile_y).vectorize(x_inner, tile_x).parallel(y);
}

これは以下のようなループネストを意味します。

for y.y_outer in [0, height/32]:
  for x.x_outer in [0, width/32]:
    for y.y_inner in [0, 31]:
      for x.x_inner in [0, 31]:
        output(...) = ...

このようなループの組み換えを手作業で行う場合、非常に根気のいる作業になりますが、Halideであればスケジュールの組み換えだけで済みます。
つまり、最適なスケジュールを探索するための処理の組み換えが低コストでできるようになるのです。

まとめ

Halideを用いることで、アルゴリズムとスケジュールを分離し、処理の組み換えを容易に行えます。
これにより、従来専門的な知識を必要とした高速化手法の適用を低コストで行えるようになります。

また、この組み換えはアーキテクチャの異なる環境向けにスケジュールを調整できるメリットを生じます。
例えば、論理十数コアのCPU SIMD命令と、数千コアのGPU演算とで、アルゴリズムは同一のまま、スケジュールだけアーキテクチャにとって最適なものに組み替えることが可能です。

Halideで扱える演算は0〜4次元のデータ構造になります。
これは画像処理だけでなく、物量が物を言うタイプのデータ加工においても活用できることを意味します。

Halideを用いた処理の記述パターンを取り上げながら、一般的な処理をHalideに落とし込む手法について解説していきます。

参考

Discussion