Extend Solver to handle more min/max cases

src/Solve.cpp

@@ -701,6 +701,8 @@ class SolveExpression : public IRMutator {
                     }
                 } else if (is_const(mul_a->b, -1)) {
                     expr = mutate(Opp::make(mul_a->a, make_zero(b.type()) - b));
+                } else if (is_const(mul_a->b, 0)) {
+                    expr = mutate(Cmp::make(make_zero(b.type()), b));
                 } else if (is_negative_const(mul_a->b) && no_overflow_int(a.type())) {
                     // Restrict to no_overflow_int types: for narrow signed types
                     // (int8, int16), negate(INT_MIN) overflows back to INT_MIN,
@@ -1055,17 +1057,40 @@ class SolveForInterval : public IRVisitor {
         static string b_name = unique_name('b');
         static string c_name = unique_name('c');
 
+        // When decomposing expressions, we may end up with a condition
+        // that doesn't depend on the variable. In these cases, we still
+        // need to return a sensible interval, e.g. in `x <= 16 && -1 <= 16`,
+        // the second condition should return everything, rather than fail,
+        // and the rule for && will intersect the LHS with everything,
+        // leaving the LHS as the final result.
+        if (Expr cond = le; !expr_uses_var(cond, var)) {
+            if (can_prove(cond)) {
+                result = Interval::everything();
+            } else if (can_prove(!cond)) {
+                result = Interval::nothing();
+            } else {
+                fail();
+            }
+            return;
+        }
+
         const Variable *v = le->a.as<Variable>();
         if (!already_solved) {
             SolverResult solved = solve_expression(le, var, scope);
             if (!solved.fully_solved) {
-                // solve_expression failed; try direct max/min decomposition on the LHS.
+                // solve_expression failed; try direct max/min decomposition.
                 if (const Max *max_fallback = le->a.as<Max>()) {
                     // max(a, b) <= c <==> a <= c && b <= c
                     (max_fallback->a <= le->b && max_fallback->b <= le->b).accept(this);
                 } else if (const Min *min_fallback = le->a.as<Min>()) {
                     // min(a, b) <= c <==> a <= c || b <= c
                     (min_fallback->a <= le->b || min_fallback->b <= le->b).accept(this);
+                } else if (const Min *min_b = le->b.as<Min>()) {
+                    // c <= min(a, b) <==> c <= a && c <= b
+                    (le->a <= min_b->a && le->a <= min_b->b).accept(this);
+                } else if (const Max *max_b = le->b.as<Max>()) {
+                    // c <= max(a, b) <==> c <= a || c <= b
+                    (le->a <= max_b->a || le->a <= max_b->b).accept(this);
                 } else if (const Mul *mul_fallback = le->a.as<Mul>()) {
                     // max/min(a, b) * pos_c <= rhs <==> a*pos_c <= rhs [&&/||] b*pos_c <= rhs
                     const Max *mxf = mul_fallback->a.as<Max>();
@@ -1140,17 +1165,35 @@ class SolveForInterval : public IRVisitor {
         static string b_name = unique_name('b');
         static string c_name = unique_name('c');
 
+        // See the analogous check in visit(const LE *).
+        if (Expr cond = ge; !expr_uses_var(ge, var)) {
+            if (can_prove(cond)) {
+                result = Interval::everything();
+            } else if (can_prove(!cond)) {
+                result = Interval::nothing();
+            } else {
+                fail();
+            }
+            return;
+        }
+
         const Variable *v = ge->a.as<Variable>();
         if (!already_solved) {
             SolverResult solved = solve_expression(ge, var, scope);
             if (!solved.fully_solved) {
-                // solve_expression failed; try direct max/min decomposition on the LHS.
+                // solve_expression failed; try direct max/min decomposition.
                 if (const Max *max_fallback = ge->a.as<Max>()) {
                     // max(a, b) >= c <==> a >= c || b >= c
                     (max_fallback->a >= ge->b || max_fallback->b >= ge->b).accept(this);
                 } else if (const Min *min_fallback = ge->a.as<Min>()) {
                     // min(a, b) >= c <==> a >= c && b >= c
                     (min_fallback->a >= ge->b && min_fallback->b >= ge->b).accept(this);
+                } else if (const Min *min_b = ge->b.as<Min>()) {
+                    // c >= min(a, b) <==> c >= a || c >= b
+                    (ge->a >= min_b->a || ge->a >= min_b->b).accept(this);
+                } else if (const Max *max_b = ge->b.as<Max>()) {
+                    // c >= max(a, b) <==> c >= a && c >= b
+                    (ge->a >= max_b->a && ge->a >= max_b->b).accept(this);
                 } else if (const Mul *mul_fallback = ge->a.as<Mul>()) {
                     // max/min(a, b) * pos_c >= rhs <==> a*pos_c >= rhs [||/&&] b*pos_c >= rhs
                     const Max *mxf = mul_fallback->a.as<Max>();
@@ -1284,7 +1327,13 @@ SolverResult solve_expression(const Expr &e, const std::string &variable, const
 
 Interval solve_for_inner_interval(const Expr &c, const std::string &var) {
     SolveForInterval s(var, false);
-    c.accept(&s);
+    // SolveForInterval's structural rewrites match on the shape of a
+    // comparison's operands, so they're defeated if an operand is hidden
+    // behind a let. Inline them first. graph_substitute keeps shared
+    // subexpressions shared rather than duplicating them, and the solver
+    // caches by Expr, so this stays cheap despite dropping the explicit
+    // sharing that lets provide.
+    substitute_in_all_lets(c).accept(&s);
     internal_assert(s.result.min.defined() && s.result.max.defined())
         << "solve_for_inner_interval returned undefined Exprs: " << c << "\n";
     s.result.min = simplify(common_subexpression_elimination(s.result.min));
@@ -1298,7 +1347,8 @@ Interval solve_for_inner_interval(const Expr &c, const std::string &var) {
 
 Interval solve_for_outer_interval(const Expr &c, const std::string &var) {
     SolveForInterval s(var, true);
-    c.accept(&s);
+    // See the remark in solve_for_inner_interval.
+    substitute_in_all_lets(c).accept(&s);
     internal_assert(s.result.min.defined() && s.result.max.defined())
         << "solve_for_outer_interval returned undefined Exprs: " << c << "\n";
     s.result.min = simplify(common_subexpression_elimination(s.result.min));

test/correctness/CMakeLists.txt

@@ -272,6 +272,7 @@ tests(
     parameter_constraints.cpp
     partial_application.cpp
     partial_realization.cpp
+    partition_clamped_extent_guard.cpp
     partition_loops.cpp
     partition_loops_bug.cpp
     partition_max_filter.cpp

test/correctness/partition_clamped_extent_guard.cpp

@@ -0,0 +1,58 @@
+#include "Halide.h"
+#include <cstdio>
+
+using namespace Halide;
+using namespace Halide::Internal;
+
+// Regression test for solving `c <= min(a, b)` and its symmetric shapes in
+// solve_for_{inner,outer}_interval. This vectorized LUT gather produces a
+// clamped per-lane guard; count loop-dependent guards that should be removed
+// from the steady state.
+
+namespace {
+int count_guards(Module m) {
+    Scope<> loops;
+    int guards = 0;
+    for (const auto &f : m.functions()) {
+        visit_with(
+            f.body,
+            [&](auto *self, const For *op) {
+                ScopedBinding<> bind(loops, op->name);
+                self->visit_base(op);
+            },
+            [&](auto *self, const IfThenElse *op) {
+                if (expr_uses_vars(op->condition, loops)) {
+                    guards++;
+                }
+                self->visit_base(op);
+            });
+    }
+    return guards;
+}
+}  // namespace
+
+int main(int argc, char **argv) {
+    ImageParam input(UInt(8), 2, "input");
+    Buffer<uint8_t> table(256, "table");
+
+    Var x("x"), y("y"), xi("xi");
+
+    Func gather("gather"), out("out");
+    gather(x, y) = table(input(x, y));
+    out(x, y) = gather(x, y);
+
+    out.split(x, x, xi, 16, TailStrategy::GuardWithIf).vectorize(xi);
+    gather.compute_at(out, x).unroll(x);
+
+    Module m = out.compile_to_module({input}, "out", get_jit_target_from_environment());
+
+    if (int guards = count_guards(m); guards != 0) {
+        printf("Per-lane bounds guard was not partitioned out of the steady-state "
+               "loop: found %d loop-dependent IfThenElse guard(s), expected 0\n",
+               guards);
+        return 1;
+    }
+
+    printf("Success!\n");
+    return 0;
+}

test/correctness/solve.cpp

@@ -606,6 +606,34 @@ void test_outer_interval_max_min() {
     check_outer_interval(expr2, 2, Interval::pos_inf());
 }
 
+void test_solve_far_side_min_max() {
+    // Handle min/max on the far side of a comparison, as emitted by
+    // extent-clamped vector guards. Some decomposed terms no longer contain x.
+
+    // c <= min(a, b)  <=>  c <= a && c <= b
+    check_inner_interval(x * 8 + 7 <= min((x + 1) * 8, 100), Interval::neg_inf(), 11);
+    check_outer_interval(x * 8 + 7 <= min((x + 1) * 8, 100), Interval::neg_inf(), 11);
+
+    // c <= max(a, b)  <=>  c <= a || c <= b
+    check_inner_interval(x * 8 + 7 <= max((x - 1) * 8, 100), Interval::neg_inf(), 11);
+    check_outer_interval(x * 8 + 7 <= max((x - 1) * 8, 100), Interval::neg_inf(), 11);
+
+    // c >= min(a, b)  <=>  c >= a || c >= b
+    check_inner_interval(x * 8 >= min((x + 1) * 8, 100), 13, Interval::pos_inf());
+    check_outer_interval(x * 8 >= min((x + 1) * 8, 100), 13, Interval::pos_inf());
+
+    // c >= max(a, b)  <=>  c >= a && c >= b
+    check_inner_interval(x * 8 >= max((x - 1) * 8, 100), 13, Interval::pos_inf());
+    check_outer_interval(x * 8 >= max((x - 1) * 8, 100), 13, Interval::pos_inf());
+
+    // The same guard through a let binding.
+    Expr bound = Variable::make(Int(32), "b");
+    check_inner_interval(Let::make("b", min((x + 1) * 8, 100), x * 8 + 7 <= bound),
+                         Interval::neg_inf(), 11);
+    check_outer_interval(Let::make("b", min((x + 1) * 8, 100), x * 8 + 7 <= bound),
+                         Interval::neg_inf(), 11);
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {
@@ -639,6 +667,7 @@ int main(int argc, char **argv) {
     test_float_mul_eq_zero_divisor_not_rewritten();
     test_float_select_condition_not_simplified();
     test_outer_interval_max_min();
+    test_solve_far_side_min_max();
     std::printf("Success!\n");
     return 0;
 }