Zenn
⚗️

Halideによる高速演算のすすめ - オートスケジューラ編

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

オートスケジューラ

Halideは処理をアルゴリズムとスケジュールに分けて記述します。
スケジュールを組み替えながら最適な処理順序を探索することができますが、ジェネレータでパラメータを調整しながらひとつひとつ確認していく作業は骨の折れる作業です。

慣れてくれば大体のスケジュールのパターンは見えてきますが、少し複雑なアルゴリズムになった時はスケジュールの組み換えですら難儀することがあります。

そんなスケジュールの生成をある程度の性能で自動で行ってくれるオートスケジューラというものが追加されました。
オートスケジューラが組み立てたスケジュールを出力できるので、ひとまずオートスケジューラで処理を生成し、使いながら調整していくこともできます。

使い方

オートスケジューラは従来のジェネレータに以下の2ステップを加えるだけで使うことができます。

  1. 入出力の概算値の設定
  2. ジェネレータの引数にオートスケジューラを指定

1. 入出力の概算値の設定

従来スケジュールを記述していたところで入出力の概算値を設定します。
Halide::Bufferの入出力にはset_estimates()を使用して、Halide::RDomでいうところの(開始値, 要素数)の組を指定します。
スカラー値の入出力にはset_estimate()を使用して、実際に引数に渡されるであろう具体的な値を設定します。

#define _USE_MATH_DEFINES
#include <cmath>

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

    Halide::RDom rd;

    Halide::Func data{"data"};
    Halide::Func kernel{"kernel"};
    Halide::Func blur1{"blur1"};
    Halide::Func blur2{"blur2"};

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

    void generate() {
        auto clamped = Halide::BoundaryConditions::repeat_edge(input);
	data(x, y) = Halide::cast<float>(clamped(x, y));

        auto size = radius * 2 + 1;
        const float PI = M_PI;
        const float sigma = 1.5f;
	kernel(w) = Halide::exp(
	    -w * w / (2 * sigma * sigma)
	) / (std::sqrt(2 * PI) * sigma);

        rd = Halide::RDom(-radius, size);
        blur1(x, y) = Halide::sum(kernel(rd) * data(x + rd, y));
	blur2(x, y) = Halide::sum(kernel(rd) * blur1(x, y + rd));

        output(x, y) = Halide::cast<std::uint8_t>(blur2(x, y));
    }

    void schedule() {
        input.set_estimates({{0, 6000}, {0, 4000}});
        output.set_estimates({{0, 6000}, {0, 4000}});
	radius.set_estimate(3);
    }
};

2. ジェネレータの引数にオートスケジューラを指定

./gaussian_blur -g gaussian_blur target=host auto_schedule=true -p ${HALIDE_PATH}/lib/libautoschedule_mullapudi2016.so -s Mullapudi2016 -o . -e h,static_library,schedule

ここで注目すべきは以下の2点です。

  • auto_schedule=true
  • -p ${HALIDE_PATH}/lib/libautoschedule_mullapudi2016.so -s Mullapudi2016

まずauto_schedule=trueで生成時にオートスケジューラを有効にします。
その際、-sオプションで使用するオートスケジューラを設定する必要があります。

Halide 10.0.0では以下のオートスケジュールプラグインが付随しています（lib/参照）。

  • libautoschedule_mullapudi2016.so : Mullapudi2016
    • 最初期のオートスケジューラ。単一のスケジュールが生成されます。
  • libautoschedule_li2018.so : Li2018
  • libautoschedule_adams2019.so : Adams2019
    • 5種類のスケジュールを作成し、実行コストを比較し最良のものを生成します。

プラグインライブラリを-pオプションで指定した後、-sで使用するオートスケジューラを設定します。

以上の手順でスケジュールの自動生成が行われます。

生成スケジュールの確認

ジェネレータに自動生成された実際のスケジュールを出力させることができます。

