Partition fix with rfactor, simplify and likely predicates. (#3444)

diff --git a/src/schedule/schedule_dataflow_rewrite.cc b/src/schedule/schedule_dataflow_rewrite.cc
index 976e880..2c77192 100644 (file)
--- a/src/schedule/schedule_dataflow_rewrite.cc
+++ b/src/schedule/schedule_dataflow_rewrite.cc
@@ -735,7 +735,7 @@ Array<Tensor> Schedule::rfactor(const Tensor& tensor,
   const Reduce* reduce = compute_op->body[idx].as<Reduce>();
   CHECK(reduce) << "Can only rfactor non-inline reductions";
   predicates.push_back(reduce->condition);
-  Expr predicate = arith::ComputeReduce<ir::And>(predicates, Expr());
+  Expr predicate = likely(simplify(arith::ComputeReduce<ir::And>(predicates, Expr())));
 
   std::unordered_map<const Variable*, Expr> vsub;