Support bw invoke fw (#50260)

* support bw invoke fw * fix scale in static_backward.yaml * fix the bug in tensorrt/convert * move 'scale','sign' into ops.yaml * add scale_grad of scale in op_compat.yaml * change generated_static_op in CMakeLists.txt
PaddlePaddle · Feb 21, 2023 · d884573 · d884573
1 parent 9af23f1
commit d884573
Show file tree

Hide file tree

Showing 20 changed files with 130 additions and 327 deletions.
diff --git a/paddle/fluid/eager/CMakeLists.txt b/paddle/fluid/eager/CMakeLists.txt
@@ -77,6 +77,6 @@ cc_library(
        op_registry
        variable_helper
        memcpy
-       scale_op
+       generated_op
        autograd_meta
        hook_utils)
diff --git a/paddle/fluid/eager/tests/performance_tests/CMakeLists.txt b/paddle/fluid/eager/tests/performance_tests/CMakeLists.txt
@@ -7,7 +7,7 @@ if(NOT (NOT WITH_PYTHON AND ON_INFER))
          ${generated_deps}
          eager_scale
          scale_node
-         scale_op
+         generated_op
          matmul_v2_op
          dygraph_function
          eager_prim_api)

diff --git a/paddle/fluid/framework/CMakeLists.txt b/paddle/fluid/framework/CMakeLists.txt
@@ -1051,7 +1051,7 @@ if(WITH_PSCORE)
       heter_pipeline_trainer_test
       SRCS heter_pipeline_trainer_test.cc
       DEPS conditional_block_op
-           scale_op
+           generated_op
            heter_listen_and_serv_op
            executor
            heter_server
@@ -1068,7 +1068,7 @@ if(WITH_PSCORE)
       heter_pipeline_trainer_test
       SRCS heter_pipeline_trainer_test.cc
       DEPS conditional_block_op
-           scale_op
+           generated_op
            heter_listen_and_serv_op
            executor
            heter_server

