[MLIR][LINALG] Add canonicalization pattern in `linalg.generic` op for static shape inference.

[gprateek93]

Summary

This commit adds canonicalization pattern in linalg.generic op
for static shape inference. If any of the inputs or outputs have
static shape or is casted from a tensor of static shape, then
shapes of all the inputs and outputs can be inferred by using the
affine map of the static shape input/output.

Signed-Off-By: Prateek Gupta <[e-mail]>

[mravishankar]

Thanks for helping push on this. I don think we can simplify the implementation here by a lot.

@nicolasvasilache might have better ways of doing this for arbitrary affine exprs in indexing maps. In general you need to solve a set of linear equations to compute static dimensions. There might be a way to formulate that in affine map "arithmetic". That seems very heavy wieght (and I havent worked out how to express than using affine map computations). I am thinking of solving a subset of a common case when affine maps are all projected permutations.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

849

General note. I prefer not to use lambdas this way. Its much more readable to have a static method outside of the function.

853

Do you need to pass a copy of SmallVector<OpOperand *> here? Can you use an ArrayRef?

855

If you are going to use a lambda, I would put the loop outside of the lambda and simplify the signature of the lambda.

861

I'd refactor this a bit. First get the shape value

• If operand is result of a tensor.cast get the shape from source of the tensor.cast. If not use the shape of the operand.
• Then iterate over the shape you get in step 1 and check if individual dimensions are static or dynamic. If static, then add it to sourceMap.

866

It is not necessary that the result be an AffineDimExpr. You need to check that the it is and then bail if it is not.

876

I think the else part can be dropped after the above suggestion.

903

Same comment here as above. Move loop out of the lambda.

[gprateek93]

Refactored the code.

[gprateek93]

Added check for the operands to be RankedTensorType.

[mravishankar]

Thanks for the updates. I have one more round of comments. I think after that this should be ready to land.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

841

The startIdx part is confusing to me in the way this is written. There seems to be some handshake between source and operands here. Maybe this could work,

static void populateMap(GenericOp genericOp, ArrayRef<OpOperand *> operands,
                        ValueRange source,
                        llvm::DenseMap<AffineExpr, int64_t> &affineExprToSize) {
  for (auto it : llvm::zip(operands, source)) {
    OpOperand *opOperand = std::get<0>(it);
    Value src = std::get<1>(it);
    ...
  }
  ...
}

In any case, looking at the uses, do you even need to use source. Looks like operands is all you need.

844

Could you add an assert here to check sourceMap.isProjectedPermutation() ? I am assuming the caller verifies that the map is a projected permutation and doesnt call this function if that is the case.

857

nit: LLVM style is to not have { ..} for 1-line scopes.

864

Please do not use -1. Use ShapedType::isDynamicDim instead.

881

Same as above, you probably dont need source and startIdx. Also please move affineExprToSize "before" the parameter that is the return by reference (resultTypeVector I guess). I think there is a simpler way to do this. Ill leave some comments below where this function is called.

903

Please dont use -1 explicitly. Use ShapedType::isDynamicDim instead.

944

You can replace lines 1007 - 1024 with

if (!genericOp.hasTensorSemantics())
  return failure();

945

You also need to check that all maps of the genericOp are projected permutations. Something like

if (llvm::any_of(genericOp.getIndexingMaps(), [](AffineMap map) { return !map.isProjectedPermutation(); }))
  return failure();

969

I think both of these calls can be combined into a single call to populateMap.

979

I think this split of updateOperands is unnecessary and confusing.

For example for the inputs newResultTypes is unnecessary. We can have separation of concerns and do it in two steps.

SmallVector<Value> newOperands;
SmallVector<Type> resultTypes;
bool noChangeNeeded = true;
newOperands.reserve(genericOp.getNumInputsAndOutputs());
resultTypes.reserve(genericOp.getNumOutputs());
for (OpOperand *operand : genericOp.getInputAndOutputOpOperands()) {
  // Check for type of operand changing and insert casts if needed
  noChangeNeeded = ... // set to true if a change is needed.
  newOperands.push_back(newOperand);
  if (genericOp->isOutputTensor(arg)) {
    resultType.push_back(newOperand.getType());
  }
}
if (noChangeNeeded) return failure();
auto *newOp = ... // Create new generic op.
SmallVector<Value> replacements;
replacements.reserve(newOp->getNumResults());
for (auto it : llvm::zip(genericOp->getResults(), newOp->getResults())) {
   if (std::get<0>(it).getType() != std::get<1>(it).getType()) {
     // Insert cast.
     replacements.push_back(castResult);
   } else {
     replacements.push_back(std::get<1>(it));
   }
}
rewriter.replaceOp(genericOp, replacements);

