<commit_before>#include <FImage.h>
#include <stdint.h>

using namespace FImage;

int schedule;

Var x("x"), y("y"), tx("tx"), ty("ty"), c("c");

Func hot_pixel_suppression(Func input) {
    Expr a = max(max(input(x-2, y), input(x+2, y)),
                 max(input(x, y-2), input(x, y+2)));
    Expr b = min(min(input(x-2, y), input(x+2, y)),
                 min(input(x, y-2), input(x, y+2)));
    
    Func denoised("denoised");
    denoised(x, y) = clamp(input(x, y), a, b);

    return denoised;
}

Expr abs(Expr e) {
    return select(e < cast(e.type(), 0), -e, e);
}

Func interleave_x(Func a, Func b) {
    Func out;
    out(x, y) = select((x%2)==0, a(x/2, y), b(x/2, y));
    return out;
}

Func interleave_y(Func a, Func b) {
    Func out;
    out(x, y) = select((y%2)==0, a(x, y/2), b(x, y/2));
    return out;
}

Func deinterleave(Func raw) {
    // Deinterleave the color channels
    Func deinterleaved("deinterleaved");
    Var t;
    deinterleaved(x, y, t) = select(t == 0, raw(2*x, 2*y), 
                                    select(t == 1, raw(2*x+1, 2*y), 
                                           select(t == 2, raw(2*x, 2*y+1), raw(2*x+1, 2*y+1))));
    
    return deinterleaved;
}

Func demosaic(Func deinterleaved) {
    // These are the values we already know from the input
    // x_y = the value of channel x at a site in the input of channel y
    // gb refers to green sites in the blue rows
    // gr refers to green sites in the red rows

    // Give more convenient names to the four channels we know
    Func r_r("r_r"), g_gr("g_gr"), g_gb("g_gb"), b_b("b_b");    
    g_gr(x, y) = deinterleaved(x, y, 0);
    r_r(x, y)  = deinterleaved(x, y, 1);
    b_b(x, y)  = deinterleaved(x, y, 2);
    g_gb(x, y) = deinterleaved(x, y, 3);

    // These are the ones we need to interpolate
    Func b_r("b_r"), g_r("g_r");
    Func b_gr("b_gr"), r_gr("r_gr");
    Func b_gb("b_gb"), r_gb("r_gb");
    Func r_b("r_b"), g_b("g_b");

    // First calculate green at the red and blue sites

    // Try interpolating vertically and horizontally. Also compute
    // differences vertically and horizontally. Use interpolation in
    // whichever direction had the smallest difference.
    Expr gv_r  =    (g_gb(x, y-1) + g_gb(x, y))/2;
    Expr gvd_r = abs(g_gb(x, y-1) - g_gb(x, y));
    Expr gh_r  =    (g_gr(x+1, y) + g_gr(x, y))/2;
    Expr ghd_r = abs(g_gr(x+1, y) - g_gr(x, y));

    g_r(x, y)  = select(ghd_r < gvd_r, gh_r, gv_r);

    Expr gv_b  =    (g_gr(x, y+1) + g_gr(x, y))/2;
    Expr gvd_b = abs(g_gr(x, y+1) - g_gr(x, y));
    Expr gh_b  =    (g_gb(x-1, y) + g_gb(x, y))/2;
    Expr ghd_b = abs(g_gb(x-1, y) - g_gb(x, y));

    g_b(x, y)  = select(ghd_b < gvd_b, gh_b, gv_b);

    // Next interpolate red at gr by first interpolating, then
    // correcting using the error green would have had if we had
    // interpolated it in the same way (i.e. add the second derivative
    // of the green channel at the same place).
    Expr correction;
    correction = g_gr(x, y) - (g_r(x, y) + g_r(x-1, y))/2;
    r_gr(x, y) = correction + (r_r(x-1, y) + r_r(x, y))/2;

    // Do the same for other reds and blues at green sites
    correction = g_gr(x, y) - (g_b(x, y) + g_b(x, y-1))/2;
    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("dem");
    output(x, y) = (r(x, y), 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 a 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;
}


Func color_correct(Func input, UniformImage matrix_3200, UniformImage matrix_7000, Uniform<float> kelvin) {
    // Get a color matrix by linearly interpolating between two
    // calibrated matrices using inverse kelvin.

    Func matrix("matrix");
    Expr alpha = (1.0f/kelvin - 1.0f/3200) / (1.0f/7000 - 1.0f/3200);
    Expr val =  (matrix_3200(x, y) * alpha + matrix_7000(x, y) * (1 - alpha));
    matrix(x, y) = cast<int32_t>(val * 256.0f); // Q8.8 fixed point
    matrix.root();

    Func corrected("corrected");
    Expr ir = cast<int32_t>(input(x, y, 0));
    Expr ig = cast<int32_t>(input(x, y, 1));
    Expr ib = cast<int32_t>(input(x, y, 2));

    Expr r = matrix(3, 0) + matrix(0, 0) * ir + matrix(1, 0) * ig + matrix(2, 0) * ib;
    Expr g = matrix(3, 1) + matrix(0, 1) * ir + matrix(1, 1) * ig + matrix(2, 1) * ib;
    Expr b = matrix(3, 2) + matrix(0, 2) * ir + matrix(1, 2) * ig + matrix(2, 2) * ib;

    corrected(x, y) = cast<int16_t>((r, g, b)/256);

    return corrected;
}


Func apply_curve(Func input, Type result_type, Uniform<float> gamma, Uniform<float> contrast) {
    // copied from FCam
    Func curve("curve");

    Expr xf = clamp(cast<float>(x)/1024.0f, 0.0f, 1.0f);
    Expr g = pow(xf, 1.0f/gamma);
    Expr b = 2.0f - pow(2.0f, contrast/100.0f);
    Expr a = 2.0f - 2.0f*b; 
    Expr z = select(g > 0.5f,
                    1.0f - (a*(1.0f-g)*(1.0f-g) + b*(1.0f-g)),
                    a*g*g + b*g);

    Expr val = cast(result_type, clamp(z*256.0f, 0.0f, 255.0f));
    curve(x) = val;
    curve.root(); // It's a LUT, compute it once ahead of time.

    Func curved("curved");
    curved(x, y, c) = curve(input(x, y, c));

    return curved;
}

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("p");
    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; 
    processed(tx, ty, c) = curved(tx, ty, ci);
    processed.split(c, co, ci, 3); // bound color loop to 0-3
    if (schedule == 0) {
        // Compute in chunks over tiles, vectorized by 8
        denoised.chunk(tx).vectorize(x, 8);

        // Unroll the tuple part and compute it in the inside loop
        Var ti, tt = deinterleaved.arg(2);
        deinterleaved.split(tt, tt, ti, 4).unroll(ti).transpose(y, ti).transpose(x, ti);
        deinterleaved.chunk(tx).vectorize(x, 8);

        Var cc = corrected.arg(2);
        corrected.chunk(tx).transpose(y, cc).transpose(x, cc).vectorize(x, 4).unroll(cc, 3);
        processed.tile(tx, ty, xi, yi, 32, 32).transpose(ty, ci).transpose(tx, ci);
        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);

        // Unroll the tuple part and compute it in the inside loop
        Var ti, tt = deinterleaved.arg(2);
        deinterleaved.split(tt, tt, ti, 4).unroll(ti).transpose(y, ti).transpose(x, ti);
        deinterleaved.chunk(tx);

        Var cc = corrected.arg(2);
        corrected.chunk(tx).transpose(y, cc).transpose(x, cc).unroll(cc, 3);
        processed.tile(tx, ty, xi, yi, 128, 128).transpose(ty, ci).transpose(tx, ci);
        processed.parallel(ty);
    } else {
        denoised.root();
        deinterleaved.root();
        corrected.root();
        processed.root();
    }

    return processed;
}

int main(int argc, char **argv) {
    // The camera pipe is specialized on the 2592x1968 images that
    // come in, so we'll just use an image instead of a uniform image.
    Image<int16_t> input(2592, 1968);
    UniformImage matrix_3200(Float(32), 2, "m3200"), matrix_7000(Float(32), 2, "m7000");
    Uniform<float> color_temp("color_temp", 3200.0f);
    Uniform<float> gamma("gamma", 1.8f);
    Uniform<float> contrast("contrast", 10.0f);

    // shift things inwards to give us enough padding on the
    // boundaries so that we don't need to check bounds. We're going
    // to make a 2560x1920 output image, so shift by 16, 12
    Func clamped("in");
    clamped(x, y) = input(x+16, y+12); //Debug( input(x+16, y+12), "Loading input at: ", x, y, x+16, y+12 );
    
    // Parameterized output type, because LLVM PTX (GPU) backend does not
    // currently allow 8-bit computations
    int bit_width = atoi(argv[1]);
    Type result_type = UInt(bit_width);
    
    // Pick a schedule   
    schedule = atoi(argv[2]);
    
    // Build the pipeline
    Func processed = process(clamped, result_type, matrix_3200, matrix_7000, color_temp, gamma, contrast);

    processed.compileToFile("curved");

#if 0
    Var xo("blockidx"), yo("blockidy"), c("c"), xi("threadidx"), yi("threadidy");
    Func output("output");
    output(xo, yo, c) = processed(xo, yo, c);
    Var nilc("nilc");
    output.root()
        .split(c, nilc, c, 3) // split C into a const loop by 3 on the inside
        .tile(xo, yo, xi, yi, 32, 16).transpose(yo, c).transpose(xo, c)
        .transpose(yi, c).transpose(xi, c)
        .parallel(xo).parallel(yo).parallel(xi).parallel(yi);

    std::vector<Func> funcs = output.rhs().funcs();    

    for (size_t i = 0; i < funcs.size(); i++) {
      if (funcs[i].name() == "curve") funcs[i].root();
      else if (funcs[i].name() == "matrix") funcs[i].root();
      else if (funcs[i].name() == "dem" || funcs[i].name() == "corrected" || funcs[i].name() == "curved") {
          // inline
          continue;
      }
      else {
          printf("scheduling %s as chunk on xo\n", funcs[i].name().c_str());
          funcs[i].chunk(xo)
                      .split(x, x, xi, 1).split(y, y, yi, 1)
                      .parallel(xi).parallel(yi);
      }
      //if (funcs[i].returnType() == UInt(8)) funcs[i].vectorize(x, 16);
      //if (funcs[i].returnType() == Int(16)) funcs[i].vectorize(x, 8);
      //if (funcs[i].returnType() == Float(32)) funcs[i].vectorize(x, 4);
    }

    output.compileToFile("curved");
#endif

    return 0;
}

<commit_msg>fixed bug and reintroduced tuple in camera pipe<commit_after>#include <FImage.h>
#include <stdint.h>

using namespace FImage;

int schedule;

Var x("x"), y("y"), tx("tx"), ty("ty"), c("c");

Func hot_pixel_suppression(Func input) {
    Expr a = max(max(input(x-2, y), input(x+2, y)),
                 max(input(x, y-2), input(x, y+2)));
    Expr b = min(min(input(x-2, y), input(x+2, y)),
                 min(input(x, y-2), input(x, y+2)));
    
    Func denoised("denoised");
    denoised(x, y) = clamp(input(x, y), b, a);

    return denoised;
}

Expr abs(Expr e) {
    return select(e < cast(e.type(), 0), -e, e);
}

Func interleave_x(Func a, Func b) {
    Func out;
    out(x, y) = select((x%2)==0, a(x/2, y), b(x/2, y));
    return out;
}

Func interleave_y(Func a, Func b) {
    Func out;
    out(x, y) = select((y%2)==0, a(x, y/2), b(x, y/2));
    return out;
}

Func deinterleave(Func raw) {
    // Deinterleave the color channels
    Func deinterleaved("deinterleaved");

    deinterleaved(x, y) = (raw(2*x, 2*y),   raw(2*x+1, 2*y), 
                           raw(2*x, 2*y+1), raw(2*x+1, 2*y+1));

    return deinterleaved;
}

Func demosaic(Func deinterleaved) {
    // These are the values we already know from the input
    // x_y = the value of channel x at a site in the input of channel y
    // gb refers to green sites in the blue rows
    // gr refers to green sites in the red rows

    // Give more convenient names to the four channels we know
    Func r_r("r_r"), g_gr("g_gr"), g_gb("g_gb"), b_b("b_b");    
    g_gr(x, y) = deinterleaved(x, y, 0);
    r_r(x, y)  = deinterleaved(x, y, 1);
    b_b(x, y)  = deinterleaved(x, y, 2);
    g_gb(x, y) = deinterleaved(x, y, 3);

    // These are the ones we need to interpolate
    Func b_r("b_r"), g_r("g_r");
    Func b_gr("b_gr"), r_gr("r_gr");
    Func b_gb("b_gb"), r_gb("r_gb");
    Func r_b("r_b"), g_b("g_b");

    // First calculate green at the red and blue sites

    // Try interpolating vertically and horizontally. Also compute
    // differences vertically and horizontally. Use interpolation in
    // whichever direction had the smallest difference.
    Expr gv_r  =    (g_gb(x, y-1) + g_gb(x, y))/2;
    Expr gvd_r = abs(g_gb(x, y-1) - g_gb(x, y));
    Expr gh_r  =    (g_gr(x+1, y) + g_gr(x, y))/2;
    Expr ghd_r = abs(g_gr(x+1, y) - g_gr(x, y));

    g_r(x, y)  = select(ghd_r < gvd_r, gh_r, gv_r);

    Expr gv_b  =    (g_gr(x, y+1) + g_gr(x, y))/2;
    Expr gvd_b = abs(g_gr(x, y+1) - g_gr(x, y));
    Expr gh_b  =    (g_gb(x-1, y) + g_gb(x, y))/2;
    Expr ghd_b = abs(g_gb(x-1, y) - g_gb(x, y));

    g_b(x, y)  = select(ghd_b < gvd_b, gh_b, gv_b);

    // Next interpolate red at gr by first interpolating, then
    // correcting using the error green would have had if we had
    // interpolated it in the same way (i.e. add the second derivative
    // of the green channel at the same place).
    Expr correction;
    correction = g_gr(x, y) - (g_r(x, y) + g_r(x-1, y))/2;
    r_gr(x, y) = correction + (r_r(x-1, y) + r_r(x, y))/2;

    // Do the same for other reds and blues at green sites
    correction = g_gr(x, y) - (g_b(x, y) + g_b(x, y-1))/2;
    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("dem");
    output(x, y) = (r(x, y), 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 a 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;
}


Func color_correct(Func input, UniformImage matrix_3200, UniformImage matrix_7000, Uniform<float> kelvin) {
    // Get a color matrix by linearly interpolating between two
    // calibrated matrices using inverse kelvin.

    Func matrix("matrix");
    Expr alpha = (1.0f/kelvin - 1.0f/3200) / (1.0f/7000 - 1.0f/3200);
    Expr val =  (matrix_3200(x, y) * alpha + matrix_7000(x, y) * (1 - alpha));
    matrix(x, y) = cast<int32_t>(val * 256.0f); // Q8.8 fixed point
    matrix.root();

    Func corrected("corrected");
    Expr ir = cast<int32_t>(input(x, y, 0));
    Expr ig = cast<int32_t>(input(x, y, 1));
    Expr ib = cast<int32_t>(input(x, y, 2));

    Expr r = matrix(3, 0) + matrix(0, 0) * ir + matrix(1, 0) * ig + matrix(2, 0) * ib;
    Expr g = matrix(3, 1) + matrix(0, 1) * ir + matrix(1, 1) * ig + matrix(2, 1) * ib;
    Expr b = matrix(3, 2) + matrix(0, 2) * ir + matrix(1, 2) * ig + matrix(2, 2) * ib;

    corrected(x, y) = cast<int16_t>((r, g, b)/256);

    return corrected;
}


Func apply_curve(Func input, Type result_type, Uniform<float> gamma, Uniform<float> contrast) {
    // copied from FCam
    Func curve("curve");

    Expr xf = clamp(cast<float>(x)/1024.0f, 0.0f, 1.0f);
    Expr g = pow(xf, 1.0f/gamma);
    Expr b = 2.0f - pow(2.0f, contrast/100.0f);
    Expr a = 2.0f - 2.0f*b; 
    Expr z = select(g > 0.5f,
                    1.0f - (a*(1.0f-g)*(1.0f-g) + b*(1.0f-g)),
                    a*g*g + b*g);

    Expr val = cast(result_type, clamp(z*256.0f, 0.0f, 255.0f));
    curve(x) = val;
    curve.root(); // It's a LUT, compute it once ahead of time.

    Func curved("curved");
    curved(x, y, c) = curve(input(x, y, c));

    return curved;
}

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("p");
    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; 
    processed(tx, ty, c) = curved(tx, ty, ci);
    processed.split(c, co, ci, 3); // bound color loop to 0-3
    if (schedule == 0) {
        // Compute in chunks over tiles, vectorized by 8
        denoised.chunk(tx).vectorize(x, 8);

        // Unroll the tuple part and compute it in the inside loop
        Var ti, tt = deinterleaved.arg(2);
        deinterleaved.split(tt, tt, ti, 4).unroll(ti).transpose(y, ti).transpose(x, ti);
        deinterleaved.chunk(tx).vectorize(x, 8);

        Var cc = corrected.arg(2);
        corrected.chunk(tx).transpose(y, cc).transpose(x, cc).vectorize(x, 4).unroll(cc, 3);
        processed.tile(tx, ty, xi, yi, 32, 32).transpose(ty, ci).transpose(tx, ci);
        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);

        // Unroll the tuple part and compute it in the inside loop
        Var ti, tt = deinterleaved.arg(2);
        deinterleaved.split(tt, tt, ti, 4).unroll(ti).transpose(y, ti).transpose(x, ti);
        deinterleaved.chunk(tx);

        Var cc = corrected.arg(2);
        corrected.chunk(tx).transpose(y, cc).transpose(x, cc).unroll(cc, 3);
        processed.tile(tx, ty, xi, yi, 128, 128).transpose(ty, ci).transpose(tx, ci);
        processed.parallel(ty);
    } else {
        denoised.root();
        deinterleaved.root();
        corrected.root();
        processed.root();
    }

    return processed;
}

