[mlir] Fix region simplification bug when later blocks use prior block argument values

mlir/lib/Transforms/Utils/RegionUtils.cpp

@@ -877,6 +877,15 @@ static LogicalResult mergeIdenticalBlocks(RewriterBase &rewriter,
       if (hasNonEmptyRegion)
         continue;
 
+      // Don't allow merging if this block's arguments are used outside of the
+      // original block.
+      bool argHasExternalUsers = llvm::any_of(
+          block->getArguments(), [block](mlir::BlockArgument &arg) {
+            return arg.isUsedOutsideOfBlock(block);
+          });
+      if (argHasExternalUsers)
+        continue;
+
       // Try to add this block to an existing cluster.
       bool addedToCluster = false;
       for (auto &cluster : clusters)

mlir/test/Transforms/canonicalize-block-merge.mlir

@@ -290,3 +290,31 @@ func.func @dead_dealloc_fold_multi_use(%cond : i1) {
   memref.dealloc %a: memref<4xf32>
   return
 }
+
+// CHECK-LABEL: func @nested_loop
+func.func @nested_loop(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i1) {
+// Irreducible control-flow: enter the middle of the loop in LoopBody_entry here.
+  "test.foo_br"(%arg0, %arg4)[^LoopBody_entry] : (i32, i32) -> ()
+
+// Loop exit condition: jump to exit or LoobBody blocks
+^Loop_header:  // 2 preds: ^bb2, ^bb3
+  // Consumes the block arg from LoopBody_entry
+  // Because of this use here, we can't merge the two blocks below.
+  "test.foo_br2"(%0)[^EXIT, ^LoopBody_entry, ^LoopBody_other] : (i32) -> ()
+
+// LoopBody_entry is jumped in from the entry block (bb0) and Loop_header
+// It **dominates** the Loop_header.
+^LoopBody_entry(%0: i32):  // 2 preds: ^bb0, ^Loop_header
+ // CHECK: test.bar
+  %1 = "test.bar"(%0) : (i32) -> i32
+  cf.br ^Loop_header
+
+// Other block inside the loop, not dominating the header
+^LoopBody_other(%2: i32):  // pred: ^Loop_header
+ // CHECK: test.bar
+  %3 = "test.bar"(%2) : (i32) -> i32
+  cf.br ^Loop_header
+
+^EXIT:  // pred: ^Loop_header
+  return
+}