[mlir][linalg][bufferize][NFC] Provide default implementation of getAliasingOpOperand

[springerm]

Summary

This simplifies op interface implementations.

Depends On D115022

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td

let methods = [
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpOperands that alias with a given OpResult when
           bufferized in-place. This method will never be called on OpResults
           that do not have a tensor type.
 
+          By default, this method is the inverse of `getAliasingOpResult`. Ops
+          with a region that yield values may want to override this method to
+          return the OpOperands that are yielded by the terminator.
+
           Note: This method can return multiple OpOperands, indicating that the
           given OpResult may at runtime alias with any of the OpOperands. This
           is useful for branches and for ops such as `std.select`.
         }],
         /*retType=*/"SmallVector<OpOperand *>",
         /*methodName=*/"getAliasingOpOperand",
         /*args=*/(ins "OpResult":$opResult),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
-          // Does not have to be implemented for ops without tensor OpResults.
-          llvm_unreachable("getInplaceableOpResult not implemented");
+          assert(opResult.getType().isa<TensorType>() &&
+                 "expected OpResult with tensor type");
+          SmallVector<OpOperand *> result;
+          auto bufferizableOp =
+              cast<BufferizableOpInterface>($_op.getOperation());
+          for (OpOperand &opOperand : $_op.getOperation()->getOpOperands()) {
+            if (!opOperand.get().getType().isa<TensorType>())
+              continue;
+            if (bufferizableOp.getAliasingOpResult(opOperand) == opResult)
+              result.push_back(&opOperand);
+          }
+          return result;
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the buffer relation between the given OpOperand and its
           aliasing OpResult when bufferized in-place. Most OpOperands have an
           "equivalence" relation.
 

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

 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace arith_ext {
 
 struct ConstantOpInterface
     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
                                                     arith::ConstantOp> {
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {};
-  }
-
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto constantOp = cast<arith::ConstantOp>(op);
     if (!constantOp.getResult().getType().isa<TensorType>())
       return success();
     assert(constantOp.getType().dyn_cast<RankedTensorType>() &&
            "not a constant ranked tensor");
     auto moduleOp = constantOp->getParentOfType<ModuleOp>();

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

 struct ToMemrefOpInterface
     : public BufferizableOpInterface::ExternalModel<ToMemrefOpInterface,
                                                     bufferization::ToMemrefOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     // It is unknown whether the resulting MemRef will be read or not.
     return true;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     return success();
   }
 };
 
 /// ToTensorOp conceptually loads a tensor from a memory location. Such ops do
 /// not lower any further, and they should have disappeared by the time the
 /// input is fully bufferized.
 ///
 /// The analysis has no information about the memref that is loaded from by the
 /// ToTensorOp. We have to assume that the loaded tensor may after bufferization
 /// potentially alias with any other bufferized tensor. Since ToTensorOp and
 /// ToMemrefOp have no aliasing OpOperand/OpResult pairs, this cannot be encoded
 /// directly in the analysis. However, declaring ToTensorOp results as not
 /// writable also enforces a buffer copy and has the same effect.
 struct ToTensorOpInterface
     : public BufferizableOpInterface::ExternalModel<ToTensorOpInterface,
                                                     bufferization::ToTensorOp> {
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {};
-  }
-
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto tensorLoadOp = cast<bufferization::ToTensorOp>(op);
     state.mapBuffer(tensorLoadOp.result(), tensorLoadOp.memref());
     return success();
   }
 
   bool isWritable(Operation *op, Value value) const {

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

LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     return bufferizeLinalgOp(b, cast<LinalgOp>(op), state);
   }
 };
 
 struct InitTensorOpInterface
     : public BufferizableOpInterface::ExternalModel<InitTensorOpInterface,
                                                     linalg::InitTensorOp> {
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {};
-  }
-
   bool isMemoryWrite(Operation *op, OpResult opResult) const {
     // InitTensorOps allocate but do not write.
     return false;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
struct TiledLoopOpInterface
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     // TiledLoop alone doesn't bufferize to a memory write, one of the uses of
     // its matching bbArg may.
     auto bufferizableOp = cast<BufferizableOpInterface>(op);
     return static_cast<bool>(bufferizableOp.getAliasingOpResult(opOperand));
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    // TODO: TiledLoopOp helper method to avoid leaking impl details.
-    auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
-    return {&op->getOpOperand(tiledLoopOp.getNumControlOperands() +
-                              tiledLoopOp.getNumInputs() +
-                              opResult.getResultNumber())};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
     return tiledLoopOp.getTiedOpResult(opOperand);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }

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

struct CallOpInterface
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     // CallOpInterface alone doesn't bufferize to a memory read, 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.
     return true;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    // TODO: Can we do better?
-    return {};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     // 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.
     return OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {

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

struct ForOpInterface
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     // Tensor iter_args of scf::ForOps are always considered as a write. This is
     // to simplify the analysis.
     // TODO: Consider doing sth. like isValueWritten.
     return true;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    auto forOp = cast<scf::ForOp>(op);
-    return {&forOp.getIterOpOperands()[opResult.getResultNumber()]};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     auto forOp = cast<scf::ForOp>(op);
     if (!opOperand.get().getType().isa<RankedTensorType>())
       return OpResult();
     return forOp.getResultForOpOperand(opOperand);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {

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

struct CastOpInterface
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {&op->getOpOperand(0)};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return op->getResult(0);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
struct ExtractSliceOpInterface
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {&op->getOpOperand(0) /*source*/};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return &opOperand == &op->getOpOperand(0) /*source*/
                ? op->getResult(0)
                : OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::None;
struct InsertSliceOpInterface
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return &opOperand == &op->getOpOperand(1) /*dest*/;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {&op->getOpOperand(1) /*dest*/};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return &opOperand == &op->getOpOperand(1) /*dest*/
                ? op->getResult(0)
                : OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;

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

struct TransferWriteOpInterface
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     assert(opOperand.get().getType().isa<TensorType>() &&
            "only tensor types expected");
     return true;
   }
 
-  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
-                                                OpResult opResult) const {
-    return {&op->getOpOperand(1)};
-  }
-
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     assert(opOperand.get().getType().isa<TensorType>() &&
            "only tensor types expected");
     return op->getOpResult(0);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;