extend capability of conv_bias_bn fusion (#66568)

LLee233 · web-flow · commit 280d84a0452b · 2024-07-29T09:02:30.000+08:00
diff --git a/paddle/fluid/pir/transforms/onednn/conv2d_bn_onednn_fuse_pass.cc b/paddle/fluid/pir/transforms/onednn/conv2d_bn_onednn_fuse_pass.cc
@@ -192,8 +192,33 @@ class Conv2dBiasBnOneDNNFusePattern
     pir::Value bn_bias = op.bias();
 
     auto add_y_shape = pir::GetShapeFromValue(add_y);
-    // bias currently only support per_tensor add
-    if (add_y_shape.size() != 1 || add_y_shape[0] != 1) return false;
+    auto bn_bias_shape = pir::GetShapeFromValue(bn_bias);
+    std::vector<int64_t> add_y_new_shape{add_y_shape[0]};
+    // Support both per_tensor & per_channel addition
+    if (add_y_shape.size() != 1) {
+      size_t idx;
+      if (data_format == "NHWC") {
+        idx = add_y_shape.size() - 1;
+      } else {
+        idx = 1;
+      }
+      if (add_y_shape[idx] != bn_bias_shape[0]) return false;
+      bool is_ok = true;
+      for (size_t i = 0; i < add_y_shape.size(); i++) {
+        if (i == idx) continue;
+        if (add_y_shape[i] != 1) {
+          is_ok = false;
+          break;
+        }
+      }
+      if (!is_ok) return false;
+      // reshape add_y from [1, X, 1, 1] (NCHW) to [X]
+      add_y_new_shape[0] = add_y_shape[idx];
+    } else if (add_y_shape[0] != 1) {
+      return false;
+    }
+    paddle::dialect::ReshapeOp reshape_add_y_op =
+        rewriter.Build<paddle::dialect::ReshapeOp>(add_op.y(), add_y_new_shape);
 
     // --- deal with filter ---
     auto bn_variance_shape = pir::GetShapeFromValue(bn_variance);
@@ -240,9 +265,10 @@ class Conv2dBiasBnOneDNNFusePattern
         conv2d_filter, reshape_scale_op.out());
 
     // --- deal with bias ---
-    // (add_op.y() - bn_mean)*scale + bn_bias
+    // (add_op.y()(reshaped) - bn_mean)*scale + bn_bias
     paddle::dialect::SubtractOp sub_op_1 =
-        rewriter.Build<paddle::dialect::SubtractOp>(add_op.y(), bn_mean);
+        rewriter.Build<paddle::dialect::SubtractOp>(reshape_add_y_op.out(),
+                                                    bn_mean);
     paddle::dialect::MultiplyOp mul_bias_op =
         rewriter.Build<paddle::dialect::MultiplyOp>(sub_op_1.out(),
                                                     div_op.out());
diff --git a/test/ir/pir/fused_pass/onednn/test_conv2d_bn_onednn_fuse_pass.py b/test/ir/pir/fused_pass/onednn/test_conv2d_bn_onednn_fuse_pass.py
@@ -151,5 +151,149 @@ def setUp(self):
         self.places.append(paddle.CPUPlace())
 
 
+class TestConv2dBiasBnOneDNNPassPatternCase2(PassTest):
+    r"""
+    x_var   f_var
+      \      /
+       conv2d  add_y
+          \     /
+            add
+             |
+         BatchNorm
+             |
+    """
+
+    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=[3, 1, 28, 28], dtype='float32'
+                )
+                bias_attr = paddle.ParamAttr(
+                    learning_rate=0.0,
+                    initializer=paddle.nn.initializer.Normal(mean=0.0, std=2.0),
+                )
+                y = paddle.static.create_parameter(
+                    shape=[1, 32, 1, 1],
+                    dtype='float32',
+                    attr=bias_attr,
+                    is_bias=False,
+                )
+                conv2d = paddle.nn.Conv2D(
+                    in_channels=1,
+                    out_channels=32,
+                    kernel_size=3,
+                    padding=1,
+                    data_format='NCHW',
+                    bias_attr=False,
+                )
+                bn = paddle.nn.BatchNorm2D(
+                    num_features=32,
+                    data_format='NCHW',
+                    use_global_stats=True,
+                )
+                add_out = paddle.add(conv2d(x), y)
+                out = bn(add_out)
+                out = paddle.assign(out)
+                self.pass_attr_list = [{'conv2d_bias_bn_onednn_fuse_pass': {}}]
+                self.feeds = {
+                    "x": np.random.random((3, 1, 28, 28)).astype("float32"),
+                    "y": np.random.random((1, 32, 1, 1)).astype("float32"),
+                }
+                self.fetch_list = [out]
+                self.valid_op_map = {
+                    "onednn_op.fused_conv2d": 1,
+                    "pd_op.batch_norm_": 0,
+                }
+                return [main_prog, start_prog]
+
+    def sample_program(self):
+        pir_program = self.build_ir_program()
+        yield pir_program, False
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+
+
+class TestConv2dBiasBnOneDNNPassPatternCase3(PassTest):
+    r"""
+    x_var   f_var
+      \      /
+       conv2d  add_y
+          \     /
+            add
+             |
+         BatchNorm
+             |
+    """
+
+    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=[3, 28, 28, 1], dtype='float32'
+                )
+                bias_attr = paddle.ParamAttr(
+                    learning_rate=0.0,
+                    initializer=paddle.nn.initializer.Normal(mean=0.0, std=2.0),
+                )
+                y = paddle.static.create_parameter(
+                    shape=[1, 1, 1, 32],
+                    dtype='float32',
+                    attr=bias_attr,
+                    is_bias=False,
+                )
+                conv2d = paddle.nn.Conv2D(
+                    in_channels=1,
+                    out_channels=32,
+                    kernel_size=3,
+                    padding=1,
+                    data_format='NHWC',
+                    bias_attr=False,
+                )
+                bn = paddle.nn.BatchNorm2D(
+                    num_features=32,
+                    data_format='NHWC',
+                    use_global_stats=True,
+                )
+                add_out = paddle.add(conv2d(x), y)
+                out = bn(add_out)
+                out = paddle.assign(out)
+                self.pass_attr_list = [{'conv2d_bias_bn_onednn_fuse_pass': {}}]
+                self.feeds = {
+                    "x": np.random.random((3, 28, 28, 1)).astype("float32"),
+                    "y": np.random.random((1, 1, 1, 32)).astype("float32"),
+                }
+                self.fetch_list = [out]
+                self.valid_op_map = {
+                    "onednn_op.fused_conv2d": 1,
+                    "pd_op.batch_norm_": 0,
+                }
+                return [main_prog, start_prog]
+
+    def sample_program(self):
+        pir_program = self.build_ir_program()
+        yield pir_program, False
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+
+
 if __name__ == "__main__":
     unittest.main()
