[PIR][oneDNN] Extend Conv bias fusion capability (#66000)

LLee233 · web-flow · commit 709d1821ca2a · 2024-07-15T10:08:31.000+08:00
* extend conv+bias fusion

* revise copyright
diff --git a/paddle/fluid/pir/transforms/onednn/conv_bias_fuse_pass.cc b/paddle/fluid/pir/transforms/onednn/conv_bias_fuse_pass.cc
@@ -58,8 +58,7 @@ class ConvBiasFusePattern : public paddle::drr::DrrPatternBase {
 
     pat.Tensor("add_out") = add(pat.Tensor("conv_out"), pat.Tensor("bias"));
 
-    if (conv_name_ == paddle::dialect::Conv2dOp::name() ||
-        conv_name_ == paddle::onednn::dialect::FusedConv2dOp::name()) {
+    if (conv_name_ == paddle::dialect::Conv2dOp::name()) {
       pat.AddConstraint([&](const paddle::drr::MatchContext &match_ctx) {
         if (!pir::ValueIsPersistable(match_ctx.Tensor("bias"))) {
           return false;
@@ -96,7 +95,19 @@ class ConvBiasFusePattern : public paddle::drr::DrrPatternBase {
     }
     pat.AddConstraint([&](const paddle::drr::MatchContext &match_ctx) {
       auto bias_shape = pir::GetShapeFromValue(match_ctx.Tensor("bias"));
-      if (bias_shape.size() != 1) return false;
+      auto output_shape = pir::GetShapeFromValue(match_ctx.Tensor("conv_out"));
+      if (bias_shape.size() != 1) {
+        if (bias_shape[1] != output_shape[1]) return false;
+        bool is_ok = true;
+        for (size_t i = 0; i < bias_shape.size(); i++) {
+          if (i == 1) continue;
+          if (bias_shape[i] != 1) {
+            is_ok = false;
+            break;
+          }
+        }
+        return is_ok;
+      }
       return true;
     });
 
@@ -304,7 +315,7 @@ class FusedConvTransposeAddFusePattern : public paddle::drr::DrrPatternBase {
 
 class Conv2dBiasFusePass : public pir::PatternRewritePass {
  public:
-  Conv2dBiasFusePass() : pir::PatternRewritePass("conv2d_bias_fuse_pass", 2) {}
+  Conv2dBiasFusePass() : pir::PatternRewritePass("conv2d_bias_fuse_pass", 3) {}
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
@@ -319,7 +330,7 @@ class Conv2dBiasFusePass : public pir::PatternRewritePass {
 class Conv2dTransposeBiasFusePass : public pir::PatternRewritePass {
  public:
   Conv2dTransposeBiasFusePass()
-      : pir::PatternRewritePass("conv2d_transpose_bias_fuse_pass", 2) {}
+      : pir::PatternRewritePass("conv2d_transpose_bias_fuse_pass", 3) {}
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
@@ -331,7 +342,7 @@ class Conv2dTransposeBiasFusePass : public pir::PatternRewritePass {
 
 class Conv3dBiasFusePass : public pir::PatternRewritePass {
  public:
-  Conv3dBiasFusePass() : pir::PatternRewritePass("conv3d_bias_fuse_pass", 2) {}
+  Conv3dBiasFusePass() : pir::PatternRewritePass("conv3d_bias_fuse_pass", 3) {}
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
diff --git a/paddle/phi/kernels/onednn/conv_handler.h b/paddle/phi/kernels/onednn/conv_handler.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -125,12 +125,30 @@ class ConvOneDNNHandlerT
                 DataLayout::ONEDNN,
                 bias->layout()));
 
-        PADDLE_ENFORCE_EQ(
-            bias->dims().size(),
-            1,
-            phi::errors::InvalidArgument("Bias must only have 1 dimension, "
-                                         "i.e. X, but got dimension = %d .",
-                                         bias->dims().size()));
+        auto bias_shape = common::vectorize(bias->dims());
+        auto output_shape = common::vectorize(output->dims());
+        // layout of bias is always NCHW/NCDHW, so channel is always at 1st dim
+        if (bias_shape.size() != 1) {
+          PADDLE_ENFORCE_EQ(
+              bias_shape[1],
+              output_shape[1],
+              phi::errors::InvalidArgument(
+                  "Bias must only have 1 dimension or only bias_dims[1] == "
+                  "output_dims[1] i.e. [X] or [1, X, 1, 1], but got dimension "
+                  "== %d and failed",
+                  bias->dims().size()));
+          for (size_t i = 0; i < bias_shape.size(); i++) {
+            if (i == 1) continue;
+            PADDLE_ENFORCE_EQ(
+                bias_shape[i],
+                1,
+                phi::errors::InvalidArgument(
+                    "Bias with multiply dimensions must only have 1 dimension "
+                    "> 1, i.e. [1, X, 1, 1], but got %d-th dimension == %d .",
+                    i,
+                    bias_shape[i]));
+          }
+        }
       }
       const auto input_dims = input->dims();
       const auto data_dims =
@@ -195,6 +213,7 @@ class ConvOneDNNHandlerT
 
       if (bias) {
         auto bias_tz = common::vectorize(bias->dims());
+        if (bias_tz.size() > 1) bias_tz = {bias_tz[1]};
         dnnl::memory::desc bias_md =
             funcs::OneDNNMemDesc(bias_tz,
                                  dnnl::memory::data_type::f32,
@@ -594,8 +613,16 @@ class ConvOneDNNHandlerT
       }
       const K_Bias* bias_data = bias->data<K_Bias>();
 
+      dnnl::memory::desc bias_md = bias->mem_desc();
+      auto bias_tz = common::vectorize(bias->dims());
+      if (bias_tz.size() > 1) {
+        bias_tz = {bias_tz[1]};
+        bias_md = funcs::OneDNNMemDesc(bias_tz,
+                                       dnnl::memory::data_type::f32,
+                                       funcs::OneDNNMemoryFormat::x);
+      }
       return this->AcquireMemoryWithReorder(
-          bias->mem_desc(),
+          bias_md,
           this->fwd_pd_->bias_desc(),
           funcs::to_void_cast<K_Bias>(bias_data),
           "@bias_mem_p",
diff --git a/test/ir/pir/fused_pass/onednn/test_conv2d_bias_fuse_pass.py b/test/ir/pir/fused_pass/onednn/test_conv2d_bias_fuse_pass.py
@@ -83,6 +83,70 @@ def test_check_output(self):
         self.check_pass_correct()
 
 
+class TestConv2dAddFusePassCase2(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, 5, 5], dtype='float32'
+                )
+                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, 3, 1, 1],
+                    dtype='float32',
+                    attr=bias_attr,
+                    is_bias=False,
+                )
+                w_attr = paddle.ParamAttr(
+                    learning_rate=0.0,
+                    initializer=paddle.nn.initializer.Normal(mean=0.0, std=2.0),
+                )
+                conv2d = paddle.nn.Conv2D(
+                    in_channels=5,
+                    out_channels=3,
+                    kernel_size=[1, 1],
+                    groups=1,
+                    stride=[1, 1],
+                    padding=[1, 1, 1, 1],
+                    dilation=[1, 1],
+                    data_format='NCHW',
+                    bias_attr=False,
+                    weight_attr=w_attr,
+                )
+
+                out = paddle.add(conv2d(x), bias)
+                out = paddle.assign(out)
+                self.pass_attr_list = [{'conv2d_bias_fuse_pass': {}}]
+                self.feeds = {
+                    "x": np.random.random((5, 5, 5, 5)).astype("float32"),
+                    "bias": np.random.random((1, 3, 1, 1)).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()
+
+
 class TestConv2dAddFusePassWithAddParam(PassTest):
     def is_program_valid(self, program=None):
         return True
