[PIR][oneDNN] Add fc_onednn_enable_pass (#63518)

LLee233 · web-flow · commit 559d018a3042 · 2024-04-24T13:22:34.000+08:00
diff --git a/paddle/fluid/inference/api/paddle_pass_builder.cc b/paddle/fluid/inference/api/paddle_pass_builder.cc
@@ -630,6 +630,8 @@ const std::vector<std::string> kPirMkldnnPasses{
     "matmul_transpose_reshape_fuse_pass",
     "matmul_elementwise_add_fuse_pass",
     "matmul_activation_fuse_pass",
+    "fc_fuse_pass",
+    "fc_onednn_enable_pass",
     "softplus_activation_fuse_pass",
     "conv_elementwise_add_onednn_fuse_pass",
     "conv_activation_onednn_fuse_pass",
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -1154,10 +1154,12 @@ def GenOneDnnExtraAttrsDefaultValue(onednn_extra_args):
     STR_TEMPLATE = """  pir::Attribute attr_{attr_name} = {op_attribute_type}::get(pir::IrContext::Instance(), {attr});
 """
     ARRAY_ATTRIBUTE_TEMPLATE = """  std::vector<pir::Attribute> vec_{attr_name};
-std::vector<{cpp_type}> vec_values = {attr_valuse};
-for (size_t i = 0; i < static_cast<size_t>(vec_values.size()); i++) {{
-    {create_attribute}
-    vec_{attr_name}.push_back(attr_{attr_name});
+{{
+    std::vector<{cpp_type}> vec_values = {attr_valuse};
+    for (size_t i = 0; i < static_cast<size_t>(vec_values.size()); i++) {{
+        {create_attribute}
+        vec_{attr_name}.push_back(attr_{attr_name});
+    }}
 }}
 pir::Attribute attr_{attr_name} = pir::ArrayAttribute::get(pir::IrContext::Instance(), vec_{attr_name});
 """
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -95,7 +95,7 @@
   extra_args : str mkldnn_data_type="float32"
 
 - op : fc
-  extra_args : bool ALL_KERNELS_MUST_COMPUTE_RUNTIME_SHAPE=true, bool use_quantizer=false, str mkldnn_data_type="float32", float scale_in=1.0, float[] scale_weights={1.0f}, float scale_out=1.0, bool force_fp32_output=false
+  extra_args : bool use_quantizer=false, str mkldnn_data_type="float32", float scale_in=1.0, float[] scale_weights={1.0f}, float scale_out=1.0, bool force_fp32_output=false, str fuse_activation = "", float fuse_alpha = 0.0, float fuse_beta = 0.0, float fused_output_scale = 1.0f, int[] fused_reshape2_shape = {}
 
 - op : flatten
   extra_args : str mkldnn_data_type="float32"
diff --git a/paddle/fluid/pir/transforms/onednn/fc_onednn_enable_pass.cc b/paddle/fluid/pir/transforms/onednn/fc_onednn_enable_pass.cc
@@ -0,0 +1,110 @@
+// 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.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pir/transforms/onednn/fc_onednn_enable_pass.h"
+
+#include "paddle/fluid/pir/dialect/operator/ir/onednn_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
+#include "paddle/fluid/pir/utils/general_functions.h"
+
+#include "paddle/pir/include/pass/pass.h"
+#include "paddle/pir/include/pass/pass_registry.h"
+
+namespace {
+
+class FcOneDNNEnablePattern : public paddle::drr::DrrPatternBase {
+ public:
+  FcOneDNNEnablePattern() {}
+
+  std::string name() const override { return "FcOneDNNEnablePattern"; }
+
+  void operator()(paddle::drr::DrrPatternContext *ctx) const override {
+    paddle::drr::SourcePattern pat = ctx->SourcePattern();
+
+    const auto &fc = pat.Op(paddle::dialect::FcOp::name(),
+                            {{"in_num_col_dims", pat.Attr("in_num_col_dims")},
+                             {"activation_type", pat.Attr("activation_type")},
+                             {"padding_weights", pat.Attr("padding_weights")}});
+
+    fc({&pat.Tensor("input"), &pat.Tensor("weight"), &pat.Tensor("bias")},
+       {&pat.Tensor("Out")});
+
+    pat.AddConstraint([&](const paddle::drr::MatchContext &match_ctx) {
+      auto input_shape = pir::GetShapeFromValue(match_ctx.Tensor("input"));
+      auto input_dims = input_shape.size();
+      bool support_dims = (input_dims >= 2 || input_shape.size() <= 4);
+      constexpr size_t height_axis = 2;
+      constexpr size_t width_axis = 3;
+      bool support_size = input_dims == 4 ? (input_shape[width_axis] == 1 &&
+                                             input_shape[height_axis] == 1)
+                                          : true;
+      if (!support_dims || !support_size) return false;
+      return true;
+    });
+
+    pat.AddConstraint([&](const paddle::drr::MatchContext &match_ctx) {
+      auto act_type = match_ctx.Attr<std::string>("activation_type");
+      if (!(act_type == "" || act_type == "relu")) return false;
+      return true;
+    });
+
+    paddle::drr::ResultPattern res = pat.ResultPattern();
+
+    std::unordered_map<std::string, paddle::drr::Attribute> fused_attrs{
+        {"in_num_col_dims", pat.Attr("in_num_col_dims")},
+        {"activation_type", pat.Attr("activation_type")},
+        {"padding_weights", pat.Attr("padding_weights")},
+        {"use_quantizer", res.BoolAttr(false)},
+        {"mkldnn_data_type", res.StrAttr("float32")},
+        {"scale_in", res.Float32Attr(1.0f)},
+        {"scale_weights", res.VectorFloatAttr({1.0f})},
+        {"scale_out", res.Float32Attr(1.0f)},
+        {"force_fp32_output", res.BoolAttr(false)},
+        {"fuse_activation", res.StrAttr("")},
+        {"fuse_alpha", res.Float32Attr(0.0f)},
+        {"fuse_beta", res.Float32Attr(0.0f)},
+        {"fused_output_scale", res.Float32Attr(1.0f)},
+        {"fused_reshape2_shape", res.VectorInt32Attr({})}};
+
+    const auto &fused_fc =
+        res.Op(paddle::onednn::dialect::FcOp::name(), fused_attrs);
+
+    fused_fc({&res.Tensor("input"), &res.Tensor("weight"), &res.Tensor("bias")},
+             {&res.Tensor("Out")});
+  }
+};
+
+class FcOneDNNEnablePass : public pir::PatternRewritePass {
+ public:
+  FcOneDNNEnablePass() : pir::PatternRewritePass("fc_onednn_enable_pass", 3) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add(paddle::drr::Create<FcOneDNNEnablePattern>(context));
+    return ps;
+  }
+};
+
+}  // namespace
+
+namespace pir {
+
+std::unique_ptr<Pass> CreateFcOneDNNEnablePass() {
+  // pd_op.fc -> onednn_op.fc
+  return std::make_unique<FcOneDNNEnablePass>();
+}
+}  // namespace pir
+
+REGISTER_IR_PASS(fc_onednn_enable_pass, FcOneDNNEnablePass);
diff --git a/paddle/fluid/pir/transforms/onednn/fc_onednn_enable_pass.h b/paddle/fluid/pir/transforms/onednn/fc_onednn_enable_pass.h
@@ -0,0 +1,26 @@
+// 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.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include "paddle/pir/include/core/dll_decl.h"
+
+namespace pir {
+
+class Pass;
+
+IR_API std::unique_ptr<Pass> CreateFcOneDNNEnablePass();
+
+}  // namespace pir
diff --git a/paddle/fluid/pir/transforms/passes.h b/paddle/fluid/pir/transforms/passes.h
@@ -53,6 +53,7 @@ USE_PIR_PASS(reshape_transpose_matmul_fuse_pass);
 USE_PIR_PASS(matmul_transpose_reshape_fuse_pass);
 USE_PIR_PASS(matmul_elementwise_add_fuse_pass);
 USE_PIR_PASS(matmul_activation_fuse_pass);
+USE_PIR_PASS(fc_onednn_enable_pass);
 USE_PIR_PASS(softplus_activation_fuse_pass);
 USE_PIR_PASS(conv_elementwise_add_onednn_fuse_pass);
 USE_PIR_PASS(conv_activation_onednn_fuse_pass);
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -1131,7 +1131,7 @@
   attrs :
     {scale_in : Scale_in, scale_out : Scale_out, scale_weights : Scale_weights}
   extra :
-    attrs : [bool ALL_KERNELS_MUST_COMPUTE_RUNTIME_SHAPE = true, bool use_mkldnn = false, bool use_quantizer = false, str mkldnn_data_type = "float32", float Scale_in = 1.0f, "float[] Scale_weights = {1.0f}", float Scale_out = 1.0f, bool force_fp32_output = false]
+    attrs : [bool use_mkldnn = false, bool use_quantizer = false, str mkldnn_data_type = "float32", float Scale_in = 1.0f, "float[] Scale_weights = {1.0f}", float Scale_out = 1.0f, bool force_fp32_output = false, str fuse_activation = "" , float fuse_alpha = 0.0f, float fuse_beta = 0.0f, float fused_output_scale = 1.0f, 'int[] fused_reshape2_shape = {}']
 
 - op : feed
   outputs: {out: Out}
diff --git a/paddle/phi/kernels/fusion/onednn/fc_kernel.cc b/paddle/phi/kernels/fusion/onednn/fc_kernel.cc
@@ -549,12 +549,14 @@ void RunKernel(const phi::OneDNNContext& dev_ctx,
   const auto out_md =
       dst_memory_p->get_desc().reshape(common::vectorize(out->dims()));
 
+  std::vector<int> reshape2_shape = {};
   if (dev_ctx.HasDnnAttr("fused_reshape2_shape")) {
+    reshape2_shape = PADDLE_GET_CONST(
+        std::vector<int>, dev_ctx.GetDnnAttr("fused_reshape2_shape"));
+  }
+  if (!reshape2_shape.empty()) {
     phi::funcs::SetOutMemDescWithReshape2FuseSupport(
-        PADDLE_GET_CONST(std::vector<int>,
-                         dev_ctx.GetDnnAttr("fused_reshape2_shape")),
-        out,
-        out_md);
+        reshape2_shape, out, out_md);
   } else {
     out->set_mem_desc(out_md);
   }
diff --git a/test/ir/pir/fused_pass/onednn/test_fc_onednn_enable_pass.py b/test/ir/pir/fused_pass/onednn/test_fc_onednn_enable_pass.py
@@ -0,0 +1,103 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from pass_test import PassTest
+
+import paddle
+
+paddle.enable_static()
+
+
+class TestFcOneDNNEnablePattern(PassTest):
+    r"""
+    x     w
+     \   /
+     matmul  y
+        \   /
+         add
+          |
+        [relu]
+          |
+         out
+    """
+
+    def is_program_valid(self, program=None):
+        return True
+
+    def sample_program(self):
+        for x_shape in [[3, 2]]:
+            for w_shape in [[2, 3]]:
+                for y_shape in [[3], [1, 3]]:
+                    for with_relu in [False, True]:
+                        with paddle.pir_utils.IrGuard():
+                            start_prog = paddle.static.Program()
+                            main_prog = paddle.static.Program()
+                            with paddle.pir.core.program_guard(
+                                main_prog, start_prog
+                            ):
+                                x = paddle.static.data(
+                                    name='x', shape=x_shape, dtype='float32'
+                                )
+                                w = paddle.static.data(
+                                    name='w', shape=w_shape, dtype='float32'
+                                )
+                                y = paddle.static.data(
+                                    name='y', shape=y_shape, dtype='float32'
+                                )
+                                if with_relu:
+                                    relu_op = paddle.nn.ReLU()
+                                    out = relu_op(
+                                        paddle.add(paddle.matmul(x, w), y)
+                                    )
+                                else:
+                                    out = paddle.add(paddle.matmul(x, w), y)
+                                out = paddle.assign(out)
+                                self.pass_attr_list = [
+                                    {'fc_fuse_pass': {}},
+                                    {'fc_onednn_enable_pass': {}},
+                                ]
+                                self.feeds = {
+                                    "x": np.random.random(x_shape).astype(
+                                        "float32"
+                                    ),
+                                    "w": np.random.random(w_shape).astype(
+                                        "float32"
+                                    ),
+                                    "y": np.random.random(y_shape).astype(
+                                        "float32"
+                                    ),
+                                }
+                                self.fetch_list = [out]
+                                self.valid_op_map = {
+                                    "pd_op.add": 0,
+                                    "pd_op.relu": 0,
+                                    "pd_op.matmul": 0,
+                                    "pd_op.fc": 0,
+                                    "onednn_op.fc": 1,
+                                }
+
+                                yield [main_prog, start_prog], False
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+
+if __name__ == "__main__":
+    unittest.main()
