[mlir][transform] TrackingListener: Allow existing ops as replacements

[springerm]

Summary

The TrackingListener was unnecessarily strict. Existing ops are now allowed when updating payload ops mappings due to replaceOp in the TrackingListener.

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

 public:
   /// Create a new TrackingListener for usage in the specified transform op.
   explicit TrackingListener(TransformState &state, TransformOpInterface op)
       : TransformState::Extension(state), transformOp(op) {}
 
 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.
+  /// the same 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;
 
   /// 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.
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
 
   /// 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) {}
 
-  /// 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;
-
   void notifyOperationRemoved(Operation *op) override;
 
   void notifyOperationReplaced(Operation *op, ValueRange newValues) override;
 
-  /// Ops that were newly created during the transform.
-  DenseMap<OperationName, DenseSet<Operation *>> newOps;
-
   /// The transform op in which this TrackingListener is used.
   TransformOpInterface transformOp;
 };
 
 } // namespace transform
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.h.inc"
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H

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

transform::TrackingListener::findReplacementOp(Operation *op,
   Operation *defOp = getCommonDefiningOp(newValues);
   if (!defOp)
     return nullptr;
 
   // If the replacement op has a different type, drop the mapping.
   if (op->getName() != defOp->getName())
     return nullptr;
 
-  // If the replacement op is not a new op, drop the mapping.
-  if (!isNewOp(defOp))
-    return nullptr;
-
   return defOp;
 }
 
-bool transform::TrackingListener::isNewOp(Operation *op) const {
-  auto it = newOps.find(op->getName());
-  if (it == newOps.end())
-    return false;
-  return it->second.contains(op);
-}
-
 LogicalResult transform::TrackingListener::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     DBGS() << "Match Failure : " << diag.str() << "\n";
   });
   return failure();
 }
 
-void transform::TrackingListener::notifyOperationInserted(Operation *op) {
-  newOps[op->getName()].insert(op);
-}
-
 void transform::TrackingListener::notifyOperationRemoved(Operation *op) {
   // TODO: Walk can be removed when D144193 has landed.
   op->walk([&](Operation *op) {
-    // Keep set of new ops up-to-date.
-    auto it = newOps.find(op->getName());
-    if (it != newOps.end())
-      it->second.erase(op);
     // Remove mappings for result values.
     for (OpResult value : op->getResults())
       (void)replacePayloadValue(value, nullptr);
     // Remove mapping for op.
     (void)replacePayloadOp(op, nullptr);
   });
 }