【custom】add train pass flag (#70909)

xiaoguoguo626807 · web-flow · commit 8c590fa2a825 · 2025-01-23T09:17:16.000+08:00
* add train pass flag

* add jit path
diff --git a/paddle/common/flags.cc b/paddle/common/flags.cc
@@ -1533,6 +1533,18 @@ PHI_DEFINE_EXPORTED_bool(enable_pir_in_executor,
                          false,
                          "Enable PIR in executor");
 
+/**
+ * Using PIR API in Python
+ * Name: enable_custom_engine
+ * Since Version: 3.0.0
+ * Value Range: bool, default=false
+ * Example:
+ * Note: If True, CustomDevice can use subgraph engine optimize
+ */
+PHI_DEFINE_EXPORTED_string(enable_custom_engine,
+                           "",
+                           "Set CustomDevice subgraph engine translate pass");
+
 /**
  * Using PIR by translating legacy program to pir program
  * for dy2st mode  FLAG
diff --git a/paddle/fluid/framework/executor_cache.cc b/paddle/fluid/framework/executor_cache.cc
@@ -26,11 +26,13 @@
 #include "paddle/pir/include/core/value.h"
 #include "paddle/pir/include/pass/pass.h"
 #include "paddle/pir/include/pass/pass_manager.h"
+#include "paddle/pir/include/pass/pass_registry.h"
 
 DECLARE_FILE_SYMBOLS(print_statistics);
 
 COMMON_DECLARE_bool(pir_apply_inplace_pass);
 COMMON_DECLARE_bool(print_ir);
+COMMON_DECLARE_string(enable_custom_engine);
 
 namespace paddle::framework {
 class ProgramDesc;
@@ -176,6 +178,25 @@ std::unique_ptr<::pir::Program> ApplyIrPass(
     ::pir::Program *program,
     phi::Place place,
     const std::set<std::string> &no_need_buffer_names) {
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+  if (!FLAGS_enable_custom_engine.empty()) {
+    std::string custom_engine_translate_pass = FLAGS_enable_custom_engine;
+    std::istringstream ss(custom_engine_translate_pass);
+    std::string pass;
+    std::vector<std::string> passes;
+
+    while (std::getline(ss, pass, ',')) {
+      passes.push_back(pass);
+      VLOG(4) << "Add CustomEngine pass : " << pass;
+    }
+
+    ::pir::PassManager pass_pm(::pir::IrContext::Instance(), 3);
+    for (std::string custom_pass : passes) {
+      pass_pm.AddPass(pir::PassRegistry::Instance().Get(custom_pass));
+      pass_pm.Run(program);
+    }
+  }
+#endif
   auto ir_res = paddle::dialect::PdOpLowerToKernelPass(program, place);
 
   if (FLAGS_pir_apply_inplace_pass) {
diff --git a/paddle/fluid/framework/new_executor/standalone_executor.cc b/paddle/fluid/framework/new_executor/standalone_executor.cc
@@ -18,19 +18,20 @@
 #include "paddle/fluid/framework/new_executor/interpreter/interpreter_util.h"
 #include "paddle/fluid/framework/new_executor/pir_interpreter.h"
 #include "paddle/fluid/framework/new_executor/program_interpreter.h"
+#include "paddle/fluid/ir_adaptor/translator/translate.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/general/inplace_pass.h"
 #include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
 #include "paddle/phi/core/platform/profiler/event_tracing.h"
-
-#include "paddle/fluid/ir_adaptor/translator/translate.h"
-#include "paddle/fluid/pir/transforms/general/inplace_pass.h"
 #include "paddle/pir/include/core/program.h"
 #include "paddle/pir/include/pass/pass.h"
 #include "paddle/pir/include/pass/pass_manager.h"
+#include "paddle/pir/include/pass/pass_registry.h"
 
 COMMON_DECLARE_bool(enable_pir_in_executor);
 COMMON_DECLARE_bool(enable_pir_api);
 COMMON_DECLARE_bool(pir_apply_inplace_pass);
+COMMON_DECLARE_string(enable_custom_engine);
 
 namespace paddle::framework {
 StandaloneExecutor::StandaloneExecutor(const phi::Place& place,
@@ -111,6 +112,25 @@ StandaloneExecutor::StandaloneExecutor(const phi::Place& place,
           job->SetFetchVarName(fetch_var_names_[index]);
         }
       }
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+      if (!FLAGS_enable_custom_engine.empty()) {
+        std::string custom_engine_translate_pass = FLAGS_enable_custom_engine;
+        std::istringstream ss(custom_engine_translate_pass);
+        std::string pass;
+        std::vector<std::string> passes;
+
+        while (std::getline(ss, pass, ',')) {
+          passes.push_back(pass);
+          VLOG(4) << "Add CustomEngine pass : " << pass;
+        }
+
+        ::pir::PassManager pass_pm(::pir::IrContext::Instance(), 3);
+        for (std::string custom_pass : passes) {
+          pass_pm.AddPass(pir::PassRegistry::Instance().Get(custom_pass));
+          pass_pm.Run(base_program.get());
+        }
+      }
+#endif
       auto kernel_program =
           paddle::dialect::PdOpLowerToKernelPass(base_program.get(), place);
       std::shared_ptr<pir::Program> shared_program = std::move(kernel_program);
