llvm-project

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

9
#include "mlir/Transforms/Passes.h"
10

11
#include "mlir/IR/BuiltinOps.h"
12
#include "mlir/IR/Operation.h"
13
#include "mlir/IR/OperationSupport.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
}
123