[mlir][transform] TrackingListener: Drop mappings of tracked ops when all handles are dead

[springerm]

Summary

No replacement ops are needed for tracked ops who's handles are all dead.

Depends On: D147447

[nicolasvasilache]

Nice!

[ftynse]

Longer-term, we may want to have an analysis that finds last use of all handles instead of checking it for every handle and every op.

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h

 namespace mlir {
 namespace tensor {
 
 /// A specialized TrackingListener for transform ops that operate on tensor IR.
 /// This listener skips cast-like tensor ops when looking for payload op
 /// replacements.
 class TrackingListener : public transform::TrackingListener {
 public:
-  explicit TrackingListener(transform::TransformState &state)
-      : transform::TrackingListener(state) {}
+  using transform::TrackingListener::TrackingListener;
 
 protected:
   Operation *findReplacementOp(Operation *op,
                                ValueRange newValues) const override;
 };
 
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H

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

::llvm::function_ref<void(::mlir::OpBuilder &, ::mlir::Location,
                               ::mlir::BlockArgument, ::mlir::ValueRange)>;
 
 /// A listener that updates a TransformState based on IR modifications. This
 /// listener can be used during a greedy pattern rewrite to keep the transform
 /// state up-to-date.
 class TrackingListener : public RewriterBase::Listener,
                          public TransformState::Extension {
 public:
-  explicit TrackingListener(TransformState &state)
-      : TransformState::Extension(state) {}
+  /// 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.
   virtual Operation *findReplacementOp(Operation *op,
                                        ValueRange newValues) const;
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/Linalg/TransformOps/LinalgTransformOps.cpp

 }
 
 DiagnosedSilenceableFailure
 transform::HoistPadOp::applyToOne(tensor::PadOp target,
                                   transform::ApplyToEachResultList &results,
                                   transform::TransformState &state) {
   tensor::PadOp hoistedPadOp;
   SmallVector<GenericOp> transposeOps;
-  TrackingListener listener(state);
+  TrackingListener listener(state, *this);
   IRRewriter rewriter(target->getContext(), &listener);
   FailureOr<Value> result =
       hoistPaddingOnTensors(rewriter, target, getNumLoops(), getTranspose(),
                             hoistedPadOp, transposeOps);
   if (succeeded(result)) {
     // We need to perform our own replacement here because this API is still
     // used in patterns that "pad and hoist", for which the replacement values
     // need to be different.
transform::VectorizeOp::applyToOne(Operation *target,
   vector::TransferReadOp::getCanonicalizationPatterns(patterns, ctx);
   vector::TransferWriteOp::getCanonicalizationPatterns(patterns, ctx);
 
   patterns.add<CopyVectorizationPattern>(ctx);
 
   if (getVectorizePadding())
     linalg::populatePadOpVectorizationPatterns(patterns);
 
-  TrackingListener listener(state);
+  TrackingListener listener(state, *this);
   GreedyRewriteConfig config;
   config.listener = &listener;
   if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns), config)))
     return emitDefaultDefiniteFailure(target);
 
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }

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

op->walk([&](Operation *op) {
     // Remove mappings for result values.
     for (OpResult value : op->getResults())
       (void)replacePayloadValue(value, nullptr);
     // Remove mapping for op.
     (void)replacePayloadOp(op, nullptr);
   });
 }
 
+/// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
+/// properly dominates `b` and `b` is not inside `a`.
+static bool happensBefore(Operation *a, Operation *b) {
+  do {
+    if (a->isProperAncestor(b))
+      return false;
+    if (Operation *bAncestor = a->getBlock()->findAncestorOpInBlock(*b)) {
+      return a->isBeforeInBlock(bAncestor);
+    }
+  } while ((a = a->getParentOp()));
+  return false;
+}
+
 void transform::TrackingListener::notifyOperationReplaced(
     Operation *op, ValueRange newValues) {
   assert(op->getNumResults() == newValues.size() &&
          "invalid number of replacement values");
 
   // Replace value handles.
   for (auto [oldValue, newValue] : llvm::zip(op->getResults(), newValues))
     (void)replacePayloadValue(oldValue, newValue);
 
   // Replace op handle.
+  SmallVector<Value> opHandles;
+  if (failed(getTransformState().getHandlesForPayloadOp(op, opHandles))) {
+    // Op is not tracked.
+    return;
+  }
+  auto hasAliveUser = [&]() {
+    for (Value v : opHandles)
+      for (Operation *user : v.getUsers())
+        if (!happensBefore(user, transformOp))
+          return true;
+    return false;
+  };
+  if (!hasAliveUser()) {
+    // The op is tracked but the corresponding handles are dead.
+    (void)replacePayloadOp(op, nullptr);
+    return;
+  }
+
   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);
-  }
+  (void)replacePayloadOp(op, replacement);
 }
 
 //===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
 OperandRange transform::AlternativesOp::getSuccessorEntryOperands(
     std::optional<unsigned> index) {
for (Region &reg : getAlternatives()) {
     }
 
     // If all operations in the given alternative succeeded, no need to consider
     // the rest. Replace the original scoping operation with the clone on which
     // the transformations were performed.
     if (!failed) {
       // We will be using the clones, so cancel their scheduled deletion.
       deleteClones.release();
-      TrackingListener listener(state);
+      TrackingListener listener(state, *this);
       IRRewriter rewriter(getContext(), &listener);
       for (const auto &kvp : llvm::zip(originals, clones)) {
         Operation *original = std::get<0>(kvp);
         Operation *clone = std::get<1>(kvp);
         original->getBlock()->getOperations().insert(original->getIterator(),
                                                      clone);
         rewriter.replaceOp(original, clone->getResults());
       }