dygraph Prim code gen （node.cc)

paddle/fluid/eager/auto_code_generator/generator/CMakeLists.txt

@@ -47,8 +47,8 @@ add_custom_target(
           ${forwards_cc_path}
   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_forwards_h_path}
           ${forwards_h_path}
-  #   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_nodes_cc_path}
-  #         ${nodes_cc_path}
+  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_nodes_cc_path}
+          ${nodes_cc_path}
   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_nodes_h_path}
           ${nodes_h_path}
   VERBATIM)

paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py

@@ -403,6 +403,23 @@ def ParseYamlInplaceInfo(string):
     return inplace_map
 
 
+def ParseYamlCompositeInfo(string):
+    # example  composite: fun(args1, args2,.....)
+    fname = r'(.*?)'
+    wspace = r'\s*'
+    fargs = r'(.*?)'
+    pattern = (
+        fr'{fname}{wspace}\({wspace}{fargs}{wspace}\)'
+    )
+
+    m = re.search(pattern, string)
+    composite_fun_info = []
+    composite_fun_info.append(m.group(1))
+    func_args = m.group(2).split(",")
+    for fun_arg in func_args:
+        composite_fun_info.append(fun_arg.strip())
+
+    return composite_fun_info
 ####################
 #  Generator Base  #
 ####################
@@ -438,6 +455,7 @@ def __init__(self, forward_api_contents, namespace):
         # Special Op Attributes
         self.optional_inputs = []  # [name, ...]
         self.no_need_buffers = []  # [name, ...]
+        self.composite_func_info = [] # [func_name, input_name, ...]
         self.intermediate_outputs = []  # [name, ...]
         self.forward_inplace_map = {}  # {name : name, ...}
 
@@ -459,6 +477,13 @@ def ParseNoNeedBuffer(self):
                 name = RemoveSpecialSymbolsInName(name)
                 self.no_need_buffers.append(name.strip())
 
+    def ParseComposite(self):
+        grad_api_contents = self.grad_api_contents
+
+        if 'composite' in grad_api_contents.keys():
+            composite_str = grad_api_contents['composite']
+            self.composite_func_info = ParseYamlCompositeInfo(composite_str)
+
     def ParseDispensable(self):
         forward_api_contents = self.forward_api_contents
 

paddle/fluid/eager/auto_code_generator/generator/eager_gen.py

