155 lines · cpp
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-exception6//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_LOCATIONSNAPSHOT22#include "mlir/Transforms/Passes.h.inc"23} // namespace mlir24 25using namespace mlir;26 27/// This function generates new locations from the given IR by snapshotting the28/// 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 a30/// NameLoc with the given tag as the name, and then fused with the existing31/// 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 may49 // be elided from the printed form, e.g. the body terminators of some region50 // 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 directly58 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 the70/// IR to the given file, and using the printed locations within that file. If71/// `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 "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 the100/// 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 the107/// IR to the given file, and using the printed locations within that file. If108/// `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 the115/// 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 as117/// 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 the124/// IR to the given file, and using the printed locations within that file. If125/// `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 LocationSnapshotPass134 : 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 pass145 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} // namespace155