[mlir][linalg][bufferize][NFC] Add bufferizesToAliasOnly

[springerm]

Summary

The list of operations that do neither read nor write, but create an alias when bufferizing inplace, is getting longer. This commit adds a helper function so that we do not have to spell out the entire list each time.

Depends On D112514

[springerm]

add assertion

[springerm]

rebase

[nicolasvasilache]

interfaces FTW

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

return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
       .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
             tensor::CastOp>(
           [&](auto op) { return getAliasingOpResult(op, opOperand); })
       // All other ops, return the result of `getInplaceableOpResult`.
       .Default(
           [&](Operation *op) { return getInplaceableOpResult(opOperand); });
 }
 
+/// Return `true` if the given OpOperand does not bufferize to a memory read or
+/// write, but creates an alias when bufferized inplace.
+static bool bufferizesToAliasOnly(OpOperand &opOperand) {
+  Operation *owner = opOperand.getOwner();
+  // TODO: In the future this may need to evolve into a TypeSwitch. For all
+  // currently supported ops, the aliasing-only OpOperand is always the first
+  // one.
+  return isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+             tensor::CastOp>(owner) &&
+         &opOperand == &owner->getOpOperand(0);
+}
+
 // Predeclaration of function.
 static bool bufferizesToMemoryRead(OpOperand &opOperand);
 
 /// Return true if the given value is read by an op that bufferizes to a memory
 /// read. Also takes into account ops that create an alias but do not read by
 /// themselves (e.g., ExtractSliceOp).
 static bool isValueRead(Value value) {
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
     // Skip over all ops that create an alias but do not read.
-    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
-            tensor::CastOp>(uMaybeReading->getOwner()))
-      for (OpOperand &use : uMaybeReading->getOwner()->getResult(0).getUses())
+    if (bufferizesToAliasOnly(*uMaybeReading))
+      for (OpOperand &use : getAliasingOpResult(*uMaybeReading).getUses())
         workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
       return true;
   }
 
   return false;
 }
 
 /// Return true if `opOperand` bufferizes to a memory read.
 static bool bufferizesToMemoryRead(OpOperand &opOperand) {
   // Unknown op that returns a tensor. The inplace analysis does not support
   // it. Conservatively return true.
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
     return true;
   // Some ops alone do not bufferize to a memory read, but one of their uses
   // may.
-  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
-          tensor::CastOp>(opOperand.getOwner()))
+  if (bufferizesToAliasOnly(opOperand))
     return false;
   // scf::ForOp alone doesn't bufferize to a memory read, one of the uses of its
   // matching bbArg may.
   if (auto forOp = dyn_cast<scf::ForOp>(opOperand.getOwner()))
     return isValueRead(forOp.getRegionIterArgForOpOperand(opOperand));
   // TiledLoop alone doesn't bufferize to a memory read, one of the uses of its
   // matching bbArg may.
   if (auto tiledLoopOp = dyn_cast<TiledLoopOp>(opOperand.getOwner()))

 
 /// Return true if `opOperand` bufferizes to a memory write.
 static bool bufferizesToMemoryWrite(OpOperand &opOperand) {
   // These terminators are not writes.
   if (isa<ReturnOp, linalg::YieldOp, scf::YieldOp>(opOperand.getOwner()))
     return false;
   // Some ops alone do not bufferize to a memory write, but one of their uses
   // may.
-  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
-          tensor::CastOp>(opOperand.getOwner()))
+  if (bufferizesToAliasOnly(opOperand))
     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
#endif // NDEBUG
   else
     aliasInfo.bufferizeInPlace(result, operand);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
-/// This analysis function is used for ops where the first OpOperand aliases
-/// with the first OpResult, without creating a read or write. There are a few
-/// ops besides ExtractSliceOp that have such semantics.
+/// This analysis function is used for OpOperands that alias with an OpResult
+/// but are not inplaceable on it. E.g., ExtractSliceOp.
 ///
 /// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace:
 ///
 /// When bufferized out of place, an ExtractSliceOp lowers to alloc + copy. This
 /// cannot change the flow of information for either the source or the
 /// result buffers.
 ///
 /// When bufferized inplace, an 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
-bufferizableInPlaceAnalysisAliasOnlyOp(Operation *op,
+bufferizableInPlaceAnalysisAliasOnlyOp(OpOperand &operand,
                                        BufferizationAliasInfo &aliasInfo,
                                        const DominanceInfo &domInfo) {
-  return bufferizableInPlaceAnalysisImpl(
-      op->getOpOperand(0), op->getOpResult(0), aliasInfo, domInfo);
+  OpResult result = getAliasingOpResult(operand);
+  assert(result && "expected that the OpOperand has an aliasing OpResult");
+  return bufferizableInPlaceAnalysisImpl(operand, result, aliasInfo, domInfo);
 }
 
 /// 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();
   return bufferizableInPlaceAnalysisImpl(operand, result, aliasInfo, domInfo);
 }
 
 /// Analyze the `ops` to determine which OpResults are inplaceable. Walk ops in
 /// reverse and bufferize ops greedily. This is a good starter heuristic.
 /// ExtractSliceOps are interleaved with other ops in traversal order.
 LogicalResult mlir::linalg::inPlaceAnalysis(SmallVector<Operation *> &ops,
                                             BufferizationAliasInfo &aliasInfo,
                                             const DominanceInfo &domInfo) {
   // Walk ops in reverse for better interference analysis.
   for (Operation *op : reverse(ops)) {
-    for (OpOperand &opOperand : op->getOpOperands())
+    for (OpOperand &opOperand : op->getOpOperands()) {
       if (failed(bufferizableInPlaceAnalysis(opOperand, aliasInfo, domInfo)))
         return failure();
 
-    // Special logic to analyze ops who's OpResults are not inplaceable on an
-    // OpOperand but may create an alias.
-    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
-            tensor::CastOp>(op))
-      if (failed(
-              bufferizableInPlaceAnalysisAliasOnlyOp(op, aliasInfo, domInfo)))
+      // Special logic to analyze OpOperands that are not inplaceable on an
+      // OpResult but may create an alias.
+      if (bufferizesToAliasOnly(opOperand))
+        if (failed(bufferizableInPlaceAnalysisAliasOnlyOp(opOperand, aliasInfo,
+                                                          domInfo)))
           return failure();
     }
+  }
 
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable.
 static LogicalResult
 inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
                           const DominanceInfo &domInfo) {
auto extractOp = b.create<tensor::ExtractSliceOp>(
         initTensor.getLoc(), insertOp.dest(), insertOp.getMixedOffsets(),
         insertOp.getMixedSizes(), insertOp.getMixedStrides());
     // Uses of the InitTensorOp are replaced here, but the op is not deleted.
     // InitTensorOps without uses are ignored by the bufferization.
     initTensor.replaceAllUsesWith(extractOp.result());
     aliasInfo.createAliasInfoEntry(extractOp.result());
 
     // Run analysis on the ExtractSliceOp.
-    if (failed(bufferizableInPlaceAnalysisAliasOnlyOp(extractOp, aliasInfo,
-                                                      domInfo)))
+    if (failed(bufferizableInPlaceAnalysisAliasOnlyOp(
+            extractOp->getOpOperand(0), aliasInfo, domInfo)))
       return WalkResult::interrupt();
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }