[mlir][transform] Add optional error checking to TrackingListener

[springerm]

Summary

Derived classes can implement notifyPayloadReplacementNotFound for custom error checking.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.h

 protected:
   /// Return a replacement payload op for the given op, which is going to be
   /// replaced with the given values. By default, if all values are defined by
   /// the same newly-created op, which also has the same type as the given op,
   /// that defining op is used as a replacement.
   virtual Operation *findReplacementOp(Operation *op,
                                        ValueRange newValues) const;
 
+  /// This function is called when a tracked payload op is dropped because no
+  /// replacement op was found. Derived classes can implement this function for
+  /// custom error handling.
+  virtual void notifyPayloadReplacementNotFound(Operation *op,
+                                                ValueRange values) const {}
+
   /// Return "true" if the given op is a new op.
   bool isNewOp(Operation *op) const;
 
   /// Return the single op that defines all given values (if any).
   static Operation *getCommonDefiningOp(ValueRange values);
 
 private:
   void notifyOperationInserted(Operation *op) override;

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

op->walk([&](Operation *op) {
     (void)replacePayloadOp(op, nullptr);
   });
 }
 
 void transform::TrackingListener::notifyOperationReplaced(
     Operation *op, ValueRange newValues) {
   assert(op->getNumResults() == newValues.size() &&
          "invalid number of replacement values");
-  if (op->getNumResults() == 0)
-    return;
 
   // Replace value handles.
   for (auto [oldValue, newValue] : llvm::zip(op->getResults(), newValues))
     (void)replacePayloadValue(oldValue, newValue);
 
   // Replace op handle.
-  (void)replacePayloadOp(op, findReplacementOp(op, newValues));
+  Operation *replacement = findReplacementOp(op, newValues);
+  if (succeeded(replacePayloadOp(op, replacement))) {
+    // If the op is tracked but no replacement op was found, send a
+    // notification.
+    if (!replacement)
+      notifyPayloadReplacementNotFound(op, newValues);
+  }
 }
 
 //===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
 OperandRange transform::AlternativesOp::getSuccessorEntryOperands(
     std::optional<unsigned> index) {