[mlir][Linalg] Avoid collapsing dimensions of linalg op that arent foldable.

[mravishankar]

Summary

The collapsing dimensions transformation is limited to only those
cases where the sequence of dimensions are contiguous in all the
ranges of the indexing maps of the operation. Add this check before
applying the transformation.

[mravishankar]

Fix typo.

[guraypp]

I might be missing something but why do we disable collapsing for this case? For the example your provided we can still collapse d0, d1, d2. I put the collapsed version below. Do you think that is it wrong?

#map = affine_map<(d0, d1) -> (d1, d0)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
module {
  func.func @uncollapsable(%arg0: tensor<41x3x1x57xf32>, %arg1: tensor<3x1x57x41xf32>) -> tensor<3x1x57x41xf32> {
    %collapsed = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<41x3x1x57xf32> into tensor<41x171xf32>
    %collapsed_0 = tensor.collapse_shape %arg1 [[0, 1, 2], [3]] : tensor<3x1x57x41xf32> into tensor<171x41xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%collapsed : tensor<41x171xf32>) outs(%collapsed_0 : tensor<171x41xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<171x41xf32>
    %expanded = tensor.expand_shape %1 [[0, 1, 2], [3]] : tensor<171x41xf32> into tensor<3x1x57x41xf32>
    return %expanded : tensor<3x1x57x41xf32>
  }
}

[mravishankar]

In D150176#4329064, @guraypp wrote:

I might be missing something but why do we disable collapsing for this case? For the example your provided we can still collapse d0, d1, d2. I put the collapsed version below. Do you think that is it wrong?

#map = affine_map<(d0, d1) -> (d1, d0)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
module {
  func.func @uncollapsable(%arg0: tensor<41x3x1x57xf32>, %arg1: tensor<3x1x57x41xf32>) -> tensor<3x1x57x41xf32> {
    %collapsed = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<41x3x1x57xf32> into tensor<41x171xf32>
    %collapsed_0 = tensor.collapse_shape %arg1 [[0, 1, 2], [3]] : tensor<3x1x57x41xf32> into tensor<171x41xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%collapsed : tensor<41x171xf32>) outs(%collapsed_0 : tensor<171x41xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<171x41xf32>
    %expanded = tensor.expand_shape %1 [[0, 1, 2], [3]] : tensor<171x41xf32> into tensor<3x1x57x41xf32>
    return %expanded : tensor<3x1x57x41xf32>
  }
}

The test is trying to fold dimensions 2 and 3 (-test-linalg-elementwise-fusion-patterns=collapse-dimensions-control=2,3 -split-input-file) . Those cant be folded. This is just adding a check and bailing on illegal collapse dimensions list .

See https://github.com/llvm/llvm-project/issues/62315

[guraypp]

Thanks for explaining. I didn’t see that the pass takes arguments.

[mravishankar]

Rebase.

[mravishankar]

Rebase

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///
 ///  return %4 : tensor<16x32xf32>
 /// ```
 FailureOr<SplitReductionResult>
 splitReductionByScaling(RewriterBase &b, LinalgOp op,
                         const ControlSplitReductionFn &controlSplitReductionFn,
                         bool useAlloc = false);
 
