[mlir][tensor/linalg] Fix bug in reifyResultShapes

[springerm]

Summary

reifyResultShapes should return an IntegerAttr if and only if the corresponding dimension is static.

[nicolasvasilache]

some test plz

[springerm]

In D145702#4184086, @nicolasvasilache wrote:

some test plz

Test added as D145777.

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

SmallVector<OpFoldResult> allResultDimValues =
       makeComposedFoldedMultiResultAffineApply(
           rewriter, loc, resultShapesFromInputShapesMap,
           createFlatListOfOperandDims(b, loc));
   int64_t pos = 0;
   ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
   for (OpOperand *opOperand : getDpsInitOperands()) {
     SmallVector<OpFoldResult> shapes;
     for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) {
-      if (checkDimExpr.visit(shapeExprs[pos]))
-        shapes.push_back(createOrFoldDimOp(b, loc, opOperand->get(), dim));
-      else
-        shapes.push_back(allResultDimValues[pos]);
+      auto shapedType = opOperand->get().getType().cast<ShapedType>();
+      if (!shapedType.isDynamicDim(dim)) {
+        // Static dim: Return IntegerAttr.
+        shapes.push_back(b.getIndexAttr(shapedType.getDimSize(dim)));
+      } else {
+        // Dynamic dim: Return Value.
+        OpFoldResult ofr =
+            checkDimExpr.visit(shapeExprs[pos])
+                ? createOrFoldDimOp(b, loc, opOperand->get(), dim)
+                : allResultDimValues[pos];
+        shapes.push_back(getValueOrCreateConstantIndexOp(b, loc, ofr));
+      }
       pos++;
     }
     reifiedReturnShapes.emplace_back(std::move(shapes));
   }
   return success();
 }
 
 /// Return the index in the indexingMaps vector that corresponds to this

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

if (auto result = foldInsertAfterExtractSlice(*this))
     return result;
   return OpFoldResult();
 }
 
 LogicalResult InsertSliceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
   for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
+    if (getType().isDynamicDim(dim)) {
       reifiedReturnShapes[0][dim] =
           builder.createOrFold<tensor::DimOp>(getLoc(), getDest(), dim);
+    } else {
+      reifiedReturnShapes[0][dim] =
+          builder.getIndexAttr(getType().getDimSize(dim));
+    }
   }
   return success();
 }
 
 namespace {
 /// Pattern to rewrite a insert_slice op with constant arguments.
 ///
 /// This pattern works with both InsertSliceOp and ParallelInsertSliceOp.

 template <typename OpTy>
 static LogicalResult
 reifyResultShapesImpl(OpTy op, OpBuilder &builder,
                       ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   int64_t destRank = op.getDestRank();
   reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(destRank));
+  ShapedType resultType = op.getResult().getType().template cast<ShapedType>();
   for (auto dim : llvm::seq<int64_t>(0, destRank)) {
+    if (resultType.isDynamicDim(dim)) {
       reifiedReturnShapes[0][dim] =
           builder.createOrFold<tensor::DimOp>(op.getLoc(), op.getDest(), dim);
+    } else {
+      reifiedReturnShapes[0][dim] =
+          builder.getIndexAttr(resultType.getDimSize(dim));
+    }
   }
   return success();
 }
 
 template <typename OpTy>
 static DenseMap<int64_t, OpFoldResult> getDimAndTileMappingImpl(OpTy op) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");