Update the IRPrinter instrumentation to work on non function/module operations.

This is necessary now that the pass manager may work on different types of operations.

PiperOrigin-RevId: 269139669
This commit is contained in:
River Riddle 2019-09-14 21:56:09 -07:00 committed by TensorFlower Gardener
parent ab17c2f07c
commit 2482e743bc
2 changed files with 25 additions and 20 deletions

View File

@ -64,27 +64,33 @@ static bool isHiddenPass(Pass *pass) {
}
static void printIR(Operation *op, bool printModuleScope, raw_ostream &out) {
// Check for printing at module scope.
auto function = dyn_cast<FuncOp>(op);
if (printModuleScope && function) {
// Print the function name and a newline before the Module.
out << " (function: " << function.getName() << ")\n";
function.getParentOfType<ModuleOp>().print(out);
return;
}
// Check to see if we are printing the top-level module.
auto module = dyn_cast<ModuleOp>(op);
if (module && !op->getBlock())
return module.print(out << "\n");
// Print a newline before the IR.
out << "\n";
// Otherwise, check to see if we are not printing at module scope.
if (!printModuleScope)
return op->print(out << "\n");
// Print the given function.
if (function) {
function.print(out);
return;
}
// Otherwise, we are printing at module scope.
out << " ('" << op->getName() << "' operation";
if (auto symbolName =
op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName()))
out << ": @" << symbolName.getValue();
out << ")\n";
// Print the given module.
assert(isa<ModuleOp>(op) && "unexpected IR unit");
cast<ModuleOp>(op).print(out);
// Find the top-level module operation.
auto *topLevelOp = op;
while (auto *parentOp = topLevelOp->getParentOp())
topLevelOp = parentOp;
// Check to see if the top-level operation is actually a module in the case of
// invalid-ir.
if (auto module = dyn_cast<ModuleOp>(topLevelOp))
module.print(out);
else
topLevelOp->print(out);
}
/// Instrumentation hooks.

View File

@ -80,8 +80,7 @@ PassManagerOptions::PassManagerOptions()
printModuleScope(
"print-ir-module-scope",
llvm::cl::desc("When printing IR for print-ir-[before|after]{-all} "
"always print "
"a module IR"),
"always print the top-level module operation"),
llvm::cl::init(false)),
//===----------------------------------------------------------------===//