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

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,196 @@ 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;
+}
+
+// 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 {};
+
+// Bounds on the difference between two matched expressions, derived from the
+// facts the prover has learned. Used as (min_diff(x, y, this) >= 0) and
+// friends. 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 B, typename Prover, bool is_min>
+struct DiffBound {
+    struct pattern_tag {};
+    A a;
+    B b;
+    Prover *prover;
+
+    static_assert(has_bound_node<A>::value && has_bound_node<B>::value,
+                  "The operands of min_diff/max_diff must be wildcards, so that "
+                  "testing the predicate doesn't have to construct any IR.");
+
+    constexpr static uint32_t binds = bindings<A>::mask | bindings<B>::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), b.bound_node(state), &result);
+        } else {
+            known = prover->known_max_diff(a.bound_node(state), b.bound_node(state), &result);
+        }
+        val.u.i64 = result;
+        ty = Int(64);
+        // An unknown bound reports as overflow, failing the predicate.
+        return !known;
+    }
+};
+
+template<typename A, typename B, typename Prover>
+HALIDE_ALWAYS_INLINE auto min_diff(A &&a, B &&b, Prover *p) noexcept
+    -> DiffBound<decltype(pattern_arg(a)), decltype(pattern_arg(b)), Prover, true> {
+    assert_is_lvalue_if_expr<A>();
+    assert_is_lvalue_if_expr<B>();
+    return {pattern_arg(a), pattern_arg(b), p};
+}
+
+template<typename A, typename B, typename Prover>
+HALIDE_ALWAYS_INLINE auto max_diff(A &&a, B &&b, Prover *p) noexcept
+    -> DiffBound<decltype(pattern_arg(a)), decltype(pattern_arg(b)), Prover, false> {
+    assert_is_lvalue_if_expr<A>();
+    assert_is_lvalue_if_expr<B>();
+    return {pattern_arg(a), pattern_arg(b), p};
+}
+
+template<typename A, typename B, typename Prover, bool is_min>
+std::ostream &operator<<(std::ostream &s, const DiffBound<A, B, Prover, is_min> &op) {
+    s << (is_min ? "min_diff(" : "max_diff(") << op.a << ", " << op.b << ")";
+    return s;
+}
+
+// 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 {};
+
+// As DiffBound, but for the affine combination (ca * a - cb * b), where 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.
+template<typename A, typename CA, typename B, typename CB, typename Prover, bool is_min>
+struct ScaledDiffBound {
+    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 scaled_min_diff/scaled_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 scaled_min_diff/scaled_max_diff "
+                  "must be WildConsts.");
+
+    constexpr static uint32_t binds = bindings<A>::mask | bindings<CA>::mask | bindings<B>::mask | bindings<CB>::mask;
+
+    // This is 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);
+        // Report an unknown bound as an overflow, which fails the predicate.
+        return !known;
+    }
+};
+
+template<typename A, typename CA, typename B, typename CB, typename Prover>
+HALIDE_ALWAYS_INLINE auto scaled_min_diff(A &&a, CA &&ca, B &&b, CB &&cb, Prover *p) noexcept
+    -> ScaledDiffBound<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 scaled_max_diff(A &&a, CA &&ca, B &&b, CB &&cb, Prover *p) noexcept
+    -> ScaledDiffBound<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};
+}
+
+template<typename A, typename CA, typename B, typename CB, typename Prover, bool is_min>
+std::ostream &operator<<(std::ostream &s, const ScaledDiffBound<A, CA, B, CB, Prover, is_min> &op) {
+    s << (is_min ? "scaled_min_diff(" : "scaled_max_diff(") << op.a << ", " << op.ca << ", " << op.b << ", " << op.cb << ")";
+    return s;
+}
+
 template<typename A>
 struct IsFloat {
     struct pattern_tag {};

src/Lower.cpp

@@ -305,6 +305,11 @@ void lower_impl(const vector<Function> &output_funcs,
     s = storage_flattening(s, outputs, env, t);
     log("Lowering after storage flattening:", s);
 
+    // Every pass that reads a region or an allocation size out of the IR has
+    // now run, so from here a clamp is only worth what its value is worth, and
+    // the simplifier may use what it knows to remove a redundant one.
+    ScopedRegionsInferred regions_inferred;
+
     debug(1) << "Adding atomic mutex allocation...\n";
     s = add_atomic_mutex(s, outputs);
     log("Lowering after adding atomic mutex allocation:", s);