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

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