Allow the simplifier to use facts in its can_prove() predicates.

src/Bounds.cpp

@@ -22,6 +22,7 @@
 #include "SimplifyCorrelatedDifferences.h"
 #include "Solve.h"
 #include "StrictifyFloat.h"
+#include "Substitute.h"
 #include "Util.h"
 #include "Var.h"
 
@@ -2039,6 +2040,102 @@ class FindInnermostVar : public IRVisitor {
     }
 };
 
+// Bind each compound index in the body of an IfThenElse that its condition
+// also mentions to a let around it, and repeat the condition in terms of the
+// name inside the original one, e.g.
+//
+//   if (min_x <= x*s + r) { f(x*s + r); g(r) }
+//
+// becomes
+//
+//   let t = x*s + r in if (min_x <= x*s + r) { if (min_x <= t) { f(t); g(r) } }
+//
+// BoxesTouched bounds an index by an IfThenElse's condition one variable at a
+// time, so it can only use a condition that relates several variables if
+// they're named by one. The original condition stays outside so that it keeps
+// bounding each of those variables on its own, for the other indices they
+// appear in.
+class NameGuardedIndices : public IRMutator {
+    using IRMutator::visit;
+
+    Stmt visit(const IfThenElse *op) override {
+        // Name the indices here before recursing, so that the conditions of
+        // nested IfThenElses that bound the same index refer to the same let.
+        Stmt then_case = op->then_case;
+        vector<Expr> indices;
+        auto add_index = [&](const Expr &e) {
+            if (e.type() != Int(32) || e.as<Variable>() || is_const(e)) {
+                return;
+            }
+            for (const Expr &i : indices) {
+                if (equal(i, e)) {
+                    return;
+                }
+            }
+            indices.push_back(e);
+        };
+        // An index and each compound part of it, outermost first, so that the
+        // largest one the condition mentions gets the name.
+        auto add_index_and_parts = [&](const Expr &arg) {
+            add_index(arg);
+            visit_with(
+                arg,
+                [&](auto *self, const Add *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Sub *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Mul *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Select *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Div *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Mod *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Min *op) { add_index(op); self->visit_base(op); },
+                [&](auto *self, const Max *op) { add_index(op); self->visit_base(op); });
+        };
+        visit_with(
+            then_case,
+            [&](auto *self, const Call *call) {
+                if (call->call_type == Call::Halide || call->call_type == Call::Image) {
+                    for (const Expr &arg : call->args) {
+                        add_index_and_parts(arg);
+                    }
+                }
+                self->visit_base(call);
+            },
+            [&](auto *self, const Provide *provide) {
+                for (const Expr &arg : provide->args) {
+                    add_index_and_parts(arg);
+                }
+                self->visit_base(provide);
+            });
+
+        Expr condition = op->condition;
+        vector<pair<string, Expr>> lets;
+        for (const Expr &i : indices) {
+            string name = unique_name('t');
+            Expr var = Variable::make(i.type(), name);
+            Expr new_condition = substitute(i, var, condition);
+            if (!new_condition.same_as(condition)) {
+                condition = new_condition;
+                then_case = substitute(i, var, then_case);
+                lets.emplace_back(name, i);
+            }
+        }
+        then_case = mutate(then_case);
+        Stmt else_case = mutate(op->else_case);
+
+        if (lets.empty()) {
+            if (then_case.same_as(op->then_case) && else_case.same_as(op->else_case)) {
+                return op;
+            }
+            return IfThenElse::make(op->condition, then_case, else_case);
+        }
+        then_case = IfThenElse::make(condition, then_case);
+        Stmt stmt = IfThenElse::make(op->condition, then_case, else_case);
+        for (const auto &let : reverse_view(lets)) {
+            stmt = LetStmt::make(let.first, let.second, stmt);
+        }
+        return stmt;
+    }
+};
+
 // Place innermost vars in an IfThenElse's condition as far to the left as possible.
 class SolveIfThenElse : public IRMutator {
 protected:
@@ -3126,6 +3223,7 @@ map<string, Box> boxes_touched(const Expr &e, Stmt s, bool consider_calls, bool
     // as possible, so that BoxesTouched can prune the variable scope tighter
     // when encountering the IfThenElse.
     if (s.defined()) {
+        s = NameGuardedIndices()(s);
         s = SolveIfThenElse()(s);
     }
 

src/IRMatch.h

@@ -451,6 +451,15 @@ struct WildConst {
         return make_const_expr(val, type);
     }
 
+    // The matched value itself, no IR built. Integer constants only.
+    HALIDE_ALWAYS_INLINE
+    int64_t bound_const_int(MatcherState &state) const noexcept {
+        halide_scalar_value_t val;
+        Type type;
+        state.get_bound_const(i, val, type);
+        return val.u.i64;
+    }
+
     constexpr static bool foldable = true;
 
     [[nodiscard]] HALIDE_ALWAYS_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const noexcept {
@@ -490,6 +499,13 @@ struct Wild {
         return state.get_binding(i);
     }
 
+    // The bound node itself. Unlike make() this doesn't even touch a reference
+    // count, which lets predicates inspect what matched for free.
+    HALIDE_ALWAYS_INLINE
+    const BaseExprNode *bound_node(MatcherState &state) const noexcept {
+        return state.get_binding(i);
+    }
+
     constexpr static bool foldable = false;
 };
 
@@ -549,6 +565,12 @@ struct IntLiteral {
         return v == b.v;
     }
 
+    // The literal value itself, no IR built.
+    HALIDE_ALWAYS_INLINE
+    int64_t bound_const_int(MatcherState &state) const noexcept {
+        return v;
+    }
+
     HALIDE_ALWAYS_INLINE
     Expr make(MatcherState &state, Type type_hint) const {
         return make_const(type_hint, v);
@@ -2554,7 +2576,7 @@ struct CanProve {
     // Includes a raw call to an inlined make method, so don't inline.
     [[nodiscard]] HALIDE_NEVER_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const {
         Expr condition = a.make(state, {});
-        condition = prover->mutate(condition, nullptr);
+        condition = prover->simplify_can_prove_condition(condition);
         val.u.u64 = is_const_one(condition);
         ty = Bool(condition.type().lanes());
         return false;
@@ -2573,6 +2595,115 @@ std::ostream &operator<<(std::ostream &s, const CanProve<A, Prover> &op) {
     return s;
 }
 
+// Detects patterns that can hand back the node they matched without building
+// anything. The predicates below are restricted to these, which is what makes
+// them allocation-free: it is a compile error to ask about a derived expression
+// like min_diff(x, y + 1). Put the offset on the other side of the comparison
+// instead: min_diff(x, y) >= 1.
+template<typename A, typename = void>
+struct has_bound_node : std::false_type {};
+
+template<typename A>
+struct has_bound_node<A, std::void_t<decltype(std::declval<const A &>().bound_node(std::declval<MatcherState &>()))>>
+    : std::true_type {};
+
+// As has_bound_node, for terms whose constant reads out as a plain int64_t.
+template<typename A, typename = void>
+struct has_bound_const_int : std::false_type {};
+
+template<typename A>
+struct has_bound_const_int<A, std::void_t<decltype(std::declval<const A &>().bound_const_int(std::declval<MatcherState &>()))>>
+    : std::true_type {};
+
+// Bounds on the linear combination (ca * a - cb * b) of two matched
+// expressions, derived from the facts the prover has learned. Used as
+// (min_diff(x, y, this) >= 0) and friends. ca and cb are constants already in
+// hand (matched WildConsts, typically) that sit outside a and b's own IR, so
+// peeling can't find them. Allocation-free: ca/cb read as raw ints, a/b as raw
+// bound nodes. When nothing is known the fold reports overflow, which the
+// rewriter already treats as a failed predicate, so the rule simply doesn't
+// fire.
+template<typename A, typename CA, typename B, typename CB, typename Prover, bool is_min>
+struct LinearDiffBound {
+    struct pattern_tag {};
+    A a;
+    CA ca;
+    B b;
+    CB cb;
+    Prover *prover;
+
+    static_assert(has_bound_node<A>::value && has_bound_node<B>::value,
+                  "The a/b operands of min_diff/max_diff must be wildcards, so "
+                  "that testing the predicate doesn't have to construct any IR.");
+    static_assert(has_bound_const_int<CA>::value && has_bound_const_int<CB>::value,
+                  "The coefficient operands of linear_min_diff/linear_max_diff "
+                  "must be WildConsts or integer literals.");
+
+    constexpr static uint32_t binds = bindings<A>::mask | bindings<CA>::mask | bindings<B>::mask | bindings<CB>::mask;
+
+    // An integer-valued term of a comparison.
+    constexpr static IRNodeType min_node_type = IRNodeType::IntImm;
+    constexpr static IRNodeType max_node_type = IRNodeType::IntImm;
+    constexpr static bool canonical = true;
+
+    constexpr static bool foldable = true;
+
+    [[nodiscard]] HALIDE_ALWAYS_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const noexcept {
+        int64_t result = 0;
+        bool known;
+        if (is_min) {
+            known = prover->known_min_diff(a.bound_node(state), ca.bound_const_int(state),
+                                           b.bound_node(state), cb.bound_const_int(state), &result);
+        } else {
+            known = prover->known_max_diff(a.bound_node(state), ca.bound_const_int(state),
+                                           b.bound_node(state), cb.bound_const_int(state), &result);
+        }
+        val.u.i64 = result;
+        ty = Int(64);
+        // An unknown bound reports as overflow, failing the predicate.
+        return !known;
+    }
+};
+
+template<typename A, typename CA, typename B, typename CB, typename Prover>
+HALIDE_ALWAYS_INLINE auto linear_min_diff(A &&a, CA &&ca, B &&b, CB &&cb, Prover *p) noexcept
+    -> LinearDiffBound<decltype(pattern_arg(a)), decltype(pattern_arg(ca)), decltype(pattern_arg(b)), decltype(pattern_arg(cb)), Prover, true> {
+    assert_is_lvalue_if_expr<A>();
+    assert_is_lvalue_if_expr<B>();
+    return {pattern_arg(a), pattern_arg(ca), pattern_arg(b), pattern_arg(cb), p};
+}
+
+template<typename A, typename CA, typename B, typename CB, typename Prover>
+HALIDE_ALWAYS_INLINE auto linear_max_diff(A &&a, CA &&ca, B &&b, CB &&cb, Prover *p) noexcept
+    -> LinearDiffBound<decltype(pattern_arg(a)), decltype(pattern_arg(ca)), decltype(pattern_arg(b)), decltype(pattern_arg(cb)), Prover, false> {
+    assert_is_lvalue_if_expr<A>();
+    assert_is_lvalue_if_expr<B>();
+    return {pattern_arg(a), pattern_arg(ca), pattern_arg(b), pattern_arg(cb), p};
+}
+
+// Bounds on the plain difference (a - b).
+template<typename A, typename B, typename Prover>
+HALIDE_ALWAYS_INLINE auto min_diff(A &&a, B &&b, Prover *p) noexcept
+    -> LinearDiffBound<decltype(pattern_arg(a)), IntLiteral, decltype(pattern_arg(b)), IntLiteral, Prover, true> {
+    assert_is_lvalue_if_expr<A>();
+    assert_is_lvalue_if_expr<B>();
+    return {pattern_arg(a), IntLiteral{1}, pattern_arg(b), IntLiteral{1}, p};
+}
+
+template<typename A, typename B, typename Prover>
+HALIDE_ALWAYS_INLINE auto max_diff(A &&a, B &&b, Prover *p) noexcept
+    -> LinearDiffBound<decltype(pattern_arg(a)), IntLiteral, decltype(pattern_arg(b)), IntLiteral, Prover, false> {
+    assert_is_lvalue_if_expr<A>();
+    assert_is_lvalue_if_expr<B>();
+    return {pattern_arg(a), IntLiteral{1}, pattern_arg(b), IntLiteral{1}, p};
+}
+
+template<typename A, typename CA, typename B, typename CB, typename Prover, bool is_min>
+std::ostream &operator<<(std::ostream &s, const LinearDiffBound<A, CA, B, CB, Prover, is_min> &op) {
+    s << (is_min ? "linear_min_diff(" : "linear_max_diff(") << op.a << ", " << op.ca << ", " << op.b << ", " << op.cb << ")";
+    return s;
+}
+
 template<typename A>
 struct IsFloat {
     struct pattern_tag {};