[mlir] DialectConversion: avoid double-free when rolling back op creation

[ftynse]

Summary

Dialect conversion infrastructure may roll back op creation by erasing the
operations in the reverse order of their creation. While this guarantees uses
of values will be deleted before their definitions, this does not guarantee
that a parent operation will not be deleted before its child. (This may happen
in case of block inlining or if child operations, such as terminators, are
created in the parent's build function before the parent itself.) Handle the
parent/child relationship between ops by removing all child ops from the blocks
before erasing the parent. The child ops remain live, detached from a block,
and will be safely destroyed in their turn, which may come later than that of
the parent.

Depends On D80133

[nicolasvasilache]

Thanks!
This must not have been fun to debug..

mlir/lib/Transforms/DialectConversion.cpp

672

recursively

[ftynse]

Rebase

mlir/lib/Transforms/DialectConversion.cpp

 } // end namespace mlir
 
 RewriterState ConversionPatternRewriterImpl::getCurrentState() {
   return RewriterState(createdOps.size(), replacements.size(),
                        argReplacements.size(), blockActions.size(),
                        ignoredOps.size(), rootUpdates.size());
 }
 
+/// Detach any operations nested in the given operation from their parent
+/// blocks, and erase the given operation. This can be used when the nested
+/// operations are scheduled for erasure themselves, so deleting the regions of
+/// the given operation together with their content would result in double-free.
+/// This happens, for example, when rolling back op creation in the reverse
+/// order and if the nested ops were created before the parent op. This function
+/// does not need to collect nested ops recurisvely because it is expected to
+/// also be called for each nested op when it is about to be deleted.
+static void detachNestedAndErase(Operation *op) {
+  for (Region &region : op->getRegions())
+    for (Block &block : region.getBlocks())
+      while (!block.getOperations().empty())
+        block.getOperations().remove(block.getOperations().begin());
+  op->erase();
+}
+
 void ConversionPatternRewriterImpl::resetState(RewriterState state) {
   // Reset any operations that were updated in place.
   for (unsigned i = state.numRootUpdates, e = rootUpdates.size(); i != e; ++i)
     rootUpdates[i].resetOperation();
   rootUpdates.resize(state.numRootUpdates);
 
   // Reset any replaced arguments.
   for (BlockArgument replacedArg :
        llvm::drop_begin(argReplacements, state.numArgReplacements))
     mapping.erase(replacedArg);
   argReplacements.resize(state.numArgReplacements);
 
   // Undo any block actions.
   undoBlockActions(state.numBlockActions);
 
   // Reset any replaced operations and undo any saved mappings.
   for (auto &repl : llvm::drop_begin(replacements, state.numReplacements))
     for (auto result : repl.op->getResults())
       mapping.erase(result);
   replacements.resize(state.numReplacements);
 
   // Pop all of the newly created operations.
   while (createdOps.size() != state.numCreatedOps) {
-    createdOps.back()->erase();
+    detachNestedAndErase(createdOps.back());
     createdOps.pop_back();
   }
 
   // Pop all of the recorded ignored operations that are no longer valid.
   while (ignoredOps.size() != state.numIgnoredOperations)
     ignoredOps.pop_back();
 }
 
void ConversionPatternRewriterImpl::discardRewrites() {
   // Reset any operations that were updated in place.
   for (auto &state : rootUpdates)
     state.resetOperation();
 
   undoBlockActions();
 
   // Remove any newly created ops.
   for (auto *op : llvm::reverse(createdOps))
-    op->erase();
+    detachNestedAndErase(op);
 }
 
 void ConversionPatternRewriterImpl::applyRewrites() {
   // Apply all of the rewrites replacements requested during conversion.
   for (auto &repl : replacements) {
     for (unsigned i = 0, e = repl.newValues.size(); i != e; ++i) {
       if (auto newValue = repl.newValues[i])
         repl.op->getResult(i).replaceAllUsesWith(

mlir/test/Transforms/test-legalizer.mlir

^bb0(%arg0: i32):
     // CHECK: ^bb0(%[[ARG:.*]]: i32):
     // CHECK-NEXT: "test.return"(%[[ARG]]) : (i32)
 
     "test.return"(%arg0) : (i32) -> ()
   }) : () -> ()
   // expected-remark@+1 {{op 'std.return' is not legalizable}}
   return
 }
+
+// -----
+
+// The op in this function is attempted to be rewritten to another illegal op
+// with an attached region containing an invalid terminator. The terminator is
+// created before the parent op. The deletion should not crash when deleting
+// created ops in the inverse order, i.e. deleting the parent op and then the
+// 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
+}

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

 def ILLegalOpC : TEST_Op<"illegal_op_c">, Results<(outs I32)>;
 def ILLegalOpD : TEST_Op<"illegal_op_d">, Results<(outs I32)>;
 def ILLegalOpE : TEST_Op<"illegal_op_e">, Results<(outs I32)>;
 def ILLegalOpF : TEST_Op<"illegal_op_f">, Results<(outs I32)>;
 def LegalOpA : TEST_Op<"legal_op_a">,
   Arguments<(ins Test_LegalizerEnum:$status)>, Results<(outs I32)>;
 def LegalOpB : TEST_Op<"legal_op_b">, Results<(outs I32)>;
 
+// Check that the conversion infrastructure can properly undo the creation of
+// operations where an operation was created before its parent, in this case,
+// in the parent's builder.
+def IllegalOpTerminator : TEST_Op<"illegal_op_terminator", [Terminator]>;
+def IllegalOpWithRegion : TEST_Op<"illegal_op_with_region"> {
+  let skipDefaultBuilders = 1;
+  let builders = [OpBuilder<"OpBuilder &builder, OperationState &state",
+                  [{ Region *bodyRegion = state.addRegion();
+                     OpBuilder::InsertionGuard g(builder);
+                     Block *body = builder.createBlock(bodyRegion);
+                     builder.setInsertionPointToEnd(body);
+                     builder.create<IllegalOpTerminator>(state.location);
+                  }]>];
+}
+def IllegalOpWithRegionAnchor : TEST_Op<"illegal_op_with_region_anchor">;
+
 // Check that smaller pattern depths are chosen, i.e. prioritize more direct
 // mappings.
 def : Pat<(ILLegalOpA), (LegalOpA Test_LegalizerEnum_Success)>;
 
 def : Pat<(ILLegalOpA), (ILLegalOpB)>;
 def : Pat<(ILLegalOpB), (LegalOpA Test_LegalizerEnum_Failure)>;
 
 // Check that the higher benefit pattern is taken for multiple legalizations

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

rewriter.updateRootInPlace(op, [&] {
                  rewriter.getI64IntegerAttr(op.depth().getSExtValue() - 1));
     });
     return success();
   }
 
   /// The conversion target handles bounding the recursion of this pattern.
   bool hasBoundedRewriteRecursion() const final { return true; }
 };
+
+struct TestNestedOpCreationUndoRewrite
+    : public OpRewritePattern<IllegalOpWithRegionAnchor> {
+  using OpRewritePattern<IllegalOpWithRegionAnchor>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IllegalOpWithRegionAnchor op,
+                                PatternRewriter &rewriter) const final {
+    // rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op);
+    rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op);
+    return success();
+  };
+};
 } // namespace
 
 namespace {
 struct TestTypeConverter : public TypeConverter {
   using TypeConverter::TypeConverter;
   TestTypeConverter() { addConversion(convertType); }
 
   static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) {
void runOnOperation() override {
     mlir::OwningRewritePatternList patterns;
     populateWithGenerated(&getContext(), &patterns);
     patterns.insert<TestRegionRewriteBlockMovement, TestRegionRewriteUndo,
                     TestCreateBlock, TestCreateIllegalBlock,
                     TestUndoBlockArgReplace, TestPassthroughInvalidOp,
                     TestSplitReturnType, TestChangeProducerTypeI32ToF32,
                     TestChangeProducerTypeF32ToF64,
                     TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
-                    TestNonRootReplacement, TestBoundedRecursiveRewrite>(
-        &getContext());
+                    TestNonRootReplacement, TestBoundedRecursiveRewrite,
+                    TestNestedOpCreationUndoRewrite>(&getContext());
     patterns.insert<TestDropOpSignatureConversion>(&getContext(), converter);
     mlir::populateFuncOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
     mlir::populateCallOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
 
     // Define the conversion target used for the test.
     ConversionTarget target(getContext());