diff --git a/paddle/fluid/framework/ir/memory_optimize_pass/CMakeLists.txt b/paddle/fluid/framework/ir/memory_optimize_pass/CMakeLists.txt
@@ -76,5 +76,5 @@ cc_library(
 cc_test(
   test_reference_count_pass_last_lived_ops
   SRCS test_reference_count_pass_last_lived_ops.cc
-  DEPS parallel_executor elementwise_mul_op elementwise_add_op scale_op
+  DEPS parallel_executor elementwise_mul_op elementwise_add_op generated_op
        eigen_function)
diff --git a/paddle/fluid/jit/CMakeLists.txt b/paddle/fluid/jit/CMakeLists.txt
@@ -65,7 +65,7 @@ if(WITH_TESTING AND NOT WIN32)
       reduce_mean_op
       feed_op
       fetch_op
-      scale_op
+      generated_op
       transfer_layout_op
       jit_layer)
   cc_test(

diff --git a/paddle/fluid/operators/generator/CMakeLists.txt b/paddle/fluid/operators/generator/CMakeLists.txt
@@ -11,6 +11,8 @@ set(legacy_bw_op_yaml_file
 set(sparse_op_yaml_file ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/sparse_ops.yaml)
 set(sparse_bw_op_yaml_file
     ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/sparse_backward.yaml)
+set(static_bw_op_yaml_file
+    ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/static_backward.yaml)
 
 if(NOT PYTHONINTERP_FOUND)
   find_package(PythonInterp REQUIRED)
@@ -66,6 +68,9 @@ execute_process(
   COMMAND
     ${PYTHON_EXECUTABLE} parse_op.py --op_yaml_path ${sparse_bw_op_yaml_file}
     --output_path ./parsed_ops/sparse_backward.parsed.yaml --backward
+  COMMAND
+    ${PYTHON_EXECUTABLE} parse_op.py --op_yaml_path ${static_bw_op_yaml_file}
+    --output_path ./parsed_ops/static_backward.parsed.yaml --backward
     RESULTS_VARIABLE _results)
 foreach(_result in ${_results})
   if(${_result})
@@ -82,14 +87,24 @@ execute_process(
   COMMAND
     ${PYTHON_EXECUTABLE} cross_validate.py --forward_yaml_paths
     ./parsed_ops/ops.parsed.yaml ./parsed_ops/legacy_ops.parsed.yaml
-    ./parsed_ops/static_ops.parsed.yaml --backward_yaml_paths
-    ./parsed_ops/backward_ops.parsed.yaml
+    --backward_yaml_paths ./parsed_ops/backward_ops.parsed.yaml
     ./parsed_ops/legacy_backward_ops.parsed.yaml
   RESULT_VARIABLE _result)
 if(${_result})
   message(FATAL_ERROR "ops validation failed, exiting.")
 endif()
 
+execute_process(
+  WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/paddle/fluid/operators/generator
+  COMMAND
+    ${PYTHON_EXECUTABLE} cross_validate.py --forward_yaml_paths
+    ./parsed_ops/static_ops.parsed.yaml --backward_yaml_paths
+    ./parsed_ops/static_backward.parsed.yaml
+  RESULT_VARIABLE _result)
+if(${_result})
+  message(FATAL_ERROR "static ops validation failed, exiting.")
+endif()
+
 execute_process(
   WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/paddle/fluid/operators/generator
   COMMAND
@@ -124,8 +139,9 @@ endif()
 execute_process(
   WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/paddle/fluid/operators/generator
   COMMAND
-    ${PYTHON_EXECUTABLE} generate_static_op.py --ops_yaml_path
-    ./parsed_ops/static_ops.parsed.yaml --op_version_yaml_path
+    ${PYTHON_EXECUTABLE} generate_op.py --ops_yaml_path
+    ./parsed_ops/static_ops.parsed.yaml --backward_yaml_path
+    ./parsed_ops/static_backward.parsed.yaml --op_version_yaml_path
     ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/op_version.yaml
     --op_compat_yaml_path ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/op_compat.yaml
     --output_op_path "${generated_static_op_path}.tmp" --output_arg_map_path

diff --git a/paddle/fluid/operators/generator/generate_op.py b/paddle/fluid/operators/generator/generate_op.py
@@ -407,6 +407,7 @@ def process_invoke_op(forward_op_dict, backward_op_dict):
             invoke_op = bw_op['invoke']['func']
             args_list = bw_op['invoke']['args']
             args_index = 0
+            # backward invoke forward
             if invoke_op in forward_op_dict:
                 reuse_op = forward_op_dict[invoke_op]
                 bw_op['invoke']['func'] = reuse_op['op_name']
@@ -460,17 +461,16 @@ def parse_drop_empty_grad(op_fluid_list: list, bw_op_dict: dict):
                 for bw_name in op_op['backward'].split(',')
             ]
             for bw_name in bw_names:
-                assert (
-                    bw_name in bw_op_dict
-                ), f"backward {bw_name} is not existed"
-                for out_grad in op_op['drop_empty_grad']:
-                    assert (
-                        out_grad in bw_op_dict[bw_name]['output_dict']
-                    ), f'''
-                         {bw_name} with {out_grad} is not existed in output_dict '''
-                    bw_op_dict[bw_name]['output_dict'][out_grad][
-                        'drop_empty_grad'
-                    ] = False
+                # static_ops.yaml and ops.yaml use the common op_compat.yaml
+                if bw_name in bw_op_dict:
+                    for out_grad in op_op['drop_empty_grad']:
+                        assert (
+                            out_grad in bw_op_dict[bw_name]['output_dict']
+                        ), f'''
+                            {bw_name} with {out_grad} is not existed in output_dict '''
+                        bw_op_dict[bw_name]['output_dict'][out_grad][
+                            'drop_empty_grad'
+                        ] = False
 
 
 def main(
@@ -493,7 +493,8 @@ def main(
         op_versions = yaml.safe_load(f)
     # add op version info into op
     for op_version in op_versions:
-        forward_op_dict[op_version['op']]['version'] = op_version['version']
+        if op_version['op'] in forward_op_dict:
+            forward_op_dict[op_version['op']]['version'] = op_version['version']
 
     with open(op_compat_yaml_path, "rt") as f:
         op_fluid_map_list = yaml.safe_load(f)

diff --git a/paddle/fluid/operators/generator/generate_static_op.py b/paddle/fluid/operators/generator/generate_static_op.py
@@ -81,6 +81,7 @@ def restruct_io(op):
 
 def main(
     ops_yaml_path,
+    backward_yaml_path,
     op_compat_yaml_path,
     op_version_yaml_path,
     output_op_path,
@@ -91,6 +92,11 @@ def main(
         ops = [restruct_io(op) for op in ops]
     forward_op_dict = to_named_dict(ops)
 
+    with open(backward_yaml_path, "rt") as f:
+        backward_ops = yaml.safe_load(f)
+        backward_ops = [restruct_io(op) for op in backward_ops]
+    backward_op_dict = to_named_dict(backward_ops)
+
     with open(op_version_yaml_path, "rt") as f:
         op_versions = yaml.safe_load(f)
 
@@ -139,6 +145,11 @@ def main(
     parser.add_argument(
         '--ops_yaml_path', type=str, help="parsed static ops yaml file."
     )
+    parser.add_argument(
+        '--backward_yaml_path',
+        type=str,
+        help="parsed static backward ops yaml file.",
+    )
     parser.add_argument(
         '--op_compat_yaml_path', type=str, help="ops args compat yaml file."
     )
@@ -157,6 +168,7 @@ def main(
     args = parser.parse_args()
     main(
         args.ops_yaml_path,
+        args.backward_yaml_path,
         args.op_compat_yaml_path,
         args.op_version_yaml_path,
         args.output_op_path,

diff --git a/paddle/fluid/operators/generator/templates/operator_utils.c.j2 b/paddle/fluid/operators/generator/templates/operator_utils.c.j2
@@ -522,6 +522,31 @@ class {{name | to_pascal_case}}OpMaker : public framework::SingleGradOpMaker<T>
       true)}});
   {% endfor %}
 
+  {% for attr in invoke_op["attrs"] %}
+    {% set attr_name = attr["fluid_name"] %}
+    {% set fw_attrs = forward_op["attrs"] %}
+    {% if attr_name in forward_attr_names %}
+      {# invoke_op's attrs and fw_attr's attrs must be the same#}
+      {% set fw_attr = fw_attrs[loop.index0] %}
+      {% if fw_attr["typename"] == "IntArray" %}
+        {% if 'tensor_name' in attr or 'manual_flag' not in attr %}
+    if (this->HasInput("{{fw_attr | to_int_array_tensor_name}}")) {
+      grad_op->SetInput("{{fw_attr | to_int_array_tensor_name}}", this->Input("{{fw_attr | to_int_array_tensor_name}}"));
+    }
+        {% endif %}
+        {% if 'tensors_name' in fw_attr or 'manual_flag' not in fw_attr %}
+    if (this->HasInput("{{fw_attr | to_int_array_tensors_name}}")) {
+      grad_op->SetInput("{{fw_attr | to_int_array_tensors_name}}", this->Input("{{fw_attr | to_int_array_tensors_name}}"));
+    }
+        {% endif %}
+      {% elif fw_attr["typename"] == "Scalar" %}
+    if (this->HasInput("{{fw_attr | to_scalar_tensor_name}}")) {
+      grad_op->SetInput("{{fw_attr | to_scalar_tensor_name}}", this->Input("{{fw_attr | to_scalar_tensor_name}}"));
+    }
+      {% endif %}
+    {% endif %}
+  {% endfor %}
+
   {% for attr in invoke_op["attrs"] %}
     grad_op->SetAttr("{{attr["fluid_name"]}}", {{attr["value"]}});
   {% endfor %}

diff --git a/paddle/fluid/operators/pscore/CMakeLists.txt b/paddle/fluid/operators/pscore/CMakeLists.txt
@@ -86,7 +86,7 @@ cc_test_old(
   executor
   scope
   proto_desc
-  scale_op
+  generated_op
   eigen_function)
 
 set_source_files_properties(
@@ -100,7 +100,7 @@ cc_test_old(
   executor
   scope
   proto_desc
-  scale_op
+  generated_op
   send_and_recv_op
   ${RPC_DEPS}
   ${DISTRIBUTE_DEPS}
@@ -117,7 +117,7 @@ cc_test_old(
   executor
   scope
   proto_desc
-  scale_op
+  generated_op
   send_and_recv_op
   ${RPC_DEPS}
   ${DISTRIBUTE_DEPS}
@@ -134,14 +134,14 @@ cc_test_old(
   executor
   scope
   proto_desc
-  scale_op
+  generated_op
   heter_listen_and_serv_op
   ${RPC_DEPS}
   ${DISTRIBUTE_DEPS}
   eigen_function)
 
 #set_source_files_properties(heter_cloud_comm_cpu_test.cc PROPERTIES COMPILE_FLAGS ${DISTRIBUTE_COMPILE_FLAGS})
-#cc_test(heter_cloud_comm_cpu_test SRCS heter_cloud_comm_cpu_test.cc DEPS executor scope proto_desc scale_op heter_listen_and_serv_op ${RPC_DEPS} ${DISTRIBUTE_DEPS} eigen_function)
+#cc_test(heter_cloud_comm_cpu_test SRCS heter_cloud_comm_cpu_test.cc DEPS executor scope proto_desc generated_static_op heter_listen_and_serv_op ${RPC_DEPS} ${DISTRIBUTE_DEPS} eigen_function)
 
 set_source_files_properties(
   switch_server_test.cc PROPERTIES COMPILE_FLAGS ${DISTRIBUTE_COMPILE_FLAGS})
@@ -153,7 +153,7 @@ cc_binary(
   executor
   scope
   proto_desc
-  scale_op
+  generated_op
   heter_listen_and_serv_op
   ${RPC_DEPS}
   ${DISTRIBUTE_DEPS}

diff --git a/paddle/fluid/operators/scale_op.cc b/paddle/fluid/operators/scale_op.cc