@@ -332,6 +332,9 @@ class {} : public egr::GradNodeBase {{
 #include "paddle/fluid/eager/nan_inf_utils.h"
 #include "paddle/phi/api/include/sparse_api.h"
 #include "paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h"
+#include "paddle/fluid/prim/api/manual/backward/composite_backward_api.h"
+#include "paddle/fluid/prim/api/all.h"
+#include "paddle/fluid/prim/utils/utils.h"
 DECLARE_bool(check_nan_inf);
 {}
 """
@@ -546,6 +549,7 @@ def __init__(
         # self.forward_outputs_position_map
         # self.optional_inputs
         # self.no_need_buffers
+        # self.composite_func_info
         # self.intermediate_outputs
         # self.forward_inplace_map
         FunctionGeneratorBase.__init__(self, forward_api_contents, namespace)
@@ -871,6 +875,7 @@ def GenerateNodeCreationCodes(self, for_backward=False):
         backward_grad_outputs_map = self.backward_grad_outputs_map
         backward_attrs_list = self.backward_attrs_list
         optional_inputs = self.optional_inputs
+        is_composite_forward_api = False if self.composite_func_info == [] else True
 
         # Pass Stop Gradient Args
         pass_stop_gradient_args_str = self.GetPassStopGradientArgsList(
@@ -1056,6 +1061,8 @@ def run(self):
             self.ParseBackwardInplaceInfo()
             # Parse no_need_buffer
             self.ParseNoNeedBuffer()
+            # Parse composite
+            self.ParseComposite()
 
         # Parse optional_inputs
         self.ParseDispensable()
@@ -1826,16 +1833,19 @@ def GenerateHigherOrderNodeCreationCode(self):
         is_invoke_forward_api = IsInvokeForwardApi(
             self.grad_api_contents, self.forward_apis_dict
         )
+        is_composite_forward_api = False if self.composite_func_info == [] else True
+
         if next_node_generator is not None:
             has_higher_order_node = True
             return (
                 has_higher_order_node,
                 is_invoke_forward_api,
+                is_composite_forward_api,
                 next_grad_node_creation_str,
                 next_grad_node_out_list,
                 next_node_generator.backward_forward_inputs_map,
             )
-        elif not is_invoke_forward_api:
+        elif not is_invoke_forward_api and not is_composite_forward_api:
             next_grad_node_creation_str = f"""  if(trace_backward) {{
     PADDLE_THROW(phi::errors::Unavailable(
     \"The Op {self.backward_api_name} doesn't have any grad\"
@@ -1845,6 +1855,7 @@ def GenerateHigherOrderNodeCreationCode(self):
         return (
             has_higher_order_node,
             is_invoke_forward_api,
+            is_composite_forward_api,
             next_grad_node_creation_str,
             next_grad_node_out_list,
             None,
@@ -1942,13 +1953,15 @@ def GenerateNodeDefinition(
         self,
         has_higher_order_node,
         is_invoke_forward_api,
+        is_composite_grad_api,
         next_grad_node_creation_str,
         next_grad_node_out_list,
         backward_forward_inputs_map_next,
     ):
         namespace = self.namespace
         forward_api_name = self.forward_api_name
         backward_api_name = self.backward_api_name
+        composite_backward_api_name = self.composite_func_info[0] if is_composite_grad_api else None
         backward_forward_inputs_map = self.backward_forward_inputs_map
         backward_grad_inputs_map = self.backward_grad_inputs_map
         backward_grad_outputs_map = self.backward_grad_outputs_map
@@ -2133,6 +2146,7 @@ def GenerateNodeDefinition(
         # Grad Function Call String
         slot_num_bwd_outputs = len(self.forward_inputs_position_map.keys())
         grad_api_namespace = f"paddle::experimental::{namespace}"
+        composite_grad_api_namespace = f"paddle::prim::{namespace}"
         grad_function_prepare_str = f"""
   const auto& out_metas = OutputMeta();
   paddle::small_vector<std::vector<paddle::experimental::Tensor>, egr::kSlotSmallVectorSize> returns({slot_num_bwd_outputs});
@@ -2203,6 +2217,8 @@ def GenerateNodeDefinition(
   }}"""
 
         grad_api_args_str = ", ".join(grad_api_args)
+        composite_grad_api_args_str = ", ".join(grad_api_args)
+        composite_template_name = "<paddle::experimental::Tensor>"
 
         if is_invoke_forward_api:
             autograd_api_out = "auto"
@@ -2225,6 +2241,16 @@ def GenerateNodeDefinition(
   {out_assign_str}}} else {{
   {indent}{autograd_api_out} api_output = paddle::experimental::{self.namespace}{self.grad_api_contents['invoke']};
   {out_assign_str}{indent}}}
+  """
+        elif is_composite_grad_api:
+            grad_function_call_str = f"""
+  if (paddle::prim::PrimCommonUtils::IsPrimEnabled()) {{
+  {indent}{composite_grad_api_namespace}{composite_backward_api_name}{composite_template_name}({composite_grad_api_args_str});
+  VLOG(4) << paddle::string::Sprintf("composite api %s is called" , "{composite_backward_api_name}");
+  }}else{{
+  {indent}{grad_api_namespace}{backward_api_name}({grad_api_args_str});
+  VLOG(4) << paddle::string::Sprintf("origin api %s is called" , "{backward_api_name}");
+  }}
   """
         else:
             grad_function_call_str = f"""
@@ -2328,6 +2354,9 @@ def GenerateNodeDefinition(
             var_str += f"\n{indent}  output_str += output_{new_name}_str; "
 
         log_str = AFTER_LOG_PRINT_TEMPLATE.format(var_str)
+        # TODO Ruting modify in the future
+        # if is_composite_forward_api:
+        #     next_grad_node_creation_str = ''
 
         self.node_definition_str = GRAD_FUNCTION_TEMPLATE.format(
             grad_node_name,
@@ -2361,6 +2390,7 @@ def run(self):
         (
             has_higher_order_node,
             is_invoke_forward_api,
+            is_composite_grad_api,
             next_grad_node_creation_str,
             next_grad_node_out_list,
             backward_forward_inputs_map,
@@ -2371,6 +2401,7 @@ def run(self):
         self.GenerateNodeDefinition(
             has_higher_order_node,
             is_invoke_forward_api,
+            is_composite_grad_api,
             next_grad_node_creation_str,
             next_grad_node_out_list,
             backward_forward_inputs_map,

paddle/fluid/operators/generator/parse_utils.py

@@ -293,6 +293,7 @@ def check_op_config(op_entry, op_name):
         'intermediate',
         'no_need_buffer',
         'data_transform',
+        'composite'
     )
     infer_meta_key_set = ('func', 'param')
     kernel_key_set = ('func', 'param', 'data_type', 'layout', 'backend')

paddle/phi/api/yaml/backward.yaml

@@ -1259,6 +1259,7 @@
     param : [out]
   kernel :
     func : tanh_grad
+  composite : tanh_grad(out, out_grad, x_grad)
   backward : tanh_double_grad
   inplace : (out_grad -> x_grad)