[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())
         // make a proper determination by itself and needs to be conservative.
         .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) {
+/// in place. Note that multiple OpOperands can may potentially alias with an
+/// OpResult. E.g.: std.select in the future.
+static SmallVector<OpOperand *> getAliasingOpOperand(OpResult result) {
+  SmallVector<OpOperand *> r;
+  // Unknown ops are handled conservatively and never bufferize in-place.
   if (!hasKnownBufferizationAliasingBehavior(result.getDefiningOp()))
-    return None;
-  return TypeSwitch<Operation *, OpOperand *>(result.getDefiningOp())
-      .Case([&](tensor::CastOp op) { return &op->getOpOperand(0); })
-      .Case([&](ConstantOp op) { return nullptr; })
-      .Case([&](ExtractSliceOp op) { return &op->getOpOperand(0); })
+    return SmallVector<OpOperand *>();
+  TypeSwitch<Operation *>(result.getDefiningOp())
+      .Case([&](tensor::CastOp op) { r.push_back(&op->getOpOperand(0)); })
+      .Case([&](ExtractSliceOp op) { r.push_back(&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) {
-        return &op.getIterOpOperands()[result.getResultNumber()];
+        r.push_back(&op.getIterOpOperands()[result.getResultNumber()]);
       })
-      .Case([&](InitTensorOp op) { return nullptr; })
-      .Case([&](InsertSliceOp op) { return &op->getOpOperand(1); })
+      .Case([&](InsertSliceOp op) { r.push_back(&op->getOpOperand(1)); })
       .Case([&](LinalgOp op) {
-        return op.getOutputTensorOperands()[result.getResultNumber()];
+        r.push_back(op.getOutputTensorOperands()[result.getResultNumber()]);
       })
       .Case([&](TiledLoopOp op) {
         // TODO: TiledLoopOp helper method to avoid leaking impl details.
-        return &op->getOpOperand(op.getNumControlOperands() +
-                                 op.getNumInputs() + result.getResultNumber());
+        r.push_back(&op->getOpOperand(op.getNumControlOperands() +
+                                      op.getNumInputs() +
+                                      result.getResultNumber()));
       })
-      .Case([&](vector::TransferWriteOp op) { return &op->getOpOperand(1); })
-      .Case([&](CallOpInterface op) { return nullptr; })
+      .Case([&](vector::TransferWriteOp op) {
+        r.push_back(&op->getOpOperand(1));
+      })
+      .Case<ConstantOp, CallOpInterface, InitTensorOp>([&](auto op) {})
       .Default([&](Operation *op) {
         op->dump();
         llvm_unreachable("unexpected defining op");
-        return nullptr;
       });
+  return r;
 }
 
 /// If the an ExtractSliceOp is bufferized in-place, the source operand will
 /// alias with the result.
 static OpResult getAliasingOpResult(ExtractSliceOp op, OpOperand &opOperand) {
   if (op.source() == opOperand.get())
     return op->getResult(0);
   return OpResult();

 /// 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(
     OpOperand &operand, OpResult result, const DominanceInfo &domInfo) const {
-  assert(getAliasingOpOperand(result) == &operand &&
+#ifndef NDEBUG
+  SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
+  assert(llvm::find(opOperands, &operand) != opOperands.end() &&
          "operand and result do not match");
+#endif // NDEBUG
 
   Operation *opToBufferize = result.getDefiningOp();
   Value opResult = result;
   Value opOperand = operand.get();
 
   LDBG("----Start wouldCreateReadAfterWriteInterference\n");
   LDBG("--------consider all aliases to root read: "
        << printValueInfo(opOperand) << "\n");
#endif // NDEBUG
   return wouldCreateReadAfterWriteInterference(opToBufferize, usesRead,
                                                usesWrite, domInfo);
 }
 
 /// Return true if bufferizing `opOperand` inplace with `opResult` would create
 /// a write to a non-writable buffer.
 bool BufferizationAliasInfo::wouldCreateWriteToNonWritableBuffer(
     OpOperand &opOperand, OpResult opResult) const {
-  assert(getAliasingOpOperand(opResult) == &opOperand &&
+#ifndef NDEBUG
+  SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
+  assert(llvm::find(opOperands, &opOperand) != opOperands.end() &&
          "operand and result do not match");
+#endif // NDEBUG
 
   // Certain buffers are not writeable:
   //   1. A function bbArg that is not inplaceable or
   //   2. A constant op.
   assert(!aliasesNonWritableBuffer(opResult) &&
          "expected that opResult does not alias non-writable buffer");
   bool nonWritable = aliasesNonWritableBuffer(opOperand.get());
   if (!nonWritable)
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))
+      SmallVector<OpOperand *> operands =
+          getAliasingOpOperand(mit->v.cast<OpResult>());
+      assert(operands.size() <= 1 && "more than 1 OpOperand not supported yet");
+      if (operands.empty() || !bufferizesToMemoryWrite(*operands.front()))
         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);
-  assert(maybeOperand && "corresponding OpOperand not found");
-  Value operand = (*maybeOperand)->get();
+  SmallVector<OpOperand *> aliasingOperands = getAliasingOpOperand(result);
+  // TODO: Support multiple OpOperands.
+  assert(aliasingOperands.size() == 1 &&
+         "more than 1 OpOperand not supported yet");
+  Value operand = aliasingOperands.front()->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 =

 
 /// Determine if `operand` can be bufferized in-place with `result`. If so, set
 /// InPlaceSpec::True on the result. Otherwise, set InPlaceSpec::False on the
 /// result.
 static LogicalResult
 bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
                                 BufferizationAliasInfo &aliasInfo,
                                 const DominanceInfo &domInfo) {
-  assert(getAliasingOpOperand(result) == &operand &&
+#ifndef NDEBUG
+  SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
+  assert(llvm::find(opOperands, &operand) != opOperands.end() &&
          "operand and result do not match");
+#endif // NDEBUG
 
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
                                    << operand.getOperandNumber() << " in "
                                    << printValueInfo(result) << '\n');