From fd7469404bcd5c4142070ef17d94007a57268ff3 Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming@google.com>
Date: Thu, 20 Oct 2022 19:40:24 +0000
Subject: [PATCH] [mlir][sparse] attach bufferizableOpInterface to InsertOp

To fix D136286

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D136372
---
 .../Transforms/BufferizableOpInterfaceImpl.cpp     | 40 ++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
index aa6c6e9..7c16cf6 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -94,6 +94,45 @@ struct NewOpInterface
   }
 };
 
+struct InsertOpInterface
+    : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
+                                                    sparse_tensor::InsertOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    // Does bufferization handle realloc?
+    return false;
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // InsertOp writes to memory.
+    return true;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    // InsertOp returns an alias of its operand.
+    assert(op->getNumResults() == 1);
+    return op->getResults();
+  }
+
+  BufferRelation bufferRelation(Operation *oo, OpResult opResult,
+                                const AnalysisState &state) const {
+    // InsertOp returns the same object (realloc should not invalidate
+    // aliases).
+    return BufferRelation::Equivalent;
+  }
+
+  bool isWritable(Operation *op, Value value,
+                  const AnalysisState &state) const {
+    return true;
+  }
+};
+
 } // namespace
 } // namespace sparse_tensor
 } // namespace mlir
@@ -105,5 +144,6 @@ void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels(
         sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
         sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
         sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
+        sparse_tensor::InsertOp::attachInterface<InsertOpInterface>(*ctx);
       });
 }
-- 
2.7.4