[linalg][fusion] Disallow fusion when it would create an invalid expand_shape

[bkramer]

Summary

The input type of a linalg.generic can be less dynamic than its output
type. If this is the case moving a reshape across the generic op would
create invalid IR, as expand_shape cannot expand arbitrary dynamic
dimensions.

Check that the reshape is actually valid before creating the
expand_shape. This exposes the existing verification logic in reshape
utils and removes the incomplete custom implementation in fusion.

[nicolasvasilache]

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

660

Should this be a static helper on tensor::CollapseOp ?
It feels like it could be reused in a bunch of places.

[mravishankar]

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

789

I am not able to follow why we need this. The source of the expand_shape should be same as the original operand. We should not need an additional cast here, and also not able to follow the example below either.

[bkramer]

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

660

ExpansionInfo is local to the fusion pass though.

789

The example fuses tensor.collapse_shape %arg3 [[0, 1]] : tensor<1x1xi64> into tensor<1xi64> into the input of a generic op. The generic op also takes a tensor<?xi64> input. This means it has to expand the tensor<?xi64> into tensor<1x1xi64>, which can't be done in a single expand_shape.

This is the result of incomplete shape information. The original program would verify the runtime shape of the input to be tensor<1xi64> and then run the linalg ops.

I'll reduce the test case a bit more to make it easier to understand.

[bkramer]

Reduce test case more.

[mravishankar]

mlir/test/Dialect/Linalg/reshape_fusion.mlir

520

Ok, the issue being fixed here is not really an issue. This op as written is incorrect. (Linalg verifier does not check that)
Based on the indexing maps used here, %0 and %arg1 have to be the same shape, otherwise the op is undefined. The transformation only looks at the indexing maps, and therefore assumes the shape of %0 and %arg1 is the same. So its trying to insert a reshape operation.

[bkramer]

mlir/test/Dialect/Linalg/reshape_fusion.mlir

520

What does that mean? The shapes are the same at runtime.

[mravishankar]

mlir/test/Dialect/Linalg/reshape_fusion.mlir

520

Yeah, thats why it isnt checked as part of Linalg verifier. Would it be possible to insert a cast for the input itself?, i.e. make the op

%1 = tensor.cat %arg1 : tensor<?xi64> to tensor<1xi64>
%2 = linalg.generic
   {indexing_maps = [affine_map<(d0) -> (d0)>,
                     affine_map<(d0) -> (d0)>,
                     affine_map<(d0) -> (d0)>],
    iterator_types = ["parallel"]}
   ins(%0, %1 : tensor<1xi64>, tensor<1xi64>) {
     ...
   } -> tensor<1xi64>

If it is known to be always true at runtime then it would be preferable to add a cast during so that all the shapes are consistent.
Though I see your point about the cast. I am just vary of ad-hoc adding casts during transformations instead of adding casts while creating the operation itself, where presumably you have more immediate information of the shapes being the same.

[mravishankar]

(just marking this as needs revision to get it off my must-review list)

[herhut]

mlir/test/Dialect/Linalg/reshape_fusion.mlir

520

So would the contract be that all operands to linalg.generic or linalg operations for that matter, need to have the same static shape annotated? If that is the contract, then we should add checking for it to linalg's verifier. Otherwise we will find these issues only when transformation patterns expect this but it is not true.

Generally, patterns can only rely on static properties that are verified during verification.

[nicolasvasilache]

I don't see a particular semantic problem with either cases (fully static, partially static, fully dynamic): all are valid and are subject to the same UB as a wrong cast would be.
I think Mahesh's suggestion is better because there is no guarantee in shape reification that you'd return 1 and you may well end up with propagating dynamic information further down.
I wouldn't change the verifier behavior here as it could limit the ability to have partial foldings that I think are valid.

[herhut]

In D116600#3234038, @nicolasvasilache wrote:

I think Mahesh's suggestion is better because there is no guarantee in shape reification that you'd return 1 and you may well end up with propagating dynamic information further down.

There could well be a separate pass that inserts these casts to avoid the behavior in shape reification if that is critical to your use cases. I think it would still be better to fix shape reification to return a static value where possible. After all, if this is valid linalg IR, reification should do as good as it can with it.

Or, we declare linalg operations with same static knowledge as the canonical form and add inserting casts as a canonicalization pattern. Not ideal, because we would need to run the canonicalizer before fusion, but at least less surprising.

I wouldn't change the verifier behavior here as it could limit the ability to have partial foldings that I think are valid.

If the verifier allows this, then I think fusion should be able to handle it. We should be able to transform all valid IR, otherwise this will become frustrating to debug.

[nicolasvasilache]

We had a deeper offline discussion with @bkramer
There is a bug that needs to be fixed.
In the same way that we have a "isCastCompatible", we should have a "isShapeCollapse/Expand" compatible.
In this form, the result of that "isShapeCollapse/Expand" is false and the pattern/transformation should fail to apply.

