void check(const std::string &inner_loop_level,
const std::string &outer_loop_level) {
Buffer<float> result = outer.realize(1, 1, 1);
Module m = outer.compile_to_module({outer.infer_arguments()});
CheckLoopLevels c(inner_loop_level, outer_loop_level);
m.functions().front().body.accept(&c);
}
int rdom_wrapper_test() {
Func source("source"), g("g"), result("result");
Var x("x"), y("y");
source(x, y) = x + y;
ImageParam img(Int(32), 2, "img");
Buffer<int> buf = source.realize(200, 200);
img.set(buf);
g(x, y) = 10;
g(x, y) += 2 * img(x, x);
RDom r(0, 200, 0, 200);
g(r.x, r.y) += 3 * img(r.y, r.y);
// Make a global wrapper on 'g', so that we can schedule initialization
// and the update on the same compute level at the global wrapper
Func wrapper = g.in().compute_root();
g.compute_at(wrapper, x);
Func img_f = img;
img_f.compute_root();
// Check the call graphs.
// Expect 'wrapper' to call 'g', initialization of 'g' to call nothing
// and its update to call 'img_f' and 'g', 'img_f' to call 'img'
Module m = wrapper.compile_to_module({wrapper.infer_arguments()});
CheckCalls c;
m.functions().front().body.accept(&c);
CallGraphs expected = {
{g.name(), {img_f.name(), g.name()}},
{wrapper.name(), {g.name()}},
{img_f.name(), {img.name()}},
};
if (check_call_graphs(c.calls, expected) != 0) {
return -1;
}
Buffer<int> im = wrapper.realize(200, 200);
auto func = [](int x, int y) { return 4*x + 6* y + 10; };
if (check_image(im, func)) {
return -1;
}
return 0;
}
int main(int argc, char **argv) {
Param<float> reservoirConcentration;
Param<float> stepTime;
Param<float> layerMixConst;
Param<float> layerTimeDivisor;
Func sumDx;
Func layerMixed;
Func initialDeveloperMirrored;
ImageParam devConc(type_of<float>(),2);
Func dDevelConc;
Func developerConcentration = lambda(x,y,devConc(x,y));
dDevelConc = calcLayerMix(developerConcentration, layerMixConst, stepTime,
layerTimeDivisor, reservoirConcentration);
std::vector<Argument> ddcArgs = dDevelConc.infer_arguments();
dDevelConc.compile_to_file("calcLayerMix",ddcArgs);
return 0;
}
