[MLIR] Avoid adding extra memref::AllocOp in bufferize pass.

[wenzhicui]

Summary

Avoid adding extra memref::AllocOp in bufferize pass.

[mehdi_amini]

mlir/test/Dialect/Linalg/bufferize.mlir

49

Should there be something like CHECK-NOT: memref.alloc here? If I understand correctly you're avoiding generating a second alloc? (assuming this test case had two generated?)
Otherwise the test does not really demonstrate what I understand from the description.

[wenzhicui]

Added CHECK-NOT as suggested by the reviewer.

[wenzhicui]

mlir/test/Dialect/Linalg/bufferize.mlir

49

Thanks for the suggestion. Added the CHECK. That was the intention of this patch.

[mehdi_amini]

mlir/test/Dialect/Linalg/bufferize.mlir

49

I don't think it is satisfying actually: the test seems to be passing without the code change. Can you look into this?

[mehdi_amini]

mlir/test/Dialect/Linalg/bufferize.mlir

49

I see this IR in the test output before the code change:

18:  func @init_tensor(%arg0: tensor<?xf32>, %arg1: index) -> tensor<?xf32> { 
19:  %0 = memref.buffer_cast %arg0 : memref<?xf32> 
20:  %1 = memref.alloc(%arg1) : memref<?xf32> 
21:  linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%0 : memref<?xf32>) outs(%1 : memref<?xf32>) {

[wenzhicui]

mlir/test/Dialect/Linalg/bufferize.mlir

49

I looked into this test and I realized that the CSE/Canonicalize pass is removing the extra memref.alloc ops because those ops are not used by any other ops. Do you think this patch is still needed since we can always have a CSE/Canonicalize pass to remove the extra alloc ops?

[nicolasvasilache]

This should be reviewed by @herhut @pifon2a and DFKI friends

[herhut]

This change is not correct in general I believe.

The existing alloc is the one for the initial value (e.g. init_tensor). However, there is no guarantee that this buffer is only passed to one linalg.generic operation. It might actually be reused. Hence, we have to create a copy if the value is required (see couple lines below the change). Otherwise, we allocate a fresh result buffer and only use the shape of the provided buffer.

If it then later turns out that the passed buffer indeed is not used anywhere else, we clean this up.

[wenzhicui]

Stephan/Mehdi, thanks for the review. Will abandon this change since it is incorrect.

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

for (auto en : llvm::enumerate(linalgOp->getResultTypes())) {
     if (tensorType == nullptr) {
       linalgOp.emitOpError()
           << "tensor to buffer conversion expects ranked tensor results";
       return failure();
     }
     auto tensorShape = tensorType.getShape();
     auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
     Value resultTensor = outputs[resultIndex];
+    if (isa<memref::AllocOp>(resultTensor.getDefiningOp())) {
+      resultBuffers.push_back(resultTensor);
+      continue;
+    }
 
     // Clone output buffers whose value is actually used.
     OpOperand *tiedOpOperand = linalgOp.getOutputOperand(resultIndex);
     if (linalgOp.payloadUsesValueFromOperand(tiedOpOperand)) {
       resultBuffers.push_back(cloneMemref(loc, resultTensor, b));
       continue;
     }
 

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:         %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
-// CHECK:         %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+// CHECK-DAG:     %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
+// CHECK-DAG:     %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+// CHECK-NOT:     memref.alloc
 // CHECK:         linalg.generic
 // CHECK-SAME:    ins(%[[MEMREF]] : memref<?xf32>)
 // CHECK-SAME:    outs(%[[OUT_BUF]] : 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"]