diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -359,6 +359,7 @@ } }]; + let hasCanonicalizer = 1; let hasFolder = 1; } diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h --- a/mlir/include/mlir/IR/PatternMatch.h +++ b/mlir/include/mlir/IR/PatternMatch.h @@ -280,6 +280,10 @@ /// This method erases an operation that is known to have no uses. virtual void eraseOp(Operation *op); + /// This method erases all operations in a block; all these operations are + /// known to have no uses. + virtual void eraseBlock(Block *block); + /// Merge the operations of block 'source' into the end of block 'dest'. /// 'source's predecessors must either be empty or only contain 'dest`. /// 'argValues' is used to replace the block arguments of 'source' after diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -79,8 +79,8 @@ /// Materialize a single constant operation from a given attribute value with /// the desired resultant type. Operation *AffineDialect::materializeConstant(OpBuilder &builder, - Attribute value, Type type, - Location loc) { + Attribute value, Type type, + Location loc) { return builder.create(loc, type, value); } @@ -1566,6 +1566,25 @@ // AffineIfOp //===----------------------------------------------------------------------===// +namespace { +/// Remove else blocks that have nothing other than the terminator. +struct SimplifyDeadElse : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(AffineIfOp ifOp, + PatternRewriter &rewriter) const override { + if (ifOp.elseRegion().empty() || + ifOp.getElseBlock()->getOperations().size() > 1) + return failure(); + + rewriter.startRootUpdate(ifOp); + rewriter.eraseBlock(ifOp.getElseBlock()); + rewriter.finalizeRootUpdate(ifOp); + return success(); + } +}; +} // end anonymous namespace. + static LogicalResult verify(AffineIfOp op) { // Verify that we have a condition attribute. auto conditionAttr = @@ -1710,6 +1729,11 @@ return failure(); } +void AffineIfOp::getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context) { + results.insert(context); +} + //===----------------------------------------------------------------------===// // AffineLoadOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp --- a/mlir/lib/IR/PatternMatch.cpp +++ b/mlir/lib/IR/PatternMatch.cpp @@ -89,6 +89,15 @@ op->erase(); } +void PatternRewriter::eraseBlock(Block *block) { + for (auto &op : llvm::make_early_inc_range(*block)) { + assert(op.use_empty() && "expected 'op' to have no uses"); + notifyOperationRemoved(&op); + op.erase(); + } + block->erase(); +} + /// Merge the operations of block 'source' into the end of block 'dest'. /// 'source's predecessors must be empty or only contain 'dest`. /// 'argValues' is used to replace the block arguments of 'source' after diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir --- a/mlir/test/Transforms/canonicalize.mlir +++ b/mlir/test/Transforms/canonicalize.mlir @@ -895,3 +895,25 @@ // CHECK: return %[[C4_I16]], %[[C4]] : i16, index return %1, %2 : i16, index } + +// CHECK-LABEL: func @remove_dead_else +func @remove_dead_else(%M : memref<100 x i32>) { + affine.for %i = 0 to 100 { + affine.load %M[%i] : memref<100xi32> + affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) { + affine.for %j = 0 to 100 { + affine.load %M[%j] : memref<100xi32> + } + } else { + // Nothing + } + affine.load %M[%i] : memref<100xi32> + } + return +} +// CHECK: affine.if +// CHECK-NEXT: affine.for +// CHECK-NEXT: affine.load +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: affine.load