[mlir][linalg][bufferize][NFC] Change bufferizableInPlaceAnalysis signature

[springerm]

Summary

Move getInplaceableOpResult() call into bufferizableInPlaceAnalysis.

Depends On D110806

[springerm]

Note: The only goal of this change is to make the signature of bufferizableInPlaceAnalysis smaller. (Fewer arguments.)

[nicolasvasilache]

In D110915#3047215, @springerm wrote:

Note: The only goal of this change is to make the signature of bufferizableInPlaceAnalysis smaller. (Fewer arguments.)

Please add this comment directly in the commit message

[springerm]

rebase

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

 }
 
 /// Return the OpResult that may bufferize into the same buffer as `opOperand`
 /// when the op is bufferized inplace.
 /// Return null if no such result exists.
 static OpResult getInplaceableOpResult(VectorTransferOpInterface op,
                                        OpOperand &opOperand) {
   if (opOperand.get() != op.source() ||
-      !op.source().getType().isa<TensorType>())
+      !op.source().getType().isa<TensorType>() ||
+      isa<vector::TransferReadOp>(op))
     return OpResult();
   return op->getResult(0);
 }
 
 /// Return the OpResult that may bufferize into the same buffer as `opOperand`
 /// when the op is bufferized inplace.
 /// Return null if no such result exists.
 static OpResult getInplaceableOpResult(InsertSliceOp op, OpOperand &opOperand) {
bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
   else
     aliasInfo.bufferizeInPlace(r, s);
 
   LDBG("Done inplace analysis for extract_slice\n");
 
   return success();
 }
 
-/// Analyze the (opOperand, result) pair to determine whether the result can
-/// be bufferized inPlace. If successful, InPlaceSpec::True is set for
-/// `result`. Otherwise, InPlaceSpec::False is set for `result`.
+/// Determine if `operand` can be bufferized in-place with one of the op's
+/// results. If so, set InPlaceSpec::True on the result. Otherwise, set
+/// InPlaceSpec::False on the result.
 static LogicalResult
-bufferizableInPlaceAnalysis(OpOperand &operand, OpResult result,
+bufferizableInPlaceAnalysis(OpOperand &operand,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
+  OpResult result = getInplaceableOpResult(operand);
+  if (!result)
+    return success();
+
   Operation *op = result.getDefiningOp();
   assert(result && !isa<ExtractSliceOp>(op) &&
          "expected OpResult not coming from a ExtractSliceOp");
   (void)op;
 
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
                                    << operand.getOperandNumber() << " in "
                                    << printValueInfo(result) << '\n');
 
   // `result` must bufferize to a writeable buffer to be a candidate.
   // This means the operand  must not alias either:
   //   1. a function bbArg that is not inplaceable or
   //   2. a constant op.
   // to be considered for inplace bufferization
   bool wouldCreateAliasingWriteToNonWriteableBuffer =
       aliasInfo.aliasesNonWriteableBuffer(operand);
   if (wouldCreateAliasingWriteToNonWriteableBuffer)
     LDBG("->the corresponding buffer is not writeable\n");
   else
     LDBG("->bufferizes to writeable inplace buffer\n");
 
-  assert(result == getInplaceableOpResult(operand));
   bool foundInterference =
       wouldCreateAliasingWriteToNonWriteableBuffer ||
       aliasInfo.wouldCreateReadAfterWriteInterference(result, domInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
     // TODO: Atm, all inplace bufferizations yield equivalent tensors. Support
funcOp.walk([&](Operation *op) {
     if (none_of(op->getOperandTypes(), isaTensor) &&
         none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
   // Walk ops in reverse for better interference analysis.
   for (Operation *op : reverse(ops)) {
-    for (OpOperand &opOperand : op->getOpOperands()) {
-      if (OpResult result = getInplaceableOpResult(opOperand))
-        if (result.getType().isa<TensorType>() &&
-            failed(bufferizableInPlaceAnalysis(opOperand, result, aliasInfo,
-                                               domInfo)))
-          return failure();
-    }
+    for (OpOperand &opOperand : op->getOpOperands())
+      if (failed(bufferizableInPlaceAnalysis(opOperand, aliasInfo, domInfo)))
+
     // Special logic to analyze ExtractSliceOp.
     // Note that ExtractSliceOp analysis needs to be interleaved with other ops
     // to properly capture aliases.
     // Walk ExtractSliceOps in reverse for better clobbering analysis behavior:
     // it is easier to detect clobbers of smaller slices before larger ones.
     if (auto extractSliceOp = dyn_cast<ExtractSliceOp>(op)) {
       if (failed(
               bufferizableInPlaceAnalysis(extractSliceOp, aliasInfo, domInfo)))