[mlir][IR] Improve listener notifications for ops without results

[springerm]

Summary

RewriterBase::Listener::notifyOperationReplaced notifies observers that an op is about to be replaced with a range of values. This notification is not very useful for ops without results, because it does not specify the replacement op (and it cannot be deduced from the replacement values). It provides no additional information over the notifyOperationRemoved notification.

This revision adds an additional notification when a rewriter replaces an op with another op. By default, this notification triggers the original "op replaced with values" notification, so there is no functional change for existing code.

This new API is useful for the transform dialect, which needs to track op replacements. (Updated in a subsequent revision.)

Also includes minor documentation improvements.

mlir/include/mlir/IR/PatternMatch.h

public:
   struct Listener : public OpBuilder::Listener {
     Listener()
         : OpBuilder::Listener(ListenerBase::Kind::RewriterBaseListener) {}
 
     /// Notify the listener that the specified operation was modified in-place.
     virtual void notifyOperationModified(Operation *op) {}
 
     /// Notify the listener that the specified operation is about to be replaced
-    /// with the set of values potentially produced by new operations. This is
-    /// called before the uses of the operation have been changed.
+    /// with another operation. This is called before the uses of the old
+    /// operation have been changed.
+    ///
+    /// By default, this function calls the "operation replaced with values"
+    /// notification.
+    virtual void notifyOperationReplaced(Operation *op,
+                                         Operation *replacement) {
+      notifyOperationReplaced(op, replacement->getResults());
+    }
+
+    /// Notify the listener that the specified operation is about to be replaced
+    /// with the a range of values, potentially produced by other operations.
+    /// This is called before the uses of the operation have been changed.
     virtual void notifyOperationReplaced(Operation *op,
                                          ValueRange replacement) {}
 
-    /// This is called on an operation that a rewrite is removing, right before
-    /// the operation is deleted. At this point, the operation has zero uses.
+    /// Notify the listener that the specified operation is about to be erased.
+    /// At this point, the operation has zero uses.
     virtual void notifyOperationRemoved(Operation *op) {}
 
     /// Notify the listener that the pattern failed to match the given
     /// operation, and provide a callback to populate a diagnostic with the
     /// reason why the failure occurred. This method allows for derived
     /// listeners to optionally hook into the reason why a rewrite failed, and
     /// display it to users.
     virtual LogicalResult
void notifyOperationInserted(Operation *op) override {
       listener->notifyOperationInserted(op);
     }
     void notifyBlockCreated(Block *block) override {
       listener->notifyBlockCreated(block);
     }
     void notifyOperationModified(Operation *op) override {
       listener->notifyOperationModified(op);
     }
+    void notifyOperationReplaced(Operation *op, Operation *newOp) override {
+      listener->notifyOperationReplaced(op, newOp);
+    }
     void notifyOperationReplaced(Operation *op,
                                  ValueRange replacement) override {
       listener->notifyOperationReplaced(op, replacement);
     }
     void notifyOperationRemoved(Operation *op) override {
       listener->notifyOperationRemoved(op);
     }
     LogicalResult notifyMatchFailure(
public:
   /// `newValues` when a use is nested within the given `block`. The number of
   /// values in `newValues` is required to match the number of results of `op`.
   /// If all uses of this operation are replaced, the operation is erased.
   void replaceOpWithinBlock(Operation *op, ValueRange newValues, Block *block,
                             bool *allUsesReplaced = nullptr);
 
   /// This method replaces the results of the operation with the specified list
   /// of values. The number of provided values must match the number of results
-  /// of the operation.
+  /// of the operation. The replaced op is erased.
   virtual void replaceOp(Operation *op, ValueRange newValues);
 
+  /// This method replaces the results of the operation with the specified
+  /// new op (replacement). The number of results of the two operations must
+  /// match. The replaced op is erased.
+  virtual void replaceOp(Operation *op, Operation *newOp);
+
   /// Replaces the result op with a new op that is created without verification.
   /// The result values of the two ops must be the same types.
   template <typename OpTy, typename... Args>
   OpTy replaceOpWithNewOp(Operation *op, Args &&...args) {
     auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
-    replaceOpWithResultsOfAnotherOp(op, newOp.getOperation());
+    replaceOp(op, newOp.getOperation());
     return newOp;
   }
 
   /// This method erases an operation that is known to have no uses.
   virtual void eraseOp(Operation *op);
 
   /// This method erases all operations in a block.
   virtual void eraseBlock(Block *block);
explicit RewriterBase(MLIRContext *ctx,
       : OpBuilder(ctx, listener) {}
   explicit RewriterBase(const OpBuilder &otherBuilder)
       : OpBuilder(otherBuilder) {}
   virtual ~RewriterBase();
 
 private:
   void operator=(const RewriterBase &) = delete;
   RewriterBase(const RewriterBase &) = delete;
-
-  /// 'op' and 'newOp' are known to have the same number of results, replace the
-  /// uses of op with uses of newOp.
-  void replaceOpWithResultsOfAnotherOp(Operation *op, Operation *newOp);
 };
 
 //===----------------------------------------------------------------------===//
 // IRRewriter
 //===----------------------------------------------------------------------===//
 
 /// This class coordinates rewriting a piece of IR outside of a pattern rewrite,
 /// providing a way to keep track of the mutations made to the IR. This class

mlir/include/mlir/Transforms/DialectConversion.h

public:
   // PatternRewriter Hooks
   //===--------------------------------------------------------------------===//
 
   /// Indicate that the conversion rewriter can recover from rewrite failure.
   /// Recovery is supported via rollback, allowing for continued processing of
   /// patterns even if a failure is encountered during the rewrite step.
   bool canRecoverFromRewriteFailure() const override { return true; }
 
-  /// PatternRewriter hook for replacing the results of an operation when the
-  /// given functor returns true.
+  /// PatternRewriter hook for replacing an operation when the given functor
+  /// returns "true".
   void replaceOpWithIf(
       Operation *op, ValueRange newValues, bool *allUsesReplaced,
       llvm::unique_function<bool(OpOperand &) const> functor) override;
 
-  /// PatternRewriter hook for replacing the results of an operation.
+  /// PatternRewriter hook for replacing an operation.
   void replaceOp(Operation *op, ValueRange newValues) override;
-  using PatternRewriter::replaceOp;
+
+  /// PatternRewriter hook for replacing an operation.
+  void replaceOp(Operation *op, Operation *newOp) override;
 
   /// PatternRewriter hook for erasing a dead operation. The uses of this
   /// operation *must* be made dead by the end of the conversion process,
   /// otherwise an assert will be issued.
   void eraseOp(Operation *op) override;
 
   /// PatternRewriter hook for erase all operations in a block. This is not yet
   /// implemented for dialect conversion.

mlir/lib/Dialect/AMDGPU/Transforms/EmulateAtomics.cpp

prevLoadForCompare =
         rewriter.create<arith::BitcastOp>(loc, equivInt, prevLoad);
     atomicResForCompare =
         rewriter.create<arith::BitcastOp>(loc, equivInt, atomicRes);
   }
   Value canLeave = rewriter.create<arith::CmpIOp>(
       loc, arith::CmpIPredicate::eq, atomicResForCompare, prevLoadForCompare);
   rewriter.create<cf::CondBranchOp>(loc, canLeave, afterAtomic, ValueRange{},
                                     loopBlock, atomicRes);
-  rewriter.replaceOp(atomicOp, {});
+  rewriter.eraseOp(atomicOp);
   return success();
 }
 
 void mlir::amdgpu::populateAmdgpuEmulateAtomicsPatterns(
     ConversionTarget &target, RewritePatternSet &patterns, Chipset chipset) {
   // gfx10 has no atomic adds.
   if (chipset.majorVersion == 10 || chipset.majorVersion < 9 ||
       (chipset.majorVersion == 9 && chipset.minorVersion < 0x08)) {

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

LogicalResult matchAndRewrite(AllocLikeOp alloc,
     assert(static_cast<int64_t>(dynamicSizes.size()) ==
            newMemRefType.getNumDynamicDims());
 
     // Create and insert the alloc op for the new memref.
     auto newAlloc = rewriter.create<AllocLikeOp>(
         alloc.getLoc(), newMemRefType, dynamicSizes, alloc.getSymbolOperands(),
         alloc.getAlignmentAttr());
     // Insert a cast so we have the same type as the old alloc.
-    auto resultCast =
-        rewriter.create<CastOp>(alloc.getLoc(), alloc.getType(), newAlloc);
-
-    rewriter.replaceOp(alloc, {resultCast});
+    rewriter.replaceOpWithNewOp<CastOp>(alloc, alloc.getType(), newAlloc);
     return success();
   }
 };
 
 /// Fold alloc operations with no users or only store and dealloc uses.
 template <typename T>
 struct SimplifyDeadAlloc : public OpRewritePattern<T> {
   using OpRewritePattern<T>::OpRewritePattern;

mlir/lib/IR/PatternMatch.cpp

void RewriterBase::replaceOpWithinBlock(Operation *op, ValueRange newValues,
                                         Block *block, bool *allUsesReplaced) {
   replaceOpWithIf(op, newValues, allUsesReplaced, [block](OpOperand &use) {
     return block->getParentOp()->isProperAncestor(use.getOwner());
   });
 }
 
 /// This method replaces the results of the operation with the specified list of
 /// values. The number of provided values must match the number of results of
-/// the operation.
+/// the operation. The replaced op is erased.
 void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
 
-  // Notify the listener that we're about to remove this op.
+  // Notify the listener that we're about to replace this op.
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
     rewriteListener->notifyOperationReplaced(op, newValues);
 
   // Replace results one-by-one. Also notifies the listener of modifications.
   for (auto it : llvm::zip(op->getResults(), newValues))
     replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
 
+  // Erase the op.
+  eraseOp(op);
+}
+
+/// This method replaces the results of the operation with the specified new op
+/// (replacement). The number of results of the two operations must match. The
+/// replaced op is erased.
+void RewriterBase::replaceOp(Operation *op, Operation *newOp) {
+  assert(op && newOp && "expected non-null op");
+  assert(op->getNumResults() == newOp->getNumResults() &&
+         "ops have different number of results");
+
+  // Notify the listener that we're about to replace this op.
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
-    rewriteListener->notifyOperationRemoved(op);
-  op->erase();
+    rewriteListener->notifyOperationReplaced(op, newOp);
+
+  // Replace results one-by-one. Also notifies the listener of modifications.
+  for (auto it : llvm::zip(op->getResults(), newOp->getResults()))
+    replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
+
+  // Erase the old op.
+  eraseOp(op);
 }
 
 /// This method erases an operation that is known to have no uses. The uses of
 /// the given operation *must* be known to be dead.
 void RewriterBase::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
     rewriteListener->notifyOperationRemoved(op);

 }
 
 /// Split the operations starting at "before" (inclusive) out of the given
 /// block into a new block, and return it.
 Block *RewriterBase::splitBlock(Block *block, Block::iterator before) {
   return block->splitBlock(before);
 }
 
-/// 'op' and 'newOp' are known to have the same number of results, replace the
-/// uses of op with uses of newOp
-void RewriterBase::replaceOpWithResultsOfAnotherOp(Operation *op,
-                                                   Operation *newOp) {
-  assert(op->getNumResults() == newOp->getNumResults() &&
-         "replacement op doesn't match results of original op");
-  if (op->getNumResults() == 1)
-    return replaceOp(op, newOp->getResult(0));
-  return replaceOp(op, newOp->getResults());
-}
-
 /// Move the blocks that belong to "region" before the given position in
 /// another region.  The two regions must be different.  The caller is in
 /// charge to update create the operation transferring the control flow to the
 /// region and pass it the correct block arguments.
 void RewriterBase::inlineRegionBefore(Region &region, Region &parent,
                                       Region::iterator before) {
   parent.getBlocks().splice(before, region.getBlocks());
 }

mlir/lib/Transforms/Utils/DialectConversion.cpp

void ConversionPatternRewriter::replaceOpWithIf(
   // operation. The main change is that now an operation can be replaced
   // multiple times, in parts. The current "set" based tracking is mainly useful
   // for tracking if a replaced operation should be ignored, i.e. if all of the
   // uses will be replaced.
   llvm_unreachable(
       "replaceOpWithIf is currently not supported by DialectConversion");
 }
 
