[mlir][Linalg] Add verification checks to disallow illegal reshape ops.

[mravishankar]

Summary

The existing verification of reshape ops in linalg (linalg.reshape and
linalg.tensor_reshape) allows specification of illegal ops, where

• A dynamic dimension is expanded into multiple dynamic dimensions. This is ill-specified.
• A static dimension is expanded into dynamic dimension or viceversa,
• The product of extents of the static dimensions in the expanded type doesnt match the static dimension of the collapsed type.

Making all of these illegal. This also implies that some pessimization
in canonicalization due to incomplete semantics of the operation can
be dropped.

[mravishankar]

Updating stale comment

[hanchung]

This patch also move "invalid" tests to roundtrip.mlir file, which is not related to patch description. Could we split it into two patches?

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

1030

Remove one the.

1031–1032

[optional] I would add two more spaces for a) and b)

1033

s/collapes/collapsed

1036

Let's use OpTy for consistency.

[mravishankar]

In D93724#2469920, @hanchung wrote:

This patch also move "invalid" tests to roundtrip.mlir file, which is not related to patch description. Could we split it into two patches?

Maybe, but my stack of CLs is already big. I would rather keep this one intact...

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

1031–1032

Its indented w.r.t to 1), ....

[mravishankar]

Rebase

[mravishankar]

Rebase

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

struct CollapseReshapeOps : public OpRewritePattern<ReshapeOpTy> {
   }
 };
 } // namespace
 
 template <typename ReshapeOpTy>
 static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
                                   ArrayRef<Attribute> operands) {
   // Fold producer-consumer reshape ops that where the operand type of the
-  // producer is same as the return type of the consumer. This can only be
-  // verified if the shapes in question are static.
+  // producer is same as the return type of the consumer.
   ReshapeOpTy reshapeSrcOp =
       reshapeOp.src().template getDefiningOp<ReshapeOpTy>();
-  if (reshapeSrcOp && reshapeSrcOp.getSrcType().hasStaticShape() &&
-      reshapeOp.getResultType().hasStaticShape() &&
-      reshapeSrcOp.getSrcType() == reshapeOp.getResultType())
+  if (reshapeSrcOp && reshapeSrcOp.getSrcType() == reshapeOp.getResultType())
     return reshapeSrcOp.src();
   // Reshape of a constant can be replaced with a new constant.
   if (auto elements = operands.front().dyn_cast_or_null<DenseElementsAttr>()) {
     return elements.reshape(
         reshapeOp.getResult().getType().template cast<ShapedType>());
   }
   return nullptr;
 }
void mlir::linalg::ReshapeOp::build(OpBuilder &b, OperationState &result,
   auto maps = getSymbolLessAffineMaps(reassociation);
   build(b, result, resultType, src, attrs);
   result.addAttribute(ReshapeOp::getReassociationAttrName(),
                       b.getAffineMapArrayAttr(maps));
 }
 
 Value mlir::linalg::ReshapeOp::getViewSource() { return src(); }
 
