Skip to content

Commit 9b21866

Browse files
[mlir][linalg] Fix invalid IR in FoldInsertPadIntoFill (#74418)
`FoldInsertPadIntoFill` used to generate an invalid `tensor.insert_slice` op: ``` error: expected type to be 'tensor<?x?x?xf32>' or a rank-reduced version. (size mismatch) ``` This commit fixes tests such as `mlir/test/Dialect/Linalg/canonicalize.mlir` when verifying the IR after each pattern application (#74270).
1 parent 3a087c1 commit 9b21866

File tree

1 file changed

+9
-4
lines changed

1 file changed

+9
-4
lines changed

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

+9-4
Original file line numberDiff line numberDiff line change
@@ -716,11 +716,16 @@ struct FoldInsertPadIntoFill : public OpRewritePattern<tensor::InsertSliceOp> {
716716
rewriter, loc, addMap, {std::get<0>(p), std::get<1>(p)}));
717717
}
718718

719+
RankedTensorType srcPadType = srcPadOp.getSourceType();
719720
SmallVector<OpFoldResult, 4> newSizes;
720-
for (int i = 0, e = srcPadOp.getSourceType().getRank(); i < e; ++i) {
721-
newSizes.push_back(
722-
rewriter.create<tensor::DimOp>(loc, srcPadOp.getSource(), i)
723-
.getResult());
721+
for (int i = 0, e = srcPadType.getRank(); i < e; ++i) {
722+
if (srcPadType.isDynamicDim(i)) {
723+
newSizes.push_back(
724+
rewriter.create<tensor::DimOp>(loc, srcPadOp.getSource(), i)
725+
.getResult());
726+
} else {
727+
newSizes.push_back(rewriter.getIndexAttr(srcPadType.getDimSize(i)));
728+
}
724729
}
725730

726731
rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(

0 commit comments

Comments
 (0)