991

You can just use cast<linalg::LinalgOp>(genericOp.getOperation()).clone(...) from here https://github.com/llvm/llvm-project/blob/be1cca1b0d7ce01dff2a3cf97438f437576a1fa0/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L959

[gprateek93]

Refactored the code.

[mravishankar]

Sorry, one more use of -1 and also a if-then-else style nit. Otherwise looks good!

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

912

Nit: I think LLVM style is to do

for (...) {
  if (...) {
     ...
     continue;
  }

  // previously else-if part.
  if (...) {
    ....
    continue;
  }

  // previously else part
  ...
}

922

Please dont use the -1 directly. Used ShapedType::isDynamicSize or equivalent. In any case, I dont think you need this check here.

[gprateek93]

Refactored the code and updated one of the reshape-fusion test case.

[bondhugula]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

897–950

This whole logic is too long (about 100 lines), with too much nesting, appears messy with generic boolean variables like changeNeeded. A couple of suggestions:

1. document changeNeeded above,
2. refactor it (for eg. the part iterating over operands) into an appropriately named method.

922

Please use is_contained instead of comparing with end().

mlir/test/Dialect/Linalg/canonicalize.mlir

655

Nit: Add a comment on what these are testing -- the function names themselves aren't descriptive enough.

[gprateek93]

Refactored code.

[gprateek93]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

922

I am not sure how is_contained can be used with map iterators. Can you please elaborate? Thanks!

[bondhugula]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

922

I see - find() is fine.

[gprateek93]

Rebasing into main.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

LogicalResult matchAndRewrite(GenericOp genericOp,
     }
 
     if (returnedArgs.size() != genericOp->getNumResults())
       return failure();
     rewriter.replaceOp(genericOp, returnedArgs);
     return success();
   }
 };