+/// 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 the expanded shape should be
+///    a) static
+///    b) the product should be same as the collaped shape.
+/// 2) if a dimension in the collapes 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.
+template <typename Op>
+static LogicalResult verifyReshapeLikeShapes(Op 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> dynamicDims;
+    int64_t linearizedStaticShape = 1;
+    for (auto dim : llvm::enumerate(expandedShape.slice(
+             expandedDimStart, map.value().getNumResults()))) {
+      if (ShapedType::isDynamic(dim.value())) {
+        if (isExpandingReshape && dynamicDims) {
+          return op->emitOpError("invalid to have a single dimension (")
+                 << map.index() << ") expanded into multiple dynamic dims ("
+                 << expandedDimStart + dynamicDims.getValue() << ","
+                 << expandedDimStart + dim.index() << ")";
+        }
+        dynamicDims = dim.index();
+      } else {
+        linearizedStaticShape *= dim.value();
+      }
+    }
+    if (dynamicDims) {
+      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();
+}
+
 // Common verifier for reshape-like types. Fills `expandedType` and
 // `collapsedType` with the proper `src` or `result` type.
 template <typename Op, typename T>
 static LogicalResult verifyReshapeLikeTypes(Op op, T &expandedType,
                                             T &collapsedType) {
   expandedType = op.getSrcType();
   collapsedType = op.getResultType();
   unsigned expandedRank = expandedType.getRank();
for (auto it : llvm::enumerate(maps))
     if (it.value().getNumDims() != expandedRank)
       return op.emitOpError("expected reassociation map #")
              << it.index() << " of same rank as expanded memref("
              << expandedRank << "), but got " << it.value().getNumDims();
   int invalidIdx = 0;
   if (!isReassociationValid(maps, &invalidIdx))
     return op.emitOpError("expected reassociation map #")
            << invalidIdx << " to be valid and contiguous";
-  return success();
+  return verifyReshapeLikeShapes(op, collapsedType, expandedType, !isCollapse);
 }
 
 static LogicalResult verify(ReshapeOp op) {
   MemRefType expandedType, collapsedType;
   if (failed(verifyReshapeLikeTypes(op, expandedType, collapsedType)))
     return failure();
   auto maps = getAffineMaps(op.reassociation());
   MemRefType expectedType = computeReshapeCollapsedType(expandedType, maps);
result.addAttribute(TensorReshapeOp::getReassociationAttrName(),
                       b.getAffineMapArrayAttr(maps));
 }
 
 static LogicalResult verify(TensorReshapeOp op) {
   RankedTensorType expandedType, collapsedType;
   if (failed(verifyReshapeLikeTypes(op, expandedType, collapsedType)))
     return failure();
   auto maps = getAffineMaps(op.reassociation());
-  // TODO: expanding a ? with a non-constant is under-specified. Error
-  // out.
   RankedTensorType expectedType =
       computeTensorReshapeCollapsedType(expandedType, maps);
   if (collapsedType != expectedType)
     return op.emitOpError("expected collapsed type to be ")
            << expectedType << ", but got " << collapsedType;
   return success();
 }
 

mlir/test/Dialect/Linalg/canonicalize.mlir

 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
 // CHECK-LABEL: collapsing_tensor_reshapes
 //       CHECK:   linalg.tensor_reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
 //   CHECK-NOT:   linalg.tensor_reshape
 
 // -----
 
-// -----
-
 func @collapsing_tensor_reshapes_to_zero_dim(%arg0 : tensor<1x1x1xf32>)
                                              -> tensor<f32> {
   %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1, d2) -> (d0, d1, d2)>] :
        tensor<1x1x1xf32> into tensor<1xf32>
   %1 = linalg.tensor_reshape %0 [] : tensor<1xf32> into tensor<f32>
   return %1 : tensor<f32>
 }
 // CHECK-LABEL: collapsing_tensor_reshapes_to_zero
func @collapsing_memref_reshapes_to_zero_dim(%arg0 : memref<1x1x1xf32>)
   return %1 : memref<f32>
 }
 // CHECK-LABEL: collapsing_memref_reshapes_to_zero
 //       CHECK:   linalg.reshape %{{.*}} []
 //  CHECK-SAME:     memref<1x1x1xf32> into memref<f32>
 
 // -----
 
