C++ Func::chunk方法代码示例

Func process(Func raw, Type result_type,
             UniformImage matrix_3200, UniformImage matrix_7000, Uniform<float> color_temp, 
             Uniform<float> gamma, Uniform<float> contrast) {

    Func processed("processed");
    Var xi, yi;

    Func denoised = hot_pixel_suppression(raw);
    Func deinterleaved = deinterleave(denoised);
    Func demosaiced = demosaic(deinterleaved);
    Func corrected = color_correct(demosaiced, matrix_3200, matrix_7000, color_temp);
    Func curved = apply_curve(corrected, result_type, gamma, contrast);

    // Schedule
    Var co, ci;
    //#define USE_CI_HACK
    #ifndef USE_CI_HACK
    ci = c;
    #endif
    processed(tx, ty, c) = curved(tx, ty, ci);
    #ifdef USE_CI_HACK
    processed.split(c, co, ci, 3); // bound color loop to 0-3
    #else
    processed.bound(c, 0, 3); // bound color loop 0-3, properly
    #endif
    if (schedule == 0) {
        // Compute in chunks over tiles, vectorized by 8
        denoised.chunk(tx).vectorize(x, 8);
        deinterleaved.chunk(tx).vectorize(x, 8);
        corrected.chunk(tx).vectorize(x, 4);
        processed.tile(tx, ty, xi, yi, 32, 32).reorder(xi, yi, ci, tx, ty);
        processed.parallel(ty);
    } else if (schedule == 1) {
        // Same as above, but don't vectorize (sse is bad at interleaved 16-bit ops)
        denoised.chunk(tx);
        deinterleaved.chunk(tx);
        corrected.chunk(tx);
        processed.tile(tx, ty, xi, yi, 128, 128).reorder(xi, yi, ci, tx, ty);
        processed.parallel(ty);
    } else {
        denoised.root();
        deinterleaved.root();
        corrected.root();
        processed.root();
    }

    return processed;
}

//.........这里部分代码省略.........
    b_gr(x, y) = correction + (b_b(x, y) + b_b(x, y-1))/2;

    correction = g_gb(x, y) - (g_r(x, y) + g_r(x, y+1))/2;
    r_gb(x, y) = correction + (r_r(x, y) + r_r(x, y+1))/2;

    correction = g_gb(x, y) - (g_b(x, y) + g_b(x+1, y))/2;
    b_gb(x, y) = correction + (b_b(x, y) + b_b(x+1, y))/2;

    // Now interpolate diagonally to get red at blue and blue at
    // red. Hold onto your hats; this gets really fancy. We do the
    // same thing as for interpolating green where we try both
    // directions (in this case the positive and negative diagonals),
    // and use the one with the lowest absolute difference. But we
    // also use the same trick as interpolating red and blue at green
    // sites - we correct our interpolations using the second
    // derivative of green at the same sites.

    correction = g_b(x, y)  - (g_r(x, y) + g_r(x-1, y+1))/2;
    Expr rp_b  = correction + (r_r(x, y) + r_r(x-1, y+1))/2;
    Expr rpd_b = abs(r_r(x, y) - r_r(x-1, y+1));

    correction = g_b(x, y)  - (g_r(x-1, y) + g_r(x, y+1))/2;
    Expr rn_b  = correction + (r_r(x-1, y) + r_r(x, y+1))/2;
    Expr rnd_b = abs(r_r(x-1, y) - r_r(x, y+1));

    r_b(x, y)  = select(rpd_b < rnd_b, rp_b, rn_b);


    // Same thing for blue at red
    correction = g_r(x, y)  - (g_b(x, y) + g_b(x+1, y-1))/2;
    Expr bp_r  = correction + (b_b(x, y) + b_b(x+1, y-1))/2;
    Expr bpd_r = abs(b_b(x, y) - b_b(x+1, y-1));

    correction = g_r(x, y)  - (g_b(x+1, y) + g_b(x, y-1))/2;
    Expr bn_r  = correction + (b_b(x+1, y) + b_b(x, y-1))/2;
    Expr bnd_r = abs(b_b(x+1, y) - b_b(x, y-1));

    b_r(x, y)  =  select(bpd_r < bnd_r, bp_r, bn_r);    

    // Interleave the resulting channels
    Func r = interleave_y(interleave_x(r_gr, r_r),
                          interleave_x(r_b, r_gb));
    Func g = interleave_y(interleave_x(g_gr, g_r),
                          interleave_x(g_b, g_gb));
    Func b = interleave_y(interleave_x(b_gr, b_r),
                          interleave_x(b_b, b_gb));


    Func output;
    output(x, y, c) = select(c == 0, r(x, y), 
                             select(c == 1, g(x, y), b(x, y)));


    /* THE SCHEDULE */    
    if (schedule == 0) {
        // optimized for ARM
        // Compute these in chunks over tiles, vectorized by 8
        g_r.chunk(tx).vectorize(x, 8);
        g_b.chunk(tx).vectorize(x, 8);
        r_gr.chunk(tx).vectorize(x, 8);
        b_gr.chunk(tx).vectorize(x, 8);
        r_gb.chunk(tx).vectorize(x, 8);
        b_gb.chunk(tx).vectorize(x, 8);
        r_b.chunk(tx).vectorize(x, 8);
        b_r.chunk(tx).vectorize(x, 8);
        // These interleave in y, so unrolling them in y helps
        r.chunk(tx).vectorize(x, 8).unroll(y, 2);
        g.chunk(tx).vectorize(x, 8).unroll(y, 2);
        b.chunk(tx).vectorize(x, 8).unroll(y, 2);
    } else if (schedule == 1) {
        // optimized for X86
        // Don't vectorize, because sse is bad at 16-bit interleaving
        g_r.chunk(tx);
        g_b.chunk(tx);
        r_gr.chunk(tx);
        b_gr.chunk(tx);
        r_gb.chunk(tx);
        b_gb.chunk(tx);
        r_b.chunk(tx);
        b_r.chunk(tx);
        // These interleave in x and y, so unrolling them helps
        r.chunk(tx).unroll(x, 2).unroll(y, 2);
        g.chunk(tx).unroll(x, 2).unroll(y, 2);
        b.chunk(tx).unroll(x, 2).unroll(y, 2);
    } else {
        // Basic naive schedule
        g_r.root();
        g_b.root();
        r_gr.root();
        b_gr.root();
        r_gb.root();
        b_gb.root();
        r_b.root();
        b_r.root();
        r.root();
        g.root();
        b.root();
    }
    return output;
}