[Mlir-commits] [mlir] bf52318 - Change the PrintOpStatsPass to operate on any operation instead of just ModuleOp

[Mehdi Amini]

Author: Mehdi Amini
Date: 2020-11-03T11:15:32Z
New Revision: bf523186fb7817ba12aa1de09d6db95589621797

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

LOG: Change the PrintOpStatsPass to operate on any operation instead of just ModuleOp

This allows to use it on other operation, like a GPUModule for example.

Added: 
    

Modified: 
    mlir/include/mlir/Transforms/Passes.h
    mlir/include/mlir/Transforms/Passes.td
    mlir/lib/Transforms/OpStats.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index deb0fbbb6132..c9b90d15e79d 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/mlir/include/mlir/Transforms/Passes.h
@@ -94,7 +94,7 @@ std::unique_ptr<Pass> createStripDebugInfoPass();
 
 /// Creates a pass which prints the list of ops and the number of occurrences in
 /// the module.
-std::unique_ptr<OperationPass<ModuleOp>> createPrintOpStatsPass();
+std::unique_ptr<Pass> createPrintOpStatsPass();
 
 /// Creates a pass which inlines calls and callable operations as defined by
 /// the CallGraph.

diff  --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index dfa6867d3106..031b30812179 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -523,7 +523,7 @@ def PrintCFG : FunctionPass<"print-cfg-graph"> {
   let constructor = "mlir::createPrintCFGGraphPass()";
 }
 
-def PrintOpStats : Pass<"print-op-stats", "ModuleOp"> {
+def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";
 }

diff  --git a/mlir/lib/Transforms/OpStats.cpp b/mlir/lib/Transforms/OpStats.cpp
index 7d6491085119..0ed8ca116d40 100644
--- a/mlir/lib/Transforms/OpStats.cpp
+++ b/mlir/lib/Transforms/OpStats.cpp
@@ -36,9 +36,8 @@ struct PrintOpStatsPass : public PrintOpStatsBase<PrintOpStatsPass> {
 void PrintOpStatsPass::runOnOperation() {
   opCount.clear();
 
-  // Compute the operation statistics for each function in the module.
-  for (auto &op : getOperation())
-    op.walk([&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
+  // Compute the operation statistics for the currently visited operation.
+  getOperation()->walk([&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
   printSummary();
 }
 
@@ -81,6 +80,6 @@ void PrintOpStatsPass::printSummary() {
   }
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createPrintOpStatsPass() {
+std::unique_ptr<Pass> mlir::createPrintOpStatsPass() {
   return std::make_unique<PrintOpStatsPass>();
 }