+void ConversionPatternRewriter::replaceOp(Operation *op, Operation *newOp) {
+  assert(op && newOp && "expected non-null op");
+  replaceOp(op, newOp->getResults());
+}
+
 void ConversionPatternRewriter::replaceOp(Operation *op, ValueRange newValues) {
+  assert(op->getNumResults() == newValues.size() &&
+         "incorrect # of replacement values");
   LLVM_DEBUG({
     impl->logger.startLine()
         << "** Replace : '" << op->getName() << "'(" << op << ")\n";
   });
   impl->notifyOpReplaced(op, newValues);
 }
 
 void ConversionPatternRewriter::eraseOp(Operation *op) {

mlir/test/lib/Dialect/Test/TestPatterns.cpp

struct TestCreateBlock : public RewritePattern {
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const final {
     Region &region = *op->getParentRegion();
     Type i32Type = rewriter.getIntegerType(32);
     Location loc = op->getLoc();
     rewriter.createBlock(&region, region.end(), {i32Type, i32Type}, {loc, loc});
     rewriter.create<TerminatorOp>(loc);
-    rewriter.replaceOp(op, {});
+    rewriter.eraseOp(op);
     return success();
   }
 };
 
 /// A simple pattern that creates a block containing an invalid operation in
 /// order to trigger the block creation undo mechanism.
 struct TestCreateIllegalBlock : public RewritePattern {
   TestCreateIllegalBlock(MLIRContext *ctx)
       : RewritePattern("test.create_illegal_block", /*benefit=*/1, ctx) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const final {
     Region &region = *op->getParentRegion();
     Type i32Type = rewriter.getIntegerType(32);
     Location loc = op->getLoc();
     rewriter.createBlock(&region, region.end(), {i32Type, i32Type}, {loc, loc});
     // Create an illegal op to ensure the conversion fails.
     rewriter.create<ILLegalOpF>(loc, i32Type);
     rewriter.create<TerminatorOp>(loc);
-    rewriter.replaceOp(op, {});
+    rewriter.eraseOp(op);
     return success();
   }
 };
 
 /// A simple pattern that tests the undo mechanism when replacing the uses of a
 /// block argument.
 struct TestUndoBlockArgReplace : public ConversionPattern {
   TestUndoBlockArgReplace(MLIRContext *ctx)
TestNonRootReplacement(MLIRContext *ctx)
       : RewritePattern("test.replace_non_root", 1, ctx) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const final {
     auto resultType = *op->result_type_begin();
     auto illegalOp = rewriter.create<ILLegalOpF>(op->getLoc(), resultType);
     auto legalOp = rewriter.create<LegalOpB>(op->getLoc(), resultType);
 
-    rewriter.replaceOp(illegalOp, {legalOp});
-    rewriter.replaceOp(op, {illegalOp});
+    rewriter.replaceOp(illegalOp, legalOp);
+    rewriter.replaceOp(op, illegalOp);
     return success();
   }
 };
 
 //===----------------------------------------------------------------------===//
 // Recursive Rewrite Testing
 /// This pattern is applied to the same operation multiple times, but has a
 /// bounded recursion.