【pir】modify paddlMix model of wrong opverify （masked_multihead_attention_） (#67339)

xiaoguoguo626807 · web-flow · commit 1d42b84c1382 · 2024-08-14T10:01:28.000+08:00
* delete -1 shape check

* recommit

* recover support -1 shape

* optional output test

* recover

* modify infermeta gen
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -509,6 +509,14 @@ def __init__(self, op_yaml_item, op_compat_item, yaml_file):
 
         # parse inplace && view
         self.inplace_map = self.parse_op_inplace_info()
+        # sometime output is the inplace value of input, if input is optional, output should be optional too
+        if self.inplace_map is not None:
+            self.refine_output_optional_list(self.output_optional_list)
+            self.cross_check(
+                self.output_name_list,
+                self.output_type_list,
+                self.output_optional_list,
+            )
         self.view_map = self.parse_op_view_info()
 
         # parse data_transform
@@ -932,6 +940,15 @@ def parse_output_optional_list(self):
                 optional_list.append("false")
         return optional_list
 
+    def refine_output_optional_list(self, optional_list):
+        for i, output_info in enumerate(self.op_yaml_item['outputs']):
+            if output_info['name'] in self.inplace_map.keys():
+                input_index = self.input_name_list.index(
+                    self.inplace_map[output_info['name']]
+                )
+                if self.input_optional_list[input_index] == "true":
+                    optional_list[i] = "true"
+
     def parse_output_intermediate_list(self):
         intermediate_list = []
         for output_info in self.op_yaml_item['outputs']:
diff --git a/paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
@@ -961,6 +961,10 @@ def gen_infermeta_func_str(args, op_info):
             )
         inuse_infer_meta_args.append(f"{op_info.output_name_list[idx]}")
 
+    if op_info.inplace_map is not None:
+        for input in op_info.inplace_map.keys():
+            inuse_infer_meta_args.append(f"{op_info.inplace_map[input]}")
+
     spmd_params = []
     if args.with_distributed:
         if op_info.spmd_rule_func is not None:
