[Mlir-commits] [mlir] 00a5ae1 - [mlir][bufferization] Better propagation of bufferizesToMemoryWrite through regions

[Matthias Springer]

Author: Matthias Springer
Date: 2023-02-13T16:52:26+01:00
New Revision: 00a5ae12448e6f33e40c38afd640eafdb7dab212

URL: https://github.com/llvm/llvm-project/commit/00a5ae12448e6f33e40c38afd640eafdb7dab212
DIFF: https://github.com/llvm/llvm-project/commit/00a5ae12448e6f33e40c38afd640eafdb7dab212.diff

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

`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.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
    mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
    mlir/test/Dialect/SCF/one-shot-bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
index 6e796c8aebb0a..4f33425169a51 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
@@ -120,8 +120,7 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
              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

diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 3e882ba5ed8b0..f25a8eb0cf6c9 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -896,7 +896,8 @@ bool bufferization::detail::defaultResultBufferizesToMemoryWrite(
     if (!state
              .findValueInReverseUseDefChain(alias.opOperand->get(),
                                             isMemoryWriteInsideOp,
-                                            /*followEquivalentOnly=*/false)
+                                            /*followEquivalentOnly=*/false,
+                                            /*alwaysIncludeLeaves=*/false)
              .empty())
       return true;
   }

diff  --git a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
index 131f3066ec2c3..e37fe73d01170 100644
--- a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
@@ -1,4 +1,5 @@
 // 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 
diff erent 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
@@ -950,3 +951,31 @@ func.func @regression_cast_in_loop() -> 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
+}