adjust transfer_layout_pass level and add statistics (#64708)

yuanlehome · web-flow · commit 00a004f14220 · 2024-05-29T19:33:21.000-07:00
* adjust transfer_layout_pass level and add statistics

* update
diff --git a/paddle/fluid/pir/transforms/general/transfer_layout_pass.cc b/paddle/fluid/pir/transforms/general/transfer_layout_pass.cc
@@ -542,7 +542,7 @@ using Edge = FlowGraph::Edge;
 
 class TransferLayoutPass : public pir::Pass {
  public:
-  TransferLayoutPass() : pir::Pass("transfer_layout_pass", 4) {}
+  TransferLayoutPass() : pir::Pass("transfer_layout_pass", 3) {}
 
   bool CanApplyOn(pir::Operation* op) const override {
     if (!op->isa<pir::ModuleOp>()) {
@@ -647,7 +647,8 @@ class TransferLayoutPass : public pir::Pass {
 
     VLOG(10)
         << "-----------------------[rewrite begin]------------------------";
-
+    int64_t num_of_layout_changed_ops{0};
+    int64_t num_of_transpose_ops{0};
     while (!q.empty()) {
       auto node = q.front();
       q.pop_front();
@@ -664,6 +665,7 @@ class TransferLayoutPass : public pir::Pass {
           if (layout_transformation_iface) {
             layout_transformation_iface.RewriteByLayout(
                 op, common::DataLayout::NHWC);
+            num_of_layout_changed_ops++;
           } else {
             PADDLE_THROW(common::errors::Unimplemented(
                 "Op %s should have a specialized RewriteByLayout function",
@@ -695,6 +697,7 @@ class TransferLayoutPass : public pir::Pass {
             ((src_set.count(node) > 0) ? common::DataLayout::NHWC
                                        : common::DataLayout::NCHW);
         builder.SetInsertionPointAfter(dst_value.defining_op());
+        num_of_transpose_ops++;
         auto transpose_op =
             builder.Build<paddle::dialect::TransposeOp>(dst_value, perm);
         transpose_op->set_attribute(
@@ -735,6 +738,7 @@ class TransferLayoutPass : public pir::Pass {
             ((src_set.count(node) > 0) ? common::DataLayout::NHWC
                                        : common::DataLayout::NCHW);
         builder.SetInsertionPointAfter(value.defining_op());
+        num_of_transpose_ops++;
         auto transpose_op =
             builder.Build<paddle::dialect::TransposeOp>(value, perm);
         transpose_op->set_attribute(
@@ -749,6 +753,7 @@ class TransferLayoutPass : public pir::Pass {
         value.ReplaceUsesWithIf(transpose_op.out(), replace_uses_in_cut_set);
       }
     }
+    AddStatistics(num_of_transpose_ops, num_of_layout_changed_ops);
   }
 };
 
diff --git a/paddle/pir/include/pass/pass.h b/paddle/pir/include/pass/pass.h
@@ -167,9 +167,9 @@ class IR_API Pass {
     Set<int64_t>("__match_count__", new int64_t{match_count});
   }
 
-  void AddStatistics(int64_t match_count, int64_t all_count) {
-    Set<int64_t>("__match_count__", new int64_t{match_count});
-    Set<int64_t>("__all_count__", new int64_t{all_count});
+  void AddStatistics(int64_t match_count_1, int64_t match_count_2) {
+    Set<int64_t>("__match_count_1__", new int64_t{match_count_1});
+    Set<int64_t>("__match_count_2__", new int64_t{match_count_2});
   }
 
   void AddStatistics(const std::string& custom_log) {
diff --git a/paddle/pir/src/pass/print_statistics.cc b/paddle/pir/src/pass/print_statistics.cc
@@ -37,20 +37,14 @@ class PrintStatistics : public PassInstrumentation {
   }
 
   void RunAfterPass(Pass *pass, Operation *op) override {
-    if (pass->Has("__match_count__") && pass->Has("__all_count__")) {
-      auto match_count = pass->Get<int64_t>("__match_count__");
-      auto all_count = pass->Get<int64_t>("__all_count__");
-      PADDLE_ENFORCE_LE(match_count,
-                        all_count,
-                        phi::errors::InvalidArgument(
-                            "match_count: %d should smaller than all_count: %d",
-                            match_count,
-                            all_count));
-      if (match_count > 0) {
-        LOG(INFO) << "--- detected [" << match_count << "/" << all_count
+    if (pass->Has("__match_count_1__") && pass->Has("__match_count_2__")) {
+      auto match_count_1 = pass->Get<int64_t>("__match_count_1__");
+      auto match_count_2 = pass->Get<int64_t>("__match_count_2__");
+      if (match_count_1 > 0 || match_count_2 > 0) {
+        LOG(INFO) << "--- detected [" << match_count_1 << ", " << match_count_2
                   << "] subgraphs!";
       }
-    } else if (pass->Has("__match_count__") && !pass->Has("__all_count__")) {
+    } else if (pass->Has("__match_count__")) {
       auto match_count = pass->Get<int64_t>("__match_count__");
       if (match_count > 0) {
         LOG(INFO) << "--- detected [" << match_count << "] subgraphs!";

Original file line number	Diff line number	Diff line change
`@@ -167,9 +167,9 @@ class IR_API Pass {`
`167`	`167`	`Set<int64_t>("__match_count__", new int64_t{match_count});`
`168`	`168`	`}`
`169`	`169`
`170`		`- void AddStatistics(int64_t match_count, int64_t all_count) {`
`171`		`- Set<int64_t>("__match_count__", new int64_t{match_count});`
`172`		`- Set<int64_t>("__all_count__", new int64_t{all_count});`
	`170`	`+ void AddStatistics(int64_t match_count_1, int64_t match_count_2) {`
	`171`	`+ Set<int64_t>("__match_count_1__", new int64_t{match_count_1});`
	`172`	`+ Set<int64_t>("__match_count_2__", new int64_t{match_count_2});`
`173`	`173`	`}`
`174`	`174`
`175`	`175`	`void AddStatistics(const std::string& custom_log) {`