MLIR  16.0.0git
OpStats.cpp
Go to the documentation of this file.
1 //===- OpStats.cpp - Prints stats of operations in module -----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
10 
11 #include "mlir/IR/BuiltinOps.h"
12 #include "mlir/IR/Operation.h"
14 #include "llvm/ADT/DenseMap.h"
15 #include "llvm/Support/Format.h"
16 #include "llvm/Support/raw_ostream.h"
17 
18 namespace mlir {
19 #define GEN_PASS_DEF_PRINTOPSTATS
20 #include "mlir/Transforms/Passes.h.inc"
21 } // namespace mlir
22 
23 using namespace mlir;
24 
25 namespace {
26 struct PrintOpStatsPass : public impl::PrintOpStatsBase<PrintOpStatsPass> {
27  explicit PrintOpStatsPass(raw_ostream &os) : os(os) {}
28 
29  explicit PrintOpStatsPass(raw_ostream &os, bool printAsJSON) : os(os) {
30  this->printAsJSON = printAsJSON;
31  }
32 
33  // Prints the resultant operation statistics post iterating over the module.
34  void runOnOperation() override;
35 
36  // Print summary of op stats.
37  void printSummary();
38 
39  // Print symmary of op stats in JSON.
40  void printSummaryInJSON();
41 
42 private:
43  llvm::StringMap<int64_t> opCount;
44  raw_ostream &os;
45 };
46 } // namespace
47 
48 void PrintOpStatsPass::runOnOperation() {
49  opCount.clear();
50 
51  // Compute the operation statistics for the currently visited operation.
52  getOperation()->walk(
53  [&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
54  if (printAsJSON) {
55  printSummaryInJSON();
56  } else
57  printSummary();
58 }
59 
60 void PrintOpStatsPass::printSummary() {
61  os << "Operations encountered:\n";
62  os << "-----------------------\n";
63  SmallVector<StringRef, 64> sorted(opCount.keys());
64  llvm::sort(sorted);
65 
66  // Split an operation name from its dialect prefix.
67  auto splitOperationName = [](StringRef opName) {
68  auto splitName = opName.split('.');
69  return splitName.second.empty() ? std::make_pair("", splitName.first)
70  : splitName;
71  };
72 
73  // Compute the largest dialect and operation name.
74  size_t maxLenOpName = 0, maxLenDialect = 0;
75  for (const auto &key : sorted) {
76  auto [dialectName, opName] = splitOperationName(key);
77  maxLenDialect = std::max(maxLenDialect, dialectName.size());
78  maxLenOpName = std::max(maxLenOpName, opName.size());
79  }
80 
81  for (const auto &key : sorted) {
82  auto [dialectName, opName] = splitOperationName(key);
83 
84  // Left-align the names (aligning on the dialect) and right-align the count
85  // below. The alignment is for readability and does not affect CSV/FileCheck
86  // parsing.
87  if (dialectName.empty())
88  os.indent(maxLenDialect + 3);
89  else
90  os << llvm::right_justify(dialectName, maxLenDialect + 2) << '.';
91 
92  // Left justify the operation name.
93  os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
94  << '\n';
95  }
96 }
97 
98 void PrintOpStatsPass::printSummaryInJSON() {
99  SmallVector<StringRef, 64> sorted(opCount.keys());
100  llvm::sort(sorted);
101 
102  os << "{\n";
103 
104  for (unsigned i = 0, e = sorted.size(); i != e; ++i) {
105  const auto &key = sorted[i];
106  os << " \"" << key << "\" : " << opCount[key];
107  if (i != e - 1)
108  os << ",\n";
109  else
110  os << "\n";
111  }
112  os << "}\n";
113 }
114 
115 std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os) {
116  return std::make_unique<PrintOpStatsPass>(os);
117 }
118 
119 std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os,
120  bool printAsJSON) {
121  return std::make_unique<PrintOpStatsPass>(os, printAsJSON);
122 }
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:31
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:50
Include the generated interface declarations.
std::unique_ptr< Pass > createPrintOpStatsPass(raw_ostream &os=llvm::errs())
Creates a pass which prints the list of ops and the number of occurrences in the module.
Definition: OpStats.cpp:115