[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 ↗

(On Diff #444738)

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 ↗

(On Diff #444876)

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

68 ↗

(On Diff #444876)

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

139 ↗

(On Diff #444876)

nit: Plz reflow comment

140 ↗

(On Diff #444876)

nit: space after for

141 ↗

(On Diff #444876)

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

150 ↗

(On Diff #444876)

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

195 ↗

(On Diff #444876)

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

225 ↗

(On Diff #444876)

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

238 ↗

(On Diff #444876)

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

250 ↗

(On Diff #444876)

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

278 ↗

(On Diff #444876)

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

287 ↗

(On Diff #444876)

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 ↗

(On Diff #444876)

nit: values.

[mravishankar]

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

[mravishankar]

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

38 ↗

(On Diff #444876)

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 ↗

(On Diff #444876)

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 ↗

(On Diff #444876)

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 ↗

(On Diff #444876)

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 ↗

(On Diff #444876)

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 ↗

(On Diff #444876)

+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 ↗

(On Diff #454114)

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 ↗

(On Diff #454114)

This looks more like it belongs in a TransformUtils.h

77 ↗

(On Diff #454114)

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 ↗

(On Diff #454114)

More structure to the doc here too plz:

... by:
  1. 
  2. 
  3.

The returned pair ...

114 ↗

(On Diff #454114)

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

94 ↗

(On Diff #454114)

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});

150 ↗

(On Diff #454114)

see, it's a helper :)

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

140 ↗

(On Diff #454114)

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 ↗

(On Diff #454114)

Moved to new file TransformUtils.h and renamed ExtractSliceFromCollapseShapeHelper.

114 ↗

(On Diff #454114)

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

94 ↗

(On Diff #454114)

agreed, fixed

[christopherbate]

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

140 ↗

(On Diff #454114)

I broke it into two patterns.

[christopherbate]

Rebase

[nicolasvasilache]

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

114 ↗

(On Diff #454114)

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 ↗

(On Diff #454114)

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

16 ↗

(On Diff #457632)

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/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
   //    operations. If the producers of the source of the `tensor.extract_slice`
   //    can be tiled such that the tiled value is generated in-place, that
   //    effectively tiles + fuses the operations.
   auto addCandidateSlices = [](Operation *fusedOp,
                                std::deque<tensor::ExtractSliceOp> &candidates) {
     for (Value operand : fusedOp->getOperands())
       if (auto sliceOp = operand.getDefiningOp<tensor::ExtractSliceOp>())
         candidates.push_back(sliceOp);
+
+    if (fusedOp->getNumRegions() > 0) {
+      fusedOp->walk([&](tensor::ExtractSliceOp sliceOp) {
+        candidates.push_back(sliceOp);
+      });
+    }
   };
 
   std::deque<tensor::ExtractSliceOp> candidates;
   addCandidateSlices(tileAndFuseResult.tiledAndFusedOps.back(), candidates);
   OpBuilder::InsertionGuard g(rewriter);
   while (!candidates.empty()) {
     // 2a. Traverse the slices in BFS fashion.
     tensor::ExtractSliceOp candidateSliceOp = candidates.front();

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

 //===- TensorTilingInterface.cpp - Tiling Interface  models *- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #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/Utils/ReshapeOpsUtils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Interfaces/TilingInterface.h"
+#include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
-namespace {
+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>();
+}
 
-struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
+/// For an operation that implements the `ReifyRankedShapedTypeOpInterface`, use
+/// that interface to construct Range's for the output shape.
+static SmallVector<Range>
+getIterationDomainUsingReifyShapedTypeInterface(Operation *op, OpBuilder &b) {
+  // Materialize the output shape values.
+  ReifiedRankedShapedTypeDims reifiedShapes;
+  ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
+      dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
+  (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
 
-  SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
-    ReifiedRankedShapedTypeDims reifiedShapes;
-    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-        dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
-    (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
+  Location loc = op->getLoc();
+  Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
+  Value one = b.create<arith::ConstantIndexOp>(loc, 1);
+  // Initialize all the ranges to {zero, one, one}. All the `ub`s are
+  // overwritten.
+  SmallVector<Range> loopRanges(reifiedShapes[0].size(), {zero, one, one});
+  for (const auto &ub : enumerate(reifiedShapes[0]))
+    loopRanges[ub.index()].size = ub.value();
+  return loopRanges;
+}
+
+/// For an operation that implements `ReifyRankedShapedTypeOpInterface`, use
+/// that interface to create a new destination operand using
+/// `linalg.init_tensor`.
+static Value getDestinationOperandValue(Operation *op, OpBuilder &b) {
 
-    auto padOp = cast<PadOp>(op);
-    SmallVector<OpFoldResult> mixedSizes = getAsOpFoldResult(reifiedShapes[0]);
-    Value initTensor = b.create<linalg::InitTensorOp>(
-        op->getLoc(), mixedSizes, padOp.getResultType().getElementType());
-    return {initTensor};
+      op->getLoc(), mixedSizes,
+      op->getResultTypes()[0].cast<ShapedType>().getElementType());
+  return initTensor;
+}
+
+// The input parameters `offsets` and `sizes` 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).
+static llvm::SmallBitVector getSlicedDimensions(OpBuilder &b,
+                                                CollapseShapeOp op,
+                                                ArrayRef<OpFoldResult> offsets,
+                                                ArrayRef<OpFoldResult> sizes) {
+  SmallVector<Range> ranges =
+      getIterationDomainUsingReifyShapedTypeInterface(op, b);
+  llvm::SmallBitVector result(op.getResultType().getRank());
+  for (const auto &size : llvm::enumerate(sizes)) {
+    result[size.index()] =
+        !isEqualConstantIntOrValue(size.value(), ranges[size.index()].size);
+  }
+  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;
+}
+
+static scf::LoopNest getEmptyLoopNest(OpBuilder &b, Location loc,
+                                      ArrayRef<Value> nestLowerBounds,
+                                      ArrayRef<Value> nestUpperBounds,
+                                      Value iterArgInit) {
+  // Create the loop nest.
+  scf::LoopNest nest;
+  Value iterArg = iterArgInit;
+  Location currLoc = loc;
+  Value one = b.create<arith::ConstantIndexOp>(loc, 1);
+  for (const auto &it : llvm::zip(nestLowerBounds, nestUpperBounds)) {
+    auto loop = b.create<scf::ForOp>(
+        currLoc, std::get<0>(it), std::get<1>(it), one, 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) {
+  // 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;
+  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); })));
+      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));
+      }
+      continue;
+    }
+
+    // Case 3: A single index, but it may be sliced.
+    sliceSizes.push_back(sizes[it.index()]);
+    sliceOffsets.push_back(offsets[it.index()]);
+  }
+
+  Value tileResult = b.create<tensor::ExtractSliceOp>(
+      loc, producer, sliceOffsets, sliceSizes,
+      SmallVector<OpFoldResult>(sliceSizes.size(), b.getIndexAttr(1)));
+
+  // Collapse the dimensions back down.
+  Value collapsedResult =
+      b.create<tensor::CollapseShapeOp>(loc, tileResult, reassociationIndices);
+  return collapsedResult;
+}
+
+namespace {
+
+struct CollapseShapeOpTiling
+    : public TilingInterface::ExternalModel<CollapseShapeOpTiling,
+                                            CollapseShapeOp> {
+
+  /// Materializes the `linalg.init_tensor` destination.
+  SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
+    return {getDestinationOperandValue(op, b)};
   }
 
   SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
-    auto padOp = cast<PadOp>(op);
-    SmallVector<StringRef> iteratorTypes(padOp.getResultType().getRank(),
+    auto collapseOp = cast<CollapseShapeOp>(op);
+    SmallVector<StringRef> iteratorTypes(collapseOp.getResultType().getRank(),
                                          getParallelIteratorTypeName());
     return iteratorTypes;
   }
 
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
-    ReifiedRankedShapedTypeDims reifiedShapes;
-    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-        dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
-    (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
+    return getIterationDomainUsingReifyShapedTypeInterface(op, b);
+  }
+
+  struct LinearizedDimensionInfo {
+    int64_t outputDimension;
+    SmallVector<Value> basis;
+
+    LinearizedDimensionInfo(int64_t outputDimension,
+                            ArrayRef<int64_t> linearizedInputDims,
+                            ArrayRef<Value> inputShape)
+        : outputDimension(outputDimension) {
+      for (auto x : linearizedInputDims)
+        basis.push_back(inputShape[x]);
+    }
+  };
+
+  SmallVector<Operation *>
+  getTiledImplementation(Operation *op, OpBuilder &b, ValueRange dest,
+                         ArrayRef<OpFoldResult> offsets,
+                         ArrayRef<OpFoldResult> sizes,
+                         bool tileDestOperands) const {
+    FailureOr<Value> resultTileValue = generateResultTileValue(
+        op, b, /*resultNumber=*/0, dest, offsets, sizes, tileDestOperands);
+    if (failed(resultTileValue))
+      return {};
+    return {(*resultTileValue).getDefiningOp()};
+  }
 
+  // Produces a tile of `TilingInterface -> tensor.collapse_shape` by
+  // materializing a loop nest that assembles the tile. Any linearized
+  // dimensions are assembled by looping over the delinearized output
+  // (basically, tiling by 1) and stitching the tile together.
+  FailureOr<Value> generateResultTileValue(Operation *op, OpBuilder &b,
+                                           unsigned resultNumber,
+                                           ValueRange dest,
+                                           ArrayRef<OpFoldResult> offsets,
+                                           ArrayRef<OpFoldResult> sizes,
+                                           bool tileDestOperands) const {
+    auto collapseOp = cast<CollapseShapeOp>(op);
     Location loc = op->getLoc();
-    Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
-    Value one = b.create<arith::ConstantIndexOp>(loc, 1);
-    // Initialize all the ranges to {zero, one, one}. All the `ub`s are
-    // overwritten.
-    SmallVector<Range> loopRanges(reifiedShapes[0].size(), {zero, one, one});
-    for (const auto &ub : enumerate(reifiedShapes[0]))
-      loopRanges[ub.index()].size = ub.value();
-    return loopRanges;
+
+    // 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.
+    llvm::SmallBitVector slicedOutputDims =
+        getSlicedDimensions(b, collapseOp, offsets, sizes);
+    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();
+
+    // Any caller of this function has passed an output
+    // buffer via `dest`. That output buffer is the size of the entire
+    // collapse_shape output. We need to slice this buffer to get the tiled
+    // output buffer.
+    Value iterArgInit = b.create<tensor::ExtractSliceOp>(
+        loc, dest[0], offsets, sizes,
+        SmallVector<OpFoldResult>(offsets.size(), b.getIndexAttr(1)));
+
+    // Create the bounds for the loop nest.
+    SmallVector<Value> nestLowerBounds;
+    SmallVector<Value> nestUpperBounds;
+    for (int idx = slicedAndLinearized.find_first(); idx != -1;
+         idx = slicedAndLinearized.find_next(idx)) {
+      nestLowerBounds.push_back(getAsValue(b, loc, offsets[idx]));
+      nestUpperBounds.push_back(makeComposedAffineApply(
+          b, loc, b.getAffineDimExpr(0) + b.getAffineDimExpr(1),
+          {nestLowerBounds.back(), getAsValue(b, loc, sizes[idx])}));
+    }
+    scf::LoopNest nest =
+        getEmptyLoopNest(b, loc, nestLowerBounds, nestUpperBounds, iterArgInit);
+
+    // 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);
+    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(op);
+    return nest.loops.begin()->getResult(0);
+  }
+};
+
+struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
+
+  SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
+    return {getDestinationOperandValue(op, b)};
+  }
+
+  SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
+    auto padOp = cast<PadOp>(op);
+    SmallVector<StringRef> iteratorTypes(padOp.getResultType().getRank(),
+                                         getParallelIteratorTypeName());
+    return iteratorTypes;
+  }
+
+  SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
+    return getIterationDomainUsingReifyShapedTypeInterface(op, b);
   }
 
   SmallVector<Operation *>
   getTiledImplementation(Operation *op, OpBuilder &b, ValueRange dest,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes,
                          bool /*tileDestOperands*/) const {
     Operation *result =
Operation *tensor::bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
   }
   return createPadOfExtractSlice();
 }
 
 void mlir::tensor::registerTilingOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
     tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
+    tensor::CollapseShapeOp::attachInterface<CollapseShapeOpTiling>(*ctx);
   });
 }

mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-interface.mlir

 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //  CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]]
 //      CHECK:       %[[GENERIC_TILE:.+]] = linalg.generic
 // CHECK-SAME:           ins(%[[GEMM_TILE]] :
 // CHECK-SAME:           outs(%[[OUTS_TILE]] :
 //      CHECK:       %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]]
 //      CHECK        scf.yield %[[INSERT]]
+
+// -----
+
+func.func @collapsed_shape_generic_contraction(%arg0 : tensor<?x10x20x30x4xf32>, %arg1 : tensor<40x20x30x4xf32>,
+    %arg2 : tensor<?x10x40xf32>) -> tensor<?x10x40xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0,1], [2,3], [4]] : tensor<?x10x20x30x4xf32> into tensor<?x600x4xf32>
+  %1 = tensor.collapse_shape %arg1 [[0], [1,2], [3]] : tensor<40x20x30x4xf32> into tensor<40x600x4xf32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x10x20x30x4xf32>
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x10x20x30x4xf32>
+  %d = arith.muli %d0, %d1 : index
+  %2 = linalg.init_tensor [%d, 40] : tensor<?x40xf32>
+  %3 = linalg.generic {
+      __internal_linalg_transform__ = "collapsed_shape_generic_contraction",
+      indexing_maps = [
+        affine_map<(d0,d1,d2,d3)->(d0,d2,d3)>,
+        affine_map<(d0,d1,d2,d3)->(d1,d2,d3)>,
+        affine_map<(d0,d1,d2,d3)->(d0,d1)>
+      ],
+      iterator_types = [
+        "parallel",
+        "parallel",
+        "reduction",
+        "reduction"
+      ]
+    }
+      ins(%0, %1 : tensor<?x600x4xf32>, tensor<40x600x4xf32>)
+      outs(%2 : tensor<?x40xf32>) {
+      ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
+        %m = arith.mulf %arg3, %arg4 : f32
+        %a = arith.addf %m, %arg5 : f32
+        linalg.yield %a : f32
+      } -> tensor<?x40xf32>
+      
+  %4 = tensor.expand_shape %3 [[0, 1], [2]] : tensor<?x40xf32> into tensor<?x10x40xf32>
+  return %4 : tensor<?x10x40xf32>
+}
+
+// CHECK-DAG: #[[map0:.+]] =  affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
+// CHECK-DAG: #[[map1:.+]] =  affine_map<(d0)[s0, s1] -> (20, -d0 + s1)>
+// CHECK-DAG: #[[map2:.+]] =  affine_map<(d0)[s0, s1] -> (1, -d0 + s1)>
+// CHECK-DAG: #[[map3:.+]] =  affine_map<()[s0] -> (s0 * 10)>
+// CHECK-DAG: #[[map4:.+]] =  affine_map<(d0, d1) -> (d0 + d1)>
+// CHECK: func.func @collapsed_shape_generic_contraction(
+// CHECK-SAME:              %[[arg0:.+]]: tensor<?x10x20x30x4xf32>,
+// CHECK-SAME:              %[[arg1:.+]]: tensor<40x20x30x4xf32>, 
+// CHECK-DAG:   %[[c10:.+]] = arith.constant 10 :
+// CHECK-DAG:   %[[c0:.+]] = arith.constant 0 :
+// CHECK-DAG:   %[[c40:.+]] = arith.constant 40 :
+// CHECK-DAG:   %[[c600:.+]] = arith.constant 600 :
+// CHECK-DAG:   %[[c20:.+]] = arith.constant 20 :
+// CHECK-DAG:   %[[c1:.+]] = arith.constant 1 :
+// CHECK-DAG:   %[[c30:.+]] = arith.constant 30 :
+// CHECK:   scf.for %[[IV1:.+]] = %[[c0]] to %{{.*}} step %[[c10]] iter_args(%[[IA1:.+]] =
+// CHECK:     %[[TS1:.+]] = affine.min #[[map0]](%[[IV1]])[%[[c10]], %{{.*}}]
+// CHECK:     scf.for %[[IV2:.+]] = %[[c0]] to %[[c40]] step %[[c20]] iter_args(%[[IA2:.+]] = %[[IA1]])
+// CHECK:       %[[TS2:.+]] = affine.min #[[map1]](%[[IV2]])[%[[c20]], %[[c40]]]
+// CHECK:       scf.for %[[IV3:.+]] = %[[c0]] to %[[c600]] step %[[c1]] iter_args(%[[IA3:.+]] = %[[IA2]])
+// CHECK:         %[[TS3:.+]] = affine.min #[[map2]](%[[IV3]])[%[[c1]], %[[c600]]]
+// CHECK:         %[[D0:.+]] = tensor.dim %[[arg0]], %[[c0]] : 
+// CHECK:         %[[OFFT1:.+]] = affine.apply #[[map3]]()[%[[D0]]]
+// CHECK:         %[[INIT:.+]] = linalg.init_tensor
+// CHECK:         %[[TA_INIT:.+]] = tensor.extract_slice %[[INIT]][%[[IV1]], %[[IV3]], 0] [%[[TS1]], %[[TS3]], 4] [1, 1, 1] :
+// CHECK:         %[[OFFT1:.+]] = affine.apply #[[map4]](%[[IV1]], %[[TS1]])
+// CHECK:         %[[OFFT2:.+]] = affine.apply #[[map4]](%[[IV3]], %[[TS3]])
+// CHECK:         %[[TA:.+]] = scf.for %[[IV4:.+]] = %[[IV1]] to %[[OFFT1]] step %[[c1]] iter_args(%[[arg10:.+]] = %[[TA_INIT]])
+// CHECK:           %[[D0:.+]] = tensor.dim %[[arg0]], %[[c0]]
+// CHECK:           %[[MI1:.+]]:2 = tensor.delinearize_index %[[IV4]](%[[D0]], %[[c10]] :
+// CHECK:           scf.for %[[IV5:.+]] = %[[IV3]] to %[[OFFT2]] step %[[c1]] iter_args(%[[arg12:.+]] = %[[arg10]])
+// CHECK:             %[[MI2:.+]]:2 = tensor.delinearize_index %[[IV5]](%[[c20]], %[[c30]] :
+// CHECK:             %[[S1:.+]] = tensor.extract_slice %[[arg0]][%[[MI1]]#0, %[[MI1]]#1, %[[MI2]]#0, %[[MI2]]#1, 0] [1, 1, 1, 1, 4] [1, 1, 1, 1, 1] :
+// CHECK:             %[[S2:.+]] = tensor.collapse_shape %[[S1]] {{\[}}[0, 1], [2, 3], [4]{{\]}} :
+// CHECK:             tensor.insert_slice %[[S2]] into %[[arg12]][%[[IV4]], %[[IV5]], 0] [1, 1, 4] [1, 1, 1] :
+// CHECK:         %[[INIT:.+]] = linalg.init_tensor
+// CHECK:         %[[TB_INIT:.+]] = tensor.extract_slice %[[INIT]]
+// CHECK:         %[[OFFT1:.+]] = affine.apply #[[map4]](%[[IV3]], %[[TS3]])
+// CHECK:         %[[TB:.+]] = scf.for %[[IV4:.+]] = %[[IV3]] to %[[OFFT1]] step %[[c1]] iter_args(%[[arg10:.+]] = %[[TB_INIT]])
+// CHECK:           %[[MI1:.+]]:2 = tensor.delinearize_index %[[IV4]](%[[c20]], %[[c30]] :
+// CHECK:           %[[S1:.+]] = tensor.extract_slice %[[arg1]][%[[IV2]], %[[MI1]]#0, %[[MI1]]#1, 0] [%[[TS2]], 1, 1, 4] [1, 1, 1, 1]
+// CHECK:           %[[S2:.+]] = tensor.collapse_shape %[[S1]] {{\[}}[0], [1, 2], [3]{{\]}}
+// CHECK:           %[[D0:.+]] = tensor.dim %[[arg10]], %[[c0]] :
+// CHECK:           tensor.insert_slice %[[S2]] into %[[arg10]][0, %[[IV4]], 0] [%[[D0]], 1, 4] [1, 1, 1] :
+// CHECK:         %[[TC:.+]] = tensor.extract_slice %[[IA3]][%[[IV1]], %[[IV2]]] [%[[TS1]], %[[TS2]]] [1, 1] :
+// CHECK:         %[[TD:.+]] = linalg.generic
+// CHECK-SAME:                    ins(%[[TA]], %[[TB]] :
+// CHECK-SAME:                    outs(%[[TC]] :
+// CHECK:         tensor.insert_slice %[[TD]] into %[[IA3]][%[[IV1]], %[[IV2]]] [%[[TS1]], %[[TS2]]] [1, 1] :
+
+// -----
+
+func.func @simple_collapse_shape_fusion(%arg0: tensor<4x4x2xf32>, %arg1: tensor<4x4x2xf32>) -> tensor<16x2xf32> {
+  %0 = linalg.init_tensor [4, 4, 2] : tensor<4x4x2xf32>
+  %1 = linalg.generic {
+    indexing_maps = [
+      affine_map<(d0, d1, d2)[]->(d0, d1, d2)>,
+      affine_map<(d0, d1, d2)[]->(d0, d1, d2)>,
+      affine_map<(d0, d1, d2)[]->(d0, d1, d2)>
+    ],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %arg1: tensor<4x4x2xf32>, tensor<4x4x2xf32>)
+    outs(%0 : tensor <4x4x2xf32>){
+    ^bb0(%arg2 : f32, %arg3: f32, %accum: f32):
+      %5 = arith.addf %arg2, %arg3 : f32
+      linalg.yield %5 : f32
+    } -> tensor<4x4x2xf32>
+
+  %2 = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<4x4x2xf32> into tensor<16x2xf32>
+  %3 = linalg.init_tensor [16, 2] : tensor<16x2xf32>
+  %4 = linalg.generic {
+    indexing_maps = [
+      affine_map<(d0, d1)[]->(d0, d1)>,
+      affine_map<(d0, d1)[]->(d0, d1)>
+    ],
+    iterator_types = ["parallel", "parallel"],
+    __internal_linalg_transform__ = "simple_collapse_shape_fusion"
+  } ins(%2: tensor<16x2xf32>)
+    outs(%3 : tensor <16x2xf32>) {
+    ^bb0(%arg4 : f32, %arg5: f32):
+      %c1 = arith.constant 1.0 : f32
+      %6 = arith.addf %arg4, %c1 : f32
+      linalg.yield %6 : f32
+    }  -> tensor<16x2xf32>
+
+  return %4 : tensor<16x2xf32>
+}
+
+// CHECK-DAG: #[[map0:.+]] = affine_map<(d0)[s0, s1] -> (8, -d0 + s1)>
+// CHECK-DAG: #[[map1:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
+// CHECK-DAG: #[[map2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[map3:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: func.func @simple_collapse_shape_fusion(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: tensor<{{.*}}>)
+// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[c16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[c8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[c4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[cst:.+]] = arith.constant 1.0
+// CHECK:     scf.for %[[IV1:.+]] = %[[c0]] to %[[c16]] step %[[c8]] iter_args(%[[IA1:.+]] =
+// CHECK:       %[[TS:.+]] = affine.min #map0(%[[IV1]])[%[[c8]], %[[c16]]]
+// CHECK:       %[[INIT:.+]] = linalg.init_tensor [16, 2] :
+// CHECK:       %[[IA_INIT:.+]] = tensor.extract_slice %[[INIT]][%[[IV1]], 0] [%[[TS]], 2] [1, 1] : 
+// CHECK:       %[[UB:.+]] = affine.apply #map1(%[[IV1]], %[[TS]])
+// CHECK:       %[[TA:.+]] = scf.for %[[IV2:.+]] = %[[IV1]] to %[[UB]] step %[[c1]] iter_args(%[[IA2:.+]] =
+// CHECK:         %[[MI:.+]]:2 = tensor.delinearize_index %[[IV2]](%[[c4]], %[[c4]] :
+// CHECK:         %[[AS:.+]] = tensor.extract_slice %[[arg0]][%[[MI]]#0, %[[MI]]#1, 0] [1, 1, 2] [1, 1, 1] :
+// CHECK:         %[[BS:.+]] = tensor.extract_slice %[[arg1]][%[[MI]]#0, %[[MI]]#1, 0] [1, 1, 2] [1, 1, 1] :
+// CHECK:         %[[CS:.+]] = tensor.extract_slice %0[%[[MI]]#0, %[[MI]]#1, 0] [1, 1, 2] [1, 1, 1] : 
+// CHECK:         %[[ST:.+]] = linalg.generic {
+// CHECK-SAME:      ins(%[[AS]], %[[BS]] :
+// CHECK-SAME:      outs(%[[CS]] :
+// CHECK:         %[[TS_FLAT:.+]] = tensor.collapse_shape %[[ST]] {{\[}}[0, 1], [2]{{\]}} :
+// CHECK:         tensor.insert_slice %[[TS_FLAT]] into %[[IA2]][%[[IV2]], 0] [1, 2] [1, 1] :
+// CHECK:       %[[TB:.+]] = tensor.extract_slice %[[IA1]][%[[IV1]], 0] [%[[TS]], 2] [1, 1] : 
+// CHECK:       %[[TILE:.+]] = linalg.generic
+// CHECK-SAME:     ins(%[[TA]] :
+// CHECK-SAME:     outs(%[[TB]] :
+// CHECK:       tensor.insert_slice %[[TILE]] into %[[IA1]][%[[IV1]], 0] [%3, 2] [1, 1] :

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterface.cpp

void TestTilingInterfacePass::addTestPatterns(MLIRContext *context,
   if (testTileConsumerAndFuseProducer) {
     // 1. Tile and fuse of gemm with bias-add operation.
     addPatternForTiling<
         TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
         context, {10, 20}, "fusion", patterns);
     addPatternForTiling<
         TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
         context, {10}, "gemm_fusion", patterns);
+
+    addPatternForTiling<
+        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
+        context, {8, 0}, "simple_collapse_shape_fusion", patterns);
+    addPatternForTiling<
+        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
+        context, {10, 20, 1, 0}, "collapsed_shape_generic_contraction",
+        patterns);
+
     return;
   }
 }
 
 void TestTilingInterfacePass::runOnOperation() {
   MLIRContext *context = &getContext();
 
   RewritePatternSet tilingPatterns(context);