[mlir] print-op-stats: Print op counts for top-k locs

[springerm]

Summary

Not only print the number of operations for the each location (show only top-k), but also a breakdown of how many operations of each kind.

Depends On D109672

[springerm]

Example output:

Top locations by number of operations:
--------------------------------------
...nrolled-tensor-transfer-ops.mlir:1 offset ":67:17) , 7169
  builtin.unrealized_conversion_cast , 1280
     llvm.add                        , 512
     llvm.bitcast                    , 256
     llvm.extractvalue               , 512
     llvm.getelementptr              , 256
     llvm.insertvalue                , 256
     llvm.intr.masked.load           , 256
     llvm.mlir.constant              , 256
     llvm.mul                        , 256
      scf.if                         , 256
      scf.yield                      , 512
      std.addi                       , 256
      std.cmpi                       , 512
      std.constant                   , 513
      std.index_cast                 , 512
      std.splat                      , 768
...nrolled-tensor-transfer-ops.mlir:1 offset ":39:14) , 7168
  builtin.unrealized_conversion_cast , 1280
     llvm.add                        , 512
     llvm.bitcast                    , 256
     llvm.extractvalue               , 512
     llvm.getelementptr              , 256
     llvm.insertvalue                , 256
     llvm.intr.masked.load           , 256
     llvm.mlir.constant              , 256
     llvm.mul                        , 256
      scf.if                         , 256
      scf.yield                      , 512
      std.addi                       , 256
      std.cmpi                       , 512
      std.constant                   , 512
      std.index_cast                 , 512
      std.splat                      , 768
...nrolled-tensor-transfer-ops.mlir:1 offset ":50:14) , 7168
  builtin.unrealized_conversion_cast , 1280
     llvm.add                        , 512
     llvm.bitcast                    , 256
     llvm.extractvalue               , 512

[...]

[bondhugula]

Can you please add a commit summary?

[jpienaar]

mlir/lib/Transforms/OpStats.cpp

43

Isn't this just count per operation name? Not sure why this is marked as histogram

[springerm]

rebase

[springerm]

address comments

[springerm]

Note: This revision is separate from D109672 because it is arguably less useful; D109672 may be good enough for debugging in many cases. Maybe we only want to have D109672...

mlir/lib/Transforms/OpStats.cpp

43

Ah yes, I meant "counts".

mlir/include/mlir/Transforms/Passes.td

