Prohibit group schedule on single op

danleifeng · Oct 23, 2023 · 13ddff9 · 13ddff9
1 parent bc64b57
commit 13ddff9
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/paddle/cinn/hlir/framework/op_lowering_impl.cc b/paddle/cinn/hlir/framework/op_lowering_impl.cc
@@ -127,9 +127,14 @@ std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
   ir_sch.MergeExprs();
   VLOG(3) << "After lower, ir is: \n" << ir_sch.GetModule().GetExprs().at(0);
   auto& op_pattern_dict = Operator::GetAttrs<OpPatternKind>("OpPattern");
-  if (apply_group_schedule &&
+  bool do_group_schedule =
+      apply_group_schedule &&
       !(nodes.size() == 1 &&
-        op_pattern_dict[nodes[0]->op()] == OpPatternKind::kNonFusible)) {
+        op_pattern_dict[nodes[0]->op()] == OpPatternKind::kNonFusible);
+  if (FLAGS_cinn_new_group_scheduler) {
+    do_group_schedule = apply_group_schedule && nodes.size() > 1;
+  }
+  if (do_group_schedule) {
     DoGroupSchedule(ir_sch, group, tensor_map);
     VLOG(3) << "After group schedule, ir is: \n"
             << ir_sch.GetModule().GetExprs().at(0);