Add support for instance specific pass statistics.

Statistics are a way to keep track of what the compiler is doing and how effective various optimizations are. It is useful to see what optimizations are contributing to making a particular program run faster. Pass-instance specific statistics take this even further as you can see the effect of placing a particular pass at specific places within the pass pipeline, e.g. they could help answer questions like "what happens if I run CSE again here".

Statistics can be added to a pass by simply adding members of type 'Pass::Statistics'. This class takes as a constructor arguments: the parent pass pointer, a name, and a description. Statistics can be dumped by the pass manager in a similar manner to how pass timing information is dumped, i.e. via PassManager::enableStatistics programmatically; or -pass-statistics and -pass-statistics-display via the command line pass manager options.

Below is an example:

struct MyPass : public OperationPass<MyPass> {
  Statistic testStat{this, "testStat", "A test statistic"};

  void runOnOperation() {
    ...
    ++testStat;
    ...
  }
};

$ mlir-opt -pass-pipeline='func(my-pass,my-pass)' foo.mlir -pass-statistics

Pipeline Display:
===-------------------------------------------------------------------------===
                         ... Pass statistics report ...
===-------------------------------------------------------------------------===
'func' Pipeline
  MyPass
    (S) 15 testStat - A test statistic
  MyPass
    (S)  6 testStat - A test statistic

List Display:
===-------------------------------------------------------------------------===
                         ... Pass statistics report ...
===-------------------------------------------------------------------------===
MyPass
  (S) 21 testStat - A test statistic

PiperOrigin-RevId: 284022014

mlir/g3doc/WritingAPass.md

@@ -319,10 +319,10 @@ program has been run through the passes. This provides several benefits:
 
 ## Pass Registration
 
-Briefly shown in the example definitions of the various
-pass types is the `PassRegistration` class. This is a utility to
-register derived pass classes so that they may be created, and inspected, by
-utilities like mlir-opt. Registering a pass class takes the form:
+Briefly shown in the example definitions of the various pass types is the
+`PassRegistration` class. This is a utility to register derived pass classes so
+that they may be created, and inspected, by utilities like mlir-opt. Registering
+a pass class takes the form:
 
 ```c++
 static PassRegistration<MyPass> pass("command-line-arg", "description");
@@ -469,6 +469,76 @@ struct MyPassOptions : public PassOptions<MyPassOptions> {
 static PassRegistration<MyPass, MyPassOptions> pass("my-pass", "description");
 ```
 
