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
18namespace mlir {
19#define GEN_PASS_DEF_PRINTOPSTATS
20#include "mlir/Transforms/Passes.h.inc"
21} // namespace mlir
22
23using namespace mlir;
24
25namespace {
26struct 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
42private:
43 llvm::StringMap<int64_t> opCount;
44 raw_ostream &os;
45};
46} // namespace
47
48void 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
60void PrintOpStatsPass::printSummary() {
61 os << "Operations encountered:\n";
62 os << "-----------------------\n";
63 SmallVector<StringRef, 64> sorted(opCount.keys());
64 llvm::sort(C&: sorted);
65
66 // Split an operation name from its dialect prefix.
67 auto splitOperationName = [](StringRef opName) {
68 auto splitName = opName.split(Separator: '.');
69 return splitName.second.empty() ? std::make_pair(x: "", y&: 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(a: maxLenDialect, b: dialectName.size());
78 maxLenOpName = std::max(a: maxLenOpName, b: 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(NumSpaces: maxLenDialect + 3);
89 else
90 os << llvm::right_justify(Str: dialectName, Width: maxLenDialect + 2) << '.';
91
92 // Left justify the operation name.
93 os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
94 << '\n';
95 }
96}
97
98void PrintOpStatsPass::printSummaryInJSON() {
99 SmallVector<StringRef, 64> sorted(opCount.keys());
100 llvm::sort(C&: 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
115std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os) {
116 return std::make_unique<PrintOpStatsPass>(args&: os);
117}
118
119std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os,
120 bool printAsJSON) {
121 return std::make_unique<PrintOpStatsPass>(args&: os, args&: printAsJSON);
122}
123

source code of mlir/lib/Transforms/OpStats.cpp