[Mlir-commits] [mlir] c50feca - [mlir] Fix region simplification bug when later blocks use prior block argument values (#97960)

[llvmlistbot at llvm.org]

Author: Ben Howe
Date: 2024-09-04T21:37:14+02:00
New Revision: c50fecaaaabcf1598dc25fbde24c8352745b4ac9

URL: https://github.com/llvm/llvm-project/commit/c50fecaaaabcf1598dc25fbde24c8352745b4ac9
DIFF: https://github.com/llvm/llvm-project/commit/c50fecaaaabcf1598dc25fbde24c8352745b4ac9.diff

LOG: [mlir] Fix region simplification bug when later blocks use prior block argument values (#97960)

This fixes #94520 by ensuring that any if any block arguments are being
used outside of the original block that the block is not considered a
candidate for merging.

More details: the root cause of the issue described in #94520 was that
`^bb2` and `^bb5` were being merged despite `%4` (an argument to `^bb2`)
was being used later in `^bb7`. When the block merge occurred, that
unintentionally changed the value of `%4` for all downstream code. This
change prevents that from happening.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/RegionUtils.cpp
    mlir/test/Transforms/canonicalize-block-merge.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index 3c7523827699c8..b7d26e88230946 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/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)

diff  --git a/mlir/test/Transforms/canonicalize-block-merge.mlir b/mlir/test/Transforms/canonicalize-block-merge.mlir
index 92cfde817cf7ff..6dbfcb562e588b 100644
--- a/mlir/test/Transforms/canonicalize-block-merge.mlir
+++ b/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
+}