[PIR] translate newir in python (#58390)

zhaoyinglia · web-flow · commit e2308b7bac07 · 2023-10-26T14:50:18.000+08:00
diff --git a/paddle/fluid/framework/new_executor/standalone_executor.cc b/paddle/fluid/framework/new_executor/standalone_executor.cc
@@ -55,7 +55,7 @@ StandaloneExecutor::StandaloneExecutor(const platform::Place& place,
     const std::string& job_type = job->Type();
     std::shared_ptr<ProgramDesc> program = nullptr;
     std::shared_ptr<::pir::Program> ir_program = nullptr;
-    if (FLAGS_enable_pir_api) {
+    if (FLAGS_enable_pir_api || FLAGS_enable_new_ir_in_executor) {
       ir_program = plan_.IrProgram(job_type);
     } else {
       program = std::make_shared<ProgramDesc>(*(plan_.Program(job_type)));
@@ -80,10 +80,6 @@ StandaloneExecutor::StandaloneExecutor(const platform::Place& place,
     // TODO(phlrain) we only support cpu for now
     if (FLAGS_enable_new_ir_in_executor) {
       std::shared_ptr<::pir::Program> base_program = ir_program;
-      if (!FLAGS_enable_pir_api) {
-        VLOG(6) << "begin to translate" << std::endl;
-        base_program = paddle::TranslateLegacyProgramToProgram(*program);
-      }
       auto block = base_program->block();
       for (auto it = block->begin(); it != block->end(); ++it) {
         if ((*it)->isa<paddle::dialect::FetchOp>()) {
diff --git a/paddle/fluid/ir_adaptor/translator/translate.cc b/paddle/fluid/ir_adaptor/translator/translate.cc
@@ -34,8 +34,9 @@ std::unique_ptr<Program> TranslateLegacyProgramToProgram(
   auto program = std::make_unique<Program>(ctx);
   translator::ProgramTranslator program_translator(&legacy_program,
                                                    program.get());
+  VLOG(6) << "begin to translate";
   program_translator.Translate();
-
+  VLOG(6) << "translate done";
   return program;
 }
 
diff --git a/python/paddle/base/executor.py b/python/paddle/base/executor.py
@@ -21,7 +21,7 @@
 
 import numpy as np
 
-from ..pir import OpResult
+from ..pir import OpResult, translate_to_new_ir
 from . import compiler, core, framework, get_flags, set_flags, unique_name
 from .data_feeder import convert_dtype
 from .framework import (
@@ -1006,7 +1006,14 @@ def _get_program_and_executor(self, cached_data):
             )
         else:
             default_job = core.Job("default")
-            type_to_program = {"default": new_program.desc}
+            if get_flags("FLAGS_enable_new_ir_in_executor")[
+                'FLAGS_enable_new_ir_in_executor'
+            ]:
+                type_to_program = {
+                    "default": translate_to_new_ir(new_program.desc)
+                }
+            else:
+                type_to_program = {"default": new_program.desc}
             plan = core.Plan([default_job], type_to_program)
 
         new_exe = _StandaloneExecutor(place, plan, scope)
