[DRR] Fix drr rewrite order when result pattern op be inserted (#64784)

* fix drr rewrite insert op order

* fix

Author: yuanlehome

paddle/fluid/pir/drr/src/rewrite_pattern.cc

@@ -484,7 +484,6 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
     }
   }
 
-  bool is_one_result = result_pattern_graph.owned_op_call().size() == 1;
   // topo order visit result_pattern_graph
   GraphTopo graph_topo_visit(&result_pattern_graph);
   graph_topo_visit.WalkGraphNodesTopoOrder([&](const OpCall& op_call) {
@@ -496,8 +495,9 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
     }
 
     // set insert point
-    size_t max_input_op_index = 0UL;
-    pir::Operation* max_index_op = nullptr;
+    // 1. get result pattern max-idx of input op
+    size_t max_res_idx = 0UL;
+    pir::Operation* max_res_idx_op = nullptr;
     for (const Tensor* input : op_call.inputs()) {
       if (input->is_none()) {
         continue;
@@ -507,18 +507,16 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
         pir::Operation* ir_input_op = ir_val.defining_op();
         if (op_2_temp_program_index.count(ir_input_op) == 0) {
           // do nothing
-        } else if (max_input_op_index <
-                   op_2_temp_program_index.at(ir_input_op)) {
-          max_input_op_index = op_2_temp_program_index.at(ir_input_op);
-          max_index_op = ir_input_op;
-        } else if (max_input_op_index ==
-                   op_2_temp_program_index.at(ir_input_op)) {
-          const auto& ops_vec = temp_program[max_input_op_index];
+        } else if (max_res_idx < op_2_temp_program_index.at(ir_input_op)) {
+          max_res_idx = op_2_temp_program_index.at(ir_input_op);
+          max_res_idx_op = ir_input_op;
+        } else if (max_res_idx == op_2_temp_program_index.at(ir_input_op)) {
+          const auto& ops_vec = temp_program[max_res_idx];
           for (auto it = ops_vec.begin(); it != ops_vec.end(); it++) {
-            if (*it == max_index_op) {
+            if (*it == max_res_idx_op) {
               break;
             } else if (*it == ir_input_op) {
-              max_index_op = ir_input_op;
+              max_res_idx_op = ir_input_op;
               break;
             } else {
               // do nothing
@@ -530,51 +528,28 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
       }
     }
 
-    if (is_one_result && !source_pattern_graph.owned_op_call().empty()) {
-      // 1. get source pattern min-idx op
-      pir::Operation* min_src_idx_op = src_match_ctx.IrOperation(
-          source_pattern_graph.owned_op_call()[0].get());
-      size_t min_src_idx = op_2_temp_program_index[min_src_idx_op];
-      for (const auto& src_owned_op_call :
-           source_pattern_graph.owned_op_call()) {
-        pir::Operation* src_owned_op =
-            src_match_ctx.IrOperation(src_owned_op_call.get());
-        size_t src_owned_op_idx = op_2_temp_program_index[src_owned_op];
-        if (min_src_idx > src_owned_op_idx) {
-          min_src_idx = src_owned_op_idx;
-          min_src_idx_op = src_owned_op;
-        }
-      }
-      // 2. insert new op at point max(max_input_op_index+1, min_src_idx)
-      if (min_src_idx > max_input_op_index) {
-        rewriter.set_insertion_point(min_src_idx_op);
-        max_input_op_index = op_2_temp_program_index[min_src_idx_op];
-      } else {
-        rewriter.SetInsertionPointAfter(max_index_op);
-      }
-      VLOG(6) << "(" << op_call.name() << ") insert at idx "
-              << std::max(max_input_op_index + 1, min_src_idx);
-    } else {
-      if (max_input_op_index == 0UL) {
-        VLOG(6) << "Not found producer op for (" << op_call.name() << ")";
-        pir::Operation* source_pattern_first_op = src_match_ctx.IrOperation(
-            source_pattern_graph.owned_op_call()[0].get());
-        max_input_op_index = op_2_temp_program_index[source_pattern_first_op];
-        rewriter.set_insertion_point(source_pattern_first_op);
-      } else {
-        rewriter.SetInsertionPointAfter(max_index_op);
+    // 2. get source pattern min-idx op
+    pir::Operation* min_src_idx_op = src_match_ctx.IrOperation(
+        source_pattern_graph.owned_op_call()[0].get());
+    size_t min_src_idx = op_2_temp_program_index[min_src_idx_op];
+    for (const auto& src_owned_op_call : source_pattern_graph.owned_op_call()) {
+      pir::Operation* src_owned_op =
+          src_match_ctx.IrOperation(src_owned_op_call.get());
+      size_t src_owned_op_idx = op_2_temp_program_index[src_owned_op];
+      if (min_src_idx > src_owned_op_idx) {
+        min_src_idx = src_owned_op_idx;
+        min_src_idx_op = src_owned_op;
       }
     }
 
-    pir::Operation* new_op =
-        CreateOperation(op_call, src_match_ctx, rewriter, &res_match_ctx);
-
-    size_t new_max_input_op_index = max_input_op_index + 1;
-    op_2_temp_program_index[new_op] = new_max_input_op_index;
-    if (new_max_input_op_index >= temp_program.size()) {
-      temp_program.emplace_back();
+    // 3. insert new op at point max(max_res_idx+1, min_src_idx)
+    if (min_src_idx > max_res_idx) {
+      rewriter.set_insertion_point(min_src_idx_op);
+    } else {
+      rewriter.SetInsertionPointAfter(max_res_idx_op);
     }
-    temp_program[new_max_input_op_index].push_back(new_op);
+
+    CreateOperation(op_call, src_match_ctx, rewriter, &res_match_ctx);
   });
 
   return res_match_ctx;