-func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>) -> tensor<?x?x?x?x?xf32>
+func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>) -> tensor<?x6x4x?x5xf32>
 {
   %0 = linalg.tensor_reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
-       tensor<?x?xf32> into tensor<?x?x?xf32>
+       tensor<?x?xf32> into tensor<?x4x?xf32>
   %1 = linalg.tensor_reshape %0
          [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>,
           affine_map<(d0, d1, d2, d3, d4) -> (d2)>,
           affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>] :
-       tensor<?x?x?xf32> into tensor<?x?x?x?x?xf32>
-  return %1 : tensor<?x?x?x?x?xf32>
+       tensor<?x4x?xf32> into tensor<?x6x4x?x5xf32>
+  return %1 : tensor<?x6x4x?x5xf32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
 // CHECK-LABEL: expanding_tensor_reshapes
 //       CHECK:   linalg.tensor_reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
 //   CHECK-NOT:   linalg.tensor_reshape
 
 // -----

 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
 // CHECK-LABEL: collapsing_memref_reshapes
 //       CHECK:   linalg.reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
 //   CHECK-NOT:   linalg.reshape
 
 // -----
 
-func @expanding_memref_reshapes(%arg0 : memref<?x?xf32>) -> memref<?x?x?x?x?xf32>
+func @expanding_memref_reshapes(%arg0 : memref<?x?xf32>) -> memref<?x6x4x5x?xf32>
 {
   %0 = linalg.reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
-       memref<?x?xf32> into memref<?x?x?xf32>
+       memref<?x?xf32> into memref<?x4x?xf32>
   %1 = linalg.reshape %0
          [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>,
           affine_map<(d0, d1, d2, d3, d4) -> (d2)>,
           affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>] :
-       memref<?x?x?xf32> into memref<?x?x?x?x?xf32>
-  return %1 : memref<?x?x?x?x?xf32>
+       memref<?x4x?xf32> into memref<?x6x4x5x?xf32>
+  return %1 : memref<?x6x4x5x?xf32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
 // CHECK-LABEL: expanding_memref_reshapes
 //       CHECK:   linalg.reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
 //   CHECK-NOT:   linalg.reshape
 
 // -----
%1 = linalg.tensor_reshape %0
        tensor<3x4x4xf32> into tensor<12x4xf32>
   return %1 : tensor<12x4xf32>
 }
 // CHECK-LABEL: @fold_tensor_reshape
 //   CHECK-NOT:   linalg.tensor_reshape
 
 // -----
 
-func @no_fold_tensor_reshape(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+func @fold_tensor_reshape_dynamic(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
 {
   %0 = linalg.tensor_reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
-       tensor<?x?xf32> into tensor<?x?x?xf32>
+       tensor<?x?xf32> into tensor<?x4x?xf32>
   %1 = linalg.tensor_reshape %0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
-       tensor<?x?x?xf32> into tensor<?x?xf32>
+       tensor<?x4x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
-// CHECK-LABEL: @no_fold_tensor_reshape
-//       CHECK:   linalg.tensor_reshape
-//       CHECK:   linalg.tensor_reshape
+// CHECK-LABEL: @fold_tensor_reshape_dynamic
+//   CHECK-NOT:   linalg.tensor_reshape
 
 // -----
 
 func @fold_memref_reshape(%arg0 : memref<12x4xf32>) -> memref<12x4xf32>
 {
   %0 = linalg.reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
        memref<12x4xf32> into memref<3x4x4xf32>
   %1 = linalg.reshape %0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
        memref<3x4x4xf32> into memref<12x4xf32>
   return %1 : memref<12x4xf32>
 }
 // CHECK-LABEL: @fold_memref_reshape
 //   CHECK-NOT:   linalg.reshape
 
 // -----
 
-func @no_fold_memref_reshape(%arg0 : memref<?x?xf32>) -> memref<?x?xf32>
+func @fold_memref_reshape_dynamic(%arg0 : memref<?x?xf32>) -> memref<?x?xf32>
 {
   %0 = linalg.reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
-       memref<?x?xf32> into memref<?x?x?xf32>
+       memref<?x?xf32> into memref<?x4x?xf32>
   %1 = linalg.reshape %0
          [affine_map<(d0, d1, d2) -> (d0, d1)>,
           affine_map<(d0, d1, d2) -> (d2)>] :
-       memref<?x?x?xf32> into memref<?x?xf32>
+       memref<?x4x?xf32> into memref<?x?xf32>
   return %1 : memref<?x?xf32>
 }
-// CHECK-LABEL: @no_fold_memref_reshape
-//       CHECK:   linalg.reshape
-//       CHECK:   linalg.reshape
+// CHECK-LABEL: @fold_memref_reshape_dynamic
+//   CHECK-NOT:   linalg.reshape
 
 // -----
 
 #accesses = [
   affine_map<(i) -> (i)>
 ]
 
 #trait = {

mlir/test/Dialect/Linalg/invalid.mlir

 
 func @matching_inits(%m: memref<?x?xf32>, %t: tensor<?x?xf32>) {
   // expected-error @+1 {{expected type of operand #2 ('tensor<?x?xf32>') to match type of corresponding result ('tensor<?xf32>')}}
   %res = linalg.matmul ins(%m, %m : memref<?x?xf32>, memref<?x?xf32>)
                       outs(%t : tensor<?x?xf32>)
                         -> tensor<?xf32>
   return
 }
+
+
+// -----
+
+func @init_tensor_err(%arg0 : index, %arg1 : index)
+{
+  // expected-error @+1 {{specified type 'tensor<4x?x?x5xf32>' does not match the inferred type 'tensor<4x5x?x?xf32>'}}
+  %1 = linalg.init_tensor [4, 5, %arg0, %arg1] : tensor<4x?x?x5xf32>
+  return
+}
+
+// -----
+
+func @init_tensor_err(%arg0 : index)
+{
+  // expected-error @+1 {{expected 4 sizes values}}
+  %1 = linalg.init_tensor [4, 5, %arg0] : tensor<4x?x?x5xf32>
+  return
+}
+
+// -----
+
+func @init_tensor_err(%arg0 : index)
+{
+  // expected-error @+1 {{expected 2 dynamic sizes values}}
+  %1 = "linalg.init_tensor"(%arg0) {static_sizes = [4, -1, -1, 5]} : (index) -> tensor<4x?x?x5xf32>
+  return
+}
+
+// -----
+
+func @illegal_expanding_reshape_dynamic_tensor
+  (%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?x4x?xf32>
+{
+  // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
+  %0 = linalg.tensor_reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    tensor<?x?x?xf32> into tensor<?x?x?x4x?xf32>
+  return %0 : tensor<?x?x?x4x?xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_dynamic_memref
+  (%arg0: memref<?x?x?xf32>) -> memref<?x?x?x4x?xf32>
+{
+  // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
+  %0 = linalg.reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    memref<?x?x?xf32> into memref<?x?x?x4x?xf32>
+  return %0 : memref<?x?x?x4x?xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_static_tensor
+  (%arg0: tensor<2x3x20xf32>) -> tensor<2x3x2x4x5xf32>
+{
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = linalg.tensor_reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    tensor<2x3x20xf32> into tensor<2x3x2x4x5xf32>
+  return %0 : tensor<2x3x2x4x5xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_static_tensor
+  (%arg0: tensor<2x3x2x4x5xf32>) -> tensor<2x3x20xf32>
+{
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = linalg.tensor_reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    tensor<2x3x2x4x5xf32> into tensor<2x3x20xf32>
+  return %0 : tensor<2x3x20xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_static_memref
+  (%arg0: memref<2x3x20xf32>) -> memref<2x3x2x4x5xf32>
+{
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = linalg.reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    memref<2x3x20xf32> into memref<2x3x2x4x5xf32>
+  return %0 : memref<2x3x2x4x5xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_static_memref
+  (%arg0: memref<2x3x2x4x5xf32>) -> memref<2x3x20xf32>
+{
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = linalg.reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    memref<2x3x2x4x5xf32> into memref<2x3x20xf32>
+  return %0 : memref<2x3x20xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_mixed_tensor(%arg0 : tensor<?x?xf32>) -> tensor<?x4x5xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = linalg.tensor_reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0, d1)>,
+          affine_map<(d0, d1, d2) -> (d2)>] :
+       tensor<?x?xf32> into tensor<?x4x5xf32>
+  return %0 : tensor<?x4x5xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_mixed_tensor_2(%arg0 : tensor<?x?xf32>) -> tensor<?x4x5xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = linalg.tensor_reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0)>,
+          affine_map<(d0, d1, d2) -> (d1, d2)>] :
+       tensor<?x?xf32> into tensor<?x4x5xf32>
+  return %0 : tensor<?x4x5xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_mixed_tensor(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = linalg.tensor_reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0, d1)>,
+          affine_map<(d0, d1, d2) -> (d2)>] :
+       tensor<?x4x5xf32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_mixed_tensor_2(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = linalg.tensor_reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0)>,
+          affine_map<(d0, d1, d2) -> (d1, d2)>] :
+       tensor<?x4x5xf32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_mixed_memref(%arg0 : memref<?x?xf32>) -> memref<?x4x5xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = linalg.reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0, d1)>,
+          affine_map<(d0, d1, d2) -> (d2)>] :
+       memref<?x?xf32> into memref<?x4x5xf32>
+  return %0 : memref<?x4x5xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_mixed_memref_2(%arg0 : memref<?x?xf32>) -> memref<?x4x5xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = linalg.reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0)>,
+          affine_map<(d0, d1, d2) -> (d1, d2)>] :
+       memref<?x?xf32> into memref<?x4x5xf32>
+  return %0 : memref<?x4x5xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_mixed_memref(%arg0 : memref<?x4x5xf32>) -> memref<?x?xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = linalg.reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0, d1)>,
+          affine_map<(d0, d1, d2) -> (d2)>] :
+       memref<?x4x5xf32> into memref<?x?xf32>
+  return %0 : memref<?x?xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_mixed_memref_2(%arg0 : memref<?x4x5xf32>) -> memref<?x?xf32>
+{
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = linalg.reshape %arg0
+         [affine_map<(d0, d1, d2) -> (d0)>,
+          affine_map<(d0, d1, d2) -> (d1, d2)>] :
+       memref<?x4x5xf32> into memref<?x?xf32>
+  return %0 : memref<?x?xf32>
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

-// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt -split-input-file %s | FileCheck %s
 
 // TODO: Re-enable LLVM lowering test after IndexedGenericOp is lowered.
 //
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
 // DISABLED: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
 
 func @range(%arg0: index, %arg1: index, %arg2: index) {
   %0 = linalg.range %arg0:%arg1:%arg2 : !linalg.range

 func @reshape_dynamic(%arg0: memref<?x?x?xf32>,
                       %arg1: memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>,
                       %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>) {
   %0 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j)>,
                              affine_map<(i, j, k) -> (k)>] :
     memref<?x?x?xf32> into memref<?x?xf32>
   %r0 = linalg.reshape %0 [affine_map<(i, j, k) -> (i, j)>,
                            affine_map<(i, j, k) -> (k)>] :
-    memref<?x?xf32> into memref<?x?x?xf32>
+    memref<?x?xf32> into memref<?x4x?xf32>
   %1 = linalg.reshape %arg1 [affine_map<(i, j, k) -> (i, j)>,
                              affine_map<(i, j, k) -> (k)>] :
     memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]> into
     memref<?x?xf32, offset : 0, strides : [?, 1]>
   %r1 = linalg.reshape %1 [affine_map<(i, j, k) -> (i, j)>,
                            affine_map<(i, j, k) -> (k)>] :
     memref<?x?xf32, offset : 0, strides : [?, 1]> into
