diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp @@ -26,6 +26,34 @@ namespace sparse_tensor { namespace { +struct ConcatenateOpInterface + : public BufferizableOpInterface::ExternalModel< + ConcatenateOpInterface, sparse_tensor::ConcatenateOp> { + bool bufferizesToAllocation(Operation *op, OpResult opResult) const { + return true; + } + + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return false; + } + + SmallVector getAliasingOpResult(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return {}; + } + + bool isWritable(Operation *op, Value value, + const AnalysisState &state) const { + return true; + } +}; + struct ConvertOpInterface : public BufferizableOpInterface::ExternalModel { @@ -123,17 +151,128 @@ } }; +struct NumberOfEntriesOpInterface + : public BufferizableOpInterface::ExternalModel< + NumberOfEntriesOpInterface, sparse_tensor::NumberOfEntriesOp> { + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return false; + } + + SmallVector getAliasingOpResult(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return {}; + } +}; + +struct ToIndicesBufferOpInterface + : public BufferizableOpInterface::ExternalModel< + ToIndicesBufferOpInterface, sparse_tensor::ToIndicesBufferOp> { + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + // Potential writes into memory through the result of sparse_tensor.indices + // are not considered. + return false; + } + + SmallVector getAliasingOpResult(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return {}; + } +}; + +struct ToIndicesOpInterface + : public BufferizableOpInterface::ExternalModel< + ToIndicesOpInterface, sparse_tensor::ToIndicesOp> { + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + // Potential writes into memory through the result of sparse_tensor.indices + // are not considered. + return false; + } + + SmallVector getAliasingOpResult(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return {}; + } +}; + +struct ToPointersOpInterface + : public BufferizableOpInterface::ExternalModel< + ToPointersOpInterface, sparse_tensor::ToPointersOp> { + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + // Potential writes into memory through the result of sparse_tensor.pointers + // are not considered. + return false; + } + + SmallVector getAliasingOpResult(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return {}; + } +}; + +struct ToValuesOpInterface + : public BufferizableOpInterface::ExternalModel { + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + // Potential writes into memory through the result of sparse_tensor.values + // are not considered. + return false; + } + + SmallVector getAliasingOpResult(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return {}; + } +}; + } // namespace } // namespace sparse_tensor } // namespace mlir void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels( DialectRegistry ®istry) { - registry.addExtension( - +[](MLIRContext *ctx, sparse_tensor::SparseTensorDialect *dialect) { - sparse_tensor::ConvertOp::attachInterface(*ctx); - sparse_tensor::LoadOp::attachInterface(*ctx); - sparse_tensor::NewOp::attachInterface(*ctx); - sparse_tensor::InsertOp::attachInterface(*ctx); - }); + registry.addExtension(+[](MLIRContext *ctx, + sparse_tensor::SparseTensorDialect *dialect) { + sparse_tensor::ConcatenateOp::attachInterface(*ctx); + sparse_tensor::ConvertOp::attachInterface(*ctx); + sparse_tensor::LoadOp::attachInterface(*ctx); + sparse_tensor::NewOp::attachInterface(*ctx); + sparse_tensor::InsertOp::attachInterface(*ctx); + sparse_tensor::NumberOfEntriesOp::attachInterface< + NumberOfEntriesOpInterface>(*ctx); + sparse_tensor::ToIndicesBufferOp::attachInterface< + ToIndicesBufferOpInterface>(*ctx); + sparse_tensor::ToIndicesOp::attachInterface(*ctx); + sparse_tensor::ToPointersOp::attachInterface(*ctx); + sparse_tensor::ToValuesOp::attachInterface(*ctx); + }); }