diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp index 607b86cb8631..0a2a0cc1d5c7 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( [&] { - 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 78f67821da13..ca3de3a2d770 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 +} + diff --git a/mlir/test/Transforms/test-canonicalize.mlir b/mlir/test/Transforms/test-canonicalize.mlir index 0fc822b0a23a..8cad6b98441d 100644 --- a/mlir/test/Transforms/test-canonicalize.mlir +++ b/mlir/test/Transforms/test-canonicalize.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize))' | FileCheck %s +// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize))' | FileCheck %s --check-prefixes=CHECK,RS // RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=disabled}))' | FileCheck %s --check-prefixes=CHECK,NO-RS // CHECK-LABEL: func @remove_op_with_inner_ops_pattern @@ -80,12 +80,10 @@ func.func @test_dialect_canonicalizer() -> (i32) { // Check that the option to control region simplification actually works // CHECK-LABEL: test_region_simplify -func.func @test_region_simplify() { - // CHECK-NEXT: return - // NO-RS-NEXT: ^bb1 - // NO-RS-NEXT: return - // CHECK-NEXT: } - return -^bb1: - return +func.func @test_region_simplify(%input1 : i32, %cond : i1) -> i32 { + // RS-NEXT: "test.br"(%arg0)[^bb1] : (i32) -> () + // NO-RS-NEXT: "test.br"(%arg0, %arg0)[^bb1] : (i32, i32) -> () + "test.br"(%input1, %input1)[^bb1] : (i32, i32) -> () +^bb1(%used_arg : i32, %unused_arg : i32): + return %used_arg : i32 }