How to schedule producer for multiple consumers

[zhouwangpei123]

Hi i'm new to halide, recently i'm trying to realize a variant of guided filter algorithm with halide. Tested on Intel Core i7-7700( the size of input/output is 2040*1536 ), halide_gf_uint16_withmap takes 8ms， which is resonable. Howerver, halide_gfcoeff_uint16_withmap has the exact same flow as halide_gf_uint16_withmap, the only difference is that halide_gfcoeff_uint16_withmap only calculate the coefficients of a and b, using similar schedule as halide_gf_uint16_withmap, but it takes 234ms！ what's the problem here? how shound i schedule for halide_gfcoeff_uint16_withmap?

#include "Halide.h"

using namespace Halide;

class GuidedFilterCoeffWithMap : public Halide::Generator<GuidedFilterCoeffWithMap> {

public:

    Input<Buffer<uint16_t, 2>> input {"input"};
    Input<Buffer<uint8_t,  2>> map{"map"};
    Input<int32_t> wsz{"window_size"};
    Input<float_t> eps{"epsilon"};
    Output<Buffer<uint16_t, 2>> output_a{"output_a"};
    Output<Buffer<uint16_t, 2>> output_b{"output_b"};

    void generate() {
        Expr hwsz = wsz >> 1;
        Expr inv = 1.f / ((2*hwsz+1) * (2*hwsz+1));
        Expr e2 = eps * eps;

        input_clamped = BoundaryConditions::repeat_edge(input);
        map_clamped = BoundaryConditions::repeat_edge(map);

        RDom r1(-hwsz, wsz);
        RDom r2(-hwsz, wsz);

        output_a.dim(0).set_bounds(input.dim(0).min(), input.dim(0).extent());
        output_b.dim(0).set_bounds(input.dim(0).min(), input.dim(0).extent());
        output_a.dim(1).set_bounds(input.dim(1).min(), input.dim(1).extent());
        output_b.dim(1).set_bounds(input.dim(1).min(), input.dim(1).extent());

        sum_x(x, y) = sum(cast(Int(32), input_clamped(x+r1[0], y)));
        sum_xy(x, y) = sum(sum_x(x, y+r2[0]));

        square_sum_x(x, y) = sum(pow(cast(Int(32), input_clamped(x+r1[0], y)), 2));
        square_sum_xy(x, y) = sum(square_sum_x(x, y+r2[0]));

        Expr var_;
        average(x, y)  = cast(Int(32), cast(Float(32), sum_xy(x, y)) * inv);
        var_ = cast(Float(32), square_sum_xy(x, y)) * inv - fast_pow(average(x, y), 2);
        variance(x, y)  = cast(Int(32), 1024.f * var_ / (var_ + fast_pow(map_clamped(x, y), 2) * e2));

        w(x, y) = max((1024 - variance(x, y))>>2, 8);

        w_sum_x(x, y) = sum(w(x+r1[0], y));
        w_sum_xy(x, y) = sum(w_sum_x(x, y+r2[0]));

        wave_sum_x(x, y) = sum(w(x+r1[0], y) * average(x+r1[0], y));
        wave_sum_xy(x, y) = sum(wave_sum_x(x, y+r2[0]));

        wvar_sum_x(x, y) = sum(w(x+r1[0], y) * variance(x+r1[0], y));
        wvar_sum_xy(x, y) = sum(wvar_sum_x(x, y+r2[0]));

        output_b(x, y) = cast(UInt(16), wave_sum_xy(x, y) / w_sum_xy(x, y));
        output_a(x, y) = cast(UInt(16), wvar_sum_xy(x, y) / w_sum_xy(x, y));
    }

