[OneDNN][PIR] Fix picodet performance drop in bf16 (#68784)

zhanglirong1999 · web-flow · commit 6203aba5d3a2 · 2024-10-23T16:07:50.000+08:00
* Scale Op with bias skip bfloat16

* add scale ut for ci
diff --git a/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_placement_pass.cc b/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_placement_pass.cc
@@ -96,6 +96,15 @@ class OneDNNBf16PlacementPattern : public pir::RewritePattern {
         return false;
       }
     }
+    if (op->name() == "onednn_op.scale" || op->name() == "onednn_op.scale_") {
+      bool bias_after_scale =
+          op_attr.at("bias_after_scale").dyn_cast<pir::BoolAttribute>().data();
+      if (bias_after_scale) {
+        // If bias after scale, add quant/dequant for sacle will cause some
+        // error
+        return false;
+      }
+    }
 
     const std::vector<std::string> permitted_input_names = {
         "x", "y", "input", "residual_param", "residual_data"};
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
@@ -1082,5 +1082,44 @@ def test_check_output(self):
         self.check_pass_correct()
 
 
+class TestScalePattern(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.scale(x, scale=1e-3, bias=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.scale": 1,
+                    "pd_op.scale": 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()
