Add promise_clamped in rfactor (#8608)

Fixes #8600

Author: alexreinking

src/Func.cpp

@@ -905,18 +905,29 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         }
     }
 
+    ReductionDomain rdom{definition.schedule().rvars(), definition.predicate(), true};
+    SubstitutionMap rdom_promises;
+    for (int i = 0; i < rdom.domain().size(); i++) {
+        const auto &[var, min, extent] = rdom.domain()[i];
+        rdom_promises.emplace(var, promise_clamped(RVar(rdom, i), min, min + extent - 1));
+    }
+
     // Project the RDom into each side
     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, rdom, rvar_splits);
         for (size_t i = 0; i < preserved.size(); i++) {
             add_let(intermediate_map, preserved_rdims[i].var, preserved_vars[i]);
         }
-        intermediate_rdom.set_predicate(simplify(substitute(intermediate_map, intermediate_rdom.predicate())));
+
+        {
+            Expr pred = intermediate_rdom.predicate();
+            pred = substitute(rdom_promises, pred);
+            pred = substitute(intermediate_map, pred);
+            intermediate_rdom.set_predicate(simplify(pred));
+        }
 
         // Preserved
         std::tie(preserved_rdom, preserved_map) = project_rdom(preserved_rdims, rdom, rvar_splits);
@@ -926,7 +937,13 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
             const auto &[_, min, extent] = intermediate_rdom.domain()[i];
             intm_rdom.push(var, Interval{min, min + extent - 1});
         }
-        preserved_rdom.set_predicate(or_condition_over_domain(substitute(preserved_map, preserved_rdom.predicate()), intm_rdom));
+        {
+            Expr pred = preserved_rdom.predicate();
+            pred = substitute(rdom_promises, pred);
+            pred = substitute(preserved_map, pred);
+            pred = or_condition_over_domain(pred, intm_rdom);
+            preserved_rdom.set_predicate(pred);
+        }
     }
 
     // Intermediate func
@@ -943,10 +960,12 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     {
         vector<Expr> args = definition.args();
         args.insert(args.end(), preserved_vars.begin(), preserved_vars.end());
+        args = substitute(rdom_promises, args);
         args = substitute(intermediate_map, args);
 
         vector<Expr> values = definition.values();
         values = substitute_self_reference(values, function.name(), intm.function(), preserved_vars);
+        values = substitute(rdom_promises, values);
         values = substitute(intermediate_map, values);
         intm.function().define_update(args, values, intermediate_rdom);
 
@@ -1041,7 +1060,7 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
         }
 
         definition.args() = dim_vars_exprs;
-        definition.values() = substitute(preserved_map, prover_result.pattern.ops);
+        definition.values() = substitute(preserved_map, substitute(rdom_promises, prover_result.pattern.ops));
         definition.predicate() = preserved_rdom.predicate();
         definition.schedule().dims() = subst_dims(preserved_map, reducing_dims);
         definition.schedule().rvars() = preserved_rdom.domain();

test/correctness/rfactor.cpp

@@ -1031,6 +1031,35 @@ int inlined_rfactor_with_disappearing_rvar_test() {
     return 0;
 }
 
+// From issue: https://github.com/halide/Halide/issues/8600
+int rfactor_precise_bounds_test() {
+    Var x("x"), y("y");
+    RDom r(0, 10, 0, 10);
+
+    // Create an input with random values
+    Buffer<uint8_t> input(10, 10, "input");
+    for (int y = 0; y < 10; ++y) {
+        for (int x = 0; x < 10; ++x) {
+            input(x, y) = (rand() % 256);
+        }
+    }
+
+    Func f;
+
+    f() = 0;
+    f() += input(r.x, r.y);
+    RVar rxo, rxi, ryo, ryi;
+    Func intm = f.update()
+                    .tile(r.x, r.y, rxo, ryo, rxi, ryi, 4, 4)
+                    .rfactor({{rxi, x}, {ryi, y}});
+
+    intm.compute_root();
+
+    Buffer<int> im = f.realize();
+
+    return 0;
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {
@@ -1063,15 +1092,14 @@ int main(int argc, char **argv) {
         {"tuple rfactor test: checking output img correctness...", tuple_rfactor_test<false>},
         {"tuple specialize rdom predicate rfactor test: checking call graphs...", tuple_specialize_rdom_predicate_rfactor_test<true>},
         {"tuple specialize rdom predicate rfactor test: checking output img correctness...", tuple_specialize_rdom_predicate_rfactor_test<false>},
-        {"parallel dot product rfactor test: checking call graphs...", parallel_dot_product_rfactor_test<true>},
-        {"parallel dot product rfactor test: checking output img correctness...", parallel_dot_product_rfactor_test<false>},
         {"tuple partial reduction rfactor test: checking call graphs...", tuple_partial_reduction_rfactor_test<true>},
         {"tuple partial reduction rfactor test: checking output img correctness...", tuple_partial_reduction_rfactor_test<false>},
         {"check allocation bound test", check_allocation_bound_test},
         {"rfactor tile reorder test: checking output img correctness...", rfactor_tile_reorder_test},
         {"complex multiply rfactor test", complex_multiply_rfactor_test},
         {"argmin rfactor test", argmin_rfactor_test},
         {"inlined rfactor with disappearing rvar test", inlined_rfactor_with_disappearing_rvar_test},
+        {"rfactor bounds tests", rfactor_precise_bounds_test},
     };
 
     using Sharder = Halide::Internal::Test::Sharder;