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

[springerm]

Summary

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

Depends On D113340

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td

let methods = [
         /*args=*/(ins "OpOperand &":$opOperand),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           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
           bufferized in-place. This method will never be called on OpOperands
           that do not have a tensor type.
         }],
         /*retType=*/"OpResult",
         /*methodName=*/"getAliasingOpResult",
         /*args=*/(ins "OpOperand &":$opOperand),
         /*methodBody=*/"",

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

SetVector<Value> result =
         Operation *op = value.getDefiningOp();
         if (!op)
           return true;
         auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
         if (!bufferizableOp)
           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();
 }
 
 /// Return true if `value` is originating from an ExtractSliceOp that matches
 /// the given InsertSliceOp.
 static bool hasMatchingExtractSliceOp(const BufferizationAliasInfo &aliasInfo,