#ifndef MLIR_IR_ASMSTATE_H_
#define MLIR_IR_ASMSTATE_H_
+#include "mlir/Support/LLVM.h"
+
#include <memory>
namespace mlir {
/// 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
--- /dev/null
+//===- 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
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
//===----------------------------------------------------------------------===//
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.
}
/// 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.
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
}
/// 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;
}
}
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) {
/// 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;
/// 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() {}
//===----------------------------------------------------------------------===//
/// 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
if (caller.isa<FileLineColLoc>()) {
os << " at ";
} else {
- os << "\n at ";
+ os << newLine << " at ";
}
} else {
- os << "\n at ";
+ os << newLine << " at ";
}
} else {
os << " at ";
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());
printBlockName(pred.second);
});
}
- os << '\n';
+ os << newLine;
}
currentIndent += indentWidth;
std::prev(block->getOperations().end(), printBlockTerminator ? 0 : 1));
for (auto &op : range) {
print(&op);
- os << '\n';
+ os << newLine;
}
currentIndent -= indentWidth;
}
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,
CSE.cpp
DialectConversion.cpp
Inliner.cpp
+ LocationSnapshot.cpp
LoopCoalescing.cpp
LoopFusion.cpp
LoopInvariantCodeMotion.cpp
--- /dev/null
+//===- 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");
--- /dev/null
+// 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")