[mlir] Add support for generating debug locations from intermediate levels of the IR.
authorRiver Riddle <riddleriver@gmail.com>
Sat, 8 Feb 2020 23:01:34 +0000 (15:01 -0800)
committerRiver Riddle <riddleriver@gmail.com>
Sat, 8 Feb 2020 23:11:29 +0000 (15:11 -0800)
Summary:
This revision adds a utility to generate debug locations from the IR during compilation, by snapshotting to a output stream and using the locations that operations were dumped in that stream. The new locations may either;
* Replace the original location of the operation.

old:
   loc("original_source.cpp":1:1)
new:
   loc("snapshot_source.mlir":10:10)

* Fuse with the original locations as NamedLocs with a specific tag.

old:
    loc("original_source.cpp":1:1)
new:
    loc(fused["original_source.cpp":1:1, "snapshot"("snapshot_source.mlir":10:10)])

This feature may be used by a debugger to display the code at various different levels of the IR. It would also be able to show the different levels of IR attached to a specific source line in the original source file.

This feature may also be used to generate locations for operations generated during compilation, that don't necessarily have a user source location to attach to.

This requires changes in the printer to track the locations of operations emitted in the stream. Moving forward we need to properly(and efficiently) track the number of newlines emitted to the stream during printing.

Differential Revision: https://reviews.llvm.org/D74019

mlir/include/mlir/IR/AsmState.h
mlir/include/mlir/Transforms/LocationSnapshot.h [new file with mode: 0644]
mlir/lib/IR/AsmPrinter.cpp
mlir/lib/Transforms/CMakeLists.txt
mlir/lib/Transforms/LocationSnapshot.cpp [new file with mode: 0644]
mlir/test/Transforms/location-snapshot.mlir [new file with mode: 0644]

index c777c99..a56f8e6 100644 (file)
@@ -13,6 +13,8 @@
 #ifndef MLIR_IR_ASMSTATE_H_
 #define MLIR_IR_ASMSTATE_H_
 
