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
|
#include "Halide.h"
#include <iostream>
// Verifies that constraints on the input ImageParam propagates to the output
// function.
using namespace Halide;
using namespace Halide::Internal;
void check_int(const Expr &expr, int expected) {
if (!is_const(expr, expected)) {
std::cerr << "Found expression " << expr << "; "
<< "expected constant int " << expected << "\n";
exit(1);
}
}
constexpr int size = 10;
class CheckLoops : public IRVisitor {
public:
int count = 0;
private:
using IRVisitor::visit;
void visit(const For *op) override {
std::cout << "for(" << op->name << ", " << op->min << ", " << op->extent << ")\n";
check_int(op->min, 0);
check_int(op->extent, size);
++count;
IRVisitor::visit(op);
}
};
class Validator : public IRMutator {
using IRMutator::mutate;
Stmt mutate(const Stmt &s) override {
CheckLoops c;
s.accept(&c);
if (c.count != 1) {
std::cerr << "expected one loop, found " << c.count << "\n";
exit(1);
}
return s;
}
};
int main(int argc, char **argv) {
ImageParam input(UInt(8), 1);
input.dim(0).set_bounds(0, size);
{
Func f;
Var x;
f(x) = input(x);
// Output must have the same size as the input.
f.output_buffer().dim(0).set_bounds(input.dim(0).min(), input.dim(0).extent());
f.add_custom_lowering_pass(new Validator);
f.compile_jit();
Buffer<uint8_t> dummy(size);
dummy.fill(42);
input.set(dummy);
Buffer<uint8_t> out = f.realize({size});
if (!out.all_equal(42)) {
std::cerr << "wrong output\n";
exit(1);
}
}
{
Func f;
Var x;
f(x) = undef(UInt(8));
RDom r(input);
f(r.x) = cast<uint8_t>(42);
f.add_custom_lowering_pass(new Validator);
f.compile_jit();
Buffer<uint8_t> dummy(size);
input.set(dummy);
Buffer<uint8_t> out = f.realize({size});
if (!out.all_equal(42)) {
std::cerr << "wrong output\n";
exit(1);
}
}
std::cout << "Success!\n";
return 0;
}
|
