PaddlePaddle
diff --git a/‎paddle/fluid/inference/api/analysis_predictor.cc
Copy file name to clipboard
+31-11Lines changed: 31 additions & 11 deletions b/‎paddle/fluid/inference/api/analysis_predictor.cc
Copy file name to clipboard
+31-11Lines changed: 31 additions & 11 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/layout_transformation.cc
Copy file name to clipboard
+53Lines changed: 53 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/layout_transformation.cc
Copy file name to clipboard
+53Lines changed: 53 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/layout_transformation.hpp
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/layout_transformation.hpp
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
@@ -131,7 +131,7 @@
 #include "paddle/pir/include/pass/pass_registry.h"
 
 COMMON_DECLARE_bool(pir_apply_inplace_pass);
-
+COMMON_DECLARE_bool(enable_auto_layout_pass);
 namespace paddle {
 namespace {
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
@@ -804,6 +804,20 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
                      pass->name()) != this->config_.ir_debug_passes_.end();
   };
 
+  auto AddAutoLayoutPasses = [&](pir::PassManager &pass_manager) {
+    auto &pass_registry = pir::PassRegistry::Instance();
+    std::vector<std::string> passes = {"auto_layout_pass",
+                                       "auto_layout_simplify_pass"};
+
+    for (const auto &pass_name : passes) {
+      if (std::find(config_.deleted_passes_.begin(),
+                    config_.deleted_passes_.end(),
+                    pass_name) == config_.deleted_passes_.end()) {
+        pass_manager.AddPass(pass_registry.Get(pass_name));
+      }
+    }
+  };
+
   auto AddAutoMixedPrecisionPass = [&](pir::PassManager &pass_manager) {
     auto auto_mixed_precision_pass = ::pir::CreateAutoMixedPrecisionPass();
     if (std::find(config_.deleted_passes_.begin(),
@@ -867,10 +881,13 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
 
         if (config_.enable_gpu_mixed_) {
           AddAutoMixedPrecisionPass(fused_op_pm);
-          fused_op_pm.AddPass(
-              pir::PassRegistry::Instance().Get("transfer_layout_pass"));
+          if (FLAGS_enable_auto_layout_pass) {
+            AddAutoLayoutPasses(fused_op_pm);
+          } else {
+            fused_op_pm.AddPass(
+                pir::PassRegistry::Instance().Get("transfer_layout_pass"));
+          }
         }
-
         fused_op_pm.Run(pir_program_.get());
       }
     }
@@ -909,7 +926,6 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
           }
         }
       }
-
 #ifdef PADDLE_WITH_XPU
     } else if (config_.use_xpu()) {
       // xpu
@@ -999,12 +1015,16 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
     if (!config_.cinn_enabled()) {
       AddAutoMixedPrecisionPass(basic_pass_pm);
 
-      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));
+      if (FLAGS_enable_auto_layout_pass) {
+        AddAutoLayoutPasses(basic_pass_pm);
+      } else {
+        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));
+        }
       }
     }
   }
 
@@ -102,6 +102,32 @@ common::DataLayout PreferLayoutImpl<Conv2dOp>(pir::Operation* op) {
   return common::StringToDataLayout(data_format_attr.AsString());
 }
 
+template <>
+common::DataLayout PreferLayoutImpl<Conv2dTransposeOp>(pir::Operation* op) {
+  auto data_format_attr = op->attribute<pir::StrAttribute>("data_format");
+  if (!data_format_attr) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "op (%s) should have attribute `data_format`, but got %s",
+        op,
+        data_format_attr));
+  }
+
+  auto concrete_op = op->dyn_cast<Conv2dTransposeOp>();
+  if (auto in = concrete_op.x()) {
+    if (auto in_type = in.type()) {
+      if (in_type.isa<DenseTensorType>()) {
+        if (auto tensor_type = in_type.dyn_cast<DenseTensorType>()) {
+          if (tensor_type.dtype().isa<pir::Float16Type>()) {
+            return common::DataLayout::NHWC;
+          }
+        }
+      }
+    }
+  }
+
+  return common::StringToDataLayout(data_format_attr.AsString());
+}
+
 template <>
 bool CanBeModifiedImpl<Conv2dOp>(pir::Operation* op) {
   return false;
@@ -352,6 +378,33 @@ void RewriteByLayoutImpl<ConcatOp>(pir::Operation* op,
   RewriteByInfermeta<ConcatOp>(op, new_layout);
 }
 
+template <>
+void RewriteByLayoutImpl<ArgmaxOp>(pir::Operation* op,
+                                   common::DataLayout new_layout) {
+  auto concrete_op = op->dyn_cast<ArgmaxOp>();
+  auto axis = concrete_op.axis();
+  if (!axis || !(axis.defining_op()->isa<FullOp>())) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Argmax's axis must be processed when rewrite by layout."));
+  }
+
+  auto axis_op = axis.defining_op()->dyn_cast<FullOp>();
+  int axis_value =
+      axis_op.attribute("value").dyn_cast<ScalarAttribute>().data().to<int>();
+
+  PADDLE_ENFORCE_EQ(
+      axis_value,
+      1,
+      common::errors::InvalidArgument(
+          "Argmax's axis was expected as 1, but got %d", axis_value));
+  axis.defining_op()->set_attribute(
+      "value",
+      ScalarAttribute::get(pir::IrContext::Instance(), phi::Scalar(3)));
+
+  // infer new meta for argmax
+  RewriteByInfermeta<ArgmaxOp>(op, new_layout);
+}
+
 template <>
 void RewriteByLayoutImpl<pir::CombineOp>(pir::Operation* op,
                                          common::DataLayout new_layout) {
 
@@ -118,6 +118,9 @@ OVERLOAD_PREFER_LAYOUT(Conv2dOp);
 OVERLOAD_CAN_BE_MODIFIED(Conv2dOp);
 OVERLOAD_REWRITE_BY_LAYOUT(Conv2dOp);
 
+class Conv2dTransposeOp;
+OVERLOAD_PREFER_LAYOUT(Conv2dTransposeOp);
+
 class GroupNormOp;
 OVERLOAD_REWRITE_BY_LAYOUT(GroupNormOp);
 OVERLOAD_RELEVANT_INPUTS(GroupNormOp);
@@ -154,6 +157,9 @@ class ConcatOp;
 OVERLOAD_REWRITE_BY_LAYOUT(ConcatOp);
 OVERLOAD_RELEVANT_INPUTS(ConcatOp);
 
+class ArgmaxOp;
+OVERLOAD_REWRITE_BY_LAYOUT(ArgmaxOp);
+
 class Pool2dOp;
 OVERLOAD_RELEVANT_INPUTS(Pool2dOp);
 OVERLOAD_REWRITE_BY_LAYOUT(Pool2dOp);