[mlir] Add redundant copy removal transform

[dfki-ehna]

Summary

This pass removes unnecessary copy operations in the following scenario which is useful for removing the redundant copy operations created by the Buffer Placement.

%from = ...
...
%to = ...
... (no user for %to)
copy(%from, %to)
... (no user for %from)
dealloc %from
...
use(%to)

[pifon2a]

@dfki-ehna Thank you!

mlir/lib/Dialect/Linalg/Transforms/CopyRemoval.cpp

29 ↗

(On Diff #274049)

please, remove /// ... lines. I don't think that they add a lot here. It should be obvious that there are operations in-between.

58 ↗

(On Diff #274049)

nit: this empty line is not needed.

74 ↗

(On Diff #274049)

Operation* toFreeingOp = FindDeallocation(from.getUsers());

and the corresponding private method.

76 ↗

(On Diff #274049)

redundant comment.

78 ↗

(On Diff #274049)

redundant comment.

80 ↗

(On Diff #274049)

redundant comment.

84 ↗

(On Diff #274049)

redundant comment.

89 ↗

(On Diff #274049)

just curious, if there is a nicer way to put 4 elements in the set. Also it probably should be SmallDenseSet.

[herhut]

mlir/lib/Dialect/Linalg/Transforms/CopyRemoval.cpp

69 ↗

(On Diff #274049)

What happens if there are multiple de-allocation nodes? Something could be allocated in block A, and then flow into B and C where it gets de-allocated. I don't think we create this yet but your rewrite needs to handle this.

89 ↗

(On Diff #274049)

Also, as it is just 4 of them, you could pointer compare.

110 ↗

(On Diff #274049)

This gets more complicated with aliasing. You could have something like

%1 = alloc
%2 = memref_cast %1 /// This aliases %1 as %2
%3 = copy(%1)
write_to(%2)
read_from(%2)
free(%1)

At least leave a TODO. It would be safer to disallow any operation with memory effects between the copy and dealloc.

[rriddle]

mlir/lib/Dialect/Linalg/Transforms/CopyRemoval.cpp

15 ↗

(On Diff #274049)

nit: This include isn't necessary.

72 ↗

(On Diff #274049)

nit: llvm::is_contained

102 ↗

(On Diff #274049)

isBeforeInBlock asserts that the operations are within the same block, is this guaranteed here?

[mehdi_amini]

Can this pass be made to not be Linalg-specific?
Can you clarify the restriction on it and when is it safe or not to use?

[dfki-ehna]

Comments are addressed.

[dfki-ehna]

@mehdi_amini Is there any generic copy interface available? The other problem would be in test files that how we can specify the copy operation class for the pass. @herhut suggested that since it is going to be a dialect independent and unique copy at the end, we can start with linalg.copy and generalize it later on.
In summary, what we'd like to achieve is to remove %to value and the Copy operation, and replace all its uses with %from value. In both of these examples, if we remove %to and the copy operation, and replace it with %from, the logic of the program will change.

There are users or aliases of %to between the allocation of %to and CopyOp:
%from = alloc()
in_place_changes_to(%from)
%to = alloc()
in_place_changes_to(%to)
%1 = exp(%to)
%2 = exp(%from)
copy(%from, %to)
dealloc(%from)
%3 = exp(%to)

There are users or aliases of %from between CopyOp and the deallocation of %from:
%from = alloc()
%to = alloc()
in_place_changes_to(%from)
copy(%from, %to)
in_place_changes_to(%to)
%1 = exp(%to)
%2 = exp(%from)
dealloc(%from)
%3 = exp(%to)

We capture both of these cases to not to change the IR.

mlir/lib/Dialect/Linalg/Transforms/CopyRemoval.cpp

89 ↗

(On Diff #274049)

I Initialized the SmallDenseSet instead of inserting them into it.

102 ↗

(On Diff #274049)

Yes, line 86 to 90 guarantees that they are in the same block.

110 ↗

(On Diff #274049)

Added another check "hasMemoryEffectOp" beside checking "hasUsers".

[dfki-ehna]

Address the multiple deallocation issue.

[dfki-ehna]

mlir/lib/Dialect/Linalg/Transforms/CopyRemoval.cpp

69 ↗

(On Diff #274049)

Thanks for catching this. Fixed and an extra test case was added.

102 ↗

(On Diff #274049)

I misunderstood your point. I added an extra condition for checking if the user operations are also in the same block as the start and end operations. Thanks.

[dfki-ehna]

Add another copy removal condition and a few more test cases.

[herhut]

Thanks, please address comments and land.

mlir/include/mlir/Dialect/Linalg/Passes.td

15 ↗

(On Diff #274806)

nit: drop the the

mlir/lib/Dialect/Linalg/Transforms/CopyRemoval.cpp

67 ↗

(On Diff #274806)

Maybe SmallPtrSet?

99 ↗

(On Diff #274806)

Maybe call this ReuseCopySourceAsTarget? Not sure that is much better but it improves on case1.

149 ↗

(On Diff #274806)

Maybe call this ReuseCopyTargetAsSource?

mlir/test/Dialect/Linalg/copy-removal.mlir

30 ↗

(On Diff #274806)

I think a CHECK-NOT cannot bind a variable anyway, so why not just use {{.*}} instead of PERCENT10?

81 ↗

(On Diff #274806)

CHECK-SAME?

82 ↗

(On Diff #274806)

CHECK-NEXT here to ensure nothing else is left?

102 ↗

(On Diff #274806)

If you want to make sure there is nothing else in the function, use CHECK-NEXT. Otherwise matching the constant does not provide value.

[mehdi_amini]

In D82757#2122827, @dfki-ehna wrote:

@mehdi_amini Is there any generic copy interface available?

In general, we create interfaces when we create the first pass that needs it. This seems like a good time to do this.

The other problem would be in test files that how we can specify the copy operation class for the pass.

Not sure I understand: we have the test dialect where you can add a copy operation implementing the interface, and you can also just add the interface to the Linalg copy.

@herhut suggested that since it is going to be a dialect independent and unique copy at the end, we can start with linalg.copy and generalize it later on.

This is a fine strategy *if this is simpler*. I am not convinced you're making any simplification by doing this: you're increasing the review load and the churn and overall this is slowing down the development.
Can you look into adding an interface? This seems like a fairly easy addition to me.

[dfki-ehna]

The following points are addressed:

• The remaining comments are resolved.
• CopyOpInterface is introduced and added to linalg.copy.
• Make this pass generic for all copies.
• The pass is moved to mlir/lib/transforms.
• copy-removal.mlir is moved to mlir/test/transforms.

[herhut]

Thanks!

[rriddle]

mlir/lib/Transforms/CopyRemoval.cpp

22

nit: Add a newline before this line.

47

nit: Use isa<> if you aren't using the result.

70

nit: This seems extremely inefficient. Can you just grab the first block and check to see if any of the others don't match?

100

Please fix the case of this function, it should be camelCase

150

Same here.

173

nit: Can you move the public section to the top?

[dfki-ehna]

@rriddle Thanks for the comments. They will be fixed in the upcoming PR for the copy removal.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 namespace linalg {
 
 class ConvOp;

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_STRUCTURED_OPS
 #define LINALG_STRUCTURED_OPS
 
 include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td"
+include "mlir/Interfaces/CopyOpInterface.td"
 
 // The Linalg `NInputs` trait provides the API for ops that are known
 // to have a specified number of inputs, all passed as operands.
 // See Linalg/LinalgTraits.h for implementation details an usage.
 class NInputs<int args_in> :
   NativeOpTrait<"linalg::NInputs<" # !cast<string>(args_in) # ">::Impl"> {}
 
 // The Linalg `NOutputs` trait provides the API for ops that are known
class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
   let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
 }
 
 //===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as special configurations of generic ops.
 //===----------------------------------------------------------------------===//
 // At the moment these are not declarative and require a bunch of C++ code.
 // In the future, these should be migrated to a declarative specification.
-def CopyOp : LinalgStructured_Op<"copy", [NInputs<1>, NOutputs<1>]> {
+def CopyOp : LinalgStructured_Op<"copy", [
+    CopyOpInterface,
+    NInputs<1>,
+    NOutputs<1>
+  ]> {
   let description = [{
     Copies the data in the input view into the output view.
 
     Usage:
 
     ```mlir
     linalg.copy(%arg0, %arg1) : memref<?xf32, stride_specification>,
                                 memref<?xf32, stride_specification>
llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       auto maybeInputMap = inputPermutation();
       auto maybeOutputMap = outputPermutation();
       unsigned inputRank = getInputShapedType(0).getRank();
       unsigned outputRank = getOutputShapedType(0).getRank();
       return SmallVector<AffineMap, 8>{
           extractOrIdentityMap(maybeInputMap, inputRank, context),
           extractOrIdentityMap(maybeOutputMap, outputRank, context)};
     }
+
+    Value getSource() { return input();}
+    Value getTarget() { return output(); }
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
 def FillOp : LinalgStructured_Op<"fill", [NInputs<0>, NOutputs<1>]> {
 

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
+add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
 add_mlir_interface(SideEffectInterfaces)
 add_mlir_interface(ViewLikeInterface)
 

mlir/include/mlir/Interfaces/CopyOpInterface.h

+//===- CopyOpInterface.h - copy operations interface ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the operation interface for copy-like operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_COPYOPINTERFACE_H_
+#define MLIR_INTERFACES_COPYOPINTERFACE_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+
+#include "mlir/Interfaces/CopyOpInterface.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_COPYOPINTERFACE_H_

mlir/include/mlir/Interfaces/CopyOpInterface.td

+//===- CopyOpInterface.td - Copy operation interface -------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for copy-like operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_COPYOPINTERFACE
+#define MLIR_INTERFACES_COPYOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def CopyOpInterface : OpInterface<"CopyOpInterface"> {
+  let description = [{
+    A copy-like operation is one that copies from source value to target value.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/"Returns the source value for this copy operation",
+      /*retTy=*/"Value",
+      /*methodName=*/"getSource"
+    >,
+    InterfaceMethod<
+      /*desc=*/"Returns the target value for this copy operation",
+      /*retTy=*/"Value",
+      /*methodName=*/"getTarget"
+    >
+  ];
+}
+
+#endif // MLIR_INTERFACES_COPYOPINTERFACE

mlir/include/mlir/Transforms/Passes.h

 template <typename T> class OperationPass;
 
 /// Creates an instance of the BufferPlacement pass.
 std::unique_ptr<Pass> createBufferPlacementPass();
 
 /// Creates an instance of the Canonicalizer pass.
 std::unique_ptr<Pass> createCanonicalizerPass();
 
+/// Create a pass that removes unnecessary Copy operations.
+std::unique_ptr<Pass> createCopyRemovalPass();
+
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
 /// Creates a loop fusion pass which fuses loops. Buffers of size less than or
 /// equal to `localBufSizeThreshold` are promoted to memory space
 /// `fastMemorySpace'.
 std::unique_ptr<OperationPass<FuncOp>>
 createLoopFusionPass(unsigned fastMemorySpace = 0,

mlir/include/mlir/Transforms/Passes.td

def Canonicalizer : Pass<"canonicalize"> {
   let description = [{
     This pass performs various types of canonicalizations over a set of
     operations. See [Operation Canonicalization](Canonicalization.md) for more
     details.
   }];
   let constructor = "mlir::createCanonicalizerPass()";
 }
 
+def CopyRemoval : FunctionPass<"copy-removal"> {
+  let summary = "Remove the redundant copies from input IR";
+  let constructor = "mlir::createCopyRemovalPass()";
+}
+
 def CSE : Pass<"cse"> {
   let summary = "Eliminate common sub-expressions";
   let description = [{
     This pass implements a generalized algorithm for common sub-expression
     elimination. This pass relies on information provided by the
     `Memory SideEffect` interface to identify when it is safe to eliminate
     operations. See [Common subexpression elimination](https://en.wikipedia.org/wiki/Common_subexpression_elimination)
     for more general details on this optimization.

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   ControlFlowInterfaces.cpp
+  CopyOpInterface.cpp
   DerivedAttributeOpInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   SideEffectInterfaces.cpp
   ViewLikeInterface.cpp
   )
 
 function(add_mlir_interface_library name)
add_mlir_library(MLIR${name}
     LINK_LIBS PUBLIC
     MLIRIR
     )
 endfunction(add_mlir_interface_library)
 
 
 add_mlir_interface_library(CallInterfaces)
 add_mlir_interface_library(ControlFlowInterfaces)
+add_mlir_interface_library(CopyOpInterface)
 add_mlir_interface_library(DerivedAttributeOpInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(LoopLikeInterface)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 

mlir/lib/Interfaces/CopyOpInterface.cpp

+//===- CopyOpInterface.cpp - Copy operations interface in MLIR ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/CopyOpInterface.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// CopyOp Interface
+//===----------------------------------------------------------------------===//
+
+/// Include the definitions of the copy operation interface.
+#include "mlir/Interfaces/CopyOpInterface.cpp.inc"

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferPlacement.cpp
   Canonicalizer.cpp
+  CopyRemoval.cpp
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
   MemRefDataFlowOpt.cpp
add_mlir_library(MLIRTransforms
 
   DEPENDS
   MLIRStandardOpsIncGen
   MLIRTransformsPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineOps
   MLIRAnalysis
+  MLIRCopyOpInterface
   MLIRLinalgOps
   MLIRLoopLikeInterface
   MLIRSCF
   MLIRPass
   MLIRTransformUtils
   MLIRVector
   )

mlir/lib/Transforms/CopyRemoval.cpp

+//===- CopyRemoval.cpp - Removing the redundant copies --------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/CopyOpInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+
+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<>> {
+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) {
+    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->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)
+        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) {
+    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;
+  }
+
+  /// Input:
+  /// func(){
+  ///   %from = alloc()
+  ///   write_to(%from)
+  ///   %to = alloc()
+  ///   copy(%from,%to)
+  ///   dealloc(%from)
+  ///   return %to
+  /// }
+  ///
+  /// Output:
+  /// func(){
+  ///   %from = alloc()
+  ///   write_to(%from)
+  ///   return %from
+  /// }
+  /// Constraints:
+  /// 1) %to, copy and dealloc must all be defined and lie in the same block.
+  /// 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) {
+    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());
+    Operation *toDefiningOp = to.getDefiningOp();
+    if (!fromDefiningOp || !fromFreeingOp || !toDefiningOp ||
+        !areOpsInTheSameBlock({fromFreeingOp, toDefiningOp, copy}) ||
+        hasUsersBetween(to, toDefiningOp, copy) ||
+        hasUsersBetween(from, copy, fromFreeingOp) ||
+        hasMemoryEffectOpBetween(copy, fromFreeingOp))
+      return;
+
+    to.replaceAllUsesWith(from);
+    eraseList.insert(copy);
+    eraseList.insert(toDefiningOp);
+    eraseList.insert(fromFreeingOp);
+  }
+
+  /// Input:
+  /// func(){
+  ///   %to = alloc()
+  ///   %from = alloc()
+  ///   write_to(%from)
+  ///   copy(%from,%to)
+  ///   dealloc(%from)
+  ///   return %to
+  /// }
+  ///
+  /// Output:
+  /// func(){
+  ///   %to = alloc()
+  ///   write_to(%to)
+  ///   return %to
+  /// }
+  /// Constraints:
+  /// 1) %from, copy and dealloc must all be defined and lie in the same block.
+  /// 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) {
+    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>();
+}

mlir/test/Transforms/copy-removal.mlir

+// RUN: mlir-opt -copy-removal -split-input-file %s | FileCheck %s
+
+// All linalg copies except the linalg.copy(%1, %9) must be removed since the
+// defining operation of %1 and its DeallocOp have been defined in another block.
+
+// CHECK-LABEL: func @nested_region_control_flow_div_nested
+func @nested_region_control_flow_div_nested(%arg0: index, %arg1: index) -> memref<?x?xf32> {
+  %0 = cmpi "eq", %arg0, %arg1 : index
+  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
+  // CHECK: %{{.*}} = scf.if
+  %2 = scf.if %0 -> (memref<?x?xf32>) {
+    // CHECK: %[[PERCENT3:.*]] = scf.if
+    %3 = scf.if %0 -> (memref<?x?xf32>) {
+      %c0_0 = constant 0 : index
+      %7 = dim %1, %c0_0 : memref<?x?xf32>
+      %c1_1 = constant 1 : index
+      %8 = dim %1, %c1_1 : memref<?x?xf32>
+      %9 = alloc(%7, %8) : memref<?x?xf32>
+      // CHECK: linalg.copy({{.*}}, %[[PERCENT9:.*]])
+      linalg.copy(%1, %9) : memref<?x?xf32>, memref<?x?xf32>
+      // CHECK: scf.yield %[[PERCENT9]]
+      scf.yield %9 : memref<?x?xf32>
+    } else {
+      // CHECK: %[[PERCENT7:.*]] = alloc
+      %7 = alloc(%arg0, %arg1) : memref<?x?xf32>
+      %c0_0 = constant 0 : index
+      %8 = dim %7, %c0_0 : memref<?x?xf32>
+      %c1_1 = constant 1 : index
+      %9 = dim %7, %c1_1 : memref<?x?xf32>
+      // CHECK-NOT: %{{.*}} = alloc
+      // CHECK-NOT: linalg.copy(%[[PERCENT7]], %{{.*}})
+      // CHECK-NOT: dealloc %[[PERCENT7]]
+      %10 = alloc(%8, %9) : memref<?x?xf32>
+      linalg.copy(%7, %10) : memref<?x?xf32>, memref<?x?xf32>
+      dealloc %7 : memref<?x?xf32>
+      // CHECK: scf.yield %[[PERCENT7]]
+      scf.yield %10 : memref<?x?xf32>
+    }
+    %c0 = constant 0 : index
+    %4 = dim %3, %c0 : memref<?x?xf32>
+    %c1 = constant 1 : index
+    %5 = dim %3, %c1 : memref<?x?xf32>
+    // CHECK-NOT: %{{.*}} = alloc
+    // CHECK-NOT: linalg.copy(%[[PERCENT3]], %{{.*}})
+    // CHECK-NOT: dealloc %[[PERCENT3]]
+    %6 = alloc(%4, %5) : memref<?x?xf32>
+    linalg.copy(%3, %6) : memref<?x?xf32>, memref<?x?xf32>
+    dealloc %3 : memref<?x?xf32>
+    // CHECK: scf.yield %[[PERCENT3]]
+    scf.yield %6 : memref<?x?xf32>
+  } else {
+    // CHECK: %[[PERCENT3:.*]] = alloc
+    %3 = alloc(%arg1, %arg1) : memref<?x?xf32>
+    %c0 = constant 0 : index
+    %4 = dim %3, %c0 : memref<?x?xf32>
+    %c1 = constant 1 : index
+    %5 = dim %3, %c1 : memref<?x?xf32>
+    // CHECK-NOT: %{{.*}} = alloc
+    // CHECK-NOT: linalg.copy(%[[PERCENT3]], %{{.*}})
+    // CHECK-NOT: dealloc %[[PERCENT3]]
+    %6 = alloc(%4, %5) : memref<?x?xf32>
+    linalg.copy(%3, %6) : memref<?x?xf32>, memref<?x?xf32>
+    dealloc %3 : memref<?x?xf32>
+    // CHECK: scf.yield %[[PERCENT3]]
+    scf.yield %6 : memref<?x?xf32>
+  }
+  dealloc %1 : memref<?x?xf32>
+  return %2 : memref<?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @simple_test
+func @simple_test() -> memref<5xf32> {
+  %temp = alloc() : memref<5xf32>
+  %ret = alloc() : memref<5xf32>
+  linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
+  dealloc %ret : memref<5xf32>
+  return %temp : memref<5xf32>
+}
+// CHECK-SAME: () -> memref<5xf32>
+// CHECK-NEXT: %[[ret:.*]] = alloc()
+// CHECK-NOT: linalg.copy(%[[ret]], %{{.*}})
+// CHECK-NOT: dealloc %[[ret]]
+// CHECK: return %[[ret]]
+
+// -----
+
+// It is legal to remove the copy operation that %ret has a usage before the copy
+// operation. The allocation of %temp and the deallocation of %ret should be also
+// removed.
+
+// CHECK-LABEL: func @test_with_ret_usage_before_copy
+func @test_with_ret_usage_before_copy() -> memref<5xf32> {
+  %ret = alloc() : memref<5xf32>
+  %temp = alloc() : memref<5xf32>
+  %c0 = constant 0 : index
+  %dimension = dim %ret, %c0 : memref<5xf32>
+  linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
+  dealloc %ret : memref<5xf32>
+  return %temp : memref<5xf32>
+}
+// CHECK-NEXT: %[[ret:.*]] = alloc()
+// CHECK-NOT: %{{.*}} = alloc
+// CHECK-NEXT: %{{.*}} = constant
+// CHECK-NEXT: %[[DIM:.*]] = dim %[[ret]]
+// CHECK-NOT: linalg.copy(%[[ret]], %{{.*}})
+// CHECK-NOT: dealloc %[[ret]]
+// CHECK: return %[[ret]]
+
+// -----
+
+// It is illegal to remove a copy operation that %ret has a usage after copy
+// operation.
+
+// CHECK-LABEL: func @test_with_ret_usage_after_copy
+func @test_with_ret_usage_after_copy() -> memref<5xf32> {
+  %ret = alloc() : memref<5xf32>
+  %temp = alloc() : memref<5xf32>
+  // CHECK: linalg.copy
+  linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
+  %c0 = constant 0 : index
+  %dimension = dim %ret, %c0 : memref<5xf32>
+  dealloc %ret : memref<5xf32>
+  return %temp : memref<5xf32>
+}
+
+// -----
+
+// It is illegal to remove a copy operation that %temp has a usage before copy
+// operation.
+
+// CHECK-LABEL: func @test_with_temp_usage_before_copy
+func @test_with_temp_usage_before_copy() -> memref<5xf32> {
+  %ret = alloc() : memref<5xf32>
+  %temp = alloc() : memref<5xf32>
+  %c0 = constant 0 : index
+  %dimension = dim %temp, %c0 : memref<5xf32>
+  // CHECK: linalg.copy
+  linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
+  dealloc %ret : memref<5xf32>
+  return %temp : memref<5xf32>
+}
+
+// -----
+
+// It is legal to remove the copy operation that %temp has a usage after the copy
+// operation. The allocation of %temp and the deallocation of %ret should be also
+// removed.
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @test_with_temp_usage_after_copy
+func @test_with_temp_usage_after_copy() -> memref<5xf32> {
+  %ret = alloc() : memref<5xf32>
+  %res = alloc() : memref<5xf32>
+  %temp = alloc() : memref<5xf32>
+  linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
+  linalg.generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel"]} %temp, %res {
+  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
+    %tmp1 = exp %gen1_arg0 : f32
+    linalg.yield %tmp1 : f32
+  }: memref<5xf32>, memref<5xf32>
+  dealloc %ret : memref<5xf32>
+  return %temp : memref<5xf32>
+}
+// CHECK-NEXT: %[[ret:.*]] = alloc()
+// CHECK-NEXT: %[[res:.*]] = alloc()
+// CHECK-NOT: %{{.*}} = alloc()
+// CHECK-NOT: linalg.copy
+// CHECK-NOT: dealloc %[[ret]]
+// CHECK: return %[[ret]]
+
+// -----
+
+// CHECK-LABEL: func @make_allocation
+func @make_allocation() -> memref<5xf32> {
+  %mem = alloc() : memref<5xf32>
+  return %mem : memref<5xf32>
+}
+
+// CHECK-LABEL: func @test_with_function_call
+func @test_with_function_call() -> memref<5xf32> {
+  // CHECK-NEXT: %[[ret:.*]] = call @make_allocation() : () -> memref<5xf32>
+  %ret = call @make_allocation() : () -> (memref<5xf32>)
+  // CHECK-NOT: %{{.*}} = alloc
+  // CHECK-NOT: linalg.copy(%[[ret]], %{{.*}})
+  // CHECK-NOT: dealloc %[[ret]]
+  %temp = alloc() : memref<5xf32>
+  linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
+  dealloc %ret : memref<5xf32>
+  // CHECK: return %[[ret]]
+  return %temp : memref<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @multiple_deallocs_in_different_blocks
+func @multiple_deallocs_in_different_blocks(%cond : i1) -> memref<5xf32> {
+  // CHECK-NEXT: %[[PERCENT0:.*]] = alloc()
+  %0 = alloc() : memref<5xf32>
+  cond_br %cond, ^bb1, ^bb2
+^bb1:
+  dealloc %0 : memref<5xf32>
+  // CHECK: br ^[[BB3:.*]](%[[PERCENT0]]
+  br ^bb3(%0 : memref<5xf32>)
+^bb2:
+  // CHECK-NOT: %{{.*}} = alloc
+  // CHECK-NOT: linalg.copy(%[[PERCENT0]], %{{.*}})
+  // CHECK-NOT: dealloc %[[PERCENT0]]
+  %temp = alloc() : memref<5xf32>
+  linalg.copy(%0, %temp) : memref<5xf32>, memref<5xf32>
+  dealloc %0 : memref<5xf32>
+  // CHECK: br ^[[BB3]](%[[PERCENT0]]
+  br ^bb3(%temp : memref<5xf32>)
+^bb3(%res : memref<5xf32>):
+  return %res : memref<5xf32>
+}
+
+// -----
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @test_ReuseCopyTargetAsSource
+func @test_ReuseCopyTargetAsSource(%arg0: memref<2xf32>, %result: memref<2xf32>){
+  // CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[RES:.*]]: memref<2xf32>)
+  // CHECK-NOT: %{{.*}} = alloc
+  %temp = alloc() : memref<2xf32>
+  // CHECK-NEXT: linalg.generic
+  // CHECK-SAME: %[[ARG0]], %[[RES]]
+  // CHECK-NOT: linalg.copy(%{{.*}}, %[[RES]])
+  // CHECK-NOT: dealloc %{{.*}}
+  linalg.generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel"]} %arg0, %temp {
+  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
+    %tmp2 = exp %gen2_arg0 : f32
+    linalg.yield %tmp2 : f32
+  }: memref<2xf32>, memref<2xf32>
+  "linalg.copy"(%temp, %result) : (memref<2xf32>, memref<2xf32>) -> ()
+  dealloc %temp : memref<2xf32>
+  // CHECK: return
+  return
+}
+
+// -----
+
+// Copy operation must not be removed since an operation writes to %to value
+// before copy.
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @test_ReuseCopyTargetAsSource
+func @test_ReuseCopyTargetAsSource(%arg0: memref<2xf32>){
+  %to = alloc() : memref<2xf32>
+  %temp = alloc() : memref<2xf32>
+  linalg.generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel"]} %arg0, %temp {
+  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
+    %tmp1 = exp %gen1_arg0 : f32
+    linalg.yield %tmp1 : f32
+  }: memref<2xf32>, memref<2xf32>
+  linalg.generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel"]} %arg0, %to {
+  ^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
+}