diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index 2fe06970eb568..424349b555228 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -2804,17 +2804,19 @@ LogicalResult OperationLegalizer::legalizePatternResult( assert(impl.pendingRootUpdates.empty() && "dangling root updates"); #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS - // Check that the root was either replaced or updated in place. - auto newRewrites = llvm::drop_begin(impl.rewrites, curState.numRewrites); - auto replacedRoot = [&] { - return hasRewrite(newRewrites, op); - }; - auto updatedRootInPlace = [&] { - return hasRewrite(newRewrites, op); - }; - if (!replacedRoot() && !updatedRootInPlace()) - llvm::report_fatal_error( - "expected pattern to replace the root operation or modify it in place"); + if (impl.config.allowPatternRollback) { + // Check that the root was either replaced or updated in place. + auto newRewrites = llvm::drop_begin(impl.rewrites, curState.numRewrites); + auto replacedRoot = [&] { + return hasRewrite(newRewrites, op); + }; + auto updatedRootInPlace = [&] { + return hasRewrite(newRewrites, op); + }; + if (!replacedRoot() && !updatedRootInPlace()) + llvm::report_fatal_error("expected pattern to replace the root operation " + "or modify it in place"); + } #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS // Legalize each of the actions registered during application.