[Mlir-commits] [mlir] 007e551 - [mlir][linalg][bufferize] Add helper method isMemoryWrite to op interface

[Matthias Springer]

Author: Matthias Springer
Date: 2021-11-10T18:35:04+09:00
New Revision: 007e55133ec676f76dd63841ffe8b5eeff82478f

URL: https://github.com/llvm/llvm-project/commit/007e55133ec676f76dd63841ffe8b5eeff82478f
DIFF: https://github.com/llvm/llvm-project/commit/007e55133ec676f76dd63841ffe8b5eeff82478f.diff

LOG: [mlir][linalg][bufferize] Add helper method isMemoryWrite to op interface

This is in preparating for decoupling the SCF dialect from the analysis.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
index cff45a2fdbfa2..c2dd8b1321b4f 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
@@ -60,6 +60,35 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
           llvm_unreachable("bufferizesToMemoryWrite not implemented");
          }]
       >,
+      InterfaceMethod<
+          /*desc=*/[{
+            Return `true` if the given OpResult is a memory write. This is the
+            case if in the following cases:
+
+            * The corresponding aliasing OpOperand bufferizes to a memory write.
+            * Or: There is no corresponding aliasing OpOperand.
+
+            If the OpResult has multiple aliasing OpOperands, this method
+            returns `true` if at least one of them bufferizes to a memory write.
+          }],
+          /*retType=*/"bool",
+          /*methodName=*/"isMemoryWrite",
+          /*args=*/(ins "OpResult":$opResult),
+          /*methodBody=*/"",
+          /*defaultImplementation=*/[{
+            auto bufferizableOp =
+                cast<BufferizableOpInterface>($_op.getOperation());
+            SmallVector<OpOperand*> opOperands =
+              bufferizableOp.getAliasingOpOperand(opResult);
+            if (opOperands.empty())
+              return true;
+            return llvm::any_of(
+                opOperands,
+                [&](OpOperand *operand) {
+                  return bufferizableOp.bufferizesToMemoryWrite(*operand);
+                });
+          }]
+      >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpResult that aliases with a given OpOperand when

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 27db0e9afcf52..8db4d74453a8a 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -734,16 +734,7 @@ static Value findLastPrecedingWrite(Value value) {
           return true;
         if (isa<scf::IfOp>(op))
           return true;
-
-        SmallVector<OpOperand *> opOperands =
-            bufferizableOp.getAliasingOpOperand(value.cast<OpResult>());
-        assert(opOperands.size() <= 1 &&
-               "op with multiple aliasing OpOperands not expected");
-
-        if (opOperands.empty())
-          return true;
-
-        return bufferizesToMemoryWrite(*opOperands.front());
+        return bufferizableOp.isMemoryWrite(value.cast<OpResult>());
       });
   assert(result.size() == 1 && "expected exactly one result");
   return result.front();