int main(int argc, char **argv) {
    // The camera pipe is specialized on the 2592x1968 images that
    // come in, so we'll just use an image instead of a uniform image.
    Image<int16_t> input(2592, 1968);
    UniformImage matrix_3200(Float(32), 2, "m3200"), matrix_7000(Float(32), 2, "m7000");
    Uniform<float> color_temp("color_temp", 3200.0f);
    Uniform<float> gamma("gamma", 1.8f);
    Uniform<float> contrast("contrast", 10.0f);

    // shift things inwards to give us enough padding on the
    // boundaries so that we don't need to check bounds. We're going
    // to make a 2560x1920 output image, so shift by 16, 12
    Func clamped("in");
    clamped(x, y) = input(x+16, y+12); //Debug( input(x+16, y+12), "Loading input at: ", x, y, x+16, y+12 );
    
    // Parameterized output type, because LLVM PTX (GPU) backend does not
    // currently allow 8-bit computations
    int bit_width = atoi(argv[1]);
    Type result_type = UInt(bit_width);
    
    // Pick a schedule   
    schedule = atoi(argv[2]);
    
    // Build the pipeline
    Func processed = process(clamped, result_type, matrix_3200, matrix_7000, color_temp, gamma, contrast);

    processed.compileToFile("curved");

#if 0
    Var xo("blockidx"), yo("blockidy"), c("c"), xi("threadidx"), yi("threadidy");
    Func output("output");
    output(xo, yo, c) = processed(xo, yo, c);
    Var nilc("nilc");
    output.root()
        .split(c, nilc, c, 3) // split C into a const loop by 3 on the inside
        .tile(xo, yo, xi, yi, 32, 16).transpose(yo, c).transpose(xo, c)
        .transpose(yi, c).transpose(xi, c)
        .parallel(xo).parallel(yo).parallel(xi).parallel(yi);

    std::vector<Func> funcs = output.rhs().funcs();    

    for (size_t i = 0; i < funcs.size(); i++) {
      if (funcs[i].name() == "curve") funcs[i].root();
      else if (funcs[i].name() == "matrix") funcs[i].root();
      else if (funcs[i].name() == "dem" || funcs[i].name() == "corrected" || funcs[i].name() == "curved") {
          // inline
          continue;
      }
      else {
          printf("scheduling %s as chunk on xo\n", funcs[i].name().c_str());
          funcs[i].chunk(xo)
                      .split(x, x, xi, 1).split(y, y, yi, 1)
                      .parallel(xi).parallel(yi);
      }
      //if (funcs[i].returnType() == UInt(8)) funcs[i].vectorize(x, 16);
      //if (funcs[i].returnType() == Int(16)) funcs[i].vectorize(x, 8);
      //if (funcs[i].returnType() == Float(32)) funcs[i].vectorize(x, 4);
    }

    output.compileToFile("curved");
#endif

    return 0;
}


{'fraction_non_alphanumeric': 0.17040753791941168, 'fraction_numerical': 0.0304504366477708, 'mean_word_length': 2.4895749799518847, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 43, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5979332', 'n_tokens_mistral': 11019, 'n_tokens_neox': 10187, 'n_words': 2676}