【PIR】modify optional output as fake value of pd_api (#67465)

xiaoguoguo626807 · web-flow · commit 5649abbc3b12 · 2024-08-28T17:11:44.000+08:00
* modify optional output

* modify optional output

* add optional output gen

* add new

* modify grad optional output

* add double grad inplace option output

* recover dygraph yaml

* modify rnn ci error

* modify generate_proposal infermeta to aviod option null type output

* modify generate_proposal inferemeta to fix xxx op.py
diff --git a/paddle/fluid/pir/dialect/op_generator/api_gen.py b/paddle/fluid/pir/dialect/op_generator/api_gen.py
@@ -249,11 +249,16 @@
         optional_{name} = paddle::make_optional<std::vector<pir::Value>>(optional_{name}_slice_op.outputs());
     }}"""
 
-SET_NULL_TYPE_TEMPLATE = """
+SET_NULL_TYPE_WITH_INPLACE_TEMPLATE = """
     if (!{input}) {{
         {op_name}_op.result({index}).set_type(pir::Type());
     }}"""
 
+SET_NULL_TYPE_TEMPLATE = """
+    pir::Type {op_name}_op_result_{index}_type = {op_name}_op.result({index}).type();
+    if ({op_name}_op_result_{index}_type.isa<paddle::dialect::DenseTensorType>() && {op_name}_op_result_{index}_type.dyn_cast<paddle::dialect::DenseTensorType>().dims().size() == -1) {{
+        {op_name}_op.result({index}).set_type(pir::Type());
+    }}"""
 
 COMBINE_OP_TEMPLATE = """
     auto {op_name} = ApiBuilder::Instance().GetBuilder()->Build<pir::CombineOp>({in_name});"""
@@ -342,6 +347,27 @@ def _is_optional_input(self, op_info, input_name):
         return False
 
     def _is_optional_output(self, op_info, output_name):
+        output_optional_list = op_info.output_optional_list
+        output_name_list = op_info.output_name_list
+        intermediate_list = op_info.output_intermediate_list
+        output_index = output_name_list.index(output_name)
+        if (
+            intermediate_list[output_index] == 'false'
+            and op_info.output_optional_list[output_index] == 'true'
+        ):
+            return True
+        else:
+            return False
+
+    def _is_backward_op(self, op_info):
+        op_names = op_info.op_phi_name
+        for name in op_names:
+            if name.endswith(('_grad', '_grad_')):
+                return True
+        else:
+            return False
+
+    def _is_optional_inplace_output(self, op_info, output_name):
         op_names = op_info.op_phi_name
         for name in op_names:
             if name.endswith(('_grad', '_grad_')):
@@ -358,6 +384,15 @@ def _is_optional_output(self, op_info, output_name):
                 return True
         return False
 
+    def _need_optional_output(self, op_info, name):
+        if self._is_optional_inplace_output(op_info, name):
+            return True
+        if self._is_backward_op(op_info) and self._is_optional_output(
+            op_info, name
+        ):
+            return True
+        return False
+
     # =====================================
     # Gen declare functions
     # =====================================
@@ -433,15 +468,15 @@ def _gen_ret_type(self, op_info):
             ):
                 if intermediate == 'true':
                     continue
-                if self._is_optional_output(op_info, name):
+                if self._need_optional_output(op_info, name):
                     ret.append(OPTIONAL_VALUE_TYPE_MAP[type])
                 else:
                     ret.append(VALUE_TYPE_MAP[type])
             return 'std::tuple<{}>'.format(', '.join(ret))
         elif output_num == 1:
             index = intermediate_list.index('false')
             name = name_list[index]
-            if self._is_optional_output(op_info, name):
+            if self._need_optional_output(op_info, name):
                 return OPTIONAL_VALUE_TYPE_MAP[type_list[index]]
             else:
                 return VALUE_TYPE_MAP[type_list[index]]
@@ -510,7 +545,7 @@ def _gen_handle_optional_inputs(self, op_info):
                     ret += OPTIONAL_VALUE_INPUT_TEMPLATE.format(name=name)
         return ret
 
-    def _gen_handle_optional_outputs(self, op_info, op_name):
+    def _gen_handle_optional_inplace_outputs(self, op_info, op_name):
         name_list = op_info.output_name_list
         type_list = op_info.output_type_list
         intermediate_list = op_info.output_intermediate_list
@@ -520,7 +555,7 @@ def _gen_handle_optional_outputs(self, op_info, op_name):
         ):
             if intermediate == 'true':
                 continue
