1 //===- IRPrinting.cpp -----------------------------------------------------===//
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 "PassDetail.h"
10 #include "mlir/IR/Module.h"
11 #include "mlir/Pass/PassManager.h"
12 #include "llvm/Support/Format.h"
13 #include "llvm/Support/FormatVariadic.h"
14 #include "llvm/Support/SHA1.h"
15 
16 using namespace mlir;
17 using namespace mlir::detail;
18 
19 namespace {
20 //===----------------------------------------------------------------------===//
21 // OperationFingerPrint
22 //===----------------------------------------------------------------------===//
23 
24 /// A unique fingerprint for a specific operation, and all of it's internal
25 /// operations.
26 class OperationFingerPrint {
27 public:
28   OperationFingerPrint(Operation *topOp) {
29     llvm::SHA1 hasher;
30 
31     // Hash each of the operations based upon their mutable bits:
32     topOp->walk([&](Operation *op) {
33       //   - Operation pointer
34       addDataToHash(hasher, op);
35       //   - Attributes
36       addDataToHash(hasher, op->getMutableAttrDict());
37       //   - Blocks in Regions
38       for (Region &region : op->getRegions()) {
39         for (Block &block : region) {
40           addDataToHash(hasher, &block);
41           for (BlockArgument arg : block.getArguments())
42             addDataToHash(hasher, arg);
43         }
44       }
45       //   - Location
46       addDataToHash(hasher, op->getLoc().getAsOpaquePointer());
47       //   - Operands
48       for (Value operand : op->getOperands())
49         addDataToHash(hasher, operand);
50       //   - Successors
51       for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i)
52         addDataToHash(hasher, op->getSuccessor(i));
53     });
54     hash = hasher.result();
55   }
56 
57   bool operator==(const OperationFingerPrint &other) const {
58     return hash == other.hash;
59   }
60   bool operator!=(const OperationFingerPrint &other) const {
61     return !(*this == other);
62   }
63 
64 private:
65   template <typename T> void addDataToHash(llvm::SHA1 &hasher, const T &data) {
66     hasher.update(
67         ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(&data), sizeof(T)));
68   }
69 
70   SmallString<20> hash;
71 };
72 
73 //===----------------------------------------------------------------------===//
74 // IRPrinter
75 //===----------------------------------------------------------------------===//
76 
77 class IRPrinterInstrumentation : public PassInstrumentation {
78 public:
79   IRPrinterInstrumentation(std::unique_ptr<PassManager::IRPrinterConfig> config)
80       : config(std::move(config)) {}
81 
82 private:
83   /// Instrumentation hooks.
84   void runBeforePass(Pass *pass, Operation *op) override;
85   void runAfterPass(Pass *pass, Operation *op) override;
86   void runAfterPassFailed(Pass *pass, Operation *op) override;
87 
88   /// Configuration to use.
89   std::unique_ptr<PassManager::IRPrinterConfig> config;
90 
91   /// The following is a set of fingerprints for operations that are currently
92   /// being operated on in a pass. This field is only used when the
93   /// configuration asked for change detection.
94   DenseMap<Pass *, OperationFingerPrint> beforePassFingerPrints;
95 };
96 } // end anonymous namespace
97 
98 /// Returns true if the given pass is hidden from IR printing.
99 static bool isHiddenPass(Pass *pass) {
100   return isa<OpToOpPassAdaptor, VerifierPass>(pass);
101 }
102 
103 static void printIR(Operation *op, bool printModuleScope, raw_ostream &out,
104                     OpPrintingFlags flags) {
105   // Check to see if we are printing the top-level module.
106   auto module = dyn_cast<ModuleOp>(op);
107   if (module && !op->getBlock())
108     return module.print(out << "\n", flags);
109 
110   // Otherwise, check to see if we are not printing at module scope.
111   if (!printModuleScope)
112     return op->print(out << "\n", flags.useLocalScope());
113 
114   // Otherwise, we are printing at module scope.
115   out << " ('" << op->getName() << "' operation";
116   if (auto symbolName =
117           op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName()))
118     out << ": @" << symbolName.getValue();
119   out << ")\n";
120 
121   // Find the top-level module operation.
122   auto *topLevelOp = op;
123   while (auto *parentOp = topLevelOp->getParentOp())
124     topLevelOp = parentOp;
125 
126   // Check to see if the top-level operation is actually a module in the case of
127   // invalid-ir.
128   if (auto module = dyn_cast<ModuleOp>(topLevelOp))
129     module.print(out, flags);
130   else
131     topLevelOp->print(out, flags);
132 }
133 
134 /// Instrumentation hooks.
135 void IRPrinterInstrumentation::runBeforePass(Pass *pass, Operation *op) {
136   if (isHiddenPass(pass))
137     return;
138   // If the config asked to detect changes, record the current fingerprint.
139   if (config->shouldPrintAfterOnlyOnChange())
140     beforePassFingerPrints.try_emplace(pass, op);
141 
142   config->printBeforeIfEnabled(pass, op, [&](raw_ostream &out) {
143     out << formatv("// *** IR Dump Before {0} ***", pass->getName());
144     printIR(op, config->shouldPrintAtModuleScope(), out,
145             config->getOpPrintingFlags());
146     out << "\n\n";
147   });
148 }
149 
150 void IRPrinterInstrumentation::runAfterPass(Pass *pass, Operation *op) {
151   if (isHiddenPass(pass))
152     return;
153   // If the config asked to detect changes, compare the current fingerprint with
154   // the previous.
155   if (config->shouldPrintAfterOnlyOnChange()) {
156     auto fingerPrintIt = beforePassFingerPrints.find(pass);
157     assert(fingerPrintIt != beforePassFingerPrints.end() &&
158            "expected valid fingerprint");
159     // If the fingerprints are the same, we don't print the IR.
160     if (fingerPrintIt->second == OperationFingerPrint(op)) {
161       beforePassFingerPrints.erase(fingerPrintIt);
162       return;
163     }
164     beforePassFingerPrints.erase(fingerPrintIt);
165   }
166 
167   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
168     out << formatv("// *** IR Dump After {0} ***", pass->getName());
169     printIR(op, config->shouldPrintAtModuleScope(), out,
170             config->getOpPrintingFlags());
171     out << "\n\n";
172   });
173 }
174 
175 void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass, Operation *op) {
176   if (isa<OpToOpPassAdaptor>(pass))
177     return;
178   if (config->shouldPrintAfterOnlyOnChange())
179     beforePassFingerPrints.erase(pass);
180 
181   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
182     out << formatv("// *** IR Dump After {0} Failed ***", pass->getName());
183     printIR(op, config->shouldPrintAtModuleScope(), out,
184             OpPrintingFlags().printGenericOpForm());
185     out << "\n\n";
186   });
187 }
188 
189 //===----------------------------------------------------------------------===//
190 // IRPrinterConfig
191 //===----------------------------------------------------------------------===//
192 
193 /// Initialize the configuration.
194 PassManager::IRPrinterConfig::IRPrinterConfig(bool printModuleScope,
195                                               bool printAfterOnlyOnChange,
196                                               OpPrintingFlags opPrintingFlags)
197     : printModuleScope(printModuleScope),
198       printAfterOnlyOnChange(printAfterOnlyOnChange),
199       opPrintingFlags(opPrintingFlags) {}
200 PassManager::IRPrinterConfig::~IRPrinterConfig() {}
201 
202 /// A hook that may be overridden by a derived config that checks if the IR
203 /// of 'operation' should be dumped *before* the pass 'pass' has been
204 /// executed. If the IR should be dumped, 'printCallback' should be invoked
205 /// with the stream to dump into.
206 void PassManager::IRPrinterConfig::printBeforeIfEnabled(
207     Pass *pass, Operation *operation, PrintCallbackFn printCallback) {
208   // By default, never print.
209 }
210 
211 /// A hook that may be overridden by a derived config that checks if the IR
212 /// of 'operation' should be dumped *after* the pass 'pass' has been
213 /// executed. If the IR should be dumped, 'printCallback' should be invoked
214 /// with the stream to dump into.
215 void PassManager::IRPrinterConfig::printAfterIfEnabled(
216     Pass *pass, Operation *operation, PrintCallbackFn printCallback) {
217   // By default, never print.
218 }
219 
220 //===----------------------------------------------------------------------===//
221 // PassManager
222 //===----------------------------------------------------------------------===//
223 
224 namespace {
225 /// Simple wrapper config that allows for the simpler interface defined above.
226 struct BasicIRPrinterConfig : public PassManager::IRPrinterConfig {
227   BasicIRPrinterConfig(
228       std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
229       std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
230       bool printModuleScope, bool printAfterOnlyOnChange,
231       OpPrintingFlags opPrintingFlags, raw_ostream &out)
232       : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange,
233                         opPrintingFlags),
234         shouldPrintBeforePass(shouldPrintBeforePass),
235         shouldPrintAfterPass(shouldPrintAfterPass), out(out) {
236     assert((shouldPrintBeforePass || shouldPrintAfterPass) &&
237            "expected at least one valid filter function");
238   }
239 
240   void printBeforeIfEnabled(Pass *pass, Operation *operation,
241                             PrintCallbackFn printCallback) final {
242     if (shouldPrintBeforePass && shouldPrintBeforePass(pass, operation))
243       printCallback(out);
244   }
245 
246   void printAfterIfEnabled(Pass *pass, Operation *operation,
247                            PrintCallbackFn printCallback) final {
248     if (shouldPrintAfterPass && shouldPrintAfterPass(pass, operation))
249       printCallback(out);
250   }
251 
252   /// Filter functions for before and after pass execution.
253   std::function<bool(Pass *, Operation *)> shouldPrintBeforePass;
254   std::function<bool(Pass *, Operation *)> shouldPrintAfterPass;
255 
256   /// The stream to output to.
257   raw_ostream &out;
258 };
259 } // end anonymous namespace
260 
261 /// Add an instrumentation to print the IR before and after pass execution,
262 /// using the provided configuration.
263 void PassManager::enableIRPrinting(std::unique_ptr<IRPrinterConfig> config) {
264   if (config->shouldPrintAtModuleScope() &&
265       getContext()->isMultithreadingEnabled())
266     llvm::report_fatal_error("IR printing can't be setup on a pass-manager "
267                              "without disabling multi-threading first.");
268   addInstrumentation(
269       std::make_unique<IRPrinterInstrumentation>(std::move(config)));
270 }
271 
272 /// Add an instrumentation to print the IR before and after pass execution.
273 void PassManager::enableIRPrinting(
274     std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
275     std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
276     bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out,
277     OpPrintingFlags opPrintingFlags) {
278   enableIRPrinting(std::make_unique<BasicIRPrinterConfig>(
279       std::move(shouldPrintBeforePass), std::move(shouldPrintAfterPass),
280       printModuleScope, printAfterOnlyOnChange, opPrintingFlags, out));
281 }
282