[MLIR][Standard] Erase redundant assertions std.assert

Differential Revision: https://reviews.llvm.org/D83118

Author: frgossen

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

@@ -467,6 +467,8 @@ def AssertOp : Std_Op<"assert"> {
 
   // AssertOp is fully verified by its traits.
   let verifier = ?;
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

@@ -439,6 +439,31 @@ OpFoldResult AndOp::fold(ArrayRef<Attribute> operands) {
                                         [](APInt a, APInt b) { return a & b; });
 }
 
+//===----------------------------------------------------------------------===//
+// AssertOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct EraseRedundantAssertions : public OpRewritePattern<AssertOp> {
+  using OpRewritePattern<AssertOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AssertOp op,
+                                PatternRewriter &rewriter) const override {
+    // Erase assertion if argument is constant true.
+    if (matchPattern(op.arg(), m_One())) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+    return failure();
+  }
+};
+} // namespace
+
+void AssertOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                           MLIRContext *context) {
+  patterns.insert<EraseRedundantAssertions>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AssumeAlignmentOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Standard/canonicalize-cf.mlir

@@ -138,3 +138,26 @@ func @cond_br_pass_through_fail(%cond : i1) {
 ^bb2:
   return
 }
+
+// -----
+
+// Erase assertion if condition is known to be true at compile time.
+// CHECK-LABEL: @assert_true
+func @assert_true() {
+  // CHECK-NOT: assert
+  %true = constant true
+  assert %true, "Computer says no"
+  return
+}
+
+// -----
+
+// Keep assertion if condition unknown at compile time.
+// CHECK-LABEL: @assert
+// CHECK-SAME:  (%[[ARG:.*]]: i1)
+func @assert(%arg : i1) {
+  // CHECK: assert %[[ARG]], "Computer says no"
+  assert %arg, "Computer says no"
+  return
+}
+