From d41b3bf7c37d36f02e55662c31c14aecdcdf9b9b Mon Sep 17 00:00:00 2001
From: Will Dietz <w@wdtz.org>
Date: Tue, 6 Dec 2022 14:44:12 -0600
Subject: [PATCH] [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
---
 mlir/lib/Pass/PassStatistics.cpp          |  6 ++++--
 mlir/test/Pass/pipeline-stats-nested.mlir | 24 +++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Pass/pipeline-stats-nested.mlir

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
+    }
+  }
+}