[Mlir-commits] [mlir] 9106d35 - [mlir][bufferize] Use rewriter instead of replacing all uses directly

[Matthias Springer]

Author: Matthias Springer
Date: 2022-02-12T02:35:36+09:00
New Revision: 9106d35b91caf2f5fd13c3925579eabcd253ffb3

URL: https://github.com/llvm/llvm-project/commit/9106d35b91caf2f5fd13c3925579eabcd253ffb3
DIFF: https://github.com/llvm/llvm-project/commit/9106d35b91caf2f5fd13c3925579eabcd253ffb3.diff

LOG: [mlir][bufferize] Use rewriter instead of replacing all uses directly

This is important for compatibility with DialectConversion.

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 261a107c0e56..ee5a34bcc83b 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(