Fix deletion of operations through the rewriter in a pattern matching a consumer operation

This allows for the conversion to match `A(B()) -> C()` with a pattern matching
`A` and marking `B` for deletion.

Also add better assertions when an operation is erased while still having uses.

Differential Revision: https://reviews.llvm.org/D99442

mlir/lib/IR/Operation.cpp

@@ -182,7 +182,17 @@ Operation::Operation(Location location, OperationName name, unsigned numResults,
 // allocated via malloc.
 Operation::~Operation() {
   assert(block == nullptr && "operation destroyed but still in a block");
-
+#ifndef NDEBUG
+  if (!use_empty()) {
+    {
+      InFlightDiagnostic diag =
+          emitOpError("operation destroyed but still has uses");
+      for (Operation *user : getUsers())
+        diag.attachNote(user->getLoc()) << "- use: " << *user << "\n";
+    }
+    llvm::report_fatal_error("operation destroyed but still has uses");
+  }
+#endif
   // Explicitly run the destructors for the operands.
   if (hasOperandStorage)
     getOperandStorage().~OperandStorage();

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -916,9 +916,13 @@ void ConversionPatternRewriterImpl::applyRewrites() {
 
   // In a second pass, erase all of the replaced operations in reverse. This
   // allows processing nested operations before their parent region is
-  // destroyed.
+  // destroyed. Because we process in reverse order, producers may be deleted
-  for (auto &repl : llvm::reverse(replacements))
+  // before their users (a pattern deleting a producer and then the consumer)
+  // so we first drop all uses explicitly.
+  for (auto &repl : llvm::reverse(replacements)) {
+    repl.first->dropAllUses();
     repl.first->erase();
+  }
 
   argConverter.applyRewrites(mapping);
 
@@ -2230,13 +2234,20 @@ LogicalResult OperationConverter::convertOperations(ArrayRef<Operation *> ops) {
   // legalized.
   if (failed(finalize(rewriter)))
     return rewriterImpl.discardRewrites(), failure();
-
   // After a successful conversion, apply rewrites if this is not an analysis
   // conversion.
   if (mode == OpConversionMode::Analysis)
     rewriterImpl.discardRewrites();
-  else
+  else {
     rewriterImpl.applyRewrites();
+
+    // It is possible for a later pattern to erase an op that was originally
+    // identified as illegal and added to the trackedOps, remove it now after
+    // replacements have been computed.
+    if (trackedOps)
+      for (auto &repl : rewriterImpl.replacements)
+        trackedOps->erase(repl.first);
+  }
   return success();
 }
 

mlir/test/Transforms/test-legalizer.mlir

@@ -270,3 +270,17 @@ func @undo_child_created_before_parent() {
   // expected-remark@+1 {{op 'std.return' is not legalizable}}
   return
 }
+
+
+// -----
+
+
+// Check that a conversion pattern on `test.blackhole` can mark the producer
+// for deletion.
+// CHECK-LABEL: @blackhole
+func @blackhole() {
+  %input = "test.blackhole_producer"() : () -> (i32)
+  "test.blackhole"(%input) : (i32) -> ()
+  // expected-remark@+1 {{op 'std.return' is not legalizable}}
+  return
+}

mlir/test/lib/Dialect/Test/TestOps.td

@@ -1309,6 +1309,11 @@ def TestRecursiveRewriteOp : TEST_Op<"recursive_rewrite"> {
   let assemblyFormat = "$depth attr-dict";
 }
 
+// Test legalization pattern: this op will be erase and will also erase the
+// producer of its operand.
+def BlackHoleOp : TEST_Op<"blackhole">,
+  Arguments<(ins AnyType)>;
+
 //===----------------------------------------------------------------------===//
 // Test Type Legalization
 //===----------------------------------------------------------------------===//

mlir/test/lib/Dialect/Test/TestPatterns.cpp

@@ -491,6 +491,21 @@ struct TestNestedOpCreationUndoRewrite
     return success();
   };
 };
+
+// This pattern matches `test.blackhole` and delete this op and its producer.
+struct TestReplaceEraseOp : public OpRewritePattern<BlackHoleOp> {
+  using OpRewritePattern<BlackHoleOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BlackHoleOp op,
+                                PatternRewriter &rewriter) const final {
+    Operation *producer = op.getOperand().getDefiningOp();
+    // Always erase the user before the producer, the framework should handle
+    // this correctly.
+    rewriter.eraseOp(op);
+    rewriter.eraseOp(producer);
+    return success();
+  };
+};
 } // namespace
 
 namespace {
@@ -568,7 +583,8 @@ struct TestLegalizePatternDriver
              TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64,
              TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
              TestNonRootReplacement, TestBoundedRecursiveRewrite,
-             TestNestedOpCreationUndoRewrite>(&getContext());
+             TestNestedOpCreationUndoRewrite, TestReplaceEraseOp>(
+            &getContext());
     patterns.add<TestDropOpSignatureConversion>(&getContext(), converter);
     mlir::populateFuncOpTypeConversionPattern(patterns, converter);
     mlir::populateCallOpTypeConversionPattern(patterns, converter);