def ParallelLoopCollapsing : Pass<"parallel-loop-collapsing"> {
   ];
 }
 
 def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";
   let options = [
     Option<"topKLocs", "top-k-locs", "int64_t",
-           /*default=*/"0", "Print top-k locations by number of operations">
+           /*default=*/"0", "Print top-k locations by number of operations">,
+    Option<"topKLocsOpCounts", "top-k-locs-op-counts", "bool",
+           /*default=*/"false", "Print op counts for top-k locations">
   ];
 }
 
 def SCCP : Pass<"sccp"> {
   let summary = "Sparse Conditional Constant Propagation";
   let description = [{
     This pass implements a general algorithm for sparse conditional constant
     propagation. This algorithm detects values that are known to be constant and

mlir/lib/Transforms/OpStats.cpp

 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
 namespace {
+struct OpCounter {
+  void increment(StringRef opName) {
+    ++map[opName];
+    ++total;
+  }
+
+  llvm::StringMap<int64_t> map;
+  int64_t total;
+};
+
 struct PrintOpStatsPass : public PrintOpStatsBase<PrintOpStatsPass> {
   explicit PrintOpStatsPass(raw_ostream &os = llvm::errs()) : os(os) {}
 
   // Prints the resultant operation statistics post iterating over the module.
   void runOnOperation() override;
 
   // Print summary of op stats.
   void printSummary();
 
   // Print top-k locations by number of operations.
   void printTopLocations();
 
+  // Print op counts, sorted by op name.
+  void printOpCounts(const llvm::StringMap<int64_t> &count);
+
 private:
   llvm::StringMap<int64_t> opCount;
-  llvm::DenseMap<Location, int64_t> opCountByLoc;
+  llvm::DenseMap<Location, OpCounter> opCountByLoc;
 
   raw_ostream &os;
 };
 } // namespace
 
 void PrintOpStatsPass::runOnOperation() {
   opCount.clear();
   opCountByLoc.clear();
   llvm::DenseMap<Location, int64_t> locCount;
 
   // Compute the operation statistics for the currently visited operation.
   getOperation()->walk([&](Operation *op) {
     ++opCount[op->getName().getStringRef()];
-    ++opCountByLoc[op->getLoc()];
+    opCountByLoc[op->getLoc()].increment(op->getName().getStringRef());
   });
 
   printSummary();
 
   if (topKLocs)
     printTopLocations();
 }
 
 void PrintOpStatsPass::printSummary() {
   os << "Operations encountered:\n";
   os << "-----------------------\n";
-  SmallVector<StringRef, 64> sorted(opCount.keys());
+  printOpCounts(opCount);
+}
+
+void PrintOpStatsPass::printOpCounts(const llvm::StringMap<int64_t> &count) {
+  SmallVector<StringRef, 64> sorted(count.keys());
   llvm::sort(sorted);
 
   // Split an operation name from its dialect prefix.
   auto splitOperationName = [](StringRef opName) {
     auto splitName = opName.split('.');
     return splitName.second.empty() ? std::make_pair("", splitName.first)
                                     : splitName;
   };
for (const auto &key : sorted) {
     // below. The alignment is for readability and does not affect CSV/FileCheck
     // parsing.
     if (dialectName.empty())
       os.indent(maxLenDialect + 3);
     else
       os << llvm::right_justify(dialectName, maxLenDialect + 2) << '.';
 
     // Left justify the operation name.
-    os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
-       << '\n';
+    os << llvm::left_justify(opName, maxLenOpName) << " , "
+       << count.find(key)->second << '\n';
   }
 }
 
 void PrintOpStatsPass::printTopLocations() {
   os << '\n'
      << "Top locations by number of operations:\n"
      << "--------------------------------------\n";
 
   // Sort by value (descending).
-  using EntryTy = std::pair<Location, int64_t>;
+  using EntryTy = std::pair<Location, OpCounter>;
   std::vector<EntryTy> sorted;
   for (const auto &it : opCountByLoc)
     sorted.push_back(it);
-  llvm::sort(sorted,
-             [](EntryTy e1, EntryTy e2) { return e1.second > e2.second; });
+  llvm::sort(sorted, [](EntryTy e1, EntryTy e2) {
+    return e1.second.total > e2.second.total;
+  });
 
   // Take top entries.
-  int64_t numLocs = sorted.size() < topKLocs ? sorted.size() : topKLocs;
+  int64_t numLocs = std::min<int64_t>(sorted.size(), topKLocs);
   for (int i = 0; i < numLocs; ++i) {
     std::string buffer;
     llvm::raw_string_ostream strOs(buffer);
     sorted[i].first.print(strOs);
     std::string str = strOs.str();
     size_t locSize = str.size();
     static const int kMaxLocLen = 50;
 
     if (locSize <= kMaxLocLen)
       os << llvm::right_justify(str, kMaxLocLen + 3);
     else
       os << "..." << str.substr(locSize - kMaxLocLen);
 
-    os << " , " << sorted[i].second << '\n';
+    os << " , " << sorted[i].second.total << '\n';
+
+    if (topKLocsOpCounts)
+      printOpCounts(sorted[i].second.map);
   }
 }
 
 std::unique_ptr<Pass> mlir::createPrintOpStatsPass() {
   return std::make_unique<PrintOpStatsPass>();
 }

mlir/test/CAPI/pass.c

void testPrintPassPipeline() {
   MlirOpPassManager nestedModulePm = mlirPassManagerGetNestedUnder(
       pm, mlirStringRefCreateFromCString("builtin.module"));
   MlirOpPassManager nestedFuncPm = mlirOpPassManagerGetNestedUnder(
       nestedModulePm, mlirStringRefCreateFromCString("builtin.func"));
   MlirPass printOpStatPass = mlirCreateTransformsPrintOpStats();
   mlirOpPassManagerAddOwnedPass(nestedFuncPm, printOpStatPass);
 
   // Print the top level pass manager
-  // CHECK: Top-level: builtin.module(builtin.func(print-op-stats{top-k-locs=0}))
+  // CHECK: Top-level: builtin.module(builtin.func(print-op-stats{top-k-locs=0 top-k-locs-op-counts=false}))
   fprintf(stderr, "Top-level: ");
   mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
                         NULL);
   fprintf(stderr, "\n");
 
   // Print the pipeline nested one level down
-  // CHECK: Nested Module: builtin.func(print-op-stats{top-k-locs=0})
+  // CHECK: Nested Module: builtin.func(print-op-stats{top-k-locs=0 top-k-locs-op-counts=false})
   fprintf(stderr, "Nested Module: ");
   mlirPrintPassPipeline(nestedModulePm, printToStderr, NULL);
   fprintf(stderr, "\n");
 
   // Print the pipeline nested two levels down
   // CHECK: Nested Module>Func: print-op-stats
   fprintf(stderr, "Nested Module>Func: ");
   mlirPrintPassPipeline(nestedFuncPm, printToStderr, NULL);
status = mlirParsePassPipeline(
       mlirStringRefCreateFromCString(
           "builtin.module(builtin.func(print-op-stats), builtin.func(print-op-stats))"));
   // Expect a failure, we haven't registered the print-op-stats pass yet.
   if (mlirLogicalResultIsFailure(status)) {
     fprintf(stderr, "Unexpected failure parsing pipeline after registering the pass\n");
     exit(EXIT_FAILURE);
   }
 
-  // CHECK: Round-trip: builtin.module(builtin.func(print-op-stats{top-k-locs=0}), builtin.func(print-op-stats{top-k-locs=0}))
+  // CHECK: Round-trip: builtin.module(builtin.func(print-op-stats{top-k-locs=0 top-k-locs-op-counts=false}),
+  // CHECK-SAME: builtin.func(print-op-stats{top-k-locs=0 top-k-locs-op-counts=false}))
   fprintf(stderr, "Round-trip: ");
   mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
                         NULL);
   fprintf(stderr, "\n");
 }
 
 int main() {
   testRunPassOnModule();
   testRunPassOnNestedModule();
   testPrintPassPipeline();
   testParsePassPipeline();
   return 0;
 }

mlir/test/IR/op-stats.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -print-op-stats="top-k-locs=5" %s -o=/dev/null 2>&1 | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -print-op-stats="top-k-locs=5 top-k-locs-op-counts" %s -o=/dev/null 2>&1 | FileCheck %s
 
 func @main(tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> {
 ^bb0(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>):
   %0 = addf %arg0, %arg1 : tensor<4xf32>
   %1 = addf %arg0, %arg1 : tensor<4xf32>
   %2 = addf %arg0, %arg1 : tensor<4xf32>
   %3 = addf %arg0, %arg1 : tensor<4xf32>
   %4 = addf %arg0, %arg1 : tensor<4xf32>

mlir/test/python/pass_manager.py

except ValueError as e:
       # CHECK: ValueError exception: invalid pass pipeline 'builtin.module(builtin.func(print-op-stats))'.
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 
     # This will register the pass and round-trip should be possible now.
     import mlir.transforms
     pm = PassManager.parse("builtin.module(builtin.func(print-op-stats))")
-    # CHECK: Roundtrip: builtin.module(builtin.func(print-op-stats{top-k-locs=0}))
+    # CHECK: Roundtrip: builtin.module(builtin.func(print-op-stats{top-k-locs=0 top-k-locs-op-counts=false}))
     log("Roundtrip: ", pm)
 run(testParseSuccess)
 
 # Verify failure on unregistered pass.
 # CHECK-LABEL: TEST: testParseFail
 def testParseFail():
   with Context():
     try: