[mlir][memref] Improve type inference for rank-reducing subviews

[springerm]

Summary

The result shape of a rank-reducing subview cannot be inferred in the general case. Just the result rank is not enough. The only thing that we can infer is the layout map.

This change also improves the bufferization patterns of tensor.extract_slice and tensor.insert_slice to fully support rank-reducing operations.

Depends On D129143

[springerm]

update

[nicolasvasilache]

Nice!

LGTM conditioned on reusing helpers rather than reinventing.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

2165

Please use/adapt/evolve the helper available in BuiltinTypes.h

llvm::Optional<llvm::SmallDenseSet<unsigned>>
mlir::computeRankReductionMask(ArrayRef<int64_t> originalShape,
                               ArrayRef<int64_t> reducedShape);

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

933

Note: I went through some cleanup and doc effort in https://reviews.llvm.org/D128803.

I haven't yet gone through it for SubViewOp, could you please do it in a followup PR?
I think after that this would still be called inferRankReducedResultType and not inferCanonicalRankReducedTesultType but I am not 100% sure, I'll let you decide.

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

220

Please use std/llvm::copy_if

[springerm]

address comments

[springerm]

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

933

Yes, it should still be called inferRankReducedResultType. I added similar documentation (wrt. ambiguity) to the op, but I added it to inferRankReducedResultType instead of the op documentation, because it only affects the C++ API. I think the documentation is good, but let me know if you think something is missing.

[nicolasvasilache]

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

933

We should have a similar section as https://reviews.llvm.org/D128803 in SubViewOp on differences between inference and verification (with a special discussion for layout).

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     static Type inferResultType(MemRefType sourceMemRefType,
                                 ArrayRef<int64_t> staticOffsets,
                                 ArrayRef<int64_t> staticSizes,
                                 ArrayRef<int64_t> staticStrides);
     static Type inferResultType(MemRefType sourceMemRefType,
                                 ArrayRef<OpFoldResult> staticOffsets,
                                 ArrayRef<OpFoldResult> staticSizes,
                                 ArrayRef<OpFoldResult> staticStrides);
-    static Type inferRankReducedResultType(unsigned resultRank,
+
+    /// A rank-reducing result type can be inferred from the desired result
+    /// shape. Only the layout map is inferred.
+    ///
+    /// Note: The result shape cannot be inferred with just the result rank and
+    /// and the desired sizes. In case there are more "ones" among the sizes
+    /// than the difference in source/result rank, it is not clear which dims of
+    /// size one should be dropped.
+    static Type inferRankReducedResultType(ArrayRef<int64_t> resultShape,
                                            MemRefType sourceMemRefType,
                                            ArrayRef<int64_t> staticOffsets,
                                            ArrayRef<int64_t> staticSizes,
                                            ArrayRef<int64_t> staticStrides);
-    static Type inferRankReducedResultType(unsigned resultRank,
+    static Type inferRankReducedResultType(ArrayRef<int64_t> resultShape,
                                            MemRefType sourceMemRefType,
                                            ArrayRef<OpFoldResult> staticOffsets,
                                            ArrayRef<OpFoldResult> staticSizes,
                                            ArrayRef<OpFoldResult> staticStrides);
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
     std::array<unsigned, 3> getArrayAttrMaxRanks() {

mlir/lib/Dialect/Bufferization/Transforms/AllocTensorElimination.cpp

return eliminateAllocTensors(
                             insertSliceOp.getStrides().end());
         neededValues.push_back(insertSliceOp.getDest());
 
         return true;
       },
       /*rewriteFunc=*/
       [](OpBuilder &b, Location loc, OpOperand &operand) {
         auto insertOp = cast<tensor::InsertSliceOp>(operand.getOwner());
-        // Expand offsets, sizes and strides to the full rank to handle the
-        // rank-reducing case.
-        SmallVector<OpFoldResult> mixedOffsets = insertOp.getMixedOffsets();
-        SmallVector<OpFoldResult> mixedSizes = insertOp.getMixedSizes();
-        SmallVector<OpFoldResult> mixedStrides = insertOp.getMixedStrides();
-        OffsetSizeAndStrideOpInterface::expandToRank(
-            insertOp.getDest(), mixedOffsets, mixedSizes, mixedStrides,
-            [&](Value target, int64_t dim) -> OpFoldResult {
-              auto shapedType = target.getType().cast<ShapedType>();
-              if (shapedType.isDynamicDim(dim))
-                return b.create<tensor::DimOp>(loc, target, dim).getResult();
-              return b.getIndexAttr(shapedType.getDimSize(dim));
-            });
-        auto t = tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
-            insertOp.getSourceType().getRank(),
-            insertOp.getDest().getType().cast<RankedTensorType>(), mixedOffsets,
-            mixedSizes, mixedStrides);
         auto extractOp = b.create<tensor::ExtractSliceOp>(
-            loc, t, insertOp.getDest(), mixedOffsets, mixedSizes, mixedStrides);
+            loc, insertOp.getSourceType(), insertOp.getDest(),
+            insertOp.getMixedOffsets(), insertOp.getMixedSizes(),
+            insertOp.getMixedStrides());
         return extractOp.getResult();
       });
 }
 
 namespace {
 struct AllocTensorElimination
     : public AllocTensorEliminationBase<AllocTensorElimination> {
   AllocTensorElimination() = default;

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
   dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
                              ShapedType::kDynamicSize);
   dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides,
                              ShapedType::kDynamicStrideOrOffset);
   return SubViewOp::inferResultType(sourceMemRefType, staticOffsets,
                                     staticSizes, staticStrides);
 }
 
-Type SubViewOp::inferRankReducedResultType(unsigned resultRank,
+Type SubViewOp::inferRankReducedResultType(ArrayRef<int64_t> resultShape,
                                            MemRefType sourceRankedTensorType,
                                            ArrayRef<int64_t> offsets,
                                            ArrayRef<int64_t> sizes,
                                            ArrayRef<int64_t> strides) {
   auto inferredType =
       inferResultType(sourceRankedTensorType, offsets, sizes, strides)
           .cast<MemRefType>();
-  assert(inferredType.getRank() >= resultRank && "expected ");
-  int rankDiff = inferredType.getRank() - resultRank;
-  if (rankDiff > 0) {
-    auto shape = inferredType.getShape();
-    llvm::SmallBitVector dimsToProject =
-        getPositionsOfShapeOne(rankDiff, shape);
-    SmallVector<int64_t> projectedShape;
-    for (unsigned pos = 0, e = shape.size(); pos < e; ++pos)
-      if (!dimsToProject.test(pos))
-        projectedShape.push_back(shape[pos]);
-
-    AffineMap map =
-        getProjectedMap(inferredType.getLayout().getAffineMap(), dimsToProject);
-    inferredType =
-        MemRefType::get(projectedShape, inferredType.getElementType(), map,
-                        inferredType.getMemorySpace());
-  }
+  assert(inferredType.getRank() >= resultShape.size() && "expected ");
+  if (inferredType.getRank() == resultShape.size())
     return inferredType;
+
+  // Compute which dimensions are dropped.
+  Optional<llvm::SmallDenseSet<unsigned>> dimsToProject =
+      computeRankReductionMask(inferredType.getShape(), resultShape);
+  assert(dimsToProject.hasValue() && "invalid rank reduction");
+  llvm::SmallBitVector dimsToProjectVector(inferredType.getRank());
+  for (unsigned dim : *dimsToProject)
+    dimsToProjectVector.set(dim);
+
+  // Compute layout map and result type.
+  AffineMap map = getProjectedMap(inferredType.getLayout().getAffineMap(),
+                                  dimsToProjectVector);
+  return MemRefType::get(resultShape, inferredType.getElementType(), map,
+                         inferredType.getMemorySpace());
 }
 
-Type SubViewOp::inferRankReducedResultType(unsigned resultRank,
+Type SubViewOp::inferRankReducedResultType(ArrayRef<int64_t> resultShape,
                                            MemRefType sourceRankedTensorType,
                                            ArrayRef<OpFoldResult> offsets,
                                            ArrayRef<OpFoldResult> sizes,
                                            ArrayRef<OpFoldResult> strides) {
   SmallVector<int64_t> staticOffsets, staticSizes, staticStrides;
   SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides;
   dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets,
                              ShapedType::kDynamicStrideOrOffset);
   dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
                              ShapedType::kDynamicSize);
   dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides,
                              ShapedType::kDynamicStrideOrOffset);
   return SubViewOp::inferRankReducedResultType(
-      resultRank, sourceRankedTensorType, staticOffsets, staticSizes,
+      resultShape, sourceRankedTensorType, staticOffsets, staticSizes,
       staticStrides);
 }
+
 // Build a SubViewOp with mixed static and dynamic entries and custom result
 // type. If the type passed is nullptr, it is inferred.
 void SubViewOp::build(OpBuilder &b, OperationState &result,
                       MemRefType resultType, Value source,
                       ArrayRef<OpFoldResult> offsets,
                       ArrayRef<OpFoldResult> sizes,
                       ArrayRef<OpFoldResult> strides,
                       ArrayRef<NamedAttribute> attrs) {

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp

for (OpOperand &use : oldOp->getUses()) {
     if (!subviewUse) {
       // Save the operand to and replace outside the loop to not invalidate the
       // iterator.
       operandsToReplace.push_back(&use);
       continue;
     }
     builder.setInsertionPoint(subviewUse);
     Type newType = memref::SubViewOp::inferRankReducedResultType(
-        subviewUse.getType().getRank(), val.getType().cast<MemRefType>(),
+        subviewUse.getType().getShape(), val.getType().cast<MemRefType>(),
         extractFromI64ArrayAttr(subviewUse.static_offsets()),
         extractFromI64ArrayAttr(subviewUse.static_sizes()),
         extractFromI64ArrayAttr(subviewUse.static_strides()));
     Value newSubview = builder.create<memref::SubViewOp>(
         subviewUse->getLoc(), newType.cast<MemRefType>(), val,
         subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),
         subviewUse.getMixedStrides());
     replaceUsesAndPropagateType(subviewUse, newSubview, builder);
LogicalResult mlir::memref::multiBuffer(memref::AllocOp allocOp,
   offsets.push_back(bufferIndex);
   offsets.append(oldShape.size(), builder.getIndexAttr(0));
   strides.assign(oldShape.size() + 1, builder.getIndexAttr(1));
   sizes.push_back(builder.getIndexAttr(1));
   for (int64_t size : oldShape)
     sizes.push_back(builder.getIndexAttr(size));
   auto dstMemref =
       memref::SubViewOp::inferRankReducedResultType(
-          allocOp.getType().getRank(), newMemref, offsets, sizes, strides)
+          allocOp.getType().getShape(), newMemref, offsets, sizes, strides)
           .cast<MemRefType>();
   Value subview = builder.create<memref::SubViewOp>(loc, dstMemref, newAlloc,
                                                     offsets, sizes, strides);
   replaceUsesAndPropagateType(allocOp, subview, builder);
   allocOp.erase();
   return success();
 }

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

struct ExtractSliceOpInterface
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::None;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
+    SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
     Location loc = extractSliceOp.getLoc();
 
-    // Even if this op was decided to bufferize out-of-place, do not insert the
-    // buffer copy yet. This is done later in this function.
+    // Get source buffer.
     FailureOr<Value> srcMemref =
         getBuffer(rewriter, extractSliceOp.getSource(), options);
     if (failed(srcMemref))
       return failure();
     auto srcMemrefType = srcMemref->getType().cast<MemRefType>();
-    auto dstTensorType =
-        extractSliceOp.getResult().getType().cast<RankedTensorType>();
 
-    // Expand offsets, sizes and strides to the full rank to handle the
-    // rank-reducing case.
-    SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
-    SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
-    SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
-    OffsetSizeAndStrideOpInterface::expandToRank(
-        *srcMemref, mixedOffsets, mixedSizes, mixedStrides,
-        [&](Value target, int64_t dim) -> OpFoldResult {
-          auto shapedType = target.getType().cast<ShapedType>();
-          if (shapedType.isDynamicDim(dim))
-            return rewriter.create<memref::DimOp>(loc, target, dim).result();
-          return rewriter.getIndexAttr(shapedType.getDimSize(dim));
-        });
-    // Bufferize to subview.
-    auto subviewMemRefType = memref::SubViewOp::inferRankReducedResultType(
-                                 dstTensorType.getRank(), srcMemrefType,
-                                 mixedOffsets, mixedSizes, mixedStrides)
-                                 .cast<MemRefType>();
+    // Take a subview of the source buffer.
+    auto subviewMemRefType =
+        memref::SubViewOp::inferRankReducedResultType(
+            extractSliceOp.getType().getShape(), srcMemrefType, mixedOffsets,
+            mixedSizes, mixedStrides)
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, *srcMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
 };
struct InsertSliceOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     // insert_slice ops arise from tiling and bufferizing them out-of-place is
     // generally a deal breaker. When used with loops, this ends up cloning the
     // whole tensor on every single iteration and is a symptom of a
     // catastrophically bad scheduling decision.
     // TODO: be very loud about it or even consider failing the pass.
     auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
+    SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
     Location loc = insertSliceOp.getLoc();
+
+    // Get destination buffer.
     FailureOr<Value> dstMemref =
         getBuffer(rewriter, insertSliceOp.getDest(), options);
     if (failed(dstMemref))
       return failure();
 
-    // Expand offsets, sizes and strides to the full rank to handle the
-    // rank-reducing case.
-    SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
-    SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
-    SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
-    OffsetSizeAndStrideOpInterface::expandToRank(
-        *dstMemref, mixedOffsets, mixedSizes, mixedStrides,
-        [&](Value target, int64_t dim) -> OpFoldResult {
-          auto shapedType = target.getType().cast<ShapedType>();
-          if (shapedType.isDynamicDim(dim))
-            return rewriter.create<memref::DimOp>(loc, target, dim).result();
-          return rewriter.getIndexAttr(shapedType.getDimSize(dim));
-        });
-    // Take a subview of the dst.
+    // Take a subview of the destination buffer.
     auto dstMemrefType = dstMemref->getType().cast<MemRefType>();
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
-            insertSliceOp.getSourceType().getRank(), dstMemrefType,
+            insertSliceOp.getSourceType().getShape(), dstMemrefType,
             mixedOffsets, mixedSizes, mixedStrides)
             .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // Copy tensor. If this tensor.insert_slice has a matching
     // tensor.extract_slice, the copy operation will eventually fold away.
if (failed(destBuffer))
       return failure();
 
     // Bufferize the ParallelInsertSliceOp outside of `parallelCombiningParent`.
     rewriter.setInsertionPoint(parallelCombiningParent);
     FailureOr<Value> srcBuffer =
         getBuffer(rewriter, parallelInsertSliceOp.getSource(), options);
     if (failed(srcBuffer))
       return failure();
+
+    // Take a subview of the destination buffer.
+    auto destBufferType = destBuffer->getType().cast<MemRefType>();
+    auto subviewMemRefType =
+        memref::SubViewOp::inferRankReducedResultType(
+            parallelInsertSliceOp.getSourceType().getShape(), destBufferType,
+            parallelInsertSliceOp.getMixedOffsets(),
+            parallelInsertSliceOp.getMixedSizes(),
+            parallelInsertSliceOp.getMixedStrides())
+            .cast<MemRefType>();
     Value subview = rewriter.create<memref::SubViewOp>(
-        parallelInsertSliceOp.getLoc(), *destBuffer,
+        parallelInsertSliceOp.getLoc(), subviewMemRefType, *destBuffer,
         parallelInsertSliceOp.getMixedOffsets(),
         parallelInsertSliceOp.getMixedSizes(),
         parallelInsertSliceOp.getMixedStrides());
+
     // This memcpy will fold away if everything bufferizes in-place.
     if (failed(options.createMemCpy(rewriter, parallelInsertSliceOp.getLoc(),
                                     *srcBuffer, subview)))
       return failure();
 
     // Replace all uses of parallelIteratingOp (just the corresponding result).
     rewriter.setInsertionPointAfter(parallelIteratingOp);
     Value toTensorOp =

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

void TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
   read.replaceAllUsesWith(lastwrite.getVector());
   opToErase.push_back(read.getOperation());
 }
 
 /// Drops unit dimensions from the input MemRefType.
 static MemRefType dropUnitDims(MemRefType inputType, ArrayRef<int64_t> offsets,
                                ArrayRef<int64_t> sizes,
                                ArrayRef<int64_t> strides) {
+  SmallVector<int64_t> targetShape = llvm::to_vector(
+      llvm::make_filter_range(sizes, [](int64_t sz) { return sz != 1; }));
   Type rankReducedType = memref::SubViewOp::inferRankReducedResultType(
-      0, inputType, offsets, sizes, strides);
+      targetShape, inputType, offsets, sizes, strides);
   return canonicalizeStridedLayout(rankReducedType.cast<MemRefType>());
 }
 
 /// Creates a rank-reducing memref.subview op that drops unit dims from its
 /// input. Or just returns the input if it was already without unit dims.
 static Value rankReducingSubviewDroppingUnitDims(PatternRewriter &rewriter,
                                                  mlir::Location loc,
                                                  Value input) {

mlir/test/Dialect/Tensor/bufferize.mlir

 }
 
 // -----
 
 // CHECK-LABEL: func @tensor.insert_slice(
 //  CHECK-SAME:     %[[t1:.*]]: tensor<?x?xf32>, %[[t2:.*]]: tensor<?x10xf32>,
 //  CHECK-SAME:     %[[idx1:.*]]: index, %[[idx2:.*]]: index
 func.func @tensor.insert_slice(%t1: tensor<?x?xf32>, %t2: tensor<?x10xf32>,
                                %idx1: index, %idx2: index) -> tensor<?x?xf32> {
   // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
   // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
   // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x?xf32>
   // CHECK-DAG: %[[m2:.*]] = bufferization.to_memref %[[t2]] : memref<?x10xf32>
   // CHECK-DAG: %[[dim0:.*]] = memref.dim %[[m1]], %[[c0]]
   // CHECK-DAG: %[[dim1:.*]] = memref.dim %[[m1]], %[[c1]]
   //     CHECK: %[[alloc:.*]] = memref.alloc(%[[dim0]], %[[dim1]])
   //     CHECK: memref.copy %[[m1]], %[[alloc]]
   //     CHECK: %[[subview:.*]] = memref.subview %[[alloc]][%[[idx1]], 5] [%[[idx2]], 10] [1, 1]
   //     CHECK: memref.copy %[[m2]], %[[subview]]
   %0 = tensor.insert_slice %t2 into %t1[%idx1, 5][%idx2, 10][1, 1]
       : tensor<?x10xf32> into tensor<?x?xf32>
 
   //     CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
   //     CHECK: return %[[r]]
   return %0 : tensor<?x?xf32>
 }
 
 // -----
 
+// CHECK: #[[$MAP11:.*]] = affine_map<()[s0] -> (s0)>
+
+// CHECK-LABEL: func @tensor.insert_slice_rank_reducing_1(
+func.func @tensor.insert_slice_rank_reducing_1(
+    %t1: tensor<?x?xf32>, %f: tensor<f32>, %idx1: index, %idx2: index)
+  -> tensor<?x?xf32>
+{
+  // CHECK: %[[alloc:.*]] = memref.alloc{{.*}} : memref<?x?xf32>
+  // CHECK: memref.subview %[[alloc]][%{{.*}}, %{{.*}}] [1, 1] [1, 1] : memref<?x?xf32> to memref<f32, #[[$MAP11]]>
+  // CHECK: memref.copy {{.*}} : memref<f32> to memref<f32, #[[$MAP11]]>
+  %0 = tensor.insert_slice %f into %t1[%idx1, %idx2][1, 1][1, 1]
+      : tensor<f32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK: #[[$MAP12:.*]] = affine_map<(d0, d1, d2, d3, d4)[s0, s1, s2, s3, s4, s5] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4 + d4 * s5)>
+
+// CHECK-LABEL: func @tensor.insert_slice_rank_reducing_2(
+func.func @tensor.insert_slice_rank_reducing_2(
+    %t1: tensor<?x?x?x?x?x?x?xf32>, %t2: tensor<2x1x4x1x1xf32>, %i: index)
+  -> tensor<?x?x?x?x?x?x?xf32>
+{
+  // CHECK: %[[alloc:.*]] = memref.alloc{{.*}} : memref<?x?x?x?x?x?x?xf32>
+  // CHECK: memref.subview %[[alloc]][{{.*}}] [1, 2, 1, 4, 1, 1, 1] [1, 1, 1, 1, 1, 1, 1] : memref<?x?x?x?x?x?x?xf32> to memref<2x1x4x1x1xf32, #[[$MAP12]]>
+  // CHECK: memref.copy {{.*}} : memref<2x1x4x1x1xf32> to memref<2x1x4x1x1xf32, #[[$MAP12]]>
+  %0 = tensor.insert_slice %t2 into %t1[%i, %i, %i, %i, %i, %i, %i][1, 2, 1, 4, 1, 1, 1][1, 1, 1, 1, 1, 1, 1]
+      : tensor<2x1x4x1x1xf32> into tensor<?x?x?x?x?x?x?xf32>
+  return %0 : tensor<?x?x?x?x?x?x?xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @tensor.insert(
 //  CHECK-SAME:     %[[t1:.*]]: tensor<5xf32>, %[[idx1:.*]]: index,
 //  CHECK-SAME:     %[[f:.*]]: f32
 func.func @tensor.insert(%t1: tensor<5xf32>, %idx1: index, %f: f32) -> tensor<5xf32> {
   // CHECK-DAG: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
   // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<5xf32>
   // CHECK: memref.copy %[[m1]], %[[alloc]]
   // CHECK: memref.store %[[f]], %[[alloc]][%[[idx1]]]

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

%9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) {
       %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32>
       scf.yield %12 : tensor<1x6x8xf32>
     }
     %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
     scf.yield %10 : tensor<?x1x6x8xf32>
   }
   return %5: tensor<?x1x6x8xf32>
 }
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+
+// CHECK-LABEL: func.func @rank_reducing_parallel_insert_slice
+func.func @rank_reducing_parallel_insert_slice(%in: tensor<100xf32>, %out: tensor<200x100xf32>) {
+  %c1 = arith.constant 1 : index
+  %num_threads = arith.constant 100 : index
+
+  // CHECK: scf.foreach_thread {{.*}} {
+  %result = scf.foreach_thread (%thread_idx) in (%num_threads) -> tensor<200x100xf32> {
+      %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
+      scf.foreach_thread.perform_concurrently {
+        // CHECK: memref.subview %{{.*}}[%{{.*}}] [1] [1] : memref<100xf32, #[[$MAP0]]> to memref<1xf32, #[[$MAP0]]>
+        // CHECK: memref.subview %{{.*}}[1, %{{.*}}] [1, 1] [1, 1] : memref<200x100xf32, #[[$MAP1]]> to memref<1xf32, #[[$MAP0]]>
+        tensor.parallel_insert_slice %1 into %out[1, %thread_idx][1, 1][1, 1] :
+          tensor<1xf32> into tensor<200x100xf32>
+      }
+  }
+  // CHECK: }
+  return
+}

mlir/unittests/Dialect/MemRef/InferShapeTest.cpp

 using namespace mlir::memref;
 
 // Source memref has identity layout.
 TEST(InferShapeTest, inferRankReducedShapeIdentity) {
   MLIRContext ctx;
   OpBuilder b(&ctx);
   auto sourceMemref = MemRefType::get({10, 5}, b.getIndexType());
   auto reducedType = SubViewOp::inferRankReducedResultType(
-      /*resultRank=*/1, sourceMemref, {2, 3}, {1, 2}, {1, 1});
+      /*resultShape=*/{2}, sourceMemref, {2, 3}, {1, 2}, {1, 1});
   AffineExpr dim0;
   bindDims(&ctx, dim0);
   auto expectedType =
       MemRefType::get({2}, b.getIndexType(), AffineMap::get(1, 0, dim0 + 13));
   EXPECT_EQ(reducedType, expectedType);
 }
 
 // Source memref has non-identity layout.
 TEST(InferShapeTest, inferRankReducedShapeNonIdentity) {
   MLIRContext ctx;
   OpBuilder b(&ctx);
   AffineExpr dim0, dim1;
   bindDims(&ctx, dim0, dim1);
   auto sourceMemref = MemRefType::get({10, 5}, b.getIndexType(),
                                       AffineMap::get(2, 0, 1000 * dim0 + dim1));
   auto reducedType = SubViewOp::inferRankReducedResultType(
-      /*resultRank=*/1, sourceMemref, {2, 3}, {1, 2}, {1, 1});
+      /*resultShape=*/{2}, sourceMemref, {2, 3}, {1, 2}, {1, 1});
   auto expectedType =
       MemRefType::get({2}, b.getIndexType(), AffineMap::get(1, 0, dim0 + 2003));
   EXPECT_EQ(reducedType, expectedType);
 }
 
 TEST(InferShapeTest, inferRankReducedShapeToScalar) {
   MLIRContext ctx;
   OpBuilder b(&ctx);
   AffineExpr dim0, dim1;
   bindDims(&ctx, dim0, dim1);
   auto sourceMemref = MemRefType::get({10, 5}, b.getIndexType(),
                                       AffineMap::get(2, 0, 1000 * dim0 + dim1));
   auto reducedType = SubViewOp::inferRankReducedResultType(
-      /*resultRank=*/0, sourceMemref, {2, 3}, {1, 1}, {1, 1});
+      /*resultShape=*/{}, sourceMemref, {2, 3}, {1, 1}, {1, 1});
   auto expectedType =
       MemRefType::get({}, b.getIndexType(),
                       AffineMap::get(0, 0, b.getAffineConstantExpr(2003)));
   EXPECT_EQ(reducedType, expectedType);
 }