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