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