[Mlir-commits] [mlir] 9b21866 - [mlir][linalg] Fix invalid IR in `FoldInsertPadIntoFill` (#74418)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2023-12-19T14:17:54+09:00
New Revision: 9b21866feaea912bdb2d76060ef79da8a4905570

URL: https://github.com/llvm/llvm-project/commit/9b21866feaea912bdb2d76060ef79da8a4905570
DIFF: https://github.com/llvm/llvm-project/commit/9b21866feaea912bdb2d76060ef79da8a4905570.diff

LOG: [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).

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 9a4d5e8845b214..b68aa77fd83a1c 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>(