【PIR】Add strategy of pool and reshape (#72118)

liuruyan · web-flow · commit 985fb6bced28 · 2025-04-11T17:22:56.000+08:00
* add strategy of pool and reshape

* fix window-ci bug
diff --git a/paddle/fluid/pir/dialect/operator/interface/layout_transformation.cc b/paddle/fluid/pir/dialect/operator/interface/layout_transformation.cc
@@ -208,7 +208,22 @@ std::vector<pir::Value> RelevantOutputsImpl<ReshapeOp>(pir::Operation* op) {
 
 template <>
 bool CanBeModifiedImpl<ReshapeOp>(pir::Operation* op) {
-  return false;
+  auto concrete_op = op->dyn_cast<ReshapeOp>();
+  auto shape = concrete_op.shape();
+  auto x = concrete_op.x();
+  if (!x || !(x.defining_op()->isa<pir::ParameterOp>())) return false;
+  if (!shape || !(shape.defining_op()->isa<FullIntArrayOp>())) return false;
+
+  auto full_int_op = shape.defining_op()->dyn_cast<FullIntArrayOp>();
+  auto value_attr =
+      full_int_op.attribute("value").dyn_cast<pir::ArrayAttribute>();
+
+  std::vector<int32_t> value_int32;
+  for (size_t i = 0; i < value_attr.size(); ++i) {
+    auto attr = value_attr.at(i);
+    value_int32.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
+  }
+  return value_int32 == std::vector<int32_t>{1, -1, 1, 1};
 }
 
 template <>
@@ -297,6 +312,46 @@ void RewriteByLayoutImpl<ConcatOp>(pir::Operation* op,
   RewriteByInfermeta<ConcatOp>(op, new_layout);
 }
 
+template <>
+void RewriteByLayoutImpl<ReshapeOp>(pir::Operation* op,
+                                    common::DataLayout new_layout) {
+  auto concrete_op = op->dyn_cast<ReshapeOp>();
+
+  auto shape = concrete_op.shape();
+  if (!shape || !(shape.defining_op()->isa<FullIntArrayOp>())) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Reshape's shape must be processed when rewrite by layout."));
+  }
+
+  auto full_int_op = shape.defining_op()->dyn_cast<FullIntArrayOp>();
+  auto value_attr =
+      full_int_op.attribute("value").dyn_cast<pir::ArrayAttribute>();
+
+  PADDLE_ENFORCE_EQ(value_attr.size(),
+                    4,
+                    common::errors::InvalidArgument(
+                        "Reshape's shape size was expected as 4, but got %d",
+                        value_attr.size()));
+  std::vector<pir::Attribute> new_value_attr;
+  if (new_layout == common::DataLayout::kNHWC) {
+    new_value_attr = std::vector<pir::Attribute>{
+        value_attr[0], value_attr[2], value_attr[3], value_attr[1]};
+  } else if (new_layout == common::DataLayout::kNCHW) {
+    new_value_attr = std::vector<pir::Attribute>{
+        value_attr[0], value_attr[3], value_attr[1], value_attr[2]};
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Layout only supports NHWC and NCHW in LayoutTransformationInterface"));
+  }
+
+  shape.defining_op()->set_attribute(
+      "value",
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), new_value_attr));
+
+  // infer new meta for concat
+  RewriteByInfermeta<ReshapeOp>(op, new_layout);
+}
+
 template <>
 void RewriteByLayoutImpl<ArgmaxOp>(pir::Operation* op,
                                    common::DataLayout new_layout) {
@@ -347,5 +402,115 @@ std::vector<pir::Value> RelevantInputsImpl<Pool2dOp>(pir::Operation* op) {
   return {concrete_op.x()};
 }
 
+template <>
+common::DataLayout PreferLayoutImpl<Pool2dOp>(pir::Operation* op) {
+  auto concrete_op = op->dyn_cast<Pool2dOp>();
+  auto data_format_attr = op->attribute<pir::StrAttribute>("data_format");
+  auto origin_format = common::StringToDataLayout(data_format_attr.AsString());
+  auto input =
+      concrete_op.x().type().dyn_cast<paddle::dialect::DenseTensorType>();
+  auto full_int_op =
+      concrete_op.kernel_size().defining_op()->dyn_cast<FullIntArrayOp>();
+  if (!full_int_op || !input) return origin_format;
+
+  // get pooling type
+  std::string pool_type = concrete_op.attribute("pooling_type")
+                              .dyn_cast<pir::StrAttribute>()
+                              .AsString();
+
+  // get input dims h, w, c
+  int32_t h, w, c;
+  if (origin_format == common::DataLayout::kNHWC) {
+    h = input.dims().at(1);
+    w = input.dims().at(2);
+    c = input.dims().at(3);
+  } else if (origin_format == common::DataLayout::kNCHW) {
+    h = input.dims().at(2);
+    w = input.dims().at(3);
+    c = input.dims().at(1);
+  } else {
+    return origin_format;
+  }
+
+  // get kernel_size
+  auto value_attr =
+      full_int_op.attribute("value").dyn_cast<pir::ArrayAttribute>();
+  std::vector<int64_t> kernel_size;
+  for (size_t i = 0; i < value_attr.size(); ++i) {
+    auto attr = value_attr.at(i);
+    kernel_size.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
+  }
+
+  auto PreferLayout = [](int c,
+                         int h,
+                         int w,
+                         std::vector<int64_t> kernel_size,
+                         std::string pool_type) {
+    auto AllEqual = [&](const std::vector<int64_t>& vec) {
+      return vec.empty() || std::all_of(vec.begin(), vec.end(), [&](int64_t x) {
+               return x == vec[0];
+             });
+    };
+    // TODO(liujinnan): need to test the prefer layout if kernel_size is not
+    // aligned.
+    if (!AllEqual(kernel_size)) return common::DataLayout::kNCHW;
+
+    int k = kernel_size[0];
+    // kernel size is all 1, prefer NCHW.
+    if (k == 1) return common::DataLayout::kNCHW;
+
+    if (pool_type == "max") {
+      if (h * w <= 64 * 64) {
+        if (k <= 3) return common::DataLayout::kNHWC;
+        return common::DataLayout::kNCHW;
+      } else {
+        if (c <= 16) {
+          if (k <= 5) return common::DataLayout::kNHWC;
+          return common::DataLayout::kNCHW;
+        }
+        // when c > 16, all kernel_size return NHWC
+        return common::DataLayout::kNHWC;
+      }
+    } else if (pool_type == "avg") {
+      if (h * w <= 64 * 64) {
+        if (c <= 16) {
+          if (k < 7)
+            return common::DataLayout::kNCHW;
+          else
+            return common::DataLayout::kNHWC;
+        } else if (c > 16 && c <= 32) {
+          if (k <= 7)
+            return common::DataLayout::kNHWC;
+          else
+            return common::DataLayout::kNCHW;
+        } else if (c > 32 && c <= 64) {
+          if (k < 5)
+            return common::DataLayout::kNCHW;
+          else
+            return common::DataLayout::kNHWC;
+        } else if (c > 64 && c <= 128) {
+          if (k < 7)
+            return common::DataLayout::kNHWC;
+          else
+            return common::DataLayout::kNCHW;
+        }
+        // when c > 128, all kernel_size return NHWC
+        return common::DataLayout::kNHWC;
+      } else {
+        if (c < 64) {
+          return common::DataLayout::kNHWC;
+        } else {
+          if (k < 7)
+            return common::DataLayout::kNHWC;
+          else
+            return common::DataLayout::kNCHW;
+        }
+      }
+    }
+    return common::DataLayout::kNCHW;
+  };
+  return PreferLayout(c, h, w, kernel_size, pool_type);
+}
+
 }  // namespace paddle::dialect
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::LayoutTransformationInterface)
diff --git a/paddle/fluid/pir/dialect/operator/interface/layout_transformation.hpp b/paddle/fluid/pir/dialect/operator/interface/layout_transformation.hpp
@@ -178,6 +178,7 @@ OVERLOAD_RELEVANT_INPUTS(AddGroupNormSiluOp);
 OVERLOAD_RELEVANT_OUTPUTS(AddGroupNormSiluOp);
 
 class ReshapeOp;
+OVERLOAD_REWRITE_BY_LAYOUT(ReshapeOp);
 OVERLOAD_RELEVANT_INPUTS(ReshapeOp);
 OVERLOAD_RELEVANT_OUTPUTS(ReshapeOp);
 OVERLOAD_CAN_BE_MODIFIED(ReshapeOp);
@@ -200,6 +201,7 @@ OVERLOAD_REWRITE_BY_LAYOUT(ArgmaxOp);
 
 class Pool2dOp;
 OVERLOAD_RELEVANT_INPUTS(Pool2dOp);
+OVERLOAD_PREFER_LAYOUT(Pool2dOp);
 
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/transforms/general/auto_layout_insert_pass.cc b/paddle/fluid/pir/transforms/general/auto_layout_insert_pass.cc
@@ -228,19 +228,12 @@ class AutoLayoutInsertPass : public pir::Pass {
         }
       } else if (!kOpsNchw.count(op_name) && !kOpsWithAxis.count(op_name) &&
                  IsInsertTransposeOpBefore(op)) {
-        // TODO(liujinnan):
-        // 1. Hide the special judgment method in PreferLayout of
-        // LayoutTransformationInterface.
-        // 2. Here we should do detailed experiments (`NHWC` vs `T
-        // + NCHW + T`) to determine whether to change the layout of pool2d. For
-        // example, when `hw` is large, it tends to be NCHW, and when `c` is
-        // large, it tends to be NHWC. Here, a temporary solution is to keep
-        // nchw when the pool type is maxpool.
-        if (auto pool2d_op = op->dyn_cast<paddle::dialect::Pool2dOp>()) {
-          if (pool2d_op.attribute("pooling_type")
-                  .dyn_cast<pir::StrAttribute>()
-                  .AsString() == "max")
-            continue;
+        // TODO(liujinnan): Remove the list and set the CanBeModified method for
+        // all ops.
+        if (op_name == "pd_op.pool2d" || op_name == "pd_op.reshape") {
+          auto layout_interface =
+              op->dyn_cast<paddle::dialect::LayoutTransformationInterface>();
+          if (!layout_interface.CanBeModified(op)) continue;
         }
         VLOG(4) << "enter NCHW op: " << op_name;
         DoTransposeOpOperand(op, builder);
@@ -268,6 +261,21 @@ class AutoLayoutInsertPass : public pir::Pass {
   void DoTransposeOpOperand(pir::Operation* op,
                             pir::Builder& builder) {  // NOLINT
     auto InsertTranspose = [&](pir::OpOperand* operand) {
+      // if operand defining op is reshape op, try to rewrite reshape op
+      if (operand->source().defining_op<paddle::dialect::ReshapeOp>()) {
+        auto reshape_op = operand->source()
+                              .defining_op()
+                              ->dyn_cast<paddle::dialect::ReshapeOp>();
+        auto layout_interface =
+            reshape_op
+                ->dyn_cast<paddle::dialect::LayoutTransformationInterface>();
+        if (layout_interface.CanBeModified(reshape_op)) {
+          layout_interface.RewriteByLayout(reshape_op,
+                                           common::DataLayout::NHWC);
+          operand->set_source(reshape_op->result(0));
+          return;
+        }
+      }
       auto transpose_op = builder.Build<paddle::dialect::TransposeOp>(
           operand->source(), kNchw2Nhwc_);
       transpose_op->set_attribute(
@@ -282,7 +290,8 @@ class AutoLayoutInsertPass : public pir::Pass {
     builder.set_insertion_point(op);
     // For conv2d, only transpose the input.
     if (op->isa<paddle::dialect::Conv2dOp>() ||
-        op->isa<paddle::dialect::Conv2dTransposeOp>()) {
+        op->isa<paddle::dialect::Conv2dTransposeOp>() ||
+        op->isa<paddle::dialect::DepthwiseConv2dOp>()) {
       auto inp = op->operand(0);
       if (!JudgeValue(inp.source())) return;
       InsertTranspose(&inp);
@@ -328,9 +337,9 @@ const std::set<std::string> kOpsNchw = {"pd_op.max_pool2d_with_index",
                                         "pd_op.grid_sample",
                                         "pd_op.shuffle_channel",
                                         "cf.yield",
-                                        "pd_op.reshape",
+                                        // "pd_op.reshape",
                                         "pd_op.instance_norm",
-                                        //  "pd_op.batch_norm_",
+                                        // "pd_op.batch_norm_",
                                         "pd_op.bilinear_interp",
                                         "pd_op.shape",
                                         "pd_op.shape64",
diff --git a/paddle/fluid/pir/transforms/general/auto_layout_pass.cc b/paddle/fluid/pir/transforms/general/auto_layout_pass.cc
@@ -110,7 +110,7 @@ class AutoLayoutPass : public pir::Pass {
         }
       }
     }
-    VLOG(4) << "end IsNeedAllTranspose"
+    VLOG(0) << "end IsNeedAllTranspose"
             << " conv_count_: " << conv_count_
             << " transpose_count_: " << transpose_count_
             << " transpose_scale_ * transpose_count_: "

Original file line number	Diff line number	Diff line change
`@@ -110,7 +110,7 @@ class AutoLayoutPass : public pir::Pass {`
`110`	`110`	`}`
`111`	`111`	`}`
`112`	`112`	`}`
`113`		`- VLOG(4) << "end IsNeedAllTranspose"`
	`113`	`+ VLOG(0) << "end IsNeedAllTranspose"`
`114`	`114`	`<< " conv_count_: " << conv_count_`
`115`	`115`	`<< " transpose_count_: " << transpose_count_`
`116`	`116`	`<< " transpose_scale_ * transpose_count_: "`