Repository: llvm/llvm-project
Pull request #130961: [MLIR][Linalg] Fix insert_slice fusion with rank reduction
Author: login: RoboTux
Labels: name: mlir:linalg
name: mlir
name: mlir:tensor
Created: 2025-03-12T13:28:09Z | Merged: 2025-05-23T08:32:40Z

Insert_slice fusion with a linalg producer does not account for
possible rank-reduction in the insert_slice return type. When that
happens, a tensor.cast gets generated due to the type mismatch which is
invalid for tensor with different rank. This later trips other pass.



Diff:
diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index 22ca8a99dd7db..1a4733df3f187 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -43,6 +43,11 @@ FailureOr<RankedTensorType>
 computeTransposedType(RankedTensorType rankedTensorType,
                       ArrayRef<int64_t> transposeVector);
 
+/// Create tensor.collapse_shape to drop unit dimensions in `dropDims` in tensor
+/// `src`.
+CollapseShapeOp dropGivenUnitDims(OpBuilder &b, Location loc, Value src,
+                                  const llvm::SmallBitVector &dropDims);
+
 /// A tensor.insert_slice is a cast-like operation if it merely rank-extends the
 /// source tensor or inserts the source tensor into a destination tensor with
 /// the same shape.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
index fcfb499bb1332..4fc8a17554435 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dominance.h"
@@ -26,6 +27,7 @@
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
@@ -271,12 +273,20 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Check if rank-reduction occurred as part of the extract_slice. If yes,
+  // collapse the dropped dimensions.
+  if (cast<ShapedType>(consumerType).getRank() !=
+      cast<ShapedType>(def.getType()).getRank()) {
+    llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
+    def =
+        tensor::dropGivenUnitDims(b, fusedProducer.getLoc(), def, droppedDims);
+  }
   // Canonicalizations are not guaranteed to have happened before constructing
   // `fusedProducer`. In the tensor case this can result in temporary type
   // mismatches. Insert a `tensor.cast` op to propagate the transformation
   // invariant that types are compatible.
-  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
-  Type consumerType = consumerOpOperand.get().getType();
   if (consumerType != def.getType())
     def = b.create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def);
   consumerOpOperand.set(def);
diff --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
index c3d56759a896a..11ae0108594dd 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -94,6 +94,37 @@ mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
   return transposedTensorType;
 }
 