-    memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>
+    memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]>
   %2 = linalg.reshape %arg2 [affine_map<(i, j, k) -> (i, j)>,
                              affine_map<(i, j, k) -> (k)>] :
     memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]> into
     memref<?x?xf32, offset : ?, strides : [?, 1]>
   %r2 = linalg.reshape %2 [affine_map<(i, j, k) -> (i, j)>,
                            affine_map<(i, j, k) -> (k)>] :
     memref<?x?xf32, offset : ?, strides : [?, 1]> into
-    memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>
+    memref<?x4x?xf32, offset : ?, strides : [?, ?, 1]>
   return
 }
 
 // CHECK-DAG: #[[$reshapeD01:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 // CHECK-DAG: #[[$reshapeD2:.*]] = affine_map<(d0, d1, d2) -> (d2)>
 
 // CHECK-LABEL: func @reshape
 //       CHECK:   linalg.reshape {{.*}} [#[[$reshapeD01]], #[[$reshapeD2]]]
 //  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[$reshapeD01]], #[[$reshapeD2]]]
-//  CHECK-SAME:     memref<?x?xf32> into memref<?x?x?xf32>
+//  CHECK-SAME:     memref<?x?xf32> into memref<?x4x?xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[$reshapeD01]], #[[$reshapeD2]]]
 //  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3DOFF0]]> into memref<?x?xf32, #[[$strided2DOFF0]]>
 //       CHECK:   linalg.reshape {{.*}} [#[[$reshapeD01]], #[[$reshapeD2]]]
