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

[springerm]

Summary

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.

[ingomueller-net]

Thanks for the fix! LGTM!

mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp

293

Just to spell it out for me: this is a trivial case because setListener is called a few lines after the construction of the rewriter, so no other listener can be attached already and replaced by the call, right?

[springerm]

mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp

293

Yes that's correct.

mlir/include/mlir/Transforms/OneToNTypeConversion.h

protected:
   TypeConverter *const typeConverter;
 };
 
 /// Specialization of `PatternRewriter` that `OneToNConversionPattern`s use. The
 /// class provides additional rewrite methods that are specific to 1:N type
 /// 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
   /// mapping. That type mapping must match the replaced op (i.e., the original
   /// types must be the same as the result types of the op) and the new values
   /// (i.e., the converted types must be the same as the types of the new
   /// values).
   void replaceOp(Operation *op, ValueRange newValues,

mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp

LogicalResult matchAndRewrite(OpTy op,
     if (!op.getIfCond())
       return failure();
 
     IntegerAttr constAttr;
     if (!matchPattern(op.getIfCond(), m_Constant(&constAttr))) {
       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 {
       if (constAttr.getInt())
         rewriter.updateRootInPlace(op,
                                    [&]() { op.getIfCondMutable().erase(0); });
       else
         rewriter.eraseOp(op);

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 }
 
 OpFoldResult
 mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                             AffineMap map,
                                             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());
   for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
     matchPattern(applyOp->getOperand(i), m_Constant(&constOperands[i]));
 
   // Try to fold the operation.
   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();
   }
 
   applyOp->erase();
   assert(foldResults.size() == 1 && "expected 1 folded result");
   return foldResults.front();
 }
mlir::affine::makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
                                     ArrayRef<OpFoldResult> operands) {
   return makeComposedMinMax<AffineMinOp>(b, loc, map, operands);
 }
 
 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());
   for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
     matchPattern(minMaxOp->getOperand(i), m_Constant(&constOperands[i]));
 
   // Try to fold the operation.
   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();
   }
 
   minMaxOp->erase();
   assert(foldResults.size() == 1 && "expected 1 folded result");
   return foldResults.front();
 }

mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp

OneToNConversionPattern::matchAndRewrite(Operation *op,
 
   // Create a `OneToNPatternRewriter` for the pattern, which provides additional
   // functionality.
   // TODO(ingomueller): I guess it would be better to use only one rewriter
   //                    throughout the whole pass, but that would require to
   //                    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,
                              resultMapping, convertedOperands)))
     return failure();
 
   return success();
 }