+CollapseShapeOp
+mlir::tensor::dropGivenUnitDims(OpBuilder &b, Location loc, Value src,
+                                const llvm::SmallBitVector &dropDims) {
+  auto srcType = cast<ShapedType>(src.getType());
+  int64_t rank = srcType.getRank();
+  assert(rank == static_cast<int64_t>(dropDims.size()) &&
+         "dropDims dimension does not match src tensor rank");
+  assert(llvm::all_of(
+             dropDims.set_bits(),
+             [&](unsigned dim) { return srcType.getShape()[dim] == 1; }) &&
+         "Dropping non unit dimension");
+  // Computed reassociation map for the corresponding tensor.collapse_shape.
+  SmallVector<ReassociationIndices, 2> reassocMaps;
+  // Current reassociation group to add dropped dimension to.
+
+  int64_t nextDimToGroup = 0;
+  llvm::SmallBitVector keptDims(dropDims);
+  keptDims.flip();
+  int64_t lastSetBit = keptDims.find_last();
+  for (int64_t setBit : keptDims.set_bits()) {
+    // Group consecutive dropped dimension with the next non-dropped dimension.
+    // If this is the last set dimension, also group all subsequent dropped
+    // dimension, if any.
+    int64_t upTo = setBit == lastSetBit ? rank - 1 : setBit;
+    auto seq = llvm::seq_inclusive(nextDimToGroup, upTo);
+    reassocMaps.emplace_back(llvm::make_range(seq.begin(), seq.end()));
+    nextDimToGroup = setBit + 1;
+  }
+  return b.create<tensor::CollapseShapeOp>(loc, src, reassocMaps);
+}
+
 bool mlir::tensor::isCastLikeInsertSliceOp(InsertSliceOp op) {
   llvm::SmallBitVector droppedDims = op.getDroppedDims();
   int64_t srcDim = 0;
diff --git a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
index 0f27a92c119cf..fd755a208b2c9 100644
--- a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
+++ b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
@@ -318,3 +318,81 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @rank_reduced_extract_slice(
+    %prod_in: tensor<1x6x5xf32>, %prod_weight: tensor<1x5x6xf32>,
+    %cons_in: tensor<4x6xf32>, %prod_init: tensor<1x6x6xf32>,
+    %for_iv_init: tensor<4x6xf32>, %cons_init: tensor<4x2xf32>
+) -> tensor<4x6xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c6 = arith.constant 6 : index
+  %mmul_prod = linalg.generic
+    {indexing_maps = [#map0, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+    ins(%prod_in, %prod_weight : tensor<1x6x5xf32>, tensor<1x5x6xf32>) outs(%prod_init : tensor<1x6x6xf32>) {
+  ^bb0(%in: f32, %in_1: f32, %out: f32):
+    %10 = arith.mulf %in, %in_1 : f32
+    %11 = arith.addf %out, %10 : f32
+    linalg.yield %11 : f32
+  } -> tensor<1x6x6xf32>
+  %for = scf.for %arg7 = %c0 to %c6 step %c2 iter_args(%arg6 = %for_iv_init) -> (tensor<4x6xf32>) {
+
+    // Extract slice with rank-reduced result type. When fused in the loop
+    // with sliced operands, the producer linalg must have its now sliced
+    // result be rank-reduced as well to match consumer's use type.
+    %prod_slice = tensor.extract_slice %mmul_prod[0, 0, %arg7] [1, 6, 2] [1, 1, 1] : tensor<1x6x6xf32> to tensor<6x2xf32>
+    %mmul_cons = linalg.generic
+     {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"]}
+     ins(%cons_in, %prod_slice : tensor<4x6xf32>, tensor<6x2xf32>) outs(%cons_init : tensor<4x2xf32>) {
+    ^bb0(%in: f32, %in_1: f32, %out: f32):
+      %20 = arith.mulf %in, %in_1 : f32
+      %21 = arith.addf %out, %20 : f32
+      linalg.yield %21 : f32
+    } -> tensor<4x2xf32>
+    %4 = tensor.insert_slice %mmul_cons into %arg6[0, %arg7] [4, 2] [1, 1]  : tensor<4x2xf32> into tensor<4x6xf32>
+    scf.yield %4 : tensor<4x6xf32>
+  }
+  return %for : tensor<4x6xf32>
+}
+
+//       CHECK: func @rank_reduced_extract_slice(
+//  CHECK-SAME: %[[PROD_IN:[0-9a-z]*]]: tensor<1x6x5xf32>
+//  CHECK-SAME: %[[PROD_WEIGHT:[0-9a-z]*]]: tensor<1x5x6xf32>
+//  CHECK-SAME: %[[CONS_IN:[0-9a-z]*]]: tensor<4x6xf32>
+//  CHECK-SAME: %[[PROD_INIT:[0-9a-z]*]]: tensor<1x6x6xf32>
+//  CHECK-SAME: %[[FOR_IV_INIT:[0-9a-z]*]]: tensor<4x6xf32>
+//  CHECK-SAME: %[[CONS_INIT:[0-9a-z]*]]: tensor<4x2xf32>
+
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index
+
+//  For loop right after tensor alloc & fill, no linalg.generic.
+//   CHECK-NOT: linalg.generic
+//  CHECK-NEXT: %[[FOR:.*]] = scf.for %[[I:[0-9a-z]*]] = %[[C0]] to %[[C6]] step %[[C2]] iter_args(%[[ARG_ITER:.*]] = %[[FOR_IV_INIT]])
+
+//  Producer linalg.generic now inside the loop, with tiled args sliced before
+//  it.
+//   CHECK-DAG:   %[[PROD_WEIGHT_SLICE:.*]] = tensor.extract_slice %[[PROD_WEIGHT]][0, 0, %[[I]]] [1, 5, 2] [1, 1, 1]  : tensor<1x5x6xf32> to tensor<1x5x2xf32>
+//   CHECK-DAG:   %[[PROD_INIT_SLICE:.*]] = tensor.extract_slice %[[PROD_INIT]][0, 0, %[[I]]] [1, 6, 2] [1, 1, 1]  : tensor<1x6x6xf32> to tensor<1x6x2xf32>
+//       CHECK:    %[[MMUL_PROD:.*]] = linalg.generic
+//  CHECK-SAME:        ins(%[[PROD_IN]], %[[PROD_WEIGHT_SLICE]] : tensor<1x6x5xf32>, tensor<1x5x2xf32>)
+//  CHECK-SAME:        outs(%[[PROD_INIT_SLICE]] : tensor<1x6x2xf32>)
+//
+//  Consumer uses a rank-reduced version of producer result so a collapse_shape
+//  is generated.
+//       CHECK:    %[[PROD_COLLAPSE:.*]] = tensor.collapse_shape %[[MMUL_PROD]] {{\[\[0, 1\], \[2\]\]}} : tensor<1x6x2xf32> into tensor<6x2xf32>
+//       CHECK:    %[[MMUL_CONS:.*]] = linalg.generic
+//  CHECK-SAME:        ins(%[[CONS_IN]], %[[PROD_COLLAPSE]] : tensor<4x6xf32>, tensor<6x2xf32>)
+//  CHECK-SAME:        outs(%[[CONS_INIT]] : tensor<4x2xf32>)
+//       CHECK:   %[[CONS_SLICE:.*]] = tensor.insert_slice %[[MMUL_CONS]] into %[[ARG_ITER]][0, %[[I]]] [4, 2] [1, 1] : tensor<4x2xf32> into tensor<4x6xf32>
+//       CHECK:   scf.yield %[[CONS_SLICE]] : tensor<4x6xf32>
+//       CHECK: return %[[FOR]] : tensor<4x6xf32>


Reviews:
user: login: banach-space
body: Looks good overall, I've made some minor suggestions. I am surprised we didn't hit this before.

Thanks for the contribution Thomas and apologies for the delay reviewing 🙏🏻 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T11:55:19Z
user: login: banach-space
body: Thanks for the updates, Thomas!

I've left a few smallish comments, but before addressing those, it would be good to make sure that we do indeed need `tensor.collapse_shape` (as opposed to _a rank reducing_ `tensor.extract_slice`). I am sure that you thought it through and there's a good reason to insert `collapse_shape`, but it's not obvious to me just yet 😅 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T09:24:02Z
user: login: RoboTux
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T11:23:59Z
user: login: banach-space
body: To me this makes sense, thanks for the discussion and apologies for the delay!

I've left a few comments inline, but I also have two higher-level suggestion:
* Move `dropGivenUnitDims` to mlir/lib/Dialect/Tensor/Utils.cpp -> this hook is a fairly generic helper.
* The test that's added here is quite complex - could you try simplifying it? In principle, the only thing that matters is "rank reduction", everything else (e.g. other ops) is creating noise. I've made some suggestions inline.

Let me also ping @qedawkins and @javedabsar , who are experts in this area :)
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T17:20:14Z
user: login: RoboTux
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T10:26:55Z
user: login: banach-space
body: Thanks Thomas, final suggestions from me.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T16:09:35Z
user: login: RoboTux
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T21:21:15Z
user: login: banach-space
body: Great, thanks for all the updates, LGTM!

With no comments from other reviewers, I think that this is safe to land. You may want to wait till tomorrow - just in case somebody changes their mind  and decides to review :) (given that this has been in review for a while, I suspect that you are not in a rush)

