[mlir][UB] Add ub.unreachable canonicalization

Author: matthias-springer

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

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"

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

mlir/lib/Dialect/ControlFlow/IR/CMakeLists.txt

@@ -12,4 +12,5 @@ add_mlir_dialect_library(MLIRControlFlowDialect
   MLIRControlFlowInterfaces
   MLIRIR
   MLIRSideEffectInterfaces
+  MLIRUBDialect
   )

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,14 +446,43 @@ 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,
                                                MLIRContext *context) {
   results.add<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
               SimplifyCondBranchIdenticalSuccessors,
               SimplifyCondBranchFromCondBranchOnSameCondition,
-              CondBranchTruthPropagation>(context);
+              CondBranchTruthPropagation, DropUnreachableCondBranch>(context);
 }
 
 SuccessorOperands CondBranchOp::getSuccessorOperands(unsigned index) {

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

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
+}

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
+}