diff --git a/src/IRMatch.h b/src/IRMatch.h index 6fa4cadc4eae..1cc69abc345b 100644 --- a/src/IRMatch.h +++ b/src/IRMatch.h @@ -490,6 +490,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; }; @@ -2554,7 +2561,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 +2580,175 @@ std::ostream &operator<<(std::ostream &s, const CanProve &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 +struct KnownTrue { + struct pattern_tag {}; + A a; + Prover *prover; // An existing simplifying mutator + + constexpr static uint32_t binds = bindings::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 +HALIDE_ALWAYS_INLINE auto known_true(A &&a, Prover *p) noexcept -> KnownTrue { + assert_is_lvalue_if_expr(); + return {pattern_arg(a), p}; +} + +template +std::ostream &operator<<(std::ostream &s, const KnownTrue &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 +struct has_bound_node : std::false_type {}; + +template +struct has_bound_node().bound_node(std::declval()))>> + : 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 +struct DiffBound { + struct pattern_tag {}; + A a; + B b; + Prover *prover; + + static_assert(has_bound_node::value && has_bound_node::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::mask | bindings::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), 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); + // Report an unknown bound as an overflow, which fails the predicate. + return !known; + } +}; + +template +HALIDE_ALWAYS_INLINE auto min_diff(A &&a, B &&b, Prover *p) noexcept + -> DiffBound { + assert_is_lvalue_if_expr(); + assert_is_lvalue_if_expr(); + return {pattern_arg(a), pattern_arg(b), p}; +} + +template +HALIDE_ALWAYS_INLINE auto max_diff(A &&a, B &&b, Prover *p) noexcept + -> DiffBound { + assert_is_lvalue_if_expr(); + assert_is_lvalue_if_expr(); + return {pattern_arg(a), pattern_arg(b), p}; +} + +template +std::ostream &operator<<(std::ostream &s, const DiffBound &op) { + s << (is_min ? "min_diff(" : "max_diff(") << op.a << ", " << op.b << ")"; + return s; +} + +// Do the facts say these two are equal, or that they differ? Equality is just a +// difference of zero, but inequality is a hole in the difference rather than a +// bound on it, so it gets its own predicate. +template +struct KnownComparison { + struct pattern_tag {}; + A a; + B b; + Prover *prover; + + static_assert(has_bound_node::value && has_bound_node::value, + "The operands of known_equal/known_not_equal must be wildcards, " + "so that testing the predicate doesn't have to construct any IR."); + + constexpr static uint32_t binds = bindings::mask | bindings::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; + + [[nodiscard]] HALIDE_ALWAYS_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const noexcept { + if (want_equal) { + val.u.u64 = prover->is_known_equal(a.bound_node(state), b.bound_node(state)) ? 1 : 0; + } else { + val.u.u64 = prover->is_known_not_equal(a.bound_node(state), b.bound_node(state)) ? 1 : 0; + } + ty = Bool(); + return false; + } +}; + +template +HALIDE_ALWAYS_INLINE auto known_equal(A &&a, B &&b, Prover *p) noexcept + -> KnownComparison { + assert_is_lvalue_if_expr(); + assert_is_lvalue_if_expr(); + return {pattern_arg(a), pattern_arg(b), p}; +} + +template +HALIDE_ALWAYS_INLINE auto known_not_equal(A &&a, B &&b, Prover *p) noexcept + -> KnownComparison { + assert_is_lvalue_if_expr(); + assert_is_lvalue_if_expr(); + return {pattern_arg(a), pattern_arg(b), p}; +} + +template +std::ostream &operator<<(std::ostream &s, const KnownComparison &op) { + s << (want_equal ? "known_equal(" : "known_not_equal(") << op.a << ", " << op.b << ")"; + return s; +} + template struct IsFloat { struct pattern_tag {}; diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 18129614aca7..198f9dfd13c0 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -84,7 +84,115 @@ void Simplify::found_buffer_reference(const string &name, size_t dimensions) { } } +namespace { + +// Rewrite (a - b) as (a' - b') + offset by stripping constant terms off either +// side, so that a fact about x and y + 3 and a query about x and y meet at the +// same pair. Walks the existing nodes; builds nothing. +void peel_constant_offsets(const BaseExprNode *&a, const BaseExprNode *&b, int64_t &offset) { + // Peels one constant term off e if there is one, returning whether it did. + // The constant is added to delta, which the caller applies with the sign + // appropriate to the side e is on. + auto peel_one = [](const BaseExprNode *&e, int64_t &delta) { + if (e->node_type == IRNodeType::Add) { + const Add *add = (const Add *)e; + if (const IntImm *i = add->b.as()) { + if (add_would_overflow(64, delta, i->value)) { + return false; + } + delta += i->value; + e = add->a.get(); + return true; + } else if (const IntImm *i = add->a.as()) { + if (add_would_overflow(64, delta, i->value)) { + return false; + } + delta += i->value; + e = add->b.get(); + return true; + } + } else if (e->node_type == IRNodeType::Sub) { + const Sub *sub = (const Sub *)e; + if (const IntImm *i = sub->b.as()) { + if (sub_would_overflow(64, delta, i->value)) { + return false; + } + delta -= i->value; + e = sub->a.get(); + return true; + } + } + return false; + }; + + // A constant on the left of the difference adds to the offset; one on the + // right subtracts from it, so accumulate it negated and subtract at the end. + int64_t from_a = 0, from_b = 0; + while (peel_one(a, from_a)) { + } + while (peel_one(b, from_b)) { + } + if (!sub_would_overflow(64, from_a, from_b)) { + offset = from_a - from_b; + } else { + offset = 0; + } +} + +} // namespace + +void Simplify::ScopedFact::learn_difference(const Expr &a, const Expr &b, + const ConstantInterval &diff, bool invert) { + // Differences are only meaningful where they can't wrap. + if (!simplify->no_overflow_int(a.type()) || a.type() != b.type()) { + return; + } + + const BaseExprNode *pa = a.get(), *pb = b.get(); + int64_t offset = 0; + peel_constant_offsets(pa, pb, offset); + + // (a - b) = (pa - pb) + offset, so the bound on the peeled pair is the + // bound we were given shifted the other way. + ConstantInterval peeled = diff - offset; + if (invert && !peeled.is_single_point()) { + // Only a single removed point is representable. + return; + } + + const uint32_t fa = Simplify::expr_fingerprint(pa), fb = Simplify::expr_fingerprint(pb); + simplify->add_difference_key(Simplify::difference_key(fa, fb)); + simplify->known_bounds.push_back( + Simplify::KnownBound{Expr(pa), Expr(pb), peeled, fa, fb, invert}); +} + 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()) { + learn_false(gt->b < gt->a); + return; + } else if (const GE *ge = fact.as()) { + learn_false(!(ge->a < ge->b)); + return; + } + + // Record what this says about the difference between the two sides. And, + // Not, and the tag intrinsic are handled by the recursion below instead. + if (const LT *lt = fact.as()) { + // !(a < b) -> a - b >= 0 + learn_difference(lt->a, lt->b, ConstantInterval::bounded_below(0), false); + } else if (const LE *le = fact.as()) { + // !(a <= b) -> a - b >= 1 + learn_difference(le->a, le->b, ConstantInterval::bounded_below(1), false); + } else if (const EQ *eq = fact.as()) { + // !(a == b) -> a - b is anything but zero + learn_difference(eq->a, eq->b, ConstantInterval::single_point(0), true); + } else if (const NE *ne = fact.as()) { + // !(a != b) -> a - b == 0 + learn_difference(ne->a, ne->b, ConstantInterval::single_point(0), false); + } + Simplify::VarInfo info; info.old_uses = info.new_uses = 0; if (const Variable *v = fact.as()) { @@ -172,6 +280,32 @@ 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()) { + learn_true(gt->b < gt->a); + return; + } else if (const GE *ge = fact.as()) { + learn_true(!(ge->a < ge->b)); + return; + } + + // Record what this says about the difference between the two sides. And, + // Not, and the tag intrinsic are handled by the recursion below instead. + if (const LT *lt = fact.as()) { + // a < b -> a - b <= -1 + learn_difference(lt->a, lt->b, ConstantInterval::bounded_above(-1), false); + } else if (const LE *le = fact.as()) { + // a <= b -> a - b <= 0 + learn_difference(le->a, le->b, ConstantInterval::bounded_above(0), false); + } else if (const EQ *eq = fact.as()) { + // a == b -> a - b == 0 + learn_difference(eq->a, eq->b, ConstantInterval::single_point(0), false); + } else if (const NE *ne = fact.as()) { + // a != b -> a - b is anything but zero + learn_difference(ne->a, ne->b, ConstantInterval::single_point(0), true); + } + Simplify::VarInfo info; info.old_uses = info.new_uses = 0; if (const Variable *v = fact.as()) { @@ -345,16 +479,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 lookup_fact(const Expr &e, + const std::set &truths, + const std::set &falsehoods) { + if (const Not *n = e.as()) { + auto known = lookup_fact(n->a, truths, falsehoods); + return known ? std::make_optional(!*known) : known; + } else if (const GT *gt = e.as()) { + return lookup_fact(gt->b < gt->a, truths, falsehoods); + } else if (const GE *ge = e.as()) { + 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()) { + // 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()) { + // 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 T substitute_facts_impl(const T &t, const std::set &truths, const std::set &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,13 +544,227 @@ Stmt Simplify::ScopedFact::substitute_facts(const Stmt &s) { return substitute_facts_impl(s, truths, falsehoods); } +namespace { + +// Intersect acc with d, reporting whether the result would be empty rather than +// constructing it. make_intersection asserts on an empty result, and empty means +// the facts contradict each other, which means this code is unreachable. We +// don't try to exploit that here; we just decline to tighten any further. +bool intersect_if_nonempty(ConstantInterval &acc, const ConstantInterval &d) { + ConstantInterval result = acc; + if (d.min_defined && (!result.min_defined || d.min > result.min)) { + result.min = d.min; + result.min_defined = true; + } + if (d.max_defined && (!result.max_defined || d.max < result.max)) { + result.max = d.max; + result.max_defined = true; + } + if (result.min_defined && result.max_defined && result.min > result.max) { + return false; + } + acc = result; + return true; +} + +// What the shape of the two sides says about (a - b) on its own, with no facts +// involved: a min is at most either of its operands, and a max is at least +// either of them. Only the immediate operands are inspected, so this stays a +// couple of pointer comparisons rather than a search. +ConstantInterval structural_difference(const BaseExprNode *a, const BaseExprNode *b) { + ConstantInterval result; + + auto is_operand_of = [](const BaseExprNode *e, const BaseExprNode *node) { + if (node->node_type == IRNodeType::Min) { + const Min *m = (const Min *)node; + return equal(*m->a.get(), *e) || equal(*m->b.get(), *e); + } else if (node->node_type == IRNodeType::Max) { + const Max *m = (const Max *)node; + return equal(*m->a.get(), *e) || equal(*m->b.get(), *e); + } + return false; + }; + + // min(p, q) - b <= 0 and max(p, q) - b >= 0, when b is one of the operands. + if (a->node_type == IRNodeType::Min && is_operand_of(b, a)) { + result = ConstantInterval::bounded_above(0); + } else if (a->node_type == IRNodeType::Max && is_operand_of(b, a)) { + result = ConstantInterval::bounded_below(0); + } else if (b->node_type == IRNodeType::Min && is_operand_of(a, b)) { + // a - min(p, q) >= 0 + result = ConstantInterval::bounded_below(0); + } else if (b->node_type == IRNodeType::Max && is_operand_of(a, b)) { + result = ConstantInterval::bounded_above(0); + } + + return result; +} + +} // namespace + +Simplify::KnownDiff Simplify::known_difference(const BaseExprNode *a, const BaseExprNode *b) { + KnownDiff result; + + // Canonicalize the query the way the facts were canonicalized when learned. + int64_t offset = 0; + peel_constant_offsets(a, b, offset); + + if (equal(*a, *b)) { + result.bounds = ConstantInterval::single_point(0); + } else { + if (a->node_type == IRNodeType::IntImm && b->node_type == IRNodeType::IntImm && + !sub_would_overflow(64, ((const IntImm *)a)->value, ((const IntImm *)b)->value)) { + // Two constants need no facts to compare. + result.bounds = ConstantInterval::single_point(((const IntImm *)a)->value - + ((const IntImm *)b)->value); + } else { + intersect_if_nonempty(result.bounds, structural_difference(a, b)); + } + } + + if (!result.bounds.is_single_point() && !known_bounds.empty()) { + // A hole only tightens the bounds once we know where the ends are, so + // collect them as we go and apply them below. There are hardly ever any. + constexpr int max_holes = 4; + int64_t holes[max_holes]; + int num_holes = 0; + + const uint32_t fa = expr_fingerprint(a), fb = expr_fingerprint(b); + // One test against the whole table before looking at any record. + if (!difference_key_present(difference_key(fa, fb))) { + result.bounds += offset; + return result; + } + for (const KnownBound &kb : known_bounds) { + // Reject on the summaries first. They live in the record, so a + // record about some other pair costs a pair of integer compares + // and never follows a pointer. + const bool same_order = (fa == kb.fingerprint_a && fb == kb.fingerprint_b); + const bool swapped = (fa == kb.fingerprint_b && fb == kb.fingerprint_a); + if (!same_order && !swapped) { + continue; + } + + ConstantInterval d; + if (same_order && equal(*a, *kb.a.get()) && equal(*b, *kb.b.get())) { + d = kb.diff; + } else if (swapped && equal(*a, *kb.b.get()) && equal(*b, *kb.a.get())) { + // We know about (b - a), and this is the other direction. + d = -kb.diff; + } else { + continue; + } + + if (kb.invert) { + if (num_holes < max_holes) { + holes[num_holes++] = d.min; + } + } else if (!intersect_if_nonempty(result.bounds, d)) { + break; + } + } + + for (int i = 0; i < num_holes; i++) { + const int64_t hole = holes[i]; + // Removing a point only narrows the bounds if it is at one end. + if (result.bounds.min_defined && result.bounds.min == hole && + !add_would_overflow(64, hole, 1)) { + result.bounds.min = hole + 1; + } + if (result.bounds.max_defined && result.bounds.max == hole && + !sub_would_overflow(64, hole, 1)) { + result.bounds.max = hole - 1; + } + // Whether the difference can be zero matters even when the hole is + // in the interior, where it can't be captured by the bounds. + if (!add_would_overflow(64, hole, offset) && hole + offset == 0) { + result.excludes_zero = true; + } + } + } + + // Undo the canonicalization: (a - b) = (peeled a - peeled b) + offset. + result.bounds += offset; + + return result; +} + +bool Simplify::is_known_equal(const BaseExprNode *a, const BaseExprNode *b) { + return known_difference(a, b).bounds.is_single_point(0); +} + +bool Simplify::is_known_not_equal(const BaseExprNode *a, const BaseExprNode *b) { + KnownDiff d = known_difference(a, b); + return d.excludes_zero || !d.bounds.contains((int64_t)0); +} + +bool Simplify::known_min_diff(const BaseExprNode *a, const BaseExprNode *b, int64_t *result) { + ConstantInterval bounds = known_difference(a, b).bounds; + if (bounds.min_defined) { + *result = bounds.min; + return true; + } + return false; +} + +bool Simplify::known_max_diff(const BaseExprNode *a, const BaseExprNode *b, int64_t *result) { + ConstantInterval bounds = known_difference(a, b).bounds; + if (bounds.max_defined) { + *result = bounds.max; + return true; + } + return false; +} + +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 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() { + if (!simplify) { + // Moved from; the object that took over owns the cleanup. + return; + } for (const auto *v : pop_list) { simplify->var_info.pop(v->name); } for (const auto *v : bounds_pop_list) { simplify->bounds_and_alignment_info.pop(v->name); } + internal_assert(simplify->known_bounds.size() >= known_bounds_size); + simplify->known_bounds.resize(known_bounds_size); + for (int i = 0; i < Simplify::difference_key_words; i++) { + simplify->difference_keys[i] = saved_difference_keys[i]; + } for (const auto &e : truths) { simplify->truths.erase(e); } diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4098f6f027e7..d7839340fb99 100644 --- a/src/Simplify_Div.cpp +++ b/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)) || diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index 94a50bebc644..6ca322df95ef 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -441,30 +441,184 @@ class Simplify : public VariadicVisitor { std::set truths, falsehoods; + /** What we know about the difference between a pair of Exprs. Every + * comparison we can learn from is a statement about (a - b): a < b means it + * is at most -1, !(a < b) means it is at least 0, a == b means it is zero. + * Because the complement of a half-line is a half-line, only the negation + * of an equality fails to be an interval, and that is always a single point + * removed, which is what invert represents. */ + struct KnownBound { + Expr a, b; + ConstantInterval diff; + // Cheap structural summaries of a and b. Equal Exprs always summarize + // to the same value, so a mismatch rules a record out without touching + // the Exprs at all. Almost every query is about a pair nothing is known + // about, so what this scan needs to be good at is saying no. + uint32_t fingerprint_a = 0, fingerprint_b = 0; + // If set, a - b is known *not* to lie in diff, which is always a single + // point. Only a != b (or !(a == b)) produces one of these. + bool invert = false; + }; + std::vector known_bounds; + + // A bit per pair key, over every record in the table. A query whose bit is + // clear cannot match anything, which is the answer almost every query gets. + // Wide enough that a few dozen facts leave it sparse: at 64 bits a typical + // table saturates and lets four queries in ten through to the scan. + static constexpr int difference_key_words = 4; + uint64_t difference_keys[difference_key_words] = {0}; + + // Summarize an Expr by its node type, plus the name or value of the leaves + // that distinguish otherwise identical-looking nodes. Deliberately ignores + // children: this only has to be equal for equal Exprs, not unique. + static uint32_t expr_fingerprint(const BaseExprNode *e) { + uint32_t h = ((uint32_t)e->node_type + 1) * 2654435761u; + if (e->node_type == IRNodeType::Variable) { + for (char c : ((const Variable *)e)->name) { + h = h * 31u + (uint32_t)(unsigned char)c; + } + } else if (e->node_type == IRNodeType::IntImm) { + h ^= (uint32_t)((const IntImm *)e)->value; + } + return h; + } + + /** What a scan of known_bounds was able to establish about (a - b). A hole + * that doesn't touch an end of the interval can't be represented in the + * bounds, so it is tracked separately when it matters, which is when the + * hole is at zero. */ + struct KnownDiff { + ConstantInterval bounds; + bool excludes_zero = false; + }; + + /** Everything the facts tell us about (a - b), without building any IR. + * The arguments are borrowed, so this is safe to call with the raw nodes a + * rewrite rule has bound to its wildcards. */ + KnownDiff known_difference(const BaseExprNode *a, const BaseExprNode *b); + + // Helpers over known_difference, for use as rewrite rule predicates. The + // diffs return false when nothing is known, so that a rule asking for a + // bound it can't get simply doesn't fire. + bool is_known_equal(const BaseExprNode *a, const BaseExprNode *b); + bool is_known_not_equal(const BaseExprNode *a, const BaseExprNode *b); + bool known_min_diff(const BaseExprNode *a, const BaseExprNode *b, int64_t *result); + bool known_max_diff(const BaseExprNode *a, const BaseExprNode *b, int64_t *result); + + // 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. Bound it. + // + // The work grows sharply with this limit -- on an adversarial nest of + // min(x, y) - min(z, w) it is roughly 0.02s at 1 or 2, 0.11s at 3 and 0.72s + // at 4 -- while no rule needs the depth: instrumenting every correctness + // test shows the deepest nesting any of them reaches is one. So this is + // already a level of headroom over anything observed. + int can_prove_depth = 0; + static constexpr int max_can_prove_depth = 2; + + // Is there anything a known_true predicate could look up? 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(); + } + + // Is there anything a min_diff or max_diff predicate could look up? Only a + // comparison of non-overflowing integers leaves a record here, so this is + // strictly narrower than has_facts: a boolean fact, or a fact about a type + // that can wrap, satisfies that one while leaving this table empty. Rules + // that ask about differences must gate on this, or they spend a lookup on + // a table that cannot answer. + bool has_difference_facts() const { + return !known_bounds.empty(); + } + + // Symmetric key for a pair. Equal summaries say only that the two nodes are + // the same kind, and xoring them throws even that away, so key those by the + // kind instead of letting every same-type pair share one bit. + HALIDE_ALWAYS_INLINE + static uint32_t difference_key(uint32_t fa, uint32_t fb) { + return fa == fb ? fa * 0x9e3779b9u : (fa ^ fb); + } + + // One bit per pair key. + HALIDE_ALWAYS_INLINE + bool difference_key_present(uint32_t key) const { + const uint32_t bit = key % (difference_key_words * 64); + return (difference_keys[bit / 64] >> (bit % 64)) & 1; + } + + HALIDE_ALWAYS_INLINE + void add_difference_key(uint32_t key) { + const uint32_t bit = key % (difference_key_words * 64); + difference_keys[bit / 64] |= (uint64_t)1 << (bit % 64); + } + + // 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; std::vector pop_list; std::vector bounds_pop_list; std::set truths, falsehoods; + // Everything in the simplifier's known_bounds from this index on was + // pushed by this scope, and is truncated away again when it ends. + size_t known_bounds_size = 0; + // Bits can't be cleared one at a time, so keep the summary from before + // this scope and put it back wholesale. + uint64_t saved_difference_keys[difference_key_words] = {0}; void learn_false(const Expr &fact); void learn_true(const Expr &fact); void learn_upper_bound(const Variable *v, int64_t val); void learn_lower_bound(const Variable *v, int64_t val); + // Record what a comparison says about the difference between its sides. + void learn_difference(const Expr &a, const Expr &b, const ConstantInterval &diff, bool invert); // Replace exprs known to be truths or falsehoods with const_true or const_false. Expr substitute_facts(const Expr &e); Stmt substitute_facts(const Stmt &s); ScopedFact(Simplify *s) - : simplify(s) { + : simplify(s), known_bounds_size(s->known_bounds.size()) { + for (int i = 0; i < difference_key_words; i++) { + saved_difference_keys[i] = s->difference_keys[i]; + } } ~ScopedFact(); // allow move but not copy ScopedFact(const ScopedFact &that) = delete; - ScopedFact(ScopedFact &&that) = default; + // Not defaulted: the moved-from object must not undo anything in its + // destructor. The containers below would be empty after a move and so + // would be harmless, but known_bounds_size would survive and truncate + // away the facts this scope had just learned. + ScopedFact(ScopedFact &&that) noexcept + : simplify(that.simplify), + pop_list(std::move(that.pop_list)), + bounds_pop_list(std::move(that.bounds_pop_list)), + truths(std::move(that.truths)), + falsehoods(std::move(that.falsehoods)), + known_bounds_size(that.known_bounds_size) { + for (int i = 0; i < difference_key_words; i++) { + saved_difference_keys[i] = that.saved_difference_keys[i]; + } + that.simplify = nullptr; + } }; // Tell the simplifier to learn from and exploit a boolean diff --git a/src/Simplify_Max.cpp b/src/Simplify_Max.cpp index 88d3ce2cbf5e..718ddfd30ac7 100644 --- a/src/Simplify_Max.cpp +++ b/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_difference_facts() && + (rewrite(max(x, y), a, min_diff(x, y, this) >= 0) || + rewrite(max(x, y), b, max_diff(x, y, this) <= 0))) || 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) || diff --git a/src/Simplify_Min.cpp b/src/Simplify_Min.cpp index 5203a0c14166..0ee489c4bb28 100644 --- a/src/Simplify_Min.cpp +++ b/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_difference_facts() && + (rewrite(min(x, y), a, max_diff(x, y, this) <= 0) || + rewrite(min(x, y), b, min_diff(x, y, this) >= 0))) || 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) || diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 981a00e6f0ee..6016b3065549 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2377,6 +2377,18 @@ void check_invariant() { } } +void check_with_assumptions(const Expr &a, const Expr &b, const std::vector &assumptions) { + Expr simpler = simplify(a, Scope(), Scope(), 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,114 @@ 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}); + + // A min is at most either of its operands and a max is at least either of + // them, which needs no facts at all. That only bounds the difference on one + // side, but knowing the two are unequal removes the endpoint, and the two + // together settle a comparison that neither settles alone. + check_with_assumptions(max(min(x, y) + 1, x), x, {min(x, y) != x}); + check_with_assumptions(min(max(x, y) - 1, x), x, {max(x, y) != x}); + + // Neither ingredient is enough by itself: without the inequality the + // difference could still be zero, and without the shape there is no bound + // for the inequality to tighten. + check_with_assumptions(max(min(x, y) + 1, x), max(min(x, y) + 1, x), {z < z + 1}); + check_with_assumptions(max(y + 1, x), max(y + 1, x), {y != x}); + + // 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(), Scope(), {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); + + // Constant offsets are peeled off both the facts and the queries, so a fact + // stated about a shifted operand still settles a predicate about the + // unshifted one, in either direction. + check_with_assumptions(max(x, y), y, {x + 1 <= y}); + check_with_assumptions(max(x, y), x, {y <= x + 0}); + check_with_assumptions(max(x + 3, y), y, {x + 4 <= y}); + check_with_assumptions(min(x, y), x, {x + 1 <= y}); + + // But an offset that leaves the order undetermined still doesn't fire. + check_with_assumptions(max(x, y), max(x, y), {x <= y + 1}); + + // 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 +2537,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");