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

src/IRMatch.h

@@ -2554,7 +2554,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 +2573,46 @@ std::ostream &operator<<(std::ostream &s, const CanProve<A, Prover> &op) {
     return s;
 }
 
+// Like can_prove, but only looks the condition up in the facts the prover
+// already knows, instead of recursively invoking it. Much cheaper, and it
+// cannot recurse, so unlike can_prove it is safe in a rule whose left-hand
+// side matches expressions the prover may construct while proving it.
+template<typename A, typename Prover>
+struct KnownTrue {
+    struct pattern_tag {};
+    A a;
+    Prover *prover;  // An existing simplifying mutator
+
+    constexpr static uint32_t binds = bindings<A>::mask;
+
+    // This rule is a boolean-valued predicate. Bools have type UIntImm.
+    constexpr static IRNodeType min_node_type = IRNodeType::UIntImm;
+    constexpr static IRNodeType max_node_type = IRNodeType::UIntImm;
+    constexpr static bool canonical = true;
+
+    constexpr static bool foldable = true;
+
+    // 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, {});
+        val.u.u64 = prover->is_known_true(condition) ? 1 : 0;
+        ty = Bool(condition.type().lanes());
+        return false;
+    }
+};
+
+template<typename A, typename Prover>
+HALIDE_ALWAYS_INLINE auto known_true(A &&a, Prover *p) noexcept -> KnownTrue<decltype(pattern_arg(a)), Prover> {
+    assert_is_lvalue_if_expr<A>();
+    return {pattern_arg(a), p};
+}
+
+template<typename A, typename Prover>
+std::ostream &operator<<(std::ostream &s, const KnownTrue<A, Prover> &op) {
+    s << "known_true(" << op.a << ")";
+    return s;
+}
+
 template<typename A>
 struct IsFloat {
     struct pattern_tag {};

src/Simplify.cpp

@@ -85,6 +85,16 @@ void Simplify::found_buffer_reference(const string &name, size_t dimensions) {
 }
 
 void Simplify::ScopedFact::learn_false(const Expr &fact) {
+    // Canonicalize the direction of comparisons, so that facts are stored in
+    // the same form the simplifier produces when it visits them.
+    if (const GT *gt = fact.as<GT>()) {
+        learn_false(gt->b < gt->a);
+        return;
+    } else if (const GE *ge = fact.as<GE>()) {
+        learn_false(!(ge->a < ge->b));
+        return;
+    }
+
     Simplify::VarInfo info;
     info.old_uses = info.new_uses = 0;
     if (const Variable *v = fact.as<Variable>()) {
@@ -172,6 +182,16 @@ void Simplify::ScopedFact::learn_lower_bound(const Variable *v, int64_t val) {
 }
 
 void Simplify::ScopedFact::learn_true(const Expr &fact) {
+    // Canonicalize the direction of comparisons, so that facts are stored in
+    // the same form the simplifier produces when it visits them.
+    if (const GT *gt = fact.as<GT>()) {
+        learn_true(gt->b < gt->a);
+        return;
+    } else if (const GE *ge = fact.as<GE>()) {
+        learn_true(!(ge->a < ge->b));
+        return;
+    }
+
     Simplify::VarInfo info;
     info.old_uses = info.new_uses = 0;
     if (const Variable *v = fact.as<Variable>()) {
@@ -345,16 +365,56 @@ void Simplify::ScopedFact::learn_true(const Expr &fact) {
 }
 
 namespace {
+// Is a boolean Expr known to be true or false? Facts are stored in the same
+// form the simplifier itself produces, so a comparison has to be canonicalized
+// the same way before looking it up.
+std::optional<bool> lookup_fact(const Expr &e,
+                                const std::set<Expr, IRDeepCompare> &truths,
+                                const std::set<Expr, IRDeepCompare> &falsehoods) {
+    if (const Not *n = e.as<Not>()) {
+        auto known = lookup_fact(n->a, truths, falsehoods);
+        return known ? std::make_optional(!*known) : known;
+    } else if (const GT *gt = e.as<GT>()) {
+        return lookup_fact(gt->b < gt->a, truths, falsehoods);
+    } else if (const GE *ge = e.as<GE>()) {
+        return lookup_fact(!(ge->a < ge->b), truths, falsehoods);
+    }
+
+    if (truths.count(e)) {
+        return true;
+    } else if (falsehoods.count(e)) {
+        return false;
+    }
+
+    // A comparison may also be settled by the other strictness of the same
+    // comparison, in either direction.
+    if (const LT *lt = e.as<LT>()) {
+        // a < b is implied by !(b <= a), and ruled out by b <= a and by b < a.
+        if (falsehoods.count(lt->b <= lt->a)) {
+            return true;
+        } else if (truths.count(lt->b <= lt->a) || truths.count(lt->b < lt->a)) {
+            return false;
+        }
+    } else if (const LE *le = e.as<LE>()) {
+        // a <= b is implied by a < b and by !(b < a), and ruled out by b < a.
+        if (truths.count(le->a < le->b) || falsehoods.count(le->b < le->a)) {
+            return true;
+        } else if (truths.count(le->b < le->a)) {
+            return false;
+        }
+    }
+
+    return std::nullopt;
+}
+
 template<typename T>
 T substitute_facts_impl(const T &t,
                         const std::set<Expr, IRDeepCompare> &truths,
                         const std::set<Expr, IRDeepCompare> &falsehoods) {
     return mutate_with(t, [&](auto *self, const Expr &e) {
         if (e.type().is_bool()) {
-            if (truths.count(e)) {
-                return make_one(e.type());
-            } else if (falsehoods.count(e)) {
-                return make_zero(e.type());
+            if (auto known = lookup_fact(e, truths, falsehoods)) {
+                return *known ? make_one(e.type()) : make_zero(e.type());
             }
         }
         return self->mutate_base(e);
@@ -370,6 +430,39 @@ Stmt Simplify::ScopedFact::substitute_facts(const Stmt &s) {
     return substitute_facts_impl(s, truths, falsehoods);
 }
 
+bool Simplify::is_known_true(const Expr &e) {
+    if (truths.empty() && falsehoods.empty()) {
+        return false;
+    }
+    auto known = lookup_fact(e, truths, falsehoods);
+    return known && *known;
+}
+
+Expr Simplify::simplify_can_prove_condition(const Expr &e) {
+    if (can_prove_depth >= max_can_prove_depth) {
+        // Too deep to safely recurse into the full simplifier. The only thing
+        // the caller does with the result is check whether it is the literal
+        // constant true, and nothing here can fold a compound expression (an
+        // And of two known-true operands stays an unfolded And, not true) --
+        // that folding is exactly the recursive work we're declining to do.
+        // So a substitute_facts tree walk can't prove anything a direct
+        // lookup of the condition itself couldn't already: skip the walk.
+        if (is_known_true(e)) {
+            return const_true(e.type().lanes(), nullptr);
+        }
+        return e;
+    }
+    ScopedValue<int> guard(can_prove_depth, can_prove_depth + 1);
+    return mutate(substitute_facts(e), nullptr);
+}
+
+Expr Simplify::substitute_facts(const Expr &e) {
+    if (truths.empty() && falsehoods.empty()) {
+        return e;
+    }
+    return substitute_facts_impl(e, truths, falsehoods);
+}
+
 Simplify::ScopedFact::~ScopedFact() {
     for (const auto *v : pop_list) {
         simplify->var_info.pop(v->name);

src/Simplify_Div.cpp

@@ -84,6 +84,18 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) {
          rewrite(select(x, c0, c1) / c2, select(x, fold(c0 / c2), fold(c1 / c2))) ||
          (!op->type.is_float() &&
           rewrite(x / x, select(x == 0, 0, 1))) ||
+
+         (no_overflow(op->type) &&
+          // Facts learned higher up in the IR may tell us which side of a max
+          // or min survives the division. Test them early on to prevents rewrites below
+          // that would make it impossible to recognize the form.
+          (has_facts() &&
+           (rewrite(max(x * c0, y) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) ||
+            rewrite(max(y, x * c0) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) ||
+            rewrite(min(x * c0, y) / c0, x, c0 > 0 && known_true(x <= y / c0, this)) ||
+            rewrite(min(y, x * c0) / c0, x, c0 > 0 && known_true(x <= y / c0, this)) ||
+            false))) ||
+
          (no_overflow(op->type) &&
           // Fold repeated division
           (rewrite((x / c0) / c2, x / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2)) ||

src/Simplify_Internal.h

@@ -441,6 +441,35 @@ class Simplify : public VariadicVisitor<Simplify, Expr, Stmt> {
 
     std::set<Expr, IRDeepCompare> truths, falsehoods;
 
+    // How deeply are we nested inside the conditions of can_prove predicates?
+    // Proving such a condition recursively invokes the simplifier on it, so a
+    // rule whose left-hand side also matches something built while proving its
+    // own predicate recurses without bound. Nesting is also expensive, and no
+    // rule currently relies on it. Bound it.
+    int can_prove_depth = 0;
+    static constexpr int max_can_prove_depth = 4;
+
+    // Is there anything in the truths/falsehoods sets that a rewrite rule could
+    // use? Used to gate rules whose predicates are only ever provable from facts
+    // learned higher up in the IR, so that we don't pay for them in the common
+    // case.
+    bool has_facts() const {
+        return !truths.empty() || !falsehoods.empty();
+    }
+
+    // Replace exprs known to be truths or falsehoods with const_true or
+    // const_false. Used to inject everything currently known into the
+    // conditions of can_prove predicates in rewrite rules.
+    Expr substitute_facts(const Expr &e);
+
+    // Simplify the condition of a can_prove predicate in a rewrite rule, using
+    // everything currently known.
+    Expr simplify_can_prove_condition(const Expr &e);
+
+    // Is a boolean Expr already known to be true? Unlike can_prove this only
+    // looks the condition up in the facts, without simplifying anything.
+    bool is_known_true(const Expr &e);
+
     struct ScopedFact {
         Simplify *simplify;
 

src/Simplify_Max.cpp

@@ -71,6 +71,10 @@ Expr Simplify::visit(const Max *op, ExprInfo *info) {
     // RHS for ExprInfo to update correctly.
     if (EVAL_IN_LAMBDA  //
         (rewrite(max(x, x), a) ||
+         // Facts learned higher up in the IR may tell us which side wins.
+         (has_facts() &&
+          (rewrite(max(x, y), a, known_true(y <= x, this)) ||
+           rewrite(max(x, y), b, known_true(x <= y, this)))) ||
          rewrite(max(x, c0), b, is_max_value(c0)) ||
          rewrite(max(x, c0), a, is_min_value(c0)) ||
          rewrite(max((x / c0) * c0, x), b, c0 > 0) ||

src/Simplify_Min.cpp

@@ -70,6 +70,10 @@ Expr Simplify::visit(const Min *op, ExprInfo *info) {
     // RHS for ExprInfo to update correctly.
     if (EVAL_IN_LAMBDA  //
         (rewrite(min(x, x), a) ||
+         // Facts learned higher up in the IR may tell us which side wins.
+         (has_facts() &&
+          (rewrite(min(x, y), a, known_true(x <= y, this)) ||
+           rewrite(min(x, y), b, known_true(y <= x, this)))) ||
          rewrite(min(x, c0), b, is_min_value(c0)) ||
          rewrite(min(x, c0), a, is_max_value(c0)) ||
          rewrite(min((x / c0) * c0, x), a, c0 > 0) ||

test/correctness/simplify.cpp

@@ -2377,6 +2377,18 @@ void check_invariant() {
     }
 }
 
+void check_with_assumptions(const Expr &a, const Expr &b, const std::vector<Expr> &assumptions) {
+    Expr simpler = simplify(a, Scope<Interval>(), Scope<ModulusRemainder>(), assumptions);
+    if (!equal(simpler, b)) {
+        std::cerr
+            << "\nSimplification failure:\n"
+            << "Input: " << a << "\n"
+            << "Output: " << simpler << "\n"
+            << "Expected output: " << b << "\n";
+        abort();
+    }
+}
+
 void check_unreachable() {
     Var x("x"), y("y");
 
@@ -2405,6 +2417,95 @@ void check_unreachable() {
           Evaluate::make(0));
 }
 
+void check_facts() {
+    Expr x = Var("x"), y = Var("y"), z = Var("z");
+
+    // A fact stated in any comparison direction should let the simplifier pick
+    // the winning side of a max or min.
+    check_with_assumptions(max(x, y), x, {x > y});
+    check_with_assumptions(max(x, y), x, {y < x});
+    check_with_assumptions(max(x, y), y, {x < y});
+    check_with_assumptions(max(x, y), y, {y > x});
+    check_with_assumptions(min(x, y), y, {x > y});
+    check_with_assumptions(min(x, y), x, {x < y});
+
+    // A non-strict fact is enough to pick a side of a max or min, and a strict
+    // fact implies the non-strict one.
+    check_with_assumptions(max(x, y), x, {x >= y});
+    check_with_assumptions(max(x, y), y, {x <= y});
+    check_with_assumptions(min(x, y), x, {x <= y});
+    check_with_assumptions(min(x, y), y, {x >= y});
+
+    // Facts about compound expressions work too.
+    check_with_assumptions(max(x + z, y * 3), x + z, {x + z > y * 3});
+    check_with_assumptions(max(max(x, y), z), z, {max(x, y) < z});
+
+    // Both branches of an if learn from the condition, in opposite directions.
+    check(IfThenElse::make(x < y, not_no_op(max(x, y)), not_no_op(max(x, y))),
+          IfThenElse::make(x < y, not_no_op(y), not_no_op(x)));
+
+    // A fact only applies where it holds.
+    check(Block::make(not_no_op(max(x, y)),
+                      IfThenElse::make(x < y, not_no_op(max(x, y)))),
+          Block::make(not_no_op(max(x, y)),
+                      IfThenElse::make(x < y, not_no_op(y))));
+
+    // A division can cancel a multiplication inside a max or min when we know
+    // which side wins after the division.
+    check_with_assumptions(max(x * 8, y) / 8, x, {x >= y / 8});
+    check_with_assumptions(max(y, x * 8) / 8, x, {x >= y / 8});
+    check_with_assumptions(min(x * 8, y) / 8, x, {x <= y / 8});
+    check_with_assumptions(min(y, x * 8) / 8, x, {x <= y / 8});
+
+    // The direction in which a fact is stated doesn't matter, on either side:
+    // both the facts and the conditions of can_prove predicates are looked up
+    // in the same canonical form.
+    check_with_assumptions(max(x * 8, y) / 8, x, {y / 8 <= x});
+    check_with_assumptions(max(x * 8, y) / 8, x, {!(x < y / 8)});
+    check_with_assumptions(min(x * 8, y) / 8, x, {y / 8 >= x});
+
+    // A strict fact settles a non-strict predicate too.
+    check_with_assumptions(max(x * 8, y) / 8, x, {x > y / 8});
+    check_with_assumptions(min(x * 8, y) / 8, x, {x < y / 8});
+
+    // Deeply nested mins and maxes must not make the work of proving the
+    // predicates of the rules above blow up.
+    Expr nest = x;
+    for (int i = 0; i < 24; i++) {
+        nest = min(max(nest + i, y - i), z * i);
+    }
+    // The result isn't interesting; what matters is that we get one at all.
+    (void)simplify(nest, Scope<Interval>(), Scope<ModulusRemainder>(), {x < y});
+
+    // can_prove-based rules (unlike the known_true ones above) recursively
+    // invoke the simplifier on their own predicate, and that predicate can be
+    // a freshly built expression rather than a piece of the original IR (e.g.
+    // min(x, y) - min(z, w) -> y - w, can_prove(x - y == z - w)) constructs a
+    // brand new subtraction). If the operands are themselves unsimplified
+    // instances of the same shape, this recurses; the depth limit must bound
+    // the work rather than let it explode.
+    Expr deep = min(Var("da"), Var("db")) - min(Var("dc"), Var("dd"));
+    for (int i = 0; i < 10; i++) {
+        Expr y = Var("dy" + std::to_string(i));
+        Expr z = Var("dz" + std::to_string(i));
+        Expr w = Var("dw" + std::to_string(i));
+        deep = min(deep, y) - min(z, w);
+    }
+    (void)simplify(deep);
+
+    // The rules above look their predicates up in the facts rather than
+    // recursively invoking the simplifier, so a fact only settles a predicate
+    // it is directly comparable to. This one needs arithmetic to connect:
+    check_with_assumptions(max(x, y), max(x, y), {x + 1 <= y});
+
+    // Without the fact, the division stays put.
+    check(max(x * 8, y) / 8, max(x * 8, y) / 8);
+
+    // Facts that don't strictly order the operands don't fire these rules.
+    check_with_assumptions(max(x, y), max(x, y), {x != y});
+    check_with_assumptions(max(x * 8, y) / 8, max(x * 8, y) / 8, {x < y / 8});
+}
+
 int main(int argc, char **argv) {
     check_invariant();
     check_casts();
@@ -2417,6 +2518,7 @@ int main(int argc, char **argv) {
     check_bitwise();
     check_lets();
     check_unreachable();
+    check_facts();
 
     // Miscellaneous cases that don't fit into one of the categories above.
     Expr x = Var("x"), y = Var("y");