[mlir][linalg][bufferize][NFC] Remove InSpaceSpec from bufferizesToMemoryWrite

[springerm]

Summary

Move functionality into a separate function isInplaceMemoryWrite.

Depends On D110982

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

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

return linalgOp.isInputTensor(&opOperand) ||
            linalgOp.isInitTensor(&opOperand);
   // All other cases are considered to bufferize to memory reads.
   // In particular, terminators are often the last use and need to be considered
   // as reads to return the proper value and avoid WAW clobbers.
   return true;
 }
 
 /// Return true if `opOperand` bufferizes to a memory write.
-/// If inPlaceSpec is different from InPlaceSpec::None, additionally require the
-/// write to match the inplace specification.
-static bool
-bufferizesToMemoryWrite(OpOperand &opOperand,
-                        InPlaceSpec inPlaceSpec = InPlaceSpec::None) {
+static bool bufferizesToMemoryWrite(OpOperand &opOperand) {
   // These terminators are not writes.
   if (isa<ReturnOp, linalg::YieldOp, scf::YieldOp>(opOperand.getOwner()))
     return false;
   // ExtractSliceOp alone doesn't bufferize to a memory write, one of its uses
   // may.
   if (isa<ExtractSliceOp>(opOperand.getOwner()))
     return false;
   // 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.
   if (auto callOp = dyn_cast<CallOpInterface>(opOperand.getOwner()))
     return true;
   // Unknown op that returns a tensor. The inplace analysis does not support
   // it. Conservatively return true.
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
     return true;
   OpResult opResult = getAliasingOpResult(opOperand);
-  // Supported op without a matching result for opOperand (e.g. ReturnOp).
-  // This does not bufferize to a write.
-  if (!opResult)
-    return false;
-  // If we have a matching OpResult, this is a write.
-  // Additionally allow to restrict to only inPlace write, if so specified.
-  return inPlaceSpec == InPlaceSpec::None ||
-         getInPlace(opResult) == inPlaceSpec;
+  // Only supported op with a matching result for opOperand bufferize to a
+  // write. E.g., ReturnOp does not bufferize to a write.
+  return static_cast<bool>(opResult);
 }
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
 
 /// Returns the relationship between the operand and the its corresponding
 /// OpResult that it may alias with.
 static BufferRelation bufferRelation(OpOperand &operand) {
   return TypeSwitch<Operation *, BufferRelation>(operand.getOwner())
       // ExtractSliceOp returns a subview of the original tensor.
       .Case([&](ExtractSliceOp op) { return BufferRelation::None; })
       // All other ops: Buffers are equivalent.
       .Default([&](Operation *op) { return BufferRelation::Equivalent; });
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
+/// Return true if opOperand has been decided to bufferize in-place.
+static bool isInplaceMemoryWrite(OpOperand &opOperand) {
+  // Ops that do not bufferize to a memory write, cannot be write in-place.
+  if (!bufferizesToMemoryWrite(opOperand))
+    return false;
+  OpResult opResult = getAliasingOpResult(opOperand);
+  return opResult && getInPlace(opResult) == InPlaceSpec::True;
+}
+
 namespace {
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 /// ExtractSliceOps have special behavior, they act as a level of indirection
 /// for bufferization. They don't create reads or writes themselves and analysis
 /// needs to look through their uses.
 /// ExtractSliceOp + InsertSliceOp have special joint behavior: they may

 
 /// Return true if the buffer to which `operand` would bufferize is equivalent
 /// to some buffer write.
 bool BufferizationAliasInfo::aliasesInPlaceWrite(Value value) const {
   LDBG("----Start aliasesInPlaceWrite\n");
   LDBG("-------for : " << printValueInfo(value) << '\n');
   for (Value v : getAliases(value)) {
     for (auto &use : v.getUses()) {
-      if (bufferizesToMemoryWrite(use, InPlaceSpec::True)) {
+      if (isInplaceMemoryWrite(use)) {
         LDBG("-----------wants to bufferize to inPlace write: "
              << printOperationInfo(use.getOwner()) << '\n');
         return true;
       }
     }
   }
   LDBG("----------->does not alias an inplace write\n");
   return false;
auto getAliasingReads = [&](DenseSet<OpOperand *> &res, Value root) {
     }
   };
 
   /// Helper function to iterate on aliases of `root` and capture the writes.
   auto getAliasingInplaceWrites = [&](DenseSet<OpOperand *> &res, Value root) {
     for (Value alias : getAliases(root)) {
       for (auto &use : alias.getUses()) {
         // Inplace write to a value that aliases root.
-        if (bufferizesToMemoryWrite(use, InPlaceSpec::True)) {
+        if (isInplaceMemoryWrite(use)) {
           LDBG("------------bufferizesToMemoryWrite: "
                << use.getOwner()->getName().getStringRef() << "\n");
           res.insert(&use);
         }
       }
     }
   };
 
bool BufferizationAliasInfo::existsInterleavedValueClobber(
   for (Value valueToClobber : {aliasingRead.get(), aliasingWrite.get()}) {
     auto leaderIt = equivalentInfo.findLeader(valueToClobber);
     for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
          ++mit) {
       Operation *candidateOp = mit->v.getDefiningOp();
       if (!candidateOp)
         continue;
       OpOperand *operand = getAliasingOpOperand(mit->v.cast<OpResult>());
+      // TODO: Should we check for isInplaceMemoryWrite instead?
       if (!operand || !bufferizesToMemoryWrite(*operand))
         continue;
       LDBG("---->clobbering candidate: " << printOperationInfo(candidateOp)
                                          << '\n');
       if (domInfo.properlyDominates(aliasingWriteOp, candidateOp) &&
           domInfo.properlyDominates(candidateOp, aliasingReadOp))
         return true;
     }