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 20 namespace mlir { 21 #define GEN_PASS_DEF_LOCATIONSNAPSHOT 22 #include "mlir/Transforms/Passes.h.inc" 23 } // namespace mlir 24 25 using 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. 32 static 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. 72 static 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 "mlir_snapshot", "tmp.mlir", 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(filepath, &error); 90 if (!outputFile) 91 return op->emitError() << error; 92 93 // Generate the intermediate locations. 94 generateLocationsFromIR(outputFile->os(), 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. 102 void 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. 109 LogicalResult 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. 118 void 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. 126 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag, 127 Operation *op, 128 OpPrintingFlags flags) { 129 return ::generateLocationsFromIR(fileName, op, flags, tag); 130 } 131 132 namespace { 133 struct 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 143 private: 144 /// build the flags from the command line arguments to the pass 145 OpPrintingFlags getFlags() { 146 OpPrintingFlags flags; 147 flags.enableDebugInfo(enableDebugInfo, printPrettyDebugInfo); 148 flags.printGenericOpForm(printGenericOpForm); 149 if (useLocalScope) 150 flags.useLocalScope(); 151 return flags; 152 } 153 }; 154 } // namespace 155