[mlir][transform] TrackingListener: Support ops without results

[springerm]

Summary

When a tracked op is replaced with another op, the internal transform dialect mapping state should be updated. Until now, this was not the case when an op with 0 results was replaced because there was no RewriterBase::Listener callback that provides the replacement op (only the replacement values).

This revision adds a new notifyOperationReplaced(Operation *op, Operation *replacement) notification to the listener infrastructure. The existing notifyOperationReplaced(Operation *, ValueRange) notification is not very useful when replacing an op with 0 results, because it provides no additional information over the notifyOperationRemoved notification (which is triggered afterwards). No changes for existing users ofthe listener API: the behavior of the existing notifyOperationReplaced does not change.

The new notification is utilized in the transform dialect to support payload op replacements of ops with 0 results.

Depends On: D152427

[springerm]

rebase

[springerm]

update

[springerm]

I have to rebase this on top of D152814.

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

protected:
   static Operation *getCommonDefiningOp(ValueRange values);
 
   /// Return the transform op in which this TrackingListener is used.
   TransformOpInterface getTransformOp() const { return transformOp; }
 
 private:
   void notifyOperationRemoved(Operation *op) override;
 
-  void notifyOperationReplaced(Operation *op, ValueRange newValues) override;
+  void notifyOperationReplaced(Operation *op, Operation *replacement) override;
+
+  void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
+
+  /// Return "true" if a payload op (that is currently in the process of being
+  /// replaced) that is mapped to given handles can be silently dropped from the
+  /// mapping. This is the case for dead/consumed handles.
+  bool canSilentlyDropPayloadOpFromMapping(ValueRange handles) const;
 
   /// The transform op in which this TrackingListener is used.
   TransformOpInterface transformOp;
 
   /// The handles that are consumed by the transform op.
   DenseSet<Value> consumedHandles;
 };
 

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 another operation. This is called before the uses of the operation
+    /// have changed.
+    ///
+    /// Note: The "op replaced with values" notification is triggered in
+    /// addition to this notification.
+    virtual void notifyOperationReplaced(Operation *op,
+                                         Operation *replacement) {}
+
+    /// 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.
+    ///
+    /// Note: The list of replacement values is empty in case the specified op
+    /// has no results.
     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.
     virtual void notifyOperationRemoved(Operation *op) {}
 
     /// Notify the listener that the pattern failed to match the given