-            if self._is_optional_output(op_info, name):
+            if self._need_optional_output(op_info, name):
                 if VECTOR_TYPE in type:
                     ret += OPTIONAL_VECTOR_VALUE_OUTPUT_TEMPLATE.format(
                         name=name,
@@ -538,16 +573,16 @@ def _gen_handle_optional_outputs(self, op_info, op_name):
     def _gen_set_null_type(self, op_info, op_name):
         name_list = op_info.output_name_list
         inplace_map = op_info.inplace_map
-        if inplace_map is None:
-            return ""
 
         ret = ""
         for i, out_name in enumerate(name_list):
-            if self._is_optional_output(op_info, out_name):
+            if self._is_optional_inplace_output(op_info, out_name):
                 in_name = inplace_map[out_name]
-                ret += SET_NULL_TYPE_TEMPLATE.format(
+                ret += SET_NULL_TYPE_WITH_INPLACE_TEMPLATE.format(
                     input=in_name, op_name=op_name, index=i
                 )
+            elif self._is_optional_output(op_info, out_name):
+                ret += SET_NULL_TYPE_TEMPLATE.format(op_name=op_name, index=i)
         return ret
 
     def _gen_in_combine(self, op_info, is_mutable_attr, is_vector_mutable_attr):
@@ -651,7 +686,7 @@ def _gen_out_split_and_ret_list(self, op_info, op_inst_name):
         ):
             if intermediate == 'true':
                 continue
-            if self._is_optional_output(op_info, name):
+            if self._need_optional_output(op_info, name):
                 ret_list.append(f'optional_{name}')
             elif VECTOR_TYPE in type:
                 split_op_name = f'{name}_split_op'
@@ -1003,7 +1038,7 @@ def _gen_one_impl(
                     ),
                     in_combine=in_combine,
                     compute_op=compute_op,
-                    handle_optional_outputs=self._gen_handle_optional_outputs(
+                    handle_optional_outputs=self._gen_handle_optional_inplace_outputs(
                         op_info, kernel_name
                     ),
                     set_null_type=self._gen_set_null_type(op_info, kernel_name),
@@ -1065,7 +1100,7 @@ def _gen_one_impl(
                 ),
                 in_combine=in_combine,
                 compute_op=compute_op,
-                handle_optional_outputs=self._gen_handle_optional_outputs(
+                handle_optional_outputs=self._gen_handle_optional_inplace_outputs(
                     op_info, op_name
                 ),
                 set_null_type=self._gen_set_null_type(op_info, op_name),
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -396,12 +396,12 @@ def process_backward_invoke_info(apis):
             api['invoke']['args'] = ', '.join(args)
 
 
-def process_optional_output_info(apis):
+def process_optional_inplace_output_info(apis):
     for api in apis:
         inputs_dict = to_named_dict(api['inputs'])
         for output in api['outputs']:
             if not api['is_fwd']:
-                output['optional'] = False
+                return
             else:
                 if (
                     api.get("inplace", None)
@@ -524,7 +524,7 @@ def gen(
     apis = extend_compat_info(apis, compats)
     apis = apis + get_inplace_api(apis)
     process_backward_invoke_info(apis)
-    process_optional_output_info(apis)
+    process_optional_inplace_output_info(apis)
 
     apis = [
         {**api, **{'class_name': to_pascal_case(api["name"]) + "Op"}}
diff --git a/paddle/fluid/primitive/codegen/templates/backend/generated/generated_static_backend.cc.j2 b/paddle/fluid/primitive/codegen/templates/backend/generated/generated_static_backend.cc.j2
@@ -81,7 +81,7 @@ auto op_res_{{i}} = std::get<{{i}}>(op_res);
       {% if outputs[i].typename == 'Tensor' and not outputs[i].optional %}
 Tensor {{outputs[i].name}}(std::make_shared<LazyTensor>(op_res_{{i}}));
       {% elif outputs[i].typename == 'Tensor' and outputs[i].optional %}
-paddle::optional<Tensor> {{outputs[i].name}};
+paddle::optional<Tensor> {{outputs[i].name}} = paddle::make_optional<Tensor>(Tensor());
 if(op_res_{{i}}){
   {{outputs[i].name}} = paddle::make_optional<Tensor>(Tensor(std::make_shared<LazyTensor>(op_res_{{i}}.get())));
 }
diff --git a/paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2 b/paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2
@@ -95,16 +95,32 @@ auto op_res = backend::{{api.name}}<LazyTensor>({{common.args(input_names, attr_
   {% if outputs|length > 1 %}
     {% for i in range(outputs|length) %}
       {% if outputs[i].typename=='Tensor' %}
+        {% if outputs[i].optional %}
+vjp_res[{{i}}][0] = std::get<{{i}}>(op_res).get();
+        {% else %}
 vjp_res[{{i}}][0] = std::get<{{i}}>(op_res);
+        {% endif %}
       {% else %}
+        {% if outputs[i].optional %}
+vjp_res[{{i}}] = std::get<{{i}}>(op_res).get();
+        {% else %}
 vjp_res[{{i}}] = std::get<{{i}}>(op_res);
+        {% endif %}
       {% endif %}
     {% endfor %}
   {% elif outputs|length == 1 %}
     {% if outputs[0].typename=='Tensor' %}
+      {% if outputs[0].optional %}
+vjp_res[0][0] = op_res.get();
+      {% else %}
 vjp_res[0][0] = op_res;
+      {% endif %}
     {% else %}
+       {% if outputs[0].optional %}
+vjp_res[0] = op_res.get();
+       {% else %}
 vjp_res[0] = op_res;
+       {% endif %}
     {% endif %}
   {% else %} {#- render nothing -#}
   {% endif %}
diff --git a/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc b/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc
@@ -166,21 +166,21 @@ std::vector<std::vector<paddle::Tensor>> fused_attention_vjp(
   // x_grad
   vjp_res[0][0] = std::get<8>(op_res);
   // ln_scale_grad
-  vjp_res[1][0] = std::get<4>(op_res);
+  vjp_res[1][0] = std::get<4>(op_res).get();
   // ln_bias_grad
-  vjp_res[2][0] = std::get<5>(op_res);
+  vjp_res[2][0] = std::get<5>(op_res).get();
   // qkv_weight_grad
   vjp_res[3][0] = std::get<9>(op_res);
   // qkv_bias_grad
-  vjp_res[4][0] = std::get<0>(op_res);
+  vjp_res[4][0] = std::get<0>(op_res).get();
   // out_linear_weight_grad
   vjp_res[5][0] = std::get<10>(op_res);
   // out_linear_bias_grad
-  vjp_res[6][0] = std::get<3>(op_res);
+  vjp_res[6][0] = std::get<3>(op_res).get();
   // ln_scale_2_grad
-  vjp_res[7][0] = std::get<6>(op_res);
+  vjp_res[7][0] = std::get<6>(op_res).get();
   // ln_bias_2_grad
-  vjp_res[8][0] = std::get<7>(op_res);
+  vjp_res[8][0] = std::get<7>(op_res).get();
   vjp_res = ConstructVjpResultByStopGradients(vjp_res, stop_gradients);
   return vjp_res;
 }
diff --git a/paddle/fluid/primitive/utils/static_utils.cc b/paddle/fluid/primitive/utils/static_utils.cc
@@ -45,7 +45,7 @@ void SetEmptyGrad(const std::vector<std::vector<Tensor>>& outputs,
                   const std::vector<std::vector<bool>>& stop_gradients) {
   for (size_t i = 0; i < outputs.size(); ++i) {
     for (size_t j = 0; j < outputs[i].size(); ++j) {
-      if (stop_gradients[i][j]) {
+      if (stop_gradients[i][j] && outputs[i][j].impl()) {
         std::static_pointer_cast<primitive::LazyTensor>(outputs[i][j].impl())
             ->set_empty();
       }
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -2597,6 +2597,9 @@ void GenerateProposalsV2InferMeta(const MetaTensor& scores,
                                   MetaTensor* rpn_rois_num) {
   rpn_rois->set_dims(common::make_ddim({-1, 4}));
   rpn_roi_probs->set_dims(common::make_ddim({-1, 1}));
+  if (rpn_rois_num) {
+    rpn_rois_num->set_dims(common::make_ddim({scores.dims()[0]}));
+  }
 }
 
 void LegacyGenerateProposalsInferMeta(const MetaTensor& scores,
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -574,7 +574,7 @@
     param : [x, x, grad_x_grad_forward]
   kernel :
     func : cos_triple_grad
-  optional: grad_out_forward, grad_x_grad_forward, grad_out_grad_grad
+  optional: grad_out_forward, grad_x_grad_forward, grad_out_grad_grad, grad_out_forward_grad
   inplace : (grad_x_grad_forward -> grad_out_forward_grad)
 
 - backward_op : cosh_grad
@@ -2928,7 +2928,7 @@
     param : [x, x, grad_x_grad_forward]
   kernel :
     func : sin_triple_grad
-  optional: grad_out_forward, grad_x_grad_forward, grad_out_grad_grad
+  optional: grad_out_forward, grad_x_grad_forward, grad_out_grad_grad, grad_out_forward_grad
   inplace : (grad_x_grad_forward -> grad_out_forward_grad)
 
 - backward_op : sinh_grad
diff --git a/paddle/phi/ops/yaml/inconsistent/static_backward.yaml b/paddle/phi/ops/yaml/inconsistent/static_backward.yaml
@@ -12,7 +12,7 @@
     param : [grad_out]
   kernel :
     func : add_double_grad
-  optional : grad_x_grad, grad_y_grad
+  optional : grad_x_grad, grad_y_grad, grad_out_grad
   backward : add_triple_grad
   inplace : (grad_x_grad -> grad_out_grad)
   composite : add_double_grad(y, grad_out, grad_x_grad, grad_y_grad, axis, grad_out_grad)
@@ -126,7 +126,7 @@
   kernel :
     func : divide_double_grad
     data_type : out
-  optional : grad_x, grad_x_grad, grad_y_grad
+  optional : grad_x, grad_x_grad, grad_y_grad, grad_out_grad
   inplace : (grad_x_grad -> grad_out_grad)
 
 - backward_op : divide_grad
@@ -328,7 +328,7 @@
     param : [x, y, grad_out]
   kernel :
     func : multiply_double_grad
-  optional : grad_x_grad, grad_y_grad
+  optional : grad_x_grad, grad_y_grad, grad_out_grad
   inplace : (grad_x_grad -> grad_out_grad)
   backward : multiply_triple_grad
   composite : multiply_double_grad(x, y, grad_out, grad_x_grad, grad_y_grad, axis, x_grad, y_grad, grad_out_grad)
@@ -457,7 +457,7 @@
     param : [grad_out]
   kernel :
     func : subtract_double_grad
-  optional : grad_x_grad, grad_y_grad
+  optional : grad_x_grad, grad_y_grad, grad_out_grad
   no_need_buffer : y, grad_out
   inplace : (grad_x_grad -> grad_out_grad)
   composite : subtract_double_grad(y, grad_out, grad_x_grad, grad_y_grad, axis, grad_out_grad)
diff --git a/test/legacy_test/test_fused_rotary_position_embedding.py b/test/legacy_test/test_fused_rotary_position_embedding.py
@@ -617,7 +617,11 @@ def test_static_time_major(self):
                 shape=(1, shape_q[0], 1, shape_q[3]),
                 dtype=self.dtype,
             )
-
+            q.stop_gradient = False
+            if v is not None:
+                v.stop_gradient = False
+            if k is not None:
+                k.stop_gradient = False
             out_q, out_k, out_v = fused_rotary_position_embedding(
                 q,
                 k,
@@ -629,6 +633,8 @@ def test_static_time_major(self):
                 time_major=True,
             )
 
+            dout = paddle.static.gradients(out_q, q)
+
         exe = paddle.static.Executor()
 
         feed = {

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ auto op_res_{{i}} = std::get<{{i}}>(op_res);`
`81`	`81`	`{% if outputs[i].typename == 'Tensor' and not outputs[i].optional %}`
`82`	`82`	`Tensor {{outputs[i].name}}(std::make_shared<LazyTensor>(op_res_{{i}}));`
`83`	`83`	`{% elif outputs[i].typename == 'Tensor' and outputs[i].optional %}`
`84`		`-paddle::optional<Tensor> {{outputs[i].name}};`
	`84`	`+paddle::optional<Tensor> {{outputs[i].name}} = paddle::make_optional<Tensor>(Tensor());`
`85`	`85`	`if(op_res_{{i}}){`
`86`	`86`	`{{outputs[i].name}} = paddle::make_optional<Tensor>(Tensor(std::make_shared<LazyTensor>(op_res_{{i}}.get())));`
`87`	`87`	`}`
Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ void SetEmptyGrad(const std::vector<std::vector<Tensor>>& outputs,`
`45`	`45`	`const std::vector<std::vector<bool>>& stop_gradients) {`
`46`	`46`	`for (size_t i = 0; i < outputs.size(); ++i) {`
`47`	`47`	`for (size_t j = 0; j < outputs[i].size(); ++j) {`
`48`		`- if (stop_gradients[i][j]) {`
	`48`	`+ if (stop_gradients[i][j] && outputs[i][j].impl()) {`
`49`	`49`	`std::static_pointer_cast<primitive::LazyTensor>(outputs[i][j].impl())`
`50`	`50`	`->set_empty();`
`51`	`51`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2597,6 +2597,9 @@ void GenerateProposalsV2InferMeta(const MetaTensor& scores,`
`2597`	`2597`	`MetaTensor* rpn_rois_num) {`
`2598`	`2598`	`rpn_rois->set_dims(common::make_ddim({-1, 4}));`
`2599`	`2599`	`rpn_roi_probs->set_dims(common::make_ddim({-1, 1}));`
	`2600`	`+ if (rpn_rois_num) {`
	`2601`	`+ rpn_rois_num->set_dims(common::make_ddim({scores.dims()[0]}));`
	`2602`	`+ }`
`2600`	`2603`	`}`
`2601`	`2604`
`2602`	`2605`	`void LegacyGenerateProposalsInferMeta(const MetaTensor& scores,`