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.
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.
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.
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.
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.
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.
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;
129   LocationSnapshotPass(OpPrintingFlags flags, StringRef fileName, StringRef tag)
130       : flags(flags) {
131     this->fileName = fileName.str();
132     this->tag = tag.str();
133   }
134 
135   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 
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 }
151 std::unique_ptr<Pass> mlir::createLocationSnapshotPass() {
152   return std::make_unique<LocationSnapshotPass>();
153 }
154