Teach unrolling to exploit conditions in enclosing ifs

src/BoundsInference.cpp

@@ -1013,11 +1013,11 @@ class BoundsInference : public IRMutator {
                     }
 
                     // Dump out the region required of each stage for debugging.
-
                     /*
                     debug(0) << "Box required of " << producer.name
                              << " by " << consumer.name
-                             << " stage " << consumer.stage << ":\n";
+                             << " stage " << consumer.stage << ":\n"
+                             << " used: " << b.used << "\n";
                     for (size_t k = 0; k < b.size(); k++) {
                         debug(0) << "  " << b[k].min << " ... " << b[k].max << "\n";
                     }

src/Simplify.cpp

@@ -355,8 +355,13 @@ Simplify::ScopedFact::~ScopedFact() {
 
 Expr simplify(const Expr &e, bool remove_dead_let_stmts,
               const Scope<Interval> &bounds,
-              const Scope<ModulusRemainder> &alignment) {
+              const Scope<ModulusRemainder> &alignment,
+              const std::vector<Expr> &assumptions) {
     Simplify m(remove_dead_let_stmts, &bounds, &alignment);
+    std::vector<Simplify::ScopedFact> facts;
+    for (const Expr &a : assumptions) {
+        facts.push_back(m.scoped_truth(a));
+    }
     Expr result = m.mutate(e, nullptr);
     if (m.in_unreachable) {
         return unreachable(e.type());
@@ -366,8 +371,13 @@ Expr simplify(const Expr &e, bool remove_dead_let_stmts,
 
 Stmt simplify(const Stmt &s, bool remove_dead_let_stmts,
               const Scope<Interval> &bounds,
-              const Scope<ModulusRemainder> &alignment) {
+              const Scope<ModulusRemainder> &alignment,
+              const std::vector<Expr> &assumptions) {
     Simplify m(remove_dead_let_stmts, &bounds, &alignment);
+    std::vector<Simplify::ScopedFact> facts;
+    for (const Expr &a : assumptions) {
+        facts.push_back(m.scoped_truth(a));
+    }
     Stmt result = m.mutate(s);
     if (m.in_unreachable) {
         return Evaluate::make(unreachable());

src/Simplify.h

@@ -13,19 +13,22 @@
 namespace Halide {
 namespace Internal {
 
-/** Perform a a wide range of simplifications to expressions and
- * statements, including constant folding, substituting in trivial
- * values, arithmetic rearranging, etc. Simplifies across let
- * statements, so must not be called on stmts with dangling or
- * repeated variable names.
+/** Perform a wide range of simplifications to expressions and statements,
+ * including constant folding, substituting in trivial values, arithmetic
+ * rearranging, etc. Simplifies across let statements, so must not be called on
+ * stmts with dangling or repeated variable names. Can optionally be passed
+ * known bounds of any variables, known alignment properties, and any other
+ * Exprs that should be assumed to be true.
  */
 // @{
 Stmt simplify(const Stmt &, bool remove_dead_code = true,
               const Scope<Interval> &bounds = Scope<Interval>::empty_scope(),
-              const Scope<ModulusRemainder> &alignment = Scope<ModulusRemainder>::empty_scope());
+              const Scope<ModulusRemainder> &alignment = Scope<ModulusRemainder>::empty_scope(),
+              const std::vector<Expr> &assumptions = std::vector<Expr>());
 Expr simplify(const Expr &, bool remove_dead_code = true,
               const Scope<Interval> &bounds = Scope<Interval>::empty_scope(),
-              const Scope<ModulusRemainder> &alignment = Scope<ModulusRemainder>::empty_scope());
+              const Scope<ModulusRemainder> &alignment = Scope<ModulusRemainder>::empty_scope(),
+              const std::vector<Expr> &assumptions = std::vector<Expr>());
 // @}
 
 /** Attempt to statically prove an expression is true using the simplifier. */

src/UnrollLoops.cpp

@@ -32,6 +32,24 @@ class UnrollLoops : public IRMutator {
         }
     }
 
+    std::vector<Expr> facts;
+    Stmt visit(const IfThenElse *op) override {
+        facts.push_back(op->condition);
+        Stmt then_case = mutate(op->then_case);
+        Stmt else_case;
+        if (op->else_case.defined()) {
+            facts.back() = simplify(!op->condition);
+            else_case = mutate(op->else_case);
+        }
+        facts.pop_back();
+        if (then_case.same_as(op->then_case) &&
+            else_case.same_as(op->else_case)) {
+            return op;
+        } else {
+            return IfThenElse::make(op->condition, then_case, else_case);
+        }
+    }
+
     Stmt visit(const For *for_loop) override {
         if (for_loop->for_type == ForType::Unrolled) {
             // Give it one last chance to simplify to an int
@@ -47,7 +65,11 @@ class UnrollLoops : public IRMutator {
                 }
                 extent = remove_likelies(extent);
                 extent = substitute_in_all_lets(extent);
-                extent = simplify(extent);
+                extent = simplify(extent,
+                                  true,
+                                  Scope<Interval>::empty_scope(),
+                                  Scope<ModulusRemainder>::empty_scope(),
+                                  facts);
                 e = extent.as<IntImm>();
             }
 

test/correctness/CMakeLists.txt

@@ -318,6 +318,7 @@ tests(GROUPS correctness
       uninitialized_read.cpp
       unique_func_image.cpp
       unroll_dynamic_loop.cpp
+      unroll_loop_with_implied_constant_bounds.cpp
       unrolled_reduction.cpp
       unsafe_dedup_lets.cpp
       unsafe_promises.cpp

test/correctness/unroll_loop_with_implied_constant_bounds.cpp

@@ -0,0 +1,45 @@
+#include "Halide.h"
+
+using namespace Halide;
+
+int main(int argc, char **argv) {
+    // This test verifies that unrolling is capable of inferring constant bounds
+    // of loops that are implied by containing if statement conditions, e.g the
+    // following structure should work:
+    /*
+      let extent = foo
+      if (foo == 7) {
+        unrolled for (x from 0 to extent) {...}
+      }
+    */
+
+    Func intermediate("intermediate");
+
+    Func output1("output1"), output2("output2");
+
+    Var x("x"), y("y"), c("c");
+
+    intermediate(x, y, c) = x + y + c;
+
+    output1(x, y, c) = intermediate(x, y, c);
+    output2(x, y, c) = intermediate(x, y, c);
+
+    Expr three_channels =
+        (output1.output_buffer().dim(2).extent() == 3 &&
+         output1.output_buffer().dim(2).min() == 0 &&
+         output2.output_buffer().dim(2).extent() == 3 &&
+         output2.output_buffer().dim(2).min() == 0);
+
+    intermediate.compute_root()
+        .specialize(three_channels)
+        .unroll(c);
+
+    Pipeline p{{output1, output2}};
+
+    // Should not throw an error in loop unrolling.
+    p.compile_jit();
+
+    printf("Success!\n");
+
+    return 0;
+}