+
+/// For each of the operand in `operands` this function maps the static sizes of
+/// dimensions to their affine dim expressions.
+static void populateMap(GenericOp genericOp, ArrayRef<OpOperand *> operands,
+                        llvm::DenseMap<AffineExpr, int64_t> &affineExprToSize) {
+  for (OpOperand *opOperand : operands) {
+    if (genericOp.isScalar(opOperand))
+      continue;
+    Value src = opOperand->get();
+    auto sourceType = src.getType().cast<RankedTensorType>();
+    auto sourceMap = genericOp.getTiedIndexingMap(opOperand);
+
+    // Get the `sourceShape` of the `sourceType`. If the operand is a result of
+    // `tensor.cast` operation and source of the cast operation has a static
+    // shape, then assign it to the `sourceShape`.
+    auto parentOp = src.getDefiningOp();
+    ArrayRef<int64_t> sourceShape = sourceType.getShape();
+    if (parentOp) {
+      if (auto castOp = dyn_cast<tensor::CastOp>(parentOp)) {
+        Value castSource = castOp.source();
+        auto castSourceType = castSource.getType().cast<RankedTensorType>();
+        if (castSourceType.hasStaticShape())
+          sourceShape = castSourceType.getShape();
+      }
+    }
+
+    // If the source shape's dimension has a static shape, map the affine dim
+    // expression to the known static size.
+    for (unsigned i = 0; i < sourceShape.size(); i++) {
+      if (sourceType.isDynamicDim(i))
+        continue;
+      if (auto affineDimExpr = sourceMap.getResult(i).dyn_cast<AffineDimExpr>())
+        affineExprToSize.try_emplace(affineDimExpr, sourceShape[i]);
+    }
+  }
+}
+
+/// Static shapes for the operands can be inferred if any one of the operands
+/// have a static shape. This can be done by referring to the affine dim
+/// expressions for the operand.
+struct InferStaticShapeOfOperands : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    if (!genericOp.hasTensorSemantics())
+      return failure();
+
+    // Maps must be projected permutations.
+    if (llvm::any_of(genericOp.getIndexingMaps(), [](AffineMap map) {
+          return !map.isProjectedPermutation();
+        }))
+      return failure();
+
+    // Maps affine dim expressions to the static size of that dimension.
+    llvm::DenseMap<AffineExpr, int64_t> affineExprToSize;
+    Location loc = genericOp.getLoc();
+
+    // For each of the affine dim expression, check if the size is known. If
+    // known add that in the map.
+    populateMap(genericOp, genericOp.getInputAndOutputOperands(),
+                affineExprToSize);
+
+    SmallVector<Value> newOperands;
+    SmallVector<Type> resultTypes;
+    bool changeNeeded = false;
+    newOperands.reserve(genericOp.getNumInputsAndOutputs());
+    resultTypes.reserve(genericOp.getNumOutputs());
+
+    // Iterate over all the operands and update the static sizes.
+    for (OpOperand *opOperand : genericOp.getInputAndOutputOperands()) {
+      Value src = opOperand->get();
+      newOperands.push_back(src);
+      if (genericOp.isScalar(opOperand))
+        continue;
+      auto sourceType = src.getType().cast<RankedTensorType>();
+      Type resultType = sourceType;
+      if (sourceType.hasStaticShape() && genericOp.isOutputTensor(opOperand)) {
+        resultTypes.push_back(resultType);
+        continue;
+      }
+      ArrayRef<int64_t> sourceShape = sourceType.getShape();
+      AffineMap sourceMap = genericOp.getTiedIndexingMap(opOperand);
+      SmallVector<int64_t> newShape;
+      // If operand is updated with new shape, `newOperandNeeded` will be
+      // true.
+      bool newOperandNeeded = false;
+      for (unsigned i = 0; i < sourceShape.size(); i++) {
+        int64_t dimShape = sourceShape[i];
+        AffineExpr dimExpr = sourceMap.getResult(i);
+        if (affineExprToSize.find(dimExpr) == affineExprToSize.end() ||
+            !sourceType.isDynamicDim(i)) {
+          newShape.push_back(dimShape);
+          continue;
+        }
+        // Dimension has a dynamic shape and corresponding affine dim
+        // expression is present in the map. So assign the size for the
+        // given affine dim expression to the dimension.
+        newShape.push_back(affineExprToSize[dimExpr]);
+        newOperandNeeded = true;
+      }
+      resultType = RankedTensorType::get(newShape, sourceType.getElementType());
+      if (newOperandNeeded) {
+        changeNeeded = true;
+        // Get the new operand value given its size and element type by
+        // casting it.
+        Value newOperand =
+            rewriter.create<tensor::CastOp>(loc, resultType, src);
+        unsigned index = opOperand->getOperandNumber();
+        newOperands[index] = newOperand;
+      }
+      if (genericOp.isOutputTensor(opOperand))
+        resultTypes.push_back(resultType);
+    }
+
+    // If the generic op has all the required static information, no
+    // canonicalization needed.
+    if (!changeNeeded)
+      return failure();
+
+    // Clone op.
+    Operation *newOp =
+        cast<linalg::LinalgOp>(genericOp.getOperation())
+            .clone(rewriter, genericOp->getLoc(), resultTypes, newOperands);
+    SmallVector<Value> replacements;
+    replacements.reserve(newOp->getNumResults());
+    for (auto it : llvm::zip(genericOp->getResults(), newOp->getResults())) {
+      Value newResult = std::get<1>(it);
+      Value oldResult = std::get<0>(it);
+      Type newType = newResult.getType();
+      Type oldType = oldResult.getType();
+      replacements.push_back(
+          (newType != oldType)
+              ? rewriter.create<tensor::CastOp>(loc, newType, newResult)
+              : newResult);
+    }
+    rewriter.replaceOp(genericOp, replacements);
+    return success();
+  }
+};
 } // namespace
 
 void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<DeduplicateGenericOpInputs, EraseIdentityGenericOp>(context);
