[Mlir-commits] [mlir] 2930642 - [mlir] Add MemoryEffects::Allocate to memref::CloneOp

[llvmlistbot at llvm.org]

Author: Butygin
Date: 2021-06-29T13:37:32+03:00
New Revision: 293064222a013055cda9240647110f5bcf8e1f31

URL: https://github.com/llvm/llvm-project/commit/293064222a013055cda9240647110f5bcf8e1f31
DIFF: https://github.com/llvm/llvm-project/commit/293064222a013055cda9240647110f5bcf8e1f31.diff

LOG: [mlir] Add MemoryEffects::Allocate to memref::CloneOp

Without it BufferDeallocationPass process only CloneOps created during pass itself and ignore all CloneOps that were already present in IR.

For our specific usecase:

```
func @dealloc_existing_clones(%arg0: memref<?x?xf64>, %arg1: memref<?x?xf64>) -> memref<?x?xf64> {
  return %arg0 : memref<?x?xf64>
}
```

Input arguments will be freed immediately after return from function and we want to prolong lifetime for the returned argument.
To achieve this we explicitly add clones to all input memrefs and expect that BufferDeallocationPass will add correct deallocs to them (unnessesary clone+dealloc pairs will be canonicalized away later).

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

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/test/Transforms/buffer-deallocation.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 6f358d834beed..c9df5fc4678ad 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -514,6 +514,8 @@ void CloneOp::getEffects(
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), output(),
                        SideEffects::DefaultResource::get());
+  effects.emplace_back(MemoryEffects::Allocate::get(), output(),
+                       SideEffects::DefaultResource::get());
 }
 
 namespace {

diff  --git a/mlir/test/Transforms/buffer-deallocation.mlir b/mlir/test/Transforms/buffer-deallocation.mlir
index 77945113e1647..7bc335a93ae50 100644
--- a/mlir/test/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Transforms/buffer-deallocation.mlir
@@ -1207,3 +1207,18 @@ func @noRegionBranchOpInterface() {
   }) : () -> (i32)
   "test.terminator"() : () -> ()
 }
+
+// -----
+
+// CHECK-LABEL: func @dealloc_existing_clones
+// CHECK: (%[[ARG0:.*]]: memref<?x?xf64>, %[[ARG1:.*]]: memref<?x?xf64>)
+// CHECK: %[[RES0:.*]] = memref.clone %[[ARG0]]
+// CHECK: %[[RES1:.*]] = memref.clone %[[ARG1]]
+// CHECK-NOT: memref.dealloc %[[RES0]]
+// CHECK: memref.dealloc %[[RES1]]
+// CHECK: return %[[RES0]]
+func @dealloc_existing_clones(%arg0: memref<?x?xf64>, %arg1: memref<?x?xf64>) -> memref<?x?xf64> {
+  %0 = memref.clone %arg0 : memref<?x?xf64> to memref<?x?xf64>
+  %1 = memref.clone %arg1 : memref<?x?xf64> to memref<?x?xf64>
+  return %0 : memref<?x?xf64>
+}