[Mlir-commits] [mlir] [MLIR] Erase unreachable blocks before applying patterns in the greedy rewriter (PR #153957)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

Operations like:

    %add = arith.addi %add, %add : i64

are legal in unreachable code. Unfortunately many patterns would be unsafe to apply on such IR and can lead to crashes or infinite loops. To avoid this we can remove unreachable blocks before attempting to apply patterns.
We may have to do this also whenever the CFG is changed by a pattern, it is left up for future work right now.

Fixes #<!-- -->153732

---
Full diff: https://github.com/llvm/llvm-project/pull/153957.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp (+12-1) 
- (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+15) 


``````````diff
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 607b86cb86315..0a2a0cc1d5c73 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -871,7 +871,18 @@ LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
 
     ctx->executeAction<GreedyPatternRewriteIteration>(
         [&] {
-          continueRewrites = processWorklist();
+          continueRewrites = false;
+
+          // Erase unreachable blocks
+          // Operations like:
+          //   %add = arith.addi %add, %add : i64
+          // are legal in unreachable code. Unfortunately many patterns would be
+          // unsafe to apply on such IR and can lead to crashes or infinite
+          // loops.
+          continueRewrites |=
+              succeeded(eraseUnreachableBlocks(rewriter, region));
+
+          continueRewrites |= processWorklist();
 
           // After applying patterns, make sure that the CFG of each of the
           // regions is kept up to date.
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 78f67821da138..c9498b147246a 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -3363,3 +3363,18 @@ func.func @bf16_fma(%arg0: vector<32x32x32xbf16>, %arg1: vector<32x32x32xbf16>,
     }
   }
 #-}
+
+// CHECK-LABEL: func @unreachable()
+// CHECK-NEXT: return
+// CHECK-NOT: arith
+  func.func @unreachable() {
+    return
+  ^unreachable:
+    %c1_i64 = arith.constant 1 : i64
+    // This self referencing operation is legal in an unreachable block.
+    // Many patterns are unsafe with respect to this kind of situation,
+    // check that we don't infinite loop here.
+    %add = arith.addi %add, %c1_i64 : i64
+    cf.br ^unreachable
+  }
+}
\ No newline at end of file

``````````

</details>


https://github.com/llvm/llvm-project/pull/153957