diff --git a/mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp --- a/mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp +++ b/mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp @@ -764,6 +764,16 @@ if (xferOp.getVectorType().getRank() <= kTargetRank) return failure(); + if (xferOp->hasAttr(kPassLabel) && !getInsertOp(xferOp)) { + // InsertOp was folded. Remove label and restart. + // This can happen if the result of a TransferReadOp is immediately + // fed into a vector.extract. The folder removes both ExtractOp and + // InsertOp, but the pass label is still there. + rewriter.updateRootInPlace(xferOp, + [&]() { xferOp->removeAttr(kPassLabel); }); + return success(); + } + ScopedContext scope(rewriter, xferOp.getLoc()); auto vec = getResultVector(xferOp, rewriter); auto vecType = vec.getType().dyn_cast(); @@ -883,6 +893,16 @@ if (xferOp.getVectorType().getRank() <= kTargetRank) return failure(); + if (xferOp->hasAttr(kPassLabel) && !getExtractOp(xferOp)) { + // ExtractOp was folded. Remove label and restart. + // This can happen if the input of a TransferWriteOp is a vector.insert. + // The folder removes both InsertOp and ExtractOp, but the pass label is + // still there. + rewriter.updateRootInPlace(xferOp, + [&]() { xferOp->removeAttr(kPassLabel); }); + return success(); + } + ScopedContext scope(rewriter, xferOp.getLoc()); auto vec = getDataVector(xferOp); auto xferVecType = xferOp.getVectorType();