Separately, there is an enabling pattern that would make the type more static by casting the inputs in accordance the semantics of the generic op.
It is unclear whether this should be a blanket canonicalization or just an enabling pattern that needs to be applied selectively to make this case fuse.
This separate enabling pattern is more generally applicable and follows similar orthogonalization principles between legality and profitability that we have been applying everywhere.

[bkramer]

• Expose the shape verification logic (interface is still a bit ugly)
• Use shared logic to stop fusion when the expand_shape would be invalid
• Remove now redundant checking in elemental fusion

[nicolasvasilache]

Can we make proper use of this new function also in the reshape op verifiers?
We should have only one logic used everywhere.
Thanks for pushing on this cleanup!

[nicolasvasilache]

Ah this already happens transitively, LGTM, thanks @bkramer !

[nicolasvasilache]

Please update the commit message to reflect your latest changes, thanks!

[mravishankar]

This makes sense to me. Thanks!

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

 /// Verify that shapes of the reshaped types using following rules
 /// 1) if a dimension in the collapsed type is static, then the corresponding
 ///    dimensions in the expanded shape should be
 ///    a) static
 ///    b) the product should be same as the collaped shape.
 /// 2) if a dimension in the collaped type is dynamic, one and only one of the
 ///    corresponding dimensions in the expanded type should be dynamic. This
 ///    rule is only needed with reshape operations that are expanding.
