diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index 84e7232d326a8..cc61bc6b6260c 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -1006,18 +1006,46 @@ class ModifyOperationRewrite : public OperationRewrite { : OperationRewrite(Kind::ModifyOperation, rewriterImpl, op), loc(op->getLoc()), attrs(op->getAttrDictionary()), operands(op->operand_begin(), op->operand_end()), - successors(op->successor_begin(), op->successor_end()) {} + successors(op->successor_begin(), op->successor_end()) { + if (OpaqueProperties prop = op->getPropertiesStorage()) { + // Make a copy of the properties. + propertiesStorage = operator new(op->getPropertiesStorageSize()); + OpaqueProperties propCopy(propertiesStorage); + op->getName().initOpProperties(propCopy, /*init=*/prop); + } + } static bool classof(const IRRewrite *rewrite) { return rewrite->getKind() == Kind::ModifyOperation; } + ~ModifyOperationRewrite() override { + assert(!propertiesStorage && + "rewrite was neither committed nor rolled back"); + } + + void commit() override { + if (propertiesStorage) { + OpaqueProperties propCopy(propertiesStorage); + op->getName().destroyOpProperties(propCopy); + operator delete(propertiesStorage); + propertiesStorage = nullptr; + } + } + void rollback() override { op->setLoc(loc); op->setAttrs(attrs); op->setOperands(operands); for (const auto &it : llvm::enumerate(successors)) op->setSuccessor(it.value(), it.index()); + if (propertiesStorage) { + OpaqueProperties propCopy(propertiesStorage); + op->copyProperties(propCopy); + op->getName().destroyOpProperties(propCopy); + operator delete(propertiesStorage); + propertiesStorage = nullptr; + } } private: @@ -1025,6 +1053,7 @@ class ModifyOperationRewrite : public OperationRewrite { DictionaryAttr attrs; SmallVector operands; SmallVector successors; + void *propertiesStorage = nullptr; }; } // namespace diff --git a/mlir/test/Transforms/test-legalizer.mlir b/mlir/test/Transforms/test-legalizer.mlir index 84fcc18ab7d37..62d776cd7573e 100644 --- a/mlir/test/Transforms/test-legalizer.mlir +++ b/mlir/test/Transforms/test-legalizer.mlir @@ -334,3 +334,15 @@ func.func @test_move_op_before_rollback() { }) : () -> () "test.return"() : () -> () } + +// ----- + +// CHECK-LABEL: func @test_properties_rollback() +func.func @test_properties_rollback() { + // CHECK: test.with_properties <{a = 32 : i64, + // expected-remark @below{{op 'test.with_properties' is not legalizable}} + test.with_properties + <{a = 32 : i64, array = array, b = "foo"}> + {modify_inplace} + "test.return"() : () -> () +} diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp index 2102a4ffabf7b..108cfe8950ef6 100644 --- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp +++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp @@ -807,6 +807,21 @@ struct TestUndoBlockErase : public ConversionPattern { } }; +/// A pattern that modifies a property in-place, but keeps the op illegal. +struct TestUndoPropertiesModification : public ConversionPattern { + TestUndoPropertiesModification(MLIRContext *ctx) + : ConversionPattern("test.with_properties", /*benefit=*/1, ctx) {} + LogicalResult + matchAndRewrite(Operation *op, ArrayRef operands, + ConversionPatternRewriter &rewriter) const final { + if (!op->hasAttr("modify_inplace")) + return failure(); + rewriter.modifyOpInPlace( + op, [&]() { cast(op).getProperties().setA(42); }); + return success(); + } +}; + //===----------------------------------------------------------------------===// // Type-Conversion Rewrite Testing @@ -1086,7 +1101,8 @@ struct TestLegalizePatternDriver TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType, TestNonRootReplacement, TestBoundedRecursiveRewrite, TestNestedOpCreationUndoRewrite, TestReplaceEraseOp, - TestCreateUnregisteredOp, TestUndoMoveOpBefore>(&getContext()); + TestCreateUnregisteredOp, TestUndoMoveOpBefore, + TestUndoPropertiesModification>(&getContext()); patterns.add(&getContext(), converter); mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns, converter);