From 007e55133ec676f76dd63841ffe8b5eeff82478f Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Wed, 10 Nov 2021 18:34:26 +0900
Subject: [PATCH] [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
---
 .../BufferizableOpInterface.td                     | 29 ++++++++++++++++++++++
 .../ComprehensiveBufferize.cpp                     | 11 +-------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
index cff45a2..c2dd8b1 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
@@ -61,6 +61,35 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
          }]
       >,
       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
           bufferized in-place. This method will never be called on OpOperands
diff --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 27db0e9..8db4d74 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();
-- 
2.7.4