[mlir][CF] Add `ub.unreachable` canonicalization #169873

matthias-springer · 2025-11-28T05:02:24Z

Basic blocks with only a ub.unreachable terminator are unreachable. This commit adds a canonicalization pattern that folds to cf.cond_br to cf.br if one of the destination branches is unreachable.

llvmbot · 2025-11-28T05:02:54Z

@llvm/pr-subscribers-mlir-ub
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-cf

Author: Matthias Springer (matthias-springer)

Changes

Basic blocks with a ub.unreachable terminator are unreachable. This commit adds a canonicalization pattern that drops all preceding operations. This commit also adds a canonicalization pattern that folds to cf.cond_br to cf.br if one of the destination branches is unreachable.

Depends on #169872.

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

8 Files Affected:

(modified) mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td (+1-1)
(modified) mlir/include/mlir/Dialect/UB/IR/UBOps.h (+4)
(modified) mlir/include/mlir/Dialect/UB/IR/UBOps.td (+1)
(modified) mlir/lib/Dialect/ControlFlow/IR/CMakeLists.txt (+1)
(modified) mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (+31-1)
(modified) mlir/lib/Dialect/UB/IR/UBOps.cpp (+26)
(modified) mlir/test/Dialect/ControlFlow/canonicalize.mlir (+25)
(modified) mlir/test/Dialect/UB/canonicalize.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
index a441fd82546e3..c9b4da44ffa01 100644
--- a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
+++ b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
@@ -22,7 +22,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 def ControlFlow_Dialect : Dialect {
   let name = "cf";
   let cppNamespace = "::mlir::cf";
-  let dependentDialects = ["arith::ArithDialect"];
+  let dependentDialects = ["arith::ArithDialect", "ub::UBDialect"];
   let description = [{
     This dialect contains low-level, i.e. non-region based, control flow
     constructs. These constructs generally represent control flow directly
diff --git a/mlir/include/mlir/Dialect/UB/IR/UBOps.h b/mlir/include/mlir/Dialect/UB/IR/UBOps.h
index 21de5cb0c182a..02081e2d6d15f 100644
--- a/mlir/include/mlir/Dialect/UB/IR/UBOps.h
+++ b/mlir/include/mlir/Dialect/UB/IR/UBOps.h
@@ -9,6 +9,10 @@
 #ifndef MLIR_DIALECT_UB_IR_OPS_H
 #define MLIR_DIALECT_UB_IR_OPS_H
 
+namespace mlir {
+class PatternRewriter;
+}
+
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
diff --git a/mlir/include/mlir/Dialect/UB/IR/UBOps.td b/mlir/include/mlir/Dialect/UB/IR/UBOps.td
index 8a354da2db10c..c1d74290ec174 100644
--- a/mlir/include/mlir/Dialect/UB/IR/UBOps.td
+++ b/mlir/include/mlir/Dialect/UB/IR/UBOps.td
@@ -84,6 +84,7 @@ def UnreachableOp : UB_Op<"unreachable", [Terminator]> {
   }];
 
   let assemblyFormat = "attr-dict";
+  let hasCanonicalizeMethod = 1;
 }
 
 #endif // MLIR_DIALECT_UB_IR_UBOPS_TD
diff --git a/mlir/lib/Dialect/ControlFlow/IR/CMakeLists.txt b/mlir/lib/Dialect/ControlFlow/IR/CMakeLists.txt
index 58551bb435c86..05a787fa53ec3 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/ControlFlow/IR/CMakeLists.txt
@@ -12,4 +12,5 @@ add_mlir_dialect_library(MLIRControlFlowDialect
   MLIRControlFlowInterfaces
   MLIRIR
   MLIRSideEffectInterfaces
+  MLIRUBDialect
   )
diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index f1da1a125e9ef..aabf8930cf78e 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -12,6 +12,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
@@ -445,6 +446,35 @@ struct CondBranchTruthPropagation : public OpRewritePattern<CondBranchOp> {
     return success(replaced);
   }
 };
