[Mlir-commits] [mlir] [MLIR][Linalg] Fix insert_slice fusion with rank reduction (PR #130961)

[Andrzej Warzyński]

================
@@ -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");
----------------
banach-space wrote:

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").

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