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

[mcourteaux]

Problem statement

While preparing #9371, I hit several dead ends trying to produce very neat IR. The reason is the simplifier cannot simplify max(x, y) = x when we give the assumption x >= y. Intuitively, one would write in Simplify_Max.cpp:

   rewrite(max(x, y), a, can_prove(x >= y, this))

However, the way can_prove(Expr, Prover) is implemented is to recursively mutate() the Expr with the Prover (i.e., this instance of Simplify). This however, does not substitute in the facts (truths), and therefore fails to "prove" that x > y.

A secondary problem with rewrite rules that use can_prove() is that they recursively invoke the simplifier, which leads potentially to infinite recursions. Specifically, Andrew stated:

Calling can_prove as the third arg to rewrite recursively invokes the simplifier on every expr that matches that LHS. If one of them itself creates an intermediate expression that also matches the LHS, you get infinite recursion.

Solution

This PR makes it possible to use facts in a can_prove(), prevents infinite recursion, and offers known_true() as a lightweight alternative.

Using facts / truths.

In order to match simple truths into expressions, whenever new facts are learned, they are canonicalized, such that lookup can do the same, and a > b matches with b < a.
The can_prove(Expr, Prover) entry-point now calls out to simplify_can_prove_condition(), which is the full power of the Simplifier at work, but using canonicalized facts.

As a bonus: substitute_facts benefits from this canonicalization and manages to substitute in more facts, even when the fact is not IR-tree-wise an exact match.

As an additional bonus: weaker forms of inequalities are also substituted in as facts if the stronger inequality is known as a fact. Example: x >= y (which is weak) is replaced by true if the fact that x > y is known.

Limiting recursion.

Implementing this, I indeed hit an infinite recursion quickly, so this PR also limits the recursion depth to 1. Experiments performed with higher recursion depth explode compile time for no measured benefit. The recursion limit of 1 means that a call to simplify() can call simplify() only once within a can_prove() but not more. This is done in the call simplify_can_prove_condition() which is exactly the new behavior of can_prove().

A lightweight alternative: known_true().

Instead of running the full simplifier for a can_prove(), we can now also write rewrite-predicates using a known_true() which is just the simple lookup in the Simplifier's fact list using the canonicalization. This naturally cannot recurse, and also does not spend time trying to simplify things when not needed. This is especially useful for when you would match max(x, y) on every such Expr and then try to can_prove(x > y), which is expensive. Instead the rewrite rule is:

   rewrite(max(x, y), a, known_true(x >= y, this))

Which is a very fast lookup, instead of a whole run through the simplifier.

A few such example rules are implemented in Simplify_Div.

Checklist

• Tests added or updated (not required for docs, CI config, or typo fixes)
• Documentation updated (if public API changed)
• Python bindings updated (if public API changed)
• Benchmarks are included here if the change is intended to affect performance.
• Commits include AI attribution where applicable (see Code of Conduct)

[codecov]

Codecov Report

❌ Patch coverage is 82.65306% with 51 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.26%. Comparing base (a3690b3) to head (260cc16).

Files with missing lines	Patch %	Lines
src/Simplify.cpp	76.81%	30 Missing and 21 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #9400      +/-   ##
==========================================
+ Coverage   70.12%   70.26%   +0.14%     
==========================================
  Files         261      261              
  Lines       79405    79692     +287     
  Branches    19362    19450      +88     
==========================================
+ Hits        55684    55999     +315     
- Misses      17896    17897       +1     
+ Partials     5825     5796      -29     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[abadams]

Was this the one that inflated the lowering time of lens_blur? Is this superseded by your approach in aligned splits take 2?

[abadams]

Some more data: yesterday in a research branch I came across a case where max(x, y) was not simplifying inside an if (x <= y) branch, and it was causing wmma ops to fail to be extracted. This is a case we need to handle, we just need to figure out how to do it without increasing compile times.

[abadams]

I think an approach to make this fast might be to define an operator< that can compare IRMatcher patterns to Exprs, so that the pattern can be looked up in the set of known facts without constructing an IR node, rather than needing to build an Expr just to do the lookup.

[mcourteaux]

Was this the one that inflated the lowering time of lens_blur? Is this superseded by your approach in aligned splits take 2?

This indeed did slow down lens blur by 10% more ore less. It's not superseded: take 2 just works around the simplification issue by using .bound_extent() and .bound_storage().

Some more data: yesterday in a research branch I came across a case where max(x, y) was not simplifying inside an if (x <= y) branch, and it was causing wmma ops to fail to be extracted. This is a case we need to handle, we just need to figure out how to do it without increasing compile times.

🤝

I think an approach to make this fast might be to define an operator< that can compare IRMatcher patterns to Exprs, so that the pattern can be looked up in the set of known facts without constructing an IR node, rather than needing to build an Expr just to do the lookup.

That sounds like a decent approach! Feel free to take over this branch!

[mcourteaux]

What if instead, we special handle the generic-form max(x, y) and min(x, y) above the chain of rewrite rules, by doing a fact_lookup_comparison(x, y) that returns an enum: NOTHING_KNOWN, LT, LE, EQ, GE, GT.

As such, no known_true() is called for every Max node, and no new Expr is constructed.

We can keep the more expensive machinery for non-trivial rules, such as the ones now in Simplify_Div, which wouldn't trigger for every Max/Min node, because the LHS of the rewrite rule is more specific:

 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)))

The existing known_true() would then call out to fact_lookup_comparison().

Later, when we learn_true(some_expr % some_other_expr == 0) for example, we can have special handling in Simplify_Mod with a fact_lookup_modrem(x, y), which can return more info specific to those kind of facts.

[mcourteaux]

Offline discussion with @abadams, with some ideas going back and forth, Andrew proposed to:

record a list of (Expr, Expr, ConstantInterval), where the interval is constant bounds on the difference between the two Exprs. Then in the predicate you could query the lower or upper bound of the difference.

As such, we can have a function that return ConstantBounds for the difference between two Expr. Then you can write predicates like this:

rewrite(max(x, y), x, min_diff(x, y, this) >= 0) // because x - y >= 0, we know that x >= y.