Error out when calling rfactor() after fusing a pure var

src/ApplySplit.cpp

@@ -11,8 +11,7 @@ using std::map;
 using std::string;
 using std::vector;
 
-vector<ApplySplitResult> apply_split(const Split &split, bool is_update, const string &prefix,
-                                     map<string, Expr> &dim_extent_alignment) {
+vector<ApplySplitResult> apply_split(const Split &split, const string &prefix, map<string, Expr> &dim_extent_alignment) {
     vector<ApplySplitResult> result;
 
     Expr outer = Variable::make(Int(32), prefix + split.outer);

src/ApplySplit.h

@@ -78,9 +78,7 @@ struct ApplySplitResult {
  * the definition (in ascending order of application), and let stmts which
  * defined the values of variables referred by the predicates and substitutions
  * (ordered from innermost to outermost let). */
-std::vector<ApplySplitResult> apply_split(
-    const Split &split, bool is_update, const std::string &prefix,
-    std::map<std::string, Expr> &dim_extent_alignment);
+std::vector<ApplySplitResult> apply_split(const Split &split, const std::string &prefix, std::map<std::string, Expr> &dim_extent_alignment);
 
 /** Compute the loop bounds of the new dimensions resulting from applying the
  * split schedules using the loop bounds of the old dimensions. */

src/Func.cpp

@@ -646,7 +646,7 @@ bool apply_split(const Split &s, vector<ReductionVariable> &rvars,
 
         rvars.insert(it + 1, {s.outer, 0, simplify((old_extent - 1 + s.factor) / s.factor)});
 
-        vector<ApplySplitResult> splits_result = apply_split(s, true, "", dim_extent_alignment);
+        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);
 
@@ -681,7 +681,7 @@ bool apply_fuse(const Split &s, vector<ReductionVariable> &rvars,
         iter_outer->extent = extent;
         rvars.erase(iter_inner);
 
-        vector<ApplySplitResult> splits_result = apply_split(s, true, "", dim_extent_alignment);
+        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);
 
@@ -705,7 +705,7 @@ bool apply_purify(const Split &s, vector<ReductionVariable> &rvars,
                  << ", deleting it from the rvars list\n";
         rvars.erase(iter);
 
-        vector<ApplySplitResult> splits_result = apply_split(s, true, "", dim_extent_alignment);
+        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);
 
@@ -725,7 +725,7 @@ bool apply_rename(const Split &s, vector<ReductionVariable> &rvars,
         debug(4) << "  Renaming " << iter->var << " into " << s.outer << "\n";
         iter->var = s.outer;
 
-        vector<ApplySplitResult> splits_result = apply_split(s, true, "", dim_extent_alignment);
+        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);
 
@@ -788,6 +788,15 @@ Func Stage::rfactor(vector<pair<RVar, Var>> preserved) {
     vector<Expr> &args = definition.args();
     vector<Expr> &values = definition.values();
 
+    // 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);
+
+    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());
+
     // Figure out which pure vars were used in this update definition.
     std::set<string> pure_vars_used;
     internal_assert(args.size() == dim_vars.size());
@@ -799,15 +808,6 @@ Func Stage::rfactor(vector<pair<RVar, Var>> preserved) {
         }
     }
 
-    // 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);
-
-    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());
-
     vector<Split> &splits = definition.schedule().splits();
     vector<Dim> &dims = definition.schedule().dims();
     vector<ReductionVariable> &rvars = definition.schedule().rvars();
@@ -870,6 +870,10 @@ Func Stage::rfactor(vector<pair<RVar, Var>> preserved) {
         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)) {
+                user_assert(!s.is_fuse())
+                    << "In schedule for " << name()
+                    << ", can't perform rfactor() after fusing " << s.outer
+                    << " and " << s.inner << "\n";
                 temp.push_back(s);
             }
         }

src/ScheduleFunctions.cpp

@@ -220,7 +220,7 @@ Stmt build_loop_nest(
         user_assert(predicated_vars.count(split.old_var) == 0)
             << "Cannot split a loop variable resulting from a split using PredicateLoads or PredicateStores.";
 
-        vector<ApplySplitResult> splits_result = apply_split(split, is_update, prefix, dim_extent_alignment);
+        vector<ApplySplitResult> splits_result = apply_split(split, prefix, dim_extent_alignment);
 
         // To ensure we substitute all indices used in call or provide,
         // we need to substitute all lets in, so we correctly guard x in

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

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;
+}

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;
+}