PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc
Copy file name to clipboard
+44-25Lines changed: 44 additions & 25 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc
Copy file name to clipboard
+44-25Lines changed: 44 additions & 25 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/utils.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/cinn/hlir/framework/pir/utils.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/utils.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/cinn/hlir/framework/pir/utils.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir_compiler.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions b/‎paddle/cinn/hlir/framework/pir_compiler.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/to_static/run_program_op_func.h
Copy file name to clipboard
+4-3Lines changed: 4 additions & 3 deletions b/‎paddle/fluid/eager/to_static/run_program_op_func.h
Copy file name to clipboard
+4-3Lines changed: 4 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/transforms/build_cinn_pass.cc
Copy file name to clipboard
+14-12Lines changed: 14 additions & 12 deletions b/‎paddle/fluid/pir/transforms/build_cinn_pass.cc
Copy file name to clipboard
+14-12Lines changed: 14 additions & 12 deletions
diff --git a/‎paddle/fluid/pybind/CMakeLists.txt
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/pybind/CMakeLists.txt
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+62-10Lines changed: 62 additions & 10 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+62-10Lines changed: 62 additions & 10 deletions
@@ -42,35 +42,59 @@ std::vector<pir::Value> GetBlockOutsideInput(
     }
   }
 
+  std::unordered_set<::pir::Value> insert_value;
   for (size_t k = 0; k < op_list.size(); ++k) {
     for (size_t i = 0; i < op_list[k]->num_operands(); ++i) {
-      if (!block_inner_output.count(op_list[k]->operand_source(i))) {
+      if (!block_inner_output.count(op_list[k]->operand_source(i)) &&
+          !insert_value.count(op_list[k]->operand_source(i))) {
         vec_res.push_back(op_list[k]->operand_source(i));
+        insert_value.insert(op_list[k]->operand_source(i));
       }
     }
   }
-
   return vec_res;
 }
 
 std::vector<pir::Value> GetBlockOutsideOutput(
-    const std::vector<pir::Operation*> op_list) {
-  std::vector<pir::Value> vec_res;
-  std::unordered_set<::pir::Value> block_inner_output;
-  for (size_t k = 0; k < op_list.size(); ++k) {
-    for (size_t i = 0; i < op_list[k]->num_operands(); ++i) {
-      block_inner_output.insert(op_list[k]->operand_source(i));
+    const std::vector<pir::Operation*> op_list,
+    const std::vector<pir::Operation*> group_all_list) {
+  assert(group_all_list.size() >= 2);
+  assert(group_all_list.back()->isa<pir::YieldOp>());
+
+  auto yeild_op = group_all_list.back()->dyn_cast<pir::YieldOp>();
+
+  std::unordered_set<pir::Value> yeild_inputs;
+  for (size_t i = 0; i < yeild_op.num_operands(); ++i) {
+    yeild_inputs.insert(yeild_op.operand_source(i));
+  }
+
+  std::unordered_set<pir::Operation*> innner_op_set(op_list.begin(),
+                                                    op_list.end());
+  std::unordered_set<pir::Operation*> outside_group_set;
+
+  for (size_t i = 0; i < group_all_list.size(); ++i) {
+    if (!innner_op_set.count(group_all_list[i])) {
+      outside_group_set.insert(group_all_list[i]);
     }
   }
 
-  for (size_t k = 0; k < op_list.size(); ++k) {
-    for (size_t i = 0; i < op_list[k]->num_results(); ++i) {
-      if (!block_inner_output.count(op_list[k]->result(i))) {
-        vec_res.push_back(op_list[k]->result(i));
+  std::vector<pir::Value> vec_res;
+
+  for (auto* op : op_list) {
+    for (size_t i = 0; i < op->num_results(); ++i) {
+      if (yeild_inputs.count(op->result(i))) {
+        vec_res.push_back(op->result(i));
+      } else {
+        for (auto it = op->result(i).use_begin(); it != op->result(i).use_end();
+             ++it) {
+          if (outside_group_set.count(it->owner())) {
+            vec_res.push_back(op->result(i));
+            break;
+          }
+        }
       }
     }
   }
-
   return vec_res;
 }
 
@@ -97,7 +121,6 @@ std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
 
   auto ir_program = std::make_unique<::pir::Program>(ctx);
   std::unordered_map<pir::Value, pir::Value> value_map;
-  std::vector<cinn::hlir::framework::PIRCompiler*> compiler_list;
 
   auto target = cinn::common::DefaultNVGPUTarget();
   auto scope = cinn::hlir::framework::BuildScope(target, *program);
@@ -121,12 +144,12 @@ std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
                         phi::errors::Unimplemented(
                             "Only support one group after group fusion"));
       for (auto group : group_list) {
-        auto ir_compiler =
-            new cinn::hlir::framework::PIRCompiler(*program, target, scope);
+        auto ir_compiler = std::make_shared<cinn::hlir::framework::PIRCompiler>(
+            *program, target, scope);
+        hlir::framework::PIRCompilerManager::Instance().insert(ir_compiler);
         auto group1 =
             std::make_shared<cinn::hlir::framework::pir::Group>(group->ops);
         auto fn_ptr_res = ir_compiler->BuildCUDAJITInfo({group1});
-        compiler_list.push_back(ir_compiler);
         std::unordered_map<std::string, ::pir::Attribute> op_attrs{
             {cinn::dialect::JitKernelOp::kAttrName,
              cinn::dialect::CUDAJITInfoAttribute::get(ctx, fn_ptr_res[0])},
@@ -140,7 +163,7 @@ std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
           vec_new_ins.push_back(value_map.at(vec_ins[i]));
         }
 
-        auto vec_outs = GetBlockOutsideOutput(group->ops);
+        auto vec_outs = GetBlockOutsideOutput(group->ops, group_op.ops());
 
         std::vector<pir::Type> vec_types;
         for (auto& out : vec_outs) {
@@ -150,11 +173,6 @@ std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
         ::pir::Operation* cinn_op =
             ::pir::Operation::Create(vec_new_ins, op_attrs, vec_types, op_info);
 
-        // for (size_t i = 0; i < vec_outs.size(); ++i) {
-        //   value_map[vec_outs[i]] = cinn_op->result(i);
-        // }
-
-        // auto yield_op = group_op.ops().back()->dyn_cast<pir::YieldOp>();
         for (size_t i = 0; i < group_op.num_results(); ++i) {
           value_map[group_op.result(i)] = cinn_op->result(i);
         }
@@ -179,8 +197,9 @@ std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
           vec_ins, (*it)->attributes(), vec_types, info1);
 
       ir_program->block()->push_back(op);
-
-      value_map[(*it)->result(0)] = op->result(0);
+      for (size_t i = 0; i < (*it)->num_results(); ++i) {
+        value_map[(*it)->result(i)] = op->result(i);
+      }
     }
   }
   return ir_program;
 
@@ -36,6 +36,14 @@ const std::unordered_map<std::string, std::string> CompatibleInfo::OP_NAMES = {
     {"pd_op.max", "reduce_max"},
     {"pd_op.add", "elementwise_add"}};
 
+const std::unordered_set<std::string> CompatibleInfo::CINN_WHITE_OPS = {
+    "subtract"};
+
+bool CompatibleInfo::IsSupportCinn(const ::pir::Operation& op) {
+  return CINN_WHITE_OPS.find(CompatibleInfo::OpName(op)) !=
+         CINN_WHITE_OPS.end();
+}
+
 std::string CompatibleInfo::OpName(const ::pir::Operation& op) {
   std::string name = op.name();
   if (OP_NAMES.count(name)) {
 
@@ -15,6 +15,7 @@
 #pragma once
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include "paddle/cinn/common/context.h"
 #include "paddle/cinn/common/type.h"
 #include "paddle/cinn/hlir/framework/op.h"
@@ -39,6 +40,13 @@ struct CompatibleInfo {
   // TODO(Aurelius): Need add name mapping logic in REGISTER_CINN_OP
   // macros or attempt to unify Op name with Paddle and CINN.
   static const std::unordered_map<std::string, std::string> OP_NAMES;
+  // NOTE(Aurelius): Some ops in CINN register different
+  // name between OpMapper and Compute/Schedule, such as
+  // 'subtract': 1. OpMapper: 'elementwise_sub'; 2. Compute/Schedule:
+  // 'subtract'.
+  static const std::unordered_set<std::string> CINN_WHITE_OPS;
+
+  static bool IsSupportCinn(const ::pir::Operation& op);
 
   static std::string OpName(const ::pir::Operation& op);
 
 
@@ -65,6 +65,23 @@ class PIRCompiler final {
 
 std::shared_ptr<Scope> BuildScope(const Target&, const ::pir::Program&);
 
+class PIRCompilerManager {
+ public:
+  static PIRCompilerManager& Instance() {
+    static PIRCompilerManager instance;
+    return instance;
+  }
+
+  void insert(const std::shared_ptr<PIRCompiler>& compiler) {
+    compilers_.push_back(compiler);
+  }
+
+  void clear() { compilers_.clear(); }
+
+ private:
+  std::vector<std::shared_ptr<PIRCompiler>> compilers_;
+};
+
 }  // namespace framework
 }  // namespace hlir
 }  // namespace cinn
@@ -232,13 +232,14 @@ inline void newir_run_program_ad_func(
   auto output_size =
       PADDLE_GET_CONST(std::vector<::pir::Value>, attrs.at("fo")).size();
   auto middles = std::vector<paddle::Tensor*>();
-  std::shared_ptr<NewIRGradNodeRunProgram> grad_node;
-  VLOG(2) << "start run run_program with require_any_grad = "
-          << require_any_grad;
+
   auto is_test = false;
   if (attrs.count("is_test")) {
     is_test = PADDLE_GET_CONST(bool, attrs.at("is_test"));
   }
+  std::shared_ptr<NewIRGradNodeRunProgram> grad_node;
+  VLOG(2) << "start run run_program with require_any_grad = "
+          << require_any_grad << ", is_test = " << is_test;
 
   if (!is_test && require_any_grad) {
     // Create GradOpNode (1 means [out_grad], 2 means [x_grad, paramx_grad])
 
@@ -14,7 +14,6 @@
 
 #include "paddle/fluid/pir/transforms/build_cinn_pass.h"
 
-#include <algorithm>
 #include <queue>
 #include <regex>
 #include <set>
@@ -136,7 +135,12 @@ bool IsSupportCinn(pir::Operation* op) {
   VLOG(4) << "The allowed Cinn Ops: " << GetDebugInfo(allow_ops);
   VLOG(4) << "The denied Cinn Ops: " << GetDebugInfo(deny_ops);
   // Strip the dialect, like pd_op.abs -> abs
-  const auto& op_name = CompatibleInfo::OpName(*op);
+  const auto op_name = CompatibleInfo::OpName(*op);
+  if (CompatibleInfo::IsSupportCinn(*op)) {
+    VLOG(4) << "Found special supported op for CINN: " << op_name;
+    return true;
+  }
+
   bool registered =
       ::cinn::frontend::OpMapperRegistry::Global()->Find(op_name) != nullptr;
 
@@ -532,12 +536,10 @@ std::vector<pir::Value> AnalysisOutputs(GroupOpsVec& group_ops) {  // NOLINT
       outputs.emplace(result);
     }
   }
-  std::vector<pir::Value> results;
-  std::set_symmetric_difference(outputs.begin(),
-                                outputs.end(),
-                                inputs.begin(),
-                                inputs.end(),
-                                std::back_inserter(results));
+  for (auto& input : inputs) {
+    outputs.erase(input);
+  }
+  std::vector<pir::Value> results(outputs.begin(), outputs.end());
   VLOG(3) << "Outputs size for GroupOp " << results.size();
   return results;
 }
@@ -562,14 +564,14 @@ void ReplaceWithGroupOp(pir::Block* block,
   for (auto* op : group_ops) {
     op->MoveTo(group_block, group_block->begin());
   }
-  // step 3: Insert YieldOp for outputs
-  builder.SetInsertionPointToEnd(group_block);
-  builder.Build<::pir::YieldOp>(outputs);
-  // step 4: Replace outputs of inner ops
+  // step 3: Replace outputs of inner ops
   std::vector<pir::OpResult> group_outs = new_group_op->results();
   for (size_t i = 0; i < outputs.size(); ++i) {
     outputs[i].ReplaceAllUsesWith(group_outs[i]);
   }
+  // step 4: Insert YieldOp for outputs
+  builder.SetInsertionPointToEnd(group_block);
+  builder.Build<::pir::YieldOp>(outputs);
 }
 
 class BuildCinnPass : public pir::Pass {
 
@@ -53,6 +53,10 @@ set(PYBIND_DEPS
     type_info
     auto_parallel)
 
+if(WITH_CINN)
+  set(PYBIND_DEPS ${PYBIND_DEPS} pd_build_cinn_pass op_with_group_merge_pass)
+endif()
+
 if(WITH_PSCORE)
   set(PYBIND_DEPS ${PYBIND_DEPS} ps_service)
   if(WITH_HETERPS)
 
@@ -54,6 +54,12 @@
 #include "paddle/utils/flags.h"
 #include "pybind11/stl.h"
 
+#ifdef PADDLE_WITH_CINN
+#include "paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.h"
+#include "paddle/cinn/hlir/framework/pir_compiler.h"
+#include "paddle/fluid/pir/transforms/build_cinn_pass.h"
+#endif
+
 namespace py = pybind11;
 using paddle::dialect::APIBuilder;
 using paddle::dialect::DenseTensorType;
@@ -69,13 +75,13 @@ using pir::Type;
 using pir::Value;
 using pybind11::return_value_policy;
 
-USE_PASS(dead_code_elimination_pass);
-USE_PASS(attention_fuse_pass);
-USE_PASS(fused_gemm_epilogue_pass);
-USE_PASS(fused_dropout_add_pass);
-USE_PASS(fused_linear_param_grad_add_pass);
-USE_PASS(inplace_pass);
-USE_PASS(replace_fetch_with_shadow_output_pass);
+USE_PIR_PASS(dead_code_elimination_pass);
+USE_PIR_PASS(attention_fuse_pass);
+USE_PIR_PASS(fused_gemm_epilogue_pass);
+USE_PIR_PASS(fused_dropout_add_pass);
+USE_PIR_PASS(fused_linear_param_grad_add_pass);
+USE_PIR_PASS(inplace_pass);
+USE_PIR_PASS(replace_fetch_with_shadow_output_pass);
 
 PHI_DECLARE_bool(print_ir);
 
@@ -1082,11 +1088,28 @@ SplitedResult SplitForwardBackward(
     if (v.impl() == nullptr) {
       return;
     }
+    // NOTE(Aurelius84): we should skip insert SetParameterOp repeatly by
+    // calling SplitForwardBackward multi-times.
+    std::string parameter_name =
+        std::string("output_") + std::to_string(counter);
+    for (auto it = forward_program->block()->rbegin();
+         it != forward_program->block()->rend();
+         ++it) {
+      auto *op = *it;
+      if (op->isa<pir::SetParameterOp>()) {
+        auto out_name =
+            op->attribute<pir::StrAttribute>("parameter_name").AsString();
+        if (out_name == parameter_name) {
+          VLOG(4) << out_name
+                  << " has been inserted SetParameterOp, skip it now.";
+          return;
+        }
+      }
+    }
+
     auto op_info = ctx->GetRegisteredOpInfo(pir::SetParameterOp::name());
     pir::AttributeMap attribute_map = {
-        {"parameter_name",
-         pir::StrAttribute::get(
-             ctx, std::string("output_") + std::to_string(counter))},
+        {"parameter_name", pir::StrAttribute::get(ctx, parameter_name)},
     };
     pir::Operation *operation = pir::Operation::Create(
         {forward_value_map[v]}, attribute_map, {}, op_info);
@@ -1362,9 +1385,38 @@ void BindUtils(pybind11::module *m) {
                 >>> print(mappings)
                 {'matmul_v2_0.tmp_0': [Value(define_op_name=pd_op.matmul, index=0, dtype=pd_op.tensor<4x4xf32>)], 'x': [Value(define_op_name=pd_op.data, index=0, dtype=pd_op.tensor<4x4xf32>)], 'tanh_0.tmp_0': [Value(define_op_name=pd_op.tanh, index=0, dtype=pd_op.tensor<4x4xf32>)], 'elementwise_add_0': [Value(define_op_name=pd_op.add, index=0, dtype=pd_op.tensor<4x4xf32>)]}
     )DOC");
+
+  m->def("clear_pir_compiler_manager", []() {
+#ifdef PADDLE_WITH_CINN
+    pybind11::gil_scoped_release release;
+    VLOG(4) << "clear PIRCompilerManager and free PIRCompiler resources.";
+    cinn::hlir::framework::PIRCompilerManager::Instance().clear();
+#endif
+  });
 }
 
+// TODO(Aurelius84): Need consider to make an agreement about
+// what a Pass should receive and return. Existed Passes have
+// mutable and immutable interface.
+std::shared_ptr<Program> ApplyPirPass(Program &forward_program) {  // NOLINT
+#ifdef PADDLE_WITH_CINN
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  pir::PassManager pass_manager(ctx);
+  pass_manager.AddPass(pir::CreateBuildCinnPass());
+  pass_manager.Run(&forward_program);
+  VLOG(3) << "after BuildCinnPass, forward_program:\n" << forward_program;
+  std::unique_ptr<pir::Program> new_program =
+      cinn::dialect::ir::CINNGroupLoweringPass(&forward_program);
+  VLOG(3) << "after CINNGroupLoweringPass, forward_program:\n" << *new_program;
+  return std::move(new_program);
+#endif
+  PADDLE_THROW(platform::errors::Unimplemented(
+      "Currently we only support CINN Pass for PIR under @to_static, please "
+      "compile PaddlePaddle with CINN"));
+}
 void BindIrPass(pybind11::module *m) {
+  m->def("apply_pir_pass", ApplyPirPass);
+
   py::class_<Pass, std::shared_ptr<Pass>> pass(*m,
                                                "Pass",
                                                R"DOC(