-/// Collapses dimensions of linalg.generic operation. It also collapses inputs
-/// before the op and expands outputs after the op.
+/// Return `true`  if a given sequence of dimensions are contiguous in the
+/// range of the specified indexing map.
+bool isDimSequencePreserved(AffineMap map, ReassociationIndicesRef dimSequence);
+/// Return `true` if all sequences of dimensions specified in `dimSequences` are
+/// contiguous in all the ranges of the `maps`.
+bool areDimSequencesPreserved(ArrayRef<AffineMap> maps,
+                              ArrayRef<ReassociationIndices> dimSequences);
+
+/// Collapses dimensions of linalg.generic operation. A precondition to
+/// calling this method is that for each list in `foldedIterationDim`, the
+/// sequence of dimensions is contiguous in domains of all `indexing_maps` of
+/// the `genericOp`. This can be checked using `areDimSequencePreserved` method.
+/// When valid, the method also collapses the operands of the op. Returns
+/// replacement values of the results of the original `genericOp` by inserting
+/// reshapes to get back values of compatible types.
 FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
     GenericOp genericOp, ArrayRef<ReassociationIndices> foldedIterationDims,
     RewriterBase &rewriter);
 
 struct LowerPackResult {
   tensor::PadOp padOp;
   tensor::ExpandShapeOp expandShapeOp;
   linalg::TransposeOp transposeOp;

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

getDomainReassociation(AffineMap indexingMap,
   // The projected permutation semantics ensures that there is no repetition of
   // the domain indices.
   return domainReassociation;
 }
 
 /// For a given `dimSequence`, check if the sequence is conserved in the
 /// `indexingMap`. `indexingMap` is expected to be a projected permutation.
 /// Non-existence of the sequence returns true as well.
-static bool isDimSequencePreserved(AffineMap indexingMap,
+bool mlir::linalg::isDimSequencePreserved(AffineMap indexingMap,
                                           ReassociationIndicesRef dimSequence) {
   assert(!dimSequence.empty() &&
          "expected non-empty list for dimension sequence");
   assert(indexingMap.isProjectedPermutation() &&
          "expected indexing map to be projected permutation");
 
   llvm::SmallDenseSet<unsigned, 4> sequenceElements;
   sequenceElements.insert(dimSequence.begin(), dimSequence.end());
 
for (const auto &expr : enumerate(indexingMap.getResults())) {
     // exist in the indexing map.
     if (sequenceElements.count(dimInMapStart))
       return false;
   }
   // 3. No element of sequence found. Return true.
   return true;
 }
 
+bool mlir::linalg::areDimSequencesPreserved(
+    ArrayRef<AffineMap> maps, ArrayRef<ReassociationIndices> dimSequences) {
+  return llvm::all_of(maps, [&](AffineMap map) {
+    return llvm::all_of(dimSequences, [&](ReassociationIndicesRef dimSequence) {
+      return isDimSequencePreserved(map, dimSequence);
+    });
+  });
+}
+
 // Return the list of dimensions of the iteration domain that can be
 // collapsed to allow for fusion with the a producer that is an expand_shape
 // operation. If all dimensions created by expansion can be collapsed in the
 // iteration space then the reshape is defunct.
 //
 // Example:
 //
 // ```mlir
public:
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     SmallVector<ReassociationIndices> collapsableIterationDims =
         controlCollapseDimension(genericOp);
     if (collapsableIterationDims.empty())
       return failure();
 
+    // Check if the specified list of dimensions to collapse is a valid list.
+    if (!areDimSequencesPreserved(genericOp.getIndexingMapsArray(),
+                                  collapsableIterationDims)) {
+      return rewriter.notifyMatchFailure(
+          genericOp, "specified dimensions cannot be collapsed");
+    }
+
     std::optional<SmallVector<Value>> replacements =
         collapseGenericOpIterationDims(genericOp, collapsableIterationDims,
                                        rewriter);
     if (!replacements) {
       return rewriter.notifyMatchFailure(genericOp,
                                          "failed to collapse dimensions");
     }
     rewriter.replaceOp(genericOp, *replacements);

mlir/test/Dialect/Linalg/collapse-dim.mlir

 // CHECK-LABEL: func @collapse_parallel
 //   CHECK-DAG:  %[[S:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<32x2x10x4096xf32> into tensor<32x2x40960xf32>
 //   CHECK-DAG:  %[[D:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<2x32x10x4096xf32> into tensor<2x32x40960xf32>
 //       CHECK:  %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]],
 //  CHECK-SAME:     iterator_types = ["parallel", "parallel", "parallel"]}
 //  CHECK-SAME:     ins(%[[S]] : tensor<32x2x40960xf32>) outs(%[[D]] : tensor<2x32x40960xf32>) {
 //       CHECK:   } -> tensor<2x32x40960xf32>
 //       CHECK:  tensor.expand_shape %[[R]] {{\[}}[0], [1], [2, 3]] : tensor<2x32x40960xf32> into tensor<2x32x10x4096xf32>
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func.func @uncollapsable(%arg0 : tensor<41x3x1x57xf32>, %arg1 : tensor<3x1x57x41xf32>) -> tensor<3x1x57x41xf32> {
+  %0 = linalg.generic {
+      indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+      ins(%arg0 : tensor<41x3x1x57xf32>) outs(%arg1 : tensor<3x1x57x41xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      linalg.yield %in : f32
+    } -> tensor<3x1x57x41xf32>
+  return %0 : tensor<3x1x57x41xf32>
+}
+// CHECK-LABEL: func @uncollapsable(
+//       CHECK:   linalg.generic
+//  CHECK-SAME:       iterator_types = ["parallel", "parallel", "parallel", "parallel"]