+
+struct DropUnreachableCondBranch : public OpRewritePattern<CondBranchOp> {
+  using OpRewritePattern<CondBranchOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CondBranchOp condbr,
+                                PatternRewriter &rewriter) const override {
+    // If the "true" destination has unreachable an unreachable terminator,
+    // always branch to the "false" destination.
+    Block *trueDest = condbr.getTrueDest();
+    Block *falseDest = condbr.getFalseDest();
+    if (llvm::hasSingleElement(*trueDest) &&
+        isa<ub::UnreachableOp>(trueDest->getTerminator())) {
+      rewriter.replaceOpWithNewOp<BranchOp>(condbr, falseDest,
+                                            condbr.getFalseOperands());
+      return success();
+    }
+
+    // If the "false" destination has unreachable an unreachable terminator,
+    // always branch to the "true" destination.
+    if (llvm::hasSingleElement(*falseDest) &&
+        isa<ub::UnreachableOp>(falseDest->getTerminator())) {
+      rewriter.replaceOpWithNewOp<BranchOp>(condbr, trueDest,
+                                            condbr.getTrueOperands());
+      return success();
+    }
+
+    return failure();
+  }
+};
 } // namespace
 
 void CondBranchOp::getCanonicalizationPatterns(RewritePatternSet &results,
@@ -452,7 +482,7 @@ void CondBranchOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
               SimplifyCondBranchIdenticalSuccessors,
               SimplifyCondBranchFromCondBranchOnSameCondition,
-              CondBranchTruthPropagation>(context);
+              CondBranchTruthPropagation, DropUnreachableCondBranch>(context);
 }
 
 SuccessorOperands CondBranchOp::getSuccessorOperands(unsigned index) {
diff --git a/mlir/lib/Dialect/UB/IR/UBOps.cpp b/mlir/lib/Dialect/UB/IR/UBOps.cpp
index ee523f9522953..419e3f9d76fb2 100644
--- a/mlir/lib/Dialect/UB/IR/UBOps.cpp
+++ b/mlir/lib/Dialect/UB/IR/UBOps.cpp
@@ -12,6 +12,7 @@
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #include "mlir/Dialect/UB/IR/UBOpsDialect.cpp.inc"
@@ -57,8 +58,33 @@ Operation *UBDialect::materializeConstant(OpBuilder &builder, Attribute value,
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// PoisonOp
+//===----------------------------------------------------------------------===//
+
 OpFoldResult PoisonOp::fold(FoldAdaptor /*adaptor*/) { return getValue(); }
 
+//===----------------------------------------------------------------------===//
+// UnreachableOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult UnreachableOp::canonicalize(UnreachableOp unreachableOp,
+                                          PatternRewriter &rewriter) {
+  Block *block = unreachableOp->getBlock();
+  if (llvm::hasSingleElement(*block))
+    return rewriter.notifyMatchFailure(
+        unreachableOp, "unreachable op is the only operation in the block");
+
+  // Erase all other operations in the block. They must be dead.
+  for (Operation &op : llvm::make_early_inc_range(*block)) {
+    if (&op == unreachableOp.getOperation())
+      continue;
+    op.dropAllUses();
+    rewriter.eraseOp(&op);
+  }
+  return success();
+}
+
 #include "mlir/Dialect/UB/IR/UBOpsInterfaces.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index 17f7d28ba59fb..75dec6dacde91 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -634,3 +634,28 @@ func.func @unsimplified_cycle_2(%c : i1) {
 ^bb7:
   cf.br ^bb6
 }
