support generating code of opmaker for backward op invoke forward op (#46912)

paddle/fluid/operators/flip_op.cc

-/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
-
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "paddle/fluid/framework/infershape_utils.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/framework/op_version_registry.h"
-#include "paddle/phi/core/infermeta_utils.h"
-#include "paddle/phi/infermeta/unary.h"
-
-namespace paddle {
-namespace operators {
-
-using framework::OpKernelType;
-
-class FlipOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
-  framework::OpKernelType GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const {
-    framework::LibraryType library = framework::LibraryType::kPlain;
-    framework::DataLayout layout = framework::DataLayout::kAnyLayout;
-    int customized_type_value =
-        framework::OpKernelType::kDefaultCustomizedTypeValue;
-    auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
-    return framework::OpKernelType(input_data_type,
-                                   ctx.GetPlace(),
-                                   layout,
-                                   library,
-                                   customized_type_value);
-  }
-};
-
-class FlipOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput("X", "(Tensor), The input tensor of flip op.");
-    AddOutput("Out", "(Tensor), The output tensor of flip op.");
-    AddAttr<std::vector<int>>("axis", "The axes to flip on.");
-    AddComment(R"DOC(
-          Flip Operator.
-          Reverse the order of a n-D tensor along given axis in axes.
-      )DOC");
-  }
-};
-
-class FlipOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
- protected:
-  std::unordered_map<std::string, std::string>& GetInputOutputWithSameType()
-      const override {
-    static std::unordered_map<std::string, std::string> m{{"X", /*->*/ "Out"}};
-    return m;
-  }
-};
-
-template <typename T>
-class FlipOpGradMaker : public framework::SingleGradOpMaker<T> {
- public:
-  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
-
- protected:
-  void Apply(GradOpPtr<T> retv) const override {
-    retv->SetType("flip");
-    retv->SetInput("X", this->OutputGrad("Out"));
-    retv->SetOutput("Out", this->InputGrad("X"));
-    retv->SetAttrMap(this->Attrs());
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-namespace plat = paddle::platform;
-DECLARE_INFER_SHAPE_FUNCTOR(flip,
-                            FlipInferShapeFunctor,
-                            PD_INFER_META(phi::FlipInferMeta));
-REGISTER_OPERATOR(flip,
-                  ops::FlipOp,
-                  ops::FlipOpMaker,
-                  ops::FlipOpInferVarType,
-                  ops::FlipOpGradMaker<paddle::framework::OpDesc>,
-                  ops::FlipOpGradMaker<paddle::imperative::OpBase>,
-                  FlipInferShapeFunctor);
-
-/* ==========================  register checkpoint ===========================*/
-REGISTER_OP_VERSION(flip).AddCheckpoint(
-    R"ROC(Upgrade flip, add new attr [axis] and delete attr [dims].)ROC",
-    paddle::framework::compatible::OpVersionDesc()
-        .NewAttr("axis",
-                 "The added attr 'axis' doesn't set default value.",
-                 paddle::none)
-        .DeleteAttr("dims", "The attr 'dims' is deleted."));

paddle/fluid/operators/unity_build_rule.cmake

@@ -108,7 +108,6 @@ register_unity_group(
 register_unity_group(
   cc
   flatten_op.cc
-  flip_op.cc
   fsp_op.cc
   gather_nd_op.cc
   gather_op.cc
@@ -423,7 +422,6 @@ register_unity_group(cu expand_v2_op.cu fake_dequantize_op.cu
                      fill_any_like_op.cu)
 register_unity_group(
   cu
-  flip_op.cu
   fsp_op.cu
   gather_nd_op.cu
   gather_op.cu

paddle/phi/api/lib/CMakeLists.txt

@@ -171,9 +171,9 @@ create or remove auto-geneated argument mappings: ${generated_argument_mapping_p
 execute_process(
   WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml
   COMMAND
-    ${PYTHON_EXECUTABLE} generator/generate_op.py --api_yaml_path
-    ./parsed_apis/api.parsed.yaml --backward_api_yaml_path
-    ./parsed_apis/backward_api.parsed.yaml --api_version_yaml_path
+    ${PYTHON_EXECUTABLE} generator/generate_op.py --ops_yaml_path
+    ./parsed_apis/api.parsed.yaml --backward_yaml_path
+    ./parsed_apis/backward_api.parsed.yaml --op_version_yaml_path
     op_version.yaml --op_compat_yaml_path op_compat.yaml --output_op_path
     "${generated_op_path}.tmp" --output_arg_map_path
     "${generated_argument_mapping_path}.tmp"

paddle/phi/api/yaml/backward.yaml

@@ -147,6 +147,12 @@
     data_type: out_grad
   no_need_buffer: x
 
+- backward_op : flip_grad
+  forward : flip (Tensor x, int[] axis) -> Tensor(out)
+  args : (Tensor out_grad, int[] axis)
+  output : Tensor(x_grad)
+  invoke : flip(out_grad, axis)
+
 - backward_op : graph_send_uv_grad
   forward : graph_send_uv (Tensor x, Tensor y, Tensor src_index, Tensor dst_index, str message_op = "ADD") -> Tensor(out)
   args: (Tensor x, Tensor y, Tensor src_index, Tensor dst_index, Tensor out_grad, str message_op = "ADD")

paddle/phi/api/yaml/generator/generate_op.py

@@ -143,6 +143,14 @@ def replace_compat_name(api_op_map, forward_api_dict, backward_api_dict):
                                                            [:-5]] + '_grad'
                     args_item['name'] = args_map[args_item['name']]
 
+            if 'invoke' in backward_api_item:
+                backward_api_item['invoke']['args'] = [
+                    args_map[param.strip()]
+                    if param.strip() in args_map else param.strip()
+                    for param in backward_api_item['invoke']['args'].split(',')
+                ]
+                continue
+
             backward_api_item['infer_meta']['param'] = [
                 args_map[param] if param in args_map else param
                 for param in backward_api_item['infer_meta']['param']
@@ -173,9 +181,9 @@ def replace_compat_name(api_op_map, forward_api_dict, backward_api_dict):
                 ]
 
 
-def main(api_yaml_path, backward_yaml_path, op_compat_yaml_path,
-         api_version_yaml_path, output_op_path, output_arg_map_path):
-    with open(api_yaml_path, "rt") as f:
+def main(ops_yaml_path, backward_yaml_path, op_compat_yaml_path,
+         op_version_yaml_path, output_op_path, output_arg_map_path):
+    with open(ops_yaml_path, "rt") as f:
         apis = yaml.safe_load(f)
         apis = [restruct_io(api) for api in apis]
     forward_api_dict = to_named_dict(apis)
@@ -185,7 +193,7 @@ def main(api_yaml_path, backward_yaml_path, op_compat_yaml_path,
         backward_apis = [restruct_io(api) for api in backward_apis]
     backward_api_dict = to_named_dict(backward_apis)
 
-    with open(api_version_yaml_path, "rt") as f:
+    with open(op_version_yaml_path, "rt") as f:
         api_versions = yaml.safe_load(f)
     # add api version info into api
     for api_version in api_versions:
@@ -201,6 +209,45 @@ def main(api_yaml_path, backward_yaml_path, op_compat_yaml_path,
 
     replace_compat_name(api_op_map, forward_api_dict, backward_api_dict)
 
+    # prepare for invoke case
+    for bw_name, bw_api in backward_api_dict.items():
+        if 'invoke' in bw_api:
+            invoke_op = bw_api['invoke']['func']
+            args_list = bw_api['invoke']['args']
+            args_index = 0
+            if invoke_op in forward_api_dict.keys():
+                reuse_op = forward_api_dict[invoke_op]
+                bw_api['invoke']['inputs'] = []
+                bw_api['invoke']['attrs'] = []
+                bw_api['invoke']['outputs'] = []
+                for input_item in reuse_op['inputs']:
+                    bw_api['invoke']['inputs'].append({
+                        'name':
+                        input_item['name'],
+                        'value':
+                        args_list[args_index]
+                    })
+                    args_index = args_index + 1
+                for attr in reuse_op['attrs']:
+                    if args_index < len(args_list):
+                        attr_value = f"this->GetAttr(\"{args_list[args_index]}\")" if args_list[
+                            args_index] in bw_api['attr_dict'] else args_list[
+                                args_index]
+                        bw_api['invoke']['attrs'].append({
+                            'name': attr['name'],
+                            'value': attr_value
+                        })
+                        args_index = args_index + 1
+                    else:
+                        break
+                for idx, output_item in enumerate(reuse_op['outputs']):
+                    bw_api['invoke']['outputs'].append({
+                        'name':
+                        output_item['name'],
+                        'value':
+                        bw_api['outputs'][idx]['name']
+                    })
+
     # fill backward field for an api if another api claims it as forward
     for name, backward_api in backward_api_dict.items():
         forward_name = backward_api["forward"]["name"]
@@ -236,18 +283,18 @@ def main(api_yaml_path, backward_yaml_path, op_compat_yaml_path,
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Generate operator file from api yaml.")
-    parser.add_argument('--api_yaml_path',
+    parser.add_argument('--ops_yaml_path',
                         type=str,
-                        help="parsed api yaml file.")
-    parser.add_argument('--backward_api_yaml_path',
+                        help="parsed ops yaml file.")
+    parser.add_argument('--backward_yaml_path',
                         type=str,
-                        help="parsed backward api yaml file.")
+                        help="parsed backward ops yaml file.")
     parser.add_argument('--op_compat_yaml_path',
                         type=str,
-                        help="api args compat yaml file.")
-    parser.add_argument('--api_version_yaml_path',
+                        help="ops args compat yaml file.")
+    parser.add_argument('--op_version_yaml_path',
                         type=str,
-                        help="api version yaml file.")
+                        help="ops version yaml file.")
     parser.add_argument("--output_op_path",
                         type=str,
                         help="path to save generated operators.")
@@ -257,6 +304,6 @@ if __name__ == "__main__":
         help="path to save generated argument mapping functions.")
 
     args = parser.parse_args()
-    main(args.api_yaml_path, args.backward_api_yaml_path,
-         args.op_compat_yaml_path, args.api_version_yaml_path,
-         args.output_op_path, args.output_arg_map_path)
+    main(args.ops_yaml_path, args.backward_yaml_path, args.op_compat_yaml_path,
+         args.op_version_yaml_path, args.output_op_path,
+         args.output_arg_map_path)

paddle/phi/api/yaml/generator/templates/op.c.j2

-{% from "operator_utils.c.j2" import op_maker, backward_op_maker, operator, register_op_with_components, register_op_version %}
+{% from "operator_utils.c.j2" import op_maker, backward_op_maker, backward_op_reused_maker, operator, register_op_with_components, register_op_version %}
 // this file is generated by paddle/phi/api/yaml/generator/generate_op.py, do not edit.
 #include <string>
 #include "paddle/fluid/framework/infershape_utils.h"
@@ -33,6 +33,8 @@ using paddle::framework::GradVarName;
 {{backward_op_maker(api, api_dict[api["forward"]["name"]])}}
 
 {{operator(api)}}
+  {% else %}
+{{backward_op_reused_maker(api, api_dict[api["forward"]["name"]], api["invoke"])}}
   {% endif %}
 {% endfor %}
 }  // namespace operators

paddle/phi/api/yaml/generator/templates/operator_utils.c.j2

@@ -352,6 +352,48 @@ class {{name | to_pascal_case}}OpMaker : public framework::SingleGradOpMaker<T>
 };
 {% endmacro %}
 
+{% macro backward_op_reused_maker(bw_op, forward_op, invoke_op) %}
+  {% set name = bw_op["op_name"] %}
+  {% set forward_input_names = bw_op["forward"]["inputs"] | map(attribute="name") | list %}
+  {% set forward_output_names = bw_op["forward"]["outputs"] | map(attribute="name") | list %}
+  {% set forward_attr_names = bw_op["forward"]["attrs"] | map(attribute="name") | list %}
+  {% set forward_input_orig_names = forward_op["inputs"] | map(attribute="name") | list %}
+  {% set forward_output_orig_names = forward_op["outputs"] | map(attribute="name") | list %}
+  {% set forward_attr_orig_names = forward_op["attrs"] | map(attribute="name") | list %}
+template <typename T>
+class {{name | to_pascal_case}}OpMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+ protected:
+  void Apply(GradOpPtr<T> grad_op) const override {
+    grad_op->SetType("{{invoke_op["func"]}}");
+
+  {% for input in invoke_op["inputs"] %}
+    grad_op->SetInput({{input["name"] | to_opmaker_name}}, this->{{extract_input_from_forward(
+      input["value"],
+      forward_input_names,
+      forward_output_names,
+      forward_input_orig_names,
+      forward_output_orig_names)}});
+  {% endfor %}
+
+  {% for output in invoke_op["outputs"] %}
+    grad_op->SetOutput({{output["name"] | to_opmaker_name}}, this->{{extract_output_from_forward(
+      output["value"],
+      forward_input_names,
+      forward_output_names,
+      forward_input_orig_names,
+      forward_output_orig_names)}});
+  {% endfor %}
+
+  {% for attr in invoke_op["attrs"] %}
+    grad_op->SetAttr("{{attr["name"]}}", {{attr["value"]}});
+  {% endfor %}
+  }
+};
+{% endmacro %}
+
 
 {% macro extract_input_from_forward(name,
   input_names, output_names,

paddle/phi/api/yaml/legacy_backward.yaml

@@ -839,12 +839,6 @@
     layout: out_grad
   inplace : (out_grad -> x_grad)
 
-- backward_op : flip_grad
-  forward : flip (Tensor x, int[] axis) -> Tensor(out)
-  args : (Tensor out_grad, int[] axis)
-  output : Tensor(x_grad)
-  invoke : flip(out_grad, axis)
-
 - backward_op : floor_grad
   forward : floor(Tensor x) -> Tensor(out)
   args : (Tensor out_grad)

paddle/phi/api/yaml/legacy_ops.yaml

@@ -962,15 +962,6 @@
   intermediate : xshape
   backward : flatten_grad
 
-- op : flip
-  args : (Tensor x, int[] axis)
-  output : Tensor
-  infer_meta :
-    func : FlipInferMeta
-  kernel :
-    func : flip
-  backward : flip_grad
-
 - op : floor
   args : (Tensor x)
   output : Tensor(out)

paddle/phi/api/yaml/op_compat.yaml

@@ -324,6 +324,12 @@
   inputs: {x: X}
   outputs: {out: Out}
 
+- op : flip
+  inputs :
+    x : X
+  outputs :
+    out : Out
+
 - op : floor
   backward : floor_grad
   extra :

paddle/phi/api/yaml/op_version.yaml

+- op : flip
+  version :
+    - checkpoint : Upgrade flip, add new attr [axis] and delete attr [dims]
+      action :
+        - add_attr : axis
+          comment : The added attr 'axis' doesn't set default value
+          default : paddle::none
+        - delete_attr : dims
+          comment : The attr 'dims' is deleted.
+
 - op : trace
   version :
     - checkpoint : Upgrade trace add a new attribute [axis2]

paddle/phi/api/yaml/ops.yaml

@@ -199,3 +199,12 @@
   kernel :
     func : trunc
   backward : trunc_grad
+
+- op : flip
+  args : (Tensor x, int[] axis)
+  output : Tensor (out)
+  infer_meta :
+    func : FlipInferMeta
+  kernel :
+    func : flip
+  backward : flip_grad