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

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/153957

>From 177c693fd2c4b7c5b9f446350598d3e2b2d18c54 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 16 Aug 2025 07:36:22 -0700
Subject: [PATCH] [MLIR] Erase unreachable blocks before applying patterns in
 the greedy rewriter

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
---
 .../Utils/GreedyPatternRewriteDriver.cpp          | 13 ++++++++++++-
 mlir/test/Dialect/Arith/canonicalize.mlir         | 15 +++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

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..4c7e742dd8296 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
+  }
+}