[mlir][Linalg] Fix incorrect logic in deciding when to fuse reshapes by linearization.

[mravishankar]

Summary

Fusion by linearization should not happen when

• The reshape is expanding and it is a consumer
• The reshape is collapsing and is a producer.

The bug introduced in this logic by some recent refactoring resulted
in a crash.
To enforce this (negetive) use case, add a test that reproduces the
error and verifies the fix.

[ThomasRaoux]

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

1492 ↗

(On Diff #354648)

Looks like this file was uploaded by mistake?

[mravishankar]

Rebase and drop unrelated BUILD file change.

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgFusionOfTensorOps : Pass<"linalg-fusion-for-tensor-ops"> {
   ];
 }
 
 def LinalgFoldReshapeOpsByLinearization :
   Pass<"linalg-fold-reshape-ops-by-linearization"> {
   let summary = "Fold TensorReshapeOps with generic/indexed generic ops by "
                 "linearization";
   let constructor = "mlir::createFoldReshapeOpsByLinearizationPass()";
+  let options = [
+    Option<"allowFoldingUnitDimReshapes", "allow-folding-unit-dim-reshapes",
+           "bool", /*default=*/"false",
+           "Allow fusing linalg.tensor_reshape ops that performs unit "
+           "dimension collapsing">
+  ];
   let dependentDialects = ["AffineDialect", "memref::MemRefDialect"];
 }
 
 def LinalgLowerTiledLoopsToSCF
     : FunctionPass<"convert-linalg-tiled-loops-to-scf"> {
   let summary = "Lower linalg tiled loops to SCF loops and parallel loops";
   let constructor = "mlir::createConvertLinalgTiledLoopsToSCFPass()";
   let dependentDialects = [

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

 /// and reshape:
 /// %1 = linalg.tensor_collapse_shape %0 [[0], [0, 1, 2]] :
 ///        tensor<?x?x4x5xf32> into tensor<?x?xf32>
 ///
 /// would be rewritten into:
 /// %0 = op ... : tensor<?x?x4x5xf32>
 /// with output index_map
 ///   `affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>`
+template <typename TensorReshapeOp>
 static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
-                                        ArrayRef<int64_t> sourceShape,
-                                        ArrayRef<AffineMap> reassociationMaps) {
+                                        TensorReshapeOp reshapeOp) {
+  constexpr bool isExpanding =
+      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value;
+  ArrayRef<int64_t> sourceShape =
+      (isExpanding ? reshapeOp.getResultType().getShape()
+                   : reshapeOp.getSrcType().getShape());
   SmallVector<AffineExpr> resultExprs;
-  resultExprs.reserve(reassociationMaps.size());
   ArrayRef<AffineExpr> sourceExprs = sourceMap.getResults();
   MLIRContext *context = sourceMap.getContext();
 
   // Compute the result exprs based on the reassociation maps.
-  for (AffineMap map : reassociationMaps) {
-    ArrayRef<AffineExpr> collapsedDims = map.getResults();
+  for (auto &indices : reshapeOp.getReassociationIndices()) {
     // Assume that they are in-order and contiguous (already checked in
     // verifier).
-    assert(!collapsedDims.empty());
-    unsigned startDim =
-        collapsedDims.front().cast<AffineDimExpr>().getPosition();
+    assert(!indices.empty());
     SmallVector<int64_t> sizes;
     SmallVector<AffineExpr> dimExprs;
-    for (auto en :
-         llvm::zip(sourceShape.slice(startDim, collapsedDims.size()),
-                   sourceExprs.slice(startDim, collapsedDims.size()))) {
+    for (auto en : llvm::zip(sourceShape.slice(indices[0], indices.size()),
+                             sourceExprs.slice(indices[0], indices.size()))) {
       if (std::get<0>(en) == 1)
         continue;
       sizes.push_back(std::get<0>(en));
       dimExprs.push_back(std::get<1>(en));
     }
     AffineExpr linearizedExpr =
         makeCanonicalStridedLayoutExpr(sizes, dimExprs, context);
     resultExprs.push_back(linearizedExpr);
   }
   return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(),
                         resultExprs, context);
 }
 
 // TensorExpandShapeOp is fusable with its consumer (i.e. reshape as a
 // producer). Fusing when operand has higher rank will require use of mods and
 // divs in the indexing maps of the fused op which would make it non-invertible.
 static bool isTensorReshapeOpFoldableByLinearization(
     TensorExpandShapeOp expandOp, AffineMap useIndexMap, bool asProducer) {
-  if (!asProducer && expandOp.getResultType().hasStaticShape())
+  if (!asProducer)
     return false;
   return useIndexMap.isPermutation();
 }
 
 // TensorCollapseShapeOp is fusable with its producer (i.e. reshape as a
 // consumer).
 static bool isTensorReshapeOpFoldableByLinearization(
     TensorCollapseShapeOp collapseOp, AffineMap useIndexMap, bool asProducer) {
-  if (asProducer && collapseOp.getSrcType().hasStaticShape())
+  if (asProducer)
     return false;
   return useIndexMap.isPermutation();
 }
 
 /// Check if the reshape operation is only expansion into/collapsing of
 /// unit-dimension.
-static bool isUnitDimExpansionOnly(ArrayRef<int64_t> expandedShape,
-                                   ArrayRef<AffineMap> reassociation) {
-  for (auto &map : reassociation) {
+template <typename TensorReshapeOp>
+static bool isUnitDimExpansionOnly(TensorReshapeOp reshapeOp) {
+  constexpr bool isExpanding =
+      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value;
+  ArrayRef<int64_t> expandedShape =
+      (isExpanding ? reshapeOp.getResultType().getShape()
+                   : reshapeOp.getSrcType().getShape());
+  for (auto &indices : reshapeOp.getReassociationIndices()) {
     unsigned numUnitDims = 0;
-    for (AffineExpr expr : map.getResults()) {
-      unsigned position = expr.cast<AffineDimExpr>().getPosition();
+    for (int64_t position : indices)
       if (expandedShape[position] == 1)
         numUnitDims++;
-    }
-    if (numUnitDims != map.getNumResults() - 1)
+    if (numUnitDims != indices.size() - 1)
       return false;
   }
   return true;
 }
 
 /// Conditions for folding a generic operation with a reshape op by expanding
 /// the iteration space dimensionality for tensor operations. These are
 /// preconditions assumed by `foldReshapeByDimExpansion` which implements the
LogicalResult matchAndRewrite(GenericOp genericOp,
     if (!genericOp.hasTensorSemantics())
       return failure();
     SmallVector<OpOperand *> inputOperands = genericOp.getInputOperands();
     for (auto en : llvm::enumerate(inputOperands)) {
       auto reshapeOp = en.value()->get().getDefiningOp<TensorReshapeOp>();
       if (!reshapeOp)
         continue;
 
-      RankedTensorType returnType = reshapeOp.getResultType();
-
       if (!isTensorReshapeOpFoldableByLinearization(
               reshapeOp, genericOp.getTiedIndexingMap(en.value()),
               /*asProducer =*/true) ||
-          (foldUnitDimReshapesOnly &&
-           !isUnitDimExpansionOnly(returnType.getShape(),
-                                   reshapeOp.getReassociationMaps())))
+          (foldUnitDimReshapesOnly && !isUnitDimExpansionOnly(reshapeOp)))
         continue;
 
       // Compute the fused operands list,
       SmallVector<Value> fusedOperands = genericOp.getInputOperands();
       fusedOperands[en.index()] = reshapeOp.src();
       SmallVector<Value> outputOperands = genericOp.getOutputOperands();
       llvm::append_range(fusedOperands, outputOperands);
 
       // Compute indexing_maps for the fused operation. The indexing_maps for
       // the operands of the consumers that arent fused are the same.
       SmallVector<AffineMap> fusedIndexMaps = genericOp.getIndexingMaps();
 
       // Accepted consumer maps are either identity or permutation.
       auto invMap = inversePermutation(fusedIndexMaps[en.index()]);
 
       // Compute the indexing map to use for the result of the producer.
-      AffineMap modifiedMap = linearizeCollapsedDims(
-          invMap, returnType.getShape(), reshapeOp.getReassociationMaps());
+      AffineMap modifiedMap = linearizeCollapsedDims(invMap, reshapeOp);
+      // The modified map cannot have symbols.
+      if (modifiedMap.getNumSymbols())
+        return failure();
       for (AffineExpr expr : modifiedMap.getResults()) {
         if (!expr.isPureAffine())
           return failure();
       }
       fusedIndexMaps[en.index()] = modifiedMap;
 
       // Further check that the resulting index maps can be fused and
       // inverted. Without this the resultant op is not legal.
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     GenericOp producer = reshapeOp.src().template getDefiningOp<GenericOp>();
     if (!producer || !producer.hasTensorSemantics() ||
         producer.getNumOutputs() != 1 ||
         !isTensorReshapeOpFoldableByLinearization(
             reshapeOp,
             producer.getTiedIndexingMap(producer.getOutputOperand(0)),
             /*asProducer =*/false) ||
-        (foldUnitDimReshapesOnly &&
-         !isUnitDimExpansionOnly(reshapeOp.getSrcType().getShape(),
-                                 reshapeOp.getReassociationMaps())))
+        (foldUnitDimReshapesOnly && !isUnitDimExpansionOnly(reshapeOp)))
       return failure();
     // The indexing_maps for the operands of the fused operation are same as
     // those for the operands of the producer.
     SmallVector<AffineMap> fusedIndexMaps = producer.getIndexingMaps();
 
     auto invMap = inversePermutation(
         producer.getTiedIndexingMap(producer.getOutputOperand(0)));
 
     // Compute the indexing map to use for the operand of the producer.
-    AffineMap modifiedMap =
-        linearizeCollapsedDims(invMap, reshapeOp.getSrcType().getShape(),
-                               reshapeOp.getReassociationMaps());
+    AffineMap modifiedMap = linearizeCollapsedDims(invMap, reshapeOp);
     for (AffineExpr expr : modifiedMap.getResults()) {
       if (!expr.isPureAffine()) {
         return rewriter.notifyMatchFailure(
             producer, "fused op indexing map is not affine");
       }
     }
     fusedIndexMaps.back() = modifiedMap;
 
struct FoldReshapeWithGenericOpByExpansion
   using OpRewritePattern<TensorExpandShapeOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(TensorExpandShapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     // Fold only if all constraints of fusing with reshape by expansion are met.
     GenericOp producer = reshapeOp.src().getDefiningOp<GenericOp>();
     if (!producer || producer.getNumOutputs() != 1 ||
         !isFusableWithReshapeByDimExpansion(producer,
                                             producer.getOutputOperand(0)) ||
-        isUnitDimExpansionOnly(reshapeOp.getResultType().getShape(),
-                               reshapeOp.getReassociationMaps()))
+        isUnitDimExpansionOnly(reshapeOp))
       return failure();
     Optional<SmallVector<Value>> replacementValues = fuseWithReshapeByExpansion(
         producer, reshapeOp, producer.getOutputOperand(0), rewriter);
     if (!replacementValues)
       return failure();
     rewriter.replaceOp(reshapeOp, replacementValues.getValue());
     return success();
   }
fuseElementwiseOps(PatternRewriter &rewriter, OpOperand *consumerOpOperand,
   return fuseElementwiseOpsImpl(producer, consumerOpOperand, controlFn,
                                 rewriter);
 }
 
 bool mlir::linalg::skipUnitDimReshape(const OpResult &producer,
                                       const OpOperand &consumer) {
   auto expandShapeOp = producer.getDefiningOp<linalg::TensorExpandShapeOp>();
   if (expandShapeOp)
-    return !isUnitDimExpansionOnly(expandShapeOp.getSrcType().getShape(),
-                                   expandShapeOp.getReassociationMaps());
+    return !isUnitDimExpansionOnly(expandShapeOp);
   auto collapseShapeOp =
       producer.getDefiningOp<linalg::TensorCollapseShapeOp>();
-  return !isUnitDimExpansionOnly(collapseShapeOp.getSrcType().getShape(),
-                                 collapseShapeOp.getReassociationMaps());
+  return !isUnitDimExpansionOnly(collapseShapeOp);
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
 public:
   FuseElementwiseOps(MLIRContext *context, ControlElementwiseOpsFusionFn &fun,
                      PatternBenefit benefit = 1)

 /// Pass to test folding of reshape ops with generic ops by linearization.
 struct FoldReshapeOpsByLinearizationPass
     : public LinalgFoldReshapeOpsByLinearizationBase<
           FoldReshapeOpsByLinearizationPass> {
   void runOnOperation() override {
     Operation *op = getOperation();
     RewritePatternSet patterns(op->getContext());
     populateFoldReshapeOpsByLinearizationPatterns(patterns);
+    if (allowFoldingUnitDimReshapes) {
+      populateFoldUnitDimsReshapeOpsByLinearizationPatterns(patterns);
+    }
     (void)applyPatternsAndFoldGreedily(op->getRegions(), std::move(patterns));
   }
 };
 
 /// Forces `outs` operands of linalg operations to use `linalg.init_tensor` if
 /// the value of the `outs` operand is not used within the op.  This is only
 /// implemented for `linalg.generic` operations for now, but should hold for all
 /// linalg structured ops.

mlir/test/Dialect/Linalg/reshape_linearization_fusion_with_unit_dims.mlir

+// RUN: mlir-opt -linalg-fold-reshape-ops-by-linearization=allow-folding-unit-dim-reshapes -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func @do_not_fold1(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?x1xf32>
+{
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32>
+  %3 = linalg.generic {
+      indexing_maps = [#map, #map, #map],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%2 : tensor<?x?xf32>) {
+      ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+        %4 = addf %arg2, %arg3 : f32
+        linalg.yield %4 : f32
+      } -> tensor<?x?xf32>
+  %4 = linalg.tensor_expand_shape %3 [[0], [1, 2]] : tensor<?x?xf32> into tensor<?x?x1xf32>
+  return %4 : tensor<?x?x1xf32>
+}
+// CHECK-LABEL: func @do_not_fold1
+//       CHECK: %[[VAL:.+]] = linalg.generic
+//       CHECK: linalg.tensor_expand_shape %[[VAL]]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func @do_not_fold2(%arg0 : tensor<?x?x1xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32>
+{
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2]] : tensor<?x?x1xf32> into tensor<?x?xf32>
+  %1 = tensor.dim %arg1, %c0 : tensor<?x?xf32>
+  %2 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %3 = linalg.init_tensor [%1, %2] : tensor<?x?xf32>
+  %4 = linalg.generic {
+      indexing_maps = [#map, #map, #map],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%3 : tensor<?x?xf32>) {
+      ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+        %4 = addf %arg2, %arg3 : f32
+        linalg.yield %4 : f32
+      } -> tensor<?x?xf32>
+  return %4 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @do_not_fold2
+//       CHECK: %[[VAL:.+]] = linalg.tensor_collapse_shape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   ins(%[[VAL]], %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)