[PIR][oneDNN] Add constraints for conv_elementwise_add_onednn_fuse_pass (#64448)

LLee233 · web-flow · commit 344b82d889ac · 2024-05-24T15:12:00.000+08:00
diff --git a/paddle/fluid/pir/transforms/onednn/conv_elementwise_add_onednn_fuse_pass.cc b/paddle/fluid/pir/transforms/onednn/conv_elementwise_add_onednn_fuse_pass.cc
@@ -74,6 +74,19 @@ class ConvElementwiseAddPattern : public paddle::drr::DrrPatternBase {
       }
       return true;
     });
+
+    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) -> bool {
+      auto conv2d_out_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("conv2d_out"));
+      auto residual_param_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("residual_param"));
+      // conv_elementwise_add_onednn_fuse_pass does not support broadcast
+      if (conv2d_out_shape != residual_param_shape) {
+        return false;
+      }
+      return true;
+    });
+
     paddle::drr::ResultPattern res = pat.ResultPattern();
 
     const auto &fused_conv2d_add =
@@ -154,6 +167,19 @@ class ConvElementwiseAddAsYPattern : public paddle::drr::DrrPatternBase {
       }
       return true;
     });
+
+    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) -> bool {
+      auto conv2d_out_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("conv2d_out"));
+      auto residual_param_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("residual_param"));
+      // conv_elementwise_add_onednn_fuse_pass does not support broadcast
+      if (conv2d_out_shape != residual_param_shape) {
+        return false;
+      }
+      return true;
+    });
+
     paddle::drr::ResultPattern res = pat.ResultPattern();
 
     const auto &fused_conv2d_add =
@@ -250,6 +276,29 @@ class FusedConvBiasElementwiseAddPattern : public paddle::drr::DrrPatternBase {
       }
       return true;
     });
+
+    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) -> bool {
+      auto fuse_activation = match_ctx.Attr<std::string>("fuse_activation");
+      auto fuse_residual_connection =
+          match_ctx.Attr<bool>("fuse_residual_connection");
+      if (!fuse_activation.empty() || fuse_residual_connection) {
+        return false;
+      }
+      return true;
+    });
+
+    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) -> bool {
+      auto conv2d_out_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("conv2d_out"));
+      auto residual_param_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("residual_param"));
+      // conv_elementwise_add_onednn_fuse_pass does not support broadcast
+      if (conv2d_out_shape != residual_param_shape) {
+        return false;
+      }
+      return true;
+    });
+
     paddle::drr::ResultPattern res = pat.ResultPattern();
 
     const auto &fused_conv2d_add =
@@ -348,6 +397,29 @@ class FusedConvBiasElementwiseAddAsYPattern
       }
       return true;
     });
+
+    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) -> bool {
+      auto fuse_activation = match_ctx.Attr<std::string>("fuse_activation");
+      auto fuse_residual_connection =
+          match_ctx.Attr<bool>("fuse_residual_connection");
+      if (!fuse_activation.empty() || fuse_residual_connection) {
+        return false;
+      }
+      return true;
+    });
+
+    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) -> bool {
+      auto conv2d_out_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("conv2d_out"));
+      auto residual_param_shape =
+          pir::GetShapeFromValue(match_ctx.Tensor("residual_param"));
+      // conv_elementwise_add_onednn_fuse_pass does not support broadcast
+      if (conv2d_out_shape != residual_param_shape) {
+        return false;
+      }
+      return true;
+    });
+
     paddle::drr::ResultPattern res = pat.ResultPattern();
 
     const auto &fused_conv2d_add =
diff --git a/test/ir/pir/fused_pass/onednn/test_conv2d_elemenwise_add_fuse_pass.py b/test/ir/pir/fused_pass/onednn/test_conv2d_elemenwise_add_fuse_pass.py
@@ -219,5 +219,82 @@ def test_check_output(self):
         self.check_pass_correct()
 
 
+class TestConv2dBiasAddFusePassasY(PassTest):
+    r"""
+         x_var   filter
+           \      /
+            conv2d   bias
+              \      /
+    residual conv2d_bias
+          \       /
+             out
+    """
+
+    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, 5, 5], dtype='float32'
+                )
+                conv2d = paddle.nn.Conv2D(
+                    in_channels=5,
+                    out_channels=1,
+                    kernel_size=[1, 1],
+                    groups=1,
+                    stride=[1, 1],
+                    padding=[1, 1, 1, 1],
+                    dilation=[1, 1],
+                    data_format='NCHW',
+                    bias_attr=False,
+                )
+
+                bias_attr = paddle.ParamAttr(
+                    learning_rate=0.0,
+                    initializer=paddle.nn.initializer.Normal(mean=0.0, std=2.0),
+                )
+                bias = paddle.static.create_parameter(
+                    shape=[1], dtype='float32', attr=bias_attr, is_bias=False
+                )
+                residual_data = paddle.static.data(
+                    name="residual_data", shape=[5, 1, 7, 7], dtype="float32"
+                )
+                conv2d_out = paddle.add(conv2d(x), bias)
+                out = paddle.add(residual_data, conv2d_out)
+                out = paddle.assign(out)
+                self.pass_attr_list = [
+                    {'conv2d_bias_fuse_pass': {}},
+                    {'conv_elementwise_add_onednn_fuse_pass': {}},
+                ]
+
+                self.feeds = {
+                    "x": np.random.random((5, 5, 5, 5)).astype("float32"),
+                    "bias": np.random.random(1).astype("float32"),
+                    "residual_data": np.random.random((5, 1, 7, 7)).astype(
+                        "float32"
+                    ),
+                }
+                self.fetch_list = [out]
+                self.valid_op_map = {
+                    "onednn_op.fused_conv2d": 1,
+                    "pd_op.conv2d": 0,
+                    "pd_op.add": 0,
+                }
+                return [main_prog, start_prog]
+
+    def sample_program(self):
+        yield self.build_ir_program(), False
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+
 if __name__ == "__main__":
     unittest.main()
