[mlir][Tensor] Add rewrites to extract slices through `tensor.collape_shape`

[christopherbate]

Summary

This change adds a set of utilities to replace the result of a
tensor.collapse_shape -> tensor.extract_slice chain with the
equivalent result formed by aggregating slices of the
tensor.collapse_shape source. In general, it is not possible to
commute extract_slice and collapse_shape if linearized dimensions
are sliced. The i-th dimension of the tensor.collapse_shape
result is a "linearized sliced dimension" if:

1. Reassociation indices of tensor.collapse_shape in the i'th position is greater than size 1 (multiple dimensions of the input are collapsed)
2. The i-th dimension is sliced by tensor.extract_slice.

We can work around this by stitching together the result of
tensor.extract_slice by iterating over any linearized sliced dimensions.
This is equivalent to "tiling" the linearized-and-sliced dimensions of
the tensor.collapse_shape operation in order to manifest the result
tile (the result of the tensor.extract_slice). The user of the
utilities must provide the mechanism to create the tiling (e.g. a loop).
In the tests, it is demonstrated how to apply the utilities using either
scf.for or scf.foreach_thread.

The below example illustrates the pattern using scf.for:

%0 = linalg.generic ... -> tensor<3x7x11x10xf32>
%1 = tensor.collapse_shape %0 [[0, 1, 2], [3]] : ... to tensor<341x10xf32>
%2 = tensor.extract_slice %1 [13, 0] [10, 10] [2, 1] : .... tensor<10x10xf32>

We can construct %2 by generating the following IR:

%dest = linalg.init_tensor() : tensor<10x10xf32>
%2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0) -> tensor<10x10xf32> {
   // Step 1: Map this output idx (%iv) to a multi-index for the input (%3):
   %linear_index = affine.apply affine_map<(d0)[]->(d0*2 + 11)>(%iv)
   %3:3 = arith.delinearize_index %iv into (3, 7, 11)
   // Step 2: Extract the slice from the input
   %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] :
         tensor<3x7x11x10xf32> to tensor<1x1x1x10xf32>
   %5 = tensor.collapse_shape %4 [[0, 1, 2], [3]] :
         tensor<1x1x1x10xf32> into tensor<1x10xf32>
   // Step 3: Insert the slice into the destination
   %6 = tensor.insert_slice %5 into %arg0 [%iv, 0] [1, 10] [1, 1] :
         tensor<1x10xf32> into tensor<10x10xf32>
   scf.yield %6 : tensor<10x10xf32>
}

The pattern was discussed in the RFC here: https://discourse.llvm.org/t/rfc-tensor-extracting-slices-from-tensor-collapse-shape/64034

[mravishankar]

High level comment, reshape-like operations should not really be part of tiling interface... They are really metadata operations and not compute operations. So I think a bit more context on the use case of this would be useful. Adding a tiling interface implementation for what should be just metadata operations leads to some unintended consequences (from experience having done exactly this for tensor.extract_slice and tensor.insert_slice during incubation in IREE).

[christopherbate]

High level comment, reshape-like operations should not really be part of tiling interface... They are really metadata operations and not compute operations. So I think a bit more context on the use case of this would be useful.

Probably the biggest reason to allow this sort of transformation as part of the normal tile-and-fuse flow is that it allows you to reduce multiple dimensions as a single dimension without blocking pulling producers into the loop nest.

I wrote this diff in two ways -- one being just adding it to tiling interface, the other adding a standalone rewrite to do the generateResultTileValue and adding a test pass, without hooking into TilingInterface. The first way ends up being less code and doesn't require extra test pass to demonstrate the end-to-end result.

I'll push up the other method if you think that's better.

[mravishankar]

Probably the biggest reason to allow this sort of transformation as part of the normal tile-and-fuse flow is that it allows you to reduce multiple dimensions as a single dimension without blocking pulling producers into the loop nest.

Totally agreed on this point. There was co-incidentally a similar discussion about this on IREE discord. Essentially you want to allow tile and fuse to work through reshapes. Currently the way tile + fuse is implemented is you use the tensor.extract_slice to find the tile of the producer that you need to compute in place. But that requires the producer to be an operation that implements the TilingInterface. If you tensor.extract_slice comes from a result of a reshape, then the fusion stops. This is probably what you are referring to as well. Instead, you can fold the reshape and tensor.extract_slice to get a tensor.extract_slice of the source of the reshape. Then the fusion proceed as normal. I think this folding pattern is the missing component. If we have that we can then figure out how to add it to the mix and get tile and fuse to work through reshapes....

[christopherbate]

Use a dedicated rewrite pattern to replace the tensor.extract_slice rather than using an implementation of TilingInterface for tensor.collapse_shape.

[christopherbate]

Totally agreed on this point. There was co-incidentally a similar discussion about this on IREE discord. Essentially you want to allow tile and fuse to work through reshapes. Currently the way tile + fuse is implemented is you use the tensor.extract_slice to find the tile of the producer that you need to compute in place. But that requires the producer to be an operation that implements the TilingInterface. If you tensor.extract_slice comes from a result of a reshape, then the fusion stops. This is probably what you are referring to as well. Instead, you can fold the reshape and tensor.extract_slice to get a tensor.extract_slice of the source of the reshape. Then the fusion proceed as normal. I think this folding pattern is the missing component.

I briefly checked out your discussion, and it appears that we are after the same thing.

If we have that we can then figure out how to add it to the mix and get tile and fuse to work through reshapes....

I updated the diff so that the pattern is implemented as a dedicated rewrite with a public method to directly produce the equivalent result of a tensor.collapse_shape -> tensor.extract_slice chain.

Maybe we need a SliceableInterface in addition to TilingInterface. Then it would be easier to hook into the tile and fuse flow if the idea is that only computation operations implement TilingInterface.

We would also need a control mechanism for tile and fuse to allow the caller to stop the greedy fusion early if desired. Were you already planning on adding a mechanism for that?

[mravishankar]

I updated the diff so that the pattern is implemented as a dedicated rewrite with a public method to directly produce the equivalent result of a tensor.collapse_shape -> tensor.extract_slice chain.Maybe we need a SliceableInterface in addition to TilingInterface. Then it would be easier to hook into the tile and fuse flow if the idea is that only computation operations implement TilingInterface. We would also need a control mechanism for tile and fuse to allow the caller to stop the greedy fusion early if desired. Were you already planning on adding a mechanism for that?

