[llvm-branch-commits] [mlir] [mlir][UB] Erase ops that precede `ub.unreachable` (PR #179104)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

<details>
<summary>Changes</summary>

Erase ops that precede `ub.unreachable` if they must progress, i.e., it's guaranteed that their next operation is reached.

Depends on #<!-- -->179039.


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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/UB/IR/UBOps.td (+1) 
- (modified) mlir/lib/Dialect/UB/IR/UBOps.cpp (+37) 
- (modified) mlir/test/Dialect/UB/canonicalize.mlir (+25) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/UB/IR/UBOps.td b/mlir/include/mlir/Dialect/UB/IR/UBOps.td
index 1bff39add691e..95c52db404912 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/UB/IR/UBOps.cpp b/mlir/lib/Dialect/UB/IR/UBOps.cpp
index 9c8ad2ddb2c6a..66dfceaffc6dd 100644
--- a/mlir/lib/Dialect/UB/IR/UBOps.cpp
+++ b/mlir/lib/Dialect/UB/IR/UBOps.cpp
@@ -59,8 +59,45 @@ 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) {
+  // Erase operations that precede ub.unreachable (if they must progress).
+  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.
+  auto it = unreachableOp->getIterator();
+  --it;
+  bool reachedBegin = false;
+  bool changed = false;
+  do {
+    Operation *op = &*it;
+    // Do not erase ops that may not progress (and ops that precede them).
+    if (!mustProgress(op))
+      break;
+    reachedBegin = it == block->begin();
+    if (!reachedBegin)
+      --it;
+    rewriter.eraseOp(op);
+    changed = true;
+    if (reachedBegin)
+      break;
+  } while (!reachedBegin);
+  return success(changed);
+}
+
 #include "mlir/Dialect/UB/IR/UBOpsInterfaces.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
diff --git a/mlir/test/Dialect/UB/canonicalize.mlir b/mlir/test/Dialect/UB/canonicalize.mlir
index c3f286e49b09b..bd8cc99449b7e 100644
--- a/mlir/test/Dialect/UB/canonicalize.mlir
+++ b/mlir/test/Dialect/UB/canonicalize.mlir
@@ -9,3 +9,28 @@ func.func @merge_poison() -> (i32, i32) {
   %1 = ub.poison : i32
   return %0, %1 : i32, i32
 }
+
+// -----
+
+// CHECK-LABEL: func @drop_ops_before_unreachable( 
+//  CHECK-NEXT:   arith.constant
+//  CHECK-NEXT:   arith.constant
+//  CHECK-NEXT:   arith.constant
+//  CHECK-NEXT:   vector.print
+//  CHECK-NEXT:   scf.for {{.*}} {
+//  CHECK-NEXT:     vector.print
+//  CHECK-NEXT:   }
+//  CHECK-NEXT:   ub.unreachable
+func.func @drop_ops_before_unreachable(%arg0: i32) {
+  %lb = arith.constant 3 : index
+  %ub = arith.constant 4 : index
+  %step = arith.constant 0 : index
+  vector.print %arg0 : i32
+  // Infinite loop that may not progress. Such ops (and everything before)
+  // is not erased.
+  scf.for %iv = %lb to %ub step %step {
+    vector.print %arg0 : i32
+  } {mustProgress = false}
+  vector.print %arg0 : i32
+  ub.unreachable
+}

``````````

</details>


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