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,OpPrintingFlags flags,StringRef tag)24 static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
25 Operation *op, 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, &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 =
48 builder.getFileLineColLoc(file, lineCol.first, lineCol.second);
49
50 // If we don't have a tag, set the location directly
51 if (!tagIdentifier) {
52 opIt->setLoc(newLoc);
53 return;
54 }
55
56 // Otherwise, build a fused location with the existing op loc.
57 opIt->setLoc(builder.getFusedLoc(
58 {opIt->getLoc(), NameLoc::get(*tagIdentifier, newLoc)}));
59 });
60 }
61
62 /// This function generates new locations from the given IR by snapshotting the
63 /// IR to the given file, and using the printed locations within that file. If
64 /// `filename` is empty, a temporary file is generated instead.
generateLocationsFromIR(StringRef fileName,Operation * op,OpPrintingFlags flags,StringRef tag)65 static LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
66 OpPrintingFlags flags,
67 StringRef tag) {
68 // If a filename wasn't provided, then generate one.
69 SmallString<32> filepath(fileName);
70 if (filepath.empty()) {
71 if (std::error_code error = llvm::sys::fs::createTemporaryFile(
72 "mlir_snapshot", "tmp.mlir", filepath)) {
73 return op->emitError()
74 << "failed to generate temporary file for location snapshot: "
75 << error.message();
76 }
77 }
78
79 // Open the output file for emission.
80 std::string error;
81 std::unique_ptr<llvm::ToolOutputFile> outputFile =
82 openOutputFile(filepath, &error);
83 if (!outputFile)
84 return op->emitError() << error;
85
86 // Generate the intermediate locations.
87 generateLocationsFromIR(outputFile->os(), filepath, op, flags, tag);
88 outputFile->keep();
89 return success();
90 }
91
92 /// This function generates new locations from the given IR by snapshotting the
93 /// IR to the given stream, and using the printed locations within that stream.
94 /// The generated locations replace the current operation locations.
generateLocationsFromIR(raw_ostream & os,StringRef fileName,Operation * op,OpPrintingFlags flags)95 void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
96 Operation *op, OpPrintingFlags flags) {
97 ::generateLocationsFromIR(os, fileName, op, flags, /*tag=*/StringRef());
98 }
99 /// This function generates new locations from the given IR by snapshotting the
100 /// IR to the given file, and using the printed locations within that file. If
101 /// `filename` is empty, a temporary file is generated instead.
generateLocationsFromIR(StringRef fileName,Operation * op,OpPrintingFlags flags)102 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, Operation *op,
103 OpPrintingFlags flags) {
104 return ::generateLocationsFromIR(fileName, op, flags, /*tag=*/StringRef());
105 }
106
107 /// This function generates new locations from the given IR by snapshotting the
108 /// IR to the given stream, and using the printed locations within that stream.
109 /// The generated locations are represented as a NameLoc with the given tag as
110 /// the name, and then fused with the existing locations.
generateLocationsFromIR(raw_ostream & os,StringRef fileName,StringRef tag,Operation * op,OpPrintingFlags flags)111 void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
112 StringRef tag, Operation *op,
113 OpPrintingFlags flags) {
114 ::generateLocationsFromIR(os, fileName, op, flags, tag);
115 }
116 /// This function generates new locations from the given IR by snapshotting the
117 /// IR to the given file, and using the printed locations within that file. If
118 /// `filename` is empty, a temporary file is generated instead.
generateLocationsFromIR(StringRef fileName,StringRef tag,Operation * op,OpPrintingFlags flags)119 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag,
120 Operation *op,
121 OpPrintingFlags flags) {
122 return ::generateLocationsFromIR(fileName, op, flags, tag);
123 }
124
125 namespace {
126 struct LocationSnapshotPass
127 : public LocationSnapshotBase<LocationSnapshotPass> {
128 LocationSnapshotPass() = default;
LocationSnapshotPass__anonabdcde090211::LocationSnapshotPass129 LocationSnapshotPass(OpPrintingFlags flags, StringRef fileName, StringRef tag)
130 : flags(flags) {
131 this->fileName = fileName.str();
132 this->tag = tag.str();
133 }
134
runOnOperation__anonabdcde090211::LocationSnapshotPass135 void runOnOperation() override {
136 Operation *op = getOperation();
137 if (failed(generateLocationsFromIR(fileName, op, OpPrintingFlags(), tag)))
138 return signalPassFailure();
139 }
140
141 /// The printing flags to use when creating the snapshot.
142 OpPrintingFlags flags;
143 };
144 } // end anonymous namespace
145
createLocationSnapshotPass(OpPrintingFlags flags,StringRef fileName,StringRef tag)146 std::unique_ptr<Pass> mlir::createLocationSnapshotPass(OpPrintingFlags flags,
147 StringRef fileName,
148 StringRef tag) {
149 return std::make_unique<LocationSnapshotPass>(flags, fileName, tag);
150 }
createLocationSnapshotPass()151 std::unique_ptr<Pass> mlir::createLocationSnapshotPass() {
152 return std::make_unique<LocationSnapshotPass>();
153 }
154