[mlir] print-op-stats: Print top-k locations

[springerm]

Summary

Add a new option for printing the top-k locations by "number of operations".

This is useful for debugging code size explosions: Those locations that result in a large number of operations after running a lowering are printed.

[springerm]

Example output:

RUN: mlir-opt -allow-unregistered-dialect -print-op-stats="top-k-locs=5" %s -o=/dev/null 2>&1 [OK]
mlir-opt -allow-unregistered-dialect -print-op-stats="top-k-locs=5" llvm-project/mlir/test/IR/op-stats.mlir -o=/dev/null 2>&1
Operations encountered:
-----------------------
  builtin.func         , 1
  builtin.module       , 1
          long_op_name , 1
      std.addf         , 6
      std.return       , 1
      xla.add          , 17

Top locations by number of operations:
--------------------------------------
...llvm/llvm-project/mlir/test/IR/op-stats.mlir":7:8) , 1
...lvm/llvm-project/mlir/test/IR/op-stats.mlir":28:9) , 1
...lvm/llvm-project/mlir/test/IR/op-stats.mlir":11:8) , 1
...llvm/llvm-project/mlir/test/IR/op-stats.mlir":9:8) , 1
...lvm/llvm-project/mlir/test/IR/op-stats.mlir":19:9) , 1

[springerm]

rebase

[jpienaar]

Is there a way to potentially compose these? E.g., it feels like (and I don't necessarily have good suggestion yet) that we have print op count, print top K locations (well which we could also just have print count per location and then defer to "grep loc\( | head -N K") and the follow up count per top loc. Basically, how much do we want to make this smarter vs relying on common unix tools? (summing per location using awk I invariably need to look up awk syntax again :) so counting is useful, but we could always emit a summary and then just have a group-by-location or also-group-by-location option [and then one would just have the global op count listed too], and then defer all else to awk/sed/grep or just searching in vim).

mlir/include/mlir/Transforms/Passes.td

ListOption<"clCollapsedIndices2", "collapsed-indices-2", "unsigned",
                "Which loop indices to combine into the position 2 loop index",
                "llvm::cl::MiscFlags::CommaSeparated">,
   ];
 }
 
 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">
+  ];
 }
 
 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
     optimistically propagates this throughout the IR. Any values proven to be

mlir/lib/Transforms/OpStats.cpp

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();
+
 private:
   llvm::StringMap<int64_t> opCount;
+  llvm::DenseMap<Location, int64_t> 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()]; });
+  getOperation()->walk([&](Operation *op) {
+    ++opCount[op->getName().getStringRef()];
+    ++opCountByLoc[op->getLoc()];
+  });
+
   printSummary();
+
+  if (topKLocs)
+    printTopLocations();
 }
 
 void PrintOpStatsPass::printSummary() {
   os << "Operations encountered:\n";
   os << "-----------------------\n";
   SmallVector<StringRef, 64> sorted(opCount.keys());
   llvm::sort(sorted);
 
else
       os << llvm::right_justify(dialectName, maxLenDialect + 2) << '.';
 
     // Left justify the operation name.
     os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
        << '\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>;
+  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; });
+
+  // Take top entries.
+  int64_t numLocs = sorted.size() < topKLocs ? 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';
+  }
+}
+
 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))
+  // CHECK: Top-level: builtin.module(builtin.func(print-op-stats{top-k-locs=0}))
   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)
+  // CHECK: Nested Module: builtin.func(print-op-stats{top-k-locs=0})
   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), builtin.func(print-op-stats))
+  // CHECK: Round-trip: builtin.module(builtin.func(print-op-stats{top-k-locs=0}), builtin.func(print-op-stats{top-k-locs=0}))
   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 %s -o=/dev/null 2>&1 | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -print-op-stats="top-k-locs=5" %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>
^bb0(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>):
   return %1 : tensor<4xf32>
 }
 
 // CHECK-LABEL: Operations encountered
 // CHECK: long_op_name , 1
 // CHECK: std.addf , 6
 // CHECK: std.return , 1
 // CHECK: xla.add , 17
+// CHECK-LABEL: Top locations by number of operations
+// CHECK: IR/op-stats.mlir{{.*}} , 1
+// CHECK: IR/op-stats.mlir{{.*}} , 1
+// CHECK: IR/op-stats.mlir{{.*}} , 1
+// CHECK: IR/op-stats.mlir{{.*}} , 1
+// CHECK: IR/op-stats.mlir{{.*}} , 1

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))
+    # CHECK: Roundtrip: builtin.module(builtin.func(print-op-stats{top-k-locs=0}))
     log("Roundtrip: ", pm)
 run(testParseSuccess)
 
 # Verify failure on unregistered pass.
 # CHECK-LABEL: TEST: testParseFail
 def testParseFail():
   with Context():
     try: