Clean up code

dropGivenUnitDims():
- move assert out of loop
- rework algorithm to make grouping more explicit and avoid complex
  nested ifs
- fix occured typo

Test: remove all tensor.empty and linalg.fill

Author: RoboTux

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

@@ -242,32 +242,30 @@ 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()) &&
+  int64_t rank = fromType.getRank();
+  assert(rank == static_cast<int64_t>(dropDims.size()) &&
          "dropDims dimension does not match from tensor rank");
+  assert(llvm::all_of(
+             dropDims.set_bits(),
+             [&](unsigned dim) { return fromType.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.
-  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");
-    else {
-      if (foundKeptDim) {
-        reassocMaps.push_back(reassocGroup);
-        reassocGroup.clear();
-      }
-      foundKeptDim = true;
-    }
-    reassocGroup.push_back(dim);
+
+  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;
   }
-  if (!reassocGroup.empty())
-    reassocMaps.push_back(reassocGroup);
   return b.create<tensor::CollapseShapeOp>(loc, from, reassocMaps);
 }
 
@@ -311,7 +309,7 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
   // Replace use.
   Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
   Type consumerType = consumerOpOperand.get().getType();
-  // Rank-reduction occured as part of the extract_slice.
+  // Rank-reduction occurred as part of the extract_slice.
   if (cast<ShapedType>(consumerType).getRank() !=
       cast<ShapedType>(def.getType()).getRank())
     def = dropGivenUnitDims(b, fusedProducer.getLoc(), def, droppedDims);

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

@@ -327,35 +327,32 @@ func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tens
 #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> {
+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
-  %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>) {
+    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>
-  %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>
+  %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>
     %3 = linalg.generic
      {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"]}
-     ins(%arg2, %2 : tensor<4x6xf32>, tensor<6x2xf32>) outs(%fill3 : tensor<4x2xf32>) {
+     ins(%arg2, %2 : tensor<4x6xf32>, tensor<6x2xf32>) outs(%arg5 : 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>
+    %4 = tensor.insert_slice %3 into %arg6[0, %arg7] [4, 2] [1, 1]  : tensor<4x2xf32> into tensor<4x6xf32>
     scf.yield %4 : tensor<4x6xf32>
   }
   return %1 : tensor<4x6xf32>
@@ -365,24 +362,22 @@ func.func @rank_reduced_extract_slice(%arg0: tensor<1x6x5xf32>, %arg1: tensor<1x
 //  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-SAME: %[[ARG3:[0-9a-z]*]]: tensor<1x6x6xf32>
+//  CHECK-SAME: %[[ARG4:[0-9a-z]*]]: tensor<4x6xf32>
+//  CHECK-SAME: %[[ARG5:[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
-//   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]])
+//   CHECK-NOT: linalg.generic
+//  CHECK-NEXT: %[[FOR:.*]] = scf.for %[[I:[0-9a-z]*]] = %[[C0]] to %[[C6]] step %[[C2]] iter_args(%[[ARG_ITER:.*]] = %[[ARG4]])
 
 //  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-DAG:   %[[PROD_SLICE:.*]] = tensor.extract_slice %[[ARG3]][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>)
@@ -392,7 +387,7 @@ func.func @rank_reduced_extract_slice(%arg0: tensor<1x6x5xf32>, %arg1: tensor<1x
 //       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-SAME:        outs(%[[ARG5]] : 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>