[CINN] Add cinn process in auto_parallel (#69362)

* refine

* fix

* refine code

Author: zhangbo9674

python/paddle/base/executor.py

@@ -1152,6 +1152,21 @@ def _get_pir_program_and_executor(self, cached_data):
         place = cached_data.place
         scope = cached_data.scope
 
+        def cinn_process(program):
+            from paddle.decomposition import decomp
+
+            if core._enable_dist_prim_all():
+                logging.info("apply decompose in executor")
+                with decomp.prim_guard():
+                    decomp.decompose_dist_program(program)
+
+            if core._enable_auto_recompute():
+                logging.info("apply auto_recompute in executor")
+                program = decomp.auto_recompute_pir_program(program, None)
+
+            apply_cinn_pass(program)
+            return program
+
         if cached_data.plan is None:
             value_map = pir.IrMapping()
             _, is_startup_program = has_fetch_operations_and_is_startup_program(
@@ -1174,6 +1189,10 @@ def _get_pir_program_and_executor(self, cached_data):
                 fetch_var_name=fetch_var_name,
             )
             default_job = core.Job("default")
+
+            if not is_startup_program and in_cinn_mode():
+                cinn_process(program)
+
             type_to_program = {"default": program}
             plan = core.Plan([default_job], type_to_program)
         else:
@@ -1200,6 +1219,11 @@ def _get_pir_program_and_executor(self, cached_data):
                     value.block.program, value, fetch_var_name + str(i), i
                 )
 
+            if in_cinn_mode():
+                for job_type in plan.job_types():
+                    ir_program = plan.ir_program(job_type)
+                    cinn_process(ir_program)
+
         new_exe = _StandaloneExecutor(place, plan, scope)
 
         data_op_infos = []
@@ -1216,18 +1240,7 @@ def _get_pir_program_and_executor(self, cached_data):
                     op.result(0).persistable,
                 )
                 data_op_infos.append(tup)
-        from paddle.decomposition import decomp
-
-        if core._enable_dist_prim_all():
-            with decomp.prim_guard():
-                decomp.decompose_dist_program(program)
 
-        if core._enable_auto_recompute():
-            logging.info("apply auto_recompute in executor")
-            program = decomp.auto_recompute_pir_program(program, None)
-
-        if in_cinn_mode():
-            apply_cinn_pass(program)
         return program, new_exe, data_op_infos
 
 

python/paddle/distributed/auto_parallel/static/engine.py

@@ -29,6 +29,7 @@
 from paddle import pir, static, utils
 from paddle.base.executor import _to_name_str
 from paddle.base.framework import auto_complete_op_role
+from paddle.decomposition import decomp
 from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.distributed.passes.pass_base import new_pass
 from paddle.distributed.passes.pass_utils import (
@@ -897,20 +898,15 @@ def _parallel_pir(self, mode):
         remove_unuseful_comm_op_pass(dense_program)
 
         if core._enable_dist_prim_all():
-            from paddle.decomposition import decomp
-
+            logging.info("apply decompose in auto parallel")
             with decomp.prim_guard():
                 decomp.decompose_dist_program(dense_program)
 
         if core._enable_auto_recompute():
-            from paddle.decomposition import decomp
-
             logging.info("apply auto_recompute in auto parallel")
             dense_program = decomp.auto_recompute_pir_program(
                 dense_program,
-                lambda op: bool(
-                    op.has_attr('op_role') and op.attrs()["op_role"] == 0
-                ),
+                lambda op: bool(op.has_attr('op_role') and op.op_role == 0),
             )
 
         if self._strategy.pipeline.enable: