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

[ezhulenev]

Summary

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

[silvas]

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

63

you might be able to just delete this code block. I think the code below can handle it.

[ezhulenev]

Do not forward memrefs to outputs when do bufferization

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

for (auto en : llvm::enumerate(linalgOp->getResultTypes())) {
     Value resultTensor = outputs[resultIndex];
 
     // Clone output buffers whose value is actually used.
     if (linalgOp.payloadUsesValueFromOutputOperandIndex(resultIndex)) {
       resultBuffers.push_back(cloneMemref(loc, resultTensor, b));
       continue;
     }
 
+    // Clone all explicitly allocated output buffers.
     if (auto alloc = resultTensor.getDefiningOp<memref::AllocOp>()) {
-      resultBuffers.push_back(resultTensor);
+      resultBuffers.push_back(cloneMemref(loc, resultTensor, b));
       continue;
     }
+
     // Allocate buffers for statically-shaped results.
     if (memrefType.hasStaticShape()) {
       resultBuffers.push_back(b.create<memref::AllocOp>(loc, memrefType));
       continue;
     }
 
     resultBuffers.push_back(b.create<memref::AllocOp>(
         loc, memrefType, getDynOperands(loc, resultTensor, b)));

mlir/test/Dialect/Linalg/bufferize.mlir

 
 #map0 = affine_map<(d0) -> (d0)>
 
 // Same as above but with linalg.init_tensor op.
 
 // 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:         %[[OUT_BUF0:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
 // CHECK:         %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
+// CHECK:         %[[OUT_BUF1:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+// CHECK:         linalg.copy(%[[OUT_BUF0]], %[[OUT_BUF1]])
 // CHECK:         linalg.generic
 // CHECK-SAME:    ins(%[[MEMREF]] : memref<?xf32>)
-// CHECK-SAME:    outs(%[[OUT_BUF]] : memref<?xf32>) {
+// CHECK-SAME:    outs(%[[OUT_BUF1]] : memref<?xf32>) {
 func @init_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> {
   %init = linalg.init_tensor [%size] : tensor<?xf32>
   %0 = linalg.generic {
     indexing_maps = [#map0, #map0],
     iterator_types = ["parallel"]
   } ins(%in : tensor<?xf32>)
     outs(%init : tensor<?xf32>) {
     ^bb0(%gen_arg1: f32, %out: f32):