    void schedule() {
        if (!auto_schedule) {

            if (get_target().has_gpu_feature()) {
                //TODO

            } else {
                const int vec = 8;

                /*Schedule 1: 234ms on Intel(R) Core(TM) i7-7700 CPU @ 3.60GHz*/
                output_a
                    .compute_root()
                    .tile(x, y, xo, yo, xi, yi, 512, 32)
                    .fuse(xo, yo, t)
                    .parallel(t)
                    .vectorize(xi, vec);
                output_b
                    .compute_root()
                    .tile(x, y, xo, yo, xi, yi, 512, 32)
                    .fuse(xo, yo, t)
                    .parallel(t)
                    .vectorize(xi, vec);
                output_b.compute_with(output_a, t);

                wvar_sum_x
                    .compute_at(output_a, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                wave_sum_x
                    .compute_at(output_a, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                w_sum_x
                    .compute_at(output_a, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                variance
                    .compute_at(output_a, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                square_sum_x
                    .compute_at(output_a, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                sum_x
                    .compute_at(output_a, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                /*Schedule 2: 70ms on Intel(R) Core(TM) i7-7700 CPU @ 3.60GHz*/
                // output_a
                //     .compute_root()
                //     .tile(x, y, xo, yo, xi, yi, 512, 32)
                //     .fuse(xo, yo, t)
                //     .parallel(t)
                //     .vectorize(xi, vec);
                // output_b
                //     .compute_root()
                //     .tile(x, y, xo, yo, xi, yi, 512, 32)
                //     .fuse(xo, yo, t)
                //     .parallel(t)
                //     .vectorize(xi, vec);
                // output_b.compute_with(output_a, t);

                // wvar_sum_x
                //     .compute_at(output_a, t)
                //     .split(x, xo, xi, 8)
                //     .vectorize(xi, vec);

                // wave_sum_x
                //     .compute_at(output_b, t)
                //     .split(x, xo, xi, 8)
                //     .vectorize(xi, vec);

                // w_sum_x
                //     .compute_root()
                //     .tile(x, y, xo, yo, xi, yi, 512, 32)
                //     .fuse(xo, yo, t)
                //     .parallel(t)
                //     .vectorize(xi, vec);

                // variance
                //     .compute_root()
                //     .tile(x, y, xo, yo, xi, yi, 512, 32)
                //     .fuse(xo, yo, t)
                //     .parallel(t)
                //     .vectorize(xi, vec);

                // square_sum_x
                //     .compute_at(variance, t)
                //     .split(x, xo, xi, 8)
                //     .vectorize(xi, vec);

                // sum_x
                //     .compute_root()
                //     .tile(x, y, xo, yo, xi, yi, 512, 32)
                //     .fuse(xo, yo, t)
                //     .parallel(t)
                //     .vectorize(xi, vec);               

            }

        }

        // Estimates (for autoscheduler; ignored otherwise)
        input.dim(0).set_estimate(0, 4080);
        input.dim(1).set_estimate(0, 3072);
        map.dim(0).set_estimate(0, 4080);
        map.dim(1).set_estimate(0, 3072);
        wsz.set_estimate(3);
        eps.set_estimate(0.7f);
        output_a.dim(0).set_estimate(0, 4080);
        output_a.dim(1).set_estimate(0, 3072);
        output_b.dim(0).set_estimate(0, 4080);
        output_b.dim(1).set_estimate(0, 3072);
    }

private:
    Var x, y, xo, yo, xi, yi, t, xi_i, xi_o;
    Func sum_x{"sum_x"}, sum_xy{"sum_xy"};
    Func square_sum_x{"square_sum_x"}, square_sum_xy{"square_sum_xy"};
    Func w{"w"}, w_sum_x{"w_sum_x"}, w_sum_xy{"w_sum_xy"};
    Func wave_sum_x{"wave_sum_x"}, wave_sum_xy{"wave_sum_xy"};
    Func wvar_sum_x{"wvar_sum_x"}, wvar_sum_xy{"wvar_sum_xy"};
    Func average{"average"}, variance{"variance"};
    Func input_clamped{"input_clamped"}, map_clamped{"map_clamped"};
};

HALIDE_REGISTER_GENERATOR(GuidedFilterCoeffWithMap, halide_gfcoeff_uint16_withmap)


class GuidedFilterWithMap : public Halide::Generator<GuidedFilterWithMap> {

public:

    Input<Buffer<uint16_t, 2>> input {"input"};
    Input<Buffer<uint8_t,  2>> map{"map"};
    Input<int32_t> wsz{"window_size"};
    Input<float_t> eps{"epsilon"};
    Output<Buffer<uint16_t, 2>> output{"output"};

    void generate() {
        Expr hwsz = wsz >> 1;
        Expr inv = 1.f / ((2*hwsz+1) * (2*hwsz+1));
        Expr e2 = eps * eps;

        input_clamped = BoundaryConditions::repeat_edge(input);
        map_clamped = BoundaryConditions::repeat_edge(map);

        RDom r1(-hwsz, wsz);
        RDom r2(-hwsz, wsz);

        sum_x(x, y) = sum(cast(Int(32), input_clamped(x+r1[0], y)));
        sum_xy(x, y) = sum(sum_x(x, y+r2[0]));

        square_sum_x(x, y) = sum(pow(cast(Int(32), input_clamped(x+r1[0], y)), 2));
        square_sum_xy(x, y) = sum(square_sum_x(x, y+r2[0]));

        Expr var_;
        average(x, y)  = cast(Int(32), cast(Float(32), sum_xy(x, y)) * inv);
        var_ = cast(Float(32), square_sum_xy(x, y)) * inv - fast_pow(average(x, y), 2);
        variance(x, y)  = cast(Int(32), 1024.f * var_ / (var_ + fast_pow(map_clamped(x, y), 2) * e2));

        w(x, y) = max((1024 - variance(x, y))>>2, 8);

        w_sum_x(x, y) = sum(w(x+r1[0], y));
        w_sum_xy(x, y) = sum(w_sum_x(x, y+r2[0]));

        wave_sum_x(x, y) = sum(w(x+r1[0], y) * average(x+r1[0], y));
        wave_sum_xy(x, y) = sum(wave_sum_x(x, y+r2[0]));

        wvar_sum_x(x, y) = sum(w(x+r1[0], y) * variance(x+r1[0], y));
        wvar_sum_xy(x, y) = sum(wvar_sum_x(x, y+r2[0]));

        bimg(x, y) = wave_sum_xy(x, y) / w_sum_xy(x, y);
        aimg(x, y) = wvar_sum_xy(x, y) / w_sum_xy(x, y);

        output(x, y) = cast(UInt(16), clamp((aimg(x, y)*input_clamped(x, y) + (1024-aimg(x, y))*bimg(x,y) + (1<<9)) >> 10, 0, 32767));
    }

    /* 8ms on Intel(R) Core(TM) i7-7700 CPU @ 3.60GHz*/
    void schedule() {
        if (!auto_schedule) {

            if (get_target().has_gpu_feature()) {
                //TODO

            } else {
                const int vec = 8;
                output
                    .compute_root()
                    .tile(x, y, xo, yo, xi, yi, 128, 32)
                    .fuse(xo, yo, t)
                    .parallel(t)
                    .vectorize(xi, vec);

                sum_x
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);
                square_sum_x
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                w_sum_x
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);
                wave_sum_x
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);
                wvar_sum_x
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);

                average
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);
                variance
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);
                w
                    .compute_at(output, t)
                    .split(x, xo, xi, 8)
                    .vectorize(xi, vec);
            }

        }

        // Estimates (for autoscheduler; ignored otherwise)
        input.dim(0).set_estimate(0, 4080);
        input.dim(1).set_estimate(0, 3072);
        map.dim(0).set_estimate(0, 4080);
        map.dim(1).set_estimate(0, 3072);
        wsz.set_estimate(3);
        eps.set_estimate(0.7f);
        output.dim(0).set_estimate(0, 4080);
        output.dim(1).set_estimate(0, 3072);
    }

private:
    Var x, y, xo, yo, xi, yi, t, xi_i, xi_o;
    Func sum_x{"sum_x"}, sum_xy{"sum_xy"};
    Func square_sum_x{"square_sum_x"}, square_sum_xy{"square_sum_xy"};
    Func w{"w"}, w_sum_x{"w_sum_x"}, w_sum_xy{"w_sum_xy"};
    Func wave_sum_x{"wave_sum_x"}, wave_sum_xy{"wave_sum_xy"};
    Func wvar_sum_x{"wvar_sum_x"}, wvar_sum_xy{"wvar_sum_xy"};
    Func aimg{"aimg"}, bimg{"bimg"};
    Func average{"average"}, variance{"variance"};
    Func input_clamped{"input_clamped"}, map_clamped{"map_clamped"};
};

HALIDE_REGISTER_GENERATOR(GuidedFilterWithMap, halide_gf_uint16_withmap)