[Mlir-commits] [mlir] abe3e5b - [mlir] Add support for generating debug locations from intermediate levels of the IR.

River Riddle llvmlistbot at llvm.org
Sat Feb 8 15:16:21 PST 2020


Author: River Riddle
Date: 2020-02-08T15:11:29-08:00
New Revision: abe3e5babdad001100b285d188e63f76fe21aea4

URL: https://github.com/llvm/llvm-project/commit/abe3e5babdad001100b285d188e63f76fe21aea4
DIFF: https://github.com/llvm/llvm-project/commit/abe3e5babdad001100b285d188e63f76fe21aea4.diff

LOG: [mlir] Add support for generating debug locations from intermediate levels of the IR.

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

Added: 
    mlir/include/mlir/Transforms/LocationSnapshot.h
    mlir/lib/Transforms/LocationSnapshot.cpp
    mlir/test/Transforms/location-snapshot.mlir

Modified: 
    mlir/include/mlir/IR/AsmState.h
    mlir/lib/IR/AsmPrinter.cpp
    mlir/lib/Transforms/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AsmState.h b/mlir/include/mlir/IR/AsmState.h
index c777c9933bb5..a56f8e63ee29 100644
--- a/mlir/include/mlir/IR/AsmState.h
+++ b/mlir/include/mlir/IR/AsmState.h
@@ -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
index 000000000000..e4c911ae07d6
--- /dev/null
+++ b/mlir/include/mlir/Transforms/LocationSnapshot.h
@@ -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

diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index fb03e3844f80..f8abf9a3e95b 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -157,6 +157,24 @@ bool OpPrintingFlags::shouldPrintGenericOpForm() const {
 /// Return if the printer should use local scope when dumping the IR.
 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 @@ class AliasState {
   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 @@ class AsmStateImpl {
   /// 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 @@ class AsmStateImpl {
 
   /// 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 @@ class ModulePrinter {
 
   /// 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 @@ class OperationPrinter : public ModulePrinter, private OpAsmPrinter {
 
 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,

diff  --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt
index 9dda0bbf467b..7d671003e7f0 100644
--- a/mlir/lib/Transforms/CMakeLists.txt
+++ b/mlir/lib/Transforms/CMakeLists.txt
@@ -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
index 000000000000..3d1ce828f1f2
--- /dev/null
+++ b/mlir/lib/Transforms/LocationSnapshot.cpp
@@ -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
index 000000000000..4153998a5a90
--- /dev/null
+++ b/mlir/test/Transforms/location-snapshot.mlir
@@ -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")


        


More information about the Mlir-commits mailing list