[Inference] Fix some model run in mixed precision (#60955)

yuanlehome · web-flow · commit 6e0fb7a149b6 · 2024-01-19T19:58:36.000+08:00
* fix some model run in mixed model

* fix

* update

* update
diff --git a/paddle/fluid/framework/ir/CMakeLists.txt b/paddle/fluid/framework/ir/CMakeLists.txt
@@ -111,7 +111,7 @@ pass_library(conv2d_trans_filter_dilations_nxn_to_1x1_pass inference)
 pass_library(preln_residual_bias_fuse_pass inference)
 pass_library(constant_folding_pass inference)
 pass_library(auto_mixed_precision_pass inference)
-pass_library(fused_conv2d_add_act_layout_transfer_pass inference)
+pass_library(transfer_layout_pass inference)
 pass_library(transfer_layout_elim_pass inference)
 pass_library(relu6_fuse_pass inference)
 pass_library(silu_fuse_pass inference)
diff --git a/paddle/fluid/framework/ir/map_op_to_another_pass.cc b/paddle/fluid/framework/ir/map_op_to_another_pass.cc
@@ -30,6 +30,7 @@ void MapOp2AnotherPass::ApplyImpl(ir::Graph* graph) const {
 
   int found_count = 0;
   std::unordered_map<std::string, std::string> replaced_map{
+      {"conv2d", "conv2d"},
       {"depthwise_conv2d", "conv2d"},
       {"flatten_contiguous_range", "reshape2"},
   };
@@ -61,6 +62,10 @@ void MapOp2AnotherPass::ApplyImpl(ir::Graph* graph) const {
         op_desc->SetAttr("use_cudnn", true);
 #endif
       }
+    } else if (op_type == "conv2d") {
+      op_desc->SetType(replaced_map[op_type]);
+      op_desc->RemoveAttr("use_cudnn");
+      op_desc->SetAttr("use_cudnn", true);
     }
     op_desc->Flush();
     ++found_count;
diff --git a/paddle/fluid/framework/ir/transfer_layout_elim_pass.cc b/paddle/fluid/framework/ir/transfer_layout_elim_pass.cc
@@ -343,8 +343,12 @@ void TransferLayoutElimPass::ApplyImpl(ir::Graph *graph) const {
     }
     if (!modify) break;
   }
-  LOG(INFO) << "move down " << move_down_count << " transfer_layout";
-  LOG(INFO) << "eliminate " << elim_count << " pair of transfer_layout";
+  if (move_down_count > 0) {
+    LOG(INFO) << "move down " << move_down_count << " transfer_layout";
+  }
+  if (elim_count > 0) {
+    LOG(INFO) << "eliminate " << elim_count << " pair of transfer_layout";
+  }
 }
 
 }  // namespace ir
diff --git a/paddle/fluid/framework/ir/transfer_layout_pass.cc b/paddle/fluid/framework/ir/transfer_layout_pass.cc
@@ -12,7 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/framework/ir/fused_conv2d_add_act_layout_transfer_pass.h"
+#include "paddle/fluid/framework/ir/transfer_layout_pass.h"
+#include <algorithm>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -99,7 +100,7 @@ void InsertLayoutTransOp(ir::Graph *graph,
 
 }  // namespace
 
-void FusedConv2dAddActLayoutTransferPass::ApplyImpl(ir::Graph *graph) const {
+void TransferLayoutPass::ApplyImpl(ir::Graph *graph) const {
   PADDLE_ENFORCE_NOT_NULL(
       graph,
       platform::errors::PreconditionNotMet("graph should not be nullptr."));
@@ -118,12 +119,12 @@ void FusedConv2dAddActLayoutTransferPass::ApplyImpl(ir::Graph *graph) const {
                     true,
                     platform::errors::InvalidArgument(
                         "the graph should be main graph when applying "
-                        "fused_conv2d_add_act_layout_transfer_pass"));
+                        "transfer_layout_pass"));
 
   PADDLE_ENFORCE_NOT_NULL(
       scope,
       platform::errors::Fatal("scope must not be nullptr when applying "
-                              "fused_conv2d_add_act_layout_transfer_pass"));
+                              "transfer_layout_pass"));
 
   // Not support multiple block now.
   std::unordered_map<ir::Node *, ir::Node *> cache;
@@ -132,7 +133,7 @@ void FusedConv2dAddActLayoutTransferPass::ApplyImpl(ir::Graph *graph) const {
   auto *block_desc = (*iter)->Op()->Block();
 
   // Process multiple fused_conv2d_add_act shares weight.
-  std::unordered_set<std::string> weights_shape_nhwc;
+  std::unordered_map<std::string, std::vector<ir::Node *>> weights_shared;
 
   // Used to control the insertion of transfer_layout op.
   std::unordered_set<ir::Node *> vars_shape_nhwc;
@@ -179,55 +180,87 @@ void FusedConv2dAddActLayoutTransferPass::ApplyImpl(ir::Graph *graph) const {
     if (op_node->Op()->Type() != target_op_type) {
       continue;
     }
-    auto filter_name = op_node->Op()->Input("Filter").front();
-    if (weights_shape_nhwc.count(filter_name)) {
-      continue;
-    }
+
     auto data_format =
         op_node->Op()->GetAttrIfExists<std::string>("data_format");
-
     if (data_format != "NCHW") {
       continue;
     }
 
+    auto filter_names = op_node->Op()->Input("Filter");
+    for (const auto &filter_name : filter_names) {
+      weights_shared[filter_name].push_back(op_node);
+    }
+
     if (cuDNNIsValid(op_node) || CutlassIsValid(op_node)) {
       valid_ops.insert(op_node);
+    }
+  }
+
+  // The target operators that share weights either run nhwc or not at all.
+  for (auto *op_node : op_nodes) {
+    CHECK_EQ(op_node->IsOp(), true);
+    if (valid_ops.count(op_node)) {
+      auto filter_names = op_node->Op()->Input("Filter");
+      for (const auto &filter_name : filter_names) {
+        if (std::any_of(weights_shared[filter_name].begin(),
+                        weights_shared[filter_name].end(),
+                        [&valid_ops](ir::Node *node) {
+                          return valid_ops.count(node) == 0;
+                        })) {
+          for (auto *node : weights_shared[filter_name]) {
+            valid_ops.erase(node);
+          }
+        }
+      }
+    }
+  }
+
+  // Insert transfer_layout op
+  for (auto *op_node : op_nodes) {
+    CHECK_EQ(op_node->IsOp(), true);
+
+    if (valid_ops.count(op_node)) {
       auto *op_desc = op_node->Op();
 
       if (CutlassIsValid(op_node)) {
         // fused_conv2d_add_act must have this attribute because of signature.
         if (!op_desc->HasAttr("fuse_alpha")) {
           op_desc->SetAttr("fuse_alpha", 0.f);
         }
+      } else if (cuDNNIsValid(op_node)) {
+        if (op_desc->HasAttr("use_cudnn")) {
+          op_desc->SetAttr("use_cudnn", true);
+        }
       }
       op_desc->SetAttr("data_format", std::string{"NHWC"});
       op_desc->Flush();
 
       // transfer weights
       auto filter_names = op_desc->Input("Filter");
       for (const auto &filter_name : filter_names) {
-        if (weights_shape_nhwc.count(filter_name) == 0) {
-          weights_shape_nhwc.insert(filter_name);
-          auto *filter_var = scope->FindLocalVar(filter_name);
-          auto *filter_tensor = filter_var->GetMutable<phi::DenseTensor>();
-          phi::DenseTensor temp_tensor;
-
-          framework::TransDataLayout(phi::DataLayout::kNCHW,
-                                     phi::DataLayout::kNHWC,
-                                     phi::CPUPlace{},
-                                     *filter_tensor,
-                                     &temp_tensor);
-          *filter_tensor = temp_tensor;
-
-          auto op_inputs = op_node->inputs;
-          for (auto *in_var_node : op_inputs) {
-            CHECK_EQ(in_var_node->IsVar(), true);
-            if (in_var_node->Var()->Persistable() &&
-                in_var_node->Var()->Name() == filter_name) {
-              auto from_shape = in_var_node->Var()->GetShape();
-              in_var_node->Var()->SetShape(
-                  {from_shape[0], from_shape[2], from_shape[3], from_shape[1]});
-            }
+        auto *filter_var = scope->FindLocalVar(filter_name);
+        auto *filter_tensor = filter_var->GetMutable<phi::DenseTensor>();
+        if (filter_tensor->layout() == phi::DataLayout::kNHWC) {
+          continue;
+        }
+        phi::DenseTensor temp_tensor;
+
+        framework::TransDataLayout(phi::DataLayout::kNCHW,
+                                   phi::DataLayout::kNHWC,
+                                   phi::CPUPlace{},
+                                   *filter_tensor,
+                                   &temp_tensor);
+        *filter_tensor = temp_tensor;
+
+        auto op_inputs = op_node->inputs;
+        for (auto *in_var_node : op_inputs) {
+          CHECK_EQ(in_var_node->IsVar(), true);
+          if (in_var_node->Var()->Persistable() &&
+              in_var_node->Var()->Name() == filter_name) {
+            auto from_shape = in_var_node->Var()->GetShape();
+            in_var_node->Var()->SetShape(
+                {from_shape[0], from_shape[2], from_shape[3], from_shape[1]});
           }
         }
       }
@@ -243,14 +276,8 @@ void FusedConv2dAddActLayoutTransferPass::ApplyImpl(ir::Graph *graph) const {
             {from_shape[0], from_shape[2], from_shape[3], from_shape[1]});
         vars_shape_nhwc.insert(out_var_node);
       }
-    }
-  }
-
-  // Insert transfer_layout op
-  for (auto *op_node : op_nodes) {
-    CHECK_EQ(op_node->IsOp(), true);
 
-    if (valid_ops.count(op_node)) {
+      // Insert transfer_layout for intermidiate var.
       auto op_inputs = op_node->inputs;
       for (auto *in_var_node : op_inputs) {
         CHECK_EQ(in_var_node->IsVar(), true);
@@ -283,11 +310,11 @@ void FusedConv2dAddActLayoutTransferPass::ApplyImpl(ir::Graph *graph) const {
       }
     }
   }
+  AddStatis(static_cast<int>(valid_ops.size()));
 }
 
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle
 
-REGISTER_PASS(fused_conv2d_add_act_layout_transfer_pass,
-              paddle::framework::ir::FusedConv2dAddActLayoutTransferPass);
+REGISTER_PASS(transfer_layout_pass, paddle::framework::ir::TransferLayoutPass);
diff --git a/paddle/fluid/framework/ir/transfer_layout_pass.h b/paddle/fluid/framework/ir/transfer_layout_pass.h
@@ -20,10 +20,10 @@ namespace paddle {
 namespace framework {
 namespace ir {
 
-class FusedConv2dAddActLayoutTransferPass : public FusePassBase {
+class TransferLayoutPass : public FusePassBase {
  public:
-  FusedConv2dAddActLayoutTransferPass() = default;
-  virtual ~FusedConv2dAddActLayoutTransferPass() = default;
+  TransferLayoutPass() = default;
+  virtual ~TransferLayoutPass() = default;
 
  protected:
   void ApplyImpl(ir::Graph* graph) const override;
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -802,7 +802,7 @@ bool AnalysisPredictor::PrepareExecutor() {
 
         //----------------------------------------------------------------------------------------------//
         // Operator fusion pass
-        gpu_pm.AddPass(::pir::CreateConv2dBnFusePass());
+        // gpu_pm.AddPass(::pir::CreateConv2dBnFusePass());
         gpu_pm.AddPass(::pir::CreateConv2dAddActFusePass());
         gpu_pm.AddPass(::pir::CreateConv2dAddFusePass());
         gpu_pm.AddPass(::pir::CreateMultiHeadMatmulFusePass());
@@ -1789,6 +1789,12 @@ void AnalysisPredictor::PrepareArgument() {
         model_precision_ == phi::DataType::FLOAT32) {
       argument_->SetEnableIrOptim(true);
       pass_builder->ClearPasses();
+      pass_builder->AppendPass("map_op_to_another_pass");
+      pass_builder->AppendPass("simplify_with_basic_ops_pass");
+      pass_builder->AppendPass("is_test_pass");
+      if (!FLAGS_enable_pir_in_executor) {
+        pass_builder->AppendPass("constant_folding_pass");
+      }
       pass_builder->AppendPass("auto_mixed_precision_pass");
       if (!FLAGS_enable_pir_in_executor) {
         pass_builder->AppendPass("inplace_op_var_pass");
diff --git a/paddle/fluid/inference/api/paddle_pass_builder.cc b/paddle/fluid/inference/api/paddle_pass_builder.cc
@@ -194,7 +194,7 @@ const std::vector<std::string> kGpuLowerPrecisionPasses{
     "conv_elementwise_add_act_fuse_pass",
     "conv_elementwise_add2_act_fuse_pass",
     "conv_elementwise_add_fuse_pass",
-    "fused_conv2d_add_act_layout_transfer_pass",
+    "transfer_layout_pass",
     "multihead_matmul_fuse_pass_v2",
     "fused_multi_transformer_encoder_pass",
     "fused_multi_transformer_decoder_pass",
@@ -304,10 +304,10 @@ GpuPassStrategy::GpuPassStrategy() : PassStrategy({}) {
         "conv_elementwise_add_act_fuse_pass",   //
         "conv_elementwise_add2_act_fuse_pass",  //
 #endif
-        "conv_elementwise_add_fuse_pass",             //
-#endif                                                //
-        "transpose_flatten_concat_fuse_pass",         //
-        "fused_conv2d_add_act_layout_transfer_pass",  //
+        "conv_elementwise_add_fuse_pass",      //
+#endif                                         //
+        "transpose_flatten_concat_fuse_pass",  //
+        "transfer_layout_pass",                //
         "transfer_layout_elim_pass",
         "auto_mixed_precision_pass",  //
         "identity_op_clean_pass",  // should be after auto_mixed_precision_pass.
diff --git a/paddle/fluid/pir/drr/ir_operation_factory.cc b/paddle/fluid/pir/drr/ir_operation_factory.cc
@@ -71,7 +71,7 @@ void OperationFactory::RegisterManualOpCreator() {
       });
 }
 
-static pir::Attribute CreateIrAttribute(const std::any& obj) {
+pir::Attribute CreateIrAttribute(const std::any& obj) {
   if (obj.type() == typeid(bool)) {
     return IrAttrbuteCreator<bool>()(std::any_cast<bool>(obj));
   } else if (obj.type() == typeid(int32_t)) {
diff --git a/paddle/fluid/pir/transforms/fusion/conv2d_bn_fuse_pass.cc b/paddle/fluid/pir/transforms/fusion/conv2d_bn_fuse_pass.cc
@@ -68,9 +68,8 @@ class Conv2dBnFusePattern
         rewriter.Build<paddle::dialect::DivideOp>(
             bn_scale.dyn_cast<pir::OpResult>(), sqrt_op.out());
     // reshape scale
-    phi::DDim conv2d_filter_shape = pir::GetShapeFromValue(conv2d_filter);
-    phi::DDim bn_scale_shape =
-        bn_scale.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+    auto conv2d_filter_shape = pir::GetShapeFromValue(conv2d_filter);
+    auto bn_scale_shape = pir::GetShapeFromValue(bn_scale);
     std::vector<int64_t> bn_scale_new_shape(conv2d_filter_shape.size(), 1);
     bn_scale_new_shape[0] = bn_scale_shape[0];
     paddle::dialect::ReshapeOp reshape_scale_op =
@@ -96,8 +95,7 @@ class Conv2dBnFusePattern
         rewriter.Build<paddle::dialect::SubtractOp>(
             bn_bias.dyn_cast<pir::OpResult>(), mul_bias_op.out());
     // reshape new bias
-    phi::DDim new_conv2d_out_shape =
-        pir::GetShapeFromValue(new_conv2d_op.out());
+    auto new_conv2d_out_shape = pir::GetShapeFromValue(new_conv2d_op.out());
     std::vector<int64_t> new_bias_new_shape(new_conv2d_out_shape.size(), 1);
     std::string data_format =
         new_conv2d_op.attribute<pir::StrAttribute>("data_format").AsString();
diff --git a/paddle/fluid/pir/transforms/fusion/fc_elementwise_layernorm_fuse_pass.cc b/paddle/fluid/pir/transforms/fusion/fc_elementwise_layernorm_fuse_pass.cc
@@ -52,7 +52,7 @@ class FcElementwiseLayerNormFusePattern : public paddle::drr::DrrPatternBase {
       int64_t layer_norm_x = 1;
       auto fc_out_dims = pir::GetShapeFromValue(match_ctx.Tensor("fc_out"));
       auto w_dims = pir::GetShapeFromValue(match_ctx.Tensor("w"));
-      for (int i = match_ctx.Attr<int>("begin_norm_axis");
+      for (size_t i = match_ctx.Attr<int>("begin_norm_axis");
            i < fc_out_dims.size();
            i++) {
         layer_norm_x *= fc_out_dims.at(i);
@@ -121,7 +121,7 @@ class FcElementwiseLayerNormFuse2Pattern : public paddle::drr::DrrPatternBase {
       int64_t layer_norm_x = 1;
       auto fc_out_dims = pir::GetShapeFromValue(match_ctx.Tensor("fc_out"));
       auto w_dims = pir::GetShapeFromValue(match_ctx.Tensor("w"));
-      for (int i = match_ctx.Attr<int>("begin_norm_axis");
+      for (size_t i = match_ctx.Attr<int>("begin_norm_axis");
            i < fc_out_dims.size();
            i++) {
         layer_norm_x *= fc_out_dims.at(i);
diff --git a/paddle/fluid/pir/transforms/fusion/fc_fuse_pass.cc b/paddle/fluid/pir/transforms/fusion/fc_fuse_pass.cc
@@ -60,7 +60,7 @@ class MatmulAddPattern : public paddle::drr::DrrPatternBase {
     const auto &in_num_col_dims_attr =
         res.ComputeAttr([](const paddle::drr::MatchContext &match_ctx) -> int {
           auto x_dims = pir::GetShapeFromValue(match_ctx.Tensor("x"));
-          return x_dims.size() - 1;
+          return static_cast<int>(x_dims.size()) - 1;
         });
     const auto &false_attr = res.BoolAttr(false);
 
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -37,6 +37,7 @@
 #include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_parser.h"
 #include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_util.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
+#include "paddle/fluid/pir/transforms/transform_general_functions.h"
 #include "paddle/fluid/platform/place.h"
 #include "paddle/phi/api/lib/data_transform.h"
 #include "paddle/phi/api/lib/kernel_dispatch.h"
@@ -77,18 +78,6 @@ pir::Type ConvertOpTypeToKernelType(pir::IrContext* ctx,
       "Not support op type %s in ConvertOpTypeToKernelType.", op_type));
 }
 
-std::vector<int64_t> GetValueShape(const pir::Value& value) {
-  if (value.type().isa<DenseTensorType>()) {
-    return phi::vectorize(value.type().dyn_cast<DenseTensorType>().dims());
-  } else if (value.type().isa<SelectedRowsType>()) {
-    return phi::vectorize(value.type().dyn_cast<SelectedRowsType>().dims());
-  } else {
-    PADDLE_THROW(phi::errors::InvalidArgument(
-        "Currently, we can only get shape for dense "
-        "tensor."));
-  }
-}
-
 static const std::vector<pir::Type> InferMetaByValue(
     pir::Operation* op,
     const std::vector<pir::Value>& input_values,
@@ -225,7 +214,7 @@ static bool NeedFallBackFromGPUDNN2GPU(pir::Operation* op,
       use_cudnn = false;
     }
 
-    auto shape = GetValueShape(op->operand_source(0));
+    auto shape = pir::GetShapeFromValue(op->operand_source(0));
     if (shape[1] == 3) {
       use_cudnn = false;
     }
diff --git a/paddle/fluid/pir/transforms/transform_general_functions.cc b/paddle/fluid/pir/transforms/transform_general_functions.cc
diff --git a/paddle/fluid/pir/transforms/transform_general_functions.h b/paddle/fluid/pir/transforms/transform_general_functions.h
diff --git a/paddle/phi/kernels/transfer_layout_kernel.cc b/paddle/phi/kernels/transfer_layout_kernel.cc

Original file line number	Diff line number	Diff line change
`@@ -343,8 +343,12 @@ void TransferLayoutElimPass::ApplyImpl(ir::Graph *graph) const {`
`343`	`343`	`}`
`344`	`344`	`if (!modify) break;`
`345`	`345`	`}`
`346`		`- LOG(INFO) << "move down " << move_down_count << " transfer_layout";`
`347`		`- LOG(INFO) << "eliminate " << elim_count << " pair of transfer_layout";`
	`346`	`+ if (move_down_count > 0) {`
	`347`	`+ LOG(INFO) << "move down " << move_down_count << " transfer_layout";`
	`348`	`+ }`
	`349`	`+ if (elim_count > 0) {`
	`350`	`+ LOG(INFO) << "eliminate " << elim_count << " pair of transfer_layout";`
	`351`	`+ }`
`348`	`352`	`}`
`349`	`353`
`350`	`354`	`} // namespace ir`
Original file line number	Diff line number	Diff line change
`@@ -71,7 +71,7 @@ void OperationFactory::RegisterManualOpCreator() {`
`71`	`71`	`});`
`72`	`72`	`}`
`73`	`73`
`74`		`-static pir::Attribute CreateIrAttribute(const std::any& obj) {`
	`74`	`+pir::Attribute CreateIrAttribute(const std::any& obj) {`
`75`	`75`	`if (obj.type() == typeid(bool)) {`
`76`	`76`	`return IrAttrbuteCreator<bool>()(std::any_cast<bool>(obj));`
`77`	`77`	`} else if (obj.type() == typeid(int32_t)) {`