[OneDNN][PIR] Fix vit model fail for matmul op (#68132)

zhanglirong1999 · web-flow · commit 14901ba2115e · 2024-09-19T09:49:01.000+08:00
* Fix vit model fail for transpose and matmul op

* add ut for transpose

* add warning

* revert tramspose op
diff --git a/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_pass.cc b/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_pass.cc
@@ -883,8 +883,9 @@ class CpuBfloat16PatternThree_one : public paddle::drr::DrrPatternBase {
       if (mkldnn_data_type != "bfloat16") {
         return false;
       }
+      // For fused_matmul, it name residual_data as residual_param
       const std::vector<std::string> permitted_input_names = {
-          "x", "y", "input", "residual_param"};
+          "x", "y", "input", "residual_param", "residual_data"};
       auto op_info =
           pir::IrContext::Instance()->GetRegisteredOpInfo(bfloat16_ops_);
       paddle::dialect::OpYamlInfoParser yaml_parser(
diff --git a/test/ir/pir/fused_pass/onednn/test_cpu_bfloat16_pir_pass.py b/test/ir/pir/fused_pass/onednn/test_cpu_bfloat16_pir_pass.py
@@ -1048,5 +1048,44 @@ def test_check_output(self):
         self.check_pass_correct()
 
 
+class TestTransposePattern(PassTest):
+    def is_program_valid(self, program=None):
+        return True
+
+    def build_ir_program(self):
+        with paddle.pir_utils.IrGuard():
+            main_prog = paddle.static.Program()
+            start_prog = paddle.static.Program()
+            with paddle.pir.core.program_guard(main_prog, start_prog):
+                x = paddle.static.data(name='x', shape=[5, 5], dtype='float32')
+                out = paddle.transpose(x, [1, 0])
+                out = paddle.assign(out)
+                self.pass_attr_list = [
+                    {'onednn_placement_pass': {}},
+                    {'cpu_bfloat16_placement_pass': {}},
+                    {'cpu_bfloat16_pass': {}},
+                    {'cpu_bfloat16_type_placement_pass': {}},
+                ]
+                self.feeds = {
+                    "x": np.random.random((5, 5)).astype("float32"),
+                }
+                self.fetch_list = [out]
+                self.valid_op_map = {
+                    "onednn_op.transpose": 1,
+                    "pd_op.transpose": 0,
+                }
+                return [main_prog, start_prog]
+
+    def sample_program(self):
+        yield self.build_ir_program(), False
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+        self.skip_accuracy_verification = True
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+
 if __name__ == "__main__":
     unittest.main()
