From 934e2f695e18033fdf740ab45bb5b0dfdf453c19 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Wed, 15 Sep 2021 17:26:29 +0900
Subject: [PATCH] [mlir][linalg] ComprehensiveBufferize: Do not copy
 InitTensorOp results

E.g.:

```
%2 = memref.alloc() {alignment = 128 : i64} : memref<256x256xf32>
%3 = memref.alloc() {alignment = 128 : i64} : memref<256x256xf32>

// ... (%3 is not written to)

linalg.copy(%3, %2) : memref<256x256xf32>, memref<256x256xf32>
vector.transfer_write %11, %2[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, memref<256x256xf32>
```

Avoid copies of %3 if %3 came directly from an InitTensorOp.

Differential Revision: https://reviews.llvm.org/D109742
---
 mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
index edf74b382d2d..b165ec752491 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
@@ -2190,7 +2190,8 @@ static LogicalResult bufferize(OpBuilder &b, VectorTransferOpInterface op,
     newInputBuffer = createNewAllocDeallocPairForShapedValue(
         b, loc, writeOp.source(), aliasInfo);
     Value v = lookup(bvm, writeOp.source());
-    b.create<CopyOp>(loc, v, newInputBuffer);
+    if (!isInitTensorOp(writeOp.source()))
+      b.create<CopyOp>(loc, v, newInputBuffer);
   } else {
     // InPlace write will result in memref.tensor_load(x) which must
     // canonicalize away with one of it uses.