[Mlir-commits] [mlir] 3f1e827 - [mlir] Linalg : do not forward memrefs to outputs when do bufferization

[Eugene Zhulenev]

Author: Eugene Zhulenev
Date: 2021-04-21T16:39:06-07:00
New Revision: 3f1e827abd7fb893f7a33f467126d5d736ffa8d4

URL: https://github.com/llvm/llvm-project/commit/3f1e827abd7fb893f7a33f467126d5d736ffa8d4
DIFF: https://github.com/llvm/llvm-project/commit/3f1e827abd7fb893f7a33f467126d5d736ffa8d4.diff

LOG: [mlir] Linalg : do not forward memrefs to outputs when do bufferization

Example:
```
%0 = linalg.init_tensor : tensor<...>
%1 = linalg.generic ... outs(%0: tensor<...>)
%2 = linalg.generic ... outs(%0: tensor<...>)
```

Memref allocated as a result of `init_tensor` bufferization can be incorrectly overwritten by the second linalg.generic operation

Reviewed By: silvas

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
    mlir/test/Dialect/Linalg/bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
index 3ab86beb9367..892942eeaef0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
@@ -59,10 +59,6 @@ allocateBuffersForResults(Location loc, LinalgOp linalgOp, ValueRange outputs,
       continue;
     }
 
-    if (auto alloc = resultTensor.getDefiningOp<memref::AllocOp>()) {
-      resultBuffers.push_back(resultTensor);
-      continue;
-    }
     // Allocate buffers for statically-shaped results.
     if (memrefType.hasStaticShape()) {
       resultBuffers.push_back(b.create<memref::AllocOp>(loc, memrefType));

diff  --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir
index b9a4362f5e34..757b7a18fa9e 100644
--- a/mlir/test/Dialect/Linalg/bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/bufferize.mlir
@@ -45,8 +45,8 @@ func @basic(%arg0: tensor<4xf32>) -> tensor<4xf32> {
 // CHECK: #map = affine_map<(d0) -> (d0)>
 // CHECK-LABEL: func @init_tensor(
 // CHECK-SAME:      %[[IN:.*]]: tensor<?xf32>, %[[SIZE:.*]]: index)
-// CHECK:         %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
 // CHECK:         %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
+// CHECK:         %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
 // CHECK:         linalg.generic
 // CHECK-SAME:    ins(%[[MEMREF]] : memref<?xf32>)
 // CHECK-SAME:    outs(%[[OUT_BUF]] : memref<?xf32>) {