[Inference] Fix pir amp pass order (#69205)

yuanlehome · web-flow · commit 98bbdfc4f403 · 2024-11-06T22:07:04.000-08:00
* fix pir amp pass order

* update

* update
diff --git a/paddle/fluid/framework/feed_fetch_method.cc b/paddle/fluid/framework/feed_fetch_method.cc
@@ -19,7 +19,6 @@ limitations under the License. */
 #include "glog/logging.h"
 
 COMMON_DECLARE_bool(enable_pir_in_executor);
-COMMON_DECLARE_bool(enable_pir_api);
 
 namespace phi {
 class DenseTensor;
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -123,6 +123,7 @@
 #include "paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.h"
 #include "paddle/fluid/pir/transforms/general/remove_shadow_feed_pass.h"
 #include "paddle/fluid/pir/transforms/general/replace_fetch_with_shadow_output_pass.h"
+#include "paddle/fluid/pir/transforms/general/transfer_layout_pass.h"
 #include "paddle/fluid/pir/transforms/passes.h"
 #include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
 #include "paddle/fluid/pir/utils/general_functions.h"
@@ -135,7 +136,6 @@
 #include "paddle/pir/include/pass/pass_registry.h"
 
 COMMON_DECLARE_bool(pir_apply_inplace_pass);
-COMMON_DECLARE_bool(enable_pir_api);
 
 namespace paddle {
 namespace {
@@ -874,10 +874,6 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
           if (std::find(config_.deleted_passes_.begin(),
                         config_.deleted_passes_.end(),
                         gpu_pass) == config_.deleted_passes_.end()) {
-            if (!config_.enable_gpu_mixed_ &&
-                gpu_pass == "auto_mixed_precision_pass") {
-              continue;
-            }
             pass_pm.AddPass(pir::PassRegistry::Instance().Get(gpu_pass));
           }
         }
@@ -937,22 +933,11 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
     for (const auto &pass : pass_pm.passes()) {
       pass->SetNotOwned(pir::Pass::kParamScopeAttr, sub_scope_);
       pass->SetNotOwned(pir::Pass::kPlaceAttr, &place_);
+      pass->Set("enable_gpu_mixed", new bool(config_.enable_gpu_mixed_));
       if (pass->name() == "matmul_add_act_fuse_pass" ||
           pass->name() == "conv2d_add_act_fuse_pass" ||
           pass->name() == "conv2d_add_fuse_pass") {
         pass->Set("use_cutlass", new bool(config_.use_cutlass_));
-      } else if (pass->name() == "auto_mixed_precision_pass") {
-        pass->Set("mixed_precision_mode",
-                  new phi::DataType(
-                      paddle::ConvertPrecision(config_.mixed_precision_mode_)));
-        pass->Set("enable_low_precision_io",
-                  new bool(config_.enable_low_precision_io_));
-        pass->Set(
-            "mixed_black_list",
-            new std::unordered_set<std::string>(config_.mixed_black_list_));
-        pass->Set(
-            "mixed_white_list",
-            new std::unordered_set<std::string>(config_.mixed_white_list_));
       }
     }
 
@@ -979,13 +964,34 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
   // Apply some basic passes required by the framework
   ::pir::PassManager basic_pass_pm(::pir::IrContext::Instance(),
                                    config_.pm_opt_level_);
-  auto common_subexpression_elimination_pass =
-      ::pir::CreateCommonSubexpressionEliminationPass();
-  if (std::find(config_.deleted_passes_.begin(),
-                config_.deleted_passes_.end(),
-                common_subexpression_elimination_pass->name()) ==
-      config_.deleted_passes_.end()) {
-    basic_pass_pm.AddPass(std::move(common_subexpression_elimination_pass));
+  if (config_.enable_gpu_mixed_) {
+    auto auto_mixed_precision_pass = ::pir::CreateAutoMixedPrecisionPass();
+    if (std::find(config_.deleted_passes_.begin(),
+                  config_.deleted_passes_.end(),
+                  auto_mixed_precision_pass->name()) ==
+        config_.deleted_passes_.end()) {
+      auto_mixed_precision_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place_);
+      auto_mixed_precision_pass->Set("mixed_precision_mode",
+                                     new phi::DataType(paddle::ConvertPrecision(
+                                         config_.mixed_precision_mode_)));
+      auto_mixed_precision_pass->Set(
+          "enable_low_precision_io",
+          new bool(config_.enable_low_precision_io_));
+      auto_mixed_precision_pass->Set(
+          "mixed_black_list",
+          new std::unordered_set<std::string>(config_.mixed_black_list_));
+      auto_mixed_precision_pass->Set(
+          "mixed_white_list",
+          new std::unordered_set<std::string>(config_.mixed_white_list_));
+      basic_pass_pm.AddPass(std::move(auto_mixed_precision_pass));
+    }
+    auto transfer_layout_pass = ::pir::CreateTransferLayoutPass();
+    if (std::find(config_.deleted_passes_.begin(),
+                  config_.deleted_passes_.end(),
+                  transfer_layout_pass->name()) ==
+        config_.deleted_passes_.end()) {
+      basic_pass_pm.AddPass(std::move(transfer_layout_pass));
+    }
   }
   auto params_sync_among_devices_pass =
       ::pir::CreateParamsSyncAmongDevicesPass();
@@ -1016,6 +1022,14 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
                                             sub_scope_);
     basic_pass_pm.AddPass(std::move(dead_code_elimination_pass));
   }
+  auto common_subexpression_elimination_pass =
+      ::pir::CreateCommonSubexpressionEliminationPass();
+  if (std::find(config_.deleted_passes_.begin(),
+                config_.deleted_passes_.end(),
+                common_subexpression_elimination_pass->name()) ==
+      config_.deleted_passes_.end()) {
+    basic_pass_pm.AddPass(std::move(common_subexpression_elimination_pass));
+  }
   auto replace_fetch_with_shadow_output_pass =
       ::pir::CreateReplaceFetchWithShadowOutputPass();
   if (std::find(config_.deleted_passes_.begin(),
diff --git a/paddle/fluid/inference/api/paddle_pass_builder.cc b/paddle/fluid/inference/api/paddle_pass_builder.cc
@@ -587,7 +587,6 @@ const std::vector<std::string> kPirGpuPasses{
     "delete_weight_dequant_linear_op_pass",
     "map_op_to_another_pass",
     "identity_op_clean_pass",
-    "auto_mixed_precision_pass",
     // Operator fusion pass
     "silu_fuse_pass",
     "conv2d_bn_fuse_pass",
@@ -606,9 +605,7 @@ const std::vector<std::string> kPirGpuPasses{
     "matmul_transpose_fuse_pass",
     "transpose_flatten_concat_fuse_pass",
     "remove_redundant_transpose_pass",
-    "transfer_layout_pass",
     "horizontal_fuse_pass",
-    "auto_mixed_precision_pass",
 };
 
 const std::vector<std::string> kPirXpuPasses{
diff --git a/paddle/fluid/pir/transforms/general/group_norm_silu_fuse_pass.cc b/paddle/fluid/pir/transforms/general/group_norm_silu_fuse_pass.cc
@@ -57,7 +57,13 @@ After the pass is applied:GPU
 namespace {
 
 class GroupNormSiluPattern : public paddle::drr::DrrPatternBase {
+ private:
+  const bool enable_gpu_mixed_;
+
  public:
+  explicit GroupNormSiluPattern(bool enable_gpu_mixed)
+      : enable_gpu_mixed_(enable_gpu_mixed) {}
+
   std::string name() const override { return "GroupNormSiluPattern"; }
 
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
@@ -75,9 +81,9 @@ class GroupNormSiluPattern : public paddle::drr::DrrPatternBase {
     silu({&pat.Tensor("Y")}, {&pat.Tensor("Out")});
 
 #ifdef PADDLE_WITH_CUDA
-    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) {
+    pat.AddConstraint([this](const paddle::drr::MatchContext &match_ctx) {
       auto x_dtype = pir::GetDataTypeFromValue(match_ctx.Tensor("X"));
-      if (!x_dtype.isa<pir::Float16Type>() &&
+      if (!this->enable_gpu_mixed_ && !x_dtype.isa<pir::Float16Type>() &&
           !x_dtype.isa<pir::BFloat16Type>()) {
         return false;
       }
@@ -121,7 +127,12 @@ class GroupNormSiluFusePass : public pir::PatternRewritePass {
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
-    ps.Add(paddle::drr::Create<GroupNormSiluPattern>(context));
+    bool enable_gpu_mixed = false;
+    if (Has("enable_gpu_mixed")) {
+      enable_gpu_mixed = Get<bool>("enable_gpu_mixed");
+    }
+    ps.Add(
+        paddle::drr::Create<GroupNormSiluPattern>(context, enable_gpu_mixed));
     return ps;
   }
 };
diff --git a/paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.h b/paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.h
@@ -15,7 +15,6 @@
 #pragma once
 
 #include <memory>
-#include "paddle/phi/common/place.h"
 #include "paddle/pir/include/core/dll_decl.h"
 
 namespace paddle {
diff --git a/paddle/fluid/pir/transforms/gpu/add_norm_fuse_pass.cc b/paddle/fluid/pir/transforms/gpu/add_norm_fuse_pass.cc
@@ -288,10 +288,15 @@ class AddGroupNormFusePattern : public paddle::drr::DrrPatternBase {
  private:
   const bool extra_add_;
   const bool trans_extra_add_;
+  const bool enable_gpu_mixed_;
 
  public:
-  AddGroupNormFusePattern(bool extra_add, bool trans_extra_add)
-      : extra_add_(extra_add), trans_extra_add_{trans_extra_add} {}
+  AddGroupNormFusePattern(bool extra_add,
+                          bool trans_extra_add,
+                          bool enable_gpu_mixed)
+      : extra_add_(extra_add),
+        trans_extra_add_{trans_extra_add},
+        enable_gpu_mixed_(enable_gpu_mixed) {}
 
   uint32_t benefit() const override { return extra_add_ ? 4 : 3; }
   std::string name() const override { return "AddGroupNormFusePattern"; }
@@ -318,9 +323,9 @@ class AddGroupNormFusePattern : public paddle::drr::DrrPatternBase {
               ? add1(pat.Tensor("any_tensor"), pat.Tensor("add_out"))
               : add1(pat.Tensor("add_out"), pat.Tensor("any_tensor"));
     }
-    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) {
+    pat.AddConstraint([this](const paddle::drr::MatchContext &match_ctx) {
       auto x_dtype = pir::GetDataTypeFromValue(match_ctx.Tensor("x"));
-      if (!x_dtype.isa<pir::Float16Type>() &&
+      if (!this->enable_gpu_mixed_ && !x_dtype.isa<pir::Float16Type>() &&
           !x_dtype.isa<pir::BFloat16Type>()) {
         return false;
       }
@@ -346,7 +351,13 @@ class AddGroupNormFusePattern : public paddle::drr::DrrPatternBase {
 };
 
 class AddGroupNormWithActPattern : public paddle::drr::DrrPatternBase {
+ private:
+  const bool enable_gpu_mixed_;
+
  public:
+  explicit AddGroupNormWithActPattern(bool enable_gpu_mixed)
+      : enable_gpu_mixed_(enable_gpu_mixed) {}
+
   uint32_t benefit() const override { return 2; }
   std::string name() const override { return "AddGroupNormWithActPattern"; }
 
@@ -368,9 +379,9 @@ class AddGroupNormWithActPattern : public paddle::drr::DrrPatternBase {
                             &pat.Tensor("mean_out_0"),
                             &pat.Tensor("variance_out_0")});
     pat.Tensor("silu_out") = silu(pat.Tensor("group_out"));
-    pat.AddConstraint([](const paddle::drr::MatchContext &match_ctx) {
+    pat.AddConstraint([this](const paddle::drr::MatchContext &match_ctx) {
       auto x_dtype = pir::GetDataTypeFromValue(match_ctx.Tensor("x"));
-      if (!x_dtype.isa<pir::Float16Type>() &&
+      if (!this->enable_gpu_mixed_ && !x_dtype.isa<pir::Float16Type>() &&
           !x_dtype.isa<pir::BFloat16Type>()) {
         return false;
       }
@@ -404,6 +415,12 @@ class AddNormFusePass : public pir::PatternRewritePass {
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
+
+    bool enable_gpu_mixed = false;
+    if (Has("enable_gpu_mixed")) {
+      enable_gpu_mixed = Get<bool>("enable_gpu_mixed");
+    }
+
     // x-pow-mean-scale->rsqrt-
     //                          mul--
     // x-----------------------
@@ -437,14 +454,15 @@ class AddNormFusePass : public pir::PatternRewritePass {
     //           add-group_norm ----> add_group_norm_silu
     // residual-
     ps.Add(paddle::drr::Create<AddGroupNormFusePattern>(
-        context, !extra_add, true));
-    ps.Add(
-        paddle::drr::Create<AddGroupNormFusePattern>(context, extra_add, true));
+        context, !extra_add, true, enable_gpu_mixed));
     ps.Add(paddle::drr::Create<AddGroupNormFusePattern>(
-        context, extra_add, false));
+        context, extra_add, true, enable_gpu_mixed));
+    ps.Add(paddle::drr::Create<AddGroupNormFusePattern>(
+        context, extra_add, false, enable_gpu_mixed));
 
     // add_group_norm_silu-silu --->add_group_norm_silu
-    ps.Add(paddle::drr::Create<AddGroupNormWithActPattern>(context));
+    ps.Add(paddle::drr::Create<AddGroupNormWithActPattern>(context,
+                                                           enable_gpu_mixed));
     // group-silu->add_group_norm_silu moved to group_norm_silu_fuse_pass
     return ps;
   }
