[Mlir-commits] [mlir] 553cebd - [mlir][Vector] Use a RewriterBase for IR rewrites in VectorTransferOpTransforms

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2023-03-25T01:48:50-07:00
New Revision: 553cebde0669aa98bb0624d9203b008c6e02f40e

URL: https://github.com/llvm/llvm-project/commit/553cebde0669aa98bb0624d9203b008c6e02f40e
DIFF: https://github.com/llvm/llvm-project/commit/553cebde0669aa98bb0624d9203b008c6e02f40e.diff

LOG: [mlir][Vector] Use a RewriterBase for IR rewrites in VectorTransferOpTransforms

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
    mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
    mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
index 52a4c9cc368d8..d8a08784618da 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
@@ -104,7 +104,7 @@ LogicalResult splitFullAndPartialTransfer(
 
 /// Implements transfer op write to read forwarding and dead transfer write
 /// optimizations.
-void transferOpflowOpt(Operation *rootOp);
+void transferOpflowOpt(RewriterBase &rewriter, Operation *rootOp);
 
 } // namespace vector
 } // namespace mlir

diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
index b0690f63422d9..cbfea2fe62203 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -45,16 +45,18 @@ namespace {
 
 class TransferOptimization {
 public:
-  TransferOptimization(Operation *op) : dominators(op), postDominators(op) {}
+  TransferOptimization(RewriterBase &rewriter, Operation *op)
+      : rewriter(rewriter), dominators(op), postDominators(op) {}
   void deadStoreOp(vector::TransferWriteOp);
   void storeToLoadForwarding(vector::TransferReadOp);
   void removeDeadOp() {
     for (Operation *op : opToErase)
-      op->erase();
+      rewriter.eraseOp(op);
     opToErase.clear();
   }
 
 private:
+  RewriterBase &rewriter;
   bool isReachable(Operation *start, Operation *dest);
   DominanceInfo dominators;
   PostDominanceInfo postDominators;
@@ -724,8 +726,9 @@ class RewriteScalarWrite : public OpRewritePattern<vector::TransferWriteOp> {
 };
 } // namespace
 
-void mlir::vector::transferOpflowOpt(Operation *rootOp) {
-  TransferOptimization opt(rootOp);
+void mlir::vector::transferOpflowOpt(RewriterBase &rewriter,
+                                     Operation *rootOp) {
+  TransferOptimization opt(rewriter, rootOp);
   // Run store to load forwarding first since it can expose more dead store
   // opportunity.
   rootOp->walk([&](vector::TransferReadOp read) {

diff  --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index 4d177e7c21184..c535cd4d14513 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -325,7 +325,10 @@ struct TestVectorTransferOpt
   StringRef getDescription() const final {
     return "Test optimization transformations for transfer ops";
   }
-  void runOnOperation() override { transferOpflowOpt(getOperation()); }
+  void runOnOperation() override {
+    IRRewriter rewriter(&getContext());
+    transferOpflowOpt(rewriter, getOperation());
+  }
 };
 
 struct TestVectorTransferCollapseInnerMostContiguousDims