[mlir][bufferization] Better propagation of bufferizesToMemoryWrite through regions

[springerm]

Summary

bufferizesToMemoryWrite(OpResult) looks for OpOperands that bufferize to memory writes inside the region of the defining op (if it has one). Currently, if the reverse use-def chain stops at any value inside of the region, the OpResult is considered to bufferize to a memory write.

It is always safe to have false positives among bufferizesToMemoryWrite, so the previous implementation is also correct. However, it can lead to additional buffer copies.

Depends On: D141684

[stevenvar]

LGTM, but might need someone more knowledgeable to confirm that the changes are safe.

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

659

What was the reasoning before for returning the leaves (even when they did not satisfy the condition) when using this function?

After this change, are there any uses for /*alwaysIncludeLeaves=*/true left?

[springerm]

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

659

There is no particular reason. We just never thought about doing this optimization. The reason why you were suddenly seeing copies was due to another change that fixed conflict detection in loops. (There should have been a copy there before.)

We may not have any uses for /*alwaysIncludeLeaves=*/true at the moment, but I have a few local changes that I haven't sent out that may need it again.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
 
           2. At least one aliasing OpOperand bufferizes to a memory write.
 
              Example: `tensor.insert %f into %t[...] : tensor<?xf32>`
              The destination OpOperand bufferizes to a memory write, so the
              result also bufferizes to a memory write.
 
           3. At least one aliasing OpOperand's value is defined inside the
-             defining op of the given OpResult and it is a memory write or the
-             reverse SSA use-def chain ends in the defining op.
+             defining op of the given OpResult and it is a memory write.
 
              According to this rule, an aliasing OpOperand value that is defined
              inside this op and is bufferizing to a memory write makes the given
              OpResult bufferize to a memory write.
 
              Example:
              ```
              %r = scf.if ... -> tensor<?xf32> {

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

#endif // NDEBUG
   if (failed(memrefType))
     return failure();
   ensureToMemrefOpIsValid(value, *memrefType);
   return rewriter
       .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
 /// Return the buffer type for a given Value (tensor) after bufferization.
 FailureOr<BaseMemRefType>
 bufferization::getBufferType(Value value, const BufferizationOptions &options) {
   DenseMap<Value, BaseMemRefType> fixedTypes;
   return getBufferType(value, options, fixedTypes);
 }
 
 /// Return the buffer type for a given Value (tensor) after bufferization.
 FailureOr<BaseMemRefType> bufferization::getBufferType(
auto isMemoryWriteInsideOp = [&](Value v) {
     if (!opResult.getDefiningOp()->isAncestor(op))
       return false;
     return state.bufferizesToMemoryWrite(v);
   };
   for (AliasingOpOperand alias : opOperands) {
     if (!state
              .findValueInReverseUseDefChain(alias.opOperand->get(),
                                             isMemoryWriteInsideOp,
-                                            /*followEquivalentOnly=*/false)
+                                            /*followEquivalentOnly=*/false,
+                                            /*alwaysIncludeLeaves=*/false)
              .empty())
       return true;
   }
   return false;
 }
 
 // Compute the AliasingOpOperandList for a given OpResult based on
 // getAliasingOpResults.

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -drop-equivalent-buffer-results -buffer-deallocation -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -drop-equivalent-buffer-results -split-input-file | FileCheck %s --check-prefix=CHECK-NO-DEALLOC-PASS
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -buffer-deallocation -split-input-file -o /dev/null
func.func @regression_cast_in_loop() -> tensor<2xindex> {
   ^bb0(%arg0: tensor<2xindex>):
     %cast = tensor.cast %0 : tensor<2xindex> to tensor<?xindex>
     %inserted = tensor.insert %c0 into %cast[%c0] : tensor<?xindex>
     %cast_0 = tensor.cast %inserted : tensor<?xindex> to tensor<2xindex>
     scf.yield %cast_0 : tensor<2xindex>
   }
   return %1 : tensor<2xindex>
 }
+
+// -----
+
+// This test does not compute anything meaningful but it tests that
+// bufferizesToMemoryWrite is correctly propagated through regions.
+
+// CHECK-NO-DEALLOC-PASS-LABEL: func @elide_copy_of_non_writing_scf_if(
+func.func @elide_copy_of_non_writing_scf_if(%c: i1, %p1: index, %p2: index, %f: f32)
+  -> (tensor<10xf32>, f32)
+{
+  %r = scf.if %c -> tensor<10xf32> {
+    // CHECK-NO-DEALLOC-PASS: memref.alloc
+    %t1 = bufferization.alloc_tensor() : tensor<10xf32>
+    scf.yield %t1 : tensor<10xf32>
+  } else {
+    // CHECK-NO-DEALLOC-PASS: memref.alloc
+    %t2 = bufferization.alloc_tensor() : tensor<10xf32>
+    scf.yield %t2 : tensor<10xf32>
+  }
+
+  // No copy should be inserted because %r does not bufferize to a memory write.
+  // I.e., %r does not have defined contents and the copy can be elided.
+  // CHECK-NO-DEALLOC-PASS-NOT: memref.alloc
+  // CHECK-NO-DEALLOC-PASS-NOT: memref.copy
+  %r2 = tensor.insert %f into %r[%p1] : tensor<10xf32>
+  %r3 = tensor.extract %r[%p2] : tensor<10xf32>
+  return %r2, %r3 : tensor<10xf32>, f32
+}