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 markAllAnalysesPreserved();
59}
60
61void PrintOpStatsPass::printSummary() {
62 os << "Operations encountered:\n";
63 os << "-----------------------\n";
64 SmallVector<StringRef, 64> sorted(opCount.keys());
65 llvm::sort(C&: sorted);
66
67 // Split an operation name from its dialect prefix.
68 auto splitOperationName = [](StringRef opName) {
69 auto splitName = opName.split(Separator: '.');
70 return splitName.second.empty() ? std::make_pair(x: "", y&: splitName.first)
71 : splitName;
72 };
73
74 // Compute the largest dialect and operation name.
75 size_t maxLenOpName = 0, maxLenDialect = 0;
76 for (const auto &key : sorted) {
77 auto [dialectName, opName] = splitOperationName(key);
78 maxLenDialect = std::max(a: maxLenDialect, b: dialectName.size());
79 maxLenOpName = std::max(a: maxLenOpName, b: opName.size());
80 }
81
82 for (const auto &key : sorted) {
83 auto [dialectName, opName] = splitOperationName(key);
84
85 // Left-align the names (aligning on the dialect) and right-align the count
86 // below. The alignment is for readability and does not affect CSV/FileCheck
87 // parsing.
88 if (dialectName.empty())
89 os.indent(NumSpaces: maxLenDialect + 3);
90 else
91 os << llvm::right_justify(Str: dialectName, Width: maxLenDialect + 2) << '.';
92
93 // Left justify the operation name.
94 os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
95 << '\n';
96 }
97}
98
99void PrintOpStatsPass::printSummaryInJSON() {
100 SmallVector<StringRef, 64> sorted(opCount.keys());
101 llvm::sort(C&: sorted);
102
103 os << "{\n";
104
105 for (unsigned i = 0, e = sorted.size(); i != e; ++i) {
106 const auto &key = sorted[i];
107 os << " \"" << key << "\" : " << opCount[key];
108 if (i != e - 1)
109 os << ",\n";
110 else
111 os << "\n";
112 }
113 os << "}\n";
114}
115
116std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os) {
117 return std::make_unique<PrintOpStatsPass>(args&: os);
118}
119
120std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os,
121 bool printAsJSON) {
122 return std::make_unique<PrintOpStatsPass>(args&: os, args&: printAsJSON);
123}
124

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