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