[mlir] Add redundant copy removal transform
authorEhsan Toosi <ehsan.nadjaran_toosi@dfki.de>
Thu, 25 Jun 2020 15:02:11 +0000 (17:02 +0200)
committerEhsan Toosi <ehsan.nadjaran_toosi@dfki.de>
Fri, 3 Jul 2020 13:36:25 +0000 (15:36 +0200)
This pass removes redundant dialect-independent Copy operations in different
situations like the following:

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

Differential Revision: https://reviews.llvm.org/D82757

12 files changed:
mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
mlir/include/mlir/Interfaces/CMakeLists.txt
mlir/include/mlir/Interfaces/CopyOpInterface.h [new file with mode: 0644]
mlir/include/mlir/Interfaces/CopyOpInterface.td [new file with mode: 0644]
mlir/include/mlir/Transforms/Passes.h
mlir/include/mlir/Transforms/Passes.td
mlir/lib/Interfaces/CMakeLists.txt
mlir/lib/Interfaces/CopyOpInterface.cpp [new file with mode: 0644]
mlir/lib/Transforms/CMakeLists.txt
mlir/lib/Transforms/CopyRemoval.cpp [new file with mode: 0644]
mlir/test/Transforms/copy-removal.mlir [new file with mode: 0644]

index 8061568..f89965a 100644 (file)
@@ -22,6 +22,7 @@
 #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"
index 067d7c2..4dd6521 100644 (file)
@@ -17,6 +17,7 @@
 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.
@@ -56,7 +57,11 @@ class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
 //===----------------------------------------------------------------------===//
 // 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.
 
@@ -141,6 +146,9 @@ def CopyOp : LinalgStructured_Op<"copy", [NInputs<1>, NOutputs<1>]> {
           extractOrIdentityMap(maybeInputMap, inputRank, context),
           extractOrIdentityMap(maybeOutputMap, outputRank, context)};
     }
+
+    Value getSource() { return input();}
+    Value getTarget() { return output(); }
   }];
   let verifier = [{ return ::verify(*this); }];
 
index ff3577d..51f3f8a 100644 (file)
@@ -1,5 +1,6 @@
 add_mlir_interface(CallInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
+add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
diff --git a/mlir/include/mlir/Interfaces/CopyOpInterface.h b/mlir/include/mlir/Interfaces/CopyOpInterface.h
new file mode 100644 (file)
index 0000000..d6dc409
--- /dev/null
@@ -0,0 +1,24 @@
+//===- 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_
diff --git a/mlir/include/mlir/Interfaces/CopyOpInterface.td b/mlir/include/mlir/Interfaces/CopyOpInterface.td
new file mode 100644 (file)
index 0000000..658474d
--- /dev/null
@@ -0,0 +1,37 @@
+//===- 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
index f54333d..d528723 100644 (file)
@@ -32,6 +32,9 @@ 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();
 
index 4cc7d63..9e0d5c4 100644 (file)
@@ -174,6 +174,11 @@ def Canonicalizer : Pass<"canonicalize"> {
   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 = [{
index d0eea1c..19b4e0a 100644 (file)
@@ -1,6 +1,7 @@
 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   ControlFlowInterfaces.cpp
+  CopyOpInterface.cpp
   DerivedAttributeOpInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
@@ -26,6 +27,7 @@ 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)
diff --git a/mlir/lib/Interfaces/CopyOpInterface.cpp b/mlir/lib/Interfaces/CopyOpInterface.cpp
new file mode 100644 (file)
index 0000000..8e6132c
--- /dev/null
@@ -0,0 +1,18 @@
+//===- 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"
index f391a26..3c6b393 100644 (file)
@@ -3,6 +3,7 @@ add_subdirectory(Utils)
 add_mlir_library(MLIRTransforms
   BufferPlacement.cpp
   Canonicalizer.cpp
+  CopyRemoval.cpp
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
@@ -30,6 +31,7 @@ add_mlir_library(MLIRTransforms
   LINK_LIBS PUBLIC
   MLIRAffineOps
   MLIRAnalysis
+  MLIRCopyOpInterface
   MLIRLinalgOps
   MLIRLoopLikeInterface
   MLIRSCF
diff --git a/mlir/lib/Transforms/CopyRemoval.cpp b/mlir/lib/Transforms/CopyRemoval.cpp
new file mode 100644 (file)
index 0000000..28648e0
--- /dev/null
@@ -0,0 +1,191 @@
+//===- 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>();
+}
diff --git a/mlir/test/Transforms/copy-removal.mlir b/mlir/test/Transforms/copy-removal.mlir
new file mode 100644 (file)
index 0000000..f750dab
--- /dev/null
@@ -0,0 +1,285 @@
+// 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
+}