diff --git a/mlir/include/mlir/IR/AsmState.h b/mlir/include/mlir/IR/AsmState.h --- 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 namespace mlir { @@ -32,8 +34,14 @@ /// 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>; + + /// 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 --- /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 + +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 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 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -158,6 +158,24 @@ 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 @@ 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 @@ } /// 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 @@ 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 @@ } /// 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 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 @@ /// 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 interfaces; @@ -794,11 +822,15 @@ /// 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(op)) {} +AsmState::AsmState(Operation *op, LocationMap *locationMap) + : impl(std::make_unique(op, locationMap)) {} AsmState::~AsmState() {} //===----------------------------------------------------------------------===// @@ -868,6 +900,9 @@ /// 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 @@ if (caller.isa()) { os << " at "; } else { - os << "\n at "; + os << newLine << " at "; } } else { - os << "\n at "; + os << newLine << " at "; } } else { os << " at "; @@ -1921,14 +1956,17 @@ 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 @@ printBlockName(pred.second); }); } - os << '\n'; + os << newLine; } currentIndent += indentWidth; @@ -2075,7 +2113,7 @@ 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::printRegion(Region ®ion, bool printEntryBlockArgs, bool printBlockTerminators) { - os << " {\n"; + os << " {" << newLine; if (!region.empty()) { auto *entryBlock = ®ion.front(); print(entryBlock, printEntryBlockArgs && entryBlock->getNumArguments() != 0, diff --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt --- a/mlir/lib/Transforms/CMakeLists.txt +++ b/mlir/lib/Transforms/CMakeLists.txt @@ -7,6 +7,7 @@ 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 --- /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 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 &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 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 { +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 fileName{ + *this, "filename", + llvm::cl::desc("The filename to print the generated IR.")}; + Option 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 mlir::createLocationSnapshotPass(OpPrintingFlags flags, + StringRef fileName, + StringRef tag) { + return std::make_unique(flags, fileName, tag); +} + +static PassRegistration + 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 --- /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")