[mlir][linalg][bufferize][NFC] Merge bufferizableInPlaceAnalysis overloads

[springerm]

Summary

Depends On D110916

[nicolasvasilache]

I used to have something similar with a single function here but found it counterintuitive.
I am on the fence about this change: on one hand code removal is always welcome, OTOH the branching inside the function is tricky.

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

2394

bool isExtractSliceOp = isa<..> and reuse this ?

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

static LogicalResult bufferize(OpBuilder &b, tensor::ExtractOp extractOp,
   Value l = b.create<memref::LoadOp>(loc, srcMemref, extractOp.indices());
   extractOp.replaceAllUsesWith(l);
   return success();
 }
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
+/// 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.
 ///
 /// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace.
 /// ===========================================================
 ///
 /// When bufferized out of place, a ExtractSlice lowers to alloc + copy. This
 /// cannot change the flow of information for either the source or the
 /// result buffers.
 ///
 /// When bufferized inplace, a ExtractSliceOp does not by itself create any read
 /// or write from memory. Instead, it has the effect of merging the alias sets
 /// of the source and the result buffers.
 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult
-bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
+bufferizableInPlaceAnalysis(OpOperand &operand,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
-  LDBG('\n');
-  LDBG("Inplace analysis for extract_slice: "
-       << printOperationInfo(extractSliceOp) << '\n');
-
-  // If `extractSliceOp` were to be bufferized inplace, it cannot end up
-  // aliasing a write into a non-writeable buffer.
-  bool wouldCreateAliasingWriteToNonWriteableBuffer =
-      aliasInfo.aliasesInPlaceWrite(extractSliceOp.result()) &&
-      aliasInfo.aliasesNonWriteableBuffer(extractSliceOp->getOpOperand(0));
-
-  if (wouldCreateAliasingWriteToNonWriteableBuffer)
-    LDBG("->the corresponding buffer is not writeable\n");
-  else
-    LDBG("->bufferizes to writeable inplace buffer\n");
-
-  // In any of extractSliceOp.result's aliases, can we find 2 such that we hit
-  // an interfering write?
-  OpResult r = extractSliceOp->getResult(0);
-  OpOperand &s = extractSliceOp->getOpOperand(0);
-  bool foundInterference =
-      wouldCreateAliasingWriteToNonWriteableBuffer ||
-      aliasInfo.wouldCreateReadAfterWriteInterference(r, domInfo);
-  if (foundInterference)
-    aliasInfo.bufferizeOutOfPlace(r);
-  else
-    aliasInfo.bufferizeInPlace(r, s);
-
-  LDBG("Done inplace analysis for extract_slice\n");
-
-  return success();
+  Operation *op = operand.getOwner();
+  bool isExtractSliceOp = false;
+  OpResult result;
+  if (isa<ExtractSliceOp>(op) && (operand.getOperandNumber() == 0)) {
+    result = op->getOpResult(0);
+    isExtractSliceOp = true;
+  } else {
+    result = getInplaceableOpResult(operand);
   }
 
-/// 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,
-                            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);
+      aliasInfo.aliasesNonWriteableBuffer(operand) &&
+      (!isExtractSliceOp || aliasInfo.aliasesInPlaceWrite(result));
+
   if (wouldCreateAliasingWriteToNonWriteableBuffer)
     LDBG("->the corresponding buffer is not writeable\n");
   else
     LDBG("->bufferizes to writeable inplace buffer\n");
 
   bool foundInterference =
       wouldCreateAliasingWriteToNonWriteableBuffer ||
       aliasInfo.wouldCreateReadAfterWriteInterference(result, domInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
-    // TODO: Atm, all inplace bufferizations yield equivalent tensors. Support
-    // more cases on a per-need basis.
     aliasInfo.bufferizeInPlace(
-        result, operand, BufferizationAliasInfo::BufferRelation::Equivalent);
+        result, operand,
+        isExtractSliceOp ? BufferizationAliasInfo::BufferRelation::None
+                         : BufferizationAliasInfo::BufferRelation::Equivalent);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable:
 ///   1. First, analyze InsertSliceOp greedily: we almost never want to
funcOp.walk([&](Operation *op) {
     // No tensors => no buffers.
     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 (Operation *op : reverse(ops))
     for (OpOperand &opOperand : op->getOpOperands())
       if (failed(bufferizableInPlaceAnalysis(opOperand, aliasInfo, domInfo)))
         return failure();
 
-    // 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)))
-        return failure();
-      continue;
-    }
-  }
-
   LDBG("End InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for functions.
 //===----------------------------------------------------------------------===//