[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

30

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

59

nit: this empty line is not needed.

75

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

and the corresponding private method.

77

redundant comment.

79

redundant comment.

81

redundant comment.

85

redundant comment.

90

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

70

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.

90

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

111

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

16

nit: This include isn't necessary.

73

nit: llvm::is_contained

103

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

90

I Initialized the SmallDenseSet instead of inserting them into it.

103

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

111

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

[dfki-ehna]

Address the multiple deallocation issue.

[dfki-ehna]

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

70

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

103

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

nit: drop the the

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

68

Maybe SmallPtrSet?

100

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

150

Maybe call this ReuseCopyTargetAsSource?

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

31

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

82

CHECK-SAME?

83

CHECK-NEXT here to ensure nothing else is left?

103

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 ↗

(On Diff #275386)

nit: Add a newline before this line.

47 ↗

(On Diff #275386)

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

70 ↗

(On Diff #275386)

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 ↗

(On Diff #275386)

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

150 ↗

(On Diff #275386)

Same here.

173 ↗

(On Diff #275386)

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/Passes.h

 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<ModuleOp>>
 createConvertLinalgOnTensorsToBuffersPass();
 
+/// Create a pass that removes unnecessary Linalg Copy operations.
+std::unique_ptr<Pass> createLinalgCopyRemovalPass();
+
 /// Patterns for fusing linalg operation on tensors.
 void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
                                            OwningRewritePatternList &patterns);
 
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors.
 void populateLinalgFoldUnitExtentDimsPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_PASSES_H_

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

 //===-- Passes.td - Linalg pass definition file ------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_PASSES
 #define MLIR_DIALECT_LINALG_PASSES
 
 include "mlir/Pass/PassBase.td"
 
+def LinalgCopyRemoval : FunctionPass<"linalg-copy-removal"> {
+  let summary = "Remove the redundant copies from the input IR";
+  let constructor = "mlir::createLinalgCopyRemovalPass()";
+}
+
 def LinalgFoldUnitExtentDims : FunctionPass<"linalg-fold-unit-extent-dims"> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [
     Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
             /*default=*/"false",
       "Only folds the one-trip loops from Linalg ops on tensors "
       "(for testing purposes only)">

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
+  CopyRemoval.cpp
   DropUnitDims.cpp
   Fusion.cpp
   Hoisting.cpp
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
   TensorsToBuffers.cpp
   Tiling.cpp

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

+//===- CopyRemoval.cpp - Removing the redundant Linalg 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the pass for removing redundant Linalg.CopyOp.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace MemoryEffects;
+
+//===----------------------------------------------------------------------===//
+// CopyRemovalPass
+//===----------------------------------------------------------------------===//
+/// This pass removes the redundant `to` value and its copy operation and
+/// replaces all of it uses with `from` value.
+///
+/// Input:
+/// %from = ...
+/// %to = ...
+/// use(%from)
+/// copy(%from,%to)
+/// use(%to)
+/// dealloc(%from)
+/// use(%to)
+///
+/// Output:
+/// %from = ...
+/// use(%from)
+/// use(%from)
+/// use(%from)
+///
+/// Constraints:
+/// 1) `from`, `to`, copy(%from,%to), and dealloc(%from) must all be defined and
+/// be in the same block for safe removing without considering CFG.
+/// 2) This transformation cannot be applied if there is a single user 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 of `from`
+/// value between the copy operation and the deallocation of `from`.
+class CopyRemovalPass : public PassWrapper<CopyRemovalPass, OperationPass<>> {
+private:
+  Operation *findDeallocation(Value value) {
+    auto valueUsers = value.getUsers();
+    auto it = llvm::find_if(valueUsers, [&](Operation *op) {
+      auto effects = dyn_cast<MemoryEffectOpInterface>(op);
+      return effects && effects.hasEffect<Free>() &&
+             llvm::is_contained(op->getOperands(), value);
+    });
+    return (it == valueUsers.end() ? nullptr : *it);
+  }
+
+  /// Returns true if any operations start and end operations has memory effect.
+  bool hasMemoryEffectOp(Operation *start, Operation *end) {
+    Operation *op = start->getNextNode();
+    while (op->isBeforeInBlock(end)) {
+      auto effects = dyn_cast<MemoryEffectOpInterface>(op);
+      if (effects)
+        return true;
+      op = op->getNextNode();
+    }
+    return false;
+  };
+
+public:
+  void runOnOperation() override {
+    SmallVector<Operation *, 2> eraseList;
+    getOperation()->walk([&](linalg::CopyOp copyOp) {
+      Value from = copyOp.input();
+      Value to = copyOp.output();
+
+      Operation *fromFreeingOp = findDeallocation(from);
+      Operation *fromDefiningOp = from.getDefiningOp();
+      Operation *toDefiningOp = to.getDefiningOp();
+      Operation *copy = copyOp.getOperation();
+      if (!fromDefiningOp || !fromFreeingOp || !toDefiningOp)
+        return;
+
+      // Check if these operations are all in the same block.
+      llvm::SmallDenseSet<Block *, 4> blocks = {
+          fromDefiningOp->getBlock(), fromFreeingOp->getBlock(),
+          toDefiningOp->getBlock(), copy->getBlock()};
+      if (blocks.size() > 1)
+        return;
+
+      // Return true if `val` value has at least a user between `start` and
+      // `end` operations.
+      auto hasUsers = [](Value val, Operation *start, Operation *end) -> bool {
+        return llvm::any_of(val.getUsers(), [&](Operation *op) {
+          return op->isBeforeInBlock(end) && start->isBeforeInBlock(op);
+        });
+      };
+
+      // If `to` value has a single user between the defining operation of `to`
+      // and `copy` operation, it is illegal to remove the current copy
+      // operation. The same rule applies to `from` value if it has a user
+      // between copy operation and the deallocation operation of `to` value.
+      if (hasUsers(to, toDefiningOp, copy) ||
+          hasUsers(from, copy, fromFreeingOp))
+        return;
+
+      // TODO: This is a strong constraint. Using alias analysis would help for
+      // a better pruning of the cases.
+      if (hasMemoryEffectOp(copy, fromFreeingOp))
+        return;
+
+      // Copy operation can be safely removed and the uses of `to' value
+      // can be replaced with the `from` value.
+      to.replaceAllUsesWith(from);
+      copy->erase();
+      eraseList.push_back(toDefiningOp);
+      eraseList.push_back(fromFreeingOp);
+    });
+
+    for (Operation *op : eraseList)
+      op->erase();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// CopyRemovalPass construction
+//===----------------------------------------------------------------------===//
+
+std::unique_ptr<Pass> mlir::createLinalgCopyRemovalPass() {
+  return std::make_unique<CopyRemovalPass>();
+}

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

+// RUN: mlir-opt -linalg-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: %[[PERCENT10]] = alloc
+      // CHECK-NOT: linalg.copy(%[[PERCENT7]], %[[PERCENT10]])
+      // 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: %[[PERCENT6]] = alloc
+    // CHECK-NOT: linalg.copy(%[[PERCENT3]], %[[PERCENT6]])
+    // 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: %[[PERCENT6]] = alloc
+    // CHECK-NOT: linalg.copy(%[[PERCENT3]], %[[PERCENT6]])
+    // 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: () -> memref<5xf32>
+// CHECK: %[[ret:.*]] = alloc()
+// 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: %[[ret:.*]] = alloc()
+// CHECK: %{{.*}} = constant
+// CHECK: %[[DIM:.*]] = dim %[[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>
+  %temp = alloc() : memref<5xf32>
+  %res = 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: %[[ret:.*]] = alloc()
+// CHECK: %[[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: %[[ret:.*]] = call @make_allocation() : () -> memref<5xf32>
+  %ret = call @make_allocation() : () -> (memref<5xf32>)
+  // CHECK-NOT: %[[temp]] = alloc
+  // CHECK-NOT: linalg.copy(%[[ret]], %[[temp]])
+  // 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>
+}