backward_api call backward_api_base (#71205)

xiaoguoguo626807 · web-flow · commit 274b32ed4dbc · 2025-02-21T14:07:47.000+08:00
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -38,6 +38,7 @@
     ParseYamlForwardFromBackward,
     ParseYamlInplaceInfo,
     ReadBwdFile,
+    ReadFwdFile,
     RemoveConstAndReference,
     core_ops_args_info,
     core_ops_args_type_info,
@@ -1829,7 +1830,34 @@ def GenerateForwardDefinitionAndDeclaration(self, is_inplaced, grad_flag):
         api_out_type = "auto"
         if is_inplaced and len(forward_outputs_position_map) == 1:
             api_out_type = "auto&"
-        forward_call_str = f"{indent}{api_out_type} api_result = paddle::experimental::{namespace}{function_name}({inputs_call_args_str});"
+
+        is_invoke_api = IsInvokeForwardApi(
+            self.forward_api_contents, self.forward_apis_dict
+        )
+        if is_invoke_api:
+            invoke_api_name = (
+                self.forward_api_contents['invoke'].split('(')[0].strip()
+            )
+            invoke_api_str = self.forward_api_contents['invoke'].replace(
+                invoke_api_name,
+                GetDygraphForwardFunctionName(invoke_api_name),
+                1,
+            )
+
+            if self.is_forward_only:
+                forward_call_str = f"""// Invoke API Call
+  {indent}{api_out_type} api_result = paddle::experimental::{self.namespace}{self.forward_api_contents['invoke']};"""
+            else:
+                forward_call_str = f"""
+  // Invoke API Call
+  if (require_any_grad) {{
+  {indent}{api_out_type} api_result = {invoke_api_str};
+  }} else {{
+  {indent}{api_out_type} api_result = paddle::experimental::{self.namespace}{self.forward_api_contents['invoke']};
+  }}
+"""
+        else:
+            forward_call_str = f"{indent}{api_out_type} api_result = paddle::experimental::{namespace}{function_name}({inputs_call_args_str});"
         num_outputs = len(forward_outputs_position_map) - len(
             intermediate_outputs
         )
@@ -3145,7 +3173,13 @@ def __init__(
         GeneratorBase.__init__(self, api_yaml_path, fw_ops)
 
         self.backward_yaml_path = backward_yaml_path
-        self.bw_ops = bw_ops
+        if bw_ops is None:
+            if self.backward_yaml_path is not None:
+                self.backward_api_list = ReadFwdFile(self.backward_yaml_path)
+            else:
+                self.backward_api_list = []
+        else:
+            self.backward_api_list = bw_ops
         self.grad_api_dict = {}
 
         self.forward_declaration_str = ""
@@ -3161,12 +3195,13 @@ def CollectIsForwardOnly(self, forward_api_contents):
 
     def ParseYamlContents(self):
         self.ParseForwardYamlContents()
-
         backward_yaml_path = self.backward_yaml_path
 
         # string api is forward_only, no backward_yaml respectively
         if backward_yaml_path is not None:
-            self.grad_api_dict = ReadBwdFile(backward_yaml_path, self.bw_ops)
+            self.grad_api_dict = ReadBwdFile(
+                backward_yaml_path, self.backward_api_list
+            )
 
     def GetBackwardAPIContents(self, forward_api_contents):
         grad_api_dict = self.grad_api_dict
@@ -3187,13 +3222,19 @@ def GenerateCode(self, grad_flag=False):
             op_string = 'backward_op'
         else:
             op_string = 'op'
+
         forward_api_list = self.forward_api_list
         forward_apis_dict = {}
         for api_item in forward_api_list:
-            forward_apis_dict[api_item[op_string]] = api_item
+            forward_apis_dict[
+                api_item['op' if 'op' in api_item else 'backward_op']
+            ] = api_item
         namespace = self.namespace
 
-        for forward_api_contents in forward_api_list:
+        true_forward_api_list = (
+            self.backward_api_list if grad_flag else self.forward_api_list
+        )
+        for forward_api_contents in true_forward_api_list:
             if forward_api_contents[op_string] in black_ops_list:
                 continue
             if op_string == 'backward_op' and (
@@ -3420,7 +3461,7 @@ def GenerateForwardHFile(filepath, forward_function_declaration_str, grad_flag):
         backward_yaml_path = backward_yaml_paths[i]
         if backward_yaml_path.endswith('/backward.yaml'):
             generator_grad = DygraphForwardAndNodesGenerator(
-                backward_yaml_paths[i], backward_yaml_paths[i], all_bw, all_bw
+                backward_yaml_path, backward_yaml_path, all_ops, all_bw
             )
         elif backward_yaml_path.endswith('/dygraph_backward.yaml'):
             continue
diff --git a/paddle/phi/api/generator/api_base.py b/paddle/phi/api/generator/api_base.py
@@ -127,6 +127,73 @@ def get_input_tensor_args(self, inplace_flag=False):
                 input_args.append(self.inputs['input_info'][name] + ' ' + name)
         return input_args
 
+    # funcs backward_api.h will use
+    def get_grad_outputs_define(self, inplace_flag=False):
+        define_string = ""
+        for i, out_type in enumerate(self.outputs['types']):
+            out_name = self.outputs['names'][i].split('@')[0]
+            if out_type == "std::vector<Tensor>":
+                if inplace_flag and out_name in self.inplace_map:
+                    out_name = self.inplace_map[out_name]
+                    define_string = " "
+                else:
+                    define_string += "    " + out_type + " " + out_name + ";\n"
+
+                vec_tensor_string = f"""    std::vector<Tensor*> {out_name}_x;
+    for (size_t i = 0; i < {out_name}.size(); i++){{
+        {out_name}_x.push_back(&({out_name}[i]));
+    }}"""
+                define_string += vec_tensor_string
+            else:
+                if inplace_flag and out_name in self.inplace_map:
+                    continue
+                define_string += "    " + out_type + " " + out_name + ";\n"
+        return define_string
+
+    def get_grad_api_call_args(self, inplace_flag):
+        args = []
+        for name in self.inputs['names']:
+            name = name.split('@')[0]
+            args.append(name)
+
+        for name in self.attrs['names']:
+            args.append(name)
+
+        for i, name in enumerate(self.outputs['names']):
+            name = name.split('@')[0]
+            out_type = self.outputs['types'][i]
+            if out_type == "std::vector<Tensor>":
+                if inplace_flag and name in self.inplace_map:
+                    name = self.inplace_map[name]
+
+                out_string = name + "_x"
+            else:
+                if inplace_flag and name in self.inplace_map:
+                    name = self.inplace_map[name]
+                    if name in self.optional_vars:
+                        out_string = name + ".get_ptr()"
+                    else:
+                        out_string = "&" + name
+                else:
+                    out_string = "&" + name
+
+            args.append(out_string)
+        return ", ".join(args)
+
+    def get_grad_output(self, inplace_flag):
+        args = []
+        for i, name in enumerate(self.outputs['names']):
+            name = name.split('@')[0]
+            if inplace_flag and name in self.inplace_map:
+                args.append("std::ref(" + self.inplace_map[name] + ")")
+            else:
+                args.append(name)
+
+        if len(args) == 1:
+            return args[0]
+        else:
+            return f"""std::make_tuple({", ".join(args)})"""
+
     def get_declare_args(self, inplace_flag=False):
         declare_args = self.get_input_tensor_args(inplace_flag)
         for name in self.attrs['names']:
diff --git a/paddle/phi/api/generator/api_gen.py b/paddle/phi/api/generator/api_gen.py
@@ -37,19 +37,6 @@
     "std::vector<Tensor>": "paddle::optional<std::vector<Tensor>>",
 }
 
-manual_impl = '''
-
-PADDLE_API Tensor embedding_grad(const Tensor& x, const Tensor& weight, const Tensor& out_grad, int64_t padding_idx, bool sparse) {
-  Tensor weight_grad;
-  embedding_grad_impl(x, weight, out_grad, padding_idx, sparse, &weight_grad);
-  return weight_grad;
-}
-
-PADDLE_API std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>> cudnn_lstm_grad(const Tensor& x, const Tensor& init_h, const Tensor& init_c, const paddle::optional<std::vector<Tensor>>& weight_list, const paddle::optional<Tensor>& sequence_length, const Tensor& out, const Tensor& reserve, const Tensor& state_out, const Tensor& out_grad, const Tensor& last_h_grad, const Tensor& last_c_grad, float dropout_prob, bool is_bidirec, int hidden_size, int num_layers, bool is_test, int seed) {
-  return cudnn_lstm_grad_impl(x, init_h, init_c, weight_list, sequence_length, out, reserve, state_out, out_grad, last_h_grad, last_c_grad, dropout_prob, is_bidirec, hidden_size, num_layers, is_test, seed) ;
-}
-'''
-
 
 class ForwardAPI(BaseAPI):
     def __init__(self, api_item_yaml):
@@ -428,6 +415,72 @@ def reset_view_after_fallback(
         return remap_code
 
 
+class BackwardAPI(ForwardAPI):
+
+    def gene_base_api_code(self, inplace_flag=False):
+        api_func_name = self.get_api_func_name()
+        if inplace_flag and api_func_name[-1] != '_':
+            inplace_name = api_func_name + '_'
+        else:
+            inplace_name = api_func_name
+        api_code = f"""
+PADDLE_API {self.get_return_type(inplace_flag)} {inplace_name}({self.get_define_args(inplace_flag)}) {{
+{self.get_grad_outputs_define(inplace_flag)}
+    {api_func_name}({self.get_grad_api_call_args(inplace_flag)});
+    return {self.get_grad_output(inplace_flag)};
+}}
+"""
+        return api_code
+
+    def gene_api_code(self):
+        if not self.is_base_api and not self.is_only_composite_api:
+            invoke_func_name = self.invoke.split('(')[0]
+            if (not invoke_func_name.endswith("_grad")) and (
+                not invoke_func_name.endswith('_impl')
+            ):
+                return ""
+
+        if self.is_only_composite_api:
+            return ""
+
+        api_code = self.gene_base_api_code()
+        if self.is_base_api and len(self.inplace_map) > 0:
+            if self.api[-1] == '_':
+                api_code = ""
+            api_code = api_code + self.gene_base_api_code_for_inplace()
+        return api_code
+
+    def gene_api_declaration(self):
+        if not self.is_base_api and not self.is_only_composite_api:
+            invoke_func_name = self.invoke.split('(')[0]
+            if (not invoke_func_name.endswith("_grad")) and (
+                not invoke_func_name.endswith('_impl')
+            ):
+                return ""
+
+        if self.is_only_composite_api:
+            return ""
+
+        api_declaration = ""
+        api_func_name = self.get_api_func_name()
+        if api_func_name[-1] != '_':
+            api_declaration = f"""
+PADDLE_API {self.get_return_type()} {api_func_name}({self.get_declare_args()});
+"""
+
+        if self.is_base_api and len(self.inplace_map) > 0:
+            if api_func_name[-1] != '_':
+                api_func_name += '_'
+            api_declaration = (
+                api_declaration
+                + f"""
+PADDLE_API {self.get_return_type(inplace_flag=True)} {api_func_name}({self.get_declare_args(inplace_flag=True)});
+"""
+            )
+
+        return api_declaration
+
+
 def header_include():
     return """
 #include <tuple>
@@ -441,12 +494,13 @@ def header_include():
 
 def source_include(header_file_path):
     return f"""
-#include "{header_file_path}"
+
 #include <memory>
 
 #include "glog/logging.h"
 #include "paddle/common/flags.h"
 
+{header_file_path}
 #include "paddle/phi/api/lib/api_custom_impl.h"
 #include "paddle/phi/api/lib/api_gen_utils.h"
 #include "paddle/phi/api/lib/api_registry.h"
@@ -536,11 +590,20 @@ def generate_api(
     header_file.write(header_include())
     header_file.write(namespace[0])
 
-    include_header_file = (
-        "paddle/phi/api/include/fused_api.h"
-        if is_fused_ops_yaml is True
-        else "paddle/phi/api/include/api.h"
-    )
+    if not grad_flag:
+        include_header_file = (
+            '#include "paddle/phi/api/include/fused_api.h"'
+            if is_fused_ops_yaml is True
+            else '#include "paddle/phi/api/include/api.h"'
+        )
+    else:
+        include_header_file = (
+            '#include "paddle/phi/api/backward/fused_backward_api.h" \n'
+            '#include "paddle/phi/api/backward/fused_backward_api_base.h" '
+            if is_fused_ops_yaml is True
+            else '#include "paddle/phi/api/backward/backward_api.h" \n'
+            '#include "paddle/phi/api/backward/backward_api_base.h" '
+        )
     # not all fused ops support dygraph
     if is_fused_ops_yaml is True:
         new_apis = [
@@ -555,8 +618,11 @@ def generate_api(
     source_file.write(namespace[0])
 
     for api in apis:
+        if not grad_flag:
+            forward_api = ForwardAPI(api)
+        else:
+            forward_api = BackwardAPI(api)
 
-        forward_api = ForwardAPI(api)
         if forward_api.api in backward_api_black_list:
             continue
         if forward_api.is_dygraph_api and not is_fused_ops_yaml:
@@ -569,10 +635,7 @@ def generate_api(
             forward_api.is_dygraph_api = True
 
         header_file.write(forward_api.gene_api_declaration())
-        if forward_api.api not in ["embedding_grad", "cudnn_lstm_grad"]:
-            source_file.write(forward_api.gene_api_code())
-    if not is_fused_ops_yaml and grad_flag:
-        source_file.write(manual_impl)
+        source_file.write(forward_api.gene_api_code())
 
     header_file.write(namespace[1])
     source_file.write(namespace[1])
diff --git a/paddle/phi/api/generator/backward_api_gen.py b/paddle/phi/api/generator/backward_api_gen.py
@@ -118,6 +118,10 @@ def gene_api_declaration(self):
                 not invoke_func_name.endswith('_impl')
             ):
                 return ""
+
+        if self.is_only_composite_api:
+            return ""
+
         api_func_name = self.get_api_func_name()
         api_declaration = f"""
 PADDLE_API void {api_func_name}({self.get_declare_args()});
diff --git a/paddle/phi/api/generator/dist_api_gen.py b/paddle/phi/api/generator/dist_api_gen.py
diff --git a/paddle/phi/api/generator/sparse_bw_api_gen.py b/paddle/phi/api/generator/sparse_bw_api_gen.py
