From 2f4f8bcd55a8db3884251fa1af095cf707943e7d Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Tue, 5 Nov 2024 12:47:22 -0500
Subject: [PATCH 01/39] Add failing rfactor tests

---
 test/error/CMakeLists.txt                     |  2 ++
 .../rfactor_after_var_and_rvar_fusion.cpp     | 25 ++++++++++++++++++
 test/error/rfactor_fused_var_and_rvar.cpp     | 26 +++++++++++++++++++
 3 files changed, 53 insertions(+)
 create mode 100644 test/error/rfactor_after_var_and_rvar_fusion.cpp
 create mode 100644 test/error/rfactor_fused_var_and_rvar.cpp

diff --git a/test/error/CMakeLists.txt b/test/error/CMakeLists.txt
index 0478c3b11087..5272b2717de7 100644
--- a/test/error/CMakeLists.txt
+++ b/test/error/CMakeLists.txt
@@ -94,6 +94,8 @@ tests(GROUPS error
       require_fail.cpp
       reuse_var_in_schedule.cpp
       reused_args.cpp
+      rfactor_after_var_and_rvar_fusion.cpp
+      rfactor_fused_var_and_rvar.cpp
       rfactor_inner_dim_non_commutative.cpp
       round_up_and_blend_race.cpp
       run_with_large_stack_throws.cpp
diff --git a/test/error/rfactor_after_var_and_rvar_fusion.cpp b/test/error/rfactor_after_var_and_rvar_fusion.cpp
new file mode 100644
index 000000000000..8b94fbdd1b16
--- /dev/null
+++ b/test/error/rfactor_after_var_and_rvar_fusion.cpp
@@ -0,0 +1,25 @@
+#include "Halide.h"
+
+using namespace Halide;
+
+int main(int argc, char **argv) {
+    Func f{"f"};
+    RDom r({{0, 5}, {0, 5}, {0, 5}}, "r");
+    Var x{"x"}, y{"y"};
+    f(x, y) = 0;
+    f(x, y) += r.x + r.y + r.z;
+
+    RVar rxy{"rxy"}, yrz{"yrz"};
+    Var z{"z"};
+
+    // Error: In schedule for f.update(0), can't perform rfactor() after fusing y and r$z
+    f.update()
+        .fuse(r.x, r.y, rxy)
+        .fuse(r.z, y, yrz)
+        .rfactor(rxy, z);
+
+    f.print_loop_nest();
+
+    printf("Success!\n");
+    return 0;
+}
\ No newline at end of file
diff --git a/test/error/rfactor_fused_var_and_rvar.cpp b/test/error/rfactor_fused_var_and_rvar.cpp
new file mode 100644
index 000000000000..a167ca543c47
--- /dev/null
+++ b/test/error/rfactor_fused_var_and_rvar.cpp
@@ -0,0 +1,26 @@
+#include "Halide.h"
+
+using namespace Halide;
+
+int main(int argc, char **argv) {
+    Func f{"f"};
+    RDom r({{0, 5}, {0, 5}, {0, 5}}, "r");
+    Var x{"x"}, y{"y"};
+    f(x, y) = 0;
+    f(x, y) += r.x + r.y + r.z;
+
+    RVar rxy{"rxy"}, yrz{"yrz"}, yr{"yr"};
+    Var z{"z"};
+
+    // Error: In schedule for f.update(0), can't perform rfactor() after fusing r$z and y
+    f.update()
+        .fuse(r.x, r.y, rxy)
+        .fuse(y, r.z, yrz)
+        .fuse(rxy, yrz, yr)
+        .rfactor(yr, z);
+
+    f.print_loop_nest();
+
+    printf("Success!\n");
+    return 0;
+}
\ No newline at end of file

From 47cdbccefe49c898317269506d43045775946701 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Thu, 21 Nov 2024 16:20:34 -0500
Subject: [PATCH 02/39] Rewrite rfactor()

---
 src/Func.cpp     | 796 +++++++++++++++++++++--------------------------
 src/Func.h       |   2 +-
 src/Substitute.h |  10 +
 3 files changed, 369 insertions(+), 439 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index c243e6950f3f..9b03c64aa5ae 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -1,6 +1,7 @@
 #include <algorithm>
 #include <cstring>
 #include <iostream>
+#include <unordered_set>
 #include <utility>
 
 #ifdef _MSC_VER
@@ -609,522 +610,441 @@ Expr substitute_self_reference(Expr val, const string &func, const Function &sub
     return val;
 }
 
-// Substitute the occurrence of 'name' in 'exprs' with 'value'.
-void substitute_var_in_exprs(const string &name, const Expr &value, vector<Expr> &exprs) {
-    for (auto &expr : exprs) {
-        expr = substitute(name, value, expr);
-    }
+}  // anonymous namespace
+
+Func Stage::rfactor(const RVar &r, const Var &v) {
+    definition.schedule().touched() = true;
+    return rfactor({{r, v}});
 }
 
-void apply_split_result(const vector<pair<string, Expr>> &bounds_let_stmts,
-                        const vector<ApplySplitResult> &splits_result,
-                        vector<Expr> &predicates, vector<Expr> &args,
-                        vector<Expr> &values) {
+// Helpers for rfactor implementation
+namespace {
 
-    for (const auto &res : splits_result) {
-        switch (res.type) {
-        case ApplySplitResult::Substitution:
-        case ApplySplitResult::LetStmt:
-            // Apply substitutions to the list of predicates, args, and values.
-            // Make sure we substitute in all the let stmts as well since we are
-            // not going to add them to the exprs.
-            substitute_var_in_exprs(res.name, res.value, predicates);
-            substitute_var_in_exprs(res.name, res.value, args);
-            substitute_var_in_exprs(res.name, res.value, values);
-            break;
-        default:
-            internal_assert(res.type == ApplySplitResult::Predicate);
-            predicates.push_back(res.value);
-            break;
-        }
+struct DimHash {
+    std::size_t operator()(const Dim &s) const noexcept {
+        return std::hash<std::string>{}(s.var);
     }
-
-    // Make sure we substitute in all the let stmts from 'bounds_let_stmts'
-    // since we are not going to add them to the exprs.
-    for (const auto &let : bounds_let_stmts) {
-        substitute_var_in_exprs(let.first, let.second, predicates);
-        substitute_var_in_exprs(let.first, let.second, args);
-        substitute_var_in_exprs(let.first, let.second, values);
+};
+struct DimEq {
+    bool operator()(const Dim &lhs, const Dim &rhs) const noexcept {
+        return lhs.var == rhs.var;
     }
-}
+};
+using DimSet = std::unordered_set<Dim, DimHash, DimEq>;
 
-/** Apply split directives on the reduction variables. Remove the old RVar from
- * the list and add the split result (inner and outer RVars) to the list. Add
- * new predicates corresponding to the TailStrategy to the RDom predicate list. */
-bool apply_split(const Split &s, vector<ReductionVariable> &rvars,
-                 vector<Expr> &predicates, vector<Expr> &args,
-                 vector<Expr> &values, map<string, Expr> &dim_extent_alignment) {
-    internal_assert(s.split_type == Split::SplitVar);
-    const auto it = std::find_if(rvars.begin(), rvars.end(),
-                                 [&s](const ReductionVariable &rv) { return (s.old_var == rv.var); });
+template<class T, class = void>
+struct has_name : std::false_type {};
 
-    Expr old_max, old_min, old_extent;
+template<class T>
+struct has_name<T, std::void_t<decltype(&T::name)>> : std::true_type {};
 
-    if (it != rvars.end()) {
-        debug(4) << "  Splitting " << it->var << " into " << s.outer << " and " << s.inner << "\n";
+template<typename T>
+std::optional<T> match_by_dim_name(const std::string &dim, const std::vector<T> &items) {
+    const auto has_v = std::find_if(items.begin(), items.end(), [&dim](auto &x) {
+        if constexpr (has_name<T>::value) {
+            return var_name_match(dim, x.name());
+        } else {
+            return var_name_match(dim, x.var);
+        }
+    });
+    return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
+}
 
-        old_max = simplify(it->min + it->extent - 1);
-        old_min = it->min;
-        old_extent = it->extent;
+template<typename T>
+std::optional<T> find_by_var_name(const std::vector<T> &items, const std::string &name) {
+    const auto has_v = std::find_if(items.begin(), items.end(), [&name](auto &x) {
+        if constexpr (has_name<T>::value) {
+            return var_name_match(x.name(), name);
+        } else {
+            return var_name_match(x.var, name);
+        }
+    });
+    return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
+}
 
-        it->var = s.inner;
-        it->min = 0;
-        it->extent = s.factor;
+template<typename T>
+std::optional<T> find_by_var_name(const std::vector<T> &items, const VarOrRVar &v) {
+    return find_by_var_name(items, v.name());
+}
 
-        rvars.insert(it + 1, {s.outer, 0, simplify((old_extent - 1 + s.factor) / s.factor)});
+std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const std::vector<Dim> &dims) {
+    if (!prover_result.associative()) {
+        return "can't perform rfactor() because we can't prove associativity of the operator";
+    }
 
-        vector<ApplySplitResult> splits_result = apply_split(s, "", dim_extent_alignment);
-        vector<pair<string, Expr>> bounds_let_stmts = compute_loop_bounds_after_split(s, "");
-        apply_split_result(bounds_let_stmts, splits_result, predicates, args, values);
+    DimSet is_rfactored;
+    for (const auto &[rv, v] : preserved) {
+        // Check that the RVar are in the dims list
+        const auto &rv_dim = find_by_var_name(dims, rv);
+        if (!(rv_dim && rv_dim->is_rvar())) {
+            std::stringstream s;
+            s << "can't perform rfactor() on " << rv.name() << " since it is not in the reduction domain";
+            return s.str();
+        }
+        is_rfactored.insert(*rv_dim);
 
-        return true;
+        // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
+        const auto &v_dim = find_by_var_name(dims, v);
+        if (v_dim) {
+            std::stringstream s;
+            s << "can't rename the rvars " << rv.name() << " into " << v.name()
+              << ", since it is already used in this Func's schedule elsewhere.";
+            return s.str();
+        }
     }
-    return false;
-}
 
-/** Apply fuse directives on the reduction variables. Remove the
- * fused RVars from the list and add the fused RVar to the list. */
-bool apply_fuse(const Split &s, vector<ReductionVariable> &rvars,
-                vector<Expr> &predicates, vector<Expr> &args,
-                vector<Expr> &values, map<string, Expr> &dim_extent_alignment) {
-    internal_assert(s.split_type == Split::FuseVars);
-    const auto &iter_outer = std::find_if(rvars.begin(), rvars.end(),
-                                          [&s](const ReductionVariable &rv) { return (s.outer == rv.var); });
-    const auto &iter_inner = std::find_if(rvars.begin(), rvars.end(),
-                                          [&s](const ReductionVariable &rv) { return (s.inner == rv.var); });
+    // If the operator is associative but non-commutative, rfactor() on inner
+    // dimensions (excluding the outer dimensions) is not valid.
+    if (!prover_result.commutative()) {
+        std::optional<Dim> last_rvar;
+        for (const auto &d : reverse_view(dims)) {
+            if (is_rfactored.count(d) && last_rvar && !is_rfactored.count(*last_rvar)) {
+                std::stringstream s;
+                s << "can't rfactor an inner dimension " << d.var
+                  << " without rfactoring the outer dimensions, since the "
+                  << "operator is non-commutative.";
+                return s.str();
+            }
+            if (d.is_rvar()) {
+                last_rvar = d;
+            }
+        }
+    }
 
-    Expr inner_min, inner_extent, outer_min, outer_extent;
-    if ((iter_outer != rvars.end()) && (iter_inner != rvars.end())) {
-        debug(4) << "  Fusing " << s.outer << " and " << s.inner << " into " << s.old_var << "\n";
+    return std::nullopt;
+}
 
-        inner_min = iter_inner->min;
-        inner_extent = iter_inner->extent;
-        outer_min = iter_outer->min;
-        outer_extent = iter_outer->extent;
+template<typename T>
+std::vector<T> &operator+=(std::vector<T> &base, const std::vector<T> &other) {
+    base.insert(base.end(), other.begin(), other.end());
+    return base;
+}
 
-        Expr extent = iter_outer->extent * iter_inner->extent;
-        iter_outer->var = s.old_var;
-        iter_outer->min = 0;
-        iter_outer->extent = extent;
-        rvars.erase(iter_inner);
+template<typename T, typename U>
+std::vector<T> &operator+=(std::vector<T> &base, const std::vector<U> &other) {
+    std::copy(other.begin(), other.end(), std::back_inserter(base));
+    return base;
+}
 
-        vector<ApplySplitResult> splits_result = apply_split(s, "", dim_extent_alignment);
-        vector<pair<string, Expr>> bounds_let_stmts = compute_loop_bounds_after_split(s, "");
-        apply_split_result(bounds_let_stmts, splits_result, predicates, args, values);
+template<typename T>
+std::vector<T> operator+(const std::vector<T> &base, const std::vector<T> &other) {
+    std::vector<T> merged = base;
+    merged += other;
+    return merged;
+}
 
-        return true;
-    }
-    return false;
+template<typename T, typename U>
+std::vector<T> operator+(const std::vector<T> &base, const std::vector<U> &other) {
+    std::vector<T> merged = base;
+    merged += other;
+    return merged;
 }
 
-/** Apply purify directives on the reduction variables and predicates. Purify
- * replace a RVar with a Var, thus, the RVar needs to be removed from the list.
- * Any reference to the RVar in the predicates will be replaced with reference
- * to a Var. */
-bool apply_purify(const Split &s, vector<ReductionVariable> &rvars,
-                  vector<Expr> &predicates, vector<Expr> &args,
-                  vector<Expr> &values, map<string, Expr> &dim_extent_alignment) {
-    internal_assert(s.split_type == Split::PurifyRVar);
-    const auto &iter = std::find_if(rvars.begin(), rvars.end(),
-                                    [&s](const ReductionVariable &rv) { return (s.old_var == rv.var); });
-    if (iter != rvars.end()) {
-        debug(4) << "  Purify RVar " << iter->var << " into Var " << s.outer
-                 << ", deleting it from the rvars list\n";
-        rvars.erase(iter);
+template<typename U, typename V>
+std::vector<U> copy_convert(const std::vector<V> &vec) {
+    return std::vector<U>{} + vec;
+}
 
-        vector<ApplySplitResult> splits_result = apply_split(s, "", dim_extent_alignment);
-        vector<pair<string, Expr>> bounds_let_stmts = compute_loop_bounds_after_split(s, "");
-        apply_split_result(bounds_let_stmts, splits_result, predicates, args, values);
+using SubstitutionMap = std::map<std::string, Expr>;
 
-        return true;
+// This is a helper function for building up a substitution map that
+// corresponds to pushing down a nest of lets. The lets should be fed
+// to this function from innermost to outermost. This is equivalent to
+// building a let-nest as a one-hole context and then simplifying.
+void add_let(SubstitutionMap &proj, const std::string &name, const Expr &value) {
+    for (auto &[_, e] : proj) {
+        e = substitute(name, value, e);
     }
-    return false;
+    proj.emplace(name, value);
 }
 
-/** Apply rename directives on the reduction variables. */
-bool apply_rename(const Split &s, vector<ReductionVariable> &rvars,
-                  vector<Expr> &predicates, vector<Expr> &args,
-                  vector<Expr> &values, map<string, Expr> &dim_extent_alignment) {
-    internal_assert(s.split_type == Split::RenameVar);
-    const auto &iter = std::find_if(rvars.begin(), rvars.end(),
-                                    [&s](const ReductionVariable &rv) { return (s.old_var == rv.var); });
-    if (iter != rvars.end()) {
-        debug(4) << "  Renaming " << iter->var << " into " << s.outer << "\n";
-        iter->var = s.outer;
-
-        vector<ApplySplitResult> splits_result = apply_split(s, "", dim_extent_alignment);
-        vector<pair<string, Expr>> bounds_let_stmts = compute_loop_bounds_after_split(s, "");
-        apply_split_result(bounds_let_stmts, splits_result, predicates, args, values);
-
-        return true;
+void rebind(SubstitutionMap &proj, const std::string &name, const Expr &value) {
+    for (auto &[_, e] : proj) {
+        e = substitute(name, value, e);
+    }
+    if (!proj.count(name)) {
+        proj.emplace(name, value);
     }
-    return false;
 }
 
-/** Apply scheduling directives (e.g. split, fuse, etc.) on the reduction
- * variables. */
-bool apply_split_directive(const Split &s, vector<ReductionVariable> &rvars,
-                           vector<Expr> &predicates, vector<Expr> &args,
-                           vector<Expr> &values) {
-    map<string, Expr> dim_extent_alignment;
-    for (const ReductionVariable &rv : rvars) {
-        dim_extent_alignment[rv.var] = rv.extent;
-    }
+std::pair<ReductionDomain, SubstitutionMap> project_rdom(const std::vector<Dim> &dims, const Definition &def) {
+    const auto &rdom = def.schedule().rvars();
+    const auto &splits = def.schedule().splits();
+    const auto &predicate = def.predicate();
 
-    vector<pair<string, Expr>> rvar_bounds;
-    for (const ReductionVariable &rv : rvars) {
-        rvar_bounds.emplace_back(rv.var + ".loop_min", rv.min);
-        rvar_bounds.emplace_back(rv.var + ".loop_max", simplify(rv.min + rv.extent - 1));
-        rvar_bounds.emplace_back(rv.var + ".loop_extent", rv.extent);
+    // Compute new bounds for the RDom by walking backwards through the splits.
+    SubstitutionMap bounds_projection{};
+    for (const Split &split : reverse_view(splits)) {
+        for (const auto &[name, value] : compute_loop_bounds_after_split(split, "")) {
+            add_let(bounds_projection, name, value);
+        }
+    }
+    for (const ReductionVariable &rv : rdom) {
+        add_let(bounds_projection, rv.var + ".loop_min", rv.min);
+        add_let(bounds_projection, rv.var + ".loop_max", rv.min + rv.extent - 1);
+        add_let(bounds_projection, rv.var + ".loop_extent", rv.extent);
     }
 
-    bool found = false;
-    switch (s.split_type) {
-    case Split::SplitVar:
-        found = apply_split(s, rvars, predicates, args, values, dim_extent_alignment);
-        break;
-    case Split::FuseVars:
-        found = apply_fuse(s, rvars, predicates, args, values, dim_extent_alignment);
-        break;
-    case Split::PurifyRVar:
-        found = apply_purify(s, rvars, predicates, args, values, dim_extent_alignment);
-        break;
-    case Split::RenameVar:
-        found = apply_rename(s, rvars, predicates, args, values, dim_extent_alignment);
-        break;
+    // Build the new RDom
+    std::vector<ReductionVariable> new_rvars;
+    for (const Dim &dim : dims) {
+        Expr new_min = substitute(bounds_projection, Variable::make(Int(32), dim.var + ".loop_min"));
+        Expr new_extent = substitute(bounds_projection, Variable::make(Int(32), dim.var + ".loop_extent"));
+        new_rvars.push_back(ReductionVariable{dim.var, simplify(new_min), simplify(new_extent)});
     }
+    ReductionDomain new_rdom{new_rvars};
+    new_rdom.where(predicate);
 
-    if (found) {
-        for (const auto &let : rvar_bounds) {
-            substitute_var_in_exprs(let.first, let.second, predicates);
-            substitute_var_in_exprs(let.first, let.second, args);
-            substitute_var_in_exprs(let.first, let.second, values);
+    // Compute the mapping of old dimensions to the projected RDom values.
+    SubstitutionMap dim_projection{};
+    SubstitutionMap dim_extent_alignment{};
+    for (const ReductionVariable &rv : rdom) {
+        add_let(dim_projection, rv.var, Variable::make(Int(32), rv.var));
+        dim_extent_alignment[rv.var] = rv.extent;
+    }
+    for (const Split &split : splits) {
+        for (const auto &result : apply_split(split, "", dim_extent_alignment)) {
+            switch (result.type) {
+            case ApplySplitResult::Substitution:
+            case ApplySplitResult::SubstitutionInCalls:
+            case ApplySplitResult::SubstitutionInProvides:
+            case ApplySplitResult::LetStmt:
+                add_let(dim_projection, result.name, result.value);
+                break;
+            case ApplySplitResult::PredicateCalls:
+            case ApplySplitResult::PredicateProvides:
+            case ApplySplitResult::Predicate:
+                new_rdom.where(substitute(bounds_projection, result.value));
+                break;
+            case ApplySplitResult::BlendProvides:
+                // TODO: what to do here? BlendProvides is not included in the above checks.
+                break;
+            }
         }
     }
-    return found;
-}
-
-}  // anonymous namespace
+    for (auto &[_, e] : dim_projection) {
+        e = substitute(bounds_projection, e);
+    }
+    for (const ReductionVariable &rv : new_rdom.domain()) {
+        rebind(dim_projection, rv.var, Variable::make(Int(32), rv.var, new_rdom));
+    }
 
-Func Stage::rfactor(const RVar &r, const Var &v) {
-    definition.schedule().touched() = true;
-    return rfactor({{r, v}});
+    return std::make_pair(new_rdom, dim_projection);
 }
 
-Func Stage::rfactor(vector<pair<RVar, Var>> preserved) {
+}  // namespace
+
+Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     user_assert(!definition.is_init()) << "rfactor() must be called on an update definition\n";
 
     definition.schedule().touched() = true;
 
-    const string &func_name = function.name();
-    vector<Expr> &args = definition.args();
-    vector<Expr> &values = definition.values();
-
-    // Figure out which pure vars were used in this update definition.
-    std::set<string> pure_vars_used;
-    internal_assert(args.size() == dim_vars.size());
-    for (size_t i = 0; i < args.size(); i++) {
-        if (const Internal::Variable *var = args[i].as<Variable>()) {
-            if (var->name == dim_vars[i].name()) {
-                pure_vars_used.insert(var->name);
-            }
-        }
-    }
-
     // Check whether the operator is associative and determine the operator and
     // its identity for each value in the definition if it is a Tuple
-    const auto &prover_result = prove_associativity(func_name, args, values);
+    const auto &prover_result = prove_associativity(function.name(), definition.args(), definition.values());
 
-    user_assert(prover_result.associative())
-        << "Failed to call rfactor() on " << name()
-        << " since it can't prove associativity of the operator\n";
-    internal_assert(prover_result.size() == values.size());
+    const auto &error = rfactor_validate_args(preserved, prover_result, definition.schedule().dims());
+    user_assert(!error) << "In schedule for " << name() << ": " << *error << "\n"
+                        << dump_argument_list();
 
-    vector<Split> &splits = definition.schedule().splits();
-    vector<Dim> &dims = definition.schedule().dims();
-    vector<ReductionVariable> &rvars = definition.schedule().rvars();
-    vector<Expr> predicates = definition.split_predicate();
-
-    Scope<string> scope;  // Contains list of RVars lifted to the intermediate Func
-    vector<string> rvars_removed;
+    // sort preserved by the dimension ordering
+    vector<RVar> preserved_rvars;
+    vector<Var> preserved_vars;
+    vector<Dim> preserved_rdims;
+    vector<Dim> intermediate_rdims;
+    {
+        std::unordered_map<std::string, int> dim_ordering;
+        for (int i = 0; i < definition.schedule().dims().size(); i++) {
+            dim_ordering.emplace(definition.schedule().dims()[i].var, i);
+        }
 
-    vector<bool> is_rfactored(dims.size(), false);
-    for (const pair<RVar, Var> &i : preserved) {
-        const RVar &rv = i.first;
-        const Var &v = i.second;
-        {
-            // Check that the RVar are in the dims list
-            const auto &iter = std::find_if(dims.begin(), dims.end(),
-                                            [&rv](const Dim &dim) { return var_name_match(dim.var, rv.name()); });
-            user_assert((iter != dims.end()) && (*iter).is_rvar())
-                << "In schedule for " << name()
-                << ", can't perform rfactor() on " << rv.name()
-                << " since it is not in the reduction domain\n"
-                << dump_argument_list();
-            is_rfactored[iter - dims.begin()] = true;
+        std::vector<std::tuple<RVar, Var, Dim>> preserved_with_dims;
+        for (const auto &[rv, v] : preserved) {
+            const std::optional<Dim> rdim = find_by_var_name(definition.schedule().dims(), rv);
+            internal_assert(rdim);
+            preserved_with_dims.emplace_back(rv, v, *rdim);
         }
-        {
-            // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
-            const auto &iter = std::find_if(dims.begin(), dims.end(),
-                                            [&v](const Dim &dim) { return var_name_match(dim.var, v.name()); });
-            user_assert(iter == dims.end())
-                << "In schedule for " << name()
-                << ", can't rename the rvars " << rv.name() << " into " << v.name()
-                << ", since it is already used in this Func's schedule elsewhere.\n"
-                << dump_argument_list();
+
+        std::sort(preserved_with_dims.begin(), preserved_with_dims.end(), [&](const auto &lhs, const auto &rhs) {
+            return dim_ordering.at(std::get<2>(lhs).var) < dim_ordering.at(std::get<2>(rhs).var);
+        });
+
+        for (const auto &[rv, v, dim] : preserved_with_dims) {
+            preserved_rvars.push_back(rv);
+            preserved_vars.push_back(v);
+            preserved_rdims.push_back(dim);
         }
-    }
 
-    // If the operator is associative but non-commutative, rfactor() on inner
-    // dimensions (excluding the outer dimensions) is not valid.
-    if (!prover_result.commutative()) {
-        int last_rvar = -1;
-        for (int i = dims.size() - 1; i >= 0; --i) {
-            if ((last_rvar != -1) && is_rfactored[i]) {
-                user_assert(is_rfactored[last_rvar])
-                    << "In schedule for " << name()
-                    << ", can't rfactor an inner dimension " << dims[i].var
-                    << " without rfactoring the outer dimensions, since the "
-                    << "operator is non-commutative.\n"
-                    << dump_argument_list();
-            }
-            if (dims[i].is_rvar()) {
-                last_rvar = i;
+        for (const Dim &dim : definition.schedule().dims()) {
+            if (dim.is_rvar() && !match_by_dim_name(dim.var, preserved_rvars)) {
+                intermediate_rdims.push_back(dim);
             }
         }
     }
 
-    // We need to apply the split directives on the reduction vars, so that we can
-    // correctly lift the RVars not in 'rvars_kept' and distribute the RVars to the
-    // intermediate and merge Funcs.
-    {
-        vector<Split> temp;
-        for (const Split &s : splits) {
-            // If it's already applied, we should remove it from the split list.
-            if (!apply_split_directive(s, rvars, predicates, args, values)) {
-                temp.push_back(s);
-            }
-        }
-        splits = temp;
-    }
-
-    // Reduction domain of the intermediate update definition
-    vector<ReductionVariable> intm_rvars;
-    for (const auto &rv : rvars) {
-        const auto &iter = std::find_if(preserved.begin(), preserved.end(),
-                                        [&rv](const pair<RVar, Var> &pair) { return var_name_match(rv.var, pair.first.name()); });
-        if (iter == preserved.end()) {
-            intm_rvars.push_back(rv);
-            scope.push(rv.var, rv.var);
-        }
-    }
-    RDom intm_rdom(intm_rvars);
-
-    // Sort the Rvars kept and their Vars replacement based on the RVars of
-    // the reduction domain AFTER applying the split directives, so that we
-    // can have a consistent args order for the update definition of the
-    // intermediate and new merge Funcs.
-    std::sort(preserved.begin(), preserved.end(),
-              [&](const pair<RVar, Var> &lhs, const pair<RVar, Var> &rhs) {
-                  const auto &iter_lhs = std::find_if(rvars.begin(), rvars.end(),
-                                                      [&lhs](const ReductionVariable &rv) { return var_name_match(rv.var, lhs.first.name()); });
-                  const auto &iter_rhs = std::find_if(rvars.begin(), rvars.end(),
-                                                      [&rhs](const ReductionVariable &rv) { return var_name_match(rv.var, rhs.first.name()); });
-                  return iter_lhs < iter_rhs;
-              });
-    // The list of RVars to keep in the new update definition
-    vector<RVar> rvars_kept(preserved.size());
-    // List of pure Vars to replace the RVars in the intermediate's update definition
-    vector<Var> vars_rename(preserved.size());
-    for (size_t i = 0; i < preserved.size(); ++i) {
-        const auto &val = preserved[i];
-        rvars_kept[i] = val.first;
-        vars_rename[i] = val.second;
-    }
-
-    // List of RVars for the new reduction domain. Any RVars not in 'rvars_kept'
-    // are removed from the RDom
+    // Intermediate func
+    Func intm(function.name() + "_intm");
+
+    // Intermediate pure definition
     {
-        vector<ReductionVariable> temp;
-        for (const auto &rv : rvars) {
-            const auto &iter = std::find_if(rvars_kept.begin(), rvars_kept.end(),
-                                            [&rv](const RVar &rvar) { return var_name_match(rv.var, rvar.name()); });
-            if (iter != rvars_kept.end()) {
-                temp.push_back(rv);
-            } else {
-                rvars_removed.push_back(rv.var);
-            }
-        }
-        rvars.swap(temp);
+        intm(dim_vars + preserved_vars) = Tuple(prover_result.pattern.identities);
     }
-    RDom f_rdom(rvars);
-
-    // Init definition of the intermediate Func
 
-    // Compute args of the init definition of the intermediate Func.
-    // Replace the RVars, which are in 'rvars_kept', with the specified new pure
-    // Vars. Also, add the pure Vars of the original init definition as part of
-    // the args.
-    // For example, if we have the following Func f:
-    //   f(x, y) = 10
-    //   f(r.x, r.y) += h(r.x, r.y)
-    // Calling f.update(0).rfactor({{r.y, u}}) will generate the following
-    // intermediate Func:
-    //   f_intm(x, y, u) = 0
-    //   f_intm(r.x, u, u) += h(r.x, u)
+    // Intermediate update definition
+    {
+        auto [intermediate_rdom, intermediate_map] = project_rdom(intermediate_rdims, definition);
+        for (int i = 0; i < preserved.size(); i++) {
+            rebind(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
+        }
+        for (const auto &var : dim_vars) {
+            intermediate_map.erase(var.name());
+        }
 
-    vector<Var> init_args;
-    init_args.insert(init_args.end(), dim_vars.begin(), dim_vars.end());
-    init_args.insert(init_args.end(), vars_rename.begin(), vars_rename.end());
+        intermediate_rdom.set_predicate(simplify(substitute(intermediate_map, intermediate_rdom.predicate())));
 
-    vector<Expr> init_vals(values.size());
-    for (size_t i = 0; i < init_vals.size(); ++i) {
-        init_vals[i] = prover_result.pattern.identities[i];
-    }
+        vector<Expr> args = definition.args() + preserved_vars;
+        vector<Expr> values;
+        for (const auto &val : definition.values()) {
+            values.push_back(substitute_self_reference(val, function.name(), intm.function(), preserved_vars));
+        }
+        args = substitute(intermediate_map, args);
+        values = substitute(intermediate_map, values);
+        intm.function().define_update(args, values, intermediate_rdom);
 
-    Func intm(func_name + "_intm");
-    intm(init_args) = Tuple(init_vals);
+        // Intermediate schedule
+        intm.function().update(0).schedule() = definition.schedule().get_copy();
 
-    // Args of the update definition of the intermediate Func
-    vector<Expr> update_args(args.size() + vars_rename.size());
+        auto &intm_dims = intm.function().update(0).schedule().dims();
 
-    // We need to substitute the reference to the old RDom's RVars with
-    // the new RDom's RVars. Also, substitute the reference to RVars which
-    // are in 'rvars_kept' with their corresponding new pure Vars
-    map<string, Expr> substitution_map;
-    for (size_t i = 0; i < intm_rvars.size(); ++i) {
-        substitution_map[intm_rvars[i].var] = intm_rdom[i];
-    }
-    for (size_t i = 0; i < vars_rename.size(); i++) {
-        update_args[i + args.size()] = vars_rename[i];
-        RVar rvar_kept = rvars_kept[i];
-        // Find the full name of rvar_kept in rvars
-        const auto &iter = std::find_if(rvars.begin(), rvars.end(),
-                                        [&rvar_kept](const ReductionVariable &rv) { return var_name_match(rv.var, rvar_kept.name()); });
-        substitution_map[iter->var] = vars_rename[i];
-    }
-    for (size_t i = 0; i < args.size(); i++) {
-        Expr arg = substitute(substitution_map, args[i]);
-        update_args[i] = arg;
-    }
+        // Replace rvar dims IN the preserved list with their Vars in the INTERMEDIATE Func
+        for (auto &dim : intm_dims) {
+            const auto it = std::find_if(preserved_rvars.begin(), preserved_rvars.end(), [&](const auto &rv) {
+                return var_name_match(dim.var, rv.name());
+            });
+            if (it != preserved_rvars.end()) {
+                const auto offset = it - preserved_rvars.begin();
+                const auto &var = preserved_vars[offset];
+                const auto &pure_dim = find_by_var_name(intm.function().definition().schedule().dims(), var);
+                internal_assert(pure_dim);
+                dim = *pure_dim;
+            }
+        }
 
-    // Compute the predicates for the intermediate Func and the new update definition
-    for (const Expr &pred : predicates) {
-        Expr subs_pred = substitute(substitution_map, pred);
-        intm_rdom.where(subs_pred);
-        if (!expr_uses_vars(pred, scope)) {
-            // Only keep the predicate that does not depend on the lifted RVars
-            // (either explicitly or implicitly). For example, if 'rx' is split
-            // into 'rxo' and 'rxi' and 'rxo' is part of the lifted RVars, we'll
-            // ignore every predicate that depends on 'rx'
-            f_rdom.where(pred);
+        // Add factored pure dims to the INTERMEDIATE func just before outermost
+        DimSet dims;
+        dims.insert(intm_dims.begin(), intm_dims.end());
+        for (const Var &dim_v : preserved_vars) {
+            const std::optional<Dim> &dim = find_by_var_name(intm.function().definition().schedule().dims(), dim_v);
+            internal_assert(dim) << "Failed to find " << dim_v.name() << " in list of pure dims";
+            if (!dims.count(*dim)) {
+                intm_dims.insert(intm_dims.end() - 1, *dim);
+            }
         }
-    }
-    definition.predicate() = f_rdom.domain().predicate();
 
-    // The update values the intermediate Func should compute
-    vector<Expr> update_vals(values.size());
-    for (size_t i = 0; i < update_vals.size(); i++) {
-        Expr val = substitute(substitution_map, values[i]);
-        // Need to update the self-reference in the update definition to point
-        // to the new intermediate Func
-        val = substitute_self_reference(val, func_name, intm.function(), vars_rename);
-        update_vals[i] = val;
+        intm.function().update(0).schedule().rvars() = intermediate_rdom.domain();
     }
-    // There may not actually be a reference to the RDom in the args or values,
-    // so we use Function::define_update, which lets pass pass an explicit RDom.
-    intm.function().define_update(update_args, update_vals, intm_rdom.domain());
 
-    // Determine the dims and schedule of the update definition of the
-    // intermediate Func. We copy over the schedule from the original
-    // update definition (e.g. split, parallelize, vectorize, etc.)
-    intm.function().update(0).schedule().dims() = dims;
-    intm.function().update(0).schedule().splits() = splits;
+    // Preserved update definition
+    {
+        auto [preserved_rdom, _] = project_rdom(preserved_rdims, definition);
 
-    // Copy over the storage order of the original pure dims
-    vector<StorageDim> &intm_storage_dims = intm.function().schedule().storage_dims();
-    internal_assert(intm_storage_dims.size() ==
-                    function.schedule().storage_dims().size() + vars_rename.size());
-    for (size_t i = 0; i < function.schedule().storage_dims().size(); ++i) {
-        intm_storage_dims[i] = function.schedule().storage_dims()[i];
-    }
+        // Replace the current definition with calls to the intermediate func.
+        vector<Expr> dim_exprs = copy_convert<Expr>(dim_vars);
+        vector<Expr> f_load_args = dim_exprs;
+        for (const ReductionVariable &rv : preserved_rdom.domain()) {
+            f_load_args.push_back(Variable::make(Int(32), rv.var, preserved_rdom));
+        }
 
-    for (size_t i = 0; i < rvars_kept.size(); ++i) {
-        // Apply the purify directive that replaces the RVar in rvars_kept
-        // with a pure Var
-        intm.update(0).purify(rvars_kept[i], vars_rename[i]);
-    }
+        SubstitutionMap replacements;
+        for (size_t i = 0; i < preserved.size(); i++) {
+            replacements.emplace(preserved_rdims[i].var, preserved_vars[i]);
+        }
+        for (size_t i = 0; i < definition.values().size(); ++i) {
+            if (!prover_result.ys[i].var.empty()) {
+                Expr r = (definition.values().size() == 1) ? Expr(intm(f_load_args)) : Expr(intm(f_load_args)[i]);
+                replacements.emplace(prover_result.ys[i].var, r);
+            }
 
-    // Determine the dims of the new update definition
-
-    // The new update definition needs all the pure vars of the Func, but the
-    // one we're rfactoring may not have used them all. Add any missing ones to
-    // the dims list.
-
-    // Add pure Vars from the original init definition to the dims list
-    // if they are not already in the list
-    for (const Var &v : dim_vars) {
-        if (!pure_vars_used.count(v.name())) {
-            Dim d = {v.name(), ForType::Serial, DeviceAPI::None, DimType::PureVar, Partition::Auto};
-            // Insert it just before Var::outermost
-            dims.insert(dims.end() - 1, d);
-        }
-    }
+            if (!prover_result.xs[i].var.empty()) {
+                Expr prev_val = Call::make(intm.types()[i], function.name(),
+                                           dim_exprs, Call::CallType::Halide,
+                                           FunctionPtr(), i);
+                replacements.emplace(prover_result.xs[i].var, prev_val);
+            } else {
+                user_warning << "Update definition of " << name() << " at index " << i
+                             << " doesn't depend on the previous value. This isn't a"
+                             << " reduction operation\n";
+            }
+        }
 
-    // Then, we need to remove lifted RVars from the dims list
-    for (const string &rv : rvars_removed) {
-        remove(rv);
-    }
+        std::vector<Dim> reducing_dims;
+        {
+            DimSet preserved_rdim_set;
+            preserved_rdim_set.insert(preserved_rdims.begin(), preserved_rdims.end());
 
-    // Define the new update definition which refers to the intermediate Func.
-    // Using the same example as above, the new update definition is:
-    //   f(x, y) += f_intm(x, y, r.y)
+            // Remove rvar dims NOT IN the preserved list from the REDUCING Func
+            for (const auto &dim : definition.schedule().dims()) {
+                if (!dim.is_rvar() || preserved_rdim_set.count(dim)) {
+                    reducing_dims.push_back(dim);
+                }
+            }
 
-    // Args for store in the new update definition
-    vector<Expr> f_store_args(dim_vars.size());
-    for (size_t i = 0; i < f_store_args.size(); ++i) {
-        f_store_args[i] = dim_vars[i];
-    }
-
-    // Call's args to the intermediate Func in the new update definition
-    vector<Expr> f_load_args;
-    f_load_args.insert(f_load_args.end(), dim_vars.begin(), dim_vars.end());
-    for (int i = 0; i < f_rdom.dimensions(); ++i) {
-        f_load_args.push_back(f_rdom[i]);
-    }
-    internal_assert(f_load_args.size() == init_args.size());
+            // Add missing pure vars to the REDUCING func just before outermost
+            for (int i = 0; i < dim_vars.size(); i++) {
+                if (!expr_uses_var(definition.args()[i], dim_vars[i].name())) {
+                    Dim d = {dim_vars[i].name(), ForType::Serial, DeviceAPI::None, DimType::PureVar, Partition::Auto};
+                    reducing_dims.insert(reducing_dims.end() - 1, d);
+                }
+            }
+        }
 
-    // Update value of the new update definition. It loads values from
-    // the intermediate Func.
-    vector<Expr> f_values(values.size());
+        definition.args() = dim_exprs;
+        definition.predicate() = const_true();  // TODO: replace with strongest postcondition of the intermediate predicate with the eliminated rvars havoc'd
+        definition.schedule().dims() = std::move(reducing_dims);
+        definition.schedule().rvars() = preserved_rdom.domain();
+        definition.values() = substitute(replacements, prover_result.pattern.ops);
+    }
 
-    // There might be cross-dependencies between tuple elements, so we need
-    // to collect all substitutions first.
-    map<string, Expr> replacements;
-    for (size_t i = 0; i < f_values.size(); ++i) {
-        if (!prover_result.ys[i].var.empty()) {
-            Expr r = (values.size() == 1) ? Expr(intm(f_load_args)) : Expr(intm(f_load_args)[i]);
-            replacements.emplace(prover_result.ys[i].var, r);
+    // Clean up the splits lists
+    for (Stage st : {*this, intm.update(0)}) {
+        Scope<> dims;
+        for (const Var &v : st.dim_vars) {
+            dims.push(v.name());
         }
-
-        if (!prover_result.xs[i].var.empty()) {
-            Expr prev_val = Call::make(intm.types()[i], func_name,
-                                       f_store_args, Call::CallType::Halide,
-                                       FunctionPtr(), i);
-            replacements.emplace(prover_result.xs[i].var, prev_val);
-        } else {
-            user_warning << "Update definition of " << name() << " at index " << i
-                         << " doesn't depend on the previous value. This isn't a"
-                         << " reduction operation\n";
+        for (const ReductionVariable &rv : st.definition.schedule().rvars()) {
+            dims.push(rv.var);
         }
+        std::vector<Split> new_splits;
+        for (const Split &split : st.definition.schedule().splits()) {
+            switch (split.split_type) {
+            case Split::SplitVar:
+                if (dims.contains(split.old_var)) {
+                    dims.pop(split.old_var);
+                    dims.push(split.outer);
+                    dims.push(split.inner);
+                    new_splits.push_back(split);
+                }
+                break;
+            case Split::FuseVars:
+                if (dims.contains(split.outer) && dims.contains(split.inner)) {
+                    dims.pop(split.outer);
+                    dims.pop(split.inner);
+                    dims.push(split.old_var);
+                    new_splits.push_back(split);
+                }
+                break;
+            case Split::PurifyRVar:
+            case Split::RenameVar:
+                if (dims.contains(split.old_var)) {
+                    dims.pop(split.old_var);
+                    dims.push(split.outer);
+                    new_splits.push_back(split);
+                }
+                break;
+            }
+        }
+        st.definition.schedule().splits().swap(new_splits);
     }
-    for (size_t i = 0; i < f_values.size(); ++i) {
-        f_values[i] = substitute(replacements, prover_result.pattern.ops[i]);
-    }
-
-    // Update the definition
-    args.swap(f_store_args);
-    values.swap(f_values);
 
     return intm;
 }
@@ -1187,7 +1107,7 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
     bool round_up_ok = !exact;
     if (round_up_ok && !definition.is_init()) {
         // If it's the outermost split in this dimension, RoundUp
-        // is OK. Otherwise we need GuardWithIf to avoid
+        // is OK. Otherwise, we need GuardWithIf to avoid
         // recomputing values in the case where the inner split
         // factor does not divide the outer split factor.
         std::set<string> inner_vars;
@@ -1224,7 +1144,7 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
     bool predicate_loads_ok = !exact;
     if (predicate_loads_ok && tail == TailStrategy::PredicateLoads) {
         // If it's the outermost split in this dimension, PredicateLoads
-        // is OK. Otherwise we can't prove it's safe.
+        // is OK. Otherwise, we can't prove it's safe.
         std::set<string> inner_vars;
         for (const Split &s : definition.schedule().splits()) {
             switch (s.split_type) {
diff --git a/src/Func.h b/src/Func.h
index 32d8f1e58c69..a0f374eb874f 100644
--- a/src/Func.h
+++ b/src/Func.h
@@ -184,7 +184,7 @@ class Stage {
      *
      */
     // @{
-    Func rfactor(std::vector<std::pair<RVar, Var>> preserved);
+    Func rfactor(const std::vector<std::pair<RVar, Var>> &preserved);
     Func rfactor(const RVar &r, const Var &v);
     // @}
 
diff --git a/src/Substitute.h b/src/Substitute.h
index 22bdf640b7a8..55e7ac1cf0fa 100644
--- a/src/Substitute.h
+++ b/src/Substitute.h
@@ -37,6 +37,16 @@ Expr substitute(const Expr &find, const Expr &replacement, const Expr &expr);
 Stmt substitute(const Expr &find, const Expr &replacement, const Stmt &stmt);
 // @}
 
+/** Substitute a container of Exprs or Stmts out of place */
+template<typename T>
+T substitute(const std::map<std::string, Expr> &replacements, const T &container) {
+    T output;
+    std::transform(container.begin(), container.end(), std::back_inserter(output), [&](const auto &expr_or_stmt) {
+        return substitute(replacements, expr_or_stmt);
+    });
+    return output;
+}
+
 /** Substitutions where the IR may be a general graph (and not just a
  * DAG). */
 // @{

From 530bd3f46c2d132cb579463542eed842a8c5adbe Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Sun, 24 Nov 2024 09:52:29 -0500
Subject: [PATCH 03/39] Remove unused operators

---
 src/Func.cpp | 24 ++----------------------
 1 file changed, 2 insertions(+), 22 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 9b03c64aa5ae..6f2a10d36309 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -684,8 +684,7 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
         is_rfactored.insert(*rv_dim);
 
         // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
-        const auto &v_dim = find_by_var_name(dims, v);
-        if (v_dim) {
+        if (find_by_var_name(dims, v)) {
             std::stringstream s;
             s << "can't rename the rvars " << rv.name() << " into " << v.name()
               << ", since it is already used in this Func's schedule elsewhere.";
@@ -714,29 +713,10 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
     return std::nullopt;
 }
 
-template<typename T>
-std::vector<T> &operator+=(std::vector<T> &base, const std::vector<T> &other) {
-    base.insert(base.end(), other.begin(), other.end());
-    return base;
-}
-
-template<typename T, typename U>
-std::vector<T> &operator+=(std::vector<T> &base, const std::vector<U> &other) {
-    std::copy(other.begin(), other.end(), std::back_inserter(base));
-    return base;
-}
-
-template<typename T>
-std::vector<T> operator+(const std::vector<T> &base, const std::vector<T> &other) {
-    std::vector<T> merged = base;
-    merged += other;
-    return merged;
-}
-
 template<typename T, typename U>
 std::vector<T> operator+(const std::vector<T> &base, const std::vector<U> &other) {
     std::vector<T> merged = base;
-    merged += other;
+    merged.insert(merged.end(), other.begin(), other.end());
     return merged;
 }
 

From 6483c12188842263ecec1fe0c3d5a7445905df94 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Sun, 24 Nov 2024 10:22:21 -0500
Subject: [PATCH 04/39] Clean up dim/var matching helpers

---
 src/Func.cpp | 49 ++++++++++++++-----------------------------------
 1 file changed, 14 insertions(+), 35 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 6f2a10d36309..74cd1328d835 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -632,41 +632,21 @@ struct DimEq {
 };
 using DimSet = std::unordered_set<Dim, DimHash, DimEq>;
 
-template<class T, class = void>
-struct has_name : std::false_type {};
-
-template<class T>
-struct has_name<T, std::void_t<decltype(&T::name)>> : std::true_type {};
-
-template<typename T>
-std::optional<T> match_by_dim_name(const std::string &dim, const std::vector<T> &items) {
-    const auto has_v = std::find_if(items.begin(), items.end(), [&dim](auto &x) {
-        if constexpr (has_name<T>::value) {
-            return var_name_match(dim, x.name());
-        } else {
-            return var_name_match(dim, x.var);
-        }
+std::optional<RVar> find_rvar(const std::vector<RVar> &items, const Dim &dim) {
+    const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
+        return var_name_match(dim.var, x.name());
     });
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
-template<typename T>
-std::optional<T> find_by_var_name(const std::vector<T> &items, const std::string &name) {
-    const auto has_v = std::find_if(items.begin(), items.end(), [&name](auto &x) {
-        if constexpr (has_name<T>::value) {
-            return var_name_match(x.name(), name);
-        } else {
-            return var_name_match(x.var, name);
-        }
+std::optional<Dim> find_dim(const std::vector<Dim> &items, const VarOrRVar &v) {
+    const std::string &name = v.name();
+    const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
+        return var_name_match(x.var, name);
     });
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
-template<typename T>
-std::optional<T> find_by_var_name(const std::vector<T> &items, const VarOrRVar &v) {
-    return find_by_var_name(items, v.name());
-}
-
 std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const std::vector<Dim> &dims) {
     if (!prover_result.associative()) {
         return "can't perform rfactor() because we can't prove associativity of the operator";
@@ -675,7 +655,7 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
     DimSet is_rfactored;
     for (const auto &[rv, v] : preserved) {
         // Check that the RVar are in the dims list
-        const auto &rv_dim = find_by_var_name(dims, rv);
+        const auto &rv_dim = find_dim(dims, rv);
         if (!(rv_dim && rv_dim->is_rvar())) {
             std::stringstream s;
             s << "can't perform rfactor() on " << rv.name() << " since it is not in the reduction domain";
@@ -684,7 +664,7 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
         is_rfactored.insert(*rv_dim);
 
         // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
-        if (find_by_var_name(dims, v)) {
+        if (find_dim(dims, v)) {
             std::stringstream s;
             s << "can't rename the rvars " << rv.name() << " into " << v.name()
               << ", since it is already used in this Func's schedule elsewhere.";
@@ -840,7 +820,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
         std::vector<std::tuple<RVar, Var, Dim>> preserved_with_dims;
         for (const auto &[rv, v] : preserved) {
-            const std::optional<Dim> rdim = find_by_var_name(definition.schedule().dims(), rv);
+            const std::optional<Dim> rdim = find_dim(definition.schedule().dims(), rv);
             internal_assert(rdim);
             preserved_with_dims.emplace_back(rv, v, *rdim);
         }
@@ -856,7 +836,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         }
 
         for (const Dim &dim : definition.schedule().dims()) {
-            if (dim.is_rvar() && !match_by_dim_name(dim.var, preserved_rvars)) {
+            if (dim.is_rvar() && !find_rvar(preserved_rvars, dim)) {
                 intermediate_rdims.push_back(dim);
             }
         }
@@ -882,12 +862,11 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
         intermediate_rdom.set_predicate(simplify(substitute(intermediate_map, intermediate_rdom.predicate())));
 
-        vector<Expr> args = definition.args() + preserved_vars;
+        vector<Expr> args = substitute(intermediate_map, definition.args() + preserved_vars);
         vector<Expr> values;
         for (const auto &val : definition.values()) {
             values.push_back(substitute_self_reference(val, function.name(), intm.function(), preserved_vars));
         }
-        args = substitute(intermediate_map, args);
         values = substitute(intermediate_map, values);
         intm.function().define_update(args, values, intermediate_rdom);
 
@@ -904,7 +883,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             if (it != preserved_rvars.end()) {
                 const auto offset = it - preserved_rvars.begin();
                 const auto &var = preserved_vars[offset];
-                const auto &pure_dim = find_by_var_name(intm.function().definition().schedule().dims(), var);
+                const auto &pure_dim = find_dim(intm.function().definition().schedule().dims(), var);
                 internal_assert(pure_dim);
                 dim = *pure_dim;
             }
@@ -914,7 +893,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         DimSet dims;
         dims.insert(intm_dims.begin(), intm_dims.end());
         for (const Var &dim_v : preserved_vars) {
-            const std::optional<Dim> &dim = find_by_var_name(intm.function().definition().schedule().dims(), dim_v);
+            const std::optional<Dim> &dim = find_dim(intm.function().definition().schedule().dims(), dim_v);
             internal_assert(dim) << "Failed to find " << dim_v.name() << " in list of pure dims";
             if (!dims.count(*dim)) {
                 intm_dims.insert(intm_dims.end() - 1, *dim);

From 7ffdb66075196bab8e05fce07962a49ed4358c86 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Sun, 24 Nov 2024 10:25:07 -0500
Subject: [PATCH 05/39] Fix definition in PyStage.cpp

---
 python_bindings/src/halide/halide_/PyStage.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python_bindings/src/halide/halide_/PyStage.cpp b/python_bindings/src/halide/halide_/PyStage.cpp
index b412a6f2b39e..fac47fa3cf1f 100644
--- a/python_bindings/src/halide/halide_/PyStage.cpp
+++ b/python_bindings/src/halide/halide_/PyStage.cpp
@@ -14,7 +14,7 @@ void define_stage(py::module &m) {
             .def("dump_argument_list", &Stage::dump_argument_list)
             .def("name", &Stage::name)
 
-            .def("rfactor", (Func(Stage::*)(std::vector<std::pair<RVar, Var>>)) & Stage::rfactor,
+            .def("rfactor", (Func(Stage::*)(const std::vector<std::pair<RVar, Var>> &)) & Stage::rfactor,
                  py::arg("preserved"))
             .def("rfactor", (Func(Stage::*)(const RVar &, const Var &)) & Stage::rfactor,
                  py::arg("r"), py::arg("v"))

From 3fb172f2ee26b4b7ce7846d471f4864901b5dd3c Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Sun, 24 Nov 2024 10:59:33 -0500
Subject: [PATCH 06/39] Use and_condition_over_domain to predicate the reducing
 definition in rfactor

---
 src/Func.cpp | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 74cd1328d835..31657327bd5e 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -905,7 +905,8 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
     // Preserved update definition
     {
-        auto [preserved_rdom, _] = project_rdom(preserved_rdims, definition);
+        auto [preserved_rdom, preserved_map] = project_rdom(preserved_rdims, definition);
+        preserved_rdom.set_predicate(simplify(substitute(preserved_map, preserved_rdom.predicate())));
 
         // Replace the current definition with calls to the intermediate func.
         vector<Expr> dim_exprs = copy_convert<Expr>(dim_vars);
@@ -957,8 +958,13 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             }
         }
 
+        Scope<Interval> intm_rdom;
+        for (const auto &[var, min, extent] : intm.update(0).definition.schedule().rvars()) {
+            intm_rdom.push(var, Interval{min, min + extent - 1});
+        }
+
         definition.args() = dim_exprs;
-        definition.predicate() = const_true();  // TODO: replace with strongest postcondition of the intermediate predicate with the eliminated rvars havoc'd
+        definition.predicate() = !and_condition_over_domain(simplify(!preserved_rdom.predicate()), intm_rdom);
         definition.schedule().dims() = std::move(reducing_dims);
         definition.schedule().rvars() = preserved_rdom.domain();
         definition.values() = substitute(replacements, prover_result.pattern.ops);

From 27922009d78b01edbe2771e35d03e8e8d43b5315 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Sun, 24 Nov 2024 11:21:39 -0500
Subject: [PATCH 07/39] Disallow rfactor() on funcs with RVar+Var fused
 schedules

---
 src/Func.cpp                 | 45 +++++++++++++++++++++++++++++++++---
 test/common/expect_abort.cpp |  3 +++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 31657327bd5e..958a762b2120 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -647,14 +647,16 @@ std::optional<Dim> find_dim(const std::vector<Dim> &items, const VarOrRVar &v) {
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
-std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const std::vector<Dim> &dims) {
+std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const Definition &definition) {
+    const std::vector<Dim> &dims = definition.schedule().dims();
+
     if (!prover_result.associative()) {
         return "can't perform rfactor() because we can't prove associativity of the operator";
     }
 
     DimSet is_rfactored;
     for (const auto &[rv, v] : preserved) {
-        // Check that the RVar are in the dims list
+        // Check that the RVars are in the dims list
         const auto &rv_dim = find_dim(dims, rv);
         if (!(rv_dim && rv_dim->is_rvar())) {
             std::stringstream s;
@@ -690,6 +692,43 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
         }
     }
 
+    // Check that no Vars were fused into RVars
+    Scope<> rdims;
+    for (const ReductionVariable &rv : definition.schedule().rvars()) {
+        rdims.push(rv.var);
+    }
+    for (const Split &split : definition.schedule().splits()) {
+        switch (split.split_type) {
+        case Split::SplitVar:
+            if (rdims.contains(split.old_var)) {
+                rdims.pop(split.old_var);
+                rdims.push(split.outer);
+                rdims.push(split.inner);
+            }
+            break;
+        case Split::FuseVars:
+            if (rdims.contains(split.outer) || rdims.contains(split.inner)) {
+                if (!(rdims.contains(split.outer) && rdims.contains(split.inner))) {
+                    std::stringstream s;
+                    s << "Cannot rfactor a Func that has fused a Var into an RVar ("
+                      << split.outer << ", " << split.inner << ")";
+                    return s.str();
+                }
+                rdims.pop(split.outer);
+                rdims.pop(split.inner);
+                rdims.push(split.old_var);
+            }
+            break;
+        case Split::PurifyRVar:
+        case Split::RenameVar:
+            if (rdims.contains(split.old_var)) {
+                rdims.pop(split.old_var);
+                rdims.push(split.outer);
+            }
+            break;
+        }
+    }
+
     return std::nullopt;
 }
 
@@ -803,7 +842,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     // its identity for each value in the definition if it is a Tuple
     const auto &prover_result = prove_associativity(function.name(), definition.args(), definition.values());
 
-    const auto &error = rfactor_validate_args(preserved, prover_result, definition.schedule().dims());
+    const auto &error = rfactor_validate_args(preserved, prover_result, definition);
     user_assert(!error) << "In schedule for " << name() << ": " << *error << "\n"
                         << dump_argument_list();
 
diff --git a/test/common/expect_abort.cpp b/test/common/expect_abort.cpp
index cb09a7242921..fec89b0913b7 100644
--- a/test/common/expect_abort.cpp
+++ b/test/common/expect_abort.cpp
@@ -19,6 +19,9 @@ auto handler = ([]() {
                       << std::flush;
             suppress_abort = false;
             std::abort();  // We should never EXPECT an internal error
+        } catch (const Halide::Error &e) {
+            std::cerr << e.what() << "\n"
+                      << std::flush;
         } catch (const std::exception &e) {
             std::cerr << e.what() << "\n"
                       << std::flush;

From 9ccacd230597f4e1aa5da2f57fa9de7131192c14 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Sun, 24 Nov 2024 11:23:58 -0500
Subject: [PATCH 08/39] Use size_t in place of int

---
 src/Func.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 958a762b2120..9a39f82ba30c 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -853,7 +853,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     vector<Dim> intermediate_rdims;
     {
         std::unordered_map<std::string, int> dim_ordering;
-        for (int i = 0; i < definition.schedule().dims().size(); i++) {
+        for (size_t i = 0; i < definition.schedule().dims().size(); i++) {
             dim_ordering.emplace(definition.schedule().dims()[i].var, i);
         }
 
@@ -892,7 +892,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     // Intermediate update definition
     {
         auto [intermediate_rdom, intermediate_map] = project_rdom(intermediate_rdims, definition);
-        for (int i = 0; i < preserved.size(); i++) {
+        for (size_t i = 0; i < preserved.size(); i++) {
             rebind(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
         }
         for (const auto &var : dim_vars) {
@@ -989,7 +989,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             }
 
             // Add missing pure vars to the REDUCING func just before outermost
-            for (int i = 0; i < dim_vars.size(); i++) {
+            for (size_t i = 0; i < dim_vars.size(); i++) {
                 if (!expr_uses_var(definition.args()[i], dim_vars[i].name())) {
                     Dim d = {dim_vars[i].name(), ForType::Serial, DeviceAPI::None, DimType::PureVar, Partition::Auto};
                     reducing_dims.insert(reducing_dims.end() - 1, d);

From e154fc17584e789ec6f0845157b13fb7e244b853 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 08:46:41 -0500
Subject: [PATCH 09/39] Clean up uses of split_predicate()

---
 src/BoundsInference.cpp   | 24 ++++++++----------------
 src/Derivative.cpp        |  2 +-
 src/ScheduleFunctions.cpp |  3 +--
 3 files changed, 10 insertions(+), 19 deletions(-)

diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp
index 724adb993afd..aba76f7798ed 100644
--- a/src/BoundsInference.cpp
+++ b/src/BoundsInference.cpp
@@ -14,6 +14,7 @@
 
 #include <algorithm>
 #include <iterator>
+#include <numeric>
 
 namespace Halide {
 namespace Internal {
@@ -297,7 +298,6 @@ class BoundsInference : public IRMutator {
             }
 
             // Default case (no specialization)
-            vector<Expr> predicates = def.split_predicate();
             for (const ReductionVariable &rv : def.schedule().rvars()) {
                 rvars.insert(rv);
             }
@@ -308,23 +308,15 @@ class BoundsInference : public IRMutator {
             }
             vecs[1] = def.values();
 
+            vector<Expr> predicates = def.split_predicate();
             for (size_t i = 0; i < result.size(); ++i) {
                 for (const Expr &val : vecs[i]) {
-                    if (!predicates.empty()) {
-                        Expr cond_val = Call::make(val.type(),
-                                                   Internal::Call::if_then_else,
-                                                   {likely(predicates[0]), val},
-                                                   Internal::Call::PureIntrinsic);
-                        for (size_t i = 1; i < predicates.size(); ++i) {
-                            cond_val = Call::make(cond_val.type(),
-                                                  Internal::Call::if_then_else,
-                                                  {likely(predicates[i]), cond_val},
-                                                  Internal::Call::PureIntrinsic);
-                        }
-                        result[i].emplace_back(const_true(), cond_val);
-                    } else {
-                        result[i].emplace_back(const_true(), val);
-                    }
+                    Expr cond_val = std::accumulate(
+                        predicates.begin(), predicates.end(), val,
+                        [](const auto &acc, const auto &pred) {
+                            return Call::make(acc.type(), Call::if_then_else, {likely(pred), acc}, Call::PureIntrinsic);
+                        });
+                    result[i].emplace_back(const_true(), cond_val);
                 }
             }
 
diff --git a/src/Derivative.cpp b/src/Derivative.cpp
index 2520d27e290f..06451732d80c 100644
--- a/src/Derivative.cpp
+++ b/src/Derivative.cpp
@@ -1534,7 +1534,7 @@ void ReverseAccumulationVisitor::propagate_halide_function_call(
             // f(r.x) = ... && r is associative
             // => f(x) = ...
             if (var != nullptr && var->reduction_domain.defined() &&
-                var->reduction_domain.split_predicate().empty()) {
+                is_const_one(var->reduction_domain.predicate())) {
                 ReductionDomain rdom = var->reduction_domain;
                 int rvar_id = -1;
                 for (int rid = 0; rid < (int)rdom.domain().size(); rid++) {
diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp
index c7a257dd085e..74f0fa738e42 100644
--- a/src/ScheduleFunctions.cpp
+++ b/src/ScheduleFunctions.cpp
@@ -178,7 +178,6 @@ Stmt build_loop_nest(
     const auto &dims = func.args();
     const auto &func_s = func.schedule();
     const auto &stage_s = def.schedule();
-    const auto &predicates = def.split_predicate();
 
     // We'll build it from inside out, starting from the body,
     // then wrapping it in for loops.
@@ -306,7 +305,7 @@ Stmt build_loop_nest(
     }
 
     // Put all the reduction domain predicates into the containers vector.
-    for (Expr pred : predicates) {
+    for (Expr pred : def.split_predicate()) {
         pred = qualify(prefix, pred);
         // Add a likely qualifier if there isn't already one
         if (Call::as_intrinsic(pred, {Call::likely, Call::likely_if_innermost})) {

From ed3c4059c862d417a6bd65dc34da33f3c5728c70 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 08:54:20 -0500
Subject: [PATCH 10/39] Clean out some excess std:: qualifications

---
 src/Func.cpp | 47 +++++++++++++++++++++++++----------------------
 1 file changed, 25 insertions(+), 22 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 9a39f82ba30c..719a128c8f8f 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -36,8 +36,11 @@ namespace Halide {
 
 using std::map;
 using std::ofstream;
+using std::optional;
 using std::pair;
 using std::string;
+using std::tuple;
+using std::unordered_map;
 using std::vector;
 
 using namespace Internal;
@@ -632,23 +635,23 @@ struct DimEq {
 };
 using DimSet = std::unordered_set<Dim, DimHash, DimEq>;
 
-std::optional<RVar> find_rvar(const std::vector<RVar> &items, const Dim &dim) {
+optional<RVar> find_rvar(const vector<RVar> &items, const Dim &dim) {
     const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
         return var_name_match(dim.var, x.name());
     });
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
-std::optional<Dim> find_dim(const std::vector<Dim> &items, const VarOrRVar &v) {
-    const std::string &name = v.name();
+optional<Dim> find_dim(const vector<Dim> &items, const VarOrRVar &v) {
+    const string &name = v.name();
     const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
         return var_name_match(x.var, name);
     });
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
-std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const Definition &definition) {
-    const std::vector<Dim> &dims = definition.schedule().dims();
+optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const Definition &definition) {
+    const vector<Dim> &dims = definition.schedule().dims();
 
     if (!prover_result.associative()) {
         return "can't perform rfactor() because we can't prove associativity of the operator";
@@ -677,7 +680,7 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
     // If the operator is associative but non-commutative, rfactor() on inner
     // dimensions (excluding the outer dimensions) is not valid.
     if (!prover_result.commutative()) {
-        std::optional<Dim> last_rvar;
+        optional<Dim> last_rvar;
         for (const auto &d : reverse_view(dims)) {
             if (is_rfactored.count(d) && last_rvar && !is_rfactored.count(*last_rvar)) {
                 std::stringstream s;
@@ -733,31 +736,31 @@ std::optional<std::string> rfactor_validate_args(const vector<pair<RVar, Var>> &
 }
 
 template<typename T, typename U>
-std::vector<T> operator+(const std::vector<T> &base, const std::vector<U> &other) {
-    std::vector<T> merged = base;
+vector<T> operator+(const vector<T> &base, const vector<U> &other) {
+    vector<T> merged = base;
     merged.insert(merged.end(), other.begin(), other.end());
     return merged;
 }
 
 template<typename U, typename V>
-std::vector<U> copy_convert(const std::vector<V> &vec) {
-    return std::vector<U>{} + vec;
+vector<U> copy_convert(const vector<V> &vec) {
+    return vector<U>{} + vec;
 }
 
-using SubstitutionMap = std::map<std::string, Expr>;
+using SubstitutionMap = std::map<string, Expr>;
 
 // This is a helper function for building up a substitution map that
 // corresponds to pushing down a nest of lets. The lets should be fed
 // to this function from innermost to outermost. This is equivalent to
 // building a let-nest as a one-hole context and then simplifying.
-void add_let(SubstitutionMap &proj, const std::string &name, const Expr &value) {
+void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
     for (auto &[_, e] : proj) {
         e = substitute(name, value, e);
     }
     proj.emplace(name, value);
 }
 
-void rebind(SubstitutionMap &proj, const std::string &name, const Expr &value) {
+void rebind(SubstitutionMap &proj, const string &name, const Expr &value) {
     for (auto &[_, e] : proj) {
         e = substitute(name, value, e);
     }
@@ -766,7 +769,7 @@ void rebind(SubstitutionMap &proj, const std::string &name, const Expr &value) {
     }
 }
 
-std::pair<ReductionDomain, SubstitutionMap> project_rdom(const std::vector<Dim> &dims, const Definition &def) {
+pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const Definition &def) {
     const auto &rdom = def.schedule().rvars();
     const auto &splits = def.schedule().splits();
     const auto &predicate = def.predicate();
@@ -785,7 +788,7 @@ std::pair<ReductionDomain, SubstitutionMap> project_rdom(const std::vector<Dim>
     }
 
     // Build the new RDom
-    std::vector<ReductionVariable> new_rvars;
+    vector<ReductionVariable> new_rvars;
     for (const Dim &dim : dims) {
         Expr new_min = substitute(bounds_projection, Variable::make(Int(32), dim.var + ".loop_min"));
         Expr new_extent = substitute(bounds_projection, Variable::make(Int(32), dim.var + ".loop_extent"));
@@ -828,7 +831,7 @@ std::pair<ReductionDomain, SubstitutionMap> project_rdom(const std::vector<Dim>
         rebind(dim_projection, rv.var, Variable::make(Int(32), rv.var, new_rdom));
     }
 
-    return std::make_pair(new_rdom, dim_projection);
+    return {new_rdom, dim_projection};
 }
 
 }  // namespace
@@ -852,14 +855,14 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     vector<Dim> preserved_rdims;
     vector<Dim> intermediate_rdims;
     {
-        std::unordered_map<std::string, int> dim_ordering;
+        unordered_map<string, int> dim_ordering;
         for (size_t i = 0; i < definition.schedule().dims().size(); i++) {
             dim_ordering.emplace(definition.schedule().dims()[i].var, i);
         }
 
-        std::vector<std::tuple<RVar, Var, Dim>> preserved_with_dims;
+        vector<tuple<RVar, Var, Dim>> preserved_with_dims;
         for (const auto &[rv, v] : preserved) {
-            const std::optional<Dim> rdim = find_dim(definition.schedule().dims(), rv);
+            const optional<Dim> rdim = find_dim(definition.schedule().dims(), rv);
             internal_assert(rdim);
             preserved_with_dims.emplace_back(rv, v, *rdim);
         }
@@ -932,7 +935,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         DimSet dims;
         dims.insert(intm_dims.begin(), intm_dims.end());
         for (const Var &dim_v : preserved_vars) {
-            const std::optional<Dim> &dim = find_dim(intm.function().definition().schedule().dims(), dim_v);
+            const optional<Dim> &dim = find_dim(intm.function().definition().schedule().dims(), dim_v);
             internal_assert(dim) << "Failed to find " << dim_v.name() << " in list of pure dims";
             if (!dims.count(*dim)) {
                 intm_dims.insert(intm_dims.end() - 1, *dim);
@@ -976,7 +979,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             }
         }
 
-        std::vector<Dim> reducing_dims;
+        vector<Dim> reducing_dims;
         {
             DimSet preserved_rdim_set;
             preserved_rdim_set.insert(preserved_rdims.begin(), preserved_rdims.end());
@@ -1018,7 +1021,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         for (const ReductionVariable &rv : st.definition.schedule().rvars()) {
             dims.push(rv.var);
         }
-        std::vector<Split> new_splits;
+        vector<Split> new_splits;
         for (const Split &split : st.definition.schedule().splits()) {
             switch (split.split_type) {
             case Split::SplitVar:

From 22add4880bba2b6c5f2d89f78968a8e316e50d73 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 10:36:47 -0500
Subject: [PATCH 11/39] Use dim_match instead of var_name_match

---
 src/Func.cpp | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 719a128c8f8f..6ae66636b7b5 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -637,15 +637,14 @@ using DimSet = std::unordered_set<Dim, DimHash, DimEq>;
 
 optional<RVar> find_rvar(const vector<RVar> &items, const Dim &dim) {
     const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
-        return var_name_match(dim.var, x.name());
+        return dim_match(dim, x);
     });
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
 optional<Dim> find_dim(const vector<Dim> &items, const VarOrRVar &v) {
-    const string &name = v.name();
     const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
-        return var_name_match(x.var, name);
+        return dim_match(x, v);
     });
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
@@ -920,7 +919,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         // Replace rvar dims IN the preserved list with their Vars in the INTERMEDIATE Func
         for (auto &dim : intm_dims) {
             const auto it = std::find_if(preserved_rvars.begin(), preserved_rvars.end(), [&](const auto &rv) {
-                return var_name_match(dim.var, rv.name());
+                return dim_match(dim, rv);
             });
             if (it != preserved_rvars.end()) {
                 const auto offset = it - preserved_rvars.begin();

From 1355d3e8d9cdab4fba10f31b562cdb6f6ffcf2f6 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 10:48:03 -0500
Subject: [PATCH 12/39] Use unordered_set directly instead of DimSet

---
 src/Func.cpp | 35 ++++++++++++++---------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 6ae66636b7b5..a4a9363494dd 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -41,6 +41,7 @@ using std::pair;
 using std::string;
 using std::tuple;
 using std::unordered_map;
+using std::unordered_set;
 using std::vector;
 
 using namespace Internal;
@@ -623,18 +624,6 @@ Func Stage::rfactor(const RVar &r, const Var &v) {
 // Helpers for rfactor implementation
 namespace {
 
-struct DimHash {
-    std::size_t operator()(const Dim &s) const noexcept {
-        return std::hash<std::string>{}(s.var);
-    }
-};
-struct DimEq {
-    bool operator()(const Dim &lhs, const Dim &rhs) const noexcept {
-        return lhs.var == rhs.var;
-    }
-};
-using DimSet = std::unordered_set<Dim, DimHash, DimEq>;
-
 optional<RVar> find_rvar(const vector<RVar> &items, const Dim &dim) {
     const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
         return dim_match(dim, x);
@@ -656,7 +645,7 @@ optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved,
         return "can't perform rfactor() because we can't prove associativity of the operator";
     }
 
-    DimSet is_rfactored;
+    unordered_set<string> is_rfactored;
     for (const auto &[rv, v] : preserved) {
         // Check that the RVars are in the dims list
         const auto &rv_dim = find_dim(dims, rv);
@@ -665,7 +654,7 @@ optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved,
             s << "can't perform rfactor() on " << rv.name() << " since it is not in the reduction domain";
             return s.str();
         }
-        is_rfactored.insert(*rv_dim);
+        is_rfactored.insert(rv_dim->var);
 
         // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
         if (find_dim(dims, v)) {
@@ -681,7 +670,7 @@ optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved,
     if (!prover_result.commutative()) {
         optional<Dim> last_rvar;
         for (const auto &d : reverse_view(dims)) {
-            if (is_rfactored.count(d) && last_rvar && !is_rfactored.count(*last_rvar)) {
+            if (is_rfactored.count(d.var) && last_rvar && !is_rfactored.count(last_rvar->var)) {
                 std::stringstream s;
                 s << "can't rfactor an inner dimension " << d.var
                   << " without rfactoring the outer dimensions, since the "
@@ -931,12 +920,14 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         }
 
         // Add factored pure dims to the INTERMEDIATE func just before outermost
-        DimSet dims;
-        dims.insert(intm_dims.begin(), intm_dims.end());
+        unordered_set<string> dims;
+        for (const auto &dim : intm_dims) {
+            dims.insert(dim.var);
+        }
         for (const Var &dim_v : preserved_vars) {
             const optional<Dim> &dim = find_dim(intm.function().definition().schedule().dims(), dim_v);
             internal_assert(dim) << "Failed to find " << dim_v.name() << " in list of pure dims";
-            if (!dims.count(*dim)) {
+            if (!dims.count(dim->var)) {
                 intm_dims.insert(intm_dims.end() - 1, *dim);
             }
         }
@@ -980,12 +971,14 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
         vector<Dim> reducing_dims;
         {
-            DimSet preserved_rdim_set;
-            preserved_rdim_set.insert(preserved_rdims.begin(), preserved_rdims.end());
+            unordered_set<string> preserved_rdim_set;
+            for (const auto &dim : preserved_rdims) {
+                preserved_rdim_set.insert(dim.var);
+            }
 
             // Remove rvar dims NOT IN the preserved list from the REDUCING Func
             for (const auto &dim : definition.schedule().dims()) {
-                if (!dim.is_rvar() || preserved_rdim_set.count(dim)) {
+                if (!dim.is_rvar() || preserved_rdim_set.count(dim.var)) {
                     reducing_dims.push_back(dim);
                 }
             }

From a201afea818589783a25eb17eddde31e0b19b71d Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 10:57:24 -0500
Subject: [PATCH 13/39] Compute preserved rdims set earlier to drop find_rvar

---
 src/Func.cpp | 23 +++++++----------------
 1 file changed, 7 insertions(+), 16 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index a4a9363494dd..a17bc3ee4057 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -624,13 +624,6 @@ Func Stage::rfactor(const RVar &r, const Var &v) {
 // Helpers for rfactor implementation
 namespace {
 
-optional<RVar> find_rvar(const vector<RVar> &items, const Dim &dim) {
-    const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
-        return dim_match(dim, x);
-    });
-    return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
-}
-
 optional<Dim> find_dim(const vector<Dim> &items, const VarOrRVar &v) {
     const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) {
         return dim_match(x, v);
@@ -841,6 +834,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     vector<RVar> preserved_rvars;
     vector<Var> preserved_vars;
     vector<Dim> preserved_rdims;
+    unordered_set<string> preserved_rdims_set;
     vector<Dim> intermediate_rdims;
     {
         unordered_map<string, int> dim_ordering;
@@ -848,14 +842,15 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             dim_ordering.emplace(definition.schedule().dims()[i].var, i);
         }
 
-        vector<tuple<RVar, Var, Dim>> preserved_with_dims;
+        using PreservedData = tuple<RVar, Var, Dim>;
+        vector<PreservedData> preserved_with_dims;
         for (const auto &[rv, v] : preserved) {
             const optional<Dim> rdim = find_dim(definition.schedule().dims(), rv);
             internal_assert(rdim);
             preserved_with_dims.emplace_back(rv, v, *rdim);
         }
 
-        std::sort(preserved_with_dims.begin(), preserved_with_dims.end(), [&](const auto &lhs, const auto &rhs) {
+        std::sort(preserved_with_dims.begin(), preserved_with_dims.end(), [&](const PreservedData &lhs, const PreservedData &rhs) {
             return dim_ordering.at(std::get<2>(lhs).var) < dim_ordering.at(std::get<2>(rhs).var);
         });
 
@@ -863,10 +858,11 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             preserved_rvars.push_back(rv);
             preserved_vars.push_back(v);
             preserved_rdims.push_back(dim);
+            preserved_rdims_set.insert(dim.var);
         }
 
         for (const Dim &dim : definition.schedule().dims()) {
-            if (dim.is_rvar() && !find_rvar(preserved_rvars, dim)) {
+            if (dim.is_rvar() && !preserved_rdims_set.count(dim.var)) {
                 intermediate_rdims.push_back(dim);
             }
         }
@@ -971,14 +967,9 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
         vector<Dim> reducing_dims;
         {
-            unordered_set<string> preserved_rdim_set;
-            for (const auto &dim : preserved_rdims) {
-                preserved_rdim_set.insert(dim.var);
-            }
-
             // Remove rvar dims NOT IN the preserved list from the REDUCING Func
             for (const auto &dim : definition.schedule().dims()) {
-                if (!dim.is_rvar() || preserved_rdim_set.count(dim.var)) {
+                if (!dim.is_rvar() || preserved_rdims_set.count(dim.var)) {
                     reducing_dims.push_back(dim);
                 }
             }

From 7a8c68679f9128c35c762266e91e682301e40c7c Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 11:35:50 -0500
Subject: [PATCH 14/39] Hoist projection code into common block

---
 src/Func.cpp | 53 ++++++++++++++++++++++++++--------------------------
 1 file changed, 27 insertions(+), 26 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index a17bc3ee4057..9e21ff4137ba 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -868,6 +868,29 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         }
     }
 
+    // Project the RDom into each side
+    ReductionDomain intermediate_rdom, preserved_rdom;
+    SubstitutionMap intermediate_map, preserved_map;
+    {
+        // Intermediate
+        std::tie(intermediate_rdom, intermediate_map) = project_rdom(intermediate_rdims, definition);
+        for (size_t i = 0; i < preserved.size(); i++) {
+            rebind(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
+        }
+        for (const auto &var : dim_vars) {
+            intermediate_map.erase(var.name());
+        }
+        intermediate_rdom.set_predicate(simplify(substitute(intermediate_map, intermediate_rdom.predicate())));
+
+        // Preserved
+        std::tie(preserved_rdom, preserved_map) = project_rdom(preserved_rdims, definition);
+        Scope<Interval> intm_rdom;
+        for (const auto &[var, min, extent] : intermediate_rdom.domain()) {
+            intm_rdom.push(var, Interval{min, min + extent - 1});
+        }
+        preserved_rdom.set_predicate(!and_condition_over_domain(simplify(substitute(preserved_map, !preserved_rdom.predicate())), intm_rdom));
+    }
+
     // Intermediate func
     Func intm(function.name() + "_intm");
 
@@ -878,16 +901,6 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
     // Intermediate update definition
     {
-        auto [intermediate_rdom, intermediate_map] = project_rdom(intermediate_rdims, definition);
-        for (size_t i = 0; i < preserved.size(); i++) {
-            rebind(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
-        }
-        for (const auto &var : dim_vars) {
-            intermediate_map.erase(var.name());
-        }
-
-        intermediate_rdom.set_predicate(simplify(substitute(intermediate_map, intermediate_rdom.predicate())));
-
         vector<Expr> args = substitute(intermediate_map, definition.args() + preserved_vars);
         vector<Expr> values;
         for (const auto &val : definition.values()) {
@@ -933,9 +946,6 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
     // Preserved update definition
     {
-        auto [preserved_rdom, preserved_map] = project_rdom(preserved_rdims, definition);
-        preserved_rdom.set_predicate(simplify(substitute(preserved_map, preserved_rdom.predicate())));
-
         // Replace the current definition with calls to the intermediate func.
         vector<Expr> dim_exprs = copy_convert<Expr>(dim_vars);
         vector<Expr> f_load_args = dim_exprs;
@@ -943,21 +953,17 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             f_load_args.push_back(Variable::make(Int(32), rv.var, preserved_rdom));
         }
 
-        SubstitutionMap replacements;
-        for (size_t i = 0; i < preserved.size(); i++) {
-            replacements.emplace(preserved_rdims[i].var, preserved_vars[i]);
-        }
         for (size_t i = 0; i < definition.values().size(); ++i) {
             if (!prover_result.ys[i].var.empty()) {
                 Expr r = (definition.values().size() == 1) ? Expr(intm(f_load_args)) : Expr(intm(f_load_args)[i]);
-                replacements.emplace(prover_result.ys[i].var, r);
+                add_let(preserved_map, prover_result.ys[i].var, r);
             }
 
             if (!prover_result.xs[i].var.empty()) {
                 Expr prev_val = Call::make(intm.types()[i], function.name(),
                                            dim_exprs, Call::CallType::Halide,
                                            FunctionPtr(), i);
-                replacements.emplace(prover_result.xs[i].var, prev_val);
+                add_let(preserved_map, prover_result.xs[i].var, prev_val);
             } else {
                 user_warning << "Update definition of " << name() << " at index " << i
                              << " doesn't depend on the previous value. This isn't a"
@@ -983,16 +989,11 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             }
         }
 
-        Scope<Interval> intm_rdom;
-        for (const auto &[var, min, extent] : intm.update(0).definition.schedule().rvars()) {
-            intm_rdom.push(var, Interval{min, min + extent - 1});
-        }
-
         definition.args() = dim_exprs;
-        definition.predicate() = !and_condition_over_domain(simplify(!preserved_rdom.predicate()), intm_rdom);
+        definition.predicate() = preserved_rdom.predicate();
         definition.schedule().dims() = std::move(reducing_dims);
         definition.schedule().rvars() = preserved_rdom.domain();
-        definition.values() = substitute(replacements, prover_result.pattern.ops);
+        definition.values() = substitute(preserved_map, prover_result.pattern.ops);
     }
 
     // Clean up the splits lists

From 7819855678b08bc51278a9ab7ce798e8ae31c4e7 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 12:13:56 -0500
Subject: [PATCH 15/39] Use structured bindings

---
 src/Func.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 9e21ff4137ba..a36f80a6cc45 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -762,10 +762,10 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
             add_let(bounds_projection, name, value);
         }
     }
-    for (const ReductionVariable &rv : rdom) {
-        add_let(bounds_projection, rv.var + ".loop_min", rv.min);
-        add_let(bounds_projection, rv.var + ".loop_max", rv.min + rv.extent - 1);
-        add_let(bounds_projection, rv.var + ".loop_extent", rv.extent);
+    for (const auto &[var, min, extent] : rdom) {
+        add_let(bounds_projection, var + ".loop_min", min);
+        add_let(bounds_projection, var + ".loop_max", min + extent - 1);
+        add_let(bounds_projection, var + ".loop_extent", extent);
     }
 
     // Build the new RDom

From 9b7a4a22609e268148ef03abe6f65a808cc34f2b Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 12:37:36 -0500
Subject: [PATCH 16/39] Drop rebind() as add_let() was equivalent

---
 src/Func.cpp | 19 +++++--------------
 1 file changed, 5 insertions(+), 14 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index a36f80a6cc45..dc77fb89be46 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -741,15 +741,6 @@ void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
     proj.emplace(name, value);
 }
 
-void rebind(SubstitutionMap &proj, const string &name, const Expr &value) {
-    for (auto &[_, e] : proj) {
-        e = substitute(name, value, e);
-    }
-    if (!proj.count(name)) {
-        proj.emplace(name, value);
-    }
-}
-
 pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const Definition &def) {
     const auto &rdom = def.schedule().rvars();
     const auto &splits = def.schedule().splits();
@@ -788,9 +779,6 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
     for (const Split &split : splits) {
         for (const auto &result : apply_split(split, "", dim_extent_alignment)) {
             switch (result.type) {
-            case ApplySplitResult::Substitution:
-            case ApplySplitResult::SubstitutionInCalls:
-            case ApplySplitResult::SubstitutionInProvides:
             case ApplySplitResult::LetStmt:
                 add_let(dim_projection, result.name, result.value);
                 break;
@@ -799,6 +787,9 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
             case ApplySplitResult::Predicate:
                 new_rdom.where(substitute(bounds_projection, result.value));
                 break;
+            case ApplySplitResult::Substitution:
+            case ApplySplitResult::SubstitutionInCalls:
+            case ApplySplitResult::SubstitutionInProvides:
             case ApplySplitResult::BlendProvides:
                 // TODO: what to do here? BlendProvides is not included in the above checks.
                 break;
@@ -809,7 +800,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
         e = substitute(bounds_projection, e);
     }
     for (const ReductionVariable &rv : new_rdom.domain()) {
-        rebind(dim_projection, rv.var, Variable::make(Int(32), rv.var, new_rdom));
+        add_let(dim_projection, rv.var, Variable::make(Int(32), rv.var, new_rdom));
     }
 
     return {new_rdom, dim_projection};
@@ -875,7 +866,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         // Intermediate
         std::tie(intermediate_rdom, intermediate_map) = project_rdom(intermediate_rdims, definition);
         for (size_t i = 0; i < preserved.size(); i++) {
-            rebind(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
+            add_let(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
         }
         for (const auto &var : dim_vars) {
             intermediate_map.erase(var.name());

From ce9af71e2c77411f1fd48e018c5e9694ca45ab2a Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 12:55:29 -0500
Subject: [PATCH 17/39] More cleaning

---
 src/Func.cpp                                     | 10 +++++-----
 test/error/rfactor_after_var_and_rvar_fusion.cpp |  2 +-
 test/error/rfactor_fused_var_and_rvar.cpp        |  2 +-
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index dc77fb89be46..88ad5e78986e 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -772,9 +772,9 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
     // Compute the mapping of old dimensions to the projected RDom values.
     SubstitutionMap dim_projection{};
     SubstitutionMap dim_extent_alignment{};
-    for (const ReductionVariable &rv : rdom) {
-        add_let(dim_projection, rv.var, Variable::make(Int(32), rv.var));
-        dim_extent_alignment[rv.var] = rv.extent;
+    for (const auto &[var, _, extent] : rdom) {
+        add_let(dim_projection, var, Variable::make(Int(32), var));
+        dim_extent_alignment[var] = extent;
     }
     for (const Split &split : splits) {
         for (const auto &result : apply_split(split, "", dim_extent_alignment)) {
@@ -791,7 +791,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
             case ApplySplitResult::SubstitutionInCalls:
             case ApplySplitResult::SubstitutionInProvides:
             case ApplySplitResult::BlendProvides:
-                // TODO: what to do here? BlendProvides is not included in the above checks.
+                // The lets returned by ApplySplit are sufficient
                 break;
             }
         }
@@ -799,7 +799,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
     for (auto &[_, e] : dim_projection) {
         e = substitute(bounds_projection, e);
     }
-    for (const ReductionVariable &rv : new_rdom.domain()) {
+    for (const auto &rv : new_rdom.domain()) {
         add_let(dim_projection, rv.var, Variable::make(Int(32), rv.var, new_rdom));
     }
 
diff --git a/test/error/rfactor_after_var_and_rvar_fusion.cpp b/test/error/rfactor_after_var_and_rvar_fusion.cpp
index 8b94fbdd1b16..acda4e4bb6fb 100644
--- a/test/error/rfactor_after_var_and_rvar_fusion.cpp
+++ b/test/error/rfactor_after_var_and_rvar_fusion.cpp
@@ -22,4 +22,4 @@ int main(int argc, char **argv) {
 
     printf("Success!\n");
     return 0;
-}
\ No newline at end of file
+}
diff --git a/test/error/rfactor_fused_var_and_rvar.cpp b/test/error/rfactor_fused_var_and_rvar.cpp
index a167ca543c47..64a79c269690 100644
--- a/test/error/rfactor_fused_var_and_rvar.cpp
+++ b/test/error/rfactor_fused_var_and_rvar.cpp
@@ -23,4 +23,4 @@ int main(int argc, char **argv) {
 
     printf("Success!\n");
     return 0;
-}
\ No newline at end of file
+}

From 576654dbc055fddf69315ee66a97ac6a8a7dba3a Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 16:56:27 -0500
Subject: [PATCH 18/39] Remove not-helpful-enough helpers

---
 src/Func.cpp | 51 +++++++++++++++++++++++++--------------------------
 1 file changed, 25 insertions(+), 26 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 88ad5e78986e..a8c00746f410 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -607,11 +607,14 @@ class SubstituteSelfReference : public IRMutator {
 /** Substitute all self-reference calls to 'func' with 'substitute' which
  * args (LHS) is the old args (LHS) plus 'new_args' in that order.
  * Expect this method to be called on the value (RHS) of an update definition. */
-Expr substitute_self_reference(Expr val, const string &func, const Function &substitute,
-                               const vector<Var> &new_args) {
+vector<Expr> substitute_self_reference(const vector<Expr> &values, const string &func,
+                                       const Function &substitute, const vector<Var> &new_args) {
     SubstituteSelfReference subs(func, substitute, new_args);
-    val = subs.mutate(val);
-    return val;
+    vector<Expr> result;
+    for (const auto &val : values) {
+        result.push_back(subs.mutate(val));
+    }
+    return result;
 }
 
 }  // anonymous namespace
@@ -716,18 +719,6 @@ optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved,
     return std::nullopt;
 }
 
-template<typename T, typename U>
-vector<T> operator+(const vector<T> &base, const vector<U> &other) {
-    vector<T> merged = base;
-    merged.insert(merged.end(), other.begin(), other.end());
-    return merged;
-}
-
-template<typename U, typename V>
-vector<U> copy_convert(const vector<V> &vec) {
-    return vector<U>{} + vec;
-}
-
 using SubstitutionMap = std::map<string, Expr>;
 
 // This is a helper function for building up a substitution map that
@@ -821,6 +812,12 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     user_assert(!error) << "In schedule for " << name() << ": " << *error << "\n"
                         << dump_argument_list();
 
+    const vector<Expr> dim_vars_exprs = [&] {
+        vector<Expr> result;
+        result.insert(result.end(), dim_vars.begin(), dim_vars.end());
+        return result;
+    }();
+
     // sort preserved by the dimension ordering
     vector<RVar> preserved_rvars;
     vector<Var> preserved_vars;
@@ -887,16 +884,19 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
     // Intermediate pure definition
     {
-        intm(dim_vars + preserved_vars) = Tuple(prover_result.pattern.identities);
+        vector<Expr> args = dim_vars_exprs;
+        args.insert(args.end(), preserved_vars.begin(), preserved_vars.end());
+        intm(args) = Tuple(prover_result.pattern.identities);
     }
 
     // Intermediate update definition
     {
-        vector<Expr> args = substitute(intermediate_map, definition.args() + preserved_vars);
-        vector<Expr> values;
-        for (const auto &val : definition.values()) {
-            values.push_back(substitute_self_reference(val, function.name(), intm.function(), preserved_vars));
-        }
+        vector<Expr> args = definition.args();
+        args.insert(args.end(), preserved_vars.begin(), preserved_vars.end());
+        args = substitute(intermediate_map, args);
+
+        vector<Expr> values = definition.values();
+        values = substitute_self_reference(values, function.name(), intm.function(), preserved_vars);
         values = substitute(intermediate_map, values);
         intm.function().define_update(args, values, intermediate_rdom);
 
@@ -938,8 +938,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     // Preserved update definition
     {
         // Replace the current definition with calls to the intermediate func.
-        vector<Expr> dim_exprs = copy_convert<Expr>(dim_vars);
-        vector<Expr> f_load_args = dim_exprs;
+        vector<Expr> f_load_args = dim_vars_exprs;
         for (const ReductionVariable &rv : preserved_rdom.domain()) {
             f_load_args.push_back(Variable::make(Int(32), rv.var, preserved_rdom));
         }
@@ -952,7 +951,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
 
             if (!prover_result.xs[i].var.empty()) {
                 Expr prev_val = Call::make(intm.types()[i], function.name(),
-                                           dim_exprs, Call::CallType::Halide,
+                                           dim_vars_exprs, Call::CallType::Halide,
                                            FunctionPtr(), i);
                 add_let(preserved_map, prover_result.xs[i].var, prev_val);
             } else {
@@ -980,7 +979,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             }
         }
 
-        definition.args() = dim_exprs;
+        definition.args() = dim_vars_exprs;
         definition.predicate() = preserved_rdom.predicate();
         definition.schedule().dims() = std::move(reducing_dims);
         definition.schedule().rvars() = preserved_rdom.domain();

From fd04843efea14a68f8b249a841610efa97161feb Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Mon, 25 Nov 2024 17:29:22 -0500
Subject: [PATCH 19/39] Reorganize update definitions to mirror each other

---
 src/Func.cpp | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index a8c00746f410..52fbd6225579 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -830,15 +830,14 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             dim_ordering.emplace(definition.schedule().dims()[i].var, i);
         }
 
-        using PreservedData = tuple<RVar, Var, Dim>;
-        vector<PreservedData> preserved_with_dims;
+        vector<tuple<RVar, Var, Dim>> preserved_with_dims;
         for (const auto &[rv, v] : preserved) {
             const optional<Dim> rdim = find_dim(definition.schedule().dims(), rv);
             internal_assert(rdim);
             preserved_with_dims.emplace_back(rv, v, *rdim);
         }
 
-        std::sort(preserved_with_dims.begin(), preserved_with_dims.end(), [&](const PreservedData &lhs, const PreservedData &rhs) {
+        std::sort(preserved_with_dims.begin(), preserved_with_dims.end(), [&](const auto &lhs, const auto &rhs) {
             return dim_ordering.at(std::get<2>(lhs).var) < dim_ordering.at(std::get<2>(rhs).var);
         });
 
@@ -901,9 +900,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         intm.function().define_update(args, values, intermediate_rdom);
 
         // Intermediate schedule
-        intm.function().update(0).schedule() = definition.schedule().get_copy();
-
-        auto &intm_dims = intm.function().update(0).schedule().dims();
+        vector<Dim> intm_dims = definition.schedule().dims();
 
         // Replace rvar dims IN the preserved list with their Vars in the INTERMEDIATE Func
         for (auto &dim : intm_dims) {
@@ -924,14 +921,16 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         for (const auto &dim : intm_dims) {
             dims.insert(dim.var);
         }
-        for (const Var &dim_v : preserved_vars) {
-            const optional<Dim> &dim = find_dim(intm.function().definition().schedule().dims(), dim_v);
-            internal_assert(dim) << "Failed to find " << dim_v.name() << " in list of pure dims";
+        for (const Var &var : preserved_vars) {
+            const optional<Dim> &dim = find_dim(intm.function().definition().schedule().dims(), var);
+            internal_assert(dim) << "Failed to find " << var.name() << " in list of pure dims";
             if (!dims.count(dim->var)) {
                 intm_dims.insert(intm_dims.end() - 1, *dim);
             }
         }
 
+        intm.function().update(0).schedule() = definition.schedule().get_copy();
+        intm.function().update(0).schedule().dims() = std::move(intm_dims);
         intm.function().update(0).schedule().rvars() = intermediate_rdom.domain();
     }
 
@@ -980,10 +979,10 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         }
 
         definition.args() = dim_vars_exprs;
+        definition.values() = substitute(preserved_map, prover_result.pattern.ops);
         definition.predicate() = preserved_rdom.predicate();
         definition.schedule().dims() = std::move(reducing_dims);
         definition.schedule().rvars() = preserved_rdom.domain();
-        definition.values() = substitute(preserved_map, prover_result.pattern.ops);
     }
 
     // Clean up the splits lists

From a57d6700d492292fc6e480b36c15ccd05c560ccc Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Tue, 3 Dec 2024 09:58:50 -0800
Subject: [PATCH 20/39] Update Func.cpp

---
 src/Func.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/Func.cpp b/src/Func.cpp
index 52fbd6225579..e8827a16c65c 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -1,6 +1,7 @@
 #include <algorithm>
 #include <cstring>
 #include <iostream>
+#include <unordered_map>
 #include <unordered_set>
 #include <utility>
 

From 01f441590b51fb44457f1f9e525470fcc76224f3 Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Tue, 3 Dec 2024 10:21:16 -0800
Subject: [PATCH 21/39] Update Substitute.h

---
 src/Substitute.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/Substitute.h b/src/Substitute.h
index 55e7ac1cf0fa..e514fda40359 100644
--- a/src/Substitute.h
+++ b/src/Substitute.h
@@ -6,6 +6,7 @@
  * Defines methods for substituting out variables in expressions and
  * statements. */
 
+#include <algorithm>
 #include <map>
 
 #include "Expr.h"

From 4cf1d3c747f488efad7855fb80c8bc4702eb12f6 Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Tue, 10 Dec 2024 09:26:51 -0800
Subject: [PATCH 22/39] trigger buildbots


From 4bce0a78f73ebe5fda2fcedb3039f9020b87ca35 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Tue, 10 Dec 2024 17:10:31 -0500
Subject: [PATCH 23/39] Simplify the preserved predicate to eliminate outermost
 NOT.

---
 src/Func.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index e8827a16c65c..41c4a14e6af1 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -876,7 +876,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         for (const auto &[var, min, extent] : intermediate_rdom.domain()) {
             intm_rdom.push(var, Interval{min, min + extent - 1});
         }
-        preserved_rdom.set_predicate(!and_condition_over_domain(simplify(substitute(preserved_map, !preserved_rdom.predicate())), intm_rdom));
+        preserved_rdom.set_predicate(simplify(!and_condition_over_domain(simplify(substitute(preserved_map, !preserved_rdom.predicate())), intm_rdom)));
     }
 
     // Intermediate func

From d7572984ec7ff8508f8b5146f2f602670cbbeef5 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Tue, 10 Dec 2024 17:13:22 -0500
Subject: [PATCH 24/39] Assert that add_let would not shadow a binding

---
 src/Func.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 41c4a14e6af1..6ce18f1ffe5d 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -727,6 +727,9 @@ using SubstitutionMap = std::map<string, Expr>;
 // to this function from innermost to outermost. This is equivalent to
 // building a let-nest as a one-hole context and then simplifying.
 void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
+    internal_assert(!proj.count(name)) << "would shadow " << name << " in let nest.\n"
+                                       << "\tPresent value: " << proj[name] << "\n"
+                                       << "\tProposed value: " << value;
     for (auto &[_, e] : proj) {
         e = substitute(name, value, e);
     }
@@ -762,12 +765,11 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
     new_rdom.where(predicate);
 
     // Compute the mapping of old dimensions to the projected RDom values.
-    SubstitutionMap dim_projection{};
     SubstitutionMap dim_extent_alignment{};
     for (const auto &[var, _, extent] : rdom) {
-        add_let(dim_projection, var, Variable::make(Int(32), var));
         dim_extent_alignment[var] = extent;
     }
+    SubstitutionMap dim_projection{};
     for (const Split &split : splits) {
         for (const auto &result : apply_split(split, "", dim_extent_alignment)) {
             switch (result.type) {

From e397d7a48e134f13811de344663e3cd7a564ddbb Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Tue, 10 Dec 2024 17:51:58 -0500
Subject: [PATCH 25/39] Improve "rvar not found" error message.

---
 src/Func.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 6ce18f1ffe5d..90ced77a112a 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -648,7 +648,9 @@ optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved,
         const auto &rv_dim = find_dim(dims, rv);
         if (!(rv_dim && rv_dim->is_rvar())) {
             std::stringstream s;
-            s << "can't perform rfactor() on " << rv.name() << " since it is not in the reduction domain";
+            s << "can't perform rfactor() on " << rv.name()
+              << " since either it is not in the reduction domain, or has"
+              << " already been consumed by another scheduling directive";
             return s.str();
         }
         is_rfactored.insert(rv_dim->var);

From 70798dd715fb931e795bbc8230084b928d4f4a2b Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Tue, 10 Dec 2024 18:22:26 -0500
Subject: [PATCH 26/39] Further cleanup using the no-shadow invariant

---
 src/Func.cpp | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 90ced77a112a..e53944b37ee3 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -759,9 +759,9 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
     // Build the new RDom
     vector<ReductionVariable> new_rvars;
     for (const Dim &dim : dims) {
-        Expr new_min = substitute(bounds_projection, Variable::make(Int(32), dim.var + ".loop_min"));
-        Expr new_extent = substitute(bounds_projection, Variable::make(Int(32), dim.var + ".loop_extent"));
-        new_rvars.push_back(ReductionVariable{dim.var, simplify(new_min), simplify(new_extent)});
+        const Expr new_min = simplify(bounds_projection.at(dim.var + ".loop_min"));
+        const Expr new_extent = simplify(bounds_projection.at(dim.var + ".loop_extent"));
+        new_rvars.push_back(ReductionVariable{dim.var, new_min, new_extent});
     }
     ReductionDomain new_rdom{new_rvars};
     new_rdom.where(predicate);
@@ -776,7 +776,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
         for (const auto &result : apply_split(split, "", dim_extent_alignment)) {
             switch (result.type) {
             case ApplySplitResult::LetStmt:
-                add_let(dim_projection, result.name, result.value);
+                add_let(dim_projection, result.name, substitute(bounds_projection, result.value));
                 break;
             case ApplySplitResult::PredicateCalls:
             case ApplySplitResult::PredicateProvides:
@@ -792,13 +792,9 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
             }
         }
     }
-    for (auto &[_, e] : dim_projection) {
-        e = substitute(bounds_projection, e);
-    }
     for (const auto &rv : new_rdom.domain()) {
         add_let(dim_projection, rv.var, Variable::make(Int(32), rv.var, new_rdom));
     }
-
     return {new_rdom, dim_projection};
 }
 

From b0e9b29aa12e2039380da6888407a186f91e7d7e Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Wed, 11 Dec 2024 14:14:00 -0500
Subject: [PATCH 27/39] MAke rfactor_validate_args a private member

---
 src/Func.cpp | 175 +++++++++++++++++++++++++--------------------------
 src/Func.h   |   3 +
 2 files changed, 88 insertions(+), 90 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index e53944b37ee3..556ca19bab0f 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -635,93 +635,6 @@ optional<Dim> find_dim(const vector<Dim> &items, const VarOrRVar &v) {
     return has_v == items.end() ? std::nullopt : std::make_optional(*has_v);
 }
 
-optional<string> rfactor_validate_args(const vector<pair<RVar, Var>> &preserved, const AssociativeOp &prover_result, const Definition &definition) {
-    const vector<Dim> &dims = definition.schedule().dims();
-
-    if (!prover_result.associative()) {
-        return "can't perform rfactor() because we can't prove associativity of the operator";
-    }
-
-    unordered_set<string> is_rfactored;
-    for (const auto &[rv, v] : preserved) {
-        // Check that the RVars are in the dims list
-        const auto &rv_dim = find_dim(dims, rv);
-        if (!(rv_dim && rv_dim->is_rvar())) {
-            std::stringstream s;
-            s << "can't perform rfactor() on " << rv.name()
-              << " since either it is not in the reduction domain, or has"
-              << " already been consumed by another scheduling directive";
-            return s.str();
-        }
-        is_rfactored.insert(rv_dim->var);
-
-        // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
-        if (find_dim(dims, v)) {
-            std::stringstream s;
-            s << "can't rename the rvars " << rv.name() << " into " << v.name()
-              << ", since it is already used in this Func's schedule elsewhere.";
-            return s.str();
-        }
-    }
-
-    // If the operator is associative but non-commutative, rfactor() on inner
-    // dimensions (excluding the outer dimensions) is not valid.
-    if (!prover_result.commutative()) {
-        optional<Dim> last_rvar;
-        for (const auto &d : reverse_view(dims)) {
-            if (is_rfactored.count(d.var) && last_rvar && !is_rfactored.count(last_rvar->var)) {
-                std::stringstream s;
-                s << "can't rfactor an inner dimension " << d.var
-                  << " without rfactoring the outer dimensions, since the "
-                  << "operator is non-commutative.";
-                return s.str();
-            }
-            if (d.is_rvar()) {
-                last_rvar = d;
-            }
-        }
-    }
-
-    // Check that no Vars were fused into RVars
-    Scope<> rdims;
-    for (const ReductionVariable &rv : definition.schedule().rvars()) {
-        rdims.push(rv.var);
-    }
-    for (const Split &split : definition.schedule().splits()) {
-        switch (split.split_type) {
-        case Split::SplitVar:
-            if (rdims.contains(split.old_var)) {
-                rdims.pop(split.old_var);
-                rdims.push(split.outer);
-                rdims.push(split.inner);
-            }
-            break;
-        case Split::FuseVars:
-            if (rdims.contains(split.outer) || rdims.contains(split.inner)) {
-                if (!(rdims.contains(split.outer) && rdims.contains(split.inner))) {
-                    std::stringstream s;
-                    s << "Cannot rfactor a Func that has fused a Var into an RVar ("
-                      << split.outer << ", " << split.inner << ")";
-                    return s.str();
-                }
-                rdims.pop(split.outer);
-                rdims.pop(split.inner);
-                rdims.push(split.old_var);
-            }
-            break;
-        case Split::PurifyRVar:
-        case Split::RenameVar:
-            if (rdims.contains(split.old_var)) {
-                rdims.pop(split.old_var);
-                rdims.push(split.outer);
-            }
-            break;
-        }
-    }
-
-    return std::nullopt;
-}
-
 using SubstitutionMap = std::map<string, Expr>;
 
 // This is a helper function for building up a substitution map that
@@ -800,6 +713,90 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
 
 }  // namespace
 
+void Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &preserved, const AssociativeOp &prover_result) {
+    const vector<Dim> &dims = definition.schedule().dims();
+
+    user_assert(prover_result.associative())
+        << "In schedule for " << name() << ": can't perform rfactor() "
+        << "because we can't prove associativity of the operator\n"
+        << dump_argument_list();
+
+    unordered_set<string> is_rfactored;
+    for (const auto &[rv, v] : preserved) {
+        // Check that the RVars are in the dims list
+        const auto &rv_dim = find_dim(dims, rv);
+        user_assert(rv_dim && rv_dim->is_rvar())
+            << "In schedule for " << name() << ": can't perform rfactor() "
+            << "on " << rv.name() << " since either it is not in the reduction "
+            << "domain, or has already been consumed by another scheduling directive\n"
+            << dump_argument_list();
+
+        is_rfactored.insert(rv_dim->var);
+
+        // Check that the new pure Vars we used to rename the RVar aren't already in the dims list
+        user_assert(!find_dim(dims, v))
+            << "In schedule for " << name() << ": can't perform rfactor() "
+            << "on " << rv.name() << " because the name " << v.name()
+            << "is already used elsewhere in the Func's schedule.\n"
+            << dump_argument_list();
+    }
+
+    // If the operator is associative but non-commutative, rfactor() on inner
+    // dimensions (excluding the outer dimensions) is not valid.
+    if (!prover_result.commutative()) {
+        optional<Dim> last_rvar;
+        for (const auto &d : reverse_view(dims)) {
+            bool is_inner = is_rfactored.count(d.var) && last_rvar && !is_rfactored.count(last_rvar->var);
+            user_assert(!is_inner)
+                << "In schedule for " << name() << ": can't rfactor an inner "
+                << "dimension " << d.var << " without rfactoring the outer "
+                << "dimensions, since the operator is non-commutative.\n"
+                << dump_argument_list();
+
+            if (d.is_rvar()) {
+                last_rvar = d;
+            }
+        }
+    }
+
+    // Check that no Vars were fused into RVars
+    Scope<> rdims;
+    for (const ReductionVariable &rv : definition.schedule().rvars()) {
+        rdims.push(rv.var);
+    }
+    for (const Split &split : definition.schedule().splits()) {
+        switch (split.split_type) {
+        case Split::SplitVar:
+            if (rdims.contains(split.old_var)) {
+                rdims.pop(split.old_var);
+                rdims.push(split.outer);
+                rdims.push(split.inner);
+            }
+            break;
+        case Split::FuseVars:
+            if (rdims.contains(split.outer) || rdims.contains(split.inner)) {
+                user_assert(rdims.contains(split.outer) && rdims.contains(split.inner))
+                    << "In schedule for " << name() << ": can't rfactor an Func "
+                    << "that has fused a Var into an RVar: " << split.outer
+                    << ", " << split.inner << "\n"
+                    << dump_argument_list();
+
+                rdims.pop(split.outer);
+                rdims.pop(split.inner);
+                rdims.push(split.old_var);
+            }
+            break;
+        case Split::PurifyRVar:
+        case Split::RenameVar:
+            if (rdims.contains(split.old_var)) {
+                rdims.pop(split.old_var);
+                rdims.push(split.outer);
+            }
+            break;
+        }
+    }
+}
+
 Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     user_assert(!definition.is_init()) << "rfactor() must be called on an update definition\n";
 
@@ -809,9 +806,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     // its identity for each value in the definition if it is a Tuple
     const auto &prover_result = prove_associativity(function.name(), definition.args(), definition.values());
 
-    const auto &error = rfactor_validate_args(preserved, prover_result, definition);
-    user_assert(!error) << "In schedule for " << name() << ": " << *error << "\n"
-                        << dump_argument_list();
+    rfactor_validate_args(preserved, prover_result);
 
     const vector<Expr> dim_vars_exprs = [&] {
         vector<Expr> result;
diff --git a/src/Func.h b/src/Func.h
index a0f374eb874f..d502d2618db0 100644
--- a/src/Func.h
+++ b/src/Func.h
@@ -60,6 +60,7 @@ struct VarOrRVar {
 class ImageParam;
 
 namespace Internal {
+struct AssociativeOp;
 class Function;
 struct Split;
 struct StorageDim;
@@ -89,6 +90,8 @@ class Stage {
 
     Stage &compute_with(LoopLevel loop_level, const std::map<std::string, LoopAlignStrategy> &align);
 
+    void rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &preserved, const Internal::AssociativeOp &prover_result);
+
 public:
     Stage(Internal::Function f, Internal::Definition d, size_t stage_index)
         : function(std::move(f)), definition(std::move(d)), stage_index(stage_index) {

From 87e52428a9e2fd459e57d345c0f34bd08491b05c Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Wed, 11 Dec 2024 14:26:10 -0500
Subject: [PATCH 28/39] Partition splits lists while validating to drop
 clean-up step.

---
 src/Func.cpp | 58 ++++++++++++++--------------------------------------
 src/Func.h   |  3 ++-
 2 files changed, 17 insertions(+), 44 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 556ca19bab0f..5d33ab5a0d6d 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -713,7 +713,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
 
 }  // namespace
 
-void Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &preserved, const AssociativeOp &prover_result) {
+pair<vector<Split>, vector<Split>> Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &preserved, const AssociativeOp &prover_result) {
     const vector<Dim> &dims = definition.schedule().dims();
 
     user_assert(prover_result.associative())
@@ -760,6 +760,7 @@ void Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &prese
     }
 
     // Check that no Vars were fused into RVars
+    vector<Split> var_splits, rvar_splits;
     Scope<> rdims;
     for (const ReductionVariable &rv : definition.schedule().rvars()) {
         rdims.push(rv.var);
@@ -771,6 +772,9 @@ void Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &prese
                 rdims.pop(split.old_var);
                 rdims.push(split.outer);
                 rdims.push(split.inner);
+                rvar_splits.emplace_back(split);
+            } else {
+                var_splits.emplace_back(split);
             }
             break;
         case Split::FuseVars:
@@ -784,6 +788,9 @@ void Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &prese
                 rdims.pop(split.outer);
                 rdims.pop(split.inner);
                 rdims.push(split.old_var);
+                rvar_splits.emplace_back(split);
+            } else {
+                var_splits.emplace_back(split);
             }
             break;
         case Split::PurifyRVar:
@@ -791,10 +798,14 @@ void Stage::rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &prese
             if (rdims.contains(split.old_var)) {
                 rdims.pop(split.old_var);
                 rdims.push(split.outer);
+                rvar_splits.emplace_back(split);
+            } else {
+                var_splits.emplace_back(split);
             }
             break;
         }
     }
+    return std::make_pair(std::move(var_splits), std::move(rvar_splits));
 }
 
 Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
@@ -806,7 +817,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     // its identity for each value in the definition if it is a Tuple
     const auto &prover_result = prove_associativity(function.name(), definition.args(), definition.values());
 
-    rfactor_validate_args(preserved, prover_result);
+    const auto &[var_splits, rvar_splits] = rfactor_validate_args(preserved, prover_result);
 
     const vector<Expr> dim_vars_exprs = [&] {
         vector<Expr> result;
@@ -928,6 +939,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         intm.function().update(0).schedule() = definition.schedule().get_copy();
         intm.function().update(0).schedule().dims() = std::move(intm_dims);
         intm.function().update(0).schedule().rvars() = intermediate_rdom.domain();
+        intm.function().update(0).schedule().splits() = var_splits;
     }
 
     // Preserved update definition
@@ -979,47 +991,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         definition.predicate() = preserved_rdom.predicate();
         definition.schedule().dims() = std::move(reducing_dims);
         definition.schedule().rvars() = preserved_rdom.domain();
-    }
-
-    // Clean up the splits lists
-    for (Stage st : {*this, intm.update(0)}) {
-        Scope<> dims;
-        for (const Var &v : st.dim_vars) {
-            dims.push(v.name());
-        }
-        for (const ReductionVariable &rv : st.definition.schedule().rvars()) {
-            dims.push(rv.var);
-        }
-        vector<Split> new_splits;
-        for (const Split &split : st.definition.schedule().splits()) {
-            switch (split.split_type) {
-            case Split::SplitVar:
-                if (dims.contains(split.old_var)) {
-                    dims.pop(split.old_var);
-                    dims.push(split.outer);
-                    dims.push(split.inner);
-                    new_splits.push_back(split);
-                }
-                break;
-            case Split::FuseVars:
-                if (dims.contains(split.outer) && dims.contains(split.inner)) {
-                    dims.pop(split.outer);
-                    dims.pop(split.inner);
-                    dims.push(split.old_var);
-                    new_splits.push_back(split);
-                }
-                break;
-            case Split::PurifyRVar:
-            case Split::RenameVar:
-                if (dims.contains(split.old_var)) {
-                    dims.pop(split.old_var);
-                    dims.push(split.outer);
-                    new_splits.push_back(split);
-                }
-                break;
-            }
-        }
-        st.definition.schedule().splits().swap(new_splits);
+        definition.schedule().splits() = var_splits;
     }
 
     return intm;
diff --git a/src/Func.h b/src/Func.h
index d502d2618db0..5a34cd803500 100644
--- a/src/Func.h
+++ b/src/Func.h
@@ -90,7 +90,8 @@ class Stage {
 
     Stage &compute_with(LoopLevel loop_level, const std::map<std::string, LoopAlignStrategy> &align);
 
-    void rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &preserved, const Internal::AssociativeOp &prover_result);
+    std::pair<std::vector<Internal::Split>, std::vector<Internal::Split>>
+    rfactor_validate_args(const std::vector<std::pair<RVar, Var>> &preserved, const Internal::AssociativeOp &prover_result);
 
 public:
     Stage(Internal::Function f, Internal::Definition d, size_t stage_index)

From c4fef0ceedecc592f8f3c02d531cb06bbd9d2866 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Wed, 11 Dec 2024 14:32:40 -0500
Subject: [PATCH 29/39] Remove PurifyRVar split type.

---
 src/ApplySplit.cpp        |  4 ----
 src/Deserialization.cpp   |  2 --
 src/Func.cpp              | 48 ---------------------------------------
 src/Func.h                |  1 -
 src/Inline.cpp            |  2 --
 src/Schedule.h            |  7 +-----
 src/ScheduleFunctions.cpp |  1 -
 src/Serialization.cpp     |  2 --
 src/halide_ir.fbs         |  1 -
 9 files changed, 1 insertion(+), 67 deletions(-)

diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp
index b6491f063fba..e6c897ec6234 100644
--- a/src/ApplySplit.cpp
+++ b/src/ApplySplit.cpp
@@ -157,7 +157,6 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
         }
     } break;
     case Split::RenameVar:
-    case Split::PurifyRVar:
         result.emplace_back(prefix + split.old_var, outer, ApplySplitResult::Substitution);
         result.emplace_back(prefix + split.old_var, outer, ApplySplitResult::LetStmt);
         break;
@@ -201,9 +200,6 @@ vector<std::pair<string, Expr>> compute_loop_bounds_after_split(const Split &spl
         let_stmts.emplace_back(prefix + split.outer + ".loop_max", old_var_max);
         let_stmts.emplace_back(prefix + split.outer + ".loop_extent", old_var_extent);
         break;
-    case Split::PurifyRVar:
-        // Do nothing for purify
-        break;
     }
 
     return let_stmts;
diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp
index 81ba48c81eae..095c12d31bf6 100644
--- a/src/Deserialization.cpp
+++ b/src/Deserialization.cpp
@@ -368,8 +368,6 @@ Split::SplitType Deserializer::deserialize_split_type(Serialize::SplitType split
         return Split::SplitType::RenameVar;
     case Serialize::SplitType::FuseVars:
         return Split::SplitType::FuseVars;
-    case Serialize::SplitType::PurifyRVar:
-        return Split::SplitType::PurifyRVar;
     default:
         user_error << "unknown split type " << (int)split_type << "\n";
         return Split::SplitType::SplitVar;
diff --git a/src/Func.cpp b/src/Func.cpp
index 5d33ab5a0d6d..e0c06cf747a6 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -431,7 +431,6 @@ void check_for_race_conditions_in_split_with_blend(const StageSchedule &sched) {
             }
             break;
         case Split::RenameVar:
-        case Split::PurifyRVar:
             if (parallel.count(split.outer)) {
                 parallel.insert(split.old_var);
             }
@@ -454,7 +453,6 @@ void check_for_race_conditions_in_split_with_blend(const StageSchedule &sched) {
             }
             break;
         case Split::RenameVar:
-        case Split::PurifyRVar:
             if (parallel.count(split.old_var)) {
                 parallel.insert(split.outer);
             }
@@ -793,7 +791,6 @@ pair<vector<Split>, vector<Split>> Stage::rfactor_validate_args(const std::vecto
                 var_splits.emplace_back(split);
             }
             break;
-        case Split::PurifyRVar:
         case Split::RenameVar:
             if (rdims.contains(split.old_var)) {
                 rdims.pop(split.old_var);
@@ -1068,7 +1065,6 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
                 }
                 break;
             case Split::RenameVar:
-            case Split::PurifyRVar:
                 if (inner_vars.count(s.old_var)) {
                     inner_vars.insert(s.outer);
                 }
@@ -1103,7 +1099,6 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
                 }
                 break;
             case Split::RenameVar:
-            case Split::PurifyRVar:
                 if (inner_vars.count(s.old_var)) {
                     inner_vars.insert(s.outer);
                 }
@@ -1165,7 +1160,6 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
                     }
                     break;
                 case Split::RenameVar:
-                case Split::PurifyRVar:
                     if (it != descends_from_shiftinwards_outer.end()) {
                         descends_from_shiftinwards_outer[s.outer] = it->second;
                     }
@@ -1352,46 +1346,6 @@ void Stage::specialize_fail(const std::string &message) {
     s.failure_message = message;
 }
 
-Stage &Stage::purify(const VarOrRVar &old_var, const VarOrRVar &new_var) {
-    user_assert(old_var.is_rvar && !new_var.is_rvar)
-        << "In schedule for " << name()
-        << ", can't rename " << (old_var.is_rvar ? "RVar " : "Var ") << old_var.name()
-        << " to " << (new_var.is_rvar ? "RVar " : "Var ") << new_var.name()
-        << "; purify must take a RVar as old_Var and a Var as new_var\n";
-
-    debug(4) << "In schedule for " << name() << ", purify RVar "
-             << old_var.name() << " to Var " << new_var.name() << "\n";
-
-    StageSchedule &schedule = definition.schedule();
-
-    // Replace the old dimension with the new dimensions in the dims list
-    bool found = false;
-    string old_name, new_name = new_var.name();
-    vector<Dim> &dims = schedule.dims();
-
-    for (size_t i = 0; (!found) && i < dims.size(); i++) {
-        if (dim_match(dims[i], old_var)) {
-            found = true;
-            old_name = dims[i].var;
-            dims[i].var = new_name;
-            dims[i].dim_type = DimType::PureVar;
-        }
-    }
-
-    if (!found) {
-        user_error
-            << "In schedule for " << name()
-            << ", could not find rename dimension: "
-            << old_var.name()
-            << "\n"
-            << dump_argument_list();
-    }
-
-    Split split = {old_name, new_name, "", 1, false, TailStrategy::RoundUp, Split::PurifyRVar};
-    definition.schedule().splits().push_back(split);
-    return *this;
-}
-
 void Stage::remove(const string &var) {
     debug(4) << "In schedule for " << name() << ", remove " << var << "\n";
 
@@ -1469,7 +1423,6 @@ void Stage::remove(const string &var) {
             }
             break;
         case Split::RenameVar:
-        case Split::PurifyRVar:
             debug(4) << "    replace/rename " << split.old_var
                      << " into " << split.outer << "\n";
             if (should_remove(split.outer)) {
@@ -1558,7 +1511,6 @@ Stage &Stage::rename(const VarOrRVar &old_var, const VarOrRVar &new_var) {
             break;
         case Split::SplitVar:
         case Split::RenameVar:
-        case Split::PurifyRVar:
             if (split.inner == old_name) {
                 split.inner = new_name;
                 found = true;
diff --git a/src/Func.h b/src/Func.h
index 5a34cd803500..96dc099ee92f 100644
--- a/src/Func.h
+++ b/src/Func.h
@@ -82,7 +82,6 @@ class Stage {
     void split(const std::string &old, const std::string &outer, const std::string &inner,
                const Expr &factor, bool exact, TailStrategy tail);
     void remove(const std::string &var);
-    Stage &purify(const VarOrRVar &old_name, const VarOrRVar &new_name);
 
     const std::vector<Internal::StorageDim> &storage_dims() const {
         return function.schedule().storage_dims();
diff --git a/src/Inline.cpp b/src/Inline.cpp
index 54399cf77b76..31b6efcdf749 100644
--- a/src/Inline.cpp
+++ b/src/Inline.cpp
@@ -76,8 +76,6 @@ void validate_schedule_inlined_function(Function f) {
                          << split.inner << " because "
                          << f.name() << " is scheduled inline.\n";
 
-            break;
-        case Split::PurifyRVar:
             break;
         }
     }
diff --git a/src/Schedule.h b/src/Schedule.h
index ea2692752a9e..906dbe6c7b5e 100644
--- a/src/Schedule.h
+++ b/src/Schedule.h
@@ -332,8 +332,7 @@ struct Split {
 
     enum SplitType { SplitVar = 0,
                      RenameVar,
-                     FuseVars,
-                     PurifyRVar };
+                     FuseVars };
 
     // If split_type is Rename, then this is just a renaming of the
     // old_var to the outer and not a split. The inner var should
@@ -341,10 +340,6 @@ struct Split {
     // the same list as splits so that ordering between them is
     // respected.
 
-    // If split type is Purify, this replaces the old_var RVar to
-    // the outer Var. The inner var should be ignored, and factor
-    // should be one.
-
     // If split_type is Fuse, then this does the opposite of a
     // split, it joins the outer and inner into the old_var.
     SplitType split_type;
diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp
index 74f0fa738e42..0cf090b3d5a5 100644
--- a/src/ScheduleFunctions.cpp
+++ b/src/ScheduleFunctions.cpp
@@ -2248,7 +2248,6 @@ bool validate_schedule(Function f, const Stmt &s, const Target &target, bool is_
                 }
                 break;
             case Split::RenameVar:
-            case Split::PurifyRVar:
                 if (parallel_vars.count(split.outer)) {
                     parallel_vars.insert(split.old_var);
                 }
diff --git a/src/Serialization.cpp b/src/Serialization.cpp
index 27bea2f5cfa3..ad1eff6ebb96 100644
--- a/src/Serialization.cpp
+++ b/src/Serialization.cpp
@@ -338,8 +338,6 @@ Serialize::SplitType Serializer::serialize_split_type(const Split::SplitType &sp
         return Serialize::SplitType::RenameVar;
     case Split::SplitType::FuseVars:
         return Serialize::SplitType::FuseVars;
-    case Split::SplitType::PurifyRVar:
-        return Serialize::SplitType::PurifyRVar;
     default:
         user_error << "Unsupported split type\n";
         return Serialize::SplitType::SplitVar;
diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs
index efc465cbee82..499488ce8b95 100644
--- a/src/halide_ir.fbs
+++ b/src/halide_ir.fbs
@@ -548,7 +548,6 @@ enum SplitType: ubyte {
     SplitVar,
     RenameVar,
     FuseVars,
-    PurifyRVar,
 }
 
 table Split {

From 20e3bb73518f92445c8c7ff282ede79c3cf03fee Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Wed, 11 Dec 2024 14:46:44 -0500
Subject: [PATCH 30/39] Only process RVar splits to eliminate intermediate map
 cleanup

---
 src/Func.cpp | 21 ++++++++-------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index e0c06cf747a6..532fed44481b 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -649,11 +649,7 @@ void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
     proj.emplace(name, value);
 }
 
-pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const Definition &def) {
-    const auto &rdom = def.schedule().rvars();
-    const auto &splits = def.schedule().splits();
-    const auto &predicate = def.predicate();
-
+pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const ReductionDomain &rdom, const vector<Split> &splits) {
     // Compute new bounds for the RDom by walking backwards through the splits.
     SubstitutionMap bounds_projection{};
     for (const Split &split : reverse_view(splits)) {
@@ -661,7 +657,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
             add_let(bounds_projection, name, value);
         }
     }
-    for (const auto &[var, min, extent] : rdom) {
+    for (const auto &[var, min, extent] : rdom.domain()) {
         add_let(bounds_projection, var + ".loop_min", min);
         add_let(bounds_projection, var + ".loop_max", min + extent - 1);
         add_let(bounds_projection, var + ".loop_extent", extent);
@@ -675,11 +671,11 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
         new_rvars.push_back(ReductionVariable{dim.var, new_min, new_extent});
     }
     ReductionDomain new_rdom{new_rvars};
-    new_rdom.where(predicate);
+    new_rdom.where(rdom.predicate());
 
     // Compute the mapping of old dimensions to the projected RDom values.
     SubstitutionMap dim_extent_alignment{};
-    for (const auto &[var, _, extent] : rdom) {
+    for (const auto &[var, _, extent] : rdom.domain()) {
         dim_extent_alignment[var] = extent;
     }
     SubstitutionMap dim_projection{};
@@ -863,18 +859,17 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     ReductionDomain intermediate_rdom, preserved_rdom;
     SubstitutionMap intermediate_map, preserved_map;
     {
+        ReductionDomain rdom{definition.schedule().rvars(), definition.predicate(), true};
+
         // Intermediate
-        std::tie(intermediate_rdom, intermediate_map) = project_rdom(intermediate_rdims, definition);
+        std::tie(intermediate_rdom, intermediate_map) = project_rdom(intermediate_rdims, rdom, rvar_splits);
         for (size_t i = 0; i < preserved.size(); i++) {
             add_let(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
         }
-        for (const auto &var : dim_vars) {
-            intermediate_map.erase(var.name());
-        }
         intermediate_rdom.set_predicate(simplify(substitute(intermediate_map, intermediate_rdom.predicate())));
 
         // Preserved
-        std::tie(preserved_rdom, preserved_map) = project_rdom(preserved_rdims, definition);
+        std::tie(preserved_rdom, preserved_map) = project_rdom(preserved_rdims, rdom, rvar_splits);
         Scope<Interval> intm_rdom;
         for (const auto &[var, min, extent] : intermediate_rdom.domain()) {
             intm_rdom.push(var, Interval{min, min + extent - 1});

From 724f0f79fc1cf958e743bf68b40456a376c0439d Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Wed, 11 Dec 2024 15:24:19 -0500
Subject: [PATCH 31/39] Fix comment in ApplySplit.cpp

---
 src/ApplySplit.cpp | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp
index e6c897ec6234..22c3425c02a4 100644
--- a/src/ApplySplit.cpp
+++ b/src/ApplySplit.cpp
@@ -166,10 +166,7 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
 }
 
 vector<std::pair<string, Expr>> compute_loop_bounds_after_split(const Split &split, const string &prefix) {
-    // Define the bounds on the split dimensions using the bounds
-    // on the function args. If it is a purify, we should use the bounds
-    // from the dims instead.
-
+    // Define the bounds on the split dimensions using the bounds on the function args.
     vector<std::pair<string, Expr>> let_stmts;
 
     Expr old_var_extent = Variable::make(Int(32), prefix + split.old_var + ".loop_extent");

From 06548a9207307f262dfc12aa260693e2515c45cc Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Wed, 11 Dec 2024 15:40:29 -0500
Subject: [PATCH 32/39] Factor out weaken_condition_under_domain and document
 logic

---
 src/Func.cpp  | 2 +-
 src/Solve.cpp | 4 ++++
 src/Solve.h   | 7 +++++++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 532fed44481b..6c522c6a13e6 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -874,7 +874,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         for (const auto &[var, min, extent] : intermediate_rdom.domain()) {
             intm_rdom.push(var, Interval{min, min + extent - 1});
         }
-        preserved_rdom.set_predicate(simplify(!and_condition_over_domain(simplify(substitute(preserved_map, !preserved_rdom.predicate())), intm_rdom)));
+        preserved_rdom.set_predicate(weaken_condition_under_domain(substitute(preserved_map, preserved_rdom.predicate()), intm_rdom));
     }
 
     // Intermediate func
diff --git a/src/Solve.cpp b/src/Solve.cpp
index 10e6232e379d..e69228781095 100644
--- a/src/Solve.cpp
+++ b/src/Solve.cpp
@@ -1239,6 +1239,10 @@ Expr and_condition_over_domain(const Expr &e, const Scope<Interval> &varying) {
     return simplify(bounds.min);
 }
 
+Expr weaken_condition_under_domain(const Expr &c, const Scope<Interval> &varying) {
+    return simplify(!and_condition_over_domain(simplify(!c), varying));
+}
+
 // Testing code
 
 namespace {
diff --git a/src/Solve.h b/src/Solve.h
index ff5124e508c6..b591aafe1790 100644
--- a/src/Solve.h
+++ b/src/Solve.h
@@ -47,6 +47,13 @@ Interval solve_for_inner_interval(const Expr &c, const std::string &variable);
  * 'and' over the vector lanes, and return a scalar result. */
 Expr and_condition_over_domain(const Expr &c, const Scope<Interval> &varying);
 
+/** Take a conditional that includes variables that vary over some
+ * domain, and convert it to a weaker (less frequently false) condition
+ * that doesn't depend on those variables. Formally, the input expr
+ * implies the output expr. Note that this function might be unable to
+ * provide a better response than simply const_true(). */
+Expr weaken_condition_under_domain(const Expr &c, const Scope<Interval> &varying);
+
 void solve_test();
 
 }  // namespace Internal

From f6a3fac6f3eef3a4894effb786e5a53e45887835 Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Wed, 11 Dec 2024 15:02:58 -0800
Subject: [PATCH 33/39] trigger buildbots


From 22a28d1fd6c1e45a73f1a68bf3dc2b86f8e7fcc9 Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Thu, 12 Dec 2024 09:35:55 -0800
Subject: [PATCH 34/39] Remove unused var

---
 src/ConstantBounds.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/ConstantBounds.cpp b/src/ConstantBounds.cpp
index 11d1a42133a9..164678e2554c 100644
--- a/src/ConstantBounds.cpp
+++ b/src/ConstantBounds.cpp
@@ -73,7 +73,6 @@ ConstantInterval bounds_helper(const Expr &e,
             ScopedBinding bind(scope, op->name, recurse(op->value));
             return recurse(op->body);
         } else if (const Call *op = e.as<Call>()) {
-            ConstantInterval result;
             if (op->is_intrinsic(Call::abs)) {
                 return abs(recurse(op->args[0]));
             } else if (op->is_intrinsic(Call::absd)) {

From 8b87a09512e4dff78e2bf3f626b37a20dfd747b3 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Fri, 13 Dec 2024 02:27:56 -0500
Subject: [PATCH 35/39] Rename weaken_condition_under_domain ~>
 or_condition_over_domain

---
 src/Func.cpp  | 2 +-
 src/Solve.cpp | 2 +-
 src/Solve.h   | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 6c522c6a13e6..53076cd9704a 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -874,7 +874,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         for (const auto &[var, min, extent] : intermediate_rdom.domain()) {
             intm_rdom.push(var, Interval{min, min + extent - 1});
         }
-        preserved_rdom.set_predicate(weaken_condition_under_domain(substitute(preserved_map, preserved_rdom.predicate()), intm_rdom));
+        preserved_rdom.set_predicate(or_condition_over_domain(substitute(preserved_map, preserved_rdom.predicate()), intm_rdom));
     }
 
     // Intermediate func
diff --git a/src/Solve.cpp b/src/Solve.cpp
index e69228781095..0dd27f9be114 100644
--- a/src/Solve.cpp
+++ b/src/Solve.cpp
@@ -1239,7 +1239,7 @@ Expr and_condition_over_domain(const Expr &e, const Scope<Interval> &varying) {
     return simplify(bounds.min);
 }
 
-Expr weaken_condition_under_domain(const Expr &c, const Scope<Interval> &varying) {
+Expr or_condition_over_domain(const Expr &c, const Scope<Interval> &varying) {
     return simplify(!and_condition_over_domain(simplify(!c), varying));
 }
 
diff --git a/src/Solve.h b/src/Solve.h
index b591aafe1790..4d06fda47d6b 100644
--- a/src/Solve.h
+++ b/src/Solve.h
@@ -52,7 +52,7 @@ Expr and_condition_over_domain(const Expr &c, const Scope<Interval> &varying);
  * that doesn't depend on those variables. Formally, the input expr
  * implies the output expr. Note that this function might be unable to
  * provide a better response than simply const_true(). */
-Expr weaken_condition_under_domain(const Expr &c, const Scope<Interval> &varying);
+Expr or_condition_over_domain(const Expr &c, const Scope<Interval> &varying);
 
 void solve_test();
 

From faeede6c410033f17c6b5e5be7dbdd4d0e62c53b Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Fri, 13 Dec 2024 02:55:24 -0500
Subject: [PATCH 36/39] Adding comment on missing pure vars

---
 src/Func.cpp | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 53076cd9704a..6d6187db13be 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -969,7 +969,17 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
                 }
             }
 
-            // Add missing pure vars to the REDUCING func just before outermost
+            // Add missing pure vars to the REDUCING func just before outermost.
+            // This is necessary whenever the update does not reference one of the
+            // pure variables. For instance, factoring a histogram (clamps elided):
+            //     g(x) = 0; g(f(r.x, r.y)) += 1;
+            //     Func intm = g.rfactor(r.y, u);
+            // Here we generate an intermediate func intm that looks like:
+            //     intm(x, u) = 0; intm(f(r.x, u), u) += 1;
+            // And we need the reducing func to be:
+            //     g(x) += intm(x, r.y);
+            // But x was not referenced in the original update definition, so that
+            // dimension is added here.
             for (size_t i = 0; i < dim_vars.size(); i++) {
                 if (!expr_uses_var(definition.args()[i], dim_vars[i].name())) {
                     Dim d = {dim_vars[i].name(), ForType::Serial, DeviceAPI::None, DimType::PureVar, Partition::Auto};

From 2a3b0df6ad205ebffe6df7e0e8c03bae7bd61d30 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Fri, 13 Dec 2024 03:16:01 -0500
Subject: [PATCH 37/39] Add more comments inside project_rdom

---
 src/Func.cpp | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 6d6187db13be..7ea846490ed3 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -650,7 +650,11 @@ void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
 }
 
 pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const ReductionDomain &rdom, const vector<Split> &splits) {
-    // Compute new bounds for the RDom by walking backwards through the splits.
+    // The bounds projections maps expressions that reference the old RDom
+    // bounds to expressions that reference the new RDom bounds (from dims).
+    // We call this a projection because we are computing the symbolic image
+    // of the N-dimensional RDom (dimensionality including splits) in the
+    // M < N - dimensional result.
     SubstitutionMap bounds_projection{};
     for (const Split &split : reverse_view(splits)) {
         for (const auto &[name, value] : compute_loop_bounds_after_split(split, "")) {
@@ -663,7 +667,7 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
         add_let(bounds_projection, var + ".loop_extent", extent);
     }
 
-    // Build the new RDom
+    // Build the new RDom from the bounds_projection.
     vector<ReductionVariable> new_rvars;
     for (const Dim &dim : dims) {
         const Expr new_min = simplify(bounds_projection.at(dim.var + ".loop_min"));
@@ -673,12 +677,17 @@ pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, con
     ReductionDomain new_rdom{new_rvars};
     new_rdom.where(rdom.predicate());
 
-    // Compute the mapping of old dimensions to the projected RDom values.
+    // Compute a mapping from old dimensions to equivalent values using only
+    // the new dimensions. For example, if we have an RDom {{0, 20}} and we
+    // split r.x by 2 into r.xo and r.xi, then this map will contain:
+    //     r.x ~> 2 * r.xo + r.xi
+    // Certain split tail cases can place additional predicates on the RDom.
+    // These are handled here, too.
+    SubstitutionMap dim_projection{};
     SubstitutionMap dim_extent_alignment{};
     for (const auto &[var, _, extent] : rdom.domain()) {
         dim_extent_alignment[var] = extent;
     }
-    SubstitutionMap dim_projection{};
     for (const Split &split : splits) {
         for (const auto &result : apply_split(split, "", dim_extent_alignment)) {
             switch (result.type) {

From 57476b110274c237081eb285bec72c46fe0731e0 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Fri, 13 Dec 2024 03:36:02 -0500
Subject: [PATCH 38/39] Add more detailed documentation to add_let.

---
 src/Func.cpp | 42 ++++++++++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 7ea846490ed3..b6fc534b3dbc 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -635,18 +635,40 @@ optional<Dim> find_dim(const vector<Dim> &items, const VarOrRVar &v) {
 
 using SubstitutionMap = std::map<string, Expr>;
 
-// This is a helper function for building up a substitution map that
-// corresponds to pushing down a nest of lets. The lets should be fed
-// to this function from innermost to outermost. This is equivalent to
-// building a let-nest as a one-hole context and then simplifying.
-void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
-    internal_assert(!proj.count(name)) << "would shadow " << name << " in let nest.\n"
-                                       << "\tPresent value: " << proj[name] << "\n"
-                                       << "\tProposed value: " << value;
-    for (auto &[_, e] : proj) {
+/** This is a helper function for building up a substitution map that
+ * corresponds to pushing down a nest of lets. The lets should be fed
+ * to this function from innermost to outermost. This is equivalent to
+ * building a let-nest as a one-hole context and then simplifying.
+ *
+ * This looks like it might be quadratic or worse, and technically it is,
+ * but this isn't a problem for the way it is used inside rfactor. There
+ * are only a few uses:
+ *
+ *   1. Remapping preserved RVars to new RVars
+ *   2. Remapping factored RVars to new Vars
+ *   3. Filling the holes in the associative template
+ *   4. Accumulating the lets from ApplySplit
+ *
+ * These are naturally bounded by O(#splits + #dims) which is quite small
+ * in practice. Classes (1) and (2) cannot blow up expressions since they
+ * simply rename variables. Class (3) cannot blow up expressions either
+ * since nothing else can refer to the holes. That leaves only class (1).
+ * Fortunately, the lets generated by splits are benign. Split factors can't
+ * refer to RVars, and we won't see the consumed RVars in another split. So
+ * in total, this avoids any sort of exponentially sized substitution.
+ *
+ * @param subst The existing let nest (represented by a SubstitutionMap).
+ * @param name The name to bind, cannot already exist in the nest.
+ * @param value The value to bind. Will be substituted into nested values.
+ */
+void add_let(SubstitutionMap &subst, const string &name, const Expr &value) {
+    internal_assert(!subst.count(name)) << "would shadow " << name << " in let nest.\n"
+                                        << "\tPresent value: " << subst[name] << "\n"
+                                        << "\tProposed value: " << value;
+    for (auto &[_, e] : subst) {
         e = substitute(name, value, e);
     }
-    proj.emplace(name, value);
+    subst.emplace(name, value);
 }
 
 pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const ReductionDomain &rdom, const vector<Split> &splits) {

From 15a4dbfd14528e80ff677f5ce770eb2584dc8130 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Fri, 13 Dec 2024 04:24:59 -0500
Subject: [PATCH 39/39] Remove references to purify in comments.

---
 src/ScheduleFunctions.cpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp
index 0cf090b3d5a5..02f5553f748c 100644
--- a/src/ScheduleFunctions.cpp
+++ b/src/ScheduleFunctions.cpp
@@ -412,8 +412,7 @@ Stmt build_loop_nest(
     }
 
     // Define the bounds on the split dimensions using the bounds
-    // on the function args. If it is a purify, we should use the bounds
-    // from the dims instead.
+    // on the function args.
     for (const Split &split : reverse_view(splits)) {
         vector<std::pair<string, Expr>> let_stmts = compute_loop_bounds_after_split(split, prefix);
         for (const auto &let_stmt : let_stmts) {
@@ -2228,7 +2227,7 @@ bool validate_schedule(Function f, const Stmt &s, const Target &target, bool is_
         //
         // However, there are four types of Split, and the concept of a child var varies across them:
         // - For a vanilla split, inner and outer are the children and old_var is the parent.
-        // - For rename and purify, the outer is the child and the inner is meaningless.
+        // - For rename, the outer is the child and the inner is meaningless.
         // - For fuse, old_var is the child and inner/outer are the parents.
         //
         // (@abadams comments: "I acknowledge that this is gross and should be refactored.")
