[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

654

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

795

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

654

ExpansionInfo is local to the fusion pass though.

795

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/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

 //===- ElementwiseOpFusion.cpp - Implementation of linalg Fusion ---------===///
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the linalg dialect Fusion on tensors operations pass.
 //
 //===----------------------------------------------------------------------===//
+#include <numeric>
 #include <utility>
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
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());
 }
 
+static RankedTensorType getCollapsedType(RankedTensorType originalType,
+                                         AffineMap indexingMap,
+                                         const ExpansionInfo &expansionInfo) {
+  SmallVector<int64_t> expandedShape;
+  for (AffineExpr expr : indexingMap.getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    auto dimExpansion = expansionInfo.getExpandedShapeOfDim(dim);
+    expandedShape.push_back(std::accumulate(
+        dimExpansion.begin(), dimExpansion.end(), 1,
+        [](int64_t acc, int64_t i) { return std::max<int64_t>(acc * i, -1); }));
+  }
+  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,
fuseWithReshapeByExpansion(GenericOp genericOp, Operation *reshapeOp,
   if (failed(isGenericOpExpandable(genericOp, expansionInfo, rewriter)))
     return llvm::None;
 
   SmallVector<AffineMap, 4> expandedOpIndexingMaps = llvm::to_vector<4>(
       llvm::map_range(genericOp.getIndexingMaps(), [&](AffineMap m) {
         return getIndexingMapInExpandedOp(rewriter, m, expansionInfo);
       }));
 
+  // The input type to the generic can be more dynamic than the input of the
+  // reshape. Insert a cast if necessary to make them match.
+  auto maybeCast = [&](Value op, Type t) {
+    return op.getType() != t
+               ? rewriter.create<tensor::CastOp>(genericOp.getLoc(), t, op)
+               : op;
+  };
   SmallVector<Value> expandedOpOperands;
   expandedOpOperands.reserve(genericOp.getNumInputs());
   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);
       RankedTensorType expandedOperandType =
           getExpandedType(opOperand->get().getType().cast<RankedTensorType>(),
                           indexingMap, expansionInfo);
+      RankedTensorType collapsedOperandType =
+          getCollapsedType(opOperand->get().getType().cast<RankedTensorType>(),
+                           indexingMap, expansionInfo);
       if (expandedOperandType != opOperand->get().getType()) {
         // Reshape the operand to get the right type.
         SmallVector<ReassociationIndices> reassociation =
             getReassociationForExpansion(indexingMap, expansionInfo);
         expandedOpOperands.push_back(rewriter.create<tensor::ExpandShapeOp>(
-            genericOp.getLoc(), expandedOperandType, opOperand->get(),
-            reassociation));
+            genericOp.getLoc(), expandedOperandType,
+            maybeCast(opOperand->get(), collapsedOperandType), reassociation));
         continue;
       }
     }
     expandedOpOperands.push_back(opOperand->get());
   }
 
   Location loc = genericOp.getLoc();
   SmallVector<Value> outputs;
   for (OpOperand *opOperand : genericOp.getOutputOperands()) {
     AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
     RankedTensorType expandedOutputType =
         getExpandedType(opOperand->get().getType().cast<RankedTensorType>(),
                         indexingMap, expansionInfo);
+    RankedTensorType collapsedOperandType =
+        getCollapsedType(opOperand->get().getType().cast<RankedTensorType>(),
+                         indexingMap, expansionInfo);
     if (expandedOutputType != opOperand->get().getType()) {
       SmallVector<ReassociationIndices> reassociation =
           getReassociationForExpansion(indexingMap, expansionInfo);
       outputs.push_back(rewriter.create<tensor::ExpandShapeOp>(
-          genericOp.getLoc(), expandedOutputType, opOperand->get(),
-          reassociation));
+          genericOp.getLoc(), expandedOutputType,
+          maybeCast(opOperand->get(), collapsedOperandType), reassociation));
     }
   }
 
   // The iterator types of the expanded op are all parallel.
   SmallVector<StringRef> iteratorTypes(expansionInfo.getExpandedOpNumDims(),
                                        getParallelIteratorTypeName());
 
   TypeRange resultTypes = ValueRange(outputs).getTypes();

mlir/test/Dialect/Linalg/reshape_fusion.mlir

 //         FOLDUNITDIM: func @unit_dim_reshape_expansion_full
 //    FOLDUNITDIM-SAME:   %[[ARG0:.+]]: tensor<1x?x1x2x1x4xf32>
 //    FOLDUNITDIM-SAME:   %[[ARG1:.+]]: tensor<?x2x4xf32>
 //     FOLDUNITDIM-DAG:   %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG1]]
 //         FOLDUNITDIM:   linalg.generic
 //    FOLDUNITDIM-SAME:     ins(%[[ARG0]], %[[RESHAPE]] : tensor<1x?x1x2x1x4xf32>, tensor<1x?x1x2x1x4xf32>)
 //    FOLDUNITDIM-SAME:     outs(%{{.+}} : tensor<1x?x1x2x1x4xf32>)
 
+// -----
+
+func @result_less_dynamic(%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 @result_less_dynamic
+// CHECK-SAME:     %[[ARG1:[^ ]+]]: tensor<?xi64>
+//      CHECK:   %[[CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?xi64> to tensor<1xi64>
+//      CHECK:   tensor.expand_shape %[[CAST]] {{\[}}[0, 1]] : tensor<1xi64> into tensor<1x1xi64>