+LogicalResult reshapeLikeShapesAreCompatible(
+    function_ref<LogicalResult(const Twine &)> emitError,
+    ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
+    ArrayRef<ReassociationIndices> reassociationMaps, bool isExpandingReshape);
+
 template <typename OpTy>
 static LogicalResult verifyReshapeLikeShapes(OpTy op, ShapedType collapsedType,
                                              ShapedType expandedType,
                                              bool isExpandingReshape) {
-  ArrayRef<int64_t> collapsedShape = collapsedType.getShape();
-  ArrayRef<int64_t> expandedShape = expandedType.getShape();
-  unsigned expandedDimStart = 0;
-  for (auto map : llvm::enumerate(op.getReassociationMaps())) {
-    Optional<int64_t> dynamicShape;
-    int64_t linearizedStaticShape = 1;
-    for (auto dim : llvm::enumerate(expandedShape.slice(
-             expandedDimStart, map.value().getNumResults()))) {
-      if (ShapedType::isDynamic(dim.value())) {
-        if (isExpandingReshape && dynamicShape) {
-          return op->emitOpError("invalid to have a single dimension (")
-                 << map.index() << ") expanded into multiple dynamic dims ("
-                 << expandedDimStart + dynamicShape.getValue() << ","
-                 << expandedDimStart + dim.index() << ")";
-        }
-        dynamicShape = dim.index();
-      } else {
-        linearizedStaticShape *= dim.value();
-      }
-    }
-    if (dynamicShape) {
-      if (!ShapedType::isDynamic(collapsedShape[map.index()])) {
-        return op->emitOpError("expected dimension ")
-               << map.index()
-               << " of collapsed type to be dynamic since one or more of the "
-                  "corresponding dimensions in the expanded type is dynamic";
-      }
-    } else {
-      if (collapsedShape[map.index()] != linearizedStaticShape) {
-        return op->emitOpError("expected dimension ")
-               << map.index() << " of collapsed type to be static value of "
-               << linearizedStaticShape << " ";
-      }
-    }
-    expandedDimStart += map.value().getNumResults();
-  }
-  return success();
+  return reshapeLikeShapesAreCompatible(
+      [&](const Twine &msg) { return op->emitOpError(msg); },
+      collapsedType.getShape(), expandedType.getShape(),
+      op.getReassociationIndices(), isExpandingReshape);
 }
 
 /// Pattern to collapse producer/consumer reshape ops that are both collapsing
 /// dimensions or are both expanding dimensions.
 template <typename ReshapeOpTy>
 struct CollapseReshapeOps : public OpRewritePattern<ReshapeOpTy> {
   using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,

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

 /// requires the shape of the expanded dimensions to be static (at least all but
 /// the most significant). For now check that these are all statically sized.
 /// Note that this could be extended to handle dynamic case, but the
 /// implementation below uses `affine.apply` which seems to have issues when the
 /// shapes are not static.
 LogicalResult isGenericOpExpandable(GenericOp genericOp,
                                     const ExpansionInfo &expansionInfo,
                                     PatternRewriter &rewriter) {
-  // Current reshape only supports expansion of a dynamic dim when only one of
-  // the expanded dims are dynamic.
-  for (const auto &originalShape :
-       llvm::enumerate(expansionInfo.getOriginalShape()))
-    if (ShapedType::isDynamic(originalShape.value())) {
-      // All but one of the expanded dims must be static.
-      bool foundDynamicExpandedDim = false;
-      for (auto expandedShape :
-           expansionInfo.getExpandedShapeOfDim(originalShape.index())) {
-        if (ShapedType::isDynamic(expandedShape)) {
-          if (foundDynamicExpandedDim) {
-            return rewriter.notifyMatchFailure(
-                genericOp,
-                "cannot expanded dynamic dims into multiple dynamic dims");
-          }
-          foundDynamicExpandedDim = true;
-        }
-      }
-      if (!foundDynamicExpandedDim) {
-        return rewriter.notifyMatchFailure(
-            genericOp, "dynamic dim expansion needs at least one dynamic dim "
-                       "in result shape");
-      }
-    }
-
   if (!genericOp.hasIndexSemantics())
     return success();
   for (unsigned i : llvm::seq<unsigned>(0, expansionInfo.getOrigOpNumDims())) {
     ArrayRef<int64_t> expandedShape = expansionInfo.getExpandedShapeOfDim(i);
     if (expandedShape.size() == 1)
       continue;
     for (int64_t shape : expandedShape.drop_front()) {
       if (ShapedType::isDynamic(shape)) {
static RankedTensorType getExpandedType(RankedTensorType originalType,
   for (AffineExpr expr : indexingMap.getResults()) {
     unsigned dim = expr.cast<AffineDimExpr>().getPosition();
     auto dimExpansion = expansionInfo.getExpandedShapeOfDim(dim);
     expandedShape.append(dimExpansion.begin(), dimExpansion.end());
   }
   return RankedTensorType::get(expandedShape, originalType.getElementType());
 }
 
 /// Returns the reassociation maps to use in the `tensor.expand_shape`
 /// operation to convert the operands of the original operation to operands of
 /// the expanded operation. The same method is used to compute the
 /// `tensor.collapse_shape` used to collapse the result of the expanded
 /// op to get the value that can replace all uses of the results of the original
 /// op.
 static SmallVector<ReassociationIndices>
 getReassociationForExpansion(AffineMap indexingMap,
                              const ExpansionInfo &expansionInfo) {
fuseWithReshapeByExpansion(GenericOp genericOp, Operation *reshapeOp,
   for (OpOperand *opOperand : genericOp.getInputOperands()) {
     if (opOperand == fusableOpOperand) {
       expandedOpOperands.push_back(isExpanding ? expandingReshapeOp.src()
                                                : collapsingReshapeOp.src());
       continue;
     }
     if (genericOp.isInputTensor(opOperand)) {
       AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
+      auto opOperandType = opOperand->get().getType().cast<RankedTensorType>();
       RankedTensorType expandedOperandType =
-          getExpandedType(opOperand->get().getType().cast<RankedTensorType>(),
-                          indexingMap, expansionInfo);
+          getExpandedType(opOperandType, indexingMap, expansionInfo);
       if (expandedOperandType != opOperand->get().getType()) {
         // Reshape the operand to get the right type.
         SmallVector<ReassociationIndices> reassociation =
             getReassociationForExpansion(indexingMap, expansionInfo);
+        if (failed(reshapeLikeShapesAreCompatible(
+                [&](const Twine &msg) {
+                  return rewriter.notifyMatchFailure(genericOp, msg);
+                },
+                opOperandType.getShape(), expandedOperandType.getShape(),
+                reassociation,
+                /*isExpandingReshape=*/true)))
+          return llvm::None;
         expandedOpOperands.push_back(rewriter.create<tensor::ExpandShapeOp>(
             genericOp.getLoc(), expandedOperandType, opOperand->get(),
             reassociation));
         continue;
       }
     }
     expandedOpOperands.push_back(opOperand->get());
   }
 
   Location loc = genericOp.getLoc();
   SmallVector<Value> outputs;
   for (OpOperand *opOperand : genericOp.getOutputOperands()) {
     AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
+    auto opOperandType = opOperand->get().getType().cast<RankedTensorType>();
     RankedTensorType expandedOutputType =
-        getExpandedType(opOperand->get().getType().cast<RankedTensorType>(),
-                        indexingMap, expansionInfo);
+        getExpandedType(opOperandType, indexingMap, expansionInfo);
     if (expandedOutputType != opOperand->get().getType()) {
       SmallVector<ReassociationIndices> reassociation =
           getReassociationForExpansion(indexingMap, expansionInfo);
+      if (failed(reshapeLikeShapesAreCompatible(
+              [&](const Twine &msg) {
+                return rewriter.notifyMatchFailure(genericOp, msg);
+              },
+              opOperandType.getShape(), expandedOutputType.getShape(),
+              reassociation,
+              /*isExpandingReshape=*/true)))
+        return llvm::None;
       outputs.push_back(rewriter.create<tensor::ExpandShapeOp>(
           genericOp.getLoc(), expandedOutputType, opOperand->get(),
           reassociation));
     }
   }
 
   // The iterator types of the expanded op are all parallel.
   SmallVector<StringRef> iteratorTypes(expansionInfo.getExpandedOpNumDims(),

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

bool mlir::isReassociationValid(ArrayRef<AffineMap> reassociation,
   }
   if (nextExpectedDim != nDims) {
     if (invalidIndex)
       *invalidIndex = reassociation.size() - 1;
     return false;
   }
   return true;
 }
+
+LogicalResult mlir::reshapeLikeShapesAreCompatible(
+    function_ref<LogicalResult(const Twine &)> emitError,
+    ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
+    ArrayRef<ReassociationIndices> reassociationMaps, bool isExpandingReshape) {
+  unsigned expandedDimStart = 0;
+  for (const auto &map : llvm::enumerate(reassociationMaps)) {
+    Optional<int64_t> dynamicShape;
+    int64_t linearizedStaticShape = 1;
+    for (const auto &dim : llvm::enumerate(
+             expandedShape.slice(expandedDimStart, map.value().size()))) {
+      if (ShapedType::isDynamic(dim.value())) {
+        if (isExpandingReshape && dynamicShape) {
+          return emitError("invalid to have a single dimension (" +
+                           Twine(map.index()) +
+                           ") expanded into multiple dynamic dims (" +
+                           Twine(expandedDimStart + dynamicShape.getValue()) +
+                           "," + Twine(expandedDimStart + dim.index()) + ")");
+        }
+        dynamicShape = dim.index();
+      } else {
+        linearizedStaticShape *= dim.value();
+      }
+    }
+    if (dynamicShape) {
+      if (!ShapedType::isDynamic(collapsedShape[map.index()])) {
+        return emitError(
+            "expected dimension " + Twine(map.index()) +
+            " of collapsed type to be dynamic since one or more of the "
+            "corresponding dimensions in the expanded type is dynamic");
+      }
+    } else {
+      if (collapsedShape[map.index()] != linearizedStaticShape) {
+        return emitError("expected dimension " + Twine(map.index()) +
+                         " of collapsed type to be static value of " +
+                         Twine(linearizedStaticShape));
+      }
+    }
+    expandedDimStart += map.value().size();
+  }
+  return success();
+}

mlir/test/Dialect/Linalg/reshape_fusion.mlir

 func @no_fuse_dynamic_dims(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {
   %c0 = arith.constant 0 : index
   %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
   %1 = tensor.dim %0, %c0 : tensor<?xf32>
   %2 = linalg.init_tensor [%1] : tensor<?xf32>
   %3 = linalg.generic {
     indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]}
     ins(%0 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {
       ^bb0(%arg1 : f32, %arg2: f32):
         %4 = arith.addf %arg1, %arg1 : f32
         linalg.yield %4 : f32
     } -> tensor<?xf32>
   return %3 : tensor<?xf32>
 }
 //      CHECK: func @no_fuse_dynamic_dims
 // CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?xf32>
 //      CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
 //      CHECK:   %[[GENERIC:.+]] = linalg.generic
 // CHECK-SAME:       ins(%[[RESHAPE]] : tensor<?xf32>)
 //      CHECK:   return %[[GENERIC]]
+
+// -----
+
+func @no_fuse_mismatched_dynamism(%arg0: tensor<1x1xi64>, %arg1: tensor<?xi64>) -> tensor<1xi64> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<1x1xi64> into tensor<1xi64>
+  %1 = linalg.init_tensor [1] : tensor<1xi64>
+  %2 = linalg.generic
+    {indexing_maps = [affine_map<(d0) -> (d0)>,
+                      affine_map<(d0) -> (d0)>,
+                      affine_map<(d0) -> (d0)>],
+     iterator_types = ["parallel"]}
+    ins(%0, %arg1 : tensor<1xi64>, tensor<?xi64>)
+    outs(%1 : tensor<1xi64>) {
+  ^bb0(%arg4: i64, %arg5: i64, %arg6: i64):  // no predecessors
+    %3 = arith.addi %arg4, %arg5 : i64
+    linalg.yield %3 : i64
+  } -> tensor<1xi64>
+  return %2 : tensor<1xi64>
+}
+
+//      CHECK: func @no_fuse_mismatched_dynamism
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<1x1xi64>
+// CHECK-SAME:     %[[ARG1:.+]]: tensor<?xi64>
+//      CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
+//      CHECK:   %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME:       ins(%[[RESHAPE]], %[[ARG1]] : tensor<1xi64>, tensor<?xi64>)
+//      CHECK:   return %[[GENERIC]]