[Mlir-commits] [mlir] 7fef8d6 - [mlir][sparse] implement bufferizableOpInterface for sparse_tensor.pack operation

[Peiming Liu]

Author: Peiming Liu
Date: 2023-02-03T23:55:59Z
New Revision: 7fef8d69cc7e0623ff00835f1a71de94d54157ed

URL: https://github.com/llvm/llvm-project/commit/7fef8d69cc7e0623ff00835f1a71de94d54157ed
DIFF: https://github.com/llvm/llvm-project/commit/7fef8d69cc7e0623ff00835f1a71de94d54157ed.diff

LOG: [mlir][sparse] implement bufferizableOpInterface for sparse_tensor.pack operation

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143281

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
index cc3555ebc7ed..d2e065739a97 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -122,6 +122,34 @@ struct NewOpInterface
   }
 };
 
+struct PackOpInterface
+    : public BufferizableOpInterface::ExternalModel<PackOpInterface,
+                                                    sparse_tensor::PackOp> {
+  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<RankedTensorType>()));
+    // 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<InsertOpInterface,
                                                     sparse_tensor::InsertOp> {