Improve argmax patterns for rfactor (#8863)

The `argmax` and `argmin` patterns in AssociativeOpsTable assumed that
comparisons would be normalized to `y < x` (or `x < y`), but the Solve
module could flip these. It was also possible an earlier pass would
negate the patterns such that `y >= x` (or `x >= y`) with the branches
reversed could be encountered instead.

To solve this issue, we extend Solve to move the variable of interest
to the `true` branch of a `select`. We also extend the table to cover
all four comparison cases.

We also substantially refactor Solve.cpp to use a helper to track and
update the ambient tree traversal state (`uses_var` and `failed`).

Author: alexreinking

src/AssociativeOpsTable.cpp

@@ -177,12 +177,20 @@ void populate_ops_table_double_general_mul(const vector<Type> &types, vector<Ass
 
 void populate_ops_table_double_general_max(const vector<Type> &types, vector<AssociativePattern> &table) {
     declare_vars_double(types);
+    // Argmax
+    table.push_back({{max(x0, y0), select(x0 > y0, x1, y1)}, {tmin_0, zero_1}, true});
+    table.push_back({{max(x0, y0), select(x0 >= y0, x1, y1)}, {tmin_0, zero_1}, true});
     table.push_back({{max(x0, y0), select(y0 < x0, x1, y1)}, {tmin_0, zero_1}, true});
+    table.push_back({{max(x0, y0), select(y0 <= x0, x1, y1)}, {tmin_0, zero_1}, true});
 }
 
 void populate_ops_table_double_general_min(const vector<Type> &types, vector<AssociativePattern> &table) {
     declare_vars_double(types);
+    // Argmin
     table.push_back({{min(x0, y0), select(x0 < y0, x1, y1)}, {tmax_0, zero_1}, true});
+    table.push_back({{min(x0, y0), select(x0 <= y0, x1, y1)}, {tmax_0, zero_1}, true});
+    table.push_back({{min(x0, y0), select(y0 > x0, x1, y1)}, {tmax_0, zero_1}, true});
+    table.push_back({{min(x0, y0), select(y0 >= x0, x1, y1)}, {tmax_0, zero_1}, true});
 }
 
 void populate_ops_table_double_general_sub(const vector<Type> &types, vector<AssociativePattern> &table) {
@@ -348,21 +356,21 @@ const vector<AssociativePattern> &get_ops_table(const vector<Expr> &exprs) {
         types[i] = exprs[i].type();
     }
 
-    {
+    const vector<AssociativePattern> &table = [&]() -> decltype(auto) {
         // get_ops_table_helper() lazily initializes the table, so ensure
         // that multiple threads can't try to do so at the same time.
         static std::mutex ops_table_lock;
         std::scoped_lock lock_guard(ops_table_lock);
 
-        const vector<AssociativePattern> &table = get_ops_table_helper(types, exprs[0].node_type(), exprs.size());
-        debug(7) << "Table size: " << table.size() << "\n";
-        for (const auto &p : table) {
-            debug(7) << p;
-        }
-        return table;
+        return get_ops_table_helper(types, exprs[0].node_type(), exprs.size());
+    }();
+
+    debug(5) << "Table size: " << table.size() << "\n";
+    for (const auto &p : table) {
+        debug(5) << p;
     }
-    debug(5) << "Returning empty table\n";
-    return empty;
+
+    return table;
 }
 
 }  // namespace Internal

src/CSE.cpp

@@ -526,7 +526,7 @@ void cse_test() {
         check(e, correct);
     }
 
-    debug(0) << "common_subexpression_elimination test passed\n";
+    std::cout << "common_subexpression_elimination test passed\n";
 }
 
 }  // namespace Internal

src/IREquality.cpp

@@ -608,7 +608,7 @@ void ir_equality_test() {
     e2 = e2 * e2 + e2;
     check_not_equal(e1, e2);
 
-    debug(0) << "ir_equality_test passed\n";
+    std::cout << "ir_equality_test passed\n";
 }
 
 }  // namespace Internal