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

[mehdi_amini]

Summary

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.

[jpienaar]

mlir/lib/Transforms/Utils/DialectConversion.cpp

919

So if we did a topological order before deletion this would not be needed? (Not saying we should just to double check I understood what would have avoided the need for drop all uses)

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

500

And the failure was specific to this order? Or either order if erase would have worked? (If the latter it would be good to document it here too even though it is an implementation detail, without the context of this change the test is difficult to interpret)

[mehdi_amini]

mlir/lib/Transforms/Utils/DialectConversion.cpp

919

Yes this is correct. Right now the usual expectation is that the developer should delete all users before deleting the producer.
But because of this reverse here it kind of break the usual idiom...

Right now a user could use the rewriter to mark eraseOp on the producer before the users, but I don't think that's desirable.

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

500

Yes the failure is specific to this order.
As mentioned in the order comment thread, when you manipulate the IR directly you need to delete the consumers before the producer.
However because calling eraseOp is "scheduling for later", the rewrite is doing it in reverse order.

Writing the code with the eraseOp in the other direction would work without my change in the rewriter right now.

[ftynse]

mlir/lib/Transforms/Utils/DialectConversion.cpp

923

Would it make sense to also have a check under NDEBUG that all users of this operation are also scheduled to be dropped? Otherwise be could be dropping operands of operations that remain live and won't catch it (previously, we would catch that in the destructor of operations because there were remaining uses).

[mehdi_amini]

Check if trackedOps is null before using it

[rriddle]

mlir/lib/IR/Operation.cpp

186

!use_empty()?

[mehdi_amini]

mlir/lib/Transforms/Utils/DialectConversion.cpp

923

Actually we already check this below and fails the legalization in this case.

2350

Here @ftynse FYI

[mehdi_amini]

Address river's comment

mlir/lib/IR/Operation.cpp

: location(location), numResults(numResults), numSuccs(numSuccessors),
       attrs(attributes) {
   assert(attributes && "unexpected null attribute dictionary");
 }
 
 // Operations are deleted through the destroy() member because they are
 // allocated via malloc.
 Operation::~Operation() {
   assert(block == nullptr && "operation destroyed but still in a block");
-
+#ifndef NDEBUG
+  if (!getUses().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();
 
   // Explicitly run the destructors for the successors.
   for (auto &successor : getBlockOperands())
     successor.~BlockOperand();
 

mlir/lib/Transforms/Utils/DialectConversion.cpp

for (BlockArgument arg : argReplacements) {
     arg.replaceUsesWithIf(repl, [&](OpOperand &operand) {
       Operation *user = operand.getOwner();
       return user->getBlock() != replBlock || replOp->isBeforeInBlock(user);
     });
   }
 
   // In a second pass, erase all of the replaced operations in reverse. This
   // allows processing nested operations before their parent region is
-  // destroyed.
-  for (auto &repl : llvm::reverse(replacements))
+  // destroyed. Because we process in reverse order, producers may be deleted
+  // 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);
 
   // Now that the ops have been erased, also erase dangling blocks.
   eraseDanglingBlocks();
 }
 
 //===----------------------------------------------------------------------===//
for (auto *op : toConvert)
     if (failed(convert(rewriter, op)))
       return rewriterImpl.discardRewrites(), failure();
 
   // Now that all of the operations have been converted, finalize the conversion
   // process to ensure any lingering conversion artifacts are cleaned up and
   // 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
     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.
+  for (auto &repl : rewriterImpl.replacements)
+    trackedOps->erase(repl.first);
+
   return success();
 }
 
 LogicalResult
 OperationConverter::finalize(ConversionPatternRewriter &rewriter) {
   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
 
   // Legalize converted block arguments.
LogicalResult OperationConverter::legalizeErasedResult(
   });
   if (liveUserIt != result.user_end()) {
     InFlightDiagnostic diag = op->emitError("failed to legalize operation '")
                               << op->getName() << "' marked as erased";
     diag.attachNote(liveUserIt->getLoc())
         << "found live user of result #" << result.getResultNumber() << ": "
         << *liveUserIt;
     return failure();
   }
   return success();
 }
 
 /// Finds a user of the given value, or of any other value that the given value
 /// replaced, that was not replaced in the conversion process.
 static Operation *
 findLiveUserOfReplaced(Value value, ConversionPatternRewriterImpl &rewriterImpl,
                        const BlockAndValueMapping &inverseMapping) {

mlir/test/Transforms/test-legalizer.mlir

 // child op.
 // CHECK-LABEL: @undo_child_created_before_parent
 func @undo_child_created_before_parent() {
   // expected-remark@+1 {{is not legalizable}}
   "test.illegal_op_with_region_anchor"() : () -> ()
   // 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

 def : Pat<(TestRewriteOp $input), (replaceWithValue $input)>;
 
 // Check that patterns can specify bounded recursion when rewriting.
 def TestRecursiveRewriteOp : TEST_Op<"recursive_rewrite"> {
   let arguments = (ins I64Attr:$depth);
   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
 //===----------------------------------------------------------------------===//
 
 def TestRegionBuilderOp : TEST_Op<"region_builder">;
 def TestReturnOp : TEST_Op<"return", [ReturnLike, Terminator]> {
   let arguments = (ins Variadic<AnyType>);
   let builders = [OpBuilder<(ins),

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

struct TestNestedOpCreationUndoRewrite
 
   LogicalResult matchAndRewrite(IllegalOpWithRegionAnchor op,
                                 PatternRewriter &rewriter) const final {
     // rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op);
     rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op);
     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();
+    rewriter.eraseOp(op);
+    rewriter.eraseOp(producer);
+    return success();
+  };
+};
 } // namespace
 
 namespace {
 struct TestTypeConverter : public TypeConverter {
   using TypeConverter::TypeConverter;
   TestTypeConverter() {
     addConversion(convertType);
     addArgumentMaterialization(materializeCast);
void runOnOperation() override {
     populateWithGenerated(patterns);
     patterns
         .add<TestRegionRewriteBlockMovement, TestRegionRewriteUndo,
              TestCreateBlock, TestCreateIllegalBlock, TestUndoBlockArgReplace,
              TestUndoBlockErase, TestPassthroughInvalidOp, TestSplitReturnType,
              TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64,
              TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
              TestNonRootReplacement, TestBoundedRecursiveRewrite,
-             TestNestedOpCreationUndoRewrite>(&getContext());
+             TestNestedOpCreationUndoRewrite, TestReplaceEraseOp>(
+            &getContext());
     patterns.add<TestDropOpSignatureConversion>(&getContext(), converter);
     mlir::populateFuncOpTypeConversionPattern(patterns, converter);
     mlir::populateCallOpTypeConversionPattern(patterns, converter);
 
     // Define the conversion target used for the test.
     ConversionTarget target(getContext());
     target.addLegalOp<ModuleOp, ModuleTerminatorOp>();
     target.addLegalOp<LegalOpA, LegalOpB, TestCastOp, TestValidOp,