struct ForwardingListener : public RewriterBase::Listener {
     }
     void notifyBlockCreated(Block *block) override {
       listener->notifyBlockCreated(block);
     }
     void notifyOperationModified(Operation *op) override {
       listener->notifyOperationModified(op);
     }
     void notifyOperationReplaced(Operation *op,
+                                 Operation *replacement) override {
+      listener->notifyOperationReplaced(op, replacement);
+    }
+    void notifyOperationReplaced(Operation *op,
                                  ValueRange replacement) override {
       listener->notifyOperationReplaced(op, replacement);
     }
     void notifyOperationRemoved(Operation *op) override {
       listener->notifyOperationRemoved(op);
     }
     LogicalResult notifyMatchFailure(
         Location loc,

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

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");
+bool transform::TrackingListener::canSilentlyDropPayloadOpFromMapping(
+    ValueRange handles) const {
+  // Check if any of the handles is consumed by the current transform op.
+  //
+  // Note: If a handle was consumed, there shouldn't be any alive handles
+  // mapped to the payload op. So it is not really necessary to check for
+  // consumed handles here. However, in case there are indeed alive handles
+  // mapped to the payload op (which is invalid IR) and a replacement op could
+  // not be found, we want to fail with a nicer error message: "op uses a handle
+  // invalidated..." instead of "could not find replacement op". This nicer
+  // error is produced later.
+  if (llvm::any_of(handles, [&](Value handle) {
+        return consumedHandles.contains(handle);
+      }))
+    return true;
 
-  // Replace value handles.
-  for (auto [oldValue, newValue] : llvm::zip(op->getResults(), newValues))
-    (void)replacePayloadValue(oldValue, newValue);
+  // Check if the payload op is mapped to any alive handle.
+  for (Value handle : handles)
+    for (Operation *user : handle.getUsers())
+      if (user != transformOp && !happensBefore(user, transformOp))
+        return false;
 
-  // Replace op handle.
+  return true;
+}
+
+void transform::TrackingListener::notifyOperationReplaced(
+    Operation *op, Operation *replacement) {
+  assert(op->getNumResults() == replacement->getNumResults() &&
+         "invalid replacement");
+
+  // Nothing to do if op is not tracked.
   SmallVector<Value> opHandles;
-  if (failed(getTransformState().getHandlesForPayloadOp(op, opHandles))) {
-    // Op is not tracked.
+  if (failed(getTransformState().getHandlesForPayloadOp(op, opHandles)))
+    return;
+
+  // Drop the op from the mapping if there are no more alive handles.
+  if (canSilentlyDropPayloadOpFromMapping(opHandles)) {
+    (void)replacePayloadOp(op, nullptr);
     return;
   }
 
-  // Helper function to check if the current transform op consumes any handle
-  // that is mapped to `op`.
-  //
-  // Note: If a handle was consumed, there shouldn't be any alive users, so it
-  // is not really necessary to check for consumed handles. However, in case
-  // there are indeed alive handles that were consumed (which is invalid IR) and
-  // a replacement op could not be found, we want to fail with a nicer error
-  // message: "op uses a handle invalidated..." instead of "could not find
-  // replacement op". This nicer error is produced later.
-  auto handleWasConsumed = [&] {
-    return llvm::any_of(opHandles,
-                        [&](Value h) { return consumedHandles.contains(h); });
-  };
+  // Handle ops with 0 results here. All other replacements are handled in the
+  // other `notifyOperationReplaced` overload.
+  if (op->getNumResults() != 0)
+    return;
 
-  // Helper function to check if the handle is alive.
-  auto hasAliveUser = [&]() {
-    for (Value v : opHandles) {
-      for (Operation *user : v.getUsers())
-        if (user != transformOp && !happensBefore(user, transformOp))
-          return true;
-    }
-    return false;
-  };
+  // Replace op handle if both ops have the same name.
+  if (op->getName() == replacement->getName()) {
+    (void)replacePayloadOp(op, replacement);
+    return;
 
-  if (!hasAliveUser() || handleWasConsumed()) {
-    // The op is tracked but the corresponding handles are dead or were
-    // consumed. Drop the op form the mapping.
+  // The replacement op has a different name. This case is handled in the other
+  // `notifyOperationReplaced` overload.
+}
+
+void transform::TrackingListener::notifyOperationReplaced(
+    Operation *op, ValueRange replacement) {
+  assert(op->getNumResults() == replacement.size() &&
+         "invalid number of replacement values");
+
+  // Replace value handles.
+  for (auto [oldValue, newValue] : llvm::zip(op->getResults(), replacement))
+    (void)replacePayloadValue(oldValue, newValue);
+
+  // Nothing more to do if op is not tracked.
+  SmallVector<Value> opHandles;
+  if (failed(getTransformState().getHandlesForPayloadOp(op, opHandles)))
+    return;
+
+  // Drop the op from the mapping if there are no more alive handles.
+  if (canSilentlyDropPayloadOpFromMapping(opHandles)) {
     (void)replacePayloadOp(op, nullptr);
     return;
   }
 
-  FailureOr<Operation *> replacement = findReplacementOp(op, newValues);
+  // Replace op handle.
+  FailureOr<Operation *> replacementOp = findReplacementOp(op, replacement);
   // If the op is tracked but no replacement op was found, send a
   // notification.
-  if (failed(replacement)) {
-    notifyPayloadReplacementNotFound(op, newValues);
+  if (failed(replacementOp)) {
+    notifyPayloadReplacementNotFound(op, replacement);
     (void)replacePayloadOp(op, nullptr);
     return;
   }
 
-  (void)replacePayloadOp(op, *replacement);
+  (void)replacePayloadOp(op, *replacementOp);
 }
 
 transform::ErrorCheckingTrackingListener::~ErrorCheckingTrackingListener() {
   // The state of the ErrorCheckingTrackingListener must be checked and reset
   // if there was an error. This is to prevent errors from accidentally being
   // missed.
   assert(status.succeeded() && "listener state was not checked");
 }

mlir/lib/IR/PatternMatch.cpp

 }
 
 /// '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));
+
+  // Notify the listener that we're about to replace this op.
+  if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
+    rewriteListener->notifyOperationReplaced(op, newOp);
+
   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,

mlir/test/Dialect/Transform/test-interpreter.mlir

 // -----
 
 // CHECK-LABEL: func @test_tracked_rewrite() {
 //  CHECK-NEXT:   transform.test_dummy_payload_op  {new_op} : () -> i1
 //  CHECK-NEXT:   transform.test_dummy_payload_op  {new_op} : () -> i1
 //  CHECK-NEXT:   return
 //  CHECK-NEXT: }
 func.func @test_tracked_rewrite() {
-  %0 = transform.test_dummy_payload_op {replace_me} : () -> (i1)
+  %0 = transform.test_dummy_payload_op {replace_me, test_replace_op} : () -> (i1)
   %1 = transform.test_dummy_payload_op {erase_me} : () -> (i1)
-  %2 = transform.test_dummy_payload_op {replace_me} : () -> (i1)
+  %2 = transform.test_dummy_payload_op {replace_me, test_replace_op} : () -> (i1)
   func.return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["transform.test_dummy_payload_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-remark @below {{2 iterations}}
   transform.test_tracked_rewrite %0 : (!transform.any_op) -> ()
   // One replacement op (test.drop_mapping) is dropped from the mapping.
   // expected-remark @below {{2}}
   test_print_number_of_associated_payload_ir_ops %0 : !transform.any_op
 }
 
+// -----
+
+// CHECK-LABEL: func @test_tracked_rewrite_no_results() {
+//  CHECK-NEXT:   transform.test_dummy_payload_op  {new_op} : () -> ()
+//  CHECK-NEXT:   transform.test_dummy_payload_op  {new_op} : () -> ()
+//  CHECK-NEXT:   return
+//  CHECK-NEXT: }
+func.func @test_tracked_rewrite_no_results() {
+  transform.test_dummy_payload_op {replace_me} : () -> ()
+  transform.test_dummy_payload_op {erase_me} : () -> ()
+  transform.test_dummy_payload_op {replace_me} : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["transform.test_dummy_payload_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below {{2 iterations}}
+  transform.test_tracked_rewrite %0 : (!transform.any_op) -> ()
+  // One replacement op is dropped from the mapping.
+  // expected-remark @below {{2}}
+  test_print_number_of_associated_payload_ir_ops %0 : !transform.any_op
+}
 
 // -----
 
 // CHECK-LABEL: func @test_annotation()
 //  CHECK-NEXT:   "test.annotate_me"()
 //  CHECK-SAME:                        broadcast_attr = 2 : i64
 //  CHECK-SAME:                        new_attr = 1 : i32
 //  CHECK-SAME:                        unit_attr

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

for (Operation *op : state.getPayloadOps(getIn())) {
       if (op->hasAttr("erase_me")) {
         rewriter.eraseOp(op);
         continue;
       }
       if (!op->hasAttr("replace_me")) {
         continue;
       }
 
+      // RewriterBase::replaceOp test, which triggers an "op replaced with
+      // values" notification.
+      if (op->hasAttr("test_replace_op")) {
         SmallVector<NamedAttribute> attributes;
         attributes.emplace_back(rewriter.getStringAttr("new_op"),
                                 rewriter.getUnitAttr());
         OperationState opState(op->getLoc(), op->getName().getIdentifier(),
                                /*operands=*/ValueRange(),
                                /*types=*/op->getResultTypes(), attributes);
         Operation *newOp = rewriter.create(opState);
         rewriter.replaceOp(op, newOp->getResults());
+        continue;
+      }
+
+      // Test RewriterBase::replaceOpWithNewOp, which triggers an "op replaced
+      // with op" notification.
+      Operation *newOp = rewriter.replaceOpWithNewOp<TestDummyPayloadOp>(
+          op, op->getResultTypes(), op->getOperands());
+      newOp->setAttr(rewriter.getStringAttr("new_op"), rewriter.getUnitAttr());
     }
   }
 
   emitRemark() << numIterations << " iterations";
   return DiagnosedSilenceableFailure::success();
 }
 
 namespace {