add group_norm_silu_xpu_fuse (#62689)

linkk08 · web-flow · commit 49c7cc7f1225 · 2024-04-07T14:14:24.000+08:00
diff --git a/paddle/fluid/framework/ir/CMakeLists.txt b/paddle/fluid/framework/ir/CMakeLists.txt
@@ -301,6 +301,8 @@ if(WITH_XPU)
                ${XPU_PASS_DEPS})
   pass_library(add_layernorm_xpu_fuse_pass inference DIR xpu DEPS
                ${XPU_PASS_DEPS})
+  pass_library(group_norm_silu_xpu_fuse_pass inference DIR xpu DEPS
+               ${XPU_PASS_DEPS})
   pass_library(xpu_delete_cast_op_pass inference DIR xpu DEPS ${XPU_PASS_DEPS})
   pass_library(fold_interp_outsize_fuse_pass inference DIR xpu DEPS
                ${XPU_PASS_DEPS})
diff --git a/paddle/fluid/framework/ir/pass.cc b/paddle/fluid/framework/ir/pass.cc
@@ -68,6 +68,7 @@ static const std::vector<std::string> xpu_support_subgraph_passes = {
     "constant_folding_pass",
     "delete_elementwise_mul_op_pass",
     "generate_sequence_xpu_fuse_pass",
+    "group_norm_silu_xpu_fuse_pass",
     "embedding_with_eltwise_add_xpu_fuse_pass",
     "multi_encoder_xpu_fuse_pass",
     "multi_encoder_xpu_adaptive_seqlen_fuse_pass",
diff --git a/paddle/fluid/framework/ir/xpu/group_norm_silu_xpu_fuse_pass.cc b/paddle/fluid/framework/ir/xpu/group_norm_silu_xpu_fuse_pass.cc
@@ -0,0 +1,208 @@
+// 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 <string>
+
+#include "glog/logging.h"
+
+#include "paddle/fluid/framework/ir/fuse_pass_base.h"
+#include "paddle/fluid/framework/ir/graph_pattern_detector.h"
+#include "paddle/fluid/framework/ir/pass.h"
+#include "paddle/fluid/framework/ir/xpu/pass_utils.h"
+#include "paddle/fluid/framework/ir/xpu/quant_utils.h"
+#include "paddle/fluid/framework/op_version_registry.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace phi {
+class DenseTensor;
+}  // namespace phi
+
+namespace paddle {
+namespace framework {
+class Scope;
+}  // namespace framework
+}  // namespace paddle
+
+namespace paddle {
+namespace framework {
+namespace ir {
+namespace patterns {
+
+/*
+fuse gn + activation block in to xpu_ele_fusion op
+For example:
+graph:
+                      X
+              Scale   |   Bias
+                   \  |  /
+                  group norm
+                   /  |  \
+                  /   |   \
+            variance  |   mean
+                      |
+                     silu
+                      |
+                    output
+------------------------------------------------------
+After the pass is applied:
+                      X
+              Scale   |   Bias
+                   \  |  /
+                gn_silu_fusion
+                      |
+                     Out
+*/
+struct GroupNormalizeSiluXPUPattern : public PatternBase {
+  GroupNormalizeSiluXPUPattern(PDPattern* pattern,
+                               const std::string& name_scope);
+  // declare operator node's name
+  PATTERN_DECL_NODE(gn);
+  PATTERN_DECL_NODE(silu);
+  // declare variable node's name
+  PATTERN_DECL_NODE(gn_x);
+  PATTERN_DECL_NODE(gn_bias);
+  PATTERN_DECL_NODE(gn_scale);
+  PATTERN_DECL_NODE(gn_y);
+  PATTERN_DECL_NODE(gn_mean);
+  PATTERN_DECL_NODE(gn_variance);
+  PATTERN_DECL_NODE(silu_out);
+};
+
+GroupNormalizeSiluXPUPattern::GroupNormalizeSiluXPUPattern(
+    PDPattern* pattern, const std::string& name_scope)
+    : PatternBase(pattern, name_scope, name_scope) {
+  auto gn = pattern->NewNode(gn_repr())->assert_is_op("group_norm");
+  auto gn_x = pattern->NewNode(gn_x_repr())
+                  ->assert_is_op_input("group_norm", "X")
+                  ->AsInput();
+  auto gn_bias = pattern->NewNode(gn_bias_repr())
+                     ->assert_is_op_input("group_norm", "Bias")
+                     ->assert_is_persistable_var()
+                     ->AsInput();
+  auto gn_scale = pattern->NewNode(gn_scale_repr())
+                      ->assert_is_op_input("group_norm", "Scale")
+                      ->assert_is_persistable_var()
+                      ->AsInput();
+  auto gn_y = pattern->NewNode(gn_y_repr())
+                  ->assert_is_op_output("group_norm", "Y")
+                  ->assert_is_op_input("silu", "X")
+                  ->assert_has_n_outputs(1);
+  auto gn_mean = pattern->NewNode(gn_mean_repr())
+                     ->assert_is_op_output("group_norm", "Mean")
+                     ->assert_has_n_outputs(0);
+  auto gn_variance = pattern->NewNode(gn_variance_repr())
+                         ->assert_is_op_output("group_norm", "Variance")
+                         ->assert_has_n_outputs(0);
+  gn->LinksFrom({gn_x, gn_bias, gn_scale})
+      .LinksTo({gn_y, gn_mean, gn_variance});
+
+  auto silu = pattern->NewNode(silu_repr())->assert_is_op("silu");
+  auto silu_out = pattern->NewNode(silu_out_repr())
+                      ->AsOutput()
+                      ->assert_is_op_output("silu", "Out");
+  silu->LinksFrom({gn_y}).LinksTo({silu_out});
+}
+
+}  // namespace patterns
+
+class GroupNormalizeSiluXPUFusePass : public FusePassBase {
+ protected:
+  void ApplyImpl(ir::Graph* graph) const override;
+
+ private:
+  void FuseGroupNormalizeSilu(ir::Graph* graph) const;
+
+  const std::string name_scope_{"group_norm_silu_xpu_fuse_pass"};
+};
+
+void GroupNormalizeSiluXPUFusePass::ApplyImpl(ir::Graph* graph) const {
+  PADDLE_ENFORCE_NOT_NULL(
+      graph, platform::errors::PreconditionNotMet("graph should not be null."));
+  Init(name_scope_, graph);
+
+  FuseGroupNormalizeSilu(graph);
+}
+
+void GroupNormalizeSiluXPUFusePass::FuseGroupNormalizeSilu(
+    ir::Graph* graph) const {
+  GraphPatternDetector gpd;
+  patterns::GroupNormalizeSiluXPUPattern pattern(gpd.mutable_pattern(),
+                                                 name_scope_);
+
+  int found_subgraph_count = 0;
+  auto handler = [&](const GraphPatternDetector::subgraph_t& subgraph,
+                     Graph* graph) {
+    VLOG(4) << "handle GroupNormalizeSiluXPUFusePass fuse";
+    // declare operator node's name
+    GET_IR_NODE(gn);
+    GET_IR_NODE(silu);
+    // declare variable node's name
+    GET_IR_NODE(gn_x);
+    GET_IR_NODE(gn_bias);
+    GET_IR_NODE(gn_scale);
+    GET_IR_NODE(gn_y);
+    GET_IR_NODE(gn_mean);
+    GET_IR_NODE(gn_variance);
+    GET_IR_NODE(silu_out);
+
+    auto* block = gn->Op()->Block();
+    auto* scope = param_scope();
+    PADDLE_ENFORCE_NOT_NULL(
+        scope, platform::errors::InvalidArgument("Scope cannot be nullptr."));
+    // delete useless node
+    std::unordered_set<const Node*> delete_nodes;
+
+    float eps = PADDLE_GET_CONST(float, gn->Op()->GetAttr("epsilon"));
+    int groups = PADDLE_GET_CONST(int, gn->Op()->GetAttr("groups"));
+
+    std::string fused_op_out_name;
+    fused_op_out_name = silu_out->Name();
+    // Generate add_layernorm fused op
+    framework::OpDesc fused_op_desc(block);
+
+    fused_op_desc.SetType("group_norm_silu_xpu");
+    // set attrs for fused op
+    fused_op_desc.SetInput("x", {gn_x->Name()});
+    fused_op_desc.SetInput("bias", {gn_bias->Name()});
+    fused_op_desc.SetInput("scale", {gn_scale->Name()});
+    fused_op_desc.SetAttr("epsilon", eps);
+    fused_op_desc.SetAttr("groups", groups);
+    fused_op_desc.SetOutput("out", {fused_op_out_name});
+    // relink fused op
+    auto* fused_op = graph->CreateOpNode(&fused_op_desc);
+    IR_NODE_LINK_TO(gn_x, fused_op);
+    IR_NODE_LINK_TO(gn_bias, fused_op);
+    IR_NODE_LINK_TO(gn_scale, fused_op);
+    IR_NODE_LINK_TO(fused_op, silu_out);
+
+    delete_nodes.insert({gn, silu, gn_y, gn_mean, gn_variance});
+    GraphSafeRemoveNodes(graph, delete_nodes);
+    found_subgraph_count++;
+  };
+
+  gpd(graph, handler);
+  AddStatis(found_subgraph_count);
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(group_norm_silu_xpu_fuse_pass,
+              paddle::framework::ir::GroupNormalizeSiluXPUFusePass);
+
+REGISTER_PASS_CAPABILITY(group_norm_silu_xpu_fuse_pass)
+    .AddCombination(
+        paddle::framework::compatible::OpVersionComparatorCombination().EQ(
+            "group_norm_silu_xpu", 0));
diff --git a/paddle/fluid/inference/api/paddle_pass_builder.cc b/paddle/fluid/inference/api/paddle_pass_builder.cc
@@ -538,6 +538,7 @@ XpuPassStrategy::XpuPassStrategy() : PassStrategy({}) {
       "cast_embedding_trans_ids_to_int32_pass",
       "delete_elementwise_mul_op_pass",
       "generate_sequence_xpu_fuse_pass",
+      "group_norm_silu_xpu_fuse_pass",
       "embedding_with_eltwise_add_xpu_fuse_pass",
       "qk_qkv_attention_xpu_fuse_pass",
       "multi_encoder_xpu_fuse_pass",
diff --git a/paddle/phi/api/yaml/fused_ops.yaml b/paddle/phi/api/yaml/fused_ops.yaml
@@ -375,6 +375,15 @@
     func : generate_sequence_xpu
     data_type : dtype
 
+- op : group_norm_silu_xpu
+  args : (Tensor x, Tensor scale, Tensor bias, int groups, float epsilon)
+  output : Tensor(out)
+  infer_meta :
+    func : GroupNormalizeSiluXPUInferMeta
+  kernel :
+    func : group_norm_silu_xpu
+    data_type : x
+
 - op : layer_norm_act_xpu
   args : (Tensor x, Tensor scale, Tensor bias, int begin_norm_axis, float epsilon, int act_type, float act_param)
   output : Tensor(out)
diff --git a/paddle/phi/backends/xpu/xpu1_op_list.cc b/paddle/phi/backends/xpu/xpu1_op_list.cc
@@ -154,6 +154,7 @@ XPUOpMap& get_kl1_ops() {
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT32})},
+      {"group_norm_silu_xpu", XPUKernelSet({phi::DataType::FLOAT32})},
       {"hard_switch_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"hard_switch", XPUKernelSet({phi::DataType::FLOAT32})},
       {"index_select",
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -549,6 +549,8 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::FLOAT32})},
       {"grid_sampler_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"grid_sampler", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"group_norm_silu_xpu",
+       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"hard_sigmoid_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"hard_sigmoid",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -523,6 +523,8 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32})},
       {"grid_sampler_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"group_norm_silu_xpu",
+       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"hard_sigmoid_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"hard_sigmoid",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
diff --git a/paddle/phi/infermeta/fusion.cc b/paddle/phi/infermeta/fusion.cc
@@ -116,6 +116,20 @@ void AddLayernormXPUInferMeta(const MetaTensor& x,
   out->share_lod(x);
 }
 
+void GroupNormalizeSiluXPUInferMeta(const MetaTensor& x,
+                                    const MetaTensor& scale,
+                                    const MetaTensor& bias,
+                                    int groups,
+                                    float epsilon,
+                                    MetaTensor* out) {
+  auto x_dims = x.dims();
+  auto out_dims = x_dims;
+  out->set_dims(out_dims);
+  out->set_dtype(x.dtype());
+  out->set_layout(x.layout());
+  out->share_lod(x);
+}
+
 void FusedMultiTransformerInferMeta(
     const MetaTensor& x,
     const std::vector<const MetaTensor*>& ln_scales,
diff --git a/paddle/phi/infermeta/fusion.h b/paddle/phi/infermeta/fusion.h
@@ -70,6 +70,13 @@ void AddLayernormXPUInferMeta(const MetaTensor& x,
                               float epsilon,
                               MetaTensor* out);
 
+void GroupNormalizeSiluXPUInferMeta(const MetaTensor& x,
+                                    const MetaTensor& scale,
+                                    const MetaTensor& bias,
+                                    int groups,
+                                    float epsilon,
+                                    MetaTensor* out);
+
 void BlockMultiheadAttentionInferMeta(const MetaTensor& qkv,
                                       const MetaTensor& key_cache,
                                       const MetaTensor& value_cache,
diff --git a/paddle/phi/kernels/fusion/xpu/group_norm_silu_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/group_norm_silu_xpu_kernel.cc
@@ -0,0 +1,66 @@
+// 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/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/common_shape.h"
+
+namespace phi {
+namespace fusion {
+
+template <typename T, typename Context>
+void GroupNormalizeSiluXPUKernel(const Context& ctx,
+                                 const DenseTensor& x,
+                                 const DenseTensor& scale,
+                                 const DenseTensor& bias,
+                                 int groups,
+                                 float epsilon,
+                                 DenseTensor* out) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
+
+  auto* in_data = reinterpret_cast<const XPUType*>(x.data<T>());
+  auto* scale_data = reinterpret_cast<const float*>(scale.data<float>());
+  auto* bias_data = reinterpret_cast<const float*>(bias.data<float>());
+  auto* out_data = reinterpret_cast<XPUType*>(ctx.template Alloc<T>(out));
+  int n = static_cast<int>(x.dims()[0]);
+  int c = static_cast<int>(x.dims()[1]);
+  int h = static_cast<int>(x.dims()[2]);
+  int w = static_cast<int>(x.dims()[3]);
+
+  int r = xpu::group_norm_silu_fusion<XPUType>(ctx.x_context(),
+                                               in_data,
+                                               out_data,
+                                               n,
+                                               c,
+                                               h,
+                                               w,
+                                               groups,
+                                               epsilon,
+                                               scale_data,
+                                               bias_data,
+                                               nullptr,
+                                               nullptr,
+                                               true);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "group_norm_silu_fusion");
+}
+
+}  // namespace fusion
+}  // namespace phi
+
+PD_REGISTER_KERNEL(group_norm_silu_xpu,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::fusion::GroupNormalizeSiluXPUKernel,
+                   float,
+                   phi::dtype::float16) {}
diff --git a/test/ir/inference/test_xpu_group_norm_silu_pass.py b/test/ir/inference/test_xpu_group_norm_silu_pass.py
