[Mlir-commits] [mlir] 2fd0ea9 - [mlir][linalg][bufferize] CallOps do not bufferize to memory writes

[Matthias Springer]

Author: Matthias Springer
Date: 2021-12-01T18:47:28+09:00
New Revision: 2fd0ea960c8a3e3a05fc0c320533b486edfa01f9

URL: https://github.com/llvm/llvm-project/commit/2fd0ea960c8a3e3a05fc0c320533b486edfa01f9
DIFF: https://github.com/llvm/llvm-project/commit/2fd0ea960c8a3e3a05fc0c320533b486edfa01f9.diff

LOG: [mlir][linalg][bufferize] CallOps do not bufferize to memory writes

However, since CallOps have no aliasing OpResults, their OpOperands always bufferize out-of-place.

This change removes `bufferizesToMemoryWrite` from `CallOpInterface`. This method was called, but its return value did not matter.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
index ca0454549bcab..7d845da3ce53e 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
@@ -42,9 +42,16 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given OpOperand bufferizes to a memory write.
+
           This method will never be called on OpOperands that do not have a
           tensor type.
 
+          This method will never be called on OpOperands that do not have an
+          aliasing OpResult. Intuitively, it does not make sense for an
+          OpOperand to bufferize to a memory write without returning an aliasing
+          tensor, because the write would have no visible effect outside of the
+          op.
+
           Note: It is always safe to consider an OpOperand as a memory write,
           even if it does actually not write; however, this can introduce
           unnecessary out-of-place bufferization decisions. The analysis of
@@ -57,6 +64,8 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
+          // Does not have to be implemented for OpOperands that do not have an
+          // aliasing OpResult.
           llvm_unreachable("bufferizesToMemoryWrite not implemented");
          }]
       >,

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index eb37a544de60e..6cbc308281e22 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -239,11 +239,15 @@ static std::string printValueInfo(Value value, bool prefix) {
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
                                  const BufferizationAliasInfo &aliasInfo) {
-  // Ops that do not bufferize to a memory write, cannot be write in-place.
+  // OpOperands without an aliasing OpResult do not write.
+  OpResult opResult = getAliasingOpResult(opOperand);
+  if (!opResult)
+    return false;
+  // OpOperands that do not bufferize to a memory write do not write in-place.
   if (!bufferizesToMemoryWrite(opOperand))
     return false;
-  OpResult opResult = getAliasingOpResult(opOperand);
-  return opResult && aliasInfo.isInPlace(opResult);
+  // Check current bufferization decisions.
+  return aliasInfo.isInPlace(opResult);
 }
 
 /// Return true if, under current bufferization decisions, the buffer of `value`

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
index c98cc1de60af1..b55395b69e46d 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
@@ -447,14 +447,6 @@ struct CallOpInterface
     return true;
   }
 
-  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
-    // CallOpInterface alone doesn't bufferize to a memory write, one of the
-    // uses of the matching bbArg may. It is the responsibility of the caller to
-    // inspect bbArgs. In the absence of a BufferizationAliasInfo, we need to be
-    // conservative.
-    return true;
-  }
-
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
     // TODO: Can we do better?

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp
index 811824202091c..2b2ccb3c4bb9a 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp
@@ -44,6 +44,17 @@ struct ExecuteRegionOpInterface
     return true;
   }
 
+  // TODO: For better bufferization results, this could return `true` only if
+  // there is a memory write in the region.
+  bool isMemoryWrite(Operation *op, OpResult opResult) const {
+    // Similar to scf.if, results of this op are always considered memory writes
+    // in the analysis. This is a useful pattern for all ops that have tensor
+    // OpResults but no tensor OpOperands. By default, `isMemoryWrite` is
+    // implemented in terms of `bufferizesToMemoryWrite`, which does not work on
+    // ops without OpOperands.
+    return true;
+  }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     // TODO: Add bufferization support when needed. scf.execute_region should be

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
index 4a8d2e48162d0..616d84545e5c2 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
@@ -1470,3 +1470,25 @@ func @scf_if_out_of_place3(%t1: tensor<?xf32> {linalg.inplaceable = true},
   return %r, %v2 : tensor<?xf32>, vector<10xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @some_use
+func @some_use(%A : tensor<?xf32> {linalg.inplaceable = true},
+               %v : vector<5xf32>) -> (tensor<?xf32>) {
+  %idx = arith.constant 0 : index
+  //      CHECK: vector.transfer_write
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+  %0 = vector.transfer_write %v, %A[%idx] : vector<5xf32>, tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+
+// CHECK-LABEL: func @main_func
+func @main_func(%A : tensor<?xf32> {linalg.inplaceable = true},
+                %v : vector<5xf32>) -> (tensor<?xf32>) {
+  // Function calls always bufferize out-of-place at the moment.
+  //      CHECK: call
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]
+  %0 = call @some_use(%A, %v) : (tensor<?xf32>, vector<5xf32>) -> (tensor<?xf32>)
+  return %0 : tensor<?xf32>
+}