[mlir][tensor] Fix ReifyRankedShapedTypeOpInterface impl. of reshape ops

[springerm]

Summary

reifyResultShapes should return an Attribute if and only if the respective dimension is static.

This fixes #64256.

[kuhar]

LGTM

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

 /// For reshape op compute the shape at dimension `dimIndex` of the output in
 /// terms of shape of the `src`, when the reshape op is a collapsing
 /// operation. It is the product of the shape of the collapsed dimensions of the
 /// `src`.
 static OpFoldResult getCollapsedOutputDimFromInputShape(
     OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociationMap) {
   if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
+    // Static dimension: return Attribute.
     return builder.getIndexAttr(dstStaticShape[dimIndex]);
   }
   AffineMap map = reassociationMap[dimIndex];
   unsigned startPos =
       map.getResults().front().cast<AffineDimExpr>().getPosition();
   unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
   AffineExpr expr;
   SmallVector<OpFoldResult> dynamicDims;
   for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
     dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim));
     AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
     expr = (expr ? expr * currExpr : currExpr);
   }
-  return affine::makeComposedFoldedAffineApply(
+
+  // Dynamic dimension: return Value.
+  return affine::makeComposedAffineApply(
              builder, loc, AffineMap::get(0, endPos - startPos + 1, expr),
-      dynamicDims);
+             dynamicDims)
+      ->getResult(0);
 }
 
 /// Given the `src` of a collapsing reshape op and its reassociation maps,
 /// compute the shape of the result of the reshape.
 static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
         return getCollapsedOutputDimFromInputShape(
             builder, loc, dim, src, dstStaticShape, reassociation);
       }));
 }
 
 /// For an expanding reshape op, compute the value for a dimension of the output
 /// from the shape of the input.
 static OpFoldResult getExpandedOutputDimFromInputShape(
     OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation,
     llvm::DenseMap<int64_t, int64_t> &expandedDimToCollapsedDim) {
   if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
+    // Static dimension: return Attribute.
     return builder.getIndexAttr(dstStaticShape[dimIndex]);
   }
   unsigned sourceDimPos = expandedDimToCollapsedDim[dimIndex];
   unsigned startPos = reassociation[sourceDimPos]
                           .getResults()
                           .front()
                           .cast<AffineDimExpr>()
                           .getPosition();
if (d.index() + startPos == static_cast<unsigned>(dimIndex))
       continue;
     assert(!ShapedType::isDynamic(d.value()) &&
            "single dimension cannot be expanded into multiple dynamic "
            "dimensions");
     linearizedStaticDim *= d.value();
   }
   OpFoldResult sourceDim =
       builder.create<tensor::DimOp>(loc, src, sourceDimPos).getResult();
-  return affine::makeComposedFoldedAffineApply(
+
+  // Dynamic dimension: return Value.
+  return affine::makeComposedAffineApply(
              builder, loc,
              AffineMap::get(
-          0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
-      sourceDim);
+                 0, 1,
+                 builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
+             sourceDim)
+      ->getResult(0);
 }
 
 /// Given the `src` of an expanding reshape op, the reassociation maps and the
 /// result type, compute the shape of the result of the reshape.
 static SmallVector<OpFoldResult, 4> getExpandedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim =

mlir/test/Dialect/Tensor/resolve-shaped-type-result-dims.mlir

func.func @extract_slice_rank_reduced_6(%arg0 : tensor<?x?x?xf32>, %arg1 : index,
   %2 = tensor.dim %0, %c2 : tensor<?x1x?xf32>
   return %1, %2 : index, index
 }
 // CHECK-LABEL: func @extract_slice_rank_reduced_6(
 //  CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?x?xf32>
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9_]+]]: index
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9_]+]]: index
 //       CHECK:   return %[[ARG1]], %[[ARG2]]
+
+// -----
+
+func.func @collapse_shape() -> index {
+  %c0 = arith.constant 0 : index
+  %c7 = arith.constant 7 : index
+  %c1_i16 = arith.constant 1 : i16
+  %generated = tensor.generate %c7 {
+  ^bb0(%arg3: index, %arg4: index):
+    tensor.yield %c1_i16 : i16
+  } : tensor<?x22xi16>
+  %collapsed = tensor.collapse_shape %generated [[0, 1]] : tensor<?x22xi16> into tensor<?xi16>
+  %d0 = tensor.dim %collapsed, %c0 : tensor<?xi16>
+  return %d0 : index
+}
+// CHECK-LABEL: func @collapse_shape(
+//       CHECK:   %[[c154:.*]] = arith.constant 154 : index
+//       CHECK:   return %[[c154]]