Address comments

- rename collapseTo to better reflect its usage
- assert it only collapse unit dimensions
- rename ReassociationIndices-using variables to reassocGroup and
  reassocMaps, the same terminology used in tensor.collapse_shape
  documentation
- use more representative test with comments to better explain what the
  patch does

Author: RoboTux

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

@@ -236,37 +236,39 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
-/// Create tensor.collapse_shape to drop dimensions in `dropDims` in tensor
+/// Create tensor.collapse_shape to drop unit dimensions in `dropDims` in tensor
 /// `from`.
 static tensor::CollapseShapeOp
-collapseTo(OpBuilder &b, Location loc, Value from,
-           const llvm::SmallBitVector &dropDims) {
+dropGivenUnitDims(OpBuilder &b, Location loc, Value from,
+                  const llvm::SmallBitVector &dropDims) {
   auto fromType = cast<ShapedType>(from.getType());
-  assert(fromType.getRank() == dropDims.size() &&
+  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> reassocIdxsVec;
-  // Current reassociation indices to add dropped dimension to.
-  ReassociationIndices reassocIdxs;
+  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 reassocIdxs until another non dropped dimension is found.
+  // 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)) {
+    if (dropDims.test(dim))
+      assert(fromType.getShape()[dim] == 1 && "Dropping non unit dimension");
+    else {
       if (foundKeptDim) {
-        reassocIdxsVec.push_back(reassocIdxs);
-        reassocIdxs.clear();
+        reassocMaps.push_back(reassocGroup);
+        reassocGroup.clear();
       }
       foundKeptDim = true;
     }
-    reassocIdxs.push_back(dim);
+    reassocGroup.push_back(dim);
   }
-  if (!reassocIdxs.empty())
-    reassocIdxsVec.push_back(reassocIdxs);
-  return b.create<tensor::CollapseShapeOp>(loc, from, reassocIdxsVec);
+  if (!reassocGroup.empty())
+    reassocMaps.push_back(reassocGroup);
+  return b.create<tensor::CollapseShapeOp>(loc, from, reassocMaps);
 }
 
 FailureOr<FusionInfo>
@@ -312,7 +314,7 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
   // 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);
+    def = 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

mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir

@@ -321,47 +321,78 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
 
 // -----
 
-func.func @rank_reduced_extract_slice(%cond : i1) -> tensor<6x2xf32> {
+#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
-  %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>
+  %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>
+    %3 = linalg.generic
+     {indexing_maps = [#map3, #map4, #map5], 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 %if : tensor<6x2xf32>
+  return %1 : tensor<4x6xf32>
 }
 
 //       CHECK: func @rank_reduced_extract_slice(
-//  CHECK-SAME: %[[COND:[0-9a-z]*]]: i1
-
-//       CHECK: %[[EMPTY_PROD:.*]] = tensor.empty() : tensor<6x6x1x1x1x1xf32>
-//       CHECK: %[[FILL_PROD:.*]] = linalg.generic
-//  CHECK-SAME:     outs(%[[EMPTY_PROD]] : tensor<6x6x1x1x1x1xf32>)
+//  CHECK-SAME: %[[ARG0:[0-9a-z]*]]: tensor<1x6x5xf32>
+//  CHECK-SAME: %[[ARG1:[0-9a-z]*]]: tensor<1x5x6xf32>
+//  CHECK-SAME: %[[ARG2:[0-9a-z]*]]: tensor<4x6xf32>
 
-//       CHECK: %[[EMPTY_CONS:.*]] = tensor.empty() : tensor<6x2xf32>
-//       CHECK: %[[EXTRACT_SLICE_CONS:.*]] = tensor.extract_slice %[[EMPTY_PROD]][0, 0, 0, 0, 0, 0] [6, 2, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<6x6x1x1x1x1xf32> to tensor<6x2x1x1x1x1xf32>
-
-//       CHECK: %[[FILL_CONS:.*]] = linalg.generic
-//  CHECK-SAME:     outs(%[[EXTRACT_SLICE_CONS]] : tensor<6x2x1x1x1x1xf32>)
-//       CHECK: %[[CONS_COLLAPSE:.*]] = tensor.collapse_shape %[[FILL_CONS]] {{\[\[0\], \[1, 2, 3, 4, 5\]\]}} : tensor<6x2x1x1x1x1xf32> into tensor<6x2xf32>
-//       CHECK: %[[ADD1_CONS:.*]] = linalg.generic
-//  CHECK-SAME:     ins(%[[CONS_COLLAPSE]] : tensor<6x2xf32>)
-//  CHECK-SAME:     outs(%[[EMPTY_CONS]] : tensor<6x2xf32>)
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index
+//   CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+//       CHECK: %[[EMPTY_PROD:.*]] = tensor.empty() : tensor<1x6x6xf32>
+//  CHECK-NEXT: %[[FILL_PROD:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY_PROD]] : tensor<1x6x6xf32>) -> tensor<1x6x6xf32>
+//  CHECK-NEXT: %[[EMPTY_FOR:.*]] = tensor.empty() : tensor<4x6xf32>
+//  CHECK-NEXT: %[[EMPTY_CONS:.*]] = tensor.empty() : tensor<4x2xf32>
+//  CHECK-NEXT: %[[FILL_CONS:.*]] = linalg.fill ins(%[[CST]] : f32)
+
+//  For loop right after tensor alloc & fill, no linalg.generic.
+//  CHECK-NEXT: %[[FOR:.*]] = scf.for %[[I:[0-9a-z]*]] = %[[C0]] to %[[C6]] step %[[C2]] iter_args(%[[ARG_ITER:.*]] = %[[EMPTY_FOR]])
+
+//  Producer linalg.generic now inside the loop, with tiled args sliced before
+//  it.
+//   CHECK-DAG:   %[[ARG1_SLICE:.*]] = tensor.extract_slice %[[ARG1]][0, 0, %[[I]]] [1, 5, 2] [1, 1, 1]  : tensor<1x5x6xf32> to tensor<1x5x2xf32>
+//   CHECK-DAG:   %[[PROD_SLICE:.*]] = tensor.extract_slice %[[FILL_PROD]][0, 0, %[[I]]] [1, 6, 2] [1, 1, 1]  : tensor<1x6x6xf32> to tensor<1x6x2xf32>
+//       CHECK:    %[[MMUL_PROD:.*]] = linalg.generic
+//  CHECK-SAME:        ins(%[[ARG0]], %[[ARG1_SLICE]] : tensor<1x6x5xf32>, tensor<1x5x2xf32>)
+//  CHECK-SAME:        outs(%[[PROD_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(%[[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>