[mlir][linalg][bufferize] tensor::CastOp is an alias-only op

[springerm]

Summary

tensor::CastOp by itself does not bufferize to memory read/write.

Depends On D112512

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

 static OpResult getInplaceableOpResult(InsertSliceOp op, OpOperand &opOperand) {
   if (&opOperand != &op->getOpOperand(1) /*dest*/)
     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(tensor::CastOp op,
-                                       OpOperand &opOperand) {
-  return op->getResult(0);
-}
-
-/// Return the OpResult that may bufferize into the same buffer as `opOperand`
-/// when the op is bufferized inplace.
 /// The inplace analysis uses this information along with interfering read
 /// analysis to determine which op results reuse the same buffer as some
 /// operand.
 static OpResult getInplaceableOpResult(OpOperand &opOperand) {
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
       // clang-format off
         // Ops that perform destructive updates on operand(s) to produce
         // result(s).
-        .Case<tensor::CastOp,
-              scf::ForOp,
+        .Case<scf::ForOp,
               InsertSliceOp,
               LinalgOp,
               TiledLoopOp,
               VectorTransferOpInterface>(
             [&](auto op) { return getInplaceableOpResult(op, opOperand); })
-        // ExtractSliceOp is special, when bufferized inplace it just returns an
-        // alias to its operand. Its result is never inplaceable on its operand.
-        .Case([&](ExtractSliceOp op) { return OpResult(); })
+        // Some ops just return an alias to an operand when bufferized inplace.
+        // Such OpResults are never inplaceable on an OpOperand.
+        .Case<ExtractSliceOp, tensor::CastOp>(
+            [] (auto op) { return OpResult(); })
         // CallOpInterface is special, it needs to wait for the callee to be
         // bufferized and needs to inspect the BufferAliasInfo object. It can't
         // 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
 }

 /// 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 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)); })
       .Case([&](scf::IfOp op) { populateAliasingOpOperands(op, result, r); })
+      .Case<ExtractSliceOp, tensor::CastOp>(
+          [&](auto 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) {
         r.push_back(&op.getIterOpOperands()[result.getResultNumber()]);
       })
       .Case([&](InsertSliceOp op) { r.push_back(&op->getOpOperand(1)); })
       .Case([&](LinalgOp op) {
TypeSwitch<Operation *>(result.getDefiningOp())
         llvm_unreachable("unexpected defining op");
       });
   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())
+  if (&op->getOpOperand(0) == &opOperand)
+    return op->getResult(0);
+  return OpResult();
+}
+
+/// If the a tensor::CastOp is bufferized in-place, the source operand will
+/// alias with the result.
+static OpResult getAliasingOpResult(tensor::CastOp op, OpOperand &opOperand) {
+  if (&op->getOpOperand(0) == &opOperand)
     return op->getResult(0);
   return OpResult();
 }
 
 /// Determine which OpResult will alias with `opOperand` if the op is bufferized
 /// in place. This is a superset of `getInplaceableOpResult`.
 /// TODO: in the future this may need to evolve towards a list of OpResult.
 static OpResult getAliasingOpResult(OpOperand &opOperand) {
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
-      // ExtractSliceOp is different: its result is not inplaceable on op.source
-      // but when bufferized inplace, the result is an aliasing subregion of
-      // op.source.
-      .Case(
-          [&](ExtractSliceOp op) { return getAliasingOpResult(op, opOperand); })
+      // Some ops are different: Their result is not inplaceable on an OpOperand
+      // but when bufferized inplace, their result is aliasing (a subregion of)
+      // an OpOperand.
+      .Case<ExtractSliceOp, tensor::CastOp>(
+          [&](auto op) { return getAliasingOpResult(op, opOperand); })
       // All other ops, return the result of `getInplaceableOpResult`.
       .Default(
           [&](Operation *op) { return getInplaceableOpResult(opOperand); });
 }
 
 // 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 ExtractSliceOps. These do not read by themselves but just
-    // add a new alias.
-    if (auto extractSliceOp =
-            dyn_cast<ExtractSliceOp>(uMaybeReading->getOwner()))
-      for (OpOperand &use : extractSliceOp.result().getUses())
+    // Skip over all ops that create an alias but do not read.
+    if (isa<ExtractSliceOp, tensor::CastOp>(uMaybeReading->getOwner()))
+      for (OpOperand &use : uMaybeReading->getOwner()->getResult(0).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;
-  // ExtractSliceOp alone doesn't bufferize to a memory read, one of its uses
+  // Some ops alone do not bufferize to a memory read, but one of their uses
   // may.
-  if (isa<ExtractSliceOp>(opOperand.getOwner()))
+  if (isa<ExtractSliceOp, tensor::CastOp>(opOperand.getOwner()))
     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()))
static bool bufferizesToMemoryRead(OpOperand &opOperand) {
   return true;
 }
 
 /// 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;
-  // ExtractSliceOp alone doesn't bufferize to a memory write, one of its uses
+  // Some ops alone do not bufferize to a memory write, but one of their uses
   // may.
-  if (isa<ExtractSliceOp>(opOperand.getOwner()))
+  if (isa<ExtractSliceOp, tensor::CastOp>(opOperand.getOwner()))
     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.
 ///
-/// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace.
-/// ===========================================================
+/// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace:
 ///
-/// When bufferized out of place, a ExtractSlice lowers to alloc + copy. This
+/// 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, 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.
+/// 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
-bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
+bufferizableInPlaceAnalysisAliasOnlyOp(Operation *op,
                                        BufferizationAliasInfo &aliasInfo,
                                        const DominanceInfo &domInfo) {
-  return bufferizableInPlaceAnalysisImpl(extractSliceOp->getOpOperand(0),
-                                         extractSliceOp->getOpResult(0),
-                                         aliasInfo, domInfo);
+  return bufferizableInPlaceAnalysisImpl(
+      op->getOpOperand(0), op->getOpResult(0), 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,
LogicalResult mlir::linalg::inPlaceAnalysis(SmallVector<Operation *> &ops,
   }
 
   // Walk ops in reverse for better interference analysis.
   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))
+    // Special logic to analyze ops who's OpResults are not inplaceable on an
+    // OpOperand but may create an alias.
+    if (isa<ExtractSliceOp, tensor::CastOp>(op))
       if (failed(
-              bufferizableInPlaceAnalysis(extractSliceOp, aliasInfo, domInfo)))
+              bufferizableInPlaceAnalysisAliasOnlyOp(op, aliasInfo, domInfo)))
         return failure();
   }
 
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable.
 static LogicalResult
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(bufferizableInPlaceAnalysis(extractOp, aliasInfo, domInfo)))
+    if (failed(bufferizableInPlaceAnalysisAliasOnlyOp(extractOp, aliasInfo,
+                                                      domInfo)))
       return WalkResult::interrupt();
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }