MLIR  20.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"
13 #include "mlir/Pass/Pass.h"
15 #include "llvm/Support/FileSystem.h"
16 #include "llvm/Support/ToolOutputFile.h"
17 #include <optional>
18 
19 namespace mlir {
20 #define GEN_PASS_DEF_LOCATIONSNAPSHOT
21 #include "mlir/Transforms/Passes.h.inc"
22 } // namespace mlir
23 
24 using namespace mlir;
25 
26 /// This function generates new locations from the given IR by snapshotting the
27 /// IR to the given stream, and using the printed locations within that stream.
28 /// If a 'tag' is non-empty, the generated locations are represented as a
29 /// NameLoc with the given tag as the name, and then fused with the existing
30 /// locations. Otherwise, the existing locations are replaced.
31 static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
32  Operation *op, const OpPrintingFlags &flags,
33  StringRef tag) {
34  // Print the IR to the stream, and collect the raw line+column information.
35  AsmState::LocationMap opToLineCol;
36  AsmState state(op, flags, &opToLineCol);
37  op->print(os, state);
38 
39  Builder builder(op->getContext());
40  std::optional<StringAttr> tagIdentifier;
41  if (!tag.empty())
42  tagIdentifier = builder.getStringAttr(tag);
43 
44  // Walk and generate new locations for each of the operations.
45  StringAttr file = builder.getStringAttr(fileName);
46  op->walk([&](Operation *opIt) {
47  // Check to see if this operation has a mapped location. Some operations may
48  // be elided from the printed form, e.g. the body terminators of some region
49  // operations.
50  auto it = opToLineCol.find(opIt);
51  if (it == opToLineCol.end())
52  return;
53  const std::pair<unsigned, unsigned> &lineCol = it->second;
54  auto newLoc = FileLineColLoc::get(file, lineCol.first, lineCol.second);
55 
56  // If we don't have a tag, set the location directly
57  if (!tagIdentifier) {
58  opIt->setLoc(newLoc);
59  return;
60  }
61 
62  // Otherwise, build a fused location with the existing op loc.
63  opIt->setLoc(builder.getFusedLoc(
64  {opIt->getLoc(), NameLoc::get(*tagIdentifier, newLoc)}));
65  });
66 }
67 
68 /// This function generates new locations from the given IR by snapshotting the
69 /// IR to the given file, and using the printed locations within that file. If
70 /// `filename` is empty, a temporary file is generated instead.
71 static LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
72  OpPrintingFlags flags,
73  StringRef tag) {
74  // If a filename wasn't provided, then generate one.
75  SmallString<32> filepath(fileName);
76  if (filepath.empty()) {
77  if (std::error_code error = llvm::sys::fs::createTemporaryFile(
78  "mlir_snapshot", "tmp.mlir", filepath)) {
79  return op->emitError()
80  << "failed to generate temporary file for location snapshot: "
81  << error.message();
82  }
83  }
84 
85  // Open the output file for emission.
86  std::string error;
87  std::unique_ptr<llvm::ToolOutputFile> outputFile =
88  openOutputFile(filepath, &error);
89  if (!outputFile)
90  return op->emitError() << error;
91 
92  // Generate the intermediate locations.
93  generateLocationsFromIR(outputFile->os(), filepath, op, flags, tag);
94  outputFile->keep();
95  return success();
96 }
97 
98 /// This function generates new locations from the given IR by snapshotting the
99 /// IR to the given stream, and using the printed locations within that stream.
100 /// The generated locations replace the current operation locations.
101 void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
102  Operation *op, OpPrintingFlags flags) {
103  ::generateLocationsFromIR(os, fileName, op, flags, /*tag=*/StringRef());
104 }
105 /// This function generates new locations from the given IR by snapshotting the
106 /// IR to the given file, and using the printed locations within that file. If
107 /// `filename` is empty, a temporary file is generated instead.
108 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, Operation *op,
109  OpPrintingFlags flags) {
110  return ::generateLocationsFromIR(fileName, op, flags, /*tag=*/StringRef());
111 }
112 
113 /// This function generates new locations from the given IR by snapshotting the
114 /// IR to the given stream, and using the printed locations within that stream.
115 /// The generated locations are represented as a NameLoc with the given tag as
116 /// the name, and then fused with the existing locations.
117 void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
118  StringRef tag, Operation *op,
119  OpPrintingFlags flags) {
120  ::generateLocationsFromIR(os, fileName, op, flags, tag);
121 }
122 /// This function generates new locations from the given IR by snapshotting the
123 /// IR to the given file, and using the printed locations within that file. If
124 /// `filename` is empty, a temporary file is generated instead.
125 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag,
126  Operation *op,
127  OpPrintingFlags flags) {
128  return ::generateLocationsFromIR(fileName, op, flags, tag);
129 }
130 
131 namespace {
132 struct LocationSnapshotPass
133  : public impl::LocationSnapshotBase<LocationSnapshotPass> {
134  LocationSnapshotPass() = default;
135  LocationSnapshotPass(OpPrintingFlags flags, StringRef fileName, StringRef tag)
136  : flags(flags) {
137  this->fileName = fileName.str();
138  this->tag = tag.str();
139  }
140 
141  void runOnOperation() override {
142  Operation *op = getOperation();
143  if (failed(generateLocationsFromIR(fileName, op, OpPrintingFlags(), tag)))
144  return signalPassFailure();
145  }
146 
147  /// The printing flags to use when creating the snapshot.
148  OpPrintingFlags flags;
149 };
150 } // namespace
151 
153  StringRef fileName,
154  StringRef tag) {
155  return std::make_unique<LocationSnapshotPass>(flags, fileName, tag);
156 }
157 std::unique_ptr<Pass> mlir::createLocationSnapshotPass() {
158  return std::make_unique<LocationSnapshotPass>();
159 }
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:302
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.
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:793
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< Pass > createLocationSnapshotPass(OpPrintingFlags flags, StringRef fileName="", StringRef tag="")
Create a pass to generate new locations by snapshotting the IR to the given file, and using the print...
std::unique_ptr< llvm::ToolOutputFile > openOutputFile(llvm::StringRef outputFilename, std::string *errorMessage=nullptr)
Open the file specified by its name for writing.