1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
//===- OpStats.cpp - Prints stats of operations in module -----------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "mlir/Transforms/Passes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Format.h"
#include "llvm/Support/raw_ostream.h"
namespace mlir {
#define GEN_PASS_DEF_PRINTOPSTATS
#include "mlir/Transforms/Passes.h.inc"
} // namespace mlir
using namespace mlir;
namespace {
struct PrintOpStatsPass : public impl::PrintOpStatsBase<PrintOpStatsPass> {
explicit PrintOpStatsPass(raw_ostream &os) : os(os) {}
explicit PrintOpStatsPass(raw_ostream &os, bool printAsJSON) : os(os) {
this->printAsJSON = printAsJSON;
}
// Prints the resultant operation statistics post iterating over the module.
void runOnOperation() override;
// Print summary of op stats.
void printSummary();
// Print symmary of op stats in JSON.
void printSummaryInJSON();
private:
llvm::StringMap<int64_t> opCount;
raw_ostream &os;
};
} // namespace
void PrintOpStatsPass::runOnOperation() {
opCount.clear();
// Compute the operation statistics for the currently visited operation.
getOperation()->walk(
[&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
if (printAsJSON) {
printSummaryInJSON();
} else
printSummary();
}
void PrintOpStatsPass::printSummary() {
os << "Operations encountered:\n";
os << "-----------------------\n";
SmallVector<StringRef, 64> sorted(opCount.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;
};
// Compute the largest dialect and operation name.
size_t maxLenOpName = 0, maxLenDialect = 0;
for (const auto &key : sorted) {
auto [dialectName, opName] = splitOperationName(key);
maxLenDialect = std::max(maxLenDialect, dialectName.size());
maxLenOpName = std::max(maxLenOpName, opName.size());
}
for (const auto &key : sorted) {
auto [dialectName, opName] = splitOperationName(key);
// Left-align the names (aligning on the dialect) and right-align the count
// 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';
}
}
void PrintOpStatsPass::printSummaryInJSON() {
SmallVector<StringRef, 64> sorted(opCount.keys());
llvm::sort(sorted);
os << "{\n";
for (unsigned i = 0, e = sorted.size(); i != e; ++i) {
const auto &key = sorted[i];
os << " \"" << key << "\" : " << opCount[key];
if (i != e - 1)
os << ",\n";
else
os << "\n";
}
os << "}\n";
}
std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os) {
return std::make_unique<PrintOpStatsPass>(os);
}
std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os,
bool printAsJSON) {
return std::make_unique<PrintOpStatsPass>(os, printAsJSON);
}
|