+## Pass Statistics
+
+Statistics are a way to keep track of what the compiler is doing and how
+effective various transformations are. It is often useful to see what effect
+specific transformations have on a particular program, and how often they
+trigger. Pass statistics are instance specific which allow for taking this a
+step further as you are able to see the effect of placing a particular
+transformation at specific places within the pass pipeline. For example, they
+help answer questions like `What happens if I run CSE again here?`.
+
+Statistics can be added to a pass by using the 'Pass::Statistic' class. This
+class takes as a constructor arguments: the parent pass, a name, and a
+description. This class acts like an unsigned integer, and may be incremented
+and updated accordingly. These statistics use the same infrastructure as
+[`llvm::Statistic`](http://llvm.org/docs/ProgrammersManual.html#the-statistic-class-stats-option)
+and thus have similar usage constraints. Collected statistics can be dumped by
+the [pass manager](#pass-manager) programmatically via
+`PassManager::enableStatistics`; or via `-pass-statistics` and
+`-pass-statistics-display` on the command line.
+
+An example is shown below:
+
+```c++
+struct MyPass : public OperationPass<MyPass> {
+  Statistic testStat{this, "testStat", "A test statistic"};
+
+  void runOnOperation() {
+    ...
+
+    // Update our statistic after some invariant was hit.
+    ++testStat;
+
+    ...
+  }
+};
+```
+
+The collected statistics may be aggregated in two types of views:
+
+A pipeline view that models the structure of the pass manager, this is the
+default view:
+
+```shell
+$ mlir-opt -pass-pipeline='func(my-pass,my-pass)' foo.mlir -pass-statistics
+
+===-------------------------------------------------------------------------===
+                         ... Pass statistics report ...
+===-------------------------------------------------------------------------===
+'func' Pipeline
+  MyPass
+    (S) 15 testStat - A test statistic
+  VerifierPass
+  MyPass
+    (S)  6 testStat - A test statistic
+  VerifierPass
+VerifierPass
+```
+
+And a list view that aggregates all instances of a specific pass together:
+
+```shell
+$ mlir-opt -pass-pipeline='func(my-pass, my-pass)' foo.mlir -pass-statistics -pass-statistics-display=list
+
+===-------------------------------------------------------------------------===
+                         ... Pass statistics report ...
+===-------------------------------------------------------------------------===
+MyPass
+  (S) 21 testStat - A test statistic
+```
+
 ## Pass Instrumentation
 
 MLIR provides a customizable framework to instrument pass execution and analysis

mlir/include/mlir/Pass/Pass.h

@@ -23,6 +23,7 @@
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/PointerIntPair.h"
+#include "llvm/ADT/Statistic.h"
 
 namespace mlir {
 namespace detail {
@@ -76,6 +77,28 @@ class Pass {
   /// pass to be to be round-trippable to the textual format.
   virtual void printAsTextualPipeline(raw_ostream &os);
 
+  /// This class represents a single pass statistic. This statistic functions
+  /// similarly to an unsigned integer value, and may be updated and incremented
+  /// accordingly. This class can be used to provide additional information
+  /// about the transformations and analyses performed by a pass.
+  class Statistic : public llvm::Statistic {
+  public:
+    /// The statistic is initialized by the pass owner, a name, and a
+    /// description.
+    Statistic(Pass *owner, const char *name, const char *description);
+
+    /// Assign the statistic to the given value.
+    Statistic &operator=(unsigned value);
+
+  private:
+    /// Hide some of the details of llvm::Statistic that we don't use.
+    using llvm::Statistic::getDebugType;
+  };
+
+  /// Returns the main statistics for this pass instance.
+  ArrayRef<Statistic *> getStatistics() const { return statistics; }
+  MutableArrayRef<Statistic *> getStatistics() { return statistics; }
+
 protected:
   explicit Pass(const PassID *passID,
                 llvm::Optional<StringRef> opName = llvm::None)
@@ -125,6 +148,9 @@ class Pass {
   /// The current execution state for the pass.
   llvm::Optional<detail::PassExecutionState> passState;
 
+  /// The set of statistics held by this pass.
+  std::vector<Statistic *> statistics;
+
   /// Allow access to 'clone' and 'run'.
   friend class OpPassManager;
 };

mlir/include/mlir/Pass/PassManager.h

@@ -21,6 +21,9 @@
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator.h"
+
+#include <vector>
 
 namespace llvm {
 class Any;
@@ -54,6 +57,13 @@ class OpPassManager {
   ~OpPassManager();
   OpPassManager &operator=(const OpPassManager &rhs);
 
+  /// Iterator over the passes in this pass manager.
+  using pass_iterator =
+      llvm::pointee_iterator<std::vector<std::unique_ptr<Pass>>::iterator>;
+  pass_iterator begin();
+  pass_iterator end();
+  llvm::iterator_range<pass_iterator> getPasses() { return {begin(), end()}; }
+
   /// Run the held passes over the given operation.
   LogicalResult run(Operation *op, AnalysisManager am);
 
@@ -93,6 +103,9 @@ class OpPassManager {
   /// the correctness of per-pass overrides of Pass::printAsTextualPipeline.
   void printAsTextualPipeline(raw_ostream &os);
 
+  /// Merge the pass statistics of this class into 'other'.
+  void mergeStatisticsInto(OpPassManager &other);
+
 private:
   OpPassManager(OperationName name, bool disableThreads, bool verifyPasses);
 
@@ -107,10 +120,10 @@ class OpPassManager {
 // PassManager
 //===----------------------------------------------------------------------===//
 
-/// An enum describing the different display modes for the pass timing
-/// information within the pass manager.
-enum class PassTimingDisplayMode {
-  // In this mode the results are displayed in a list sorted by total time,
+/// An enum describing the different display modes for the information within
+/// the pass manager.
+enum class PassDisplayMode {
+  // In this mode the results are displayed in a list sorted by total,
   // with each pass/analysis instance aggregated into one unique result.
   List,
 
@@ -162,13 +175,23 @@ class PassManager : public OpPassManager {
   /// Note: Timing should be enabled after all other instrumentations to avoid
   /// any potential "ghost" timing from other instrumentations being
   /// unintentionally included in the timing results.
-  void enableTiming(
-      PassTimingDisplayMode displayMode = PassTimingDisplayMode::Pipeline);
+  void enableTiming(PassDisplayMode displayMode = PassDisplayMode::Pipeline);
+
+  /// Prompts the pass manager to print the statistics collected for each of the
+  /// held passes after each call to 'run'.
+  void
+  enableStatistics(PassDisplayMode displayMode = PassDisplayMode::Pipeline);
 
 private:
+  /// Dump the statistics of the passes within this pass manager.
+  void dumpStatistics();
+
   /// Flag that specifies if pass timing is enabled.
   bool passTiming : 1;
 
+  /// Flag that specifies if pass statistics should be dumped.
+  Optional<PassDisplayMode> passStatisticsMode;
+
   /// A manager for pass instrumentations.
   std::unique_ptr<PassInstrumentor> instrumentor;
 

mlir/lib/Pass/Pass.cpp

@@ -216,6 +216,11 @@ OpPassManager &OpPassManager::operator=(const OpPassManager &rhs) {
 
 OpPassManager::~OpPassManager() {}
 
+OpPassManager::pass_iterator OpPassManager::begin() {
+  return impl->passes.begin();
+}
+OpPassManager::pass_iterator OpPassManager::end() { return impl->passes.end(); }
+
 /// Run all of the passes in this manager over the current operation.
 LogicalResult OpPassManager::run(Operation *op, AnalysisManager am) {
   // Run each of the held passes.
@@ -341,6 +346,17 @@ void OpToOpPassAdaptorBase::mergeInto(OpToOpPassAdaptorBase &rhs) {
                        });
 }
 
+/// Returns the adaptor pass name.
+std::string OpToOpPassAdaptorBase::getName() {
+  std::string name = "Pipeline Collection : [";
+  llvm::raw_string_ostream os(name);
+  interleaveComma(getPassManagers(), os, [&](OpPassManager &pm) {
+    os << '\'' << pm.getOpName() << '\'';
+  });
+  os << ']';
+  return os.str();
+}
+
 OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr)
     : OpToOpPassAdaptorBase(std::move(mgr)) {}
 
@@ -560,9 +576,15 @@ LogicalResult PassManager::run(ModuleOp module) {
 
   // If reproducer generation is enabled, run the pass manager with crash
   // handling enabled.
-  if (crashReproducerFileName)
-    return runWithCrashRecovery(*this, am, module, *crashReproducerFileName);
-  return OpPassManager::run(module, am);
+  LogicalResult result =
+      crashReproducerFileName
+          ? runWithCrashRecovery(*this, am, module, *crashReproducerFileName)
+          : OpPassManager::run(module, am);
+
+  // Dump all of the pass statistics if necessary.
+  if (passStatisticsMode)
+    dumpStatistics();
+  return result;
 }
 
 /// Disable support for multi-threading within the pass manager.

mlir/lib/Pass/PassDetail.h

@@ -48,6 +48,9 @@ class OpToOpPassAdaptorBase {
   /// Returns the pass managers held by this adaptor.
   MutableArrayRef<OpPassManager> getPassManagers() { return mgrs; }
 
+  /// Returns the adaptor pass name.
+  std::string getName();
+
 protected:
   // A set of adaptors to run.
   SmallVector<OpPassManager, 1> mgrs;
@@ -75,6 +78,11 @@ class OpToOpPassAdaptorParallel
   /// Run the held pipeline over all operations.
   void runOnOperation() override;
 
+  /// Return the async pass managers held by this parallel adaptor.
+  MutableArrayRef<SmallVector<OpPassManager, 1>> getParallelPassManagers() {
+    return asyncExecutors;
+  }
+
 private:
   // A set of executors, cloned from the main executor, that run asynchronously
   // on different threads.

mlir/lib/Pass/PassManagerOptions.cpp

@@ -69,14 +69,30 @@ struct PassManagerOptions {
   llvm::cl::opt<bool> passTiming{
       "pass-timing",
       llvm::cl::desc("Display the execution times of each pass")};
-  llvm::cl::opt<PassTimingDisplayMode> passTimingDisplayMode{
+  llvm::cl::opt<PassDisplayMode> passTimingDisplayMode{
       "pass-timing-display",
       llvm::cl::desc("Display method for pass timing data"),
-      llvm::cl::init(PassTimingDisplayMode::Pipeline),
+      llvm::cl::init(PassDisplayMode::Pipeline),
       llvm::cl::values(
-          clEnumValN(PassTimingDisplayMode::List, "list",
+          clEnumValN(PassDisplayMode::List, "list",
                      "display the results in a list sorted by total time"),
-          clEnumValN(PassTimingDisplayMode::Pipeline, "pipeline",
+          clEnumValN(PassDisplayMode::Pipeline, "pipeline",
+                     "display the results with a nested pipeline view"))};
+
+  //===--------------------------------------------------------------------===//
+  // Pass Statistics
+  //===--------------------------------------------------------------------===//
+  llvm::cl::opt<bool> passStatistics{
+      "pass-statistics", llvm::cl::desc("Display the statistics of each pass")};
+  llvm::cl::opt<PassDisplayMode> passStatisticsDisplayMode{
+      "pass-statistics-display",
+      llvm::cl::desc("Display method for pass statistics"),
+      llvm::cl::init(PassDisplayMode::Pipeline),
+      llvm::cl::values(
+          clEnumValN(
+              PassDisplayMode::List, "list",
+              "display the results in a merged list sorted by pass name"),
+          clEnumValN(PassDisplayMode::Pipeline, "pipeline",
                      "display the results with a nested pipeline view"))};
 
   /// Add a pass timing instrumentation if enabled by 'pass-timing' flags.
@@ -146,6 +162,10 @@ void mlir::applyPassManagerCLOptions(PassManager &pm) {
   if ((*options)->disableThreads)
     pm.disableMultithreading();
 
+  // Enable statistics dumping.
+  if ((*options)->passStatistics)
+    pm.enableStatistics((*options)->passStatisticsDisplayMode);
+
   // Add the IR printing instrumentation.
   (*options)->addPrinterInstrumentation(pm);