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

source code of mlir/lib/Transforms/LocationSnapshot.cpp