[mlir][NFC] Polish copy removal transform

[dfki-ehna]

Summary

Address a few remaining comments in copy removal transform.

[bondhugula]

Is this an NFC patch?

mlir/lib/Transforms/CopyRemoval.cpp

81–85

Please !operations.empty() (doesn't make a difference for ArrayRef but in general).

[herhut]

Maybe mark it as NFC. Looks good with @bondhugula's comment addressed.

[dfki-ehna]

Address the comments.

mlir/lib/Transforms/CopyRemoval.cpp

 using namespace mlir;
 using namespace MemoryEffects;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // CopyRemovalPass
 //===----------------------------------------------------------------------===//
+
 /// This pass removes the redundant Copy operations. Additionally, it
 /// removes the leftover definition and deallocation operations by erasing the
 /// copy operation.
 class CopyRemovalPass : public PassWrapper<CopyRemovalPass, OperationPass<>> {
+public:
+  void runOnOperation() override {
+    getOperation()->walk([&](CopyOpInterface copyOp) {
+      reuseCopySourceAsTarget(copyOp);
+      reuseCopyTargetAsSource(copyOp);
+    });
+    for (Operation *op : eraseList)
+      op->erase();
+  }
+
 private:
   /// List of operations that need to be removed.
   DenseSet<Operation *> eraseList;
 
   /// Returns the deallocation operation for `value` in `block` if it exists.
   Operation *getDeallocationInBlock(Value value, Block *block) {
+    assert(block && "Block cannot be null");
     auto valueUsers = value.getUsers();
     auto it = llvm::find_if(valueUsers, [&](Operation *op) {
       auto effects = dyn_cast<MemoryEffectOpInterface>(op);
       return effects && op->getBlock() == block && effects.hasEffect<Free>();
     });
     return (it == valueUsers.end() ? nullptr : *it);
   }
 
   /// Returns true if an operation between start and end operations has memory
   /// effect.
   bool hasMemoryEffectOpBetween(Operation *start, Operation *end) {
+    assert((start || end) && "Start and end operations cannot be null");
     assert(start->getBlock() == end->getBlock() &&
            "Start and end operations should be in the same block.");
     Operation *op = start->getNextNode();
     while (op->isBeforeInBlock(end)) {
-      auto effects = dyn_cast<MemoryEffectOpInterface>(op);
-      if (effects)
+      if (isa<MemoryEffectOpInterface>(op))
         return true;
       op = op->getNextNode();
     }
     return false;
   };
 
   /// Returns true if `val` value has at least a user between `start` and
   /// `end` operations.
   bool hasUsersBetween(Value val, Operation *start, Operation *end) {
+    assert((start || end) && "Start and end operations cannot be null");
     Block *block = start->getBlock();
     assert(block == end->getBlock() &&
            "Start and end operations should be in the same block.");
     return llvm::any_of(val.getUsers(), [&](Operation *op) {
       return op->getBlock() == block && start->isBeforeInBlock(op) &&
              op->isBeforeInBlock(end);
     });
   };
 
   bool areOpsInTheSameBlock(ArrayRef<Operation *> operations) {
-    llvm::SmallPtrSet<Block *, 4> blocks;
-    for (Operation *op : operations)
-      blocks.insert(op->getBlock());
-    return blocks.size() == 1;
+    assert(!operations.empty() &&
+           "The operations list should contain at least a single operation");
+    Block *block = operations.front()->getBlock();
+    return llvm::none_of(
+        operations, [&](Operation *op) { return block != op->getBlock(); });
   }
 
   /// Input:
   /// func(){
   ///   %from = alloc()
   ///   write_to(%from)
   ///   %to = alloc()
   ///   copy(%from,%to)
private:
   /// 2) This transformation cannot be applied if there is a single user/alias
   /// of `to` value between the defining operation of `to` and the copy
   /// operation.
   /// 3) This transformation cannot be applied if there is a single user/alias
   /// of `from` value between the copy operation and the deallocation of `from`.
   /// TODO: Alias analysis is not available at the moment. Currently, we check
   /// if there are any operations with memory effects between copy and
   /// deallocation operations.
-  void ReuseCopySourceAsTarget(CopyOpInterface copyOp) {
+  void reuseCopySourceAsTarget(CopyOpInterface copyOp) {
     if (eraseList.count(copyOp))
       return;
 
     Value from = copyOp.getSource();
     Value to = copyOp.getTarget();
 
     Operation *copy = copyOp.getOperation();
     Operation *fromDefiningOp = from.getDefiningOp();
private:
   /// 2) This transformation cannot be applied if there is a single user/alias
   /// of `to` value between the defining operation of `from` and the copy
   /// operation.
   /// 3) This transformation cannot be applied if there is a single user/alias
   /// of `from` value between the copy operation and the deallocation of `from`.
   /// TODO: Alias analysis is not available at the moment. Currently, we check
   /// if there are any operations with memory effects between copy and
   /// deallocation operations.
-  void ReuseCopyTargetAsSource(CopyOpInterface copyOp) {
+  void reuseCopyTargetAsSource(CopyOpInterface copyOp) {
     if (eraseList.count(copyOp))
       return;
 
     Value from = copyOp.getSource();
     Value to = copyOp.getTarget();
 
     Operation *copy = copyOp.getOperation();
     Operation *fromDefiningOp = from.getDefiningOp();
     Operation *fromFreeingOp = getDeallocationInBlock(from, copy->getBlock());
     if (!fromDefiningOp || !fromFreeingOp ||
         !areOpsInTheSameBlock({fromFreeingOp, fromDefiningOp, copy}) ||
         hasUsersBetween(to, fromDefiningOp, copy) ||
         hasUsersBetween(from, copy, fromFreeingOp) ||
         hasMemoryEffectOpBetween(copy, fromFreeingOp))
       return;
 
     from.replaceAllUsesWith(to);
     eraseList.insert(copy);
     eraseList.insert(fromDefiningOp);
     eraseList.insert(fromFreeingOp);
   }
-
-public:
-  void runOnOperation() override {
-    getOperation()->walk([&](CopyOpInterface copyOp) {
-      ReuseCopySourceAsTarget(copyOp);
-      ReuseCopyTargetAsSource(copyOp);
-    });
-    for (Operation *op : eraseList)
-      op->erase();
-  }
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // CopyRemovalPass construction
 //===----------------------------------------------------------------------===//
+
 std::unique_ptr<Pass> mlir::createCopyRemovalPass() {
   return std::make_unique<CopyRemovalPass>();
 }