[MLIR] Add pattern rewriter util to erase block; add affine.if pattern to remove dead else blocks

[bondhugula]

Summary

Add a pattern rewriter utility to erase blocks (while notifying the
pattern rewriting driver of the erased ops). Use this to remove trivial
else blocks in affine.if ops.

Signed-off-by: Uday Bondhugula <[e-mail]>

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

285

Can you add proper handling for DialectConversion?

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1580

Please use has_single_element instead. size() is O(N).

1583

nit:

rewriter.updateRootInPlace(ifOp, [&] {
  rewriter.eraseBlock(if.getElseBlock());
});

?

mlir/lib/IR/PatternMatch.cpp

94

This seems wrong, any non-trivially empty block will likely fail this check. Did you want to iterate the block in reverse?

96

Why not just replace this loop with 'eraseOp'?

[bondhugula]

mlir/include/mlir/IR/PatternMatch.h

285

Sorry, I didn't understand. What does dialect conversion need from eraseBlock? and which it doesn't from eraseOp?

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1580

Thanks.

1583

Thanks, but I tend to prefer the way it is - I find it much simpler to read given the single line action.

mlir/lib/IR/PatternMatch.cpp

94

Sorry, I didn't understand - what's a non-trivially empty block? AFAICT, the block is either empty of ops or it isn't.

96

Sure.

[bondhugula]

Address review comments.

[bondhugula]

mlir/lib/IR/PatternMatch.cpp

94

It looks like I was looking at the wrong line / code context - fixed, thanks.

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

285

DialectConversion doesn't allow performing the IR mutation until after the conversion is successful. In this case, you can't erase the block directly because this action may need to be undone later.

mlir/lib/IR/PatternMatch.cpp

94

If you have something like:

^bb:
%cst = constant 0 : i64
foo.yield %cst : i64

This will assert when trying to erase the block. Why should we assert that all ops have no uses?

[bondhugula]

fix eraseBlock iteration order

[bondhugula]

mlir/include/mlir/IR/PatternMatch.h

285

Sure - but could you provide a pointer on what handling you need here? Is it trivial? It looks like one shouldn't use eraseBlock in such cases, but rely on a lower level approach?

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

285

I'm fine with adding a llvm_unreachable in the ConversionPatternRewriter for now.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1580

I think you could just do: *ifOp.getElseBlock()

[bondhugula]

Address review comments and rebase

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1580

Thanks.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let extraClassDeclaration = [{
     }
     OpBuilder getElseBodyBuilder() {
       assert(!elseRegion().empty() && "Unexpected empty 'else' region.");
       Block &body = elseRegion().front();
       return OpBuilder(&body, std::prev(body.end()));
     }
   }];
 
+  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineMinMaxOpBase<string mnemonic, list<OpTrait> traits = []> :
     Op<Affine_Dialect, mnemonic, traits> {
   let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
   let results = (outs Index);
 

mlir/include/mlir/IR/PatternMatch.h

public:
   void replaceOpWithNewOp(Operation *op, Args &&... args) {
     auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
     replaceOpWithResultsOfAnotherOp(op, newOp.getOperation());
   }
 
   /// 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.
+  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
   /// merging.
   virtual void mergeBlocks(Block *source, Block *dest,
                            ValueRange argValues = llvm::None);
 
   /// Split the operations starting at "before" (inclusive) out of the given

mlir/include/mlir/Transforms/DialectConversion.h

public:
   void replaceOp(Operation *op, ValueRange newValues) override;
   using PatternRewriter::replaceOp;
 
   /// PatternRewriter hook for erasing a dead operation. The uses of this
   /// operation *must* be made dead by the end of the conversion process,
   /// otherwise an assert will be issued.
   void eraseOp(Operation *op) override;
 
+  /// PatternRewriter hook for erase all operations in a block. This is not yet
+  /// implemented for dialect conversion.
+  void eraseBlock(Block *block) override;
+
   /// PatternRewriter hook for creating a new block with the given arguments.
   Block *createBlock(Region *parent, Region::iterator insertPt = {},
                      TypeRange argTypes = llvm::None) override;
 
   /// PatternRewriter hook for splitting a block into two parts.
   Block *splitBlock(Block *block, Block::iterator before) override;
 
   /// PatternRewriter hook for merging a block into another.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
                 >();
   addInterfaces<AffineInlinerInterface>();
 }
 
 /// 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) {
   return builder.create<ConstantOp>(loc, type, value);
 }
 
 /// A utility function to check if a given region is attached to a function.
 static bool isFunctionRegion(Region *region) {
   return llvm::isa<FuncOp>(region->getParentOp());
 }
 
