Upgrade pir op lowerer into new group scheduler strategy (#59810)

BiynXu · Aurelius84 · web-flow · commit c9ecd0707c5e · 2023-12-09T17:46:43.000+08:00
* [PIR+CINN]Upgrade into New GroupScheduler Strategy

upload debug code

fix UT and switch C++ into new strategy

* fix conflict

* fix do_op_schedule

* remove code

* fix env

* fix UT

* [CINN] fix group schedule of keep dim reduction

* fix group schedule name bug

---------

Co-authored-by: Aurelius84 &lt;zhangliujie@baidu.com&gt;
diff --git a/paddle/cinn/ast_gen_ius/ast_gen.cc b/paddle/cinn/ast_gen_ius/ast_gen.cc
@@ -21,6 +21,8 @@
 #include "paddle/cinn/lang/compute.h"
 #include "paddle/cinn/optim/replace_var_with_expr.h"
 
+PD_DECLARE_bool(cinn_new_group_scheduler);
+
 namespace cinn {
 namespace ast_gen_ius {
 
@@ -92,11 +94,15 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     // for same axis so we re-create objects
     std::vector<Var> axis_vars = cinn::common::GenDefaultAxis(axis_len);
     for (int i = 0; i < shape.size(); ++i) {
+      if (FLAGS_cinn_new_group_scheduler && shape[i] == Expr(1)) {
+        optim::ReplaceVarWithExpr(&init_body, axis[i], Expr(0));
+        continue;
+      }
       block_vars.push_back(Var(Expr(0),
                                shape[i],
                                cinn::UniqName("i" + std::to_string(i)),
                                /*is_reduce = */ false));
-      optim::ReplaceVarWithExpr(&init_body, axis[i], block_vars[i]);
+      optim::ReplaceVarWithExpr(&init_body, axis[i], block_vars.back());
       axis_vars[i]->is_reduce_axis = false;
       if (shape[i] == Expr(1)) {
         iter_values.push_back(Expr(0));
@@ -120,6 +126,10 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     // for same axis so we re-create objects
     std::vector<Var> reduce_axis_vars = cinn::common::GenDefaultAxis(axis_len);
     for (int i = 0; i < shape.size(); ++i) {
+      if (FLAGS_cinn_new_group_scheduler && shape[i] == Expr(1)) {
+        optim::ReplaceVarWithExpr(&reduce_body, axis[i], Expr(0));
+        continue;
+      }
       reduce_block_vars.push_back(Var(Expr(0),
                                       shape[i],
                                       cinn::UniqName("i" + std::to_string(i)),
@@ -142,12 +152,20 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
       reduce_axis_var->is_reduce_axis = true;
       reduce_iter_values.push_back(reduce_axis_var);
     }
+
+    int non_zero_axis_size = 0;
     for (int i = 0; i < axis.size(); ++i) {
-      optim::ReplaceVarWithExpr(&reduce_body, axis[i], reduce_block_vars[i]);
-    }
-    for (int i = axis.size(); i < reduce_block_vars.size(); ++i) {
+      if (FLAGS_cinn_new_group_scheduler && shape[i] == Expr(1)) {
+        continue;
+      }
       optim::ReplaceVarWithExpr(
-          &reduce_body, reduce_axis[i - axis.size()], reduce_block_vars[i]);
+          &reduce_body, axis[i], reduce_block_vars[non_zero_axis_size]);
+      ++non_zero_axis_size;
+    }
+    for (int i = non_zero_axis_size; i < reduce_block_vars.size(); ++i) {
+      optim::ReplaceVarWithExpr(&reduce_body,
+                                reduce_axis[i - non_zero_axis_size],
+                                reduce_block_vars[i]);
     }
 
     reduce_body = ir::ScheduleBlockRealize::Make(
@@ -166,6 +184,9 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     // Put the two parts together
     ir::Expr body = ir::Block::Make({init_body, reduce_body});
     for (int i = static_cast<int>(axis_len) - 1; i >= 0; --i) {
+      if (shape[i] == Expr(1)) {
+        continue;
+      }
       ir::Var loop_var = axis[i];
       ir::Expr loop_extent = shape[i];
       body = ir::For::Make(
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -36,6 +36,7 @@ PD_DECLARE_bool(cinn_use_cuda_vectorize);
 PD_DECLARE_bool(cinn_enable_map_expr);
 PD_DECLARE_bool(cinn_enable_map_expr_schedule);
 PD_DECLARE_bool(cinn_bucket_compile);
+PD_DECLARE_bool(cinn_new_group_scheduler);
 
 namespace cinn {
 namespace hlir {
@@ -659,6 +660,20 @@ ir::Expr OpLowererImpl::DoGroupSchedule(
     const GroupPtr& group,
     const std::unordered_map<::pir::Value, ir::Tensor>& tensor_map,
     const std::unordered_map<std::string, ir::Tensor>& tmp_tensor_info) {
+  if (FLAGS_cinn_new_group_scheduler) {
+    VLOG(3) << "using StaticShapeGroupScheduler to schedule group.";
+    std::unordered_set<std::string> output_tensor_names;
+    std::transform(
+        group->output_ops.begin(),
+        group->output_ops.end(),
+        std::inserter(output_tensor_names, output_tensor_names.begin()),
+        [&](::pir::Operation* op) { return ValueName(op->result(0)); });
+    std::unique_ptr<ir::GroupScheduler> group_scheduler =
+        ir::GroupScheduler::Make(
+            &ir_sch, output_tensor_names, target_, /* is_dy_shape = */ false);
+    group_scheduler->Schedule();
+    return ir_sch.GetModule().GetExprs().at(0);
+  }
   // topological order.
   auto ops_set = group->OpSet();
   auto v_consumers = BuildVirtualConsumer(group);
diff --git a/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.cc b/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.cc
@@ -256,16 +256,19 @@ void StaticShapeGroupScheduler::DoLoopAlignment() {
     std::vector<ir::Expr> indices =
         reduce_loads.begin()->As<ir::Load>()->indices;
     for (ir::Expr index : indices) {
+      if (index.is_constant()) continue;
       CHECK_NOTNULL(index.as_var());
       int idx = 0;
       bool is_reduce_var = false;
-      for (const ir::Var& iter_var : master_iter_vars) {
+      for (int iter_idx = 0; iter_idx < master_iter_vars.size(); ++iter_idx) {
+        auto& iter_var = master_iter_vars[iter_idx];
         if (iter_var->name == index.as_var_ref()->name) {
           is_reduce_var = iter_var->is_reduce_axis;
           break;
         }
         ++idx;
       }
+      if (master_iter_values[idx].is_constant()) continue;
       std::vector<ir::Var> loop_vars_in_order;
       ir::ir_utils::CollectIRNodesInOrder(
           master_iter_values[idx], [&](const ir::Expr* x) {
diff --git a/paddle/cinn/ir/schedule/factorize_reduction.h b/paddle/cinn/ir/schedule/factorize_reduction.h
@@ -33,7 +33,7 @@ namespace ir {
 Tensor CreateRFTensor(const Tensor& original_tensor,
                       const Expr& rf_loop,
                       int rf_axis) {
-  std::string name = cinn::common::UniqName(original_tensor->name + "_rf");
+  std::string name = original_tensor->name + "_rf";
   std::vector<Expr> new_shape = original_tensor->shape;
   new_shape.insert(new_shape.begin() + rf_axis, rf_loop.As<For>()->extent);
   Tensor rf_tensor = _Tensor_::Make(name,
diff --git a/paddle/cinn/ir/schedule/impl/compute_location.cc b/paddle/cinn/ir/schedule/impl/compute_location.cc
@@ -135,6 +135,8 @@ void StScheduleImpl::SimpleComputeAt(const Expr& block, const Expr& loop) {
   std::vector<Var> replaced_var;
   std::vector<Expr> substitute_expr;
   for (int i = 0; i < loops.size(); ++i) {
+    VLOG(3) << i << "-th loop is:\n " << loops[i];
+    VLOG(3) << i << "-th block_loop:\n" << block_loops[i];
     CHECK_EQ(GetLoopExtent(loops[i]), GetLoopExtent(block_loops[i]));
     if (block_loops[i].As<ir::For>()->bind_info().valid() &&
         !loops[i].As<ir::For>()->bind_info().valid()) {
diff --git a/paddle/cinn/ir/schedule/ir_schedule_util.h b/paddle/cinn/ir/schedule/ir_schedule_util.h
@@ -1061,7 +1061,9 @@ struct FindBlocksVisitor {
     if (expr->As<ir::For>()) {
       Visit(&(expr->As<ir::For>()->body));
     } else if (expr->As<ir::ScheduleBlockRealize>()) {
-      if (!expr->As<ir::ScheduleBlockRealize>()->iter_values.empty()) {
+      if (expr->As<ir::ScheduleBlockRealize>()
+              ->schedule_block.As<ScheduleBlock>()
+              ->name.substr(0, 4) != "root") {
         auto* schedule_block = expr->As<ir::ScheduleBlockRealize>()
                                    ->schedule_block.As<ir::ScheduleBlock>();
         if (block_name_.empty() || schedule_block->name == block_name_) {
diff --git a/test/cpp/pir/cinn/CMakeLists.txt b/test/cpp/pir/cinn/CMakeLists.txt
@@ -3,14 +3,12 @@ add_subdirectory(adt)
 if(WITH_TESTING AND WITH_CINN)
   paddle_test(test_pir_compiler SRCS pir_compiler_test.cc DEPS pir_compiler
               cinn_runtime_dialect)
-  set_tests_properties(test_pir_compiler PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(test_jit_instruction SRCS jit_instruction_test.cc DEPS
               cinn_runtime_dialect pir_compiler)
-  set_tests_properties(test_jit_instruction PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(
-    dialect_convert_test
+    test_dialect_convert
     SRCS
     dialect_convert_test.cc
     DEPS
@@ -19,10 +17,9 @@ if(WITH_TESTING AND WITH_CINN)
     op_dialect_vjp
     cinn_op_dialect
     pir)
-  set_tests_properties(dialect_convert_test PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(
-    add_broadcast_to_elementwise_test
+    test_add_broadcast_to_elementwise
     SRCS
     add_broadcast_to_elementwise_test.cc
     DEPS
@@ -32,7 +29,6 @@ if(WITH_TESTING AND WITH_CINN)
     cinn_op_dialect
     add_broadcast_to_elementwise_pass
     pir)
-  set_tests_properties(dialect_convert_test PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(
     test_sub_graph_extract
@@ -45,17 +41,15 @@ if(WITH_TESTING AND WITH_CINN)
     op_dialect_vjp
     pir_transforms
     pir)
-  set_tests_properties(test_sub_graph_extract PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(
-    ir_op_fusion_test
+    test_ir_op_fusion
     SRCS
     ir_op_fusion_test.cc
     DEPS
     op_with_group_merge_pass
     cinn_op_dialect
     pir)
-  set_tests_properties(ir_op_fusion_test PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(
     test_pir_all_path
@@ -67,7 +61,6 @@ if(WITH_TESTING AND WITH_CINN)
     cinn_op_dialect
     pd_to_cinn_pass
     add_broadcast_to_elementwise_pass)
-  set_tests_properties(test_pir_all_path PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(
     test_group_op
@@ -79,13 +72,34 @@ if(WITH_TESTING AND WITH_CINN)
     op_with_group_merge_pass
     cinn_op_dialect
     pir_transforms)
-  set_tests_properties(test_group_op PROPERTIES LABELS "RUN_TYPE=CINN")
 
   paddle_test(test_pir_build_cinn_pass SRCS build_cinn_pass_test.cc DEPS
               pir_transforms pir)
-  set_tests_properties(test_pir_build_cinn_pass PROPERTIES LABELS
-                                                           "RUN_TYPE=CINN")
 
   paddle_test(test_compilation_task SRCS compilation_task_test.cc DEPS pir)
-  set_tests_properties(test_compilation_task PROPERTIES LABELS "RUN_TYPE=CINN")
+
+  # DO NOT forget add test name here, otherwise it will not be executed in
+  # CINN CI.
+  set(cinn_unit_tests
+      test_pir_compiler
+      test_jit_instruction
+      test_dialect_convert
+      test_add_broadcast_to_elementwise
+      test_sub_graph_extract
+      test_ir_op_fusion
+      test_pir_all_path
+      test_group_op
+      test_pir_build_cinn_pass
+      test_compilation_task)
+
+  foreach(test_name ${cinn_unit_tests})
+    get_property(
+      env
+      TEST ${test_name}
+      PROPERTY ENVIRONMENT)
+    set_property(TEST ${test_name}
+                 PROPERTY ENVIRONMENT "FLAGS_cinn_new_group_scheduler=1" ${env})
+    set_tests_properties(${test_name} PROPERTIES LABELS "RUN_TYPE=CINN")
+  endforeach()
+
 endif()
diff --git a/test/cpp/pir/cinn/add_broadcast_to_elementwise_test.cc b/test/cpp/pir/cinn/add_broadcast_to_elementwise_test.cc
@@ -98,7 +98,7 @@ TEST(PatternRewrite, broadcast_elementwise) {
   it++;
   CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
-  CHECK_EQ(it->isa<cinn::dialect::BroadcastOp>(), true);
+  CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
   CHECK_EQ(it->isa<paddle::dialect::AddOp>(), true);
 }
@@ -124,9 +124,9 @@ TEST(PatternRewrite, broadcast_elementwise_both) {
   it++;
   CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
-  CHECK_EQ(it->isa<cinn::dialect::BroadcastOp>(), true);
+  CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
-  CHECK_EQ(it->isa<cinn::dialect::BroadcastOp>(), true);
+  CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
   CHECK_EQ(it->isa<paddle::dialect::AddOp>(), true);
 }
@@ -152,9 +152,9 @@ TEST(PatternRewrite, broadcast_elementwise_sub_both) {
   it++;
   CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
-  CHECK_EQ(it->isa<cinn::dialect::BroadcastOp>(), true);
+  CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
-  CHECK_EQ(it->isa<cinn::dialect::BroadcastOp>(), true);
+  CHECK_EQ(it->isa<paddle::dialect::FullOp>(), true);
   it++;
   CHECK_EQ(it->isa<paddle::dialect::SubtractOp>(), true);
 }
diff --git a/test/cpp/pir/cinn/pir_compiler_test.cc b/test/cpp/pir/cinn/pir_compiler_test.cc
@@ -207,7 +207,6 @@ TEST(PirCompier, CompileSoftmax) {
   executor.Run({}, true);
   auto out_tensor =
       executor.local_scope()->FindVar("out@fetch")->Get<phi::DenseTensor>();
-
   bool res0 = simple_cmp(out_tensor.data<float>()[0], 1.0 / 16);
   EXPECT_EQ(res0, true);
 }

Original file line number	Diff line number	Diff line change
`@@ -207,7 +207,6 @@ TEST(PirCompier, CompileSoftmax) {`
`207`	`207`	`executor.Run({}, true);`
`208`	`208`	`auto out_tensor =`
`209`	`209`	`executor.local_scope()->FindVar("out@fetch")->Get<phi::DenseTensor>();`
`210`		`-`
`211`	`210`	`bool res0 = simple_cmp(out_tensor.data<float>()[0], 1.0 / 16);`
`212`	`211`	`EXPECT_EQ(res0, true);`
`213`	`212`	`}`