[mlir] Fix debug output for passes that modify top-level operation. (#80022)

Make it so that when the top-level (root) operation itself is being
modified, it is also used as the root for debug output in
PatternApplicator.

Fix #80021

Author: Jezurko

mlir/lib/Rewrite/PatternApplicator.cpp

@@ -40,7 +40,11 @@ static void logImpossibleToMatch(const Pattern &pattern) {
 
 /// Log IR after pattern application.
 static Operation *getDumpRootOp(Operation *op) {
-  return op->getParentWithTrait<mlir::OpTrait::IsIsolatedFromAbove>();
+  Operation *isolatedParent =
+      op->getParentWithTrait<mlir::OpTrait::IsIsolatedFromAbove>();
+  if (isolatedParent)
+    return isolatedParent;
+  return op;
 }
 static void logSucessfulPatternApplication(Operation *op) {
   llvm::dbgs() << "// *** IR Dump After Pattern Application ***\n";