[mlir] Fix bug in copy removal

[dfki-ehna]

Summary

A crash could happen due to an illegal copy removal in cases such as

%from = scf.for ...
copy(%from, %to) 
dealloc %to

where copy removal removes the entire scf.for operation with its region, Copy, Dealloc operations. Therefore, all getDefiningOp to find the source of Copy operation are replaced with getAllocationOpInBlock. To test this crash, loop_alloc test is added. check_with_affine_dialect test is also added to check the crash issue in TensorFlow::lhlo-copy-removal (issue: https://github.com/tensorflow/mlir-hlo/issues/1). TensorFlow::lhlo-copy-removal pass is going to be replaced with mlir::copy-removal pass.

[bondhugula]

Nit: commit summary: to -> due to / during the

mlir/lib/Transforms/CopyRemoval.cpp

42

Please document this field.

44

nullptr otherwise.

mlir/test/Transforms/copy-removal.mlir

326

Nit: addition spell fix.

[dfki-ehna]

Address the comments.

[dfki-ehna]

Thanks.

[bondhugula]

@dfki-ehna Thanks for fixing the posted bug. Please do rely on @herhut or @pifon2a's review - they are I assume more familiar with prior code.

[pifon2a]

mlir/lib/Transforms/CopyRemoval.cpp

33

const std::pair<Value, Value>&

[herhut]

Can you be a bit more descriptive in the summary of what bug was fixed? If the bug was filed somewhere, reference it. Otherwise please describe a little more.

[dfki-ehna]

Address the comment.

mlir/lib/Transforms/CopyRemoval.cpp

 /// copy operation.
 class CopyRemovalPass : public PassWrapper<CopyRemovalPass, OperationPass<>> {
 public:
   void runOnOperation() override {
     getOperation()->walk([&](CopyOpInterface copyOp) {
       reuseCopySourceAsTarget(copyOp);
       reuseCopyTargetAsSource(copyOp);
     });
+    for (std::pair<Value, Value> &pair : replaceList)
+      pair.first.replaceAllUsesWith(pair.second);
     for (Operation *op : eraseList)
       op->erase();
   }
 
 private:
   /// List of operations that need to be removed.
-  DenseSet<Operation *> eraseList;
+  llvm::SmallPtrSet<Operation *, 4> eraseList;
+
+  /// List of values that need to be replaced with their counterparts.
+  llvm::SmallDenseSet<std::pair<Value, Value>, 4> replaceList;
+
+  /// Returns the allocation operation for `value` in `block` if it exists.
+  /// nullptr otherwise.
+  Operation *getAllocationOpInBlock(Value value, Block *block) {
+    assert(block && "Block cannot be null");
+    Operation *op = value.getDefiningOp();
+    if (op && op->getBlock() == block) {
+      auto effects = dyn_cast<MemoryEffectOpInterface>(op);
+      if (effects && effects.hasEffect<Allocate>())
+        return op;
+    }
+    return nullptr;
+  }
 
   /// Returns the deallocation operation for `value` in `block` if it exists.
-  Operation *getDeallocationInBlock(Value value, Block *block) {
+  /// nullptr otherwise.
+  Operation *getDeallocationOpInBlock(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);
   }
private:
   void reuseCopySourceAsTarget(CopyOpInterface copyOp) {
     if (eraseList.count(copyOp))
       return;
 
     Value from = copyOp.getSource();
     Value to = copyOp.getTarget();
 
     Operation *copy = copyOp.getOperation();
+    Block *copyBlock = copy->getBlock();
     Operation *fromDefiningOp = from.getDefiningOp();
-    Operation *fromFreeingOp = getDeallocationInBlock(from, copy->getBlock());
-    Operation *toDefiningOp = to.getDefiningOp();
+    Operation *fromFreeingOp = getDeallocationOpInBlock(from, copyBlock);
+    Operation *toDefiningOp = getAllocationOpInBlock(to, copyBlock);
     if (!fromDefiningOp || !fromFreeingOp || !toDefiningOp ||
         !areOpsInTheSameBlock({fromFreeingOp, toDefiningOp, copy}) ||
         hasUsersBetween(to, toDefiningOp, copy) ||
         hasUsersBetween(from, copy, fromFreeingOp) ||
         hasMemoryEffectOpBetween(copy, fromFreeingOp))
       return;
 
-    to.replaceAllUsesWith(from);
+    replaceList.insert({to, from});
     eraseList.insert(copy);
     eraseList.insert(toDefiningOp);
     eraseList.insert(fromFreeingOp);
   }
 
   /// Input:
   /// func(){
   ///   %to = alloc()
private:
   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());
+    Block *copyBlock = copy->getBlock();
+    Operation *fromDefiningOp = getAllocationOpInBlock(from, copyBlock);
+    Operation *fromFreeingOp = getDeallocationOpInBlock(from, copyBlock);
     if (!fromDefiningOp || !fromFreeingOp ||
         !areOpsInTheSameBlock({fromFreeingOp, fromDefiningOp, copy}) ||
         hasUsersBetween(to, fromDefiningOp, copy) ||
         hasUsersBetween(from, copy, fromFreeingOp) ||
         hasMemoryEffectOpBetween(copy, fromFreeingOp))
       return;
 
-    from.replaceAllUsesWith(to);
+    replaceList.insert({from, to});
     eraseList.insert(copy);
     eraseList.insert(fromDefiningOp);
     eraseList.insert(fromFreeingOp);
   }
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // CopyRemovalPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createCopyRemovalPass() {
   return std::make_unique<CopyRemovalPass>();
 }

