[Mlir-commits] [mlir] d41b3bf - [mlir][Pass] Fix dropped statistics with nested adaptors.

[Will Dietz]

Author: Will Dietz
Date: 2022-12-07T08:31:43-06:00
New Revision: d41b3bf7c37d36f02e55662c31c14aecdcdf9b9b

URL: https://github.com/llvm/llvm-project/commit/d41b3bf7c37d36f02e55662c31c14aecdcdf9b9b
DIFF: https://github.com/llvm/llvm-project/commit/d41b3bf7c37d36f02e55662c31c14aecdcdf9b9b.diff

LOG: [mlir][Pass] Fix dropped statistics with nested adaptors.

When running in parallel, nesting more than once caused
statistics to be dropped.

Fix by also preparing "async" pass managers before merging,
as they may also have "async" pass managers within.

Add test checking reported statistics have expected values
with and without threading enabled.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D139459

Added: 
    mlir/test/Pass/pipeline-stats-nested.mlir

Modified: 
    mlir/lib/Pass/PassStatistics.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Pass/PassStatistics.cpp b/mlir/lib/Pass/PassStatistics.cpp
index 198009942c4f..a3db6c6dfaa4 100644
--- a/mlir/lib/Pass/PassStatistics.cpp
+++ b/mlir/lib/Pass/PassStatistics.cpp
@@ -226,10 +226,12 @@ static void prepareStatistics(OpPassManager &pm) {
     MutableArrayRef<OpPassManager> nestedPms = adaptor->getPassManagers();
 
     // Merge the statistics from the async pass managers into the main nested
-    // pass managers.
+    // pass managers.  Prepare recursively before merging.
     for (auto &asyncPM : adaptor->getParallelPassManagers()) {
-      for (unsigned i = 0, e = asyncPM.size(); i != e; ++i)
+      for (unsigned i = 0, e = asyncPM.size(); i != e; ++i) {
+        prepareStatistics(asyncPM[i]);
         asyncPM[i].mergeStatisticsInto(nestedPms[i]);
+      }
     }
 
     // Prepare the statistics of each of the nested passes.

diff  --git a/mlir/test/Pass/pipeline-stats-nested.mlir b/mlir/test/Pass/pipeline-stats-nested.mlir
new file mode 100644
index 000000000000..386dc366857f
--- /dev/null
+++ b/mlir/test/Pass/pipeline-stats-nested.mlir
@@ -0,0 +1,24 @@
+// REQUIRES: asserts
+// Check that statistics in nested pipelines are not ignored and that this works with and without threading.
+// RUN: mlir-opt %s -verify-each=true -pass-pipeline='builtin.module(builtin.module(func.func(test-stats-pass,test-stats-pass)))' -mlir-pass-statistics -mlir-pass-statistics-display=list -mlir-disable-threading 2>&1 | FileCheck -check-prefix=LIST %s
+// RUN: mlir-opt %s -verify-each=true -pass-pipeline='builtin.module(builtin.module(func.func(test-stats-pass,test-stats-pass)))' -mlir-pass-statistics -mlir-pass-statistics-display=list                         2>&1 | FileCheck -check-prefix=LIST %s
+
+// RUN: mlir-opt %s -verify-each=true -pass-pipeline='builtin.module(builtin.module(func.func(test-stats-pass,test-stats-pass)))' -mlir-pass-statistics -mlir-pass-statistics-display=pipeline -mlir-disable-threading 2>&1 | FileCheck -check-prefix=PIPELINE %s
+// RUN: mlir-opt %s -verify-each=true -pass-pipeline='builtin.module(builtin.module(func.func(test-stats-pass,test-stats-pass)))' -mlir-pass-statistics -mlir-pass-statistics-display=pipeline                         2>&1 | FileCheck -check-prefix=PIPELINE %s
+
+// Check for the correct statistic values.
+// Each test-stats-pass will count two ops: the func.func and the func.return .
+//      LIST: (S) 4 num-ops
+// LIST-NEXT: (S) 4 num-ops2
+//      PIPELINE: (S) 2 num-ops
+// PIPELINE-NEXT: (S) 2 num-ops2
+//      PIPELINE: (S) 2 num-ops
+// PIPELINE-NEXT: (S) 2 num-ops2
+
+module {
+  module {
+    func.func @foo() {
+      return
+    }
+  }
+}