From 9106d35b91caf2f5fd13c3925579eabcd253ffb3 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Sat, 12 Feb 2022 02:31:00 +0900
Subject: [PATCH] [mlir][bufferize] Use rewriter instead of replacing all uses
 directly

This is important for compatibility with DialectConversion.
---
 mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 261a107..ee5a34b 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -293,14 +293,13 @@ FailureOr<Value> BufferizationState::getBuffer(
 void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
                                                   Operation *op,
                                                   ValueRange values) {
+  assert(values.size() == op->getNumResults() &&
+         "expected one value per OpResult");
   OpBuilder::InsertionGuard g(rewriter);
 
   // Replace all OpResults with the given values.
+  SmallVector<Value> replacements;
   for (OpResult opResult : op->getOpResults()) {
-    // Skip OpResult if it has no uses.
-    if (opResult.getUses().empty())
-      continue;
-
     Value replacement = values[opResult.getResultNumber()];
     if (opResult.getType().isa<TensorType>()) {
       // The OpResult is a tensor. Such values are replaced with memrefs during
@@ -315,10 +314,10 @@ void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
       replacement = rewriter.create<bufferization::ToTensorOp>(
           replacement.getLoc(), replacement);
     }
-    opResult.replaceAllUsesWith(replacement);
+    replacements.push_back(replacement);
   }
 
-  rewriter.eraseOp(op);
+  rewriter.replaceOp(op, replacements);
 }
 
 AlwaysCopyBufferizationState::AlwaysCopyBufferizationState(
-- 
2.7.4