+#include "mlir/Support/LLVM.h"
+
 #include <memory>
 
 namespace mlir {
@@ -32,8 +34,14 @@ class AsmStateImpl;
 /// parent operation cannot reuse this state.
 class AsmState {
 public:
-  /// Initialize the asm state at the level of the given operation.
-  AsmState(Operation *op);
+  /// This map represents the raw locations of operations within the output
+  /// stream. This maps the original pointer to the operation, to a pair of line
+  /// and column in the output stream.
+  using LocationMap = DenseMap<Operation *, std::pair<unsigned, unsigned>>;
+
+  /// Initialize the asm state at the level of the given operation. A location
+  /// map may optionally be provided to be populated when printing.
+  AsmState(Operation *op, LocationMap *locationMap = nullptr);
   ~AsmState();
 
   /// Return an instance of the internal implementation. Returns nullptr if the
diff --git a/mlir/include/mlir/Transforms/LocationSnapshot.h b/mlir/include/mlir/Transforms/LocationSnapshot.h
new file mode 100644 (file)
index 0000000..e4c911a
--- /dev/null
@@ -0,0 +1,64 @@
+//===- LocationSnapshot.h - Location Snapshot Utilities ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file several utility methods for snapshotting the current IR to
+// produce new debug locations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_LOCATIONSNAPSHOT_H
+#define MLIR_TRANSFORMS_LOCATIONSNAPSHOT_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <memory>
+
+namespace mlir {
+class Location;
+struct LogicalResult;
+class Operation;
+class OpPrintingFlags;
+class Pass;
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations replace the current operation locations.
+void generateLocationsFromIR(raw_ostream &os, StringRef fileName, Operation *op,
+                             OpPrintingFlags flags);
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
+                                      OpPrintingFlags flags);
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations are represented as a NameLoc with the given tag as
+/// the name, and then fused with the existing locations.
+void generateLocationsFromIR(raw_ostream &os, StringRef fileName, StringRef tag,
+                             Operation *op, OpPrintingFlags flags);
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult generateLocationsFromIR(StringRef fileName, StringRef tag,
+                                      Operation *op, OpPrintingFlags flags);
+
+/// Create a pass to generate new locations by snapshotting the IR to the given
+/// file, and using the printed locations within that file. If `filename` is
+/// empty, a temporary file is generated instead. If a 'tag' is non-empty, the
+/// generated locations are represented as a NameLoc with the given tag as the
+/// name, and then fused with the existing locations. Otherwise, the existing
+/// locations are replaced.
+std::unique_ptr<Pass> createLocationSnapshotPass(OpPrintingFlags flags,
+                                                 StringRef fileName = "",
+                                                 StringRef tag = "");
+
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_LOCATIONSNAPSHOT_H
index fb03e38..f8abf9a 100644 (file)
@@ -158,6 +158,24 @@ bool OpPrintingFlags::shouldPrintGenericOpForm() const {
 bool OpPrintingFlags::shouldUseLocalScope() const { return printLocalScope; }
 
 //===----------------------------------------------------------------------===//
+// NewLineCounter
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class is a simple formatter that emits a new line when inputted into a
+/// stream, that enables counting the number of newlines emitted. This class
+/// should be used whenever emitting newlines in the printer.
+struct NewLineCounter {
+  unsigned curLine = 1;
+};
+} // end anonymous namespace
+
+static raw_ostream &operator<<(raw_ostream &os, NewLineCounter &newLine) {
+  ++newLine.curLine;
+  return os << '\n';
+}
+
+//===----------------------------------------------------------------------===//
 // AliasState
 //===----------------------------------------------------------------------===//
 
@@ -174,14 +192,14 @@ public:
   Twine getAttributeAlias(Attribute attr) const;
 
   /// Print all of the referenced attribute aliases.
-  void printAttributeAliases(raw_ostream &os) const;
+  void printAttributeAliases(raw_ostream &os, NewLineCounter &newLine) const;
 
   /// Return a string to use as an alias for the given type, or empty if there
   /// is no alias recorded.
   StringRef getTypeAlias(Type ty) const;
 
   /// Print all of the referenced type aliases.
-  void printTypeAliases(raw_ostream &os) const;
+  void printTypeAliases(raw_ostream &os, NewLineCounter &newLine) const;
 
 private:
   /// A special index constant used for non-kind attribute aliases.
@@ -307,12 +325,13 @@ Twine AliasState::getAttributeAlias(Attribute attr) const {
 }
 
 /// Print all of the referenced attribute aliases.
-void AliasState::printAttributeAliases(raw_ostream &os) const {
+void AliasState::printAttributeAliases(raw_ostream &os,
+                                       NewLineCounter &newLine) const {
   auto printAlias = [&](StringRef alias, Attribute attr, int index) {
     os << '#' << alias;
     if (index != NonAttrKindAlias)
       os << index;
-    os << " = " << attr << '\n';
+    os << " = " << attr << newLine;
   };
 
   // Print all of the attribute kind aliases.
@@ -320,7 +339,7 @@ void AliasState::printAttributeAliases(raw_ostream &os) const {
     auto &aliasAttrsPair = kindAlias.second;
     for (unsigned i = 0, e = aliasAttrsPair.second.size(); i != e; ++i)
       printAlias(aliasAttrsPair.first, aliasAttrsPair.second[i], i);
-    os << "\n";
+    os << newLine;
   }
 
   // In a second pass print all of the remaining attribute aliases that aren't
@@ -339,11 +358,12 @@ StringRef AliasState::getTypeAlias(Type ty) const {
 }
 
 /// Print all of the referenced type aliases.
-void AliasState::printTypeAliases(raw_ostream &os) const {
+void AliasState::printTypeAliases(raw_ostream &os,
+                                  NewLineCounter &newLine) const {
   for (Type type : usedTypes) {
     auto alias = typeToAlias.find(type);
     if (alias != typeToAlias.end())
-      os << '!' << alias->second << " = type " << type << '\n';
+      os << '!' << alias->second << " = type " << type << newLine;
   }
 }
 
@@ -765,8 +785,9 @@ namespace mlir {
 namespace detail {
 class AsmStateImpl {
 public:
-  explicit AsmStateImpl(Operation *op)
-      : interfaces(op->getContext()), nameState(op, interfaces) {}
+  explicit AsmStateImpl(Operation *op, AsmState::LocationMap *locationMap)
+      : interfaces(op->getContext()), nameState(op, interfaces),
+        locationMap(locationMap) {}
 
   /// Initialize the alias state to enable the printing of aliases.
   void initializeAliases(Operation *op) {
@@ -785,6 +806,13 @@ public:
   /// Get the state used for SSA names.
   SSANameState &getSSANameState() { return nameState; }
 
+  /// Register the location, line and column, within the buffer that the given
+  /// operation was printed at.
+  void registerOperationLocation(Operation *op, unsigned line, unsigned col) {
+    if (locationMap)
+      (*locationMap)[op] = std::make_pair(line, col);
+  }
+
 private:
   /// Collection of OpAsm interfaces implemented in the context.
   DialectInterfaceCollection<OpAsmDialectInterface> interfaces;
@@ -794,11 +822,15 @@ private:
 
   /// The state used for SSA value names.
   SSANameState nameState;
+
+  /// An optional location map to be populated.
+  AsmState::LocationMap *locationMap;
 };
 } // end namespace detail
 } // end namespace mlir
 
-AsmState::AsmState(Operation *op) : impl(std::make_unique<AsmStateImpl>(op)) {}
+AsmState::AsmState(Operation *op, LocationMap *locationMap)
+    : impl(std::make_unique<AsmStateImpl>(op, locationMap)) {}
 AsmState::~AsmState() {}
 
 //===----------------------------------------------------------------------===//
@@ -868,6 +900,9 @@ protected:
 
   /// An optional printer state for the module.
   AsmStateImpl *state;
+
+  /// A tracker for the number of new lines emitted during printing.
+  NewLineCounter newLine;
 };
 } // end anonymous namespace
 
@@ -923,10 +958,10 @@ void ModulePrinter::printLocationInternal(LocationAttr loc, bool pretty) {
         if (caller.isa<FileLineColLoc>()) {
           os << " at ";
         } else {
-          os << "\n at ";
+          os << newLine << " at ";
         }
       } else {
-        os << "\n at ";
+        os << newLine << " at ";
       }
     } else {
       os << " at ";
@@ -1921,14 +1956,17 @@ private:
 
 void OperationPrinter::print(ModuleOp op) {
   // Output the aliases at the top level.
-  state->getAliasState().printAttributeAliases(os);
-  state->getAliasState().printTypeAliases(os);
+  state->getAliasState().printAttributeAliases(os, newLine);
+  state->getAliasState().printTypeAliases(os, newLine);
 
   // Print the module.
   print(op.getOperation());
 }
 
 void OperationPrinter::print(Operation *op) {
+  // Track the location of this operation.
+  state->registerOperationLocation(op, newLine.curLine, currentIndent);
+
   os.indent(currentIndent);
   printOperation(op);
   printTrailingLocation(op->getLoc());
@@ -2066,7 +2104,7 @@ void OperationPrinter::print(Block *block, bool printBlockArgs,
         printBlockName(pred.second);
       });
     }
-    os << '\n';
+    os << newLine;
   }
 
   currentIndent += indentWidth;
@@ -2075,7 +2113,7 @@ void OperationPrinter::print(Block *block, bool printBlockArgs,
       std::prev(block->getOperations().end(), printBlockTerminator ? 0 : 1));
   for (auto &op : range) {
     print(&op);
-    os << '\n';
+    os << newLine;
   }
   currentIndent -= indentWidth;
 }
@@ -2103,7 +2141,7 @@ void OperationPrinter::printSuccessorAndUseList(Operation *term,
 
 void OperationPrinter::printRegion(Region &region, bool printEntryBlockArgs,
                                    bool printBlockTerminators) {
-  os << " {\n";
+  os << " {" << newLine;
   if (!region.empty()) {
     auto *entryBlock = &region.front();
     print(entryBlock, printEntryBlockArgs && entryBlock->getNumArguments() != 0,
index 9dda0bb..7d67100 100644 (file)
@@ -7,6 +7,7 @@ add_llvm_library(MLIRTransforms
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
+  LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
diff --git a/mlir/lib/Transforms/LocationSnapshot.cpp b/mlir/lib/Transforms/LocationSnapshot.cpp
new file mode 100644 (file)
index 0000000..3d1ce82
--- /dev/null
@@ -0,0 +1,162 @@
+//===- LocationSnapshot.cpp - Location Snapshot Utilities -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/LocationSnapshot.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/FileUtilities.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/ToolOutputFile.h"
+
+using namespace mlir;
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// If a 'tag' is non-empty, the generated locations are represented as a
+/// NameLoc with the given tag as the name, and then fused with the existing
+/// locations. Otherwise, the existing locations are replaced.
+static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
+                                    Operation *op, OpPrintingFlags flags,
+                                    StringRef tag) {
+  // Print the IR to the stream, and collect the raw line+column information.
+  AsmState::LocationMap opToLineCol;
+  AsmState state(op, &opToLineCol);
+  op->print(os, state, flags);
+
+  Builder builder(op->getContext());
+  Optional<Identifier> tagIdentifier;
+  if (!tag.empty())
+    tagIdentifier = builder.getIdentifier(tag);
+
+  // Walk and generate new locations for each of the operations.
+  Identifier file = builder.getIdentifier(fileName);
+  op->walk([&](Operation *opIt) {
+    // Check to see if this operation has a mapped location. Some operations may
+    // be elided from the printed form, e.g. the body terminators of some region
+    // operations.
+    auto it = opToLineCol.find(opIt);
+    if (it == opToLineCol.end())
+      return;
+    const std::pair<unsigned, unsigned> &lineCol = it->second;
+    auto newLoc =
+        builder.getFileLineColLoc(file, lineCol.first, lineCol.second);
+
+    // If we don't have a tag, set the location directly
+    if (!tagIdentifier) {
+      opIt->setLoc(newLoc);
+      return;
+    }
+
+    // Otherwise, build a fused location with the existing op loc.
+    opIt->setLoc(builder.getFusedLoc(
+        {opIt->getLoc(), NameLoc::get(*tagIdentifier, newLoc)}));
+  });
+}
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+static LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
+                                             OpPrintingFlags flags,
+                                             StringRef tag) {
+  // If a filename wasn't provided, then generate one.
+  SmallString<32> filepath(fileName);
+  if (filepath.empty()) {
+    if (std::error_code error = llvm::sys::fs::createTemporaryFile(
+            "mlir_snapshot", "tmp.mlir", filepath)) {
+      return op->emitError()
+             << "failed to generate temporary file for location snapshot: "
+             << error.message();
+    }
+  }
+
+  // Open the output file for emission.
+  std::string error;
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      openOutputFile(filepath, &error);
+  if (!outputFile)
+    return op->emitError() << error;
+
+  // Generate the intermediate locations.
+  generateLocationsFromIR(outputFile->os(), filepath, op, flags, tag);
+  outputFile->keep();
+  return success();
+}
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations replace the current operation locations.
+void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
+                                   Operation *op, OpPrintingFlags flags) {
+  ::generateLocationsFromIR(os, fileName, op, flags, /*tag=*/StringRef());
+}
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult mlir::generateLocationsFromIR(StringRef fileName, Operation *op,
+                                            OpPrintingFlags flags) {
+  return ::generateLocationsFromIR(fileName, op, flags, /*tag=*/StringRef());
+}
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations are represented as a NameLoc with the given tag as
+/// the name, and then fused with the existing locations.
+void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
+                                   StringRef tag, Operation *op,
+                                   OpPrintingFlags flags) {
+  ::generateLocationsFromIR(os, fileName, op, flags, tag);
+}
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag,
+                                            Operation *op,
+                                            OpPrintingFlags flags) {
+  return ::generateLocationsFromIR(fileName, op, flags, tag);
+}
+
+namespace {
+class LocationSnapshotPass : public OperationPass<LocationSnapshotPass> {
+public:
+  LocationSnapshotPass() = default;
+  LocationSnapshotPass(const LocationSnapshotPass &) {}
+  LocationSnapshotPass(OpPrintingFlags flags, StringRef fileName, StringRef tag)
+      : flags(flags) {
+    this->fileName = fileName.str();
+    this->tag = tag.str();
+  }
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    if (failed(generateLocationsFromIR(fileName, op, OpPrintingFlags(), tag)))
+      return signalPassFailure();
+  }
+
+  Option<std::string> fileName{
+      *this, "filename",
+      llvm::cl::desc("The filename to print the generated IR.")};
+  Option<std::string> tag{
+      *this, "tag",
+      llvm::cl::desc("A tag to use when fusing the new locations with the "
+                     "original. If unset, the locations are replaced.")};
+
+  /// The printing flags to use when creating the snapshot.
+  OpPrintingFlags flags;
+};
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createLocationSnapshotPass(OpPrintingFlags flags,
+                                                       StringRef fileName,
+                                                       StringRef tag) {
+  return std::make_unique<LocationSnapshotPass>(flags, fileName, tag);
+}
+
+static PassRegistration<LocationSnapshotPass>
+    reg("snapshot-op-locations", "generate new locations from the current IR");
diff --git a/mlir/test/Transforms/location-snapshot.mlir b/mlir/test/Transforms/location-snapshot.mlir
new file mode 100644 (file)
index 0000000..4153998
--- /dev/null
@@ -0,0 +1,17 @@
+// RUN: mlir-opt -snapshot-op-locations='filename=%/t' -mlir-print-debuginfo %s | FileCheck %s -DFILE=%/t
+// RUN: mlir-opt -snapshot-op-locations='filename=%/t tag='tagged'' -mlir-print-debuginfo %s | FileCheck %s --check-prefix=TAG -DFILE=%/t
+
+// CHECK-LABEL: func @function
+// CHECK-NEXT: loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
+// CHECK-NEXT: loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
+// CHECK-NEXT: } loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
+
+// TAG-LABEL: func @function
+// TAG-NEXT: loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
+// TAG-NEXT: loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
+// TAG-NEXT: } loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
+
+func @function() -> i32 {
+  %1 = "foo"() : () -> i32 loc("original")
+  return %1 : i32 loc("original")
+} loc("original")