[PIR]Add Fused dropout add pass (#58272)

xiaoguoguo626807 · zyfncg · 0x45f · web-flow · commit cc6d16aa14ab · 2023-10-23T18:53:28.000+08:00
* tmp

* fix flashattn compile bug

* modify unsqueeze

* modify part complex bug

* modify

* modify cpu

* [PIR]Migrate maximum into pir

* Polish code

* add ir_grad of static_gradient

* add test

* modify bug

* modify

* test_with_pir

* close one test

* add_math_op_patch

* modify

* modify

* add mean fill_constant test

* modify cpu int32 test

* get_shape_tensor

* delete

* add default place

* add dropout unittest

* Update test/legacy_test/test_elementwise_div_op.py

* modify review comment

* modify add test

* modify logic and sub

* modify

* add fused_gemm_epilogue pass

* polish code

* udate cmake

* fused_dropout_add_pass

* modify cmake

* modify

* Update test/cpp/pir/pattern_rewrite/CMakeLists.txt

* Update paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc

Co-authored-by: Yuanle Liu &lt;yuanlehome@163.com&gt;

* Update paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc

Co-authored-by: Yuanle Liu &lt;yuanlehome@163.com&gt;

* Update paddle/fluid/pybind/pir.cc

* Update test/ir/new_ir/fused_pass/test_fused_dropout_add.py

* Update test/ir/new_ir/fused_pass/test_fused_dropout_add.py

* Update paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc

* add ci converage

---------

Co-authored-by: zyfncg &lt;zhangyunfei07@baidu.com&gt;
Co-authored-by: 0x45f &lt;wangzhen45@baidu.com&gt;
Co-authored-by: Yuanle Liu &lt;yuanlehome@163.com&gt;
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -68,7 +68,20 @@
 OPS_API_TEMPLATE = """
 {{"{name}", (PyCFunction)(void (*)(void)){name}, METH_VARARGS | METH_KEYWORDS, "C++ interface function for {name}."}},"""
 
-NEED_GEN_STATIC_ONLY_APIS = ['fetch']
+NEED_GEN_STATIC_ONLY_APIS = [
+    'fetch',
+    'fused_embedding_eltwise_layernorm',
+    'fused_fc_elementwise_layernorm',
+    'fused_multi_transformer_xpu',
+    'fused_scale_bias_relu_conv_bnstats',
+    'fusion_transpose_flatten_concat',
+    'generate_sequence_xpu',
+    'layer_norm_act_xpu',
+    'multi_encoder_xpu',
+    'multihead_matmul',
+    'squeeze_excitation_block',
+    'yolo_box_xpu',
+]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
     'add_n_',
diff --git a/paddle/fluid/pir/dialect/operator/ir/CMakeLists.txt b/paddle/fluid/pir/dialect/operator/ir/CMakeLists.txt
@@ -98,7 +98,7 @@ set(api_source_file_tmp ${api_source_file}.tmp)
 add_custom_command(
   OUTPUT ${api_header_file} ${api_source_file}
   COMMAND
-    ${PYTHON_EXECUTABLE} ${api_gen_file} --op_yaml_files ${api_gen_yaml_files}
+    ${PYTHON_EXECUTABLE} ${api_gen_file} --op_yaml_files ${op_yaml_files}
     --op_compat_yaml_file ${op_compat_yaml_file} --namespaces ${op_namespace}
     --api_def_h_file ${api_header_file_tmp} --api_def_cc_file
     ${api_source_file_tmp}
@@ -129,10 +129,9 @@ set(python_c_source_file_tmp ${python_c_source_file}.tmp)
 add_custom_command(
   OUTPUT ${python_c_header_file} ${python_c_source_file}
   COMMAND
-    ${PYTHON_EXECUTABLE} ${python_c_gen_file} --op_yaml_files
-    ${api_gen_yaml_files} --op_compat_yaml_file ${op_compat_yaml_file}
-    --namespaces "paddle,pybind" --python_c_def_h_file
-    ${python_c_header_file_tmp} --python_c_def_cc_file
+    ${PYTHON_EXECUTABLE} ${python_c_gen_file} --op_yaml_files ${op_yaml_files}
+    --op_compat_yaml_file ${op_compat_yaml_file} --namespaces "paddle,pybind"
+    --python_c_def_h_file ${python_c_header_file_tmp} --python_c_def_cc_file
     ${python_c_source_file_tmp}
   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${python_c_header_file_tmp}
           ${python_c_header_file}
@@ -160,9 +159,9 @@ set(ops_api_source_file_tmp ${ops_api_source_file}.tmp)
 add_custom_command(
   OUTPUT ${ops_api_source_file}
   COMMAND
-    ${PYTHON_EXECUTABLE} ${ops_api_gen_file} --op_yaml_files
-    ${api_gen_yaml_files} --op_compat_yaml_file ${op_compat_yaml_file}
-    --namespaces "paddle,pybind" --ops_api_file ${ops_api_source_file_tmp}
+    ${PYTHON_EXECUTABLE} ${ops_api_gen_file} --op_yaml_files ${op_yaml_files}
+    --op_compat_yaml_file ${op_compat_yaml_file} --namespaces "paddle,pybind"
+    --ops_api_file ${ops_api_source_file_tmp}
   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${ops_api_source_file_tmp}
           ${ops_api_source_file}
   COMMENT "copy_if_different ${ops_api_source_file}"
diff --git a/paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc
@@ -0,0 +1,146 @@
+// Copyright (c) 2023 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/fusion/fused_dropout_add_pass.h"
+
+#include "paddle/fluid/pir/drr/api/drr_pattern_base.h"
+#include "paddle/pir/pass/pass.h"
+#include "paddle/pir/pass/pass_registry.h"
+#include "paddle/pir/pattern_rewrite/pattern_rewrite_driver.h"
+
+namespace {
+
+class FusedDropoutAddPattern
+    : public pir::drr::DrrPatternBase<FusedDropoutAddPattern> {
+ public:
+  void operator()(pir::drr::DrrPatternContext *ctx) const override {
+    pir::drr::SourcePattern pat = ctx->SourcePattern();
+    const auto &dropout = pat.Op("pd_op.dropout",
+                                 {{"p", pat.Attr("p")},
+                                  {"is_test", pat.Attr("is_test")},
+                                  {"mode", pat.Attr("mod")},
+                                  {"seed", pat.Attr("seed")},
+                                  {"fix_seed", pat.Attr("fix_seed")}});
+    const auto &add = pat.Op("pd_op.add");
+
+    dropout({&pat.Tensor("x"), &pat.Tensor("seed_tensor")},
+            {&pat.Tensor("dropout_out"), &pat.Tensor("mask")});
+    pat.Tensor("add_out") = add(pat.Tensor("dropout_out"), pat.Tensor("y"));
+
+    pir::drr::ResultPattern res = pat.ResultPattern();
+    const auto &fused_dropout_add =
+        res.Op("pd_op.fused_dropout_add",
+               {{{"p", pat.Attr("p")},
+                 {"is_test", pat.Attr("is_test")},
+                 {"mode", pat.Attr("mod")},
+                 {"seed", pat.Attr("seed")},
+                 {"fix_seed", pat.Attr("fix_seed")}}});
+    fused_dropout_add(
+        {&res.Tensor("x"), &res.Tensor("y"), &res.Tensor("seed_tensor")},
+        {&res.Tensor("add_out"), &res.Tensor("mask")});
+  }
+};
+
+class FusedDropoutGradAddGradPattern
+    : public pir::drr::DrrPatternBase<FusedDropoutAddPattern> {
+ public:
+  void operator()(pir::drr::DrrPatternContext *ctx) const override {
+    pir::drr::SourcePattern pat = ctx->SourcePattern();
+    const auto &dropout = pat.Op("pd_op.dropout",
+                                 {{"p", pat.Attr("p")},
+                                  {"is_test", pat.Attr("is_test")},
+                                  {"mode", pat.Attr("mod")},
+                                  {"seed", pat.Attr("seed")},
+                                  {"fix_seed", pat.Attr("fix_seed")}});
+    const auto &add = pat.Op("pd_op.add");
+
+    const auto &add_grad = pat.Op("pd_op.add_grad");
+    const auto &dropout_grad = pat.Op("pd_op.dropout_grad",
+                                      {{"p", pat.Attr("p")},
+                                       {"is_test", pat.Attr("is_test")},
+                                       {"mode", pat.Attr("mod")}});
+
+    dropout({&pat.Tensor("x"), &pat.Tensor("seed_tensor")},
+            {&pat.Tensor("dropout_out"), &pat.Tensor("mask")});
+    pat.Tensor("add_out") = add(pat.Tensor("dropout_out"), pat.Tensor("y"));
+    add_grad({&pat.Tensor("dropout_out"),
+              &pat.Tensor("y"),
+              &pat.Tensor("add_out_grad")},
+             {&pat.Tensor("dropout_out_grad"), &pat.Tensor("y_grad")});
+    dropout_grad({&pat.Tensor("mask"), &pat.Tensor("dropout_out_grad")},
+                 {&pat.Tensor("x_grad")});
+
+    pir::drr::ResultPattern res = pat.ResultPattern();
+    const auto &fused_dropout_add =
+        res.Op("pd_op.fused_dropout_add",
+               {{{"p", pat.Attr("p")},
+                 {"is_test", pat.Attr("is_test")},
+                 {"mode", pat.Attr("mod")},
+                 {"seed", pat.Attr("seed")},
+                 {"fix_seed", pat.Attr("fix_seed")}}});
+
+    const auto &fused_dropout_add_grad =
+        res.Op("pd_op.fused_dropout_add_grad",
+               {{{"p", pat.Attr("p")},
+                 {"is_test", pat.Attr("is_test")},
+                 {"mode", pat.Attr("mod")},
+                 {"fix_seed", pat.Attr("fix_seed")}}});
+
+    fused_dropout_add(
+        {&res.Tensor("x"), &res.Tensor("y"), &res.Tensor("seed_tensor")},
+        {&res.Tensor("add_out"), &res.Tensor("mask")});
+    fused_dropout_add_grad({&res.Tensor("mask"), &res.Tensor("add_out_grad")},
+                           {&res.Tensor("x_grad"), &res.Tensor("y_grad")});
+  }
+};
+
+class FusedDropoutAddPass : public pir::Pass {
+ public:
+  FusedDropoutAddPass() : pir::Pass("fused_dropout_add_pass", 1) {}
+
+  bool Initialize(pir::IrContext *context) override {
+    pir::RewritePatternSet ps(context);
+
+    ps.Add(FusedDropoutAddPattern().Build(context));
+    ps.Add(FusedDropoutGradAddGradPattern().Build(context));
+    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
+    return true;
+  }
+
+  void Run(pir::Operation *op) override {
+    pir::GreedyRewriteConfig cfg;
+    cfg.use_top_down_traversal = true;
+    cfg.max_iterations = 10;
+    pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
+  }
+
+  bool CanApplyOn(pir::Operation *op) const override {
+    return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
+  }
+
+ private:
+  pir::FrozenRewritePatternSet patterns_;
+};
+
+}  // namespace
+
+namespace pir {
+
+std::unique_ptr<Pass> CreateFusedDropoutAddPass() {
+  return std::make_unique<FusedDropoutAddPass>();
+}
+
+}  // namespace pir
+
+REGISTER_IR_PASS(fused_dropout_add_pass, FusedDropoutAddPass);
diff --git a/paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.h b/paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.h
@@ -0,0 +1,26 @@
+// Copyright (c) 2023 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/core/dll_decl.h"
+
+namespace pir {
+
+class Pass;
+
+IR_API std::unique_ptr<Pass> CreateFusedDropoutAddPass();
+
+}  // namespace pir
diff --git a/paddle/fluid/pybind/CMakeLists.txt b/paddle/fluid/pybind/CMakeLists.txt
@@ -42,6 +42,7 @@ set(PYBIND_DEPS
     pd_op_dialect
     program_translator
     pd_inplace_pass
+    fusion_passes
     pir
     new_profiler
     jit_layer
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -36,6 +36,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
+#include "paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.h"
 #include "paddle/fluid/pir/transforms/inplace_pass.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/pir/core/block.h"
@@ -66,6 +67,7 @@ using pir::Value;
 using pybind11::return_value_policy;
 
 USE_PASS(dead_code_elimination_pass);
+USE_PASS(fused_dropout_add_pass);
 USE_PASS(inplace_pass);
 
 PHI_DECLARE_bool(print_ir);
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -105,12 +105,12 @@ def prepare_grad_outputs(grad_outputs, outputs, state):
             if output.shape != grad.shape:
                 raise ValueError(
                     "The shape of grad_output[%d] %s should be the same as the shape of output[%d] %s"
-                    % (i, str(output.shape), i, str(grad.shape))
+                    % (i, str(grad.shape), i, str(output.shape))
                 )
             if output.dtype != grad.dtype:
                 raise ValueError(
                     "The dtype of grad_output[%d] %s should be the same as the dtype of output[%d] %s"
-                    % (i, str(output.dtype), i, str(grad.dtype))
+                    % (i, str(grad.dtype), i, str(output.dtype))
                 )
             feedop = grad.get_defining_op()
             update_bwdop_structure(
diff --git a/python/paddle/incubate/nn/functional/fused_dropout_add.py b/python/paddle/incubate/nn/functional/fused_dropout_add.py
@@ -16,7 +16,7 @@
 from paddle import _C_ops
 from paddle.base import core
 from paddle.common_ops_import import default_main_program
-from paddle.framework import LayerHelper, in_dynamic_mode
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
 
 
 def fused_dropout_add(
@@ -84,7 +84,7 @@ def fused_dropout_add(
             "mode argument should be 'downscale_in_infer' or 'upscale_in_train'"
         )
     seed = None
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         if default_main_program().random_seed != 0:
             seed = default_main_program().random_seed
         out, seed_offset = _C_ops.fused_dropout_add(
diff --git a/test/ir/new_ir/CMakeLists.txt b/test/ir/new_ir/CMakeLists.txt
@@ -19,3 +19,5 @@ foreach(target ${TEST_IR_SYSTEM_CASES})
 endforeach()
 
 set_tests_properties(test_pd_inplace_pass PROPERTIES TIMEOUT 60)
+
+add_subdirectory(fused_pass)
diff --git a/test/ir/new_ir/fused_pass/CMakeLists.txt b/test/ir/new_ir/fused_pass/CMakeLists.txt
@@ -0,0 +1,5 @@
+file(
+  GLOB TEST_INTERP_CASES
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_INTERP_CASES "${TEST_INTERP_CASES}")
diff --git a/test/ir/new_ir/fused_pass/test_fused_dropout_add.py b/test/ir/new_ir/fused_pass/test_fused_dropout_add.py
