Update the rewrite methods of each of the DialectConversion patterns to notify the PatternRewriter that the operation is being replaced.

--

PiperOrigin-RevId: 248965082

3 files changed, 17 insertions, 17 deletions

diff --git a/mlir/lib/Dialect/QuantOps/Transforms/ConvertConst.cpp b/mlir/lib/Dialect/QuantOps/Transforms/ConvertConst.cpp
index 228e16d752e..44b1156ec28 100644
--- a/mlir/lib/Dialect/QuantOps/Transforms/ConvertConst.cpp
+++ b/mlir/lib/Dialect/QuantOps/Transforms/ConvertConst.cpp
@@ -115,7 +115,7 @@ void QuantizedConstRewrite::rewrite(Operation *op,
   auto newConstOp =
       rewriter.create<ConstantOp>(fusedLoc, newConstValueType, newConstValue);
   rewriter.replaceOpWithNewOp<StorageCastOp>(
-      op, {origConstOp}, *op->result_type_begin(), newConstOp);
+      {origConstOp}, op, *op->result_type_begin(), newConstOp);
 }
 
 void ConvertConstPass::runOnFunction() {
diff --git a/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp b/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp
index 6de9a151455..f1c43eff289 100644
--- a/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp
+++ b/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp
@@ -311,7 +311,7 @@ struct OneToOneLLVMOpLowering : public LLVMLegalizationPattern<SourceOp> {
 
     // If the operation produced 0 or 1 result, return them immediately.
     if (numResults == 0)
-      return;
+      return rewriter.replaceOp(op, llvm::None);
     if (numResults == 1)
       return rewriter.replaceOp(op, newOp.getOperation()->getResult(0));
 
@@ -542,8 +542,8 @@ struct DeallocOpLowering : public LLVMLegalizationPattern<DeallocOp> {
         rewriter, op->getLoc(), operands[0], elementPtrType, hasStaticShape);
     Value *casted = rewriter.create<LLVM::BitcastOp>(
         op->getLoc(), getVoidPtrType(), bufferPtr);
-    rewriter.create<LLVM::CallOp>(op->getLoc(), ArrayRef<Type>(),
-                                  rewriter.getFunctionAttr(freeFunc), casted);
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+        op, ArrayRef<Type>(), rewriter.getFunctionAttr(freeFunc), casted);
   }
 };
 
@@ -803,7 +803,7 @@ struct StoreOpLowering : public LoadStoreOpLowering<StoreOp> {
 
     Value *dataPtr = getDataPtr(op->getLoc(), type, operands[1],
                                 operands.drop_front(2), rewriter, getModule());
-    rewriter.create<LLVM::StoreOp>(op->getLoc(), operands[0], dataPtr);
+    rewriter.replaceOpWithNewOp<LLVM::StoreOp>(op, operands[0], dataPtr);
   }
 };
 
@@ -818,8 +818,8 @@ struct OneToOneLLVMTerminatorLowering
                ArrayRef<Block *> destinations,
                ArrayRef<ArrayRef<Value *>> operands,
                PatternRewriter &rewriter) const override {
-    rewriter.create<TargetOp>(op->getLoc(), properOperands, destinations,
-                              operands, op->getAttrs());
+    rewriter.replaceOpWithNewOp<TargetOp>(op, properOperands, destinations,
+                                          operands, op->getAttrs());
   }
 };
 
@@ -838,17 +838,15 @@ struct ReturnOpLowering : public LLVMLegalizationPattern<ReturnOp> {
 
     // If ReturnOp has 0 or 1 operand, create it and return immediately.
     if (numArguments == 0) {
-      rewriter.create<LLVM::ReturnOp>(
-          op->getLoc(), llvm::ArrayRef<Value *>(), llvm::ArrayRef<Block *>(),
+      return rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(
+          op, llvm::ArrayRef<Value *>(), llvm::ArrayRef<Block *>(),
           llvm::ArrayRef<llvm::ArrayRef<Value *>>(), op->getAttrs());
-      return;
     }
     if (numArguments == 1) {
-      rewriter.create<LLVM::ReturnOp>(
-          op->getLoc(), llvm::ArrayRef<Value *>(operands.front()),
+      return rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(
+          op, llvm::ArrayRef<Value *>(operands.front()),
           llvm::ArrayRef<Block *>(), llvm::ArrayRef<llvm::ArrayRef<Value *>>(),
           op->getAttrs());
-      return;
     }
 
     // Otherwise, we need to pack the arguments into an LLVM struct type before
@@ -861,8 +859,8 @@ struct ReturnOpLowering : public LLVMLegalizationPattern<ReturnOp> {
           op->getLoc(), packedType, packed, operands[i],
           getIntegerArrayAttr(rewriter, i));
     }
-    rewriter.create<LLVM::ReturnOp>(
-        op->getLoc(), llvm::makeArrayRef(packed), llvm::ArrayRef<Block *>(),
+    rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(
+        op, llvm::makeArrayRef(packed), llvm::ArrayRef<Block *>(),
         llvm::ArrayRef<llvm::ArrayRef<Value *>>(), op->getAttrs());
   }
 };
diff --git a/mlir/lib/Linalg/Transforms/LowerToLLVMDialect.cpp b/mlir/lib/Linalg/Transforms/LowerToLLVMDialect.cpp
index 2ecea9ca5a2..a4da12a922b 100644
--- a/mlir/lib/Linalg/Transforms/LowerToLLVMDialect.cpp
+++ b/mlir/lib/Linalg/Transforms/LowerToLLVMDialect.cpp
@@ -243,6 +243,7 @@ public:
     Value *casted = bitcast(voidPtrTy, extractvalue(elementPtrTy, operands[0],
                                                     positionAttr(rewriter, 0)));
     call(ArrayRef<Type>(), rewriter.getFunctionAttr(freeFunc), casted);
+    rewriter.replaceOp(op, llvm::None);
   }
 };
 
@@ -498,6 +499,7 @@ class StoreOpConversion : public LoadStoreOpConversion<linalg::StoreOp> {
     ArrayRef<Value *> indices = operands.drop_front(2);
     Value *ptr = obtainDataPtr(op, viewDescriptor, indices, rewriter);
     llvm_store(data, ptr);
+    rewriter.replaceOp(op, llvm::None);
   }
 };
 
@@ -578,8 +580,8 @@ public:
 
     auto fAttr = rewriter.getFunctionAttr(f);
     auto named = rewriter.getNamedAttr("callee", fAttr);
-    rewriter.create<LLVM::CallOp>(op->getLoc(), operands,
-                                  ArrayRef<NamedAttribute>{named});
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, operands,
+                                              ArrayRef<NamedAttribute>{named});
   }
 };