[CINN] comb the op lowering code (#54982)

* [CINN] comb the op lowering code

* [CINN] format code of OpLower

paddle/cinn/auto_schedule/search_space/auto_gen_rule/auto_inline_test.cc

@@ -71,7 +71,6 @@ TEST(AutoInline, SingleLoopInline) {
                      nullptr,
                      target,
                      true);
-
   VLOG(6) << "Expr after lowering:";
   VLOG(6) << funcs[0]->body;
 
@@ -170,7 +169,9 @@ TEST(AutoInline, AddReluInline) {
 
   EXPECT_EQ(graph->fusion_groups.size(), 1UL);
   std::vector<ir::LoweredFunc> funcs =
-      op_lowerer->LowerWithoutSchedule(graph->fusion_groups[0]);
+      op_lowerer->Lower(graph->fusion_groups[0],
+                        /*apply_op_schedule = */ false,
+                        /*apply_group_schedule=*/false);
 
   VLOG(6) << "Expr before auto inline: " << funcs[0]->body;
 

paddle/cinn/auto_schedule/search_space/auto_gen_rule/multi_level_tiling_test.cc

@@ -388,9 +388,9 @@ TEST_F(TestMultiLevelTiling, ReduceSum) {
 
 TEST_F(TestMultiLevelTiling, Pool2d) {
   default_input_names = {"input"};
-  default_output_names = {"var_0"};
-  std::vector<int32_t> input_shape{2, 8, 16, 16};
-  std::vector<int32_t> output_shape{2, 8, 8, 8};
+  default_output_names = {"var_0", "pad_temp_0"};
+  std::vector<std::vector<int32_t>> input_shapes{{2, 8, 16, 16}};
+  std::vector<std::vector<int32_t>> output_shapes{{2, 8, 8, 8}, {2, 8, 18, 18}};
   std::string pooling_type = "max";
   std::vector<int> ksize{3, 3};
   std::vector<int> strides{2, 2};
@@ -402,7 +402,7 @@ TEST_F(TestMultiLevelTiling, Pool2d) {
   bool adaptive = false;
   std::string padding_algorithm = "EXPLICIT";
   frontend::Program pool2d_program = tests::OpBuilder("pool2d").Build(
-      {{"input", input_shape}},
+      {{"input", input_shapes[0]}},
       {{"pool_type", pooling_type},
        {"kernel_size", ksize},
        {"stride_size", strides},
@@ -439,6 +439,7 @@ TEST_F(TestMultiLevelTiling, Pool2d) {
   std::string expected_ir = R"ROC(Expr 0 {
 {
   ScheduleBlock(root)
+  {
     {
       serial for (i, 0, 2)
       {
@@ -451,6 +452,7 @@ TEST_F(TestMultiLevelTiling, Pool2d) {
               ScheduleBlock(pad_temp_0)
               {
                 i0, i1, i2, i3 = axis.bind(i, j, k, a)
+                {
                   pad_temp_0[i, j, k, a] = select(((a < 17) and ((a >= 1) and ((k < 17) and (k >= 1)))), input[i, j, (-1 + k), (-1 + a)], -3.40282347e+38f)
                 }
               }
@@ -458,12 +460,6 @@ TEST_F(TestMultiLevelTiling, Pool2d) {
           }
         }
       }
-}
-}  // end Expr 0
-Expr 1 {
-{
-  ScheduleBlock(root_0)
-  {
       {
         thread_bind[blockIdx.x] for (i_j_k_a_fused, 0, 16)
         {
@@ -552,8 +548,9 @@ Expr 1 {
         }
       }
     }
+  }
 }
-}  // end Expr 1
+}  // end Expr 0
 )ROC";
   ASSERT_EQ(ir, expected_ir);
 
@@ -569,8 +566,8 @@ Expr 1 {
           pool2d_program, fixed_rand_seed, /* apply_manual_schedule*/ true))),
       default_input_names,
       default_output_names,
-      {input_shape},
-      {output_shape},
+      input_shapes,
+      output_shapes,
       target_);
 }
 

paddle/cinn/auto_schedule/search_space/auto_gen_rule/test_helper.cc

@@ -63,12 +63,10 @@ ir::IRSchedule TestAutoGenRuleBase::MakeIRSchedule(
       absl::flat_hash_map<std::string, hlir::framework::shape_t>>("infershape");
   hlir::framework::OpLowerer op_lowerer(dtype_dict, shape_dict, target_);
 
-  if (apply_manual_schedule) {
-    lowered_funcs_ = op_lowerer.Lower(graph->fusion_groups.front());
-  } else {
   lowered_funcs_ =
-        op_lowerer.LowerWithoutSchedule(graph->fusion_groups.front());
-  }
+      op_lowerer.Lower(graph->fusion_groups.front(),
+                       /*apply_op_schedule = */ apply_manual_schedule,
+                       /*apply_group_schedule = */ apply_manual_schedule);
   CHECK(!lowered_funcs_.empty()) << "lowered_funcs_ is empty";
 
   std::vector<Expr> bodys;

paddle/cinn/auto_schedule/task/tune_task.cc

@@ -39,7 +39,8 @@ void TuneTask::Initialize(
   op_lowerer = lower_handler;
 
   // Set lowered_funcs and analyze output names.
-  this->lowered_funcs = op_lowerer->LowerWithoutSchedule(subgraph);
+  this->lowered_funcs = op_lowerer->Lower(
+      subgraph, /*apply_op_schedule = */ false, /*apply_group_schedule=*/false);
   this->output_names = GetOutputNamesFromLoweredFunc(this->lowered_funcs);
   this->serialized_key = SerializeToString(shape_dict, dtype_dict);
 }

paddle/cinn/auto_schedule/tests/performance_comparison_test.cc

@@ -157,7 +157,9 @@ class PerformanceTester : public ::testing::Test {
 
     for (auto group : graph->fusion_groups) {
       compile_options.lowered_funcs.push_back(
-          op_lowerer->LowerWithoutSchedule(group));
+          op_lowerer->Lower(group,
+                            /*apply_op_schedule = */ false,
+                            /*apply_group_schedule=*/false));
     }
 
     VLOG(3) << "===========================No Schedule LoweredFunc "

paddle/cinn/hlir/framework/op_lowering.h

@@ -39,46 +39,132 @@ using GroupPtr = std::shared_ptr<Graph::Group>;
 using common::Target;
 
 class OpLowerer;
-typedef std::vector<Expr> (OpLowerer::*IRComputeFunction)(
-    poly::StageMap&,
-    std::vector<ir::Tensor>&,
-    std::unordered_map<std::string, ir::Tensor>&,
-    const GroupPtr&,
-    const GroupPtr&,
-    bool);
+
+typedef bool (OpLowerer::*ScheduleDetermineFunction)(Node*);
 
 class OpLowerer {
  public:
   OpLowerer(const absl::flat_hash_map<std::string, Type>&,
             const absl::flat_hash_map<std::string, shape_t>&,
             const Target&);
-  std::vector<ir::LoweredFunc> Lower(GroupPtr& group);                 // NOLINT
-  std::vector<ir::LoweredFunc> LowerWithoutSchedule(GroupPtr& group);  // NOLINT
+
+  /**
+   * @brief Lower a group to CINN IR.
+   * @param group The group to be lowered.
+   * @param apply_op_schedule Whether to schedule at Op level.
+   * @param apply_group_schedule Whether to schedule at group level.
+   * @return The lowered funcs.
+   */
+  std::vector<ir::LoweredFunc> Lower(const GroupPtr& group,
+                                     bool apply_op_schedule = true,
+                                     bool apply_group_schedule = true);
 
  private:
-  std::vector<ir::LoweredFunc> IRLowerOp(IRComputeFunction, GroupPtr&);
-  std::vector<ir::LoweredFunc> IRLowerNonFusibleOp(GroupPtr&, bool);
-  std::vector<ir::LoweredFunc> IRLowerOpWithoutSchedule(IRComputeFunction,
-                                                        GroupPtr&);
-#define DEFINE_IR_COMPUTE(type)                                \
-  std::vector<Expr> IR##type##Compute(                         \
-      poly::StageMap& stages,                                  \
-      std::vector<ir::Tensor>& func_args,                      \
-      std::unordered_map<std::string, ir::Tensor>& tensor_map, \
-      const GroupPtr& group,                                   \
-      const GroupPtr& sub_group,                               \
-      bool apply_impl_schedule = false);
-
-  // compute and schedule
-  DEFINE_IR_COMPUTE(Elementwise);
-  DEFINE_IR_COMPUTE(Reduce);
-  DEFINE_IR_COMPUTE(OutEWiseFusable);
-
-  void IRSchedule(
+  /**
+   * @brief Lower a group to CINN IR.
+   * @param group The group to be lowered.
+   * @param apply_op_schedule Whether to schedule at Op level.
+   * @param apply_group_schedule Whether to schedule at group level.
+   * @param schedule_determine_func Function used to determine which Ops to
+   * schedule.
+   * @return The lowered funcs.
+   */
+  std::vector<ir::LoweredFunc> LowerGroup(
+      const GroupPtr& group,
+      bool apply_op_schedule,
+      bool apply_group_schedule,
+      ScheduleDetermineFunction schedule_determine_func);
+
+  /**
+   * @brief Lower a group composed of CustomCall Op.
+   * @param group The group to be lowered.
+   * @return The lowered funcs.
+   */
+  std::vector<ir::LoweredFunc> LowerCustomCall(const GroupPtr& group);
+
+  /**
+   * @brief Post processing, including preparing function args and temporary
+   * variables, applying low-level optimization passes, etc.
+   * @param group The group to be lowered.
+   * @param tensor_map All tensors used for calculating the group.
+   * @param done_op_schedule Mark whether the Op level schedule has been
+   * applied.
+   * @param ir_sch The IRSchedule object of group.
+   * @param group_func_arg_tensors Tensors used as the group function arguments.
+   * @return The lowered funcs after the post processing.
+   */
+  std::vector<ir::LoweredFunc> PostProcess(
+      const GroupPtr& group,
+      const std::unordered_map<std::string, ir::Tensor>& tensor_map,
+      bool done_op_schedule,
+      ir::IRSchedule* ir_sch,
+      std::vector<ir::Tensor>* group_func_arg_tensors);
+
+  /**
+   * @brief Lower an Op set to CINN IR.
+   * Compute, Lower and optional Schedule will be performed one by one
+   * for each Op.
+   * @param nodes The Op nodes to be lowered.
+   * @param apply_op_schedule Whether to schedule at Op level.
+   * @param schedule_determine_func Function used to determine which Ops to
+   * schedule.
+   * @param group_func_arg_tensors Tensors used as the group function arguments.
+   * @param tensor_map All tensors used for calculating the group.
+   * @return The lowered func bodies of Op set.
+   */
+  std::vector<ir::Expr> LowerOps(
+      const std::vector<Node*>& nodes,
+      bool apply_op_schedule,
+      ScheduleDetermineFunction schedule_determine_func,
+      std::vector<ir::Tensor>* group_func_arg_tensors,
+      std::unordered_map<std::string, ir::Tensor>* tensor_map);
+
+  /**
+   * @brief Lower an Op to CINN IR. The Compute and Lower processes will be
+   * called sequentially.
+   * @param op_impl The Op implementation defining Compute and Schedule.
+   * @param node The Op node to be lowered.
+   * @param tensor_map All tensors used for calculating the group.
+   * @param op_func_arg_tensors Tensors used as the Op function arguments.
+   * @return The lowered func of the Op node.
+   */
+  std::vector<ir::LoweredFunc> DoOpLower(
+      std::shared_ptr<hlir::framework::OpImpl> op_impl,
+      Node* node,
+      std::unordered_map<std::string, ir::Tensor>* tensor_map,
+      std::vector<ir::Tensor>* op_func_arg_tensors);
+
+  /**
+   * @brief Apply schedule on an Op.
+   * @param op_impl The Op implementation defining Compute and Schedule.
+   * @param op_func_arg_tensors Tensors used as the Op function arguments.
+   * @param lowered_funcs The lowered funcs of an Op to be scheduled.
+   * @return The lowered func body after schedule of the Op.
+   */
+  ir::Expr DoOpSchedule(std::shared_ptr<hlir::framework::OpImpl> op_impl,
+                        const std::vector<ir::Tensor>& op_func_arg_tensors,
+                        const std::vector<ir::LoweredFunc>& lowered_funcs);
+
+  /**
+   * @brief Apply schedule on a group.
+   * @param ir_sch The IRSchedule containing the entire group's lowered func
+   * bodies.
+   * @param group The group to be scheduled.
+   * @param tensor_map All tensors used for calculating the group.
+   * @return The lowered func body after schedule of the group.
+   */
+  ir::Expr DoGroupSchedule(
       ir::IRSchedule& ir_sch,  // NOLINT
       const GroupPtr& group,
       const std::unordered_map<std::string, ir::Tensor>& tensor_map);
 
+  // Functions used to determine which Ops to schedule at op level, define a
+  // policy for each type of group.
+  inline bool ReduceScheduleDetermineFunction(Node* node);
+  inline bool ElementwiseScheduleDetermineFunction(Node* node);
+  inline bool NonFusibleScheduleDetermineFunction(Node* node);
+
+ private:
   Target target_;
   const absl::flat_hash_map<std::string, Type>& type_dict_;
   const absl::flat_hash_map<std::string, shape_t>& shape_dict_;

paddle/cinn/hlir/framework/op_lowering_util.cc

@@ -92,19 +92,19 @@ ir::Tensor GetTensor(
 
 std::vector<ir::Tensor> CollectInputTensor(
     const Node* node,
-    std::vector<ir::Tensor>& func_args,                       // NOLINT
-    std::unordered_map<std::string, ir::Tensor>& tensor_map,  // NOLINT
     const absl::flat_hash_map<std::string, Type>& type_dict,
-    const absl::flat_hash_map<std::string, shape_t>& shape_dict) {
+    const absl::flat_hash_map<std::string, shape_t>& shape_dict,
+    std::vector<ir::Tensor>* func_args,
+    std::unordered_map<std::string, ir::Tensor>* tensor_map) {
   std::vector<ir::Tensor> tensors;
   // get all input nodes
   for (auto& node_data : GetInputNodeData(node)) {
     CHECK(node_data);
     auto tensor = GetTensor(node_data, type_dict, shape_dict);
-    if (!tensor_map.count(node_data->id())) {
-      tensor_map[node_data->id()] = tensor;
+    if (!tensor_map->count(node_data->id())) {
+      (*tensor_map)[node_data->id()] = tensor;
       // record func input args
-      func_args.push_back(tensor);
+      func_args->push_back(tensor);
     }
     tensors.push_back(tensor);
   }

paddle/cinn/hlir/framework/op_lowering_util.h

@@ -31,10 +31,10 @@ ir::Tensor GetTensor(
 
 std::vector<ir::Tensor> CollectInputTensor(
     const Node* node,
-    std::vector<ir::Tensor>& func_args,                       // NOLINT
-    std::unordered_map<std::string, ir::Tensor>& tensor_map,  // NOLINT
     const absl::flat_hash_map<std::string, Type>& type_dict,
-    const absl::flat_hash_map<std::string, shape_t>& shape_dict);
+    const absl::flat_hash_map<std::string, shape_t>& shape_dict,
+    std::vector<ir::Tensor>* func_args,
+    std::unordered_map<std::string, ir::Tensor>* tensor_map);
 
 std::unordered_map<Node*, Node*> BuildVirtualConsumer(
     const GroupPtr& group,