Fix optional output error when input is none (#60269)

0x45f · web-flow · commit 1468a0a29de8 · 2023-12-25T14:12:45.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/api_gen.py b/paddle/fluid/pir/dialect/op_generator/api_gen.py
@@ -85,6 +85,7 @@
     {in_combine}
     {compute_op}
     {handle_optional_outputs}
+    {set_null_type}
     {out_split}
     {return_result}"""
 
@@ -148,6 +149,12 @@
         optional_{name} = paddle::make_optional<std::vector<pir::OpResult>>(optional_{name}_slice_op.outputs());
     }}"""
 
+SET_NULL_TYPE_TEMPLATE = """
+    if (!{input}) {{
+        {op_name}_op.result({index}).set_type(pir::Type());
+    }}"""
+
+
 COMBINE_OP_TEMPLATE = """
     auto {op_name} = ApiBuilder::Instance().GetBuilder()->Build<pir::CombineOp>({in_name});"""
 
@@ -435,6 +442,21 @@ def _gen_handle_optional_outputs(self, op_info, op_name):
                     )
         return ret
 
+    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):
+                in_name = inplace_map[out_name]
+                ret += SET_NULL_TYPE_TEMPLATE.format(
+                    input=in_name, op_name=op_name, index=i
+                )
+        return ret
+
     def _gen_in_combine(self, op_info, is_mutable_attr, is_vector_mutable_attr):
         name_list = op_info.input_name_list
         type_list = op_info.input_type_list
@@ -727,6 +749,7 @@ def _gen_one_impl(
                     handle_optional_outputs=self._gen_handle_optional_outputs(
                         op_info, kernel_name
                     ),
+                    set_null_type=self._gen_set_null_type(op_info, kernel_name),
                     out_split=out_split,
                     return_result=self._gen_return_result(ret_list),
                 )
@@ -782,6 +805,7 @@ def _gen_one_impl(
                 handle_optional_outputs=self._gen_handle_optional_outputs(
                     op_info, op_name
                 ),
+                set_null_type=self._gen_set_null_type(op_info, op_name),
                 out_split=out_split,
                 return_result=self._gen_return_result(ret_list),
             )
