[mlir][linalg] Fix invalid IR in FoldInsertPadIntoFill

[matthias-springer]

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).

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

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).

---

Full diff: https://github.com/llvm/llvm-project/pull/74418.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+9-4)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 58af9995548e9..86e5215d72ac5 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -716,11 +716,16 @@ struct FoldInsertPadIntoFill : public OpRewritePattern<tensor::InsertSliceOp> {
           rewriter, loc, addMap, {std::get<0>(p), std::get<1>(p)}));
     }
 
+    RankedTensorType srcPadType = srcPadOp.getSourceType();
     SmallVector<OpFoldResult, 4> newSizes;
-    for (int i = 0, e = srcPadOp.getSourceType().getRank(); i < e; ++i) {
-      newSizes.push_back(
-          rewriter.create<tensor::DimOp>(loc, srcPadOp.getSource(), i)
-              .getResult());
+    for (int i = 0, e = srcPadType.getRank(); i < e; ++i) {
+      if (srcPadType.isDynamicDim(i)) {
+        newSizes.push_back(
+            rewriter.create<tensor::DimOp>(loc, srcPadOp.getSource(), i)
+                .getResult());
+      } else {
+        newSizes.push_back(rewriter.getIndexAttr(srcPadType.getDimSize(i)));
+      }
     }
 
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(