mlir/lib/Transforms/LocationSnapshot.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-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
  19 namespace mlir {
  20 #define GEN_PASS_DEF_LOCATIONSNAPSHOT
  21 #include "mlir/Transforms/Passes.h.inc"
  22 } // namespace mlir
  23
  24 using 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.
  31 static 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.
  71 static 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             "mlir_snapshot", "tmp.mlir", 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(filepath, &error);
  89   if (!outputFile)
  90     return op->emitError() << error;
  91
  92   // Generate the intermediate locations.
  93   generateLocationsFromIR(outputFile->os(), 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.
 101 void 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.
 108 LogicalResult 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.
 117 void 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.
 125 LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag,
 126                                             Operation *op,
 127                                             OpPrintingFlags flags) {
 128   return ::generateLocationsFromIR(fileName, op, flags, tag);
 129 }
 130
 131 namespace {
 132 struct 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
 152 std::unique_ptr<Pass> mlir::createLocationSnapshotPass(OpPrintingFlags flags,
 153                                                        StringRef fileName,
 154                                                        StringRef tag) {
 155   return std::make_unique<LocationSnapshotPass>(flags, fileName, tag);
 156 }
 157 std::unique_ptr<Pass> mlir::createLocationSnapshotPass() {
 158   return std::make_unique<LocationSnapshotPass>();
 159 }