+
+// CHECK-LABEL: @drop_unreachable_branch_1
+//  CHECK-NEXT:   "test.foo"() : () -> ()
+//  CHECK-NEXT:   return
+func.func @drop_unreachable_branch_1(%c: i1) {
+  cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+  "test.foo"() : () -> ()
+  return
+^bb2:
+  "test.bar"() : () -> ()
+  ub.unreachable
+}
+
+// CHECK-LABEL: @drop_unreachable_branch_2
+//  CHECK-NEXT:   ub.unreachable
+func.func @drop_unreachable_branch_2(%c: i1) {
+  cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+  "test.foo"() : () -> ()
+  ub.unreachable
+^bb2:
+  "test.bar"() : () -> ()
+  ub.unreachable
+}
diff --git a/mlir/test/Dialect/UB/canonicalize.mlir b/mlir/test/Dialect/UB/canonicalize.mlir
index c3f286e49b09b..74ba9f1932384 100644
--- a/mlir/test/Dialect/UB/canonicalize.mlir
+++ b/mlir/test/Dialect/UB/canonicalize.mlir
@@ -9,3 +9,13 @@ func.func @merge_poison() -> (i32, i32) {
   %1 = ub.poison : i32
   return %0, %1 : i32, i32
 }
+
+// -----
+
+// CHECK-LABEL: func @drop_ops_before_unreachable() 
+//  CHECK-NEXT:   ub.unreachable
+func.func @drop_ops_before_unreachable() {
+  "test.foo"() : () -> ()
+  "test.bar"() : () -> ()
+  ub.unreachable
+}

mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

joker-eph · 2025-11-28T09:32:43Z

mlir/lib/Dialect/UB/IR/UBOps.cpp

+      continue;
+    op.dropAllUses();
+    rewriter.eraseOp(&op);
+  }


We should TODO that:

this assumes we don't have calls that are "no return".

this assumes that loops terminates.

this assumes that nothing interrupts the control-flow (which is fine until early-exit is added).

Maybe we should have a trait "AlwaysForwardProgress" on operations to allow this kind of transformations?

I didn't think of infinite loops. I am wondering if this canonicalization is actually safe. I'm just merging the cf.cond_br canonicalization for now, so we can discuss further. I feel like this kind of "canonicalization" may be better suited for a separate pass. (Something like -remove-dead-values maybe?) We could also make this optimization customizable (turn off/on via pass option).

llvm-ci · 2025-11-30T03:04:35Z

LLVM Buildbot has detected a new failure on builder ppc64le-flang-rhel-clang running on ppc64le-flang-rhel-test while building mlir at step 3 "clean-build-dir".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/157/builds/42632

Here is the relevant piece of the build log for the reference

Step 3 (clean-build-dir) failure: Delete failed. (failure) (timed out)
Step 4 (cmake-configure) failure: cmake (failure) (timed out)
command timed out: 1200 seconds without output running [b'cmake', b'-DLLVM_TARGETS_TO_BUILD=PowerPC', b'-DLLVM_INSTALL_UTILS=ON', b'-DCMAKE_CXX_STANDARD=17', b'-DLLVM_LIT_ARGS=-vj 256', b'-DFLANG_ENABLE_WERROR=ON', b'-DLLVM_ENABLE_ASSERTIONS=ON', b'-DCMAKE_C_COMPILER_LAUNCHER=ccache', b'-DCMAKE_CXX_COMPILER_LAUNCHER=ccache', b'-DLLVM_ENABLE_PROJECTS=mlir;clang;llvm;flang', b'-DLLVM_ENABLE_RUNTIMES=flang-rt;openmp', b'-DCMAKE_BUILD_TYPE=Release', b'-GNinja', b'../llvm-project/llvm'], attempting to kill
process killed by signal 9
program finished with exit code -1
elapsedTime=1200.786072

A new dependency was added.

Basic blocks with only a `ub.unreachable` terminator are unreachable. This commit adds a canonicalization pattern that folds to `cf.cond_br` to `cf.br` if one of the destination branches is unreachable.

A new dependency was added.

Basic blocks with only a `ub.unreachable` terminator are unreachable. This commit adds a canonicalization pattern that folds to `cf.cond_br` to `cf.br` if one of the destination branches is unreachable.

A new dependency was added.

Basic blocks with only a `ub.unreachable` terminator are unreachable. This commit adds a canonicalization pattern that folds to `cf.cond_br` to `cf.br` if one of the destination branches is unreachable.

A new dependency was added.

Basic blocks with only a `ub.unreachable` terminator are unreachable. This commit adds a canonicalization pattern that folds to `cf.cond_br` to `cf.br` if one of the destination branches is unreachable.