-//  CHECK-SAME:     memref<?x?xf32, #[[$strided2DOFF0]]> into memref<?x?x?xf32, #[[$strided3DOFF0]]>
+//  CHECK-SAME:     memref<?x?xf32, #[[$strided2DOFF0]]> into memref<?x4x?xf32, #[[$strided3DOFF0]]>
 //       CHECK:   linalg.reshape {{.*}} [#[[$reshapeD01]], #[[$reshapeD2]]]
 //  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3D]]> into memref<?x?xf32, #[[$strided2D]]>
 //       CHECK:   linalg.reshape {{.*}} [#[[$reshapeD01]], #[[$reshapeD2]]]
-//  CHECK-SAME:     memref<?x?xf32, #[[$strided2D]]> into memref<?x?x?xf32, #[[$strided3D]]>
+//  CHECK-SAME:     memref<?x?xf32, #[[$strided2D]]> into memref<?x4x?xf32, #[[$strided3D]]>
 
 func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>,
                 %ta3: tensor<?x?x?xf32>, %tb3: tensor<?x?x?xf32>, %tc3: tensor<?x?x?xf32>)
   -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
 {
   linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?x?xf32>)
                      outs(%c3: memref<?x?x?xf32>)
   linalg.batch_matmul ins(%ta3, %tb3: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
func @init_tensor(%arg0 : index, %arg1 : index)
   return
 }
 // CHECK-LABEL: func @init_tensor
 //       CHECK:   linalg.init_tensor [3, 42] : tensor<3x42xf32>
 //       CHECK:   linalg.init_tensor [4, %{{.*}}, %{{.*}}, 5] : tensor<4x?x?x5xf32>
 
 // -----
 
