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 @@ -122,6 +122,34 @@ } }; +struct PackOpInterface + : public BufferizableOpInterface::ExternalModel { + bool bufferizesToAllocation(Operation *op, OpResult opResult) const { + // PackOp reuses all the buffers instead of allocating new ones + return false; + } + + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return true; + } + + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + return false; + } + + AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand, + const AnalysisState &state) const { + assert(op->getNumResults() == 1); + assert(isUniqueCOOType(op->getResultTypes()[0].cast())); + // PackOp reuses the input tensors as data/indices instead of creating new + // ones when packing into a COO format. + return op->getResults(); + } +}; + struct InsertOpInterface : public BufferizableOpInterface::ExternalModel {