[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

26

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

32

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

69

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

140

nit: Plz reflow comment

141

nit: space after for

142

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

151

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

196

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

226

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

239

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

251

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

279

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

288

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?

352

nit: values.

[mravishankar]

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

[mravishankar]

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

39

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.

111

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.

345

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

111

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

111

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

39

+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

133

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

133

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] 
+///    %6 = tensor.insert_slice %4 into %arg0 [%iv, 0] [1, 10] [1, 1]
+/// }
+/// ```
+// clang-format on
+/// This function directly creates the materialized slice from offsets and
+/// sizes. The parameter `loopType` can be `scf.for` or `scf.foreach_thread`.
+FailureOr<Value> materializeSliceFromCollapseShape(
+    OpBuilder &builder, tensor::CollapseShapeOp op, Value dest,
+    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+    ArrayRef<OpFoldResult> strides, StringRef loopType);
+
+/// 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 {
+  StringRef loopType;
+  CreateDestTensorFn createDestTensorFn = nullptr;
+
+  SliceCollapseShapeOptions() : loopType("for") {}
+
+  SliceCollapseShapeOptions &setLoopType(StringRef scfLoopType) {
+    loopType = scfLoopType;
+    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/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;
+
+/// Return an OpFoldResult as a Value, assuming it represents an index type.
+static Value getAsValue(OpBuilder &b, Location loc, OpFoldResult ofr) {
+  Optional<int64_t> constValue = getConstantIntValue(ofr);
+  if (constValue.hasValue())
+    return b.create<arith::ConstantIndexOp>(loc, *constValue);
+  return ofr.dyn_cast<Value>();
+}
+
+/// 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 FailureOr<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());
+  if (failed(reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes)))
+    return failure();
+  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.
+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;
+}
+
+/// Create an empty loop nest using the given loop parameters. All `scf.yield`
+/// operations are inserted except for the inner-most loop's yield.
+static scf::LoopNest getEmptyLoopNest(OpBuilder &b, Location loc,
+                                      ArrayRef<Value> nestLowerBounds,
+                                      ArrayRef<Value> nestUpperBounds,
+                                      ArrayRef<Value> nestStrides,
+                                      Value iterArgInit) {
+  scf::LoopNest nest;
+  Value iterArg = iterArgInit;
+  Location currLoc = loc;
+  for (const auto &it :
+       llvm::zip(nestLowerBounds, nestUpperBounds, nestStrides)) {
+    auto loop = b.create<scf::ForOp>(
+        currLoc, std::get<0>(it), std::get<1>(it), std::get<2>(it), iterArg,
+        [&](OpBuilder &nB, Location nLoc, Value nIvs, ValueRange nIterArgs) {
+          currLoc = nLoc;
+          iterArg = nIterArgs[0];
+        });
+    b.setInsertionPointToStart(loop.getBody());
+    nest.loops.push_back(loop);
+  }
+  for (unsigned i = 0, e = nest.loops.size() - 1; i < e; i++) {
+    b.setInsertionPointToEnd(nest.loops[i].getBody());
+    b.create<scf::YieldOp>(nest.loops[i]->getLoc(),
+                           nest.loops[i + 1].getResults());
+  }
+  return nest;
+}
+
+/// Given a `tensor.collapse_shape` op and information regarding which
+/// dimensions have been collapsed and which are being tiled, as well as the
+/// multi-index elements for each delinearized index, create the
+/// tiled form of the matmul.
+static FailureOr<Value> createCollapseShapeTiledResult(
+    CollapseShapeOp collapseOp, Value producer, OpBuilder &b,
+    const llvm::SmallBitVector &slicedOutputDims,
+    const llvm::SmallBitVector &linearizedOutputDims,
+    ArrayRef<SmallVector<Value>> multiIndices, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) {
+  // Construct offsets to extract from the result of the producer op (which is
+  // the input to the `tensor.collapse_shape`).
+  SmallVector<OpFoldResult> sliceOffsets;
+  SmallVector<OpFoldResult> sliceSizes;
+  SmallVector<OpFoldResult> sliceStrides;
+  int64_t loopIdx = 0;
+  RankedTensorType srcType = collapseOp.getSrcType();
+  Location loc = collapseOp->getLoc();
+  SmallVector<ReassociationIndices> reassociationIndices =
+      collapseOp.getReassociationIndices();
+
+  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()]) {
+      sliceSizes.append(it.value().size(), b.getIndexAttr(1));
+      sliceOffsets.append(llvm::to_vector(llvm::map_range(
+          multiIndices[loopIdx++],
+          [&](Value v) -> OpFoldResult { return getAsOpFoldResult(v); })));
+      sliceStrides.append(it.value().size(), b.getIndexAttr(1));
+      continue;
+    }
+
+    // Case 2: One or possible 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))
+          sliceSizes.push_back(b.getIndexAttr(srcType.getDimSize(srcIndex)));
+        else
+          sliceSizes.push_back(b.createOrFold<tensor::DimOp>(
+              loc, producer, b.create<arith::ConstantIndexOp>(loc, srcIndex)));
+        sliceOffsets.push_back(b.getIndexAttr(0));
+        sliceStrides.push_back(b.getIndexAttr(1));
+      }
+      continue;
+    }
+
+    // Case 3: A single index, but it may be sliced.
+    sliceSizes.push_back(sizes[it.index()]);
+    sliceOffsets.push_back(offsets[it.index()]);
+    sliceStrides.push_back(strides[it.index()]);
+  }
+
+  Value tileResult = b.create<tensor::ExtractSliceOp>(
+      loc, producer, sliceOffsets, sliceSizes, sliceStrides);
+
+  // Collapse the dimensions back down.
+  Value collapsedResult =
+      b.create<tensor::CollapseShapeOp>(loc, tileResult, reassociationIndices);
+  return collapsedResult;
+}
+
+FailureOr<Value> mlir::tensor::materializeSliceFromCollapseShape(
+    OpBuilder &b, tensor::CollapseShapeOp collapseOp, Value dest,
+    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+    ArrayRef<OpFoldResult> strides, StringRef loopType) {
+  Location loc = collapseOp->getLoc();
+
+  // Materialize the source shapes.
+  SmallVector<OpFoldResult> sourceShape;
+  for (unsigned i = 0; i < collapseOp.getSrcType().getRank(); i++) {
+    Value dimIdx = b.createOrFold<arith::ConstantIndexOp>(loc, i);
+    sourceShape.push_back(
+        b.createOrFold<tensor::DimOp>(loc, collapseOp.getSrc(), dimIdx));
+  }
+
+  // Try to find which dimensions are sliced and/or linearized.
+  FailureOr<llvm::SmallBitVector> slicedOutputDims =
+      getSlicedDimensions(b, collapseOp, offsets, sizes, strides);
+  if (failed(slicedOutputDims))
+    return failure();
+  llvm::SmallBitVector linearizedOutputDims =
+      getLinearizedDimensions(collapseOp);
+
+  // If there are no sliced and linearized dimensions, then we cannot proceed.
+  llvm::SmallBitVector slicedAndLinearized =
+      *slicedOutputDims & linearizedOutputDims;
+  if (!slicedAndLinearized.any())
+    return failure();
+
+  // Create the bounds for the loop nest.
+  SmallVector<Value> nestLowerBounds;
+  SmallVector<Value> nestUpperBounds;
+  SmallVector<Value> nestStrides;
+  AffineExpr s0, s1, s2;
+  bindSymbols(b.getContext(), s0, s1, s2);
+  for (int idx = slicedAndLinearized.find_first(); idx != -1;
+       idx = slicedAndLinearized.find_next(idx)) {
+    Value strideVal = getAsValue(b, loc, strides[idx]);
+    nestLowerBounds.push_back(getAsValue(b, loc, offsets[idx]));
+    nestUpperBounds.push_back(makeComposedAffineApply(
+        b, loc, s0 + s1 * s2,
+        {nestLowerBounds.back(), getAsValue(b, loc, sizes[idx]), strideVal}));
+    nestStrides.push_back(strideVal);
+  }
+  scf::LoopNest nest = getEmptyLoopNest(b, loc, nestLowerBounds,
+                                        nestUpperBounds, nestStrides, dest);
+
+  // Create the de-linearized multi indices at the start of each loop body.
+  SmallVector<SmallVector<Value>> multiIndices;
+  SmallVector<ReassociationIndices> reassociationIndices =
+      collapseOp.getReassociationIndices();
+  for (unsigned i = 0, loopIdx = 0; i < reassociationIndices.size(); i++) {
+    if (!linearizedOutputDims[i] || !(*slicedOutputDims)[i])
+      continue;
+    assert(loopIdx < nest.loops.size());
+    auto loop = nest.loops[loopIdx++];
+    b.setInsertionPointToStart(loop.getBody());
+    Value iv = loop.getInductionVar();
+
+    SmallVector<OpFoldResult> basis;
+    for (auto idx : reassociationIndices[i]) {
+      basis.push_back(b.createOrFold<tensor::DimOp>(
+          loc, collapseOp.src(), b.create<arith::ConstantIndexOp>(loc, idx)));
+    }
+    auto delinOp = b.create<tensor::DelinearizeIndexOp>(loop->getLoc(),
+                                                        /*linear_index=*/iv,
+                                                        /*basis=*/basis);
+    multiIndices.push_back(llvm::to_vector(llvm::map_range(
+        delinOp.getResults(), [](OpResult r) -> Value { return r; })));
+  }
+
+  // Fill out the first part of the loop body - sub tile for a single
+  // iteration.
+  scf::ForOp innerLoop = nest.loops.back();
+  b.setInsertionPointToEnd(innerLoop.getBody());
+  Value iterArg = innerLoop.getRegionIterArgs()[0];
+  FailureOr<Value> tiledResult = createCollapseShapeTiledResult(
+      collapseOp, collapseOp.src(), b, *slicedOutputDims, linearizedOutputDims,
+      multiIndices, offsets, sizes, strides);
+  if (failed(tiledResult))
+    return failure();
+
+  // Insert the collapse_shape sub-tile into the iteration argument.
+  SmallVector<OpFoldResult> insertOffsets;
+  SmallVector<OpFoldResult> insertSizes;
+  for (unsigned i = 0, loopIdx = 0; i < reassociationIndices.size(); i++) {
+    // Case 1: Linearized dimensions that have been sliced. The insert size is
+    // 1, and the offset is the iv.
+    if (linearizedOutputDims[i] && (*slicedOutputDims)[i]) {
+      insertOffsets.push_back(nest.loops[loopIdx++].getInductionVar());
+      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 = iterArg.getType().cast<RankedTensorType>();
+    if (iterArgsType.isDynamicDim(i))
+      insertSizes.push_back(b.createOrFold<tensor::DimOp>(
+          loc, iterArg, b.createOrFold<arith::ConstantIndexOp>(loc, i)));
+    else
+      insertSizes.push_back(b.getIndexAttr(iterArgsType.getDimSize(i)));
+  }
+
+  Value result = b.create<tensor::InsertSliceOp>(
+      loc, *tiledResult, iterArg, insertOffsets, insertSizes,
+      /*strides=*/
+      SmallVector<OpFoldResult>(insertOffsets.size(), b.getIndexAttr(1)));
+  b.create<scf::YieldOp>(loc, result);
+  b.setInsertionPointAfter(collapseOp);
+  return nest.loops.begin()->getResult(0);
+  return success();
+}
+
+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);
+
+  // Materialize the loop nest and replace the `tensor.extract_slice` op.
+  FailureOr<Value> result = materializeSliceFromCollapseShape(
+      rewriter, collapseOp, dest, op.getMixedOffsets(), op.getMixedSizes(),
+      op.getMixedStrides(), options.loopType);
+  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
+
+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]]
+
+// -----
+
+
+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:     %[[d1:.+]] = tensor.dim %arg0, %[[c1]] : tensor<3x?x?x11xf32>
+// CHECK:     %[[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:   %[[c1:.+]] = arith.constant 1 : index
+// CHECK:   %[[c2:.+]] = arith.constant 2 : index
+// CHECK:   %[[c3:.+]] = arith.constant 3 : index
+// CHECK:   %[[c4:.+]] = arith.constant 4 : index
+// CHECK:   %[[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:     %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] : 
+// CHECK:     %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] : 
+// CHECK:     %[[multiIndex1:.+]]:3 = tensor.delinearize_index %[[iv1]](%[[c3]], %[[d1]], %[[d2]] :
+// CHECK:     %[[tile2:.+]] = scf.for %[[iv2:.+]] = %[[lb2]] to %[[ub2]] step %[[c1]] iter_args(%[[iterArg2:.+]] = %[[iterArg1]])
+// 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]] : 
+
+// -----
+
+// Verifies that a linearized dimension that is not sliced does not generate a loop.
+
+// 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)};
 };
 } // 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) {
+  RewritePatternSet patterns(rootOp->getContext());
+  auto options = tensor::SliceCollapseShapeOptions().setLoopType("for");
+  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);
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir