[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); })
         // Some ops just return an alias to an operand when bufferized inplace.
         // Such OpResults are never inplaceable on an OpOperand.
-        .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+        .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+              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
 }
 
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// 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, TensorCollapseShapeOp, TensorExpandShapeOp>(
-          [&](auto op) { r.push_back(&op->getOpOperand(0)); })
+      .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+            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) {

 /// will alias with the result.
 static OpResult getAliasingOpResult(TensorCollapseShapeOp op,
                                     OpOperand &opOperand) {
   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())
       // 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, TensorCollapseShapeOp, TensorExpandShapeOp>(
+      .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); });
 }
 
 // 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>(
-            uMaybeReading->getOwner()))
+    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+            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;
   // Some ops alone do not bufferize to a memory read, but one of their uses
   // may.
-  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
-          opOperand.getOwner()))
+  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+          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()))

 
 /// 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>(
-          opOperand.getOwner()))
+  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+          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
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 ops who's OpResults are not inplaceable on an
     // OpOperand but may create an alias.
-    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(op))
+    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp,
+            tensor::CastOp>(op))
       if (failed(
               bufferizableInPlaceAnalysisAliasOnlyOp(op, aliasInfo, domInfo)))
         return failure();
   }
 
   return success();
 }