diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index 9f8a482d6e2d2..44cf8331d55a7 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -1382,16 +1382,21 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, assert(newValues.size() == op->getNumResults()); assert(!ignoredOps.contains(op) && "operation was already replaced"); + // Check if replaced op is an unresolved materialization, i.e., an + // unrealized_conversion_cast op that was created by the conversion driver. + bool isUnresolvedMaterialization = false; + if (auto castOp = dyn_cast(op)) + if (unresolvedMaterializations.contains(castOp)) + isUnresolvedMaterialization = true; + // Create mappings for each of the new result values. for (auto [newValue, result] : llvm::zip(newValues, op->getResults())) { if (!newValue) { // This result was dropped and no replacement value was provided. - if (auto castOp = dyn_cast(op)) { - if (unresolvedMaterializations.contains(castOp)) { - // Do not create another materializations if we are erasing a - // materialization. - continue; - } + if (isUnresolvedMaterialization) { + // Do not create another materializations if we are erasing a + // materialization. + continue; } // Materialize a replacement value "out of thin air". @@ -1400,10 +1405,20 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, result.getLoc(), /*inputs=*/ValueRange(), /*outputType=*/result.getType(), /*originalType=*/Type(), currentTypeConverter); + } else { + // Make sure that the user does not mess with unresolved materializations + // that were inserted by the conversion driver. We keep track of these + // ops in internal data structures. Erasing them must be allowed because + // this can happen when the user is erasing an entire block (including + // its body). But replacing them with another value should be forbidden + // to avoid problems with the `mapping`. + assert(!isUnresolvedMaterialization && + "attempting to replace an unresolved materialization"); } - // Remap, and check for any result type changes. - mapping.map(result, newValue); + // Remap result to replacement value. + if (newValue) + mapping.map(result, newValue); } appendRewrite(op, currentTypeConverter);