Func process(Func raw, Type result_type,
ImageParam matrix_3200, ImageParam matrix_7000, Param<float> color_temp,
Param<float> gamma, Param<float> contrast, Param<int> blackLevel, Param<int> whiteLevel) {
Var yii, xi;
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, blackLevel, whiteLevel);
processed(x, y, c) = curved(x, y, c);
// Schedule
Expr out_width = processed.output_buffer().width();
Expr out_height = processed.output_buffer().height();
int strip_size = 32;
int vec = target.natural_vector_size(UInt(16));
if (target.has_feature(Target::HVX_64)) {
vec = 32;
} else if (target.has_feature(Target::HVX_128)) {
vec = 64;
}
denoised.compute_at(processed, yi).store_at(processed, yo)
.fold_storage(y, 8)
.vectorize(x, vec);
deinterleaved.compute_at(processed, yi).store_at(processed, yo)
.fold_storage(y, 4)
.vectorize(x, 2*vec, TailStrategy::RoundUp)
.reorder(c, x, y)
.unroll(c);
corrected.compute_at(processed, x)
.vectorize(x, vec)
.reorder(c, x, y)
.unroll(c);
processed.compute_root()
.split(y, yo, yi, strip_size)
.split(yi, yi, yii, 2)
.split(x, x, xi, 2*vec, TailStrategy::RoundUp)
.reorder(xi, c, yii, x, yi, yo)
.vectorize(xi, 2*vec)
.parallel(yo);
if (target.features_any_of({Target::HVX_64, Target::HVX_128})) {
processed.hexagon();
denoised.align_storage(x, vec);
deinterleaved.align_storage(x, vec);
corrected.align_storage(x, vec);
}
// We can generate slightly better code if we know the splits divide the extent.
processed
.bound(c, 0, 3)
.bound(x, 0, ((out_width)/(2*vec))*(2*vec))
.bound(y, 0, (out_height/strip_size)*strip_size);
return processed;
}