void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
   for (auto forInst : forInsts)
     ivs->push_back(forInst.getInductionVar());
 }
 
 //===----------------------------------------------------------------------===//
 // AffineIfOp
 //===----------------------------------------------------------------------===//
 
+namespace {
+/// Remove else blocks that have nothing other than the terminator.
+struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> {
+  using OpRewritePattern<AffineIfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineIfOp ifOp,
+                                PatternRewriter &rewriter) const override {
+    if (ifOp.elseRegion().empty() || !has_single_element(*ifOp.getElseBlock()))
+      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 =
       op.getAttrOfType<IntegerSetAttr>(op.getConditionAttrName());
   if (!conditionAttr)
     return op.emitOpError(
         "requires an integer set attribute named 'condition'");
 
if (operands.size() < getIntegerSet().getNumInputs() ||
       set.getNumSymbols() > getIntegerSet().getNumSymbols()) {
     setConditional(set, operands);
     return success();
   }
 
   return failure();
 }
 
+void AffineIfOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                             MLIRContext *context) {
+  results.insert<SimplifyDeadElse>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AffineLoadOp
 //===----------------------------------------------------------------------===//
 
 void AffineLoadOp::build(Builder *builder, OperationState &result,
                          AffineMap map, ValueRange operands) {
   assert(operands.size() == 1 + map.getNumInputs() && "inconsistent operands");
   result.addOperands(operands);

mlir/lib/IR/PatternMatch.cpp

 /// This method erases an operation that is known to have no uses. The uses of
 /// the given operation *must* be known to be dead.
 void PatternRewriter::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
   notifyOperationRemoved(op);
   op->erase();
 }
 
+void PatternRewriter::eraseBlock(Block *block) {
+  for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
+    assert(op.use_empty() && "expected 'op' to have no uses");
+    eraseOp(&op);
+  }
+  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
 /// merging.
 void PatternRewriter::mergeBlocks(Block *source, Block *dest,
                                   ValueRange argValues) {
   assert(llvm::all_of(source->getPredecessors(),
                       [dest](Block *succ) { return succ == dest; }) &&

mlir/lib/Transforms/DialectConversion.cpp

void ConversionPatternRewriter::eraseOp(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()
         << "** Erase   : '" << op->getName() << "'(" << op << ")\n";
   });
   SmallVector<Value, 1> nullRepls(op->getNumResults(), nullptr);
   impl->replaceOp(op, nullRepls);
 }
 
+void ConversionPatternRewriter::eraseBlock(Block *block) {
+  llvm_unreachable("erasing blocks for dialect conversion not implemented");
+}
+
 /// Apply a signature conversion to the entry block of the given region.
 Block *ConversionPatternRewriter::applySignatureConversion(
     Region *region, TypeConverter::SignatureConversion &conversion) {
   return impl->applySignatureConversion(region, conversion);
 }
 
 void ConversionPatternRewriter::replaceUsesOfBlockArgument(BlockArgument from,
                                                            Value to) {

mlir/test/Transforms/canonicalize.mlir

func @index_cast_fold() -> (i16, index) {
   %1 = index_cast %c4 : index to i16
   %c4_i16 = constant 4 : i16
   %2 = index_cast %c4_i16 : i16 to index
   // CHECK: %[[C4_I16:.*]] = constant 4 : i16
   // CHECK: %[[C4:.*]] = constant 4 : index
   // 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