A new dependency was added.

matthias-springer requested review from joker-eph and kuhar November 28, 2025 05:02

llvmbot added mlir mlir:cf mlir:ub labels Nov 28, 2025

matthias-springer force-pushed the users/matthias-springer/unreachable_canonicalization branch from bc96208 to 561b6ca Compare November 28, 2025 05:04

joker-eph approved these changes Nov 28, 2025

View reviewed changes

mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td Outdated Show resolved Hide resolved

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp Show resolved Hide resolved

joker-eph reviewed Nov 28, 2025

View reviewed changes

Base automatically changed from users/matthias-springer/ub_unreachable to main November 28, 2025 10:35

matthias-springer requested a review from antiagainst as a code owner November 28, 2025 10:35

[mlir][UB] Add ub.unreachable canonicalization

04a3be8

matthias-springer force-pushed the users/matthias-springer/unreachable_canonicalization branch from 561b6ca to 92caacb Compare November 30, 2025 02:18

remove UB canonicalization

143a40b

matthias-springer force-pushed the users/matthias-springer/unreachable_canonicalization branch from 92caacb to 143a40b Compare November 30, 2025 02:19

matthias-springer changed the title ~~[mlir][UB] Add ub.unreachable canonicalization~~ [mlir][CF] Add ub.unreachable canonicalization Nov 30, 2025

matthias-springer enabled auto-merge (squash) November 30, 2025 02:21

matthias-springer merged commit 3de11e9 into main Nov 30, 2025
10 checks passed

matthias-springer deleted the users/matthias-springer/unreachable_canonicalization branch November 30, 2025 02:36

boomanaiden154 added a commit to boomanaiden154/llvm-project that referenced this pull request Nov 30, 2025

[bazel] Port llvm#169873

f2ced05

A new dependency was added.

boomanaiden154 added a commit that referenced this pull request Nov 30, 2025

[bazel] Port #169873 (#170027)

22257e8

A new dependency was added.

aahrun pushed a commit to aahrun/llvm-project that referenced this pull request Dec 1, 2025

[bazel] Port llvm#169873 (llvm#170027)

01b465f

A new dependency was added.

augusto2112 pushed a commit to augusto2112/llvm-project that referenced this pull request Dec 3, 2025

[bazel] Port llvm#169873 (llvm#170027)

1674c55

A new dependency was added.

kcloudy0717 pushed a commit to kcloudy0717/llvm-project that referenced this pull request Dec 4, 2025

[bazel] Port llvm#169873 (llvm#170027)

5a07ebd

A new dependency was added.

honeygoyal pushed a commit to honeygoyal/llvm-project that referenced this pull request Dec 9, 2025

[bazel] Port llvm#169873 (llvm#170027)

c4045a7

A new dependency was added.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][CF] Add `ub.unreachable` canonicalization #169873

[mlir][CF] Add `ub.unreachable` canonicalization #169873

Uh oh!

matthias-springer commented Nov 28, 2025 •

edited

Loading

Uh oh!

llvmbot commented Nov 28, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

joker-eph Nov 28, 2025

Uh oh!

joker-eph Nov 28, 2025

Uh oh!

matthias-springer Nov 30, 2025 •

edited

Loading

Uh oh!

Uh oh!

llvm-ci commented Nov 30, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[mlir][CF] Add ub.unreachable canonicalization #169873

[mlir][CF] Add ub.unreachable canonicalization #169873

Uh oh!

Conversation

matthias-springer commented Nov 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Nov 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

joker-eph Nov 28, 2025

Choose a reason for hiding this comment

Uh oh!

joker-eph Nov 28, 2025

Choose a reason for hiding this comment

Uh oh!

matthias-springer Nov 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

llvm-ci commented Nov 30, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[mlir][CF] Add `ub.unreachable` canonicalization #169873

[mlir][CF] Add `ub.unreachable` canonicalization #169873

matthias-springer commented Nov 28, 2025 •

edited

Loading

llvmbot commented Nov 28, 2025 •

edited

Loading

matthias-springer Nov 30, 2025 •

edited

Loading