./gaussian_blur -g gaussian_blur target=host auto_schedule=true -p ${HALIDE_PATH}/lib/libautoschedule_mullapudi2016.so -s Mullapudi2016 -o . -e h,static_library,schedule

ジェネレータの引数に-eオプションで生成フォーマットを指定する際、scheduleを指定すると、{func_name}.schedule.hというファイルが生成されます。

先のHalideコードをMullapudi2016で自動生成させたスケジュール結果が以下のようになりました。

  • 画像を128x128にタイル化
    • X方向にベクトル化
    • Y方向に並列化
  • ぼかし処理をベクトル化

生成されたスケジュール：

inline void apply_schedule_gaussian_blur(
    ::Halide::Pipeline pipeline,
    ::Halide::Target target
) {
    using ::Halide::Func;
    using ::Halide::MemoryType;
    using ::Halide::RVar;
    using ::Halide::TailStrategy;
    using ::Halide::Var;
    Var x_i("x_i");
    Var x_i_vi("x_i_vi");
    Var x_i_vo("x_i_vo");
    Var x_o("x_o");
    Var x_vi("x_vi");
    Var x_vo("x_vo");
    Var y_i("y_i");
    Var y_o("y_o");

    Func kernel = pipeline.get_func(4);
    Func output = pipeline.get_func(9);
    Func sum = pipeline.get_func(5);
    Func sum_1 = pipeline.get_func(7);

    {
        Var w = kernel.args()[0];
        kernel
            .compute_root()
            .parallel(w);
    }
    {
        Var x = output.args()[0];
        Var y = output.args()[1];
        output
            .compute_root()
            .split(x, x_o, x_i, 128)
            .split(y, y_o, y_i, 128)
            .reorder(x_i, y_i, x_o, y_o)
            .split(x_i, x_i_vo, x_i_vi, 32)
            .vectorize(x_i_vi)
            .parallel(y_o);
    }
    {
        Var x = sum.args()[0];
        sum
            .compute_at(output, x_o)
            .split(x, x_vo, x_vi, 8)
            .vectorize(x_vi);
        sum.update(0)
            .split(x, x_vo, x_vi, 8, TailStrategy::GuardWithIf)
            .vectorize(x_vi);
    }
    {
        Var x = sum_1.args()[0];
        RVar r17$x(sum_1.update(0).get_schedule().rvars()[0].var);
        sum_1
            .compute_at(output, x_o)
            .split(x, x_vo, x_vi, 8)
            .vectorize(x_vi);
        sum_1.update(0)
            .reorder(x, r17$x, y)
            .split(x, x_vo, x_vi, 8, TailStrategy::GuardWithIf)
            .vectorize(x_vi);
    }
}

トラブルシューティング

No autoscheduler has been run for this Generator.

{proc_name}.schedule.hに// No autoscheduler has been run for this Generator.と表示された場合、ジェネレータにauto_schedule=trueオプションを設定していない場合があります。

Error: Need an estimate on dimension 0 of "output"

Halide::GeneratorInput<Halide::Buffer<T>>などに概算値が設定されていません。
スケジュールの指定の代わりに概算値を設定する必要があります。
「1. 入出力の概算値の設定」を参照して下さい。

Condition failed: expr.defined(): Missing estimate for hogehoge

Halide::GeneratorInput<T>などの入力に概算値が設定されていません。
スケジュール指定の代わりに概算値を設定する必要があります。
「1. 入出力の概算値の設定」を参照して下さい。

Error: AutoSchedule: cannot auto-schedule function "func" since it is scheduled to be computed at root等

  • Error: AutoSchedule: cannot auto-schedule function "func" since it is scheduled to be computed at root
  • Error: AutoSchedule: cannot auto-schedule function "func" since stage 0 is not serial at dim i

Halide::Funcなどの途中の演算で、既にcompute_atなどでスケジュールが設定されています。
生成時のヒントとなるスケジュール以外の指定は含めずにコンパイルすることで解消することがあります。

参考

Discussion