@@ -22,14 +22,15 @@ class AllocationInference : public IRMutator {
22
22
using IRMutator::visit;
23
23
24
24
const map<string, Function> &env;
25
+ const FuncValueBounds &func_bounds;
25
26
set<string> touched_by_extern;
26
27
27
28
void visit (const Realize *op) {
28
29
map<string, Function>::const_iterator iter = env.find (op->name );
29
30
assert (iter != env.end ());
30
31
Function f = iter->second ;
31
32
32
- Box b = box_touched (op->body , op->name );
33
+ Box b = box_touched (op->body , op->name , Scope<Interval>(), func_bounds );
33
34
34
35
if (touched_by_extern.count (f.name ())) {
35
36
// The region touched is at least the region required at this
@@ -64,7 +65,8 @@ class AllocationInference : public IRMutator {
64
65
}
65
66
66
67
public:
67
- AllocationInference (const map<string, Function> &e) : env(e) {
68
+ AllocationInference (const map<string, Function> &e, const FuncValueBounds &fb) :
69
+ env (e), func_bounds(fb) {
68
70
// Figure out which buffers are touched by extern stages
69
71
for (map<string, Function>::const_iterator iter = e.begin ();
70
72
iter != e.end (); ++iter) {
@@ -82,8 +84,10 @@ class AllocationInference : public IRMutator {
82
84
}
83
85
};
84
86
85
- Stmt allocation_bounds_inference (Stmt s, const map<string, Function> &env) {
86
- AllocationInference inf (env);
87
+ Stmt allocation_bounds_inference (Stmt s,
88
+ const map<string, Function> &env,
89
+ const FuncValueBounds &fb) {
90
+ AllocationInference inf (env, fb);
87
91
s = inf.mutate (s);
88
92
return s;
89
93
}
0 commit comments