[mlir][transform] Simplify TrackingListener test case

[springerm]

Summary

Use the default TrackingListener. No need to set up a derived listener just for the test case. This revision is in preparation of a future change that adds a TrackingRewriter infrastructure.

Depends On: D152426

[ftynse]

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

470

471

We can also create an unregistered op from OperationState using RewriterBase::create and then call RewriterBase::replaceOp. I don't have a preference here.

[springerm]

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

471

This is needed for D152438. It adds a new notifyOperationReplaced(Operation*, Operation*) overload that is only triggered by replaceOpWithNewOp<OpTy> but not replaceOp(Operation*, ValueRange).

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

func.func @definition() {
     "test.something"() : () -> ()
     return
   }
 }
 
 // -----
 
 // CHECK-LABEL: func @test_tracked_rewrite() {
-//  CHECK-NEXT:   "test.update_mapping"() {original_op = "test.replace_me"}
-//  CHECK-NEXT:   "test.drop_mapping"() {original_op = "test.replace_me"}
-//  CHECK-NEXT:   "test.update_mapping"() {original_op = "test.replace_me"}
+//  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 = "test.replace_me"() {replacement = "test.update_mapping"} : () -> (i1)
-  %1 = "test.replace_me"() {replacement = "test.drop_mapping"} : () -> (i1)
-  %2 = "test.replace_me"() {replacement = "test.update_mapping"} : () -> (i1)
+  %0 = transform.test_dummy_payload_op {replace_me} : () -> (i1)
+  %1 = transform.test_dummy_payload_op {erase_me} : () -> (i1)
+  %2 = transform.test_dummy_payload_op {replace_me} : () -> (i1)
+  func.return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
-  %0 = transform.structured.match ops{["test.replace_me"]} in %arg1 : (!transform.any_op) -> !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
 }
 
 

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

 }
 
 void mlir::test::TestTrackedRewriteOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::modifiesPayload(effects);
 }
 
-namespace {
-/// A TrackingListener for test cases. When the replacement op is
-/// "test.update_mapping", it is considered as a replacement op in the transform
-/// state mapping. Otherwise, it is not and the original op is simply removed
-/// from the mapping.
-class TestTrackingListener : public transform::TrackingListener {
-  using transform::TrackingListener::TrackingListener;
-
-protected:
-  FailureOr<Operation *>
-  findReplacementOp(Operation *op, ValueRange newValues) const override {
-    if (newValues.size() != 1)
-      return failure();
-    Operation *replacement = newValues[0].getDefiningOp();
-    if (!replacement)
-      return failure();
-    if (replacement->getName().getStringRef() != "test.update_mapping")
-      return failure();
-    return replacement;
-  }
-};
-} // namespace
+void mlir::test::TestDummyPayloadOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  for (OpResult result : getResults())
+    transform::producesHandle(result, effects);
 
 DiagnosedSilenceableFailure
 mlir::test::TestTrackedRewriteOp::apply(transform::TransformResults &results,
                                         transform::TransformState &state) {
-  TestTrackingListener listener(state, *this);
+  transform::ErrorCheckingTrackingListener listener(state, *this);
   IRRewriter rewriter(getContext(), &listener);
   int64_t numIterations = 0;
 
   // `getPayloadOps` returns an iterator that skips ops that are erased in the
   // loop body. Replacement ops are not enumerated.
   for (Operation *op : state.getPayloadOps(getIn())) {
     ++numIterations;
-    rewriter.setInsertionPointToEnd(op->getBlock());
+    (void)op;
 
     // Erase all payload ops. The outer loop should have only one iteration.
     for (Operation *op : state.getPayloadOps(getIn())) {
-      if (op->getName().getStringRef() != "test.replace_me")
+      rewriter.setInsertionPoint(op);
+      if (op->hasAttr("erase_me")) {
+        rewriter.eraseOp(op);
         continue;
-      auto replacementName = op->getAttrOfType<StringAttr>("replacement");
-      if (!replacementName)
+      }
+      if (!op->hasAttr("replace_me")) {
         continue;
+      }
+
       SmallVector<NamedAttribute> attributes;
-      attributes.emplace_back(rewriter.getStringAttr("original_op"),
-                              op->getName().getIdentifier());
-      OperationState opState(op->getLoc(), replacementName,
+      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());
     }
   }
 
   emitRemark() << numIterations << " iterations";

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

let arguments = (ins TransformHandleTypeInterface:$in,
                        UnitAttr:$has_operand_effect,
                        UnitAttr:$has_result_effect,
                        UnitAttr:$modifies_payload);
   let results = (outs TransformHandleTypeInterface:$out);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 }
 
+// This op is used as a payload op. It must be a registered op, so that it can
+// be created with "RewriterBase::replaceOpWithNewOp" (needed for a test case).
+// Since only TransformOpInterface can be injected into the transform dialect,
+// this op implements the interface, even though it is not used as a transform
+// op.
+def TestDummyPayloadOp
+  : Op<Transform_Dialect, "test_dummy_payload_op",
+      [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       TransformOpInterface]> {
+  let arguments = (ins Variadic<AnyType>:$args);
+  let results = (outs Variadic<AnyType>:$outs);
+  let assemblyFormat = "$args attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+
+  let extraClassDeclaration = [{
+    DiagnosedSilenceableFailure apply(transform::TransformResults &results,
+                                      transform::TransformState &state) {
+      llvm_unreachable("op should not be used as a transform");
+      return DiagnosedSilenceableFailure::definiteFailure();
+    }
+  }];
+}
+
 def TestTrackedRewriteOp
   : Op<Transform_Dialect, "test_tracked_rewrite",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$in);
   let results = (outs);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";