[mlir][linalg][bufferize][NFC] Simplify getAliasingOpOperand signature

[springerm]

Summary

Depends On D110941

[nicolasvasilache]

I am not sure why this is desirable?
I expect the ternary state to be useful in the future: None should be the catchall case that makes copies of everything it needs to.

[springerm]

In D110982#3042103, @nicolasvasilache wrote:

I am not sure why this is desirable?
I expect the ternary state to be useful in the future: None should be the catchall case that makes copies of everything it needs to.

My reasoning was that ops that we don't know about do not have an aliasingOpOperand. Therefore, nullptr is a valid return value. (Judging just from the function name, without reading the comment.) We would call hasKnownBufferizationAliasBehavior first when the op may be an unknown one. (It is not always the case.) This would make getAliasingOpOperand symmetrical to getAliasingOpResult, which also does no longer return Optional<>.

[springerm]

rebase

[nicolasvasilache]

I am still on the fence about this one.
Note that in the other bigger CL, I point thatwe should prob. move to returning a list now.
Can we make that API change (which may also impact getAliasingOpResult) and see where we stand once it is in?

No need to support any new case, we would just bail on any list of size != 1 for now.

[springerm]

rebase

[springerm]

rebase

[springerm]

address comments

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

return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
         .Case([&](CallOpInterface op) { return OpResult(); })
         // Other ops.
         .Default([&](Operation *op) { return OpResult(); });
   // clang-format on
 }
 
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place.
-/// Return None if the owner of `opOperand` does not have known
-/// bufferization aliasing behavior, which indicates that the op must allocate
-/// all of its tensor results.
 /// TODO: in the future this may need to evolve towards a list of OpOperand*.
-static Optional<OpOperand *> getAliasingOpOperand(OpResult result) {
+static OpOperand *getAliasingOpOperand(OpResult result) {
+  // Unknown ops are handled conservatively and never bufferize in-place.
   if (!hasKnownBufferizationAliasingBehavior(result.getDefiningOp()))
-    return None;
+    return nullptr;
   return TypeSwitch<Operation *, OpOperand *>(result.getDefiningOp())
       .Case([&](tensor::CastOp op) { return &op->getOpOperand(0); })
       .Case([&](ConstantOp op) { return &op->getOpOperand(0); })
       .Case([&](ExtractSliceOp op) { return &op->getOpOperand(0); })
       // In the case of scf::ForOp, this currently assumes the iter_args / yield
       // are 1-1. This may fail and is verified at the end.
       // TODO: update this.
       .Case([&](scf::ForOp op) {

 /// opResult when:
 ///   1. R is not known to properly dominate W (i.e. the effects of the write
 ///      may be visible from R).
 ///   2. one cannot find an intermediate clobbering write `C` to W, such that
 ///      C interleaved between W and R (i.e. W -> C -> R where -> denotes
 ///      dominance).
 bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
     OpResult result, const DominanceInfo &domInfo) const {
-  Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(result);
-  if (!maybeAliasingOperand)
+  OpOperand *aliasingOperand = getAliasingOpOperand(result);
+  if (!aliasingOperand)
     return false;
 
   Operation *opToBufferize = result.getDefiningOp();
   Value opResult = result;
-  Value opOperand = (*maybeAliasingOperand)->get();
+  Value opOperand = aliasingOperand->get();
 
   LDBG("----Start wouldCreateReadAfterWriteInterference\n");
   LDBG("--------consider all aliases to root read: "
        << printValueInfo(opOperand) << "\n");
   LDBG("--------consider all aliases to root write: "
        << printValueInfo(opResult) << "\n");
 
   /// Helper function to iterate on aliases of `root` and capture the reads.
bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
   // interferences as if it were bufferized inplace: add `maybeAliasingOperand`
   // if it is a write. This handles the case:
   //
   // ```
   //  %0 = op_to_bufferize_maybe_inplace(%1)
   //  %2 = some_alias(%1)
   //  read(%2)
   // ```
-  if (bufferizesToMemoryWrite(**maybeAliasingOperand))
-    usesWrite.insert(*maybeAliasingOperand);
+  if (bufferizesToMemoryWrite(*aliasingOperand))
+    usesWrite.insert(aliasingOperand);
   if (wouldCreateReadAfterWriteInterference(opToBufferize, usesRead, usesWrite,
                                             domInfo))
     return true;
 
   // Check if we can find any interference between writes to
   // aliases[`opOperand`] and reads to aliases[`opResult`]. This handles the
   // case:
   //
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;
-      auto maybeAliasingOperand = getAliasingOpOperand(mit->v.cast<OpResult>());
-      if (!maybeAliasingOperand || !*maybeAliasingOperand ||
-          !bufferizesToMemoryWrite(**maybeAliasingOperand))
+      OpOperand *operand = getAliasingOpOperand(mit->v.cast<OpResult>());
+      if (!operand || !bufferizesToMemoryWrite(*operand))
         continue;
       LDBG("---->clobbering candidate: " << printOperationInfo(candidateOp)
                                          << '\n');
       if (domInfo.properlyDominates(aliasingWriteOp, candidateOp) &&
           domInfo.properlyDominates(candidateOp, aliasingReadOp))
         return true;
     }
   }

 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 static Value getResultBuffer(OpBuilder &b, OpResult result,
                              const BlockAndValueMapping &bvm,
                              BufferizationAliasInfo &aliasInfo,
                              bool skipCopy = false) {
   OpBuilder::InsertionGuard guard(b);
   Operation *op = result.getOwner();
-  Optional<OpOperand *> maybeOperand = getAliasingOpOperand(result);
+  OpOperand *maybeOperand = getAliasingOpOperand(result);
   assert(maybeOperand && "corresponding OpOperand not found");
-  Value operand = (*maybeOperand)->get();
+  Value operand = maybeOperand->get();
   Value operandBuffer = lookup(bvm, operand);
   assert(operandBuffer && "operand buffer not found");
 
   // If bufferizing out-of-place, allocate a new buffer.
   if (getInPlace(result) != InPlaceSpec::True) {
     Location loc = op->getLoc();
     // Allocate the result buffer.
     Value resultBuffer =