+  results.add<DeduplicateGenericOpInputs, EraseIdentityGenericOp,
+              InferStaticShapeOfOperands>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // InitTensorOp
 //===----------------------------------------------------------------------===//
 
 void InitTensorOp::build(OpBuilder &b, OperationState &result,
                          ArrayRef<OpFoldResult> sizes, Type elementType,
                          ArrayRef<NamedAttribute> attrs) {
   SmallVector<Value, 4> dynamicSizes;
   SmallVector<int64_t, 4> staticSizes;
   dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
                              ShapedType::kDynamicSize);
   auto resultType = RankedTensorType ::get(staticSizes, elementType);
   build(b, result, resultType, dynamicSizes, b.getI64ArrayAttr(staticSizes));
   result.addAttributes(attrs);
 }
 
 LogicalResult InitTensorOp::verify() {
   RankedTensorType resultType = getType();
   SmallVector<int64_t, 4> staticSizes = llvm::to_vector<4>(llvm::map_range(
       static_sizes().cast<ArrayAttr>(),

mlir/test/Dialect/Linalg/canonicalize.mlir

func @no_fold_pad_fill_value_mismatch() -> tensor<412x276xf32> {
   %fill = linalg.fill(%f0, %init) : f32, tensor<400x273xf32> -> tensor<400x273xf32>
   // CHECK: tensor.pad
   %pad = tensor.pad %fill low[4, 1] high[8, 2] {
   ^bb0(%arg1: index, %arg2: index):
     tensor.yield %f1 : f32
   } : tensor<400x273xf32> to tensor<412x276xf32>
   return %pad : tensor<412x276xf32>
 }
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @static_input_without_cast
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+func @static_input_without_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<?x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+    %9 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<?x?x?xf32>)
+  %5 = tensor.cast %4 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %5 : tensor<2x3x4xf32>
+    //  CHECK:      %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @static_input_with_cast
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+func @static_input_with_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32>
+  %5 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %4 : tensor<2x3x4xf32>, tensor<2x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+    %9 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<?x?x?xf32>)
+  %6 = tensor.cast %5 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %6: tensor<2x3x4xf32>
+    //  CHECK:      %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @static_output_with_cast
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<?x?x?xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>, %[[ARG2:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+func @static_output_with_cast(%arg0 : tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>, %arg2: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg2, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg2, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg2, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = tensor.cast %3 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  %5 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32>
+  %6 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %5 : tensor<?x?x?xf32>, tensor<2x?x?xf32>)
+    outs(%4 : tensor<2x3x4xf32>) {
+  ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
+    %9 = arith.addf %arg3, %arg4 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<2x3x4xf32>)
+  return %6: tensor<2x3x4xf32>
+    //  CHECK:      %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[CAST_ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @cast_source
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+func @cast_source(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = tensor.cast %arg0 : tensor<2x3x4xf32> to tensor<2x?x?xf32>
+  %5 = tensor.cast %arg1 : tensor<2x3x4xf32> to tensor<2x?x?xf32>
+  %6 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%4, %5 : tensor<2x?x?xf32>, tensor<2x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+    %9 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<?x?x?xf32>)
+  %7 = tensor.cast %6 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %7: tensor<2x3x4xf32>
+    //  CHECK:      %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}

mlir/test/Dialect/Linalg/reshape_fusion.mlir

 // 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>
+func @no_fuse_mismatched_dynamism(%arg0: tensor<2x1xi64>, %arg1: tensor<?xi64>) -> tensor<2xi64> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<2x1xi64> into tensor<2xi64>
+  %1 = linalg.init_tensor [2] : tensor<2xi64>
   %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>) {
+    ins(%0, %arg1 : tensor<2xi64>, tensor<?xi64>)
+    outs(%1 : tensor<2xi64>) {
   ^bb0(%arg4: i64, %arg5: i64, %arg6: i64):  
     %3 = arith.addi %arg4, %arg5 : i64
     linalg.yield %3 : i64
-  } -> tensor<1xi64>
-  return %2 : tensor<1xi64>
+  } -> tensor<2xi64>
+  return %2 : tensor<2xi64>
 }
 
 //      CHECK: func @no_fuse_mismatched_dynamism
-// CHECK-SAME:     %[[ARG0:.+]]: tensor<1x1xi64>
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<2x1xi64>
 // CHECK-SAME:     %[[ARG1:.+]]: tensor<?xi64>
 //      CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
+//      CHECK:   %[[CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?xi64> to tensor<2xi64>
 //      CHECK:   %[[GENERIC:.+]] = linalg.generic
-// CHECK-SAME:       ins(%[[RESHAPE]], %[[ARG1]] : tensor<1xi64>, tensor<?xi64>)
+// CHECK-SAME:       ins(%[[RESHAPE]], %[[CAST]] : tensor<2xi64>, tensor<2xi64>)
 //      CHECK:   return %[[GENERIC]]