I really enjoyed the discussion, Thomas! Send more PRs :)
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T12:44:38Z

Review comments:
diff_hunk: @@ -235,6 +237,31 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// `from`.
+tensor::CollapseShapeOp collapseTo(OpBuilder &b, Location loc, Value from,
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: `static`?
created_at: 2025-03-25T11:48:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -235,6 +237,31 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// `from`.
+tensor::CollapseShapeOp collapseTo(OpBuilder &b, Location loc, Value from,
+                                   const llvm::SmallBitVector &dropDims) {
+  auto fromType = cast<ShapedType>(from.getType());
+  assert(fromType.getRank() == dropDims.size());
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: Add a message.
created_at: 2025-03-25T11:48:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -235,6 +237,31 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// `from`.
+tensor::CollapseShapeOp collapseTo(OpBuilder &b, Location loc, Value from,
+                                   const llvm::SmallBitVector &dropDims) {
+  auto fromType = cast<ShapedType>(from.getType());
+  assert(fromType.getRank() == dropDims.size());
+  SmallVector<ReassociationIndices, 2> reassocIdxsVec;
+  ReassociationIndices reassocIdxs;
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: Could you document these two?
created_at: 2025-03-25T11:51:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,66 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+func.func @rank_reduced_extract_slice(%arg0: tensor<6x6x1x1x1x1xf32>, %arg1: tensor<6x6x1x1xf32>, %arg2: tensor<4x6xf32>) -> tensor<4x6xf32> {
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: This test is a bit dense. I have a few suggestions.

1. Would it still work if you replaced the current `linalg.generic`s with something simpler (reducing the number of `arith` Ops would already be a win).
2. Could you add empty lines as in the example above?
3. Could you try using naming consistent with the example above? (e.g. `%0` -> `%gen1` or `%wide_madmul`, `%arg0` -> `%wide_src`)

Hopefully that's not too much work, this looks mostly good.
created_at: 2025-03-25T11:54:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,50 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+func.func @rank_reduced_extract_slice(%cond : i1) -> tensor<6x2xf32> {
+  %cst = arith.constant 0.0 : f32
+  %cst1 = arith.constant 1.0 : f32
+
+  %empty1 = tensor.empty() : tensor<6x6x1x1x1x1xf32>
+  %init1 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} outs(%empty1 : tensor<6x6x1x1x1x1xf32>) {
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: [nit] Would you be able to split this into multiple lines? I doesn't even fit in my browser 😅 There are some examples how to format `linalg.generic` higher up.
created_at: 2025-03-26T08:54:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,50 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+func.func @rank_reduced_extract_slice(%cond : i1) -> tensor<6x2xf32> {
+  %cst = arith.constant 0.0 : f32
+  %cst1 = arith.constant 1.0 : f32
+
+  %empty1 = tensor.empty() : tensor<6x6x1x1x1x1xf32>
+  %init1 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} outs(%empty1 : tensor<6x6x1x1x1x1xf32>) {
+  ^bb0(%out: f32):
+    linalg.yield %cst : f32
+  } -> tensor<6x6x1x1x1x1xf32>
+
+  %if = scf.if %cond -> tensor<6x2xf32> {
+    %extract0 = tensor.extract_slice %init1[0, 0, 0, 0, 0, 0] [6, 2, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<6x6x1x1x1x1xf32> to tensor<6x2xf32>
+
+    %init2 = tensor.empty() : tensor<6x2xf32>
+    %add1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%extract0 : tensor<6x2xf32>) outs(%init2 : tensor<6x2xf32>) {
+      ^bb0(%in: f32, %out: f32):
+        %add = arith.addf %in, %cst1 : f32
+        linalg.yield %add : f32
+    } -> tensor<6x2xf32>
+    scf.yield %add1 : tensor<6x2xf32>
+  } else {
+    %extract2 = tensor.extract_slice %init1[0, 2, 0, 0, 0, 0] [6, 2, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<6x6x1x1x1x1xf32> to tensor<6x2xf32>
+    scf.yield %extract2 : tensor<6x2xf32>
+  }
+
+  return %if : tensor<6x2xf32>
+}
+
+//       CHECK: func @rank_reduced_extract_slice(
+//  CHECK-SAME: %[[COND:[0-9a-z]*]]: i1
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: Is `COND` in any way significant to this test? I don't see it being re-used below this.
created_at: 2025-03-26T08:56:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -272,12 +308,16 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occured as part of the extract_slice.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: I have just realised ... Why not re-use rank-reducing `tensor.extract_slice` instead? It's a bit counter-intuitive to make that a rank-reducing `tensor.extract_slice` is replaced with a pair of `tensor.extract_slice` + `tensor.collapse_shape`. 
created_at: 2025-03-26T09:08:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -235,6 +237,39 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// `from`.
+static tensor::CollapseShapeOp
+collapseTo(OpBuilder &b, Location loc, Value from,
+           const llvm::SmallBitVector &dropDims) {
+  auto fromType = cast<ShapedType>(from.getType());
+  assert(fromType.getRank() == dropDims.size() &&
+         "dropDims dimension does not match from tensor rank");
+  // Computed reassociation map for the corresponding tensor.collapse_shape.
+  SmallVector<ReassociationIndices, 2> reassocIdxsVec;
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: OK, thanks for the comments! Initially I got a bit confused -  in my "head" these are reassociation groups. As in, this is _a vector of reassociation groups_. I would rename this to make it clearer (perhaps there's a better name?)
created_at: 2025-03-26T09:21:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -235,6 +237,39 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// `from`.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: Since this is only used for dropping unit dims ("rank reducing extract slice"), I would make that clear and also add an assert. Generalise is fine, but then we should test more.
created_at: 2025-03-26T09:22:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -272,12 +308,16 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occured as part of the extract_slice.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: RoboTux
body: The extract_slice is done on the producer's arguments while the expand_shape is done on its result. I did look into doing the rank-reducing on the argument since the result type is derived from the DPS init, but the pass applies to both named and generic linalg ops which makes this tricky. In the case of a generic, one of the input might be using the unit dimension in a non-dim affine expression which requires changing the affine maps to update, but since the pass works on the LinalgOp interface there is no way to change the affine map (since named ops have an implicit map). Of course one could do an if but the code becomes quite complex.

My feeling this is the work for another pass. With this one, the producing linalg is moved inside the loop, alongside the consuming generic. Then it's a matter of folding expand shape / extract slice which sounds like a simple linalg fusion that probably already exists.
created_at: 2025-05-07T11:22:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -271,12 +309,16 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occured as part of the extract_slice.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: >   // Rank-reduction occured as part of the extract_slice.

* occured -> occurred
* shouldn't we verify that it's indeed coming from `tensor.extract_slice`?

created_at: 2025-05-16T17:11:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,81 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @rank_reduced_extract_slice(%arg0: tensor<1x6x5xf32>, %arg1: tensor<1x5x6xf32>, %arg2: tensor<4x6xf32>) -> tensor<4x6xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c6 = arith.constant 6 : index
+  %cst = arith.constant 0.0 : f32
+  %init1 = tensor.empty() : tensor<1x6x6xf32>
+  %fill1 = linalg.fill ins(%cst : f32) outs(%init1 : tensor<1x6x6xf32>) -> tensor<1x6x6xf32>
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: This could be replaced with function input parameters.
created_at: 2025-05-16T17:12:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,81 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @rank_reduced_extract_slice(%arg0: tensor<1x6x5xf32>, %arg1: tensor<1x5x6xf32>, %arg2: tensor<4x6xf32>) -> tensor<4x6xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c6 = arith.constant 6 : index
+  %cst = arith.constant 0.0 : f32
+  %init1 = tensor.empty() : tensor<1x6x6xf32>
+  %fill1 = linalg.fill ins(%cst : f32) outs(%init1 : tensor<1x6x6xf32>) -> tensor<1x6x6xf32>
+  %0 = linalg.generic
+    {indexing_maps = [#map0, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+    ins(%arg0, %arg1 : tensor<1x6x5xf32>, tensor<1x5x6xf32>) outs(%fill1 : tensor<1x6x6xf32>) {
+  ^bb0(%in: f32, %in_1: f32, %out: f32):
+    %10 = arith.mulf %in, %in_1 : f32
+    %11 = arith.addf %out, %10 : f32
+    linalg.yield %11 : f32
+  } -> tensor<1x6x6xf32>
+  %init2 = tensor.empty() : tensor<4x6xf32>
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: Could be replaced with a function input param.
created_at: 2025-05-16T17:12:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,81 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @rank_reduced_extract_slice(%arg0: tensor<1x6x5xf32>, %arg1: tensor<1x5x6xf32>, %arg2: tensor<4x6xf32>) -> tensor<4x6xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c6 = arith.constant 6 : index
+  %cst = arith.constant 0.0 : f32
+  %init1 = tensor.empty() : tensor<1x6x6xf32>
+  %fill1 = linalg.fill ins(%cst : f32) outs(%init1 : tensor<1x6x6xf32>) -> tensor<1x6x6xf32>
+  %0 = linalg.generic
+    {indexing_maps = [#map0, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+    ins(%arg0, %arg1 : tensor<1x6x5xf32>, tensor<1x5x6xf32>) outs(%fill1 : tensor<1x6x6xf32>) {
+  ^bb0(%in: f32, %in_1: f32, %out: f32):
+    %10 = arith.mulf %in, %in_1 : f32
+    %11 = arith.addf %out, %10 : f32
+    linalg.yield %11 : f32
+  } -> tensor<1x6x6xf32>
+  %init2 = tensor.empty() : tensor<4x6xf32>
+  %1 = scf.for %arg4 = %c0 to %c6 step %c2 iter_args(%arg3 = %init2) -> (tensor<4x6xf32>) {
+    %2 = tensor.extract_slice %0[0, 0, %arg4] [1, 6, 2] [1, 1, 1] : tensor<1x6x6xf32> to tensor<6x2xf32>
+    %init3 = tensor.empty() : tensor<4x2xf32>
+    %fill3 = linalg.fill ins(%cst : f32) outs(%init3 : tensor<4x2xf32>) -> tensor<4x2xf32>
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: is this required?
created_at: 2025-05-16T17:13:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -234,6 +236,41 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop unit dimensions in `dropDims` in tensor
+/// `from`.
+static tensor::CollapseShapeOp
+dropGivenUnitDims(OpBuilder &b, Location loc, Value from,
+                  const llvm::SmallBitVector &dropDims) {
+  auto fromType = cast<ShapedType>(from.getType());
+  assert(fromType.getRank() == static_cast<int64_t>(dropDims.size()) &&
+         "dropDims dimension does not match from tensor rank");
+  // Computed reassociation map for the corresponding tensor.collapse_shape.
+  SmallVector<ReassociationIndices, 2> reassocMaps;
+  // Current reassociation group to add dropped dimension to.
+  ReassociationIndices reassocGroup;
+
+  bool foundKeptDim = false;
+  // Dropped dimensions might be at the beginning or end of the shape so
+  // combine all contiguous dimensions before and after a given non dropped
+  // dimension in reassocGroup until another non dropped dimension is found.
+  // When that happens, add the reassociation indices to the map.
+  for (int dim = 0; dim < fromType.getRank(); dim++) {
+    if (dropDims.test(dim))
+      assert(fromType.getShape()[dim] == 1 && "Dropping non unit dimension");
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: There is a lot of nested if/else in this loop - I suggest moving this condition out of this loop and simplifying it with e.g. `llvm::all_of`.  This way you will be able to reduce the nesting. I would move this condition to the very beginning of this function (as an early "sanity-check").
created_at: 2025-05-16T17:18:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -271,12 +309,16 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occured as part of the extract_slice.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: RoboTux
body: That function only fuses ExtractSliceOp, see a little above:
`Value inputTensor = consumerOpOperand.get();

  // Must be an extract_slice op to guarantee there are loops we can fuse into.
  auto sliceOp = inputTensor.getDefiningOp<tensor::ExtractSliceOp>();
  if (!sliceOp) {
    LLVM_DEBUG(llvm::dbgs()
               << "\nNot fusable, not an extract_slice op: " << inputTensor);
    return failure();
  }
`
created_at: 2025-05-19T18:01:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -272,12 +308,16 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occured as part of the extract_slice.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: RoboTux
body: Happy for me to close this?
created_at: 2025-05-20T10:26:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -271,12 +274,17 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occurred as part of the extract_slice.
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: [nit] ATM this comment states that something took place _before_ checking that it has indeed happened :) 
```suggestion
  // Check if rank-reduction occurred as part of the extract_slice. If yes, collapse the dims.
```
created_at: 2025-05-20T15:41:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,6 +94,39 @@ mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
   return transposedTensorType;
 }
 
+/// Create tensor.collapse_shape to drop unit dimensions in `dropDims` in tensor
+/// `from`.
path: mlir/lib/Dialect/Tensor/Utils/Utils.cpp
user: login: banach-space
body: As with code, I would avoid duplication and only document this method in Utils.h
created_at: 2025-05-20T15:42:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -254,6 +256,7 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
                << "\nNot fusable, not an extract_slice op: " << inputTensor);
     return failure();
   }
+  llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
path: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
user: login: banach-space
body: Could this be moved closer to where it's used?
created_at: 2025-05-20T15:43:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -43,6 +43,11 @@ FailureOr<RankedTensorType>
 computeTransposedType(RankedTensorType rankedTensorType,
                       ArrayRef<int64_t> transposeVector);
 
+/// Create tensor.collapse_shape to drop unit dimensions in `dropDims` in tensor
+/// `from`.
+CollapseShapeOp dropGivenUnitDims(OpBuilder &b, Location loc, Value from,
+                                  const llvm::SmallBitVector &dropDims);
path: mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
user: login: banach-space
body: [nit] `from` -> `src` to match the Op definition: https://github.com/llvm/llvm-project/blob/dc0dcab397ae3de38141e1995e4b4e5e3bb98660/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td#L1186
created_at: 2025-05-20T15:46:03Z
author_association: CONTRIBUTOR
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: Thanks for simplifying this test!

Could you introduce some more descriptive var names so that it is easier to parse? (and consistent with other tests). I would probably limit this to:
* Function input variables (e.g. `%arg0 --> `%lhs_1`, `%arg1` -> `%rhs_1`, `%arg2` -> `%lhs_2`, `%arg3`->`%out1` etc)
* High-level ops (`scf.for` + `linalg.generic`)

Hopefully that's not too much work 😅 I won't be requesting any further changes after this.
created_at: 2025-05-20T16:08:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,3 +318,76 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @rank_reduced_extract_slice(
+    %arg0: tensor<1x6x5xf32>, %arg1: tensor<1x5x6xf32>, %arg2: tensor<4x6xf32>,
+    %arg3: tensor<1x6x6xf32>, %arg4: tensor<4x6xf32>, %arg5: tensor<4x2xf32>
+) -> tensor<4x6xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c6 = arith.constant 6 : index
+  %0 = linalg.generic
+    {indexing_maps = [#map0, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+    ins(%arg0, %arg1 : tensor<1x6x5xf32>, tensor<1x5x6xf32>) outs(%arg3 : tensor<1x6x6xf32>) {
+  ^bb0(%in: f32, %in_1: f32, %out: f32):
+    %10 = arith.mulf %in, %in_1 : f32
+    %11 = arith.addf %out, %10 : f32
+    linalg.yield %11 : f32
+  } -> tensor<1x6x6xf32>
+  %1 = scf.for %arg7 = %c0 to %c6 step %c2 iter_args(%arg6 = %arg4) -> (tensor<4x6xf32>) {
+    %2 = tensor.extract_slice %0[0, 0, %arg7] [1, 6, 2] [1, 1, 1] : tensor<1x6x6xf32> to tensor<6x2xf32>
path: mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
user: login: banach-space
body: I would add some empty line and a big bold comment to highlight this rank-reducing extract_slice - this is the key part of this test.
created_at: 2025-05-20T16:09:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,6 +94,39 @@ mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
   return transposedTensorType;
 }
 
+/// Create tensor.collapse_shape to drop unit dimensions in `dropDims` in tensor
+/// `from`.
path: mlir/lib/Dialect/Tensor/Utils/Utils.cpp
user: login: RoboTux
body: Doh, I meant to remove that after moving the code there but forgot.
created_at: 2025-05-20T21:21:15Z
author_association: CONTRIBUTOR

PR comments:
user: login: llvmbot
created_at: 2025-03-12T13:28:44Z
body: 
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tensor
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Thomas Preud'homme (RoboTux)

<details>
<summary>Changes</summary>

Insert_slice fusion with a linalg producer does not account for
possible rank-reduction in the insert_slice return type. When that
happens, a tosa.cast gets generated due to the type mismatch which is
invalid for tensor with different rank. This later trips other pass.


---
Full diff: https://github.com/llvm/llvm-project/pull/130961.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp (+34-2) 
- (modified) mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir (+63) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
index 223d728b0b27d..81b204df5a0aa 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dominance.h"
@@ -26,6 +27,7 @@
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
@@ -235,6 +237,31 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
+/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// `from`.
+tensor::CollapseShapeOp collapseTo(OpBuilder &b, Location loc, Value from,
+                                   const llvm::SmallBitVector &dropDims) {
+  auto fromType = cast<ShapedType>(from.getType());
+  assert(fromType.getRank() == dropDims.size());
+  SmallVector<ReassociationIndices, 2> reassocIdxsVec;
+  ReassociationIndices reassocIdxs;
+
+  bool foundKeptDim = false;
+  for (int dim = 0; dim < fromType.getRank(); dim++) {
+    if (!dropDims.test(dim)) {
+      if (foundKeptDim) {
+        reassocIdxsVec.push_back(reassocIdxs);
+        reassocIdxs.clear();
+      }
+      foundKeptDim = true;
+    }
+    reassocIdxs.push_back(dim);
+  }
+  if (!reassocIdxs.empty())
+    reassocIdxsVec.push_back(reassocIdxs);
+  return b.create<tensor::CollapseShapeOp>(loc, from, reassocIdxsVec);
+}
+
 FailureOr<FusionInfo>
 mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
                                    OpOperand &consumerOpOperand) {
@@ -255,6 +282,7 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
                << "\nNot fusable, not an extract_slice op: " << inputTensor);
     return failure();
   }
+  llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
 
   // If producer is already in the same block as consumer, we are done.
   if (consumerOpOperand.get().getParentBlock() ==
@@ -272,12 +300,16 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
            consumerOpOperand);
 
   // Replace use.
+  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+  Type consumerType = consumerOpOperand.get().getType();
+  // Rank-reduction occured as part of the extract_slice.
+  if (cast<ShapedType>(consumerType).getRank() !=
+      cast<ShapedType>(def.getType()).getRank())
+    def = collapseTo(b, fusedProducer.getLoc(), def, droppedDims);
   // Canonicalizations are not guaranteed to have happened before constructing
   // `fusedProducer`. In the tensor case this can result in temporary type
   // mismatches. Insert a `tensor.cast` op to propagate the transformation
   // invariant that types are compatible.
-  Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
-  Type consumerType = consumerOpOperand.get().getType();
   if (consumerType != def.getType())
     def = b.create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def);
   consumerOpOperand.set(def);
diff --git a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
index 0f27a92c119cf..b4fbdfacde899 100644
--- a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
+++ b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
@@ -318,3 +318,66 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
   }
   return %for0 : tensor<64x128xf32>
 }
+
+// -----
+
+func.func @rank_reduced_extract_slice(%arg0: tensor<6x6x1x1x1x1xf32>, %arg1: tensor<6x6x1x1xf32>, %arg2: tensor<4x6xf32>) -> tensor<4x6xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c6 = arith.constant 6 : index
+  %cst = arith.constant 0.0 : f32
+  %init1 = tensor.empty() : tensor<6x6x1x1x1x1xf32>
+  %fill1 = linalg.fill ins(%cst : f32) outs(%init1 : tensor<6x6x1x1x1x1xf32>) -> tensor<6x6x1x1x1x1xf32>
+  %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d6, d5)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<6x6x1x1x1x1xf32>, tensor<6x6x1x1xf32>) outs(%fill1 : tensor<6x6x1x1x1x1xf32>) {
+  ^bb0(%in: f32, %in_1: f32, %out: f32):
+    %10 = arith.mulf %in, %in_1 : f32
+    %11 = arith.addf %out, %10 : f32
+    linalg.yield %11 : f32
+  } -> tensor<6x6x1x1x1x1xf32>
+  %init2 = tensor.empty() : tensor<4x6xf32>
+  %1 = scf.for %arg4 = %c0 to %c6 step %c2 iter_args(%arg3 = %init2) -> (tensor<4x6xf32>) {
+    %2 = tensor.extract_slice %0[0, %arg4, 0, 0, 0, 0] [6, 2, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<6x6x1x1x1x1xf32> to tensor<6x2xf32>
+    %init3 = tensor.empty() : tensor<4x2xf32>
+    %fill3 = linalg.fill ins(%cst : f32) outs(%init3 : tensor<4x2xf32>) -> tensor<4x2xf32>
+    %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg2, %2 : tensor<4x6xf32>, tensor<6x2xf32>) outs(%fill3 : tensor<4x2xf32>) {
+    ^bb0(%in: f32, %in_1: f32, %out: f32):
+      %20 = arith.mulf %in, %in_1 : f32
+      %21 = arith.addf %out, %20 : f32
+      linalg.yield %21 : f32
+    } -> tensor<4x2xf32>
+    %4 = tensor.insert_slice %3 into %arg3[0, %arg4] [4, 2] [1, 1]  : tensor<4x2xf32> into tensor<4x6xf32>
+    scf.yield %4 : tensor<4x6xf32>
+  }
+  return %1 : tensor<4x6xf32>
+}
+
+//       CHECK: func @rank_reduced_extract_slice(
+//  CHECK-SAME: %[[ARG0:[0-9a-z]*]]: tensor<6x6x1x1x1x1xf32>
+//  CHECK-SAME: %[[ARG1:[0-9a-z]*]]: tensor<6x6x1x1xf32>
+//  CHECK-SAME: %[[ARG2:[0-9a-z]*]]: tensor<4x6xf32>
+
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index
+//       CHECK: %[[EMPTY_PROD:.*]] = tensor.empty() : tensor<6x6x1x1x1x1xf32>
+//       CHECK: %[[FILL_PROD:.*]] = linalg.fill ins({{%.*}} : f32)
+//  CHECK-SAME:     outs(%[[EMPTY_PROD]] : tensor<6x6x1x1x1x1xf32>) -> tensor<6x6x1x1x1x1xf32>
+//       CHECK: %[[EMPTY_FOR:.*]] = tensor.empty() : tensor<4x6xf32>
+//       CHECK: %[[EMPTY_CONS:.*]] = tensor.empty() : tensor<4x2xf32>
+//       CHECK: %[[FILL_CONS:.*]] = linalg.fill ins({{%.*}} : f32)
+//  CHECK-SAME:     outs(%[[EMPTY_CONS]] : tensor<4x2xf32>) -> tensor<4x2xf32>
+//       CHECK: %[[FOR:.*]] = scf.for %[[I:[0-9a-z]*]] = %[[C0]] to %[[C6]] step %[[C2]] iter_args(%[[ARG_ITER:.*]] = %[[EMPTY_FOR]])
+//   CHECK-DAG:   %[[ARG0_SLICE:.*]] = tensor.extract_slice %[[ARG0]][0, %[[I]], 0, 0, 0, 0] [6, 2, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1]  : tensor<6x6x1x1x1x1xf32> to tensor<6x2x1x1x1x1xf32>
+//   CHECK-DAG:   %[[ARG1_SLICE:.*]] = tensor.extract_slice %[[ARG1]][0, %[[I]], 0, 0] [6, 2, 1, 1] [1, 1, 1, 1]  : tensor<6x6x1x1xf32> to tensor<6x2x1x1xf32>
+//   CHECK-DAG:   %[[FILL_PROD_SLICE:.*]] = tensor.extract_slice %[[FILL_PROD]][0, %[[I]], 0, 0, 0, 0] [6, 2, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<6x6x1x1x1x1xf32> to tensor<6x2x1x1x1x1xf32>
+
+//       CHECK:    %[[MMUL_PROD:.*]] = linalg.generic
+//  CHECK-SAME:        ins(%[[ARG0_SLICE]], %[[ARG1_SLICE]] : tensor<6x2x1x1x1x1xf32>, tensor<6x2x1x1xf32>)
+//  CHECK-SAME:        outs(%[[FILL_PROD_SLICE]] : tensor<6x2x1x1x1x1xf32>)
+//       CHECK:    %[[PROD_COLLAPSE:.*]] = tensor.collapse_shape %[[MMUL_PROD]] {{\[\[0\], \[1, 2, 3, 4, 5\]\]}} : tensor<6x2x1x1x1x1xf32> into tensor<6x2xf32>
+//       CHECK:    %[[MMUL_CONS:.*]] = linalg.generic
+//  CHECK-SAME:        ins(%[[ARG2]], %[[PROD_COLLAPSE]] : tensor<4x6xf32>, tensor<6x2xf32>)
+//  CHECK-SAME:        outs(%[[FILL_CONS]] : tensor<4x2xf32>)
+//       CHECK:   %[[CONS_SLICE:.*]] = tensor.insert_slice %[[MMUL_CONS]] into %[[ARG_ITER]][0, %[[I]]] [4, 2] [1, 1] : tensor<4x2xf32> into tensor<4x6xf32>
+//       CHECK:   scf.yield %[[CONS_SLICE]] : tensor<4x6xf32>
+//       CHECK: return %[[FOR]] : tensor<4x6xf32>

``````````

</details>

author_association: MEMBER
user: login: RoboTux
created_at: 2025-03-17T09:44:54Z
body: Ping?
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-03-24T21:13:02Z
body: Ping?
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-03-26T09:32:02Z
body: Thanks for the review @banach-space , I hope the new revision addresses all your comment in a satisfactory way.
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-03-26T09:46:53Z
body: > Thanks for the updates, Thomas!
> 
> I've left a few smallish comments, but before addressing those, it would be good to make sure that we do indeed need `tensor.collapse_shape` (as opposed to _a rank reducing_ `tensor.extract_slice`). I am sure that you thought it through and there's a good reason to insert `collapse_shape`, but it's not obvious to me just yet 😅

That's a very good question and I'll need to get back to you before I can answer fully. To be clear, the fusion does not insert a new extract_slice, this is done by another transform. I don't even think it was there in the earlier example and I don't know why there's a new linalg as well so I will look into understanding that. Normally there is an extract slice that might happen to collapse some unit dimension which then causes a type mismatch. This patch is to ensure we do the proper collapse_shape after the extract_slice is fused to reproduce that unit dimension collapsing and make the type match. By the way the if is because the pass requires the producing linalg and consuming linalg (the one that consume the extract_slice) to be in different block, the bug I was facing was a case of loop but I used a if here to make the code simpler.
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-03-26T15:48:17Z
body: My bad, the pass does indeed create an extract slice (or a subview if working on memref). The code that does that is in Linalg Utils and doesn't deal with dropped unit dimension. I'll rework the patch along those lines instead.
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-05-04T11:01:14Z
body: Sorry for the long pause. I've restarted working on it and should have a patch soon.
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-05-07T11:50:14Z
body: > > Thanks for the updates, Thomas!
> > I've left a few smallish comments, but before addressing those, it would be good to make sure that we do indeed need `tensor.collapse_shape` (as opposed to _a rank reducing_ `tensor.extract_slice`). I am sure that you thought it through and there's a good reason to insert `collapse_shape`, but it's not obvious to me just yet 😅
> 
> That's a very good question and I'll need to get back to you before I can answer fully. To be clear, the fusion does not insert a new extract_slice, this is done by another transform. I don't even think it was there in the earlier example and I don't know why there's a new linalg as well so I will look into understanding that. Normally there is an extract slice that might happen to collapse some unit dimension which then causes a type mismatch. This patch is to ensure we do the proper collapse_shape after the extract_slice is fused to reproduce that unit dimension collapsing and make the type match. By the way the if is because the pass requires the producing linalg and consuming linalg (the one that consume the extract_slice) to be in different block, the bug I was facing was a case of loop but I used a if here to make the code simpler.

This is answered in my other comment: https://github.com/llvm/llvm-project/pull/130961#discussion_r2077407969 but I'll rephrase it here in case that comment wasn't clear enough.

First of all, the fusion is worth doing in case of rank-reducing extract_slice even if an expand_shape is involved because the producer is moved inside the loop and operate directly on rank-reduced buffers. The only case it wouldn't be worth it is if the extract_slice was only used to do rank-reduction but that would be strange IR in my opinion, a collapse_shape should have been used instead in that case.

Second the collapse_shape is used because it would add a lot of complexity to fold the rank-reduction in the extract_slice operations of the producer's operands. This is because the producing linalg could be something like:

```
#map0 = affine_map<(d0, d1) -> (d0, d1)
#map1 = affine_map<(d0, d1) -> (d0, d0 + d1)
linalg.generic
    {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel"]}
    ins(%0, %1 : tensor<1x4xf32>, tensor <1x4xf32>) outs(%2 : tensor<1x4>) {
^bb0(%in : f32, %in1 : f32):
    %result = arith.addf %in, %in1 : f32
    linalg.yield %result : f32
}
// followed by loop with consumer
```

If the dimension #1 was used in a loop of step 2 and dimension #0 was rank-reduced, one would at the very least need to removed the dropped dimension from the indexing_maps as follows:

```
#map0 = affine_map<(d1) -> (d1)
#map1 = affine_map<(d1) -> (0 + d1)

// following is inside the loop where the consumer is
%2 = tensor.empty() : tensor<2xf32>
%3 = tensor.extract_slice %0 into %2[1,%iv][1,2][1,1] : tensor<2xf32>
%4 = tensor.empty() : tensor<2xf32>
%5 = tensor.extract_slice %1 into %4[1,%iv][1,2][1,1] : tensor<2xf32>
linalg.generic
    {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel"]}
    ins(%3, %5 : tensor<2xf32>, tensor <2xf32>) outs(%2 : tensor<2xf32>) {
^bb0(%in : f32, %in1 : f32):
    %result = arith.addf %in, %in1 : f32
    linalg.yield %result : f32
}
```

However the pass does not operate on GenericOp but on LinalgOp which does not allow modifying indexing maps since it can be an interface for a named linalg op with an implicit indexing maps. So one would need to distinguish between GenericOp and other LinalgOp. I also feel that this is a further optimization to fold the rank-reducing earlier, better suited into a dedicated pass if there isn't already one.

Does that make sense to you?
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-20T10:29:28Z
body: <!--LLVM CODE FORMAT COMMENT: {clang-format}-->

:white_check_mark: With the latest revision this PR passed the C/C++ code formatter.
author_association: NONE
user: login: RoboTux
created_at: 2025-05-20T11:01:55Z
body: Sorry I forgot to move the function into Utils.cpp, will do.
author_association: CONTRIBUTOR
user: login: RoboTux
created_at: 2025-05-20T12:32:12Z
body: > 
>     * Move `dropGivenUnitDims` to mlir/lib/Dialect/Tensor/Utils.cpp -> this hook is a fairly generic helper.

Done
author_association: CONTRIBUTOR