Thanks! I'll take a look. There is a lot to unpack here and I need to understand how this works a bit more. Maybe @nicolasvasilache might already have some ideas here.
With respect to SliceableInterface, its along the lines what Nicolas had suggested a while back to, but more in terms of a "SubSetInterface" where you can define how to extract/insert subset of data (not necessarily rectangular or contiguous). Not sure if anything has already been done on that.
For control, I think the greedy pattern is more for "demonstration". Any control mechanism will probably use the guts of things and be wrapped within the transform dialect. Its really hard IMO to provide a universal control mechanism, but one thing that can work is having a callback that allows callers to decide if a slice has to be used for fusion or not. I didnt add that cause I wasnt sure how it would interact with the transform dialect approach, but I have used such callbacks in a different context, and they did the job.

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

27

I think the arith dialect has some of this code as utilities (https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h)

[christopherbate]

Add support for using scf.foreach_thread instead of an scf.for loop nest.
Refactor the implementation to deal with added complexity of multiple loop types.

[nicolasvasilache]

Very cool stuff, thanks @christopherbate this is something that has been missing for a long time!

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

31

Can you move the helpers to include/mlir/Dialect/Utils/ReshapeOpsUtils.h for better visibility?

68

Probably worth it for this one to inspect the tensor dim and immediately produce IndexAttr(size).
Unfortunately, createOrFold always returns a Value atm.

139

nit: Plz reflow comment

140

nit: space after for

141

I think we can move to SmallVector<OpFoldResult> lbs, ubs, steps; thanks to recent additions that allow creating more folded expressions than before.

150

You now have makeComposedFoldedAffineApply that takes OpFoldResult and does the right thing under the hood.

195

Why the difference with the above ?
I personally prefer the mixed s0 + d0 * s1 form everywhere if possible.

225

I would also move this to the reshape utils for visibility, I anticipate this will become a more important use case over time.

238

nit: llvm::append_range(extractOffsets, llvm::map_range()) may be a bit nicer.

250

seemslike this is the impl you want for your getShapeDimSizes above and then you can just reuse here ?

278

Must the strides always be all 1s ?
Or should you pass them as an extra argument too?

287

I'll comment on your other PR, but I would either:

• move this op to arith and keep the operands as they are, or
• keep it in tensor but then take an actual tensor as the basis and abstract away the fact that you get the Dims as a lowering detail.

Unless you have a stronger argument for keeping the semantics as is that I don't see yet?

351

nit: values.

[mravishankar]

Left a few comments. Removing my blocker on this. Thanks!

[mravishankar]

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

38

It might be better to make the return type Optional<llvm::SmallBitVector> or FailureOr<llvm::SmallBitVector> and return errors instead of (void)-ing them this way.
Also not clear to me what is the expected behavior if the shape are dynamic. I am expecting that it would return failure that needs to be propagated back up the stack.

110

Just out of curiosity, is there a solution space where we dont use scf.for or scf.for_each here? Its unfortunate coupling of these loop constructs with the transformation.

344

I dont think you need to cast it to the interface op. The method is implemented on tensor.extract_slice itself. You should be able to call it directly.

[christopherbate]

Address comments

[christopherbate]

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

110

So we need some sort of intermediate construct that handles the discontinuous iteration space. I just chose these loops because they seem commonly used.

I'm not sure whether you're asking

1. whether the implementation can be improved to be more generic to the loop op kind or
2. whether there's a better abstraction that can be used instead of a loop.

For 1, maybe. But it's probably the same amount of effort to add a method called buildWithMyLoopOp to the implementation class.

For 2, yes, I think there is more opportunity to fill missing abstractions here. If the linear dim size is 128 and you don't insert a loop, then instead you could directly insert an "unrolled" form of 128 iterations of what you originally put in the body. Of course, in this form it is too unwieldy. One option would be to add an op equivalent to a higher-order tensor.extract_slice and/or vector.transfer_read/write: it maps some iterations space to a set of slices. That may be easier to further transform versus the loop form.

[nicolasvasilache]

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

110

Yes, we will need new types and ops to handle the discontinuity and that intercepts gather/scatter/parallel_scatter and related abstractions.
This is much longer-term though.

it maps some iterations space to a set of slices -> yes the piecewise discontinuities are important to handle right, I am unclear yet what a good abstraction for this looks like but it will be a major focus or ours for the rest of the year.

[christopherbate]

Rebase

[nicolasvasilache]

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

38

+1 to FailureOr anywhere we need to propagate errors.

[christopherbate]

Refactor to remove dependence on scf.for/scf.foreach_thread and caller provides iteration mechanism. Fix a couple bugs and improve formatting of tests.

[nicolasvasilache]

Thanks!

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

39

Could we please add some spacing and structuring to the comments?

... by the caller.

The class provides two methods:
  1. create: short description
  2. emitExtractSliceFromCollapseShapeBody: short description

Intended use case:
==============
The caller should first ... 

Then the caller should ...

Example:
=======
zzz

77

This looks more like it belongs in a TransformUtils.h

77

s/Builder/Helper

I find this class is a helper that exposes methods. Such methods take a builder to do X.
So the class itself is not a Builder, and we since that name has heavy implications in MLIR, we should avoid reusing it IMO.

114

More structure to the doc here too plz:

... by:
  1. 
  2. 
  3.

The returned pair ...

114

As in the other PR I commented on, I think inverting may be misleading, we are really propagating through the CollapseShapeOp.

In a followup, I imagine we could refactor and reuse some of the recent improvements introduced in: https://reviews.llvm.org/D128986 ?

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

392 ↗

(On Diff #454114)

this class provides ?

394 ↗

(On Diff #454114)

The name ExtractShapeExtractSlice is hard to relate to.
Also, this not a Builder.

Something like CollapseLinearizationHelper ?

409 ↗

(On Diff #454114)

Better structuring of the docc here would also help readability (e.g. isolating pieces of equations better, having paragraphs)

409 ↗

(On Diff #454114)

invert -> fold plz

411 ↗

(On Diff #454114)

Example in IR would help to relate the proper pieces, the description is a bit too abstract to follow.

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

95

With C++17 I think it becomes nicer to just:

SmallVector<Range> ranges;
ranges.reserve(...);
for (const auto & [o, s, t] : llvm::zip(...))
  ranges.push_back(Range{o, s, t});

151

see, it's a helper :)

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

140

Fine for a test but if it want to be hoisted, this should prob be 1 pattern per loop op type.
Maybe add a TODO?

[christopherbate]

Address comments

[christopherbate]

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

77

Moved to new file TransformUtils.h and renamed ExtractSliceFromCollapseShapeHelper.

114

As in the other PR I commented on, I think inverting may be misleading, we are really propagating through the CollapseShapeOp.

I don't follow how inverting is misleading.

In x -> CollapseShape -> y -> ExtractSlice -> z , each of x,y,z are different Tensors with different spaces defined by their shapes. The CollapseShape and Extractslice define maps that map one coordinate to another. The effect of this function is to take coordinates in z, which are given by the tile coordinates / loop induction variables, and map them back to coordinates in x. I tried to make this point of view explicit in the implementation by using functions invertSliceIndexing and invertCollapseShapeIndexing.

In a followup, I imagine we could refactor and reuse some of the recent improvements introduced in: https://reviews.llvm.org/D128986 ?

There are some common aspects that could be factored out to unify these index space manipulation ideas. For example, extracting an element from a collapse-shape'd memref could be phrased as inverting the index transformation of the collapse, which is exactly what thy are doing with the lienarize/delinearize. They could also use the new AffineDelinearizeOp in order to simplify some of the code.

The two implementations would be the same except for one aspect: all the complexity in this use case comes from the fact that we have to carry around information about which dimensions are sliced and/or linearized in order to make the generated code efficient. In the linked diff, they don't have to worry about the effect of the slicing aspect, as that code is only dealing with loading individual elements. Even in the body of the loop generated by this function, we are not guaranteed to just load scalars. Rather, you need to also account for the potential of a) slicing a non-linearized dimension (in which case you don't need a loop dimension) or b) taking the entirety of a linearized dimension. You need this information at basically every point in this transformation. If we could discard the slicing aspect, then it would be greatly simplified into the special case of just loading scalar elements.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

394 ↗

(On Diff #454114)

Oops, that was supposed to be CollapseShapeExtractSlice. fixed.

409 ↗

(On Diff #454114)

I changed it, but please see my other comment on the rationale for using "invert" here.

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

95

agreed, fixed

[christopherbate]

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

140

I broke it into two patterns.

[christopherbate]

Rebase

[nicolasvasilache]

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

114

I don't follow how inverting is misleading.

Ok, then I misunderstood this from the doc.
Could we improve the doc with better structuring and some IR example that makes the inversion pop up very clearly?

[christopherbate]

Expand the documentation for the main class that is added. I checked the result of the Doxygen build to make sure it looks correct with relevant hyperlinks and headings correctly generated.

Shortened some function/class names.

[christopherbate]

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

114

Ok, I've expanded the IR example to be more complicated and added a step-by-step commentary that explains what is happening, including why I'm saying invert/inversion here.
Let me know if you still want me to adjust the wording.

I've checked the output of this comment block in the Doxygen build as well to ensure all the formatting shows up correctly.

[christopherbate]

Fix typos.

[nicolasvasilache]

Looks great, thanks for all the doc !

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

42 ↗

(On Diff #457155)

nit: "the the"

[christopherbate]

Fix typo in doc

[christopherbate]

Rebase

[christopherbate]

Windows bot seems to be failing for unrelated reason

[mehdi_amini]

There is an issue of circular dependency here: you're introducing a dependency from Dialect/Utils/ to the ViewLikeInterface, but it already depends on Dialect/Utils

[mehdi_amini]

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

17

On top of the cyclic dependency mentioned elsewhere, this introduces a dependency on Linalg which isn't defined in CMake as far as I can tell: this is breaking the bots.

Also I am not sure that it is conceptually correct to have lib/Dialect/Tensor to depend on Linalg?

I'll revert for now.

[christopherbate]

In D129699#3768090, @mehdi_amini wrote:

There is an issue of circular dependency here: you're introducing a dependency from Dialect/Utils/ to the ViewLikeInterface, but it already depends on Dialect/Utils

Thanks, I should have checked for that.

[christopherbate]

Depends on https://reviews.llvm.org/D133523 to resolve circular dependency with ViewLikeInterface

[christopherbate]

Rebase on D133523, resolve linalg dependency issue.

Verified Bazel build is working and added required Bazel build file changes

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

 /// transform dialect that control is done within the transform dialect. Other
 /// use cases can inherit from this pattern and add necessary controls.
 FailureOr<Value> replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
 
 /// Populates rewrite patterns that lower `tensor.delinearize_index`.
 void populateLowerDelinearizeIndexPatterns(RewritePatternSet &patterns);
 
+/// Generates IR required to materialize a slice "through" a
+/// `tensor.collapse_shape` by creating a loop nest and populating the slice in
+/// `dest` by stitching together different slices of the source tensor. The loop
+/// nest contains one loop for each sliced output dimension that maps to
+/// multiple source dimensions.
+/// Example:
+// clang-format off
+/// ```
+/// %0 = ... -> tensor<3x7x11x10xf32>
+/// %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]
+/// %2 = tensor.extract_slice %1 [0, 0] [10, 10] [1, 1]
+/// ```
+/// The "materialized" slice is equivalent to
+/// ```
+/// %2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0 = ...
+///    %3:3 = tensor.delinearize_index %iv (3, 7, 11)
+///    %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] 
+///    %5 = tensor.collapse_shape %4 [[0, 1, 2], [3]]
+///    %6 = tensor.insert_slice %5 into %arg0 [%iv, 0] [1, 10] [1, 1]
+/// }
+/// ```
+// clang-format on
+/// This function directly creates the materialized slice from offsets and
+/// sizes. If `useForeach` is true, then an `scf.foreach_thread` operation will
+/// be used instead of an `scf.for` loop nest.
+FailureOr<Value> materializeSliceFromCollapseShape(
+    OpBuilder &builder, tensor::CollapseShapeOp op, Value dest,
+    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+    ArrayRef<OpFoldResult> strides, bool useForeach);
+
+/// A callback type that returns a destination tensor given the size and element
+/// type.
+using CreateDestTensorFn =
+    std::function<Value(OpBuilder &, Location, ArrayRef<OpFoldResult>, Type)>;
+
+/// Pattern to swap a `tensor.extract_slice` with its producer when the
+/// producer is a `tensor.collapse_shape`. The `tensor.extract_slice` is
+/// replaced by a loop nest that stitches together the desired tile by
+/// iterating over the linearized slice dimensions that cannot be represented
+/// as a rectangular slice of the source tensor (see above example).
+struct SliceCollapseShapeOptions {
+
+  bool useForeach = false;
+  CreateDestTensorFn createDestTensorFn = nullptr;
+
+  SliceCollapseShapeOptions &setUseForeach(bool useScfForeach) {
+    useForeach = useScfForeach;
+    return *this;
+  }
+
+  SliceCollapseShapeOptions &setCreateDestTensorFunc(CreateDestTensorFn fn) {
+    createDestTensorFn = std::move(fn);
+    return *this;
+  }
+};
+struct RewriteExtractSliceWithTiledCollapseShape
+    : public OpRewritePattern<tensor::ExtractSliceOp> {
+  RewriteExtractSliceWithTiledCollapseShape(MLIRContext *context,
+                                            SliceCollapseShapeOptions options)
+      : OpRewritePattern<tensor::ExtractSliceOp>(context),
+        options(std::move(options)) {}
+
+  LogicalResult matchAndRewrite(tensor::ExtractSliceOp op,
+                                PatternRewriter &rewriter) const override;
+
+  SliceCollapseShapeOptions options;
+};
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   DelinearizeIndex.cpp
+  ExtractSliceFromReshape.cpp
   SplitPadding.cpp
   SwapExtractSliceWithProducer.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen

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

+//===- ExtractSliceFromReshape.cpp - Slice reshape rewrites-------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements rewrites that replace slices of reshape results with
+// aggregated slices of the reshape source.
+//
+//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/OpDefinition.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+/// The input parameters `offsets`, `sizes`, `strides` specify a rectangular
+/// slice of the collapse_shape output. Try to find which dimensions have been
+/// sliced and which dimensions are not sliced (offset = 0, size = dim, size =
+/// 1). Note that this conservative as it cannot detect if a dynamic size
+/// corresponds to the full tensor dimension or not.
+static llvm::SmallBitVector getSlicedDimensions(
+    OpBuilder &b, CollapseShapeOp op, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) {
+  // Materialize the output shape values.
+  ReifiedRankedShapedTypeDims reifiedShapes;
+  ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
+      dyn_cast<ReifyRankedShapedTypeOpInterface>(op.getOperation());
+  (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
+  llvm::SmallBitVector result(op.getResultType().getRank());
+  for (const auto &it : llvm::enumerate(llvm::zip(offsets, sizes, strides))) {
+    auto size = std::get<1>(it.value());
+    Optional<int64_t> stride = getConstantIntValue(std::get<2>(it.value()));
+    Optional<int64_t> offset = getConstantIntValue(std::get<0>(it.value()));
+    result[it.index()] =
+        !isEqualConstantIntOrValue(size, reifiedShapes[0][it.index()]) ||
+        (!stride || *stride != 1) || (!offset || *offset != 0);
+  }
+  return result;
+}
+
+/// Determine which dimensions are linearized by a collapse shape op.
+static llvm::SmallBitVector getLinearizedDimensions(CollapseShapeOp op) {
+  llvm::SmallBitVector result(op.getResultType().getRank());
+  for (const auto &it : llvm::enumerate(op.getReassociationIndices())) {
+    result[it.index()] = it.value().size() > 1;
+  }
+  return result;
+}
+
+/// Get the dimension sizes of a value of RankedTensor type. Only dimensions at
+/// indices specified in `dimensions` are returned.
+static SmallVector<OpFoldResult>
+getShapeDimSizes(OpBuilder &b, Location loc, Value rankedTensor,
+                 ArrayRef<int64_t> dimensions) {
+  SmallVector<OpFoldResult> basis;
+  for (auto idx : dimensions) {
+    Value idxValue = b.create<arith::ConstantIndexOp>(loc, idx);
+    basis.push_back(b.createOrFold<tensor::DimOp>(loc, rankedTensor, idxValue));
+  }
+  return basis;
+}
+
+namespace {
+struct CollapseShapeSliceInternal {
+protected:
+  CollapseShapeOp collapseOp;
+  const SmallVector<ReassociationIndices> reassociationIndices;
+
+  const SmallVector<OpFoldResult> sliceOffsets;
+  const SmallVector<OpFoldResult> sliceSizes;
+  const SmallVector<OpFoldResult> sliceStrides;
+
+  const llvm::SmallBitVector linearizedOutputDims;
+  const llvm::SmallBitVector slicedOutputDims;
+  SmallVector<int64_t> tiledDimensions;
+
+public:
+  /// Construct the class using a `tensor.collapse_shape` and slice offsets,
+  /// sizes, and strides.
+  CollapseShapeSliceInternal(OpBuilder &b, CollapseShapeOp op,
+                             ArrayRef<OpFoldResult> offsets,
+                             ArrayRef<OpFoldResult> sizes,
+                             ArrayRef<OpFoldResult> strides)
+      : collapseOp(op), reassociationIndices(op.getReassociationIndices()),
+        sliceOffsets(offsets.begin(), offsets.end()),
+        sliceSizes(sizes.begin(), sizes.end()),
+        sliceStrides(strides.begin(), strides.end()),
+        linearizedOutputDims(getLinearizedDimensions(collapseOp)),
+        slicedOutputDims(getSlicedDimensions(b, op, sliceOffsets, sliceSizes,
+                                             sliceStrides)) {
+
+    for (int64_t i = 0; i < op.getResultType().getRank(); i++) {
+      if (linearizedOutputDims[i] && slicedOutputDims[i])
+        tiledDimensions.push_back(i);
+    }
+  }
+
+  /// Materialize a slice from a `tensor.collapse_shape` by using `scf.for` to
+  /// sub-tile linearized dimensions.
+  FailureOr<Value> buildWithScfFor(OpBuilder &b, Location loc, Value dest);
+
+  /// Materialize a slice from a `tensor.collapse_shape` by using
+  /// `scf.foreach_thread` to sub-tile linearized dimensions.
+  FailureOr<Value> buildWithScfForeach(OpBuilder &b, Location loc, Value dest);
+
+protected:
+  /// Fill in the body of the generated loop nest. `ivs` contains one Value for
+  /// each tiled output dimension that corresponds to a set of linearized and
+  /// sliced input dimensions. The `ivs` values should be treated as the
+  /// position within the output dimension to delinearize. `insertOffsets` and
+  /// `insertSizes` are filled with the output positions where the sub-slice
+  /// should be inserted into the destination tensor.
+  Value buildBody(OpBuilder &b, Location loc, ValueRange ivs, Value insertDest,
+                  SmallVector<OpFoldResult> &insertOffsets,
+                  SmallVector<OpFoldResult> &insertSizes);
+
+  /// Use the given `multiIndices` for the collapsed and tiled dims to
+  /// extract a subset (sub-tile) of the slice from the source of the
+  /// `tensor.collapse_op`. The subset extracted is equivalent to a tile from
+  /// the source where linearized and sliced dimensions have been tiled by 1.
+  Value extractSubTileFromSource(OpBuilder &b, Location loc,
+                                 ArrayRef<ValueRange> multiIndices);
+};
+} // namespace
+
+FailureOr<Value> CollapseShapeSliceInternal::buildWithScfFor(OpBuilder &b,
+                                                             Location loc,
+                                                             Value dest) {
+  // Create the bounds for the loop nest to be akin
+  // to "for(i = offset; i < size * stride; i += stride)"
+  SmallVector<Value> lbs;
+  SmallVector<Value> ubs;
+  SmallVector<Value> steps;
+  AffineExpr s0, s1, s2;
+  bindSymbols(b.getContext(), s0, s1, s2);
+  for (int64_t idx : tiledDimensions) {
+    Value strideVal =
+        getValueOrCreateConstantIndexOp(b, loc, sliceStrides[idx]);
+    lbs.push_back(getValueOrCreateConstantIndexOp(b, loc, sliceOffsets[idx]));
+    ubs.push_back(makeComposedAffineApply(
+        b, loc, s0 + s1 * s2,
+        {lbs.back(), getValueOrCreateConstantIndexOp(b, loc, sliceSizes[idx]),
+         strideVal}));
+    steps.push_back(strideVal);
+  }
+  scf::LoopNest nest = scf::buildLoopNest(
+      b, loc, lbs, ubs, steps, dest,
+      [this](OpBuilder &nestedBuilder, Location loc, ValueRange ivs,
+             ValueRange iterArgs) -> scf::ValueVector {
+        // Fill in the body.
+        SmallVector<OpFoldResult> insertOffsets, insertSizes;
+        Value slice = buildBody(nestedBuilder, loc, ivs, iterArgs[0],
+                                insertOffsets, insertSizes);
+
+        // Insert the sub-tile into the iteration arg.
+        Value result = nestedBuilder.create<tensor::InsertSliceOp>(
+            loc, slice, iterArgs[0], insertOffsets, insertSizes,
+            /*strides=*/
+            SmallVector<OpFoldResult>(sliceOffsets.size(),
+                                      nestedBuilder.getIndexAttr(1)));
+        return {result};
+      });
+  return nest.getResults()[0];
+}
+
+FailureOr<Value> CollapseShapeSliceInternal::buildWithScfForeach(OpBuilder &b,
+                                                                 Location loc,
+                                                                 Value dest) {
+  // For reach sliced and linearized dimension, set the number of threads to be
+  // the size of the slice in that dim.
+  SmallVector<OpFoldResult> sizes;
+  for (int64_t idx : tiledDimensions)
+    sizes.push_back(sliceSizes[idx]);
+
+  auto foreachOp = b.create<scf::ForeachThreadOp>(
+      loc, /*numThreads=*/getValueOrCreateConstantIndexOp(b, loc, sizes),
+      /*threadDimMapping=*/ArrayRef<int64_t>{},
+      [this, dest](OpBuilder &nestedBuilder, Location loc,
+                   ValueRange threadIds) {
+        // Within the body of the foreach op, create the adjusted induction
+        // variables to map the thread index iv to a position within the slice
+        // of the linearized dimension.
+        AffineExpr d0, s0, s1;
+        bindDims(nestedBuilder.getContext(), d0);
+        bindSymbols(nestedBuilder.getContext(), s0, s1);
+        SmallVector<Value> ivs;
+        for (const auto &it : llvm::enumerate(tiledDimensions)) {
+          ivs.push_back(makeComposedAffineApply(
+              nestedBuilder, loc, s0 + d0 * s1,
+              {threadIds[it.index()],
+               getValueOrCreateConstantIndexOp(nestedBuilder, loc,
+                                               sliceOffsets[it.value()]),
+               getValueOrCreateConstantIndexOp(nestedBuilder, loc,
+                                               sliceStrides[it.value()])}));
+        }
+
+        // Populate the rest of the body.
+        SmallVector<OpFoldResult> insertOffsets, insertSizes;
+        Value slice = buildBody(nestedBuilder, loc, ivs, dest, insertOffsets,
+                                insertSizes);
+
+        // Create the terminator and insert.
+        auto term = nestedBuilder.create<scf::PerformConcurrentlyOp>(loc);
+        nestedBuilder.setInsertionPointToStart(term.getBody());
+        nestedBuilder.create<tensor::ParallelInsertSliceOp>(
+            loc, slice, dest, insertOffsets, insertSizes,
+            /*strides=*/
+            SmallVector<OpFoldResult>(insertOffsets.size(),
+                                      nestedBuilder.getIndexAttr(1)));
+      });
+
+  return foreachOp->getResult(0);
+}
+
+Value CollapseShapeSliceInternal::extractSubTileFromSource(
+    OpBuilder &b, Location loc, ArrayRef<ValueRange> multiIndices) {
+  int64_t loopIdx = 0;
+  RankedTensorType srcType = collapseOp.getSrcType();
+  SmallVector<OpFoldResult> extractOffsets;
+  SmallVector<OpFoldResult> extractSizes;
+  SmallVector<OpFoldResult> extractStrides;
+  for (const auto &it : llvm::enumerate(reassociationIndices)) {
+    // Case 1: De-linearized dimensions that have also been sliced. These a
+    // size of 1 because we are iterating over these dimensions. The offset
+    // is exactly the de-linearized multi index created from the iv's.
+    if (slicedOutputDims[it.index()] && linearizedOutputDims[it.index()]) {
+      extractSizes.append(it.value().size(), b.getIndexAttr(1));
+      extractOffsets.append(llvm::to_vector(llvm::map_range(
+          multiIndices[loopIdx++],
+          [&](Value v) -> OpFoldResult { return getAsOpFoldResult(v); })));
+      extractStrides.append(it.value().size(), b.getIndexAttr(1));
+      continue;
+    }
+
+    // Case 2: One or possibly multiple combined input dimensions, but we
+    // have proven that these are not sliced. In this case we just take the
+    // full extent of each index in the list.
+    if (linearizedOutputDims[it.index()]) {
+      for (int64_t srcIndex : it.value()) {
+        if (!srcType.isDynamicDim(srcIndex))
+          extractSizes.push_back(b.getIndexAttr(srcType.getDimSize(srcIndex)));
+        else
+          extractSizes.push_back(b.createOrFold<tensor::DimOp>(
+              loc, collapseOp.src(),
+              b.create<arith::ConstantIndexOp>(loc, srcIndex)));
+        extractOffsets.push_back(b.getIndexAttr(0));
+        extractStrides.push_back(b.getIndexAttr(1));
+      }
+      continue;
+    }
+
+    // Case 3: A single index, but it may be sliced.
+    extractSizes.push_back(sliceSizes[it.index()]);
+    extractOffsets.push_back(sliceOffsets[it.index()]);
+    extractStrides.push_back(sliceStrides[it.index()]);
+  }
+  return b.create<tensor::ExtractSliceOp>(loc, collapseOp.src(), extractOffsets,
+                                          extractSizes, extractStrides);
+}
+
+/// Fill in the body of the generated loop nest. `ivs` contains one Value for
+/// each tiled output dimension that corresponds to a set of linearized and
+/// sliced input dimensions. The `ivs` value then should be treated as the
+/// position within the output tensor dimension to gather along.
+Value CollapseShapeSliceInternal::buildBody(
+    OpBuilder &b, Location loc, ValueRange ivs, Value insertDest,
+    SmallVector<OpFoldResult> &insertOffsets,
+    SmallVector<OpFoldResult> &insertSizes) {
+  // Create the de-linearized multi indices at the start of
+  // each loop body.
+  SmallVector<ValueRange> multiIndices;
+  for (const auto &it : llvm::enumerate(tiledDimensions)) {
+    SmallVector<OpFoldResult> basis =
+        getShapeDimSizes(b, collapseOp->getLoc(), collapseOp.src(),
+                         reassociationIndices[it.value()]);
+    auto delinearizeOp =
+        b.create<tensor::DelinearizeIndexOp>(loc, ivs[it.index()], basis);
+    multiIndices.push_back(delinearizeOp->getResults());
+  }
+
+  // Extract a sub-slice from the source using multi-indices.
+  Value subTileResult = extractSubTileFromSource(b, loc, multiIndices);
+
+  // Collapse the dimensions of the source slice back down.
+  Value collapsedResult = b.create<tensor::CollapseShapeOp>(
+      loc, subTileResult, reassociationIndices);
+
+  // Calculate the position where the sub-slice should be inserted into the
+  // destination tensor.
+  int64_t loopIdx = 0;
+  for (int64_t outDimIdx = 0, e = collapseOp.getResultType().getRank();
+       outDimIdx < e; outDimIdx++) {
+    // Case 1: Linearized dimensions that have been sliced.
+    // The insert size is 1, and the offset is the iv.
+    if (linearizedOutputDims[outDimIdx] && slicedOutputDims[outDimIdx]) {
+      insertOffsets.push_back(ivs[loopIdx++]);
+      insertSizes.push_back(b.getIndexAttr(1));
+      continue;
+    }
+    // Case 2: Otherwise, the insert is the full shape of
+    // the iteration argument dimension, because this output
+    // dimension is not being iterated over in the loop
+    // nest.
+    insertOffsets.push_back(b.getIndexAttr(0));
+    RankedTensorType iterArgsType =
+        insertDest.getType().cast<RankedTensorType>();
+    if (iterArgsType.isDynamicDim(outDimIdx))
+      insertSizes.push_back(b.createOrFold<tensor::DimOp>(
+          loc, insertDest,
+          b.createOrFold<arith::ConstantIndexOp>(loc, outDimIdx)));
+    else
+      insertSizes.push_back(b.getIndexAttr(iterArgsType.getDimSize(outDimIdx)));
+  }
+  return collapsedResult;
+}
+
+FailureOr<Value> mlir::tensor::materializeSliceFromCollapseShape(
+    OpBuilder &b, tensor::CollapseShapeOp op, Value dest,
+    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+    ArrayRef<OpFoldResult> strides, bool useForeach) {
+  CollapseShapeSliceInternal sliceBuilder(b, op, offsets, sizes, strides);
+  return useForeach ? sliceBuilder.buildWithScfForeach(b, op->getLoc(), dest)
+                    : sliceBuilder.buildWithScfFor(b, op.getLoc(), dest);
+}
+
+LogicalResult RewriteExtractSliceWithTiledCollapseShape::matchAndRewrite(
+    tensor::ExtractSliceOp op, PatternRewriter &rewriter) const {
+  auto collapseOp = op.getSource().getDefiningOp<tensor::CollapseShapeOp>();
+  if (!collapseOp)
+    return rewriter.notifyMatchFailure(
+        op, "producer is not a tensor.collapse_shape op");
+
+  // Materialize the output shape values.
+  ReifiedRankedShapedTypeDims reifiedShapes;
+  ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
+      cast<ReifyRankedShapedTypeOpInterface>(op.getOperation());
+  if (failed(
+          reifyShapedTypeInterface.reifyResultShapes(rewriter, reifiedShapes)))
+    return rewriter.notifyMatchFailure(op, "failed to reify result shapes");
+
+  // Create the destination tensor using the above vales.
+  Value dest;
+  Type elementType = op.getSourceType().getElementType();
+  SmallVector<OpFoldResult> outputShape = getAsOpFoldResult(reifiedShapes[0]);
+  if (options.createDestTensorFn)
+    dest = options.createDestTensorFn(rewriter, op->getLoc(), outputShape,
+                                      elementType);
+  else
+    dest = rewriter.create<linalg::InitTensorOp>(op->getLoc(), outputShape,
+                                                 elementType);
+
+  FailureOr<Value> result = materializeSliceFromCollapseShape(
+      rewriter, collapseOp, dest, op.getMixedOffsets(), op.getMixedSizes(),
+      op.getMixedStrides(), options.useForeach);
+  if (failed(result))
+    return rewriter.notifyMatchFailure(
+        op, "failed to extract slice from tensor.collapse_shape");
+  rewriter.replaceOp(op, *result);
+  return success();
+}

mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir

+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-rewrite-extract-slice-from-collapse-shape %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-rewrite-extract-slice-from-collapse-shape use-foreach" %s | FileCheck %s --check-prefix=FOREACH
+
+func.func @extract_slice_static(%input: tensor<3x5x7x11xf32>) -> tensor<20x11xf32> {
+  %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x5x7x11xf32> into tensor<105x11xf32>
+  %slice = tensor.extract_slice %collapsed [0, 0] [20, 11] [1, 1] : tensor<105x11xf32> to tensor<20x11xf32>
+  return %slice : tensor<20x11xf32>
+}
+
+// CHECK: func.func @extract_slice_static(%[[arg0:.+]]:
+// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[c20:.+]] = arith.constant 20 : index
+// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[c5:.+]] = arith.constant 5 : index
+// CHECK-DAG: %[[c7:.+]] = arith.constant 7 : index
+// CHECK-DAG: %[[init:.+]] = linalg.init_tensor [20, 11] :
+// CHECK-DAG: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c20]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])
+// CHECK:   %[[multiIndex:.+]]:3 = tensor.delinearize_index %[[iv]](%[[c3]], %[[c5]], %[[c7]]
+// CHECK:   %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] : 
+// CHECK:   %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : 
+// CHECK:   %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 11] [1, 1] : 
+// CHECK:   scf.yield %[[update]] :
+// CHECK: return %[[tile]]
+
+// FOREACH: func.func @extract_slice_static(%[[arg0:.+]]:
+// FOREACH-DAG: %[[c20:.+]] = arith.constant 20 : index
+// FOREACH-DAG: %[[c3:.+]] = arith.constant 3 : index
+// FOREACH-DAG: %[[c5:.+]] = arith.constant 5 : index
+// FOREACH-DAG: %[[c7:.+]] = arith.constant 7 : index
+// FOREACH-DAG: %[[init:.+]] = linalg.init_tensor [20, 11] :
+// FOREACH-DAG: %[[tile:.+]] = scf.foreach_thread (%[[iv:.+]]) in (%[[c20]])
+// FOREACH:   %[[multiIndex:.+]]:3 = tensor.delinearize_index %[[iv]](%[[c3]], %[[c5]], %[[c7]]
+// FOREACH:   %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] : 
+// FOREACH:   %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : 
+// FOREACH:   perform_concurrently
+// FOREACH-NEXT:   tensor.parallel_insert_slice %[[sliceFlat]] into %[[init]][%[[iv]], 0] [1, 11] [1, 1] : 
+// FOREACH: return %[[tile]]
+
+// -----
+
+
+func.func @extract_slice_static_strided(%input: tensor<3x5x7x11xf32>) -> tensor<10x5xf32> {
+  %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x5x7x11xf32> into tensor<105x11xf32>
+  %slice = tensor.extract_slice %collapsed [13, 0] [10, 5] [2, 2] : tensor<105x11xf32> to tensor<10x5xf32>
+  return %slice : tensor<10x5xf32>
+}
+
+// CHECK: func.func @extract_slice_static_strided(%[[arg0:.+]]:
+// CHECK-DAG: %[[c13:.+]] = arith.constant 13 : index
+// CHECK-DAG: %[[c33:.+]] = arith.constant 33 : index
+// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[c5:.+]] = arith.constant 5 : index
+// CHECK-DAG: %[[c7:.+]] = arith.constant 7 : index
+// CHECK: %[[init:.+]] = linalg.init_tensor [10, 5] :
+// CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c13]] to %[[c33]] step %[[c2]] iter_args(%[[iterArg:.+]] = %[[init]])
+// CHECK:   %[[multiIndex:.+]]:3 = tensor.delinearize_index %[[iv]](%[[c3]], %[[c5]], %[[c7]]
+// CHECK:   %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] : 
+// CHECK:   %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : 
+// CHECK:   %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] : 
+// CHECK:   scf.yield %[[update]] :
+// CHECK: return %[[tile]]
+
+
+// -----
+
+
+func.func @extract_slice_dynamic(%input: tensor<3x?x?x11xf32>, %offt: index, %size: index) -> tensor<?x5xf32> {
+  %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x?x?x11xf32> into tensor<?x11xf32>
+  %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 5] [2, 2] : tensor<?x11xf32> to tensor<?x5xf32>
+  return %slice : tensor<?x5xf32>
+}
+
+// CHECK: #[[map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
+// CHECK: func.func @extract_slice_dynamic(%[[arg0:.+]]: tensor<{{.*}}>, %[[lb:.+]]: index, %[[sz:.+]]: index)
+// CHECK-DAG:   %[[c2:.+]] = arith.constant 2 : index
+// CHECK-DAG:   %[[c1:.+]] = arith.constant 1 : index
+// CHECK-DAG:   %[[c3:.+]] = arith.constant 3 : index
+// CHECK:   %[[init:.+]] = linalg.init_tensor [%[[sz]], 5] : tensor<?x5xf32>
+// CHECK:   %[[ub:.+]] = affine.apply #[[map1]]()[%[[lb]], %[[sz]]]
+// CHECK:   %[[tile:.+]] = scf.for %[[iv:.+]] = %[[lb]] to %[[ub]] step %[[c2]] iter_args(%[[iterArg:.+]] = %[[init]])
+// CHECK-DAG: %[[d1:.+]] = tensor.dim %arg0, %[[c1]] : tensor<3x?x?x11xf32>
+// CHECK-DAG: %[[d2:.+]] = tensor.dim %arg0, %[[c2]] : tensor<3x?x?x11xf32>
+// CHECK:     %[[multiIndex:.+]]:3 = tensor.delinearize_index %arg3(%[[c3]], %[[d1]], %[[d2]] : index, index, index)
+// CHECK:     %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] :
+// CHECK:     %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :
+// CHECK:     %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] :
+// CHECK:     scf.yield %[[update]] :
+// CHECK:   return %[[tile]] :
+
+// -----
+
+
+func.func @extract_slice_dynamic_multidim(%input: tensor<3x?x?x11x?xf32>, %offt0: index, %size0: index, %offt1: index, %size1: index) -> tensor<?x?xf32> {
+  %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3, 4]] : tensor<3x?x?x11x?xf32> into tensor<?x?xf32>
+  %slice = tensor.extract_slice %collapsed [%offt0, %offt1] [%size0, %size1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  return %slice : tensor<?x?xf32>
+}
+
+// CHECK: #[[map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
+// CHECK: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index)
+// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[c4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[c11:.+]] = arith.constant 11 : index
+// CHECK:     %[[init:.+]] = linalg.init_tensor [%[[sz1]], %[[sz2]]] : tensor<?x?xf32>
+// CHECK:     %[[ub1:.+]] = affine.apply #[[map1:.+]]()[%[[lb1]], %[[sz1]]]
+// CHECK:     %[[ub2:.+]] = affine.apply #[[map1:.+]]()[%[[lb2]], %[[sz2]]]
+// CHECK:     %[[tile1:.+]] = scf.for %[[iv1:.+]] = %[[lb1]] to %[[ub1]] step %[[c1]] iter_args(%[[iterArg1:.+]] = %[[init]])
+// CHECK:       %[[tile2:.+]] = scf.for %[[iv2:.+]] = %[[lb2]] to %[[ub2]] step %[[c1]] iter_args(%[[iterArg2:.+]] = %[[iterArg1]])
+// CHECK-DAG:     %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] : 
+// CHECK-DAG:     %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] : 
+// CHECK:         %[[multiIndex1:.+]]:3 = tensor.delinearize_index %[[iv1]](%[[c3]], %[[d1]], %[[d2]] :
+// CHECK:         %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] :
+// CHECK:         %[[multiIndex2:.+]]:2 = tensor.delinearize_index %[[iv2]](%[[c11]], %[[d4]] :
+// CHECK:         %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : 
+// CHECK:         %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} : 
+// CHECK:         %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg2]][%[[iv1]], %[[iv2]]] [1, 1] [1, 1] : 
+// CHECK:         scf.yield %[[update]] :
+// CHECK:     scf.yield %[[tile2]] :
+// CHECK:   return %[[tile1]] : 
+
+//     FOREACH: #[[map1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//     FOREACH: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index)
+// FOREACH-DAG: %[[c1:.+]] = arith.constant 1 : index
+// FOREACH-DAG: %[[c2:.+]] = arith.constant 2 : index
+// FOREACH-DAG: %[[c3:.+]] = arith.constant 3 : index
+// FOREACH-DAG: %[[c4:.+]] = arith.constant 4 : index
+// FOREACH-DAG: %[[c11:.+]] = arith.constant 11 : index
+//     FOREACH:     %[[init:.+]] = linalg.init_tensor [%[[sz1]], %[[sz2]]] : tensor<?x?xf32>
+//     FOREACH:     %[[tile1:.+]] = scf.foreach_thread (%[[tid1:.+]], %[[tid2:.+]]) in (%[[sz1]], %[[sz2]])
+// FOREACH-DAG:       %[[iv1:.+]] = affine.apply #[[map1]](%[[tid1]])[%[[lb1]]]
+// FOREACH-DAG:       %[[iv2:.+]] = affine.apply #[[map1]](%[[tid2]])[%[[lb2]]]
+// FOREACH-DAG:       %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] : 
+// FOREACH-DAG:       %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] : 
+//     FOREACH:       %[[multiIndex1:.+]]:3 = tensor.delinearize_index %[[iv1]](%[[c3]], %[[d1]], %[[d2]] :
+//     FOREACH:       %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] :
+//     FOREACH:       %[[multiIndex2:.+]]:2 = tensor.delinearize_index %[[iv2]](%[[c11]], %[[d4]] :
+//     FOREACH:       %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : 
+//     FOREACH:       %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} : 
+//     FOREACH:       perform_concurrently
+//FOREACH-NEXT:         tensor.parallel_insert_slice %[[sliceFlat]] into %[[init]][%[[iv1]], %[[iv2]]] [1, 1] [1, 1] : 
+
+// -----
+
+// Verifies that a linearized dimension that is not sliced does not generate a loop. Note that this
+// only works for static shapes.
+
+// CHECK: @extract_slice_non_sliced_linearized_dim(%[[arg0:.+]]: tensor<{{.*}}>,
+func.func @extract_slice_non_sliced_linearized_dim(%input: tensor<3x?x?x11x2xf32>, %offt: index, %size: index) -> tensor<?x22xf32> {
+  %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3, 4]] : tensor<3x?x?x11x2xf32> into tensor<?x22xf32>  
+  %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 22] [1, 1] : tensor<?x22xf32> to tensor<?x22xf32>
+  // CHECK: scf.for
+  // CHECK-NOT: scf.for
+  // CHECK: %[[multiIndex:.+]]:3 = tensor.delinearize_index
+  // CHECK: tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0, 0] [1, 1, 1, 11, 2] [1, 1, 1, 1, 1]
+  return %slice : tensor<?x22xf32>
+}

mlir/test/lib/Dialect/Tensor/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTensorTestPasses
   TestTensorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRArithmeticDialect
+  MLIRLinalgDialect
   MLIRPass
   MLIRSCFDialect
   MLIRTensorDialect
   MLIRTensorTransforms
   MLIRTransforms
   )

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

 //===- TestTensorTransforms.cpp - Test Tensor transformation patterns -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Tensor transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestTensorTransforms
     : public PassWrapper<TestTensorTransforms, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTensorTransforms)
 
   TestTensorTransforms() = default;
   TestTensorTransforms(const TestTensorTransforms &pass) : PassWrapper(pass) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<arith::ArithmeticDialect, scf::SCFDialect>();
