[mlir][Transform] Fix dropReverseMapping early exit condition

[nicolasvasilache]

Summary

Previously, the erasure would not trigger and result in surprising behavior.

[ftynse]

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

906

Took me some time to realize "nested under" means "handles to payload operations whose ancestor is pointed to by the given handle", but I can't find a good and short name.

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

for (Operation *op : targets)
     mappings.reverse[op].push_back(value);
 
   return success();
 }
 
 void transform::TransformState::dropReverseMapping(Mappings &mappings,
                                                    Operation *op, Value value) {
   auto it = mappings.reverse.find(op);
-  if (it != mappings.reverse.end())
+  if (it == mappings.reverse.end())
     return;
 
   llvm::erase_value(it->getSecond(), value);
   if (it->getSecond().empty())
     mappings.reverse.erase(it);
 }
 
 void transform::TransformState::removePayloadOps(Value value) {

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

^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @other in %arg1 : (!pdl.operation) -> !pdl.operation
     // expected-error @below {{incompatible payload operation name}}
     %2 = transform.cast %0 : !pdl.operation to !transform.op<"test.some_op">
     transform.cast %2 : !transform.op<"test.some_op"> to !pdl.operation
   }
 }
+
+// -----
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 : !pdl.operation failures(propagate) {
+  ^bb0(%arg1: !pdl.operation):
+    %0 = pdl_match @some in %arg1 : (!pdl.operation) -> !pdl.operation
+    // here, the handles nested under are {%arg0, %arg1, %0}
+    // expected-remark @below {{3 handles nested under}}
+    transform.test_report_number_of_tracked_handles_nested_under %arg1
+    // expected-remark @below {{erased}}
+    transform.test_emit_remark_and_erase_operand %0, "erased"
+    // here, the handles nested under are only {%arg0, %arg1}
+    // expected-remark @below {{2 handles nested under}}
+    transform.test_report_number_of_tracked_handles_nested_under %arg1
+  }
+
+  pdl.pattern @some : benefit(1) {
+    %0 = pdl.operation "test.some_op"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+}
+
+"test.some_op"() : () -> ()

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

if (op->getName().getDialectNamespace() != "test") {
       diag << "expected the payload operation to belong to the 'test' dialect";
       return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
     }
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
+void mlir::test::TestReportNumberOfTrackedHandlesNestedUnder::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestReportNumberOfTrackedHandlesNestedUnder::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  int64_t count = 0;
+  for (Operation *op : state.getPayloadOps(getTarget())) {
+    op->walk([&](Operation *nested) {
+      SmallVector<Value> handles;
+      (void)state.getHandlesForPayloadOp(nested, handles);
+      count += handles.size();
+    });
+  }
+  emitRemark() << count << " handles nested under";
+  return DiagnosedSilenceableFailure::success();
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

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

: Op<Transform_Dialect, "test_copy_payload",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins Arg<PDL_Operation, "", [TransformMappingRead]>:$handle);
   let results = (outs Res<PDL_Operation, "",
       [TransformMappingAlloc, TransformMappingWrite]>:$copy);
   let cppNamespace = "::mlir::test";
   let assemblyFormat = "$handle attr-dict";
 }
 
+def TestReportNumberOfTrackedHandlesNestedUnder
+  : Op<Transform_Dialect, "test_report_number_of_tracked_handles_nested_under",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins PDL_Operation:$target);
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD