[mlir][tensor] Support more cases in MergeConsecutiveExtractSlice

[antiagainst]

Summary

This commit adds utility functions to perform general merging of
OffsetSizeAndStrideOpInterface by supporting producer rank
reducing and non-unit strides.

With it we can extend MergeConsecutiveExtractSlice to support
more cases.

[ThomasRaoux]

LGTM

[mehdi_amini]

Seems like this is failing when built with GCC as host compiler: https://buildkite.com/mlir/mlir-core/builds/25622#01835d94-e233-483d-a001-8a9dadc530b7

[christopherbate]

Seems like the utility function mergeOffsetsSizesAndStrides should be added to ViewLikeInterface or Dialect/Utils because it has greater scope. The primary use that comes to mind is that the memref "fold alias ops" (used to be "fold subviews") essential does the same thing (SubView and ExtractSlice are both implement the OffsetSizeAndStrideOpInterface). Therefore, it seems like an opportunity for code de-duplication.

Edit: I tried doing this in https://reviews.llvm.org/D134393 and run into the issue with Interface/ViewLikeInterface depending on Dialect/Utils and use of Affine ops. So the only common place it could go to be used by Memref and Tensor dialects is in affine dialect utils.

[christopherbate]

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

23

All three of these helper functions seem superfluous because the the makeComposedFoldedAffineApply does the optimization in getAffineExpr and the expressions aren't complicated enough to require the add and mul helpers. I tested the below code:

LogicalResult mlir::mergeOffsetsSizesAndStrides(
    OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
    ArrayRef<OpFoldResult> producerSizes,
    ArrayRef<OpFoldResult> producerStrides,
    const llvm::SmallBitVector &droppedProducerDims,
    ArrayRef<OpFoldResult> consumerOffsets,
    ArrayRef<OpFoldResult> consumerSizes,
    ArrayRef<OpFoldResult> consumerStrides,
    SmallVector<OpFoldResult> &combinedOffsets,
    SmallVector<OpFoldResult> &combinedSizes,
    SmallVector<OpFoldResult> &combinedStrides) {
  combinedOffsets.resize(producerOffsets.size());
  combinedSizes.resize(producerOffsets.size());
  combinedStrides.resize(producerOffsets.size());

  AffineExpr s0, s1, d0;
  bindDims(builder.getContext(), d0);
  bindSymbols(builder.getContext(), s0, s1);

  unsigned consumerPos = 0;
  for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
    if (droppedProducerDims.test(i)) {
      // For dropped dims, get the values from the producer.
      combinedOffsets[i] = producerOffsets[i];
      combinedSizes[i] = producerSizes[i];
      combinedStrides[i] = producerStrides[i];
      continue;
    }
    SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
    // The combined offset is computed as
    //    producer_offset + consumer_offset * producer_strides.
    combinedOffsets[i] = makeComposedFoldedAffineApply(
        builder, loc, d0 * s0 + s1,
        {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
    combinedSizes[i] = consumerSizes[consumerPos];
    // The combined stride is computed as
    //    consumer_stride * producer_stride.
    combinedStrides[i] = makeComposedFoldedAffineApply(
        builder, loc, d0 * s0,
        {consumerStrides[consumerPos], producerStrides[i]});

    consumerPos++;
  }
  return success();
}

[mravishankar]

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

23

Thats great! This was developed in IREE first, and that time we didnt have makeComposedFoldedAffineApply. I did suspect that those would make these methods redundant.

[antiagainst]

Thanks @mehdi_amini for the revert! I'll address it tomorrow.

Thanks @christopherbate for the suggestion! Yeah I didn't put in the ViewLikeInterface impl cpp because of the affine op dependency. Didn't know about the makeComposedFoldedAffineApply; that's pretty nice! Let me update accordingly tomorrow.

[antiagainst]

Relanded https://reviews.llvm.org/rGbd81524e7f52c38ad0bc689934343a476e545265 and uploaded https://reviews.llvm.org/D134487 to move to a more common place to be shared with memref too.

mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h

 //===- TransformsUtils.h - Tensor Transformation Utilities-------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMUTILS_H
 #define MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMUTILS_H
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 namespace tensor {
 
+/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
+/// when combining a producer slice **into** a consumer slice.
+///
+/// This function performs the following computation:
+/// - Combined offsets = producer_offsets * consumer_strides + consumer_offsets
+/// - Combined sizes = consumer_sizes
+/// - Combined strides = producer_strides * consumer_strides
+LogicalResult
+mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
+                            ArrayRef<OpFoldResult> producerOffsets,
+                            ArrayRef<OpFoldResult> producerSizes,
+                            ArrayRef<OpFoldResult> producerStrides,
+                            const llvm::SmallBitVector &droppedProducerDims,
+                            ArrayRef<OpFoldResult> consumerOffsets,
+                            ArrayRef<OpFoldResult> consumerSizes,
+                            ArrayRef<OpFoldResult> consumerStrides,
+                            SmallVector<OpFoldResult> &combinedOffsets,
+                            SmallVector<OpFoldResult> &combinedSizes,
+                            SmallVector<OpFoldResult> &combinedStrides);
+
+/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
+/// when combining a `producer` slice op **into** a `consumer` slice op.
+LogicalResult
+mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
+                            OffsetSizeAndStrideOpInterface producer,
+                            OffsetSizeAndStrideOpInterface consumer,
+                            const llvm::SmallBitVector &droppedProducerDims,
+                            SmallVector<OpFoldResult> &combinedOffsets,
+                            SmallVector<OpFoldResult> &combinedSizes,
+                            SmallVector<OpFoldResult> &combinedStrides);
+
 //===----------------------------------------------------------------------===//
 // Extract slice from `tensor.collapse_shape`
 //===----------------------------------------------------------------------===//
 
 /// This class assists with generating IR required to materialize an
 /// arbitrary-sized slice from the result of a CollapseShapeOp. In order to
 /// accomplish this, a loop nest or similar operation must be created by the
 /// caller. The purpose of the loop nest is to generate a "tiling by 1" of all

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

 //===- MergeConsecutiveInsertExtractSlicePatterns.cpp ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/TransformUtils.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
-/// Adds each corresponding pair of offsets in `offsets1` and `offsets2` and
-/// returns the results.
-static SmallVector<OpFoldResult> mergeOffsets(Location loc,
-                                              ArrayRef<OpFoldResult> offsets1,
-                                              ArrayRef<OpFoldResult> offsets2,
-                                              OpBuilder &builder) {
-  SmallVector<OpFoldResult> foldedOffsets;
-  assert(offsets1.size() == offsets2.size());
-  foldedOffsets.reserve(offsets1.size());
-
-  AffineExpr dim1, dim2;
-  bindDims(builder.getContext(), dim1, dim2);
-
-  for (const auto &pair : llvm::zip(offsets1, offsets2)) {
-    auto offset0 =
-        getValueOrCreateConstantIndexOp(builder, loc, std::get<0>(pair));
-    auto offset1 =
-        getValueOrCreateConstantIndexOp(builder, loc, std::get<1>(pair));
-    auto foldedOffset =
-        makeComposedAffineApply(builder, loc, dim1 + dim2, {offset0, offset1});
-    foldedOffsets.push_back(foldedOffset.getResult());
+/// Creates AffineExpr from `ofr`: if the OpFoldResult is a Value, creates a
+/// AffineSymbolExpr and appends it to `symbols`; otherwise creates a
+/// AffineConstantExpr.
+static AffineExpr getAffineExpr(OpFoldResult ofr,
+                                SmallVector<OpFoldResult> &symbols) {
+  if (auto attr = ofr.dyn_cast<Attribute>()) {
+    return getAffineConstantExpr(attr.cast<IntegerAttr>().getInt(),
+                                 attr.getContext());
   }
-  return foldedOffsets;
+  Value v = ofr.get<Value>();
+  AffineExpr expr = getAffineSymbolExpr(symbols.size(), v.getContext());
+  symbols.push_back(v);
+  return expr;
+}
+
+/// Builds the AffineExpr incrementally for arithmetic operations.
+static AffineExpr add(AffineExpr expr, OpFoldResult ofr,
+                      SmallVector<OpFoldResult> &symbols) {
+  return expr + getAffineExpr(ofr, symbols);
+}
+static AffineExpr mul(OpFoldResult lhs, OpFoldResult rhs,
+                      SmallVector<OpFoldResult> &symbols) {
+  return getAffineExpr(lhs, symbols) * getAffineExpr(rhs, symbols);
+}
+
+/// Converts an AffineExpr to OpFoldResult by generating an `affine.apply`
+/// op and fold it.
+static OpFoldResult getOpFoldResult(OpBuilder &builder, Location loc,
+                                    AffineExpr expr,
+                                    SmallVector<OpFoldResult> &symbols) {
+  AffineMap m = AffineMap::get(0, symbols.size(), expr);
+  return makeComposedFoldedAffineApply(builder, loc, m, symbols);
+}
+
+LogicalResult tensor::mergeOffsetsSizesAndStrides(
+    OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
+    ArrayRef<OpFoldResult> producerSizes,
+    ArrayRef<OpFoldResult> producerStrides,
+    const llvm::SmallBitVector &droppedProducerDims,
+    ArrayRef<OpFoldResult> consumerOffsets,
+    ArrayRef<OpFoldResult> consumerSizes,
+    ArrayRef<OpFoldResult> consumerStrides,
+    SmallVector<OpFoldResult> &combinedOffsets,
+    SmallVector<OpFoldResult> &combinedSizes,
+    SmallVector<OpFoldResult> &combinedStrides) {
+  combinedOffsets.resize(producerOffsets.size());
+  combinedSizes.resize(producerOffsets.size());
+  combinedStrides.resize(producerOffsets.size());
+  unsigned consumerPos = 0;
+  for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
+    if (droppedProducerDims.test(i)) {
+      // For dropped dims, get the values from the producer.
+      combinedOffsets[i] = producerOffsets[i];
+      combinedSizes[i] = producerSizes[i];
+      combinedStrides[i] = producerStrides[i];
+      continue;
+    }
+    SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
+    // The combined offset is computed as
+    //    producer_offset + consumer_offset * producer_strides.
+    combinedOffsets[i] =
+        getOpFoldResult(builder, loc,
+                        add(mul(consumerOffsets[consumerPos],
+                                producerStrides[i], offsetSymbols),
+                            producerOffsets[i], offsetSymbols),
+                        offsetSymbols);
+    combinedSizes[i] = consumerSizes[consumerPos];
+    // The combined stride is computed as
+    //    consumer_stride * producer_stride.
+    combinedStrides[i] = getOpFoldResult(
+        builder, loc,
+        mul(consumerStrides[consumerPos], producerStrides[i], strideSymbols),
+        strideSymbols);
+    consumerPos++;
+  }
+  return success();
+}
+
+LogicalResult tensor::mergeOffsetsSizesAndStrides(
+    OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
+    OffsetSizeAndStrideOpInterface consumer,
+    const llvm::SmallBitVector &droppedProducerDims,
+    SmallVector<OpFoldResult> &combinedOffsets,
+    SmallVector<OpFoldResult> &combinedSizes,
+    SmallVector<OpFoldResult> &combinedStrides) {
+  SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
+  SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
+  SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
+  SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
+  SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
+  SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
+  return tensor::mergeOffsetsSizesAndStrides(
+      builder, loc, producerOffsets, producerSizes, producerStrides,
+      droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
+      combinedOffsets, combinedSizes, combinedStrides);
 }
 
 namespace {
 /// Merges consecutive tensor.extract_slice ops into one.
 struct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractSliceOp nextOp,
                                 PatternRewriter &rewriter) const override {
     auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>();
     if (!prevOp)
       return failure();
 
-    if (!prevOp.hasUnitStride() || !nextOp.hasUnitStride())
+    SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
+    if (failed(mergeOffsetsSizesAndStrides(rewriter, nextOp.getLoc(), prevOp,
+                                           nextOp, prevOp.getDroppedDims(),
+                                           newOffsets, newSizes, newStrides)))
       return failure();
 
-    auto prevResultType = prevOp.getResult().getType().cast<ShapedType>();
-    if (prevOp.getSourceType().getRank() != prevResultType.getRank())
-      return rewriter.notifyMatchFailure(
-          prevOp, "rank-reducing producder case unimplemented");
-
-    Location loc = nextOp.getLoc();
-
-    SmallVector<OpFoldResult> prevOffsets = prevOp.getMixedOffsets();
-    SmallVector<OpFoldResult> nextOffsets = nextOp.getMixedOffsets();
-    SmallVector<OpFoldResult> foldedOffsets =
-        mergeOffsets(loc, prevOffsets, nextOffsets, rewriter);
-
-    rewriter.replaceOpWithNewOp<ExtractSliceOp>(
-        nextOp, nextOp.getType(), prevOp.getSource(), foldedOffsets,
-        nextOp.getMixedSizes(), nextOp.getMixedStrides());
+    rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(),
+                                                prevOp.getSource(), newOffsets,
+                                                newSizes, newStrides);
     return success();
   }
 };
 
 /// Merges consecutive tensor.insert_slice ops into one.
 struct MergeConsecutiveInsertSlice : public OpRewritePattern<InsertSliceOp> {
   using OpRewritePattern::OpRewritePattern;
 

mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir

 // RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-consecutive-insert-extract-slice -canonicalize -mlir-print-local-scope %s | FileCheck %s
 
 func.func @extract_slice_same_rank(
     %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<8x16x32x?xf32> {
   %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [128, 128, 128, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x128x128x?xf32>
   %1 = tensor.extract_slice %0[7, 8, 9, %offset1] [8, 16, 32, %size1] [1, 1, 1, 1] : tensor<128x128x128x?xf32> to tensor<8x16x32x?xf32>
   return %1: tensor<8x16x32x?xf32>
 }
 
 // CHECK-LABEL: func.func @extract_slice_same_rank
 //  CHECK-SAME: (%[[SOURCE:.+]]: tensor<?x?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %{{.+}}: index, %[[SIZE1:.+]]: index)
-//       CHECK:   %[[OFFSET:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%[[OFFSET0]], %[[OFFSET1]]]
+//       CHECK:   %[[OFFSET:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%[[OFFSET1]], %[[OFFSET0]]]
 //       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[SOURCE]][7, 9, 11, %[[OFFSET]]] [8, 16, 32, %[[SIZE1]]] [1, 1, 1, 1]
 //       CHECK:   return %[[EXTRACT]] : tensor<8x16x32x?xf32>
 
+// -----
+
 func.func @extract_slice_rank_reducing_consumer(
     %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<16x?xf32> {
   %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [128, 128, 128, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x128x128x?xf32>
   %1 = tensor.extract_slice %0[7, 8, 9, %offset1] [1, 16, 1, %size1] [1, 1, 1, 1] : tensor<128x128x128x?xf32> to tensor<16x?xf32>
   return %1: tensor<16x?xf32>
 }
 
 // CHECK-LABEL: func.func @extract_slice_rank_reducing_consumer
 //       CHECK:   tensor.extract_slice %{{.+}}[7, 9, 11, %{{.+}}] [1, 16, 1, %{{.+}}] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<16x?xf32>
 
+// -----
+
 func.func @extract_slice_rank_reducing_producer(
     %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<8x?xf32> {
   %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [1, 128, 1, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x?xf32>
   %1 = tensor.extract_slice %0[7, %offset1] [8, %size1] [1, 1] : tensor<128x?xf32> to tensor<8x?xf32>
   return %1: tensor<8x?xf32>
 }
 
-//   CHECK-LABEL: func.func @extract_slice_rank_reducing_producer
-// CHECK-COUNT-2:   tensor.extract_slice
+// CHECK-LABEL: func.func @extract_slice_rank_reducing_producer
+//  CHECK-SAME: (%[[SRC:.+]]: tensor<?x?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %{{.+}}: index, %[[SIZE1:.+]]: index)
+//       CHECK:   %[[OFFSET:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%[[OFFSET1]], %[[OFFSET0]]]
+//       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[SRC]][0, 8, 2, %[[OFFSET]]] [1, 8, 1, %[[SIZE1]]] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<8x?xf32>
+//       CHECK:   return %[[EXTRACT]] : tensor<8x?xf32>
+
+// -----
+
+func.func @extract_slice_non_one_stride(
+    %src: tensor<?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index, %stride0: index, %stride1: index) -> tensor<?xf32> {
+  %0 = tensor.extract_slice %src[%offset0] [%size0] [%stride0] : tensor<?xf32> to tensor<?xf32>
+  %1 = tensor.extract_slice %0[%offset1] [%size1] [%stride1] : tensor<?xf32> to tensor<?xf32>
+  return %1: tensor<?xf32>
+}
+
+// CHECK-LABEL: func.func @extract_slice_non_one_stride
+//  CHECK-SAME: (%[[SRC:.+]]: tensor<?xf32>, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %{{.+}}: index, %[[SIZE1:.+]]: index, %[[STRIDE0:.+]]: index, %[[STRIDE1:.+]]: index)
+//       CHECK:   %[[OFFSET:.+]] = affine.apply affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>()[%[[OFFSET1]], %[[STRIDE0]], %[[OFFSET0]]]
+//       CHECK:   %[[STRIDE:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%[[STRIDE1]], %[[STRIDE0]]]
+//       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[SRC]][%[[OFFSET]]] [%[[SIZE1]]] [%[[STRIDE]]] : tensor<?xf32> to tensor<?xf32>
+//       CHECK:   return %[[EXTRACT]] : tensor<?xf32>
 
 // -----
 
 func.func @insert_slice_rank_reducing(
     %dst: tensor<128x128x128x128xf32>, %mid: tensor<1x16x1xf32>, %src: tensor<16xf32>, %offset: index) -> tensor<128x128x128x128xf32> {
   %0 = tensor.insert_slice %src into %mid[0, 0, 0] [1, 16, 1] [1, 1, 1] : tensor<16xf32> into tensor<1x16x1xf32>
   %1 = tensor.insert_slice %0 into %dst[6, 7, 8, %offset] [1, 1, 16, 1] [1, 1, 1, 1] : tensor<1x16x1xf32> into tensor<128x128x128x128xf32>
   return %1: tensor<128x128x128x128xf32>
 }
 
 // CHECK-LABEL: func.func @insert_slice_rank_reducing
 //  CHECK-SAME: (%[[DST:.+]]: tensor<128x128x128x128xf32>, %{{.+}}: tensor<1x16x1xf32>, %[[SRC:.+]]: tensor<16xf32>, %[[IDX:.+]]: index)
 //       CHECK:  %[[INSERT:.+]] = tensor.insert_slice %[[SRC]] into %[[DST]][6, 7, 8, %[[IDX]]] [1, 1, 16, 1] [1, 1, 1, 1]
 //       CHECK:  return %[[INSERT]]
 
+// -----
+
 func.func @insert_slice_rank_reducing_dynamic_shape(
     %dst: tensor<128x128x128x128xf32>, %mid: tensor<1x?x1xf32>, %src: tensor<?xf32>, %offset: index, %size: index) -> tensor<128x128x128x128xf32> {
   %0 = tensor.insert_slice %src into %mid[0, 0, 0] [1, %size, 1] [1, 1, 1] : tensor<?xf32> into tensor<1x?x1xf32>
   %1 = tensor.insert_slice %0 into %dst[6, 7, 8, %offset] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<1x?x1xf32> into tensor<128x128x128x128xf32>
   return %1: tensor<128x128x128x128xf32>
 }
 
 //   CHECK-LABEL: func.func @insert_slice_rank_reducing_dynamic_shape
 // CHECK-COUNT-2:   tensor.insert_slice