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

#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);
    }
};

./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

./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

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);
    }
}