mlir/test/Transforms/copy-removal.mlir

^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
     %tmp2 = exp %gen2_arg0 : f32
     linalg.yield %tmp2 : f32
   }: memref<2xf32>, memref<2xf32>
   // CHECK: linalg.copy
   "linalg.copy"(%temp, %to) : (memref<2xf32>, memref<2xf32>) -> ()
   dealloc %temp : memref<2xf32>
   return
 }
+
+// -----
+
+// The only redundant copy is linalg.copy(%4, %5)
+
+// CHECK-LABEL: func @loop_alloc
+func @loop_alloc(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<2xf32>, %arg4: memref<2xf32>) {
+  // CHECK: %{{.*}} = alloc()
+  %0 = alloc() : memref<2xf32>
+  dealloc %0 : memref<2xf32>
+  // CHECK: %{{.*}} = alloc()
+  %1 = alloc() : memref<2xf32>
+  // CHECK: linalg.copy
+  linalg.copy(%arg3, %1) : memref<2xf32>, memref<2xf32>
+  %2 = scf.for %arg5 = %arg0 to %arg1 step %arg2 iter_args(%arg6 = %1) -> (memref<2xf32>) {
+    %3 = cmpi "eq", %arg5, %arg1 : index
+    // CHECK: dealloc
+    dealloc %arg6 : memref<2xf32>
+    // CHECK: %[[PERCENT4:.*]] = alloc()
+    %4 = alloc() : memref<2xf32>
+    // CHECK-NOT: alloc
+    // CHECK-NOT: linalg.copy
+    // CHECK-NOT: dealloc
+    %5 = alloc() : memref<2xf32>
+    linalg.copy(%4, %5) : memref<2xf32>, memref<2xf32>
+    dealloc %4 : memref<2xf32>
+    // CHECK: %[[PERCENT6:.*]] = alloc()
+    %6 = alloc() : memref<2xf32>
+    // CHECK: linalg.copy(%[[PERCENT4]], %[[PERCENT6]])
+    linalg.copy(%5, %6) : memref<2xf32>, memref<2xf32>
+    scf.yield %6 : memref<2xf32>
+  }
+  // CHECK: linalg.copy
+  linalg.copy(%2, %arg4) : memref<2xf32>, memref<2xf32>
+  dealloc %2 : memref<2xf32>
+  return
+}
+
+// -----
+
+// The linalg.copy operation can be removed in addition to alloc and dealloc
+// operations. All uses of %0 is then replaced with %arg2.
+
+// CHECK-LABEL: func @check_with_affine_dialect
+func @check_with_affine_dialect(%arg0: memref<4xf32>, %arg1: memref<4xf32>, %arg2: memref<4xf32>) {
+  // CHECK-SAME: (%[[ARG0:.*]]: memref<4xf32>, %[[ARG1:.*]]: memref<4xf32>, %[[RES:.*]]: memref<4xf32>)
+  // CHECK-NOT: alloc
+  %0 = alloc() : memref<4xf32>
+  affine.for %arg3 = 0 to 4 {
+    %5 = affine.load %arg0[%arg3] : memref<4xf32>
+    %6 = affine.load %arg1[%arg3] : memref<4xf32>
+    %7 = cmpf "ogt", %5, %6 : f32
+    // CHECK: %[[SELECT_RES:.*]] = select
+    %8 = select %7, %5, %6 : f32
+    // CHECK-NEXT: affine.store %[[SELECT_RES]], %[[RES]]
+    affine.store %8, %0[%arg3] : memref<4xf32>
+  }
+  // CHECK-NOT: linalg.copy
+  // CHECK-NOT: dealloc
+  "linalg.copy"(%0, %arg2) : (memref<4xf32>, memref<4xf32>) -> ()
+  dealloc %0 : memref<4xf32>
+  //CHECK: return
+  return
+}