diff --git a/src/Solve.cpp b/src/Solve.cpp index a9f4e69f7e3e..6df3310d8353 100644 --- a/src/Solve.cpp +++ b/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(); 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(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(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()) { + // 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()) { + // 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()) { // max/min(a, b) * pos_c <= rhs <==> a*pos_c <= rhs [&&/||] b*pos_c <= rhs const Max *mxf = mul_fallback->a.as(); @@ -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(); 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(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(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()) { + // 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()) { + // 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()) { // max/min(a, b) * pos_c >= rhs <==> a*pos_c >= rhs [||/&&] b*pos_c >= rhs const Max *mxf = mul_fallback->a.as(); @@ -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)); diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d6e1518e0079..f013c51de608 100644 --- a/test/correctness/CMakeLists.txt +++ b/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 diff --git a/test/correctness/partition_clamped_extent_guard.cpp b/test/correctness/partition_clamped_extent_guard.cpp new file mode 100644 index 000000000000..3240b94ebff1 --- /dev/null +++ b/test/correctness/partition_clamped_extent_guard.cpp @@ -0,0 +1,58 @@ +#include "Halide.h" +#include + +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 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; +} diff --git a/test/correctness/solve.cpp b/test/correctness/solve.cpp index 22120be8b685..cdfd434528f5 100644 --- a/test/correctness/solve.cpp +++ b/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; }