Bounds propagation during sliding window optimization

src/SlidingWindow.cpp

@@ -223,6 +223,7 @@ class SlidingWindowOnFunctionAndLoop : public IRMutator {
     Expr loop_min;
     set<int> &slid_dimensions;
     Scope<Expr> scope;
+    Scope<Interval> &bounds_scope;
 
     // For loops strictly between the loop being slid over and the current
     // node (not including the loop being slid over itself).
@@ -282,8 +283,8 @@ class SlidingWindowOnFunctionAndLoop : public IRMutator {
                 internal_assert(min_val && max_val);
                 Expr min_req = *min_val;
                 Expr max_req = *max_val;
-                min_req = expand_expr(min_req, scope);
-                max_req = expand_expr(max_req, scope);
+                min_req = simplify(expand_expr(min_req, scope), bounds_scope);
+                max_req = simplify(expand_expr(max_req, scope), bounds_scope);
 
                 debug(3) << func_args[i] << ":" << min_req << ", " << max_req << "\n";
                 if (expr_depends_on_var(min_req, loop_var) ||
@@ -594,7 +595,10 @@ class SlidingWindowOnFunctionAndLoop : public IRMutator {
     }
 
     Stmt visit(const LetStmt *op) override {
-        ScopedBinding<Expr> bind(scope, op->name, simplify(expand_expr(op->value, scope)));
+        ScopedBinding<Interval> bind_bounds(bounds_scope, op->name,
+                                            bounds_of_expr_in_scope(op->value, bounds_scope));
+        ScopedBinding<Expr> bind(scope, op->name, simplify(expand_expr(op->value, scope), bounds_scope));
+
         Stmt new_body = mutate(op->body);
 
         Expr value = op->value;
@@ -613,8 +617,10 @@ class SlidingWindowOnFunctionAndLoop : public IRMutator {
     }
 
 public:
-    SlidingWindowOnFunctionAndLoop(Function f, string v, Expr v_min, set<int> &slid_dimensions)
-        : func(std::move(f)), loop_var(std::move(v)), loop_min(std::move(v_min)), slid_dimensions(slid_dimensions) {
+    SlidingWindowOnFunctionAndLoop(Function f, string v, Expr v_min, set<int> &slid_dimensions,
+                                   Scope<Interval> &bounds_scope)
+        : func(std::move(f)), loop_var(std::move(v)), loop_min(std::move(v_min)),
+          slid_dimensions(slid_dimensions), bounds_scope(bounds_scope) {
     }
 
     Expr new_loop_min;
@@ -755,9 +761,16 @@ class SlidingWindow : public IRMutator {
     // Keep track of realizations we want to slide, from innermost to
     // outermost.
     list<Function> sliding;
+    Scope<Interval> bounds_scope;
 
     using IRMutator::visit;
 
+    Stmt visit(const LetStmt *op) override {
+        ScopedBinding<Interval> bind(bounds_scope, op->name,
+                                     bounds_of_expr_in_scope(op->value, bounds_scope));
+        return IRMutator::visit(op);
+    }
+
     Stmt visit(const Realize *op) override {
         // Find the args for this function
         map<string, Function>::const_iterator iter = env.find(op->name);
@@ -827,7 +840,14 @@ class SlidingWindow : public IRMutator {
 
             set<int> &slid_dims = slid_dimensions[func.name()];
             size_t old_slid_dims_size = slid_dims.size();
-            SlidingWindowOnFunctionAndLoop slider(func, name, prev_loop_min, slid_dims);
+
+            Interval min_bounds = bounds_of_expr_in_scope(loop_min, bounds_scope);
+            Interval max_bounds = bounds_of_expr_in_scope(loop_max, bounds_scope);
+            ScopedBinding<Interval> bind_bounds(bounds_scope, op->name,
+                                                Interval(min_bounds.min, max_bounds.max));
+
+            SlidingWindowOnFunctionAndLoop slider(func, name, prev_loop_min, slid_dims, bounds_scope);
+
             body = slider(body);
 
             if (func.schedule().memory_type() == MemoryType::Register &&

test/correctness/sliding_window.cpp

@@ -415,6 +415,30 @@ int main(int argc, char **argv) {
         }
     }
 
+    {
+        // Sliding a cascade of filters. Halide needs bounds propagation
+        // to prove that the innermost filters have monotonic bounds.
+        count = 0;
+        Func f1, f2, f3, f4, g;
+        f1(x) = call_counter(x, 0);
+        f2(x) = f1(0) + f1(x);
+        f3(x) = f2(0) + f2(x);
+        f4(x) = f3(0) + f3(x);
+        g(x) = f4(0) + f4(x);
+        f1.store_root().compute_at(g, x);
+        f2.store_root().compute_at(g, x);
+        f3.store_root().compute_at(g, x);
+        f4.store_root().compute_at(g, x);
+        g.bound(x, 0, 10);
+
+        g.realize({10});
+        // f1 spans x in [0, 9], so 10 calls when slid.
+        if (count != 10) {
+            printf("f1 was called %d times instead of %d times\n", count, 10);
+            return 1;
+        }
+    }
+
     printf("Success!\n");
     return 0;
 }