-func @init_tensor_err(%arg0 : index, %arg1 : index)
+func @legal_collapsing_reshape_dynamic_tensor
+  (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32>
 {
-  // expected-error @+1 {{specified type 'tensor<4x?x?x5xf32>' does not match the inferred type 'tensor<4x5x?x?xf32>'}}
-  %1 = linalg.init_tensor [4, 5, %arg0, %arg1] : tensor<4x?x?x5xf32>
-  return
-}
-
-// -----
-
-func @init_tensor_err(%arg0 : index)
-{
-  // expected-error @+1 {{expected 4 sizes values}}
-  %1 = linalg.init_tensor [4, 5, %arg0] : tensor<4x?x?x5xf32>
-  return
+  %0 = linalg.tensor_reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32>
+  return %0 : tensor<?x?x?xf32>
 }
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>
+//     CHECK: func @legal_collapsing_reshape_dynamic_tensor
+//     CHECK:   linalg.tensor_reshape %{{.+}} [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
 
 // -----
 
-func @init_tensor_err(%arg0 : index)
+func @legal_collapsing_reshape_dynamic_memref
+  (%arg0: memref<?x?x?x4x?xf32>) -> memref<?x?x?xf32>
 {
-  // expected-error @+1 {{expected 2 dynamic sizes values}}
-  %1 = "linalg.init_tensor"(%arg0) {static_sizes = [4, -1, -1, 5]} : (index) -> tensor<4x?x?x5xf32>
-  return
-}
+  %0 = linalg.reshape %arg0
+    [affine_map<(d0, d1, d2, d3, d4) -> (d0)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d1)>,
+     affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>] :
+    memref<?x?x?x4x?xf32> into memref<?x?x?xf32>
+  return %0 : memref<?x?x?xf32>
+}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>
+//     CHECK: func @legal_collapsing_reshape_dynamic_memref
+//     CHECK:   linalg.reshape %{{.+}} [#[[MAP0]], #[[MAP1]], #[[MAP2]]]