Shared bounds tracker machinery.

Makefile

@@ -460,6 +460,7 @@ SOURCE_FILES = \
   Bounds.cpp \
   BoundsInference.cpp \
   BoundSmallAllocations.cpp \
+  BoundsTracker.cpp \
   Buffer.cpp \
   Callable.cpp \
   CanonicalizeGPUVars.cpp \
@@ -665,6 +666,7 @@ HEADER_FILES = \
   Bounds.h \
   BoundsInference.h \
   BoundSmallAllocations.h \
+  BoundsTracker.h \
   Buffer.h \
   Callable.h \
   CanonicalizeGPUVars.h \

src/BoundConstantExtentLoops.cpp

@@ -1,11 +1,9 @@
 #include "BoundConstantExtentLoops.h"
-#include "Bounds.h"
-#include "CSE.h"
+#include "BoundsTracker.h"
 #include "IRMutator.h"
 #include "IROperator.h"
 #include "Simplify.h"
-#include "SimplifyCorrelatedDifferences.h"
-#include "Substitute.h"
+#include "Util.h"
 
 namespace Halide {
 namespace Internal {
@@ -15,29 +13,23 @@ class BoundLoops : public IRMutator {
 protected:
     using IRMutator::visit;
 
-    std::vector<std::pair<std::string, Expr>> lets;
+    BoundsTracker tracker;
 
     Stmt visit(const LetStmt *op) override {
-        if (is_pure(op->value)) {
-            lets.emplace_back(op->name, op->value);
-            Stmt s = IRMutator::visit(op);
-            lets.pop_back();
-            return s;
-        } else {
-            return IRMutator::visit(op);
-        }
+        auto binding = tracker.push_let(op->name, op->value);
+        return IRMutator::visit(op);
     }
 
-    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;
+        Stmt then_case, else_case;
+        {
+            auto fact = tracker.push_fact(op->condition);
+            then_case = mutate(op->then_case);
+        }
         if (op->else_case.defined()) {
-            facts.back() = simplify(!op->condition);
+            auto fact = tracker.push_fact(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;
@@ -47,6 +39,7 @@ class BoundLoops : public IRMutator {
     }
 
     Stmt visit(const For *op) override {
+        auto bind = tracker.push_for(op->name, op->min, op->max);
         Expr extent = simplify(op->extent());
         if (is_const(extent)) {
             // Nothing needs to be done
@@ -56,36 +49,46 @@ class BoundLoops : public IRMutator {
         if (op->for_type == ForType::Unrolled ||
             op->for_type == ForType::Vectorized) {
             // Give it one last chance to simplify to an int
+            extent = tracker.simplify_with_context(extent);
             Stmt body = op->body;
             const IntImm *e = extent.as<IntImm>();
 
-            if (e == nullptr) {
-                // We're about to hard fail. Get really aggressive
-                // with the simplifier.
-                extent = rewrap_used_lets(extent, lets);
-                extent = remove_likelies(extent);
-                extent = substitute_in_all_lets(extent);
-                extent = simplify(extent,
-                                  Scope<Interval>::empty_scope(),
-                                  Scope<ModulusRemainder>::empty_scope(),
-                                  facts);
-                e = extent.as<IntImm>();
-            }
-
             Expr extent_upper;
             if (e == nullptr) {
-                // Still no luck. Try taking an upper bound and
-                // injecting an if statement around the body.
-                extent_upper = find_constant_bound(extent, Direction::Upper, Scope<Interval>());
-                if (extent_upper.defined()) {
-                    e = extent_upper.as<IntImm>();
-                    body =
-                        IfThenElse::make(likely_if_innermost(Variable::make(Int(32), op->name) <=
-                                                             op->max),
-                                         body);
+                // We're about to hard fail. Get really aggressive with the
+                // simplifier: inline every enclosing let and simplify under
+                // every dominating condition.
+                debug(4) << "Trying to find a constant bound for loop " << op->name << "\n"
+                         << "Extent: " << extent << "\n";
+                Interval bounds = tracker.find_constant_bounds_aggressive(extent);
+                debug(4) << "Bounds found: [" << bounds.min << ", " << bounds.max << "]\n";
+                auto lo = bounds.has_lower_bound() ? as_const_int(bounds.min) : std::nullopt;
+                auto hi = bounds.has_upper_bound() ? as_const_int(bounds.max) : std::nullopt;
+                if (hi) {
+                    // Copy the Expr out of `bounds` before it goes out of
+                    // scope below -- otherwise e, taken as a raw pointer via
+                    // as<IntImm>(), would be left dangling into a node whose
+                    // only reference was owned by this soon-to-be-destroyed
+                    // Interval.
+                    extent_upper = bounds.max;
+                    if (lo && *lo == *hi) {
+                        // The bound is exact: no guard needed.
+                        e = extent_upper.as<IntImm>();
+                    }
                 }
             }
 
+            if (e == nullptr && extent_upper.defined()) {
+                // Still no luck getting an exact extent. Take the upper
+                // bound instead and guard the body with an if statement.
+                debug(4) << "Found an upper bound instead: " << extent_upper << "\n";
+                e = extent_upper.as<IntImm>();
+                body =
+                    IfThenElse::make(likely_if_innermost(Variable::make(Int(32), op->name) <=
+                                                         op->max),
+                                     body);
+            }
+
             if (e == nullptr && permit_failed_unroll && op->for_type == ForType::Unrolled) {
                 // Still no luck, but we're allowed to fail. Rewrite
                 // to a serial loop.

src/BoundSmallAllocations.cpp

@@ -1,9 +1,8 @@
 #include "BoundSmallAllocations.h"
-#include "Bounds.h"
+#include "BoundsTracker.h"
 #include "CodeGen_Internal.h"
 #include "IRMutator.h"
 #include "IROperator.h"
-#include "Simplify.h"
 
 namespace Halide {
 namespace Internal {
@@ -15,25 +14,25 @@ class BoundSmallAllocations : public IRMutator {
     using IRMutator::visit;
 
     // Track constant bounds
-    Scope<Interval> scope;
+    BoundsTracker tracker;
 
     template<typename LetOrLetStmt>
     auto visit_let(const LetOrLetStmt *op) -> decltype(op->body) {
         // Visit an entire chain of lets in a single method to conserve stack space.
         struct Frame {
             const LetOrLetStmt *op;
-            ScopedBinding<Interval> binding;
-            Frame(const LetOrLetStmt *op, Scope<Interval> &scope)
+            BoundsTracker::Binding binding;
+            Frame(const LetOrLetStmt *op, BoundsTracker &tracker)
                 : op(op),
-                  binding(scope, op->name, find_constant_bounds(op->value, scope)) {
+                  binding(tracker.push_let(op->name, op->value)) {
             }
         };
         std::vector<Frame> frames;
         decltype(op->body) result;
 
         do {
             result = op->body;
-            frames.emplace_back(op, scope);
+            frames.emplace_back(op, tracker);
         } while ((op = result.template as<LetOrLetStmt>()));
 
         result = mutate(result);
@@ -58,12 +57,7 @@ class BoundSmallAllocations : public IRMutator {
     DeviceAPI device_api = DeviceAPI::None;
 
     Stmt visit(const For *op) override {
-        Interval min_bounds = find_constant_bounds(op->min, scope);
-        Interval max_bounds = find_constant_bounds(op->max, scope);
-        Interval b = Interval::make_union(min_bounds, max_bounds);
-        b.min = simplify(b.min);
-        b.max = simplify(b.max);
-        ScopedBinding<Interval> bind(scope, op->name, b);
+        auto binding = tracker.push_for(op->name, op->min, op->max);
         bool new_in_thread_loop =
             in_thread_loop || op->for_type == ForType::GPUThread;
         ScopedValue<bool> old_in_thread_loop(in_thread_loop, new_in_thread_loop);
@@ -86,7 +80,7 @@ class BoundSmallAllocations : public IRMutator {
             bool changed = false;
             bool found_non_constant_extent = false;
             for (Range &r : region) {
-                Expr bound = find_constant_bound(r.extent, Direction::Upper, scope);
+                Expr bound = tracker.find_constant_bound_aggressive(r.extent, Direction::Upper);
                 // We can allow non-constant extents for now, as long as all
                 // remaining dimensions are 1 (so the stride is unused, which
                 // will be non-constant).
@@ -116,7 +110,7 @@ class BoundSmallAllocations : public IRMutator {
         for (const Expr &e : op->extents) {
             total_extent *= e;
         }
-        Expr bound = find_constant_bound(total_extent, Direction::Upper, scope);
+        Expr bound = tracker.find_constant_bound_aggressive(total_extent, Direction::Upper);
 
         if (!bound.defined() && must_be_constant(op->memory_type)) {
             user_assert(op->memory_type != MemoryType::Register)