+    registry.insert<arith::ArithmeticDialect, scf::SCFDialect,
+                    linalg::LinalgDialect>();
   }
 
   StringRef getArgument() const final {
     return "test-tensor-transform-patterns";
   }
   StringRef getDescription() const final {
     return "Test Tensor transformation patterns by applying them greedily.";
   }
 
   void runOnOperation() override;
 
   Option<bool> testSplitPaddingPatterns{
       *this, "test-split-padding-patterns",
       llvm::cl::desc("Test patterns to split tensor.pad ops"),
       llvm::cl::init(false)};
 
   Option<bool> testFoldConstantExtractSlice{
       *this, "test-fold-constant-extract-slice",
       llvm::cl::desc("Test folding arith.constant and tensor.extract_slice"),
       llvm::cl::init(false)};
+
+  Option<bool> testRewriteExtractSliceWithTiledCollapseShape{
+      *this, "test-rewrite-extract-slice-from-collapse-shape",
+      llvm::cl::desc("Test swapping tensor.extract_slice of a collapse_shape "
+                     "with loop nest"),
+      llvm::cl::init(false)};
+
+  Option<bool> useForeach{
+      *this, "use-foreach",
+      llvm::cl::desc(
+          "Use the scf.foreach_thread operation when generating loop nests for "
+          "the extract_slice of collapse_shape pattern"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applySplitPaddingPatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateSplitPaddingPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
tensor::ControlConstantExtractSliceFusionFn controlFn =
         constexpr int64_t kConstantFoldingMaxNumElements = 1024;
         return resultType.getNumElements() <= kConstantFoldingMaxNumElements;
       };
 
   tensor::populateFoldConstantExtractSlicePatterns(patterns, controlFn);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
+static void applyRewriteExtractFromCollapseShapePatterns(Operation *rootOp,
+                                                         bool useForeach) {
+  RewritePatternSet patterns(rootOp->getContext());
+  auto options = tensor::SliceCollapseShapeOptions().setUseForeach(useForeach);
+  patterns.add<tensor::RewriteExtractSliceWithTiledCollapseShape>(
+      rootOp->getContext(), options);
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
+}
+
 void TestTensorTransforms::runOnOperation() {
   Operation *rootOp = getOperation();
   if (testSplitPaddingPatterns)
     applySplitPaddingPatterns(rootOp);
   if (testFoldConstantExtractSlice)
     applyFoldConstantExtractSlicePatterns(rootOp);
+  if (testRewriteExtractSliceWithTiledCollapseShape)
+    applyRewriteExtractFromCollapseShapePatterns(rootOp, useForeach);
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir