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

third_party/mlir/lib/Pass/IRPrinting.cpp

@@ -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.
+  // Check to see if we are printing the top-level module.
-  auto function = dyn_cast<FuncOp>(op);
+  auto module = dyn_cast<ModuleOp>(op);
-  if (printModuleScope && function) {
+  if (module && !op->getBlock())
-    // Print the function name and a newline before the Module.
+    return module.print(out << "\n");
-    out << " (function: " << function.getName() << ")\n";
-    function.getParentOfType<ModuleOp>().print(out);
-    return;
-  }
 
-  // Print a newline before the IR.
+  // Otherwise, check to see if we are not printing at module scope.
-  out << "\n";
+  if (!printModuleScope)
+    return op->print(out << "\n");
 
-  // Print the given function.
+  // Otherwise, we are printing at module scope.
-  if (function) {
+  out << " ('" << op->getName() << "' operation";
-    function.print(out);
+  if (auto symbolName =
-    return;
+          op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName()))
-  }
+    out << ": @" << symbolName.getValue();
+  out << ")\n";
 
-  // Print the given module.
+  // Find the top-level module operation.
-  assert(isa<ModuleOp>(op) && "unexpected IR unit");
+  auto *topLevelOp = op;
-  cast<ModuleOp>(op).print(out);
+  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.

third_party/mlir/lib/Pass/PassManagerOptions.cpp

@@ -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 "
+                         "always print the top-level module operation"),
-                         "a module IR"),
           llvm::cl::init(false)),
 
       //===----------------------------------------------------------------===//