1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
|
#include "Halide.h"
using namespace Halide;
using namespace Halide::Internal;
class CheckLoopLevels : public IRVisitor {
public:
CheckLoopLevels(const std::string &inner_loop_level,
const std::string &outer_loop_level)
: inner_loop_level(inner_loop_level), outer_loop_level(outer_loop_level) {
}
private:
using IRVisitor::visit;
const std::string inner_loop_level, outer_loop_level;
std::string inside_for_loop;
void visit(const For *op) override {
std::string old_for_loop = inside_for_loop;
inside_for_loop = op->name;
IRVisitor::visit(op);
inside_for_loop = old_for_loop;
}
void visit(const Call *op) override {
IRVisitor::visit(op);
if (op->name == "sin_f32") {
_halide_user_assert(starts_with(inside_for_loop, inner_loop_level));
} else if (op->name == "cos_f32") {
_halide_user_assert(starts_with(inside_for_loop, outer_loop_level));
}
}
void visit(const Store *op) override {
IRVisitor::visit(op);
if (op->name.substr(0, 5) == "inner") {
_halide_user_assert(starts_with(inside_for_loop, inner_loop_level));
} else if (op->name.substr(0, 5) == "outer") {
_halide_user_assert(starts_with(inside_for_loop, outer_loop_level));
} else {
_halide_user_assert(0);
}
}
};
Var x("x"), y("y"), c("c");
struct Test {
Func inner, outer;
LoopLevel inner_compute_at, inner_store_at;
explicit Test(int i) {
// We use specific calls as proxies for verifying that compute_at
// happens where we expect: sin() for the inner function, cos()
// for the outer one; these are chosen mainly because they won't
// ever get generated incidentally by the lowering code as part of
// general code structure.
inner = Func("inner" + std::to_string(i));
inner(x, y, c) = sin(cast<float>(x + y + c));
inner.compute_at(inner_compute_at).store_at(inner_store_at);
outer = Func("outer" + std::to_string(i));
outer(x, y, c) = cos(cast<float>(inner(x, y, c)));
}
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 main(int argc, char **argv) {
// Test that LoopLevels set after being specified still take effect.
{
Test t(1);
t.inner_compute_at.set(LoopLevel(t.outer, x));
t.inner_store_at.set(LoopLevel(t.outer, x));
t.check("outer1.s0.x", "outer1.s0.x");
}
// Same as before, but using inlined() for both inner LoopLevels.
{
Test t(2);
t.inner_compute_at.set(LoopLevel::inlined());
t.inner_store_at.set(LoopLevel::inlined());
t.check("outer2.s0.x", "outer2.s0.x");
}
// Same as before, but using root() for both inner LoopLevels.
{
Test t(3);
t.inner_compute_at.set(LoopLevel::root());
t.inner_store_at.set(LoopLevel::root());
t.check("inner3.s0.x", "outer3.s0.x");
}
// Same as before, but using different store_at and compute_at()
{
Test t(4);
t.inner_compute_at.set(LoopLevel(t.outer, y));
t.inner_store_at.set(LoopLevel::root());
t.check("inner4.s0.x", "outer4.s0.x");
}
// Same as before, but using inlined for store_at() [equivalent to omitting
// the store_at() call entirely] and non-inlined for compute_at
{
Test t(5);
t.inner_compute_at.set(LoopLevel(t.outer, y));
t.inner_store_at.set(LoopLevel::inlined());
t.check("inner5.s0.x", "outer5.s0.x");
}
printf("Success!\n");
return 0;
}
|
