[mlir][linalg][bufferize] Fix copy elision in `getResultBuffer`

[springerm]

Summary

A buffer copy may not be elided if the to-be-bufferized op is reading the data.

Depends On D116453

[nicolasvasilache]

Based on our offline discussion, seems like this should be an assert ?

[springerm]

update

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

if (!aliasInfo.isInPlace(result)) {
     // not write (skipping ops that merely create aliases). E.g., InitTensorOp.
     // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
     // use-def chain, it returns that value, regardless of whether it is a
     // memory write or not.
     Value lastWrite = findLastPrecedingWrite(operand);
     if (auto bufferizableOp = options.dynCastBufferizableOp(lastWrite))
       if (!bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(), *this))
         skipCopy = true;
-    // Do not copy if the copied data is never read.
-    if (!isValueRead(result))
+    // Do not copy if the copied data is never read. (Neither by this op nor by
+    // any following op.)
+    if (!bufferizesToMemoryRead(*opOperand) && !isValueRead(result))
       skipCopy = true;
     // Do not copy if this op does not read the data, but writes it.
     if (bufferizesToMemoryWrite(*opOperand) &&
         !bufferizesToMemoryRead(*opOperand))
       skipCopy = true;
     if (!skipCopy) {
       // The copy happens right before the op that is bufferized.
       rewriter.setInsertionPoint(op);

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

%r = linalg.generic {
     outs(%t3 : tensor<?x?xf32>) {
       ^bb0(%arg0 : f32, %arg1 : f32, %arg2 : f32) :
         %add = arith.addf %arg0, %arg1 : f32
         linalg.yield %add : f32
     } -> tensor<?x?xf32>
   return %r : tensor<?x?xf32>
 }
 
+// -----
+
+#accesses = [
+  affine_map<(i) -> (i)>
+]
+#trait = {
+  indexing_maps = #accesses,
+  iterator_types = ["parallel"]
+}
+
+// CHECK-LABEL: func @op_is_reading_but_following_ops_are_not
+//  CHECK-SAME:     %[[t0:.*]]: memref<?xf32
+func @op_is_reading_but_following_ops_are_not(
+    %t0 : tensor<?xf32> {linalg.inplaceable = false},
+    %cst : f32)
+  -> tensor<?xf32>
+{
+  // Make sure that a copy is inserted here.
+  // CHECK: %[[ALLOC:.*]] = memref.alloc
+  // CHECK: linalg.copy(%[[t0]], %[[ALLOC]])
+  // CHECK: linalg.generic {{.*}} outs(%[[ALLOC]] : memref
+  %r0 =linalg.generic #trait outs (%t0 : tensor<?xf32>) {
+      ^bb(%0: f32) :
+        %a = arith.addf %cst, %0 : f32
+        linalg.yield %a : f32
+    } -> (tensor<?xf32>)
+
+  // CHECK: linalg.generic {{.*}} outs(%[[ALLOC]] : memref
+  %r1 = linalg.generic #trait outs (%r0 : tensor<?xf32>) {
+      ^bb(%0: f32) :
+        linalg.yield %cst : f32
+    } -> (tensor<?xf32>)
+
+  // CHECK: return %[[ALLOC]]
+  return %r1 : tensor<?xf32>
+}