[Mlir-commits] [mlir] 9d072bb - [mlir][NFC] Avoid `OpBuilder::setListener` when possible

[Matthias Springer]

Author: Matthias Springer
Date: 2023-07-19T09:13:38+02:00
New Revision: 9d072bbe0f0725eb22d2f2d175c64f3a9cbfb788

URL: https://github.com/llvm/llvm-project/commit/9d072bbe0f0725eb22d2f2d175c64f3a9cbfb788
DIFF: https://github.com/llvm/llvm-project/commit/9d072bbe0f0725eb22d2f2d175c64f3a9cbfb788.diff

LOG: [mlir][NFC] Avoid `OpBuilder::setListener` when possible

`setListener` is dangerous because an already registered listener may accidentally be overwritten/replaced. (A `ForwardingListener` must be used in such cases.) This change updates a few trivial call sites of `setListener`, where no forwarding listener is needed.

Differential Revision: https://reviews.llvm.org/D155599

Added: 
    

Modified: 
    mlir/include/mlir/Transforms/OneToNTypeConversion.h
    mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/OneToNTypeConversion.h b/mlir/include/mlir/Transforms/OneToNTypeConversion.h
index 47bd276eed61f5..b86a6b5f76a3bc 100644
--- a/mlir/include/mlir/Transforms/OneToNTypeConversion.h
+++ b/mlir/include/mlir/Transforms/OneToNTypeConversion.h
@@ -151,7 +151,9 @@ class RewritePatternWithConverter : public mlir::RewritePattern {
 /// conversions.
 class OneToNPatternRewriter : public PatternRewriter {
 public:
-  OneToNPatternRewriter(MLIRContext *context) : PatternRewriter(context) {}
+  OneToNPatternRewriter(MLIRContext *context,
+                        OpBuilder::Listener *listener = nullptr)
+      : PatternRewriter(context, listener) {}
 
   /// Replaces the results of the operation with the specified list of values
   /// mapped back to the original types as specified in the provided type

diff  --git a/mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp b/mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp
index b42da60f4c89c1..8c1a7d9c6b2a43 100644
--- a/mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp
+++ b/mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp
@@ -44,8 +44,7 @@ class ExpandIfCondition : public OpRewritePattern<OpTy> {
       auto ifOp = rewriter.create<scf::IfOp>(op.getLoc(), TypeRange(),
                                              op.getIfCond(), false);
       rewriter.updateRootInPlace(op, [&]() { op.getIfCondMutable().erase(0); });
-      auto thenBodyBuilder = ifOp.getThenBodyBuilder();
-      thenBodyBuilder.setListener(rewriter.getListener());
+      auto thenBodyBuilder = ifOp.getThenBodyBuilder(rewriter.getListener());
       thenBodyBuilder.clone(*op.getOperation());
       rewriter.eraseOp(op);
     } else {

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 0e8e84994963f8..81f331fd6c2400 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1278,17 +1278,16 @@ mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                             ArrayRef<OpFoldResult> operands) {
   assert(map.getNumResults() == 1 && "building affine.apply with !=1 result");
 
-  // Temporarily disconnect the listener, so that no notification is triggered
-  // if the op is folded.
+  // Create new builder without a listener, so that no notification is
+  // triggered if the op is folded.
   // TODO: OpBuilder::createOrFold should return OpFoldResults, then this
   // workaround is no longer needed.
-  OpBuilder::Listener *listener = b.getListener();
-  b.setListener(nullptr);
-  auto listenerResetter =
-      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
+  OpBuilder newBuilder(b.getContext());
+  newBuilder.setInsertionPoint(b.getInsertionBlock(), b.getInsertionPoint());
 
   // Create op.
-  AffineApplyOp applyOp = makeComposedAffineApply(b, loc, map, operands);
+  AffineApplyOp applyOp =
+      makeComposedAffineApply(newBuilder, loc, map, operands);
 
   // Get constant operands.
   SmallVector<Attribute> constOperands(applyOp->getNumOperands());
@@ -1299,7 +1298,7 @@ mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
   SmallVector<OpFoldResult> foldResults;
   if (failed(applyOp->fold(constOperands, foldResults)) ||
       foldResults.empty()) {
-    if (listener)
+    if (OpBuilder::Listener *listener = b.getListener())
       listener->notifyOperationInserted(applyOp);
     return applyOp.getResult();
   }
@@ -1347,17 +1346,15 @@ template <typename OpTy>
 static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc,
                                              AffineMap map,
                                              ArrayRef<OpFoldResult> operands) {
-  // Temporarily disconnect the listener, so that no notification is triggered
-  // if the op is folded.
+  // Create new builder without a listener, so that no notification is
+  // triggered if the op is folded.
   // TODO: OpBuilder::createOrFold should return OpFoldResults, then this
   // workaround is no longer needed.
-  OpBuilder::Listener *listener = b.getListener();
-  b.setListener(nullptr);
-  auto listenerResetter =
-      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
+  OpBuilder newBuilder(b.getContext());
+  newBuilder.setInsertionPoint(b.getInsertionBlock(), b.getInsertionPoint());
 
   // Create op.
-  auto minMaxOp = makeComposedMinMax<OpTy>(b, loc, map, operands);
+  auto minMaxOp = makeComposedMinMax<OpTy>(newBuilder, loc, map, operands);
 
   // Get constant operands.
   SmallVector<Attribute> constOperands(minMaxOp->getNumOperands());
@@ -1368,7 +1365,7 @@ static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc,
   SmallVector<OpFoldResult> foldResults;
   if (failed(minMaxOp->fold(constOperands, foldResults)) ||
       foldResults.empty()) {
-    if (listener)
+    if (OpBuilder::Listener *listener = b.getListener())
       listener->notifyOperationInserted(minMaxOp);
     return minMaxOp.getResult();
   }

diff  --git a/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp b/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp
index 8c4df4bf8fe4d6..c95aa608636d8a 100644
--- a/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp
+++ b/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp
@@ -289,9 +289,9 @@ OneToNConversionPattern::matchAndRewrite(Operation *op,
   //                    drive the pattern application ourselves, which is a lot
   //                    of additional boilerplate code. This seems to work fine,
   //                    so I leave it like this for the time being.
-  OneToNPatternRewriter oneToNPatternRewriter(rewriter.getContext());
+  OneToNPatternRewriter oneToNPatternRewriter(rewriter.getContext(),
+                                              rewriter.getListener());
   oneToNPatternRewriter.restoreInsertionPoint(rewriter.saveInsertionPoint());
-  oneToNPatternRewriter.setListener(rewriter.getListener());
 
   // Apply actual pattern.
   if (failed(matchAndRewrite(op, oneToNPatternRewriter, operandMapping,