Fit vpp for pir (#60823)

From00 · web-flow · commit 6d52468d8281 · 2024-02-02T10:57:14.000+08:00
* Fit vpp for pir

* Add debug log

* Remove unused code

* Remove unused code

* Fix CI errors

* Fix CI errors
diff --git a/paddle/fluid/framework/new_executor/instruction/instruction_base.cc b/paddle/fluid/framework/new_executor/instruction/instruction_base.cc
@@ -359,10 +359,6 @@ std::string InstructionBase::DebugStringEx(
         if (row_size >= 0) {
           ss << "row_size=" << row_size << ";";
         }
-        double ele_sum = GetDenseTensorEleSum(*scope, var_name);
-        if (!std::isnan(ele_sum)) {
-          ss << "ele_sum=" << ele_sum << ";";
-        }
       }
     }
     ++it;
@@ -390,10 +386,6 @@ std::string InstructionBase::DebugStringEx(
         if (row_size >= 0) {
           ss << "row_size=" << row_size << ";";
         }
-        double ele_sum = GetDenseTensorEleSum(*scope, var_name);
-        if (!std::isnan(ele_sum)) {
-          ss << "ele_sum=" << ele_sum << ";";
-        }
       }
     }
     ++it;
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -17,7 +17,7 @@
 from enum import Enum
 
 from paddle.base import core
-from paddle.base.framework import Operator, Parameter, Program, get_flags
+from paddle.base.framework import Parameter, Program
 from paddle.distributed.auto_parallel.static.dist_attribute import (
     OperatorDistAttr,
 )
@@ -228,43 +228,6 @@ def var_can_be_deleted(var_name, block):
     return var is not None and not var.persistable
 
 
-def prepare_ir_program(cur_prog, next_prog):
-    set_output_names = set()
-    for op in cur_prog.global_block().ops:
-        for arg_name in op.output_arg_names:
-            if var_can_be_deleted(arg_name, cur_prog.global_block()):
-                set_output_names.add(arg_name)
-
-    set_input_names = set()
-    for op in next_prog.global_block().ops:
-        for arg_name in op.input_arg_names:
-            if var_can_be_deleted(arg_name, next_prog.global_block()):
-                set_input_names.add(arg_name)
-
-    shadow_var_names = sorted(set_output_names & set_input_names)
-    for var_name in shadow_var_names:
-        shadow_op_desc = cur_prog.global_block().desc.append_op()
-        shadow_op_desc.set_type("shadow_output")
-        shadow_op_desc.set_input('x', [var_name])
-        shadow_op_desc.set_output('out', ["@EMPTY@"])
-        shadow_op_desc._set_attr("name", var_name)
-        shadow_op = Operator(cur_prog.global_block(), shadow_op_desc)
-        cur_prog.global_block().ops.append(shadow_op)
-
-        data_op_desc = next_prog.global_block().desc._prepend_op()
-        data_op_desc.set_type("data")
-        data_op_desc._set_attr("shape", [])
-        data_op_desc._set_attr("dtype", 0)
-        data_op_desc._set_attr("place", 2)  # GPUPlace
-        data_op_desc._set_attr("name", var_name)
-        data_op_desc.set_output("out", [var_name])
-        data_op = Operator(next_prog.global_block(), data_op_desc)
-        next_prog.global_block().ops.insert(0, data_op)
-
-    cur_prog._sync_with_cpp()
-    next_prog._sync_with_cpp()
-
-
 def set_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
     """
     Set `skip_gc_vars` for every job in jobs.
@@ -285,6 +248,7 @@ def set_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
                 if op.type in [
                     "c_sync_comm_stream",
                     "conditional_block",
+                    "data",
                     "nop",
                     "while",
                 ]:
@@ -319,18 +283,65 @@ def set_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
         job.set_skip_gc_vars(skip_gc_vars)
         suffixed_required_vars[micro_batch_id] |= required_vars
 
-    if get_flags("FLAGS_enable_pir_in_executor")[
-        'FLAGS_enable_pir_in_executor'
-    ]:
-        for i, type in enumerate(job_types):
-            if i == len(job_types) - 1:
-                break
-            next_type = job_types[i + 1]
-            prepare_ir_program(
-                type_to_program[type], type_to_program[next_type]
+    return type_to_program
+
+
+def shadow_var_between_sub_programs(sub_programs):
+    """
+    Add shadow_output and data op pair to share vars between sub_programs.
+    """
+    suffixed_shadow_arg_names = (
+        set()
+    )  # arg_names that are required in later sub_programs
+    for sub_program in reversed(sub_programs):
+        # step 1: parse shadow arguements
+        block = sub_program.global_block()
+        input_arg_names = set()
+        output_arg_names = set()
+        shadow_arg_names = set()
+        for op in block.ops:
+            for input_arg_name in op.input_arg_names:
+                if var_can_be_deleted(input_arg_name, block):
+                    input_arg_names.add(input_arg_name)
+                    # NOTE(Ruibiao): When translating these codes to pir, we can simplely set
+                    # `shadow_arg_names=input_arg_names-output_arg_names` since the program
+                    # in pir satifies SSA form.
+                    if input_arg_name not in output_arg_names:
+                        shadow_arg_names.add(input_arg_name)
+            for output_arg_name in op.output_arg_names:
+                output_arg_names.add(output_arg_name)
+
+        # step 2: add `shadow_output` op
+        shadow_arg_names_for_suffixed_programs = (
+            output_arg_names & suffixed_shadow_arg_names
+        )
+        for shadow_arg_name in shadow_arg_names_for_suffixed_programs:
+            block.append_op(
+                type="shadow_output",
+                inputs={"x": shadow_arg_name},
+                outputs={"out": shadow_arg_name},  # unused
+                attrs={"name": shadow_arg_name},
             )
 
-    return type_to_program
+        # step3: add `data` op
+        for shadow_arg_name in shadow_arg_names:
+            shadow_var = block.var(shadow_arg_name)
+            block._prepend_op(
+                type="data",
+                outputs={"out": shadow_arg_name},
+                attrs={
+                    "shape": shadow_var.shape,
+                    "dtype": shadow_var.dtype,
+                    "place": 2,  # GPUPlace
+                    "name": shadow_arg_name,
+                },
+            )
+
+        sub_program._sync_with_cpp()
+
+        # step4: update suffixed_shadow_arg_names
+        suffixed_shadow_arg_names -= shadow_arg_names_for_suffixed_programs
+        suffixed_shadow_arg_names |= shadow_arg_names
 
 
 def _create_param(dst_block, src_var):
diff --git a/python/paddle/distributed/passes/pipeline_pass_base.py b/python/paddle/distributed/passes/pipeline_pass_base.py
@@ -19,7 +19,7 @@
 
 from ..utils.log_utils import get_logger
 from .pass_base import PassBase
-from .pass_utils import set_skip_gc_vars
+from .pass_utils import set_skip_gc_vars, shadow_var_between_sub_programs
 
 logger = get_logger(logging.INFO)
 
@@ -58,24 +58,30 @@ def _apply_single_impl(self, main_program, startup_program, context):
         to implement two interfaces above, 'create_job_list' and 'partial_programs'.
         """
         job_types, sub_programs = self._partial_programs(main_program)
+
+        enable_pir_in_executor = paddle.framework.get_flags(
+            "FLAGS_enable_pir_in_executor"
+        )['FLAGS_enable_pir_in_executor']
+        if enable_pir_in_executor:
+            shadow_var_between_sub_programs(sub_programs)
+
         for i in range(len(job_types)):
             logger.debug(
                 f"sub_program type: {job_types[i]}, sum_program:\n{sub_programs[i]}"
             )
-        jobs = self._create_job_list()
 
+        jobs = self._create_job_list()
         type_to_program = set_skip_gc_vars(
             self.get_attr("num_micro_batches"), job_types, sub_programs, jobs
         )
 
         for type in type_to_program.keys():
-            if paddle.framework.get_flags("FLAGS_enable_pir_in_executor")[
-                'FLAGS_enable_pir_in_executor'
-            ]:
+            if enable_pir_in_executor:
                 type_to_program[type] = paddle.pir.translate_to_pir(
                     type_to_program[type].desc
                 )
             else:
                 type_to_program[type] = type_to_program[type].desc
+
         plan = core.Plan(jobs, type_to_program)
         context.set_attr("plan", plan)
diff --git a/tools/count_api_without_core_ops.py b/tools/count_api_without_core_ops.py
@@ -27,6 +27,7 @@
     "paddle.base.LoDTensor.set",  # Do not know why it should be omitted
     "paddle.base.io.ComposeNotAligned",
     "paddle.base.io.ComposeNotAligned.__init__",
+    "paddle.distributed.passes.pass_utils.shadow_var_between_sub_programs",  # append shadow_output and data op in this function
 ]
 
 
@@ -130,7 +131,6 @@ def is_primitive(instance):
 
 ErrorSet = set()
 IdSet = set()
-skiplist = []
 visited_modules = set()
 
 
@@ -153,8 +153,6 @@ def visit_all_module(mod, func):
         if member_name.startswith('_'):
             continue
         cur_name = mod_name + '.' + member_name
-        if cur_name in skiplist:
-            continue
         try:
             instance = getattr(mod, member_name)
             if inspect.ismodule(instance):
@@ -166,7 +164,7 @@ def visit_all_module(mod, func):
                 IdSet.add(instance_id)
                 visit_member(mod.__name__, instance, func)
         except:
-            if cur_name not in ErrorSet and cur_name not in skiplist:
+            if cur_name not in ErrorSet:
                 ErrorSet.add(cur_name)
 
 

Original file line number	Diff line number	Diff line change
`@@ -359,10 +359,6 @@ std::string InstructionBase::DebugStringEx(`
`359`	`359`	`if (row_size >= 0) {`
`360`	`360`	`ss << "row_size=" << row_size << ";";`
`361`	`361`	`}`
`362`		`- double ele_sum = GetDenseTensorEleSum(*scope, var_name);`
`363`		`- if (!std::isnan(ele_sum)) {`
`364`		`- ss << "ele_sum=" << ele_sum << ";";`
`365`		`- }`
`366`	`362`	`}`
`367`	`363`	`}`
`368`	`364`	`++it;`
`@@ -390,10 +386,6 @@ std::string InstructionBase::DebugStringEx(`
`390`	`386`	`if (row_size >= 0) {`
`391`	`387`	`ss << "row_size=" << row_size << ";";`
`392`	`388`	`}`
`393`		`- double ele_sum = GetDenseTensorEleSum(*scope, var_name);`
`394`		`- if (!std::isnan(ele_sum)) {`
`395`		`- ss << "ele_sum=" << ele_sum << ";";`
`396`		`- }`
`397`	`389`	`}`
`398`	`390`	`}`
`399`	`391`	`++it;`