MLIR  21.0.0git
LocationSnapshot.cpp
Go to the documentation of this file.
1 //===- LocationSnapshot.cpp - Location Snapshot Utilities -----------------===//
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 
10 
11 #include "mlir/IR/AsmState.h"
12 #include "mlir/IR/Builders.h"
14 #include "mlir/Pass/Pass.h"
16 #include "llvm/Support/FileSystem.h"
17 #include "llvm/Support/ToolOutputFile.h"
18 #include <optional>
19 
20 namespace mlir {
21 #define GEN_PASS_DEF_LOCATIONSNAPSHOT
22 #include "mlir/Transforms/Passes.h.inc"
23 } // namespace mlir
24 
25 using namespace mlir;
26 
27 /// This function generates new locations from the given IR by snapshotting the
28 /// IR to the given stream, and using the printed locations within that stream.
29 /// If a 'tag' is non-empty, the generated locations are represented as a
30 /// NameLoc with the given tag as the name, and then fused with the existing
31 /// locations. Otherwise, the existing locations are replaced.
32 static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
33  Operation *op, const OpPrintingFlags &flags,
34  StringRef tag) {
35  // Print the IR to the stream, and collect the raw line+column information.
36  AsmState::LocationMap opToLineCol;
37  AsmState state(op, flags, &opToLineCol);
38  op->print(os, state);
39 
40  Builder builder(op->getContext());
41  std::optional<StringAttr> tagIdentifier;
42  if (!tag.empty())
43  tagIdentifier = builder.getStringAttr(tag);
44 
45  // Walk and generate new locations for each of the operations.
46  StringAttr file = builder.getStringAttr(fileName);
47  op->walk([&](Operation *opIt) {
48  // Check to see if this operation has a mapped location. Some operations may
49  // be elided from the printed form, e.g. the body terminators of some region
50  // operations.
51  auto it = opToLineCol.find(opIt);
52  if (it == opToLineCol.end())
53  return;
54  const std::pair<unsigned, unsigned> &lineCol = it->second;
55  auto newLoc = FileLineColLoc::get(file, lineCol.first, lineCol.second);
56 
57  // If we don't have a tag, set the location directly
58  if (!tagIdentifier) {
59  opIt->setLoc(newLoc);
60  return;
61  }
62 
63  // Otherwise, build a fused location with the existing op loc.
64  opIt->setLoc(builder.getFusedLoc(
65  {opIt->getLoc(), NameLoc::get(*tagIdentifier, newLoc)}));
66  });
67 }
68 
69 /// This function generates new locations from the given IR by snapshotting the
70 /// IR to the given file, and using the printed locations within that file. If
71 /// `filename` is empty, a temporary file is generated instead.
72 static LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
73  OpPrintingFlags flags,
74  StringRef tag) {
75  // If a filename wasn't provided, then generate one.
76  SmallString<32> filepath(fileName);
77  if (filepath.empty()) {
78  if (std::error_code error = llvm::sys::fs::createTemporaryFile(
79  "mlir_snapshot", "tmp.mlir", filepath)) {
80  return op->emitError()
81  << "failed to generate temporary file for location snapshot: "
82  << error.message();
83  }
84  }
85 
86  // Open the output file for emission.
87  std::string error;
88  std::unique_ptr<llvm::ToolOutputFile> outputFile =
89  openOutputFile(filepath, &error);
90  if (!outputFile)
91  return op->emitError() << error;
92 
93  // Generate the intermediate locations.
94  generateLocationsFromIR(outputFile->os(), filepath, op, flags, tag);
95  outputFile->keep();
96  return success();
97 }
98 
99 /// This function generates new locations from the given IR by snapshotting the
100 /// IR to the given stream, and using the printed locations within that stream.
101 /// The generated locations replace the current operation locations.
102 void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
103  Operation *op, OpPrintingFlags flags) {
104  ::generateLocationsFromIR(os, fileName, op, flags, /*tag=*/StringRef());
105 }
106 /// This function generates new locations from the given IR by snapshotting the
107 /// IR to the given file, and using the printed locations within that file. If
108 /// `filename` is empty, a temporary file is generated instead.
109 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, Operation *op,
110  OpPrintingFlags flags) {
111  return ::generateLocationsFromIR(fileName, op, flags, /*tag=*/StringRef());
112 }
113 
114 /// This function generates new locations from the given IR by snapshotting the
115 /// IR to the given stream, and using the printed locations within that stream.
116 /// The generated locations are represented as a NameLoc with the given tag as
117 /// the name, and then fused with the existing locations.
118 void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
119  StringRef tag, Operation *op,
120  OpPrintingFlags flags) {
121  ::generateLocationsFromIR(os, fileName, op, flags, tag);
122 }
123 /// This function generates new locations from the given IR by snapshotting the
124 /// IR to the given file, and using the printed locations within that file. If
125 /// `filename` is empty, a temporary file is generated instead.
126 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag,
127  Operation *op,
128  OpPrintingFlags flags) {
129  return ::generateLocationsFromIR(fileName, op, flags, tag);
130 }
131 
132 namespace {
133 struct LocationSnapshotPass
134  : public impl::LocationSnapshotBase<LocationSnapshotPass> {
135  using impl::LocationSnapshotBase<LocationSnapshotPass>::LocationSnapshotBase;
136 
137  void runOnOperation() override {
138  Operation *op = getOperation();
139  if (failed(generateLocationsFromIR(fileName, op, getFlags(), tag)))
140  return signalPassFailure();
141  }
142 
143 private:
144  /// build the flags from the command line arguments to the pass
145  OpPrintingFlags getFlags() {
146  OpPrintingFlags flags;
147  flags.enableDebugInfo(enableDebugInfo, printPrettyDebugInfo);
148  flags.printGenericOpForm(printGenericOpForm);
149  if (useLocalScope)
150  flags.useLocalScope();
151  return flags;
152  }
153 };
154 } // namespace
static void generateLocationsFromIR(raw_ostream &os, StringRef fileName, Operation *op, const OpPrintingFlags &flags, StringRef tag)
This function generates new locations from the given IR by snapshotting the IR to the given stream,...
This class provides management for the lifetime of the state used when printing the IR.
Definition: AsmState.h:540
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:51
StringAttr getStringAttr(const Twine &bytes)
Definition: Builders.cpp:258
Location getFusedLoc(ArrayRef< Location > locs, Attribute metadata=Attribute())
Definition: Builders.cpp:29
static FileLineColLoc get(StringAttr filename, unsigned line, unsigned column)
Definition: Location.cpp:160
Set of flags used to control the behavior of the various IR print methods (e.g.
OpPrintingFlags & enableDebugInfo(bool enable=true, bool prettyForm=false)
Enable or disable printing of debug information (based on enable).
Definition: AsmPrinter.cpp:269
OpPrintingFlags & printGenericOpForm(bool enable=true)
Always print operations in the generic form.
Definition: AsmPrinter.cpp:277
OpPrintingFlags & useLocalScope(bool enable=true)
Use local scope when printing the operation.
Definition: AsmPrinter.cpp:297
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
void setLoc(Location loc)
Set the source location the operation was defined or derived from.
Definition: Operation.h:226
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
Definition: Operation.h:798
void print(raw_ostream &os, const OpPrintingFlags &flags=std::nullopt)
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:268
Include the generated interface declarations.
void generateLocationsFromIR(raw_ostream &os, StringRef fileName, Operation *op, OpPrintingFlags flags)
This function generates new locations from the given IR by snapshotting the IR to the given stream,...
std::unique_ptr< llvm::ToolOutputFile > openOutputFile(llvm::StringRef outputFilename, std::string *errorMessage=nullptr)
Open the file specified by its name for writing.