[mlir][Linalg] Add a linalg.tensor_reshape to operate on tensors

[nicolasvasilache]

Summary

This revision adds a tensor_reshape operation that operates on tensors.
In the tensor world the constraints are less stringent and we can allow more
arbitrary dynamic reshapes, as long as they are contractions.

The expansion of a dynamic dimension into multiple dynamic dimensions is under-specified and is punted on for now.

[rriddle]

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

514

nit: llvm::is_contained(...)

[bondhugula]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

79–81

This looks problematic. Why add auto in-placing / other arg dependent in-place/out-of-place semantics? Instead, always make it out of place (other passes/utilities can in-place) or have separate ops for in-place and out-of-place?

mlir/test/Dialect/Linalg/invalid.mlir

512

The error message could be confusing - they don't have to be of the same shaped type - instead 'expected source and result to both be either tensor or memref'?

[bondhugula]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

67

Update summary to include memref.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

79–81

Could we have a utility function that codifies this semantics. This is a contract that backends have to enforce AFAICS., where such a utility function would be very useful.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

@bondhugula you're right a new op makes more sense.

[mravishankar]

Looked through this with the perspective of using this for fusing linalg operation on tensors. Seems fine to me, and if something needs changing when implementing that, will send that follow up CL.

[antiagainst]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

129

Nit: I think a reassociation is represented as an affine map attribute and all the reassociations as a whole are represented as the array attribute?

131

s/sizes/dimensions/ ?

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

let builders = [OpBuilder<
     }]>];
 
   // Fully specified by traits.
   let verifier = ?;
   let assemblyFormat = "$min `:` $max `:` $step attr-dict `:` type(results)";
 }
 
 def Linalg_ReshapeOp : Linalg_Op<"reshape", [NoSideEffect]>,
-    Arguments<(ins AnyStridedMemRef:$view, AffineMapArrayAttr:$reassociation)>,
-    Results<(outs AnyStridedMemRef)> {
-  let summary = "linalg.reshape produces a new view into the operand view";
+    Arguments<(ins AnyStridedMemRefOrRankedTensor:$src,
+                   AffineMapArrayAttr:$reassociation)>,
+    Results<(outs AnyStridedMemRefOrRankedTensor)> {
+  let summary = "linalg.reshape produces a new reshaped tensor (resp. view)";
   let description = [{
-    The `linalg.reshape` op produces a new view whose sizes are a reassociation
-    of the original `view`. Depending on whether or not the reassociated
-    MemRefType is contiguous, the resulting memref may require explicit alloc
-    and copies.
+    The `linalg.reshape` op produces a new tensor (resp. view) whose sizes are a
+    reassociation of the original `src`.
 
     A reassociation is defined as a continuous grouping of dimensions and is
     represented with an affine map array attribute. In the future,
     non-continuous groupings may be allowed (i.e. permutations, reindexings
     etc).
 
-    For now, it is assumed that either:
+    In the tensor case, a reshape always produces a new tensor.
+
+    In the strided memref case, depending on whether or not the reassociated
+    MemRefType is contiguous, the resulting memref may require explicit alloc
+    and copies.
+
+    For now, in the strided memref case, it is assumed that either:
       1. a reassociation produces and consumes contiguous MemRefType or,
       2. the reshape op will be folded into its consumers (by changing the shape
          of the computations).
     All other cases are undefined behavior and a reshape op may not lower to
     LLVM if it cannot be proven statically that it does not require alloc+copy.
 
     A reshape may either collapse or expand dimensions, depending on the
     relationship between source and target memref ranks. The verification rule
     is that the reassociation maps are applied to the memref with the larger
     rank to obtain the memref with the smaller rank. In the case of a dimension
     expansion, the reassociation maps can be interpreted as inverse maps.
 
     Examples:
 
     ```mlir
     // Dimension collapse (i, j) -> i' and k -> k'
+    %b = linalg.reshape %a [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+      tensor<?x?x?xf32> into tensor<?x?xf32>
     %1 = linalg.reshape %0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
       memref<?x?x?xf32, stride_spec> into memref<?x?xf32, stride_spec_2>
     ```
 
     ```mlir
     // Dimension expansion i -> (i', j') and (k) -> (k')
+    %b = linalg.reshape %a [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+      tensor<?x?xf32> into tensor<?x?x?xf32>
     %1 = linalg.reshape %0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
       memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>
     ```
   }];
 
   let builders = [
     // Builder for a contracting reshape whose result type is computed from
-    // `view` and `reassociation`.
-    OpBuilder<"Builder *b, OperationState &result, Value view, "
+    // `src` and `reassociation`.
+    OpBuilder<"Builder *b, OperationState &result, Value src, "
     "ArrayRef<ArrayRef<AffineExpr>> reassociation, "
     "ArrayRef<NamedAttribute> attrs = {}">,
     // Builder for a reshape whose result type is passed explicitly. This may be
     // either a contracting or expanding reshape.
-    OpBuilder<"Builder *b, OperationState &result, Type resultType, Value view,"
+    OpBuilder<"Builder *b, OperationState &result, Type resultType, Value src,"
     "ArrayRef<ArrayRef<AffineExpr>> reassociation, "
     "ArrayRef<NamedAttribute> attrs = {}">];
 
   let extraClassDeclaration = [{
     static StringRef getReassociationAttrName() { return "reassociation"; }
-    MemRefType getViewType() { return view().getType().cast<MemRefType>(); }
+    ShapedType getSrcShapedType() {return src().getType().cast<ShapedType>();}
   }];
   let assemblyFormat = [{
-    $view $reassociation attr-dict `:` type($view) `into` type(results)
+    $src $reassociation attr-dict `:` type($src) `into` type(results)
   }];
   let hasFolder = 1;
 }
 
 def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view,
                    Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
     Results<(outs AnyStridedMemRef)> {

mlir/include/mlir/IR/OpBase.td

 
 def AnyStridedMemRef : StridedMemRefOf<[AnyType]>;
 
 class AnyStridedMemRefOfRank<int rank> :
   Type<And<[AnyStridedMemRef.predicate,
             MemRefRankOf<[AnyType], [rank]>.predicate]>,
        AnyStridedMemRef.description # " of rank " # rank>;
 
+def AnyStridedMemRefOrRankedTensor:
+  AnyTypeOf<[AnyStridedMemRef, AnyRankedTensor]>;
+
 // This represents a generic tuple without any constraints on element type.
 def AnyTuple : Type<IsTupleTypePred, "tuple">;
 
 // A container type that has other types embedded in it, but (unlike
 // ContainerType) can hold elements with a mix of types. Requires a call that
 // produces a list of all elements' types.
 class MixedContainerType<Type etype, Pred containerPred, code elementTypesCall,
                          string descr> :

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 //===- LinalgToLLVM.cpp - conversion from Linalg to LLVM dialect ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
explicit ReshapeOpConversion(MLIRContext *context,
                                LLVMTypeConverter &lowering_)
       : ConvertToLLVMPattern(ReshapeOp::getOperationName(), context,
                              lowering_) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto reshapeOp = cast<ReshapeOp>(op);
-    MemRefType dstType = reshapeOp.getResult().getType().cast<MemRefType>();
+    MemRefType dstType = reshapeOp.getResult().getType().dyn_cast<MemRefType>();
 
-    if (!dstType.hasStaticShape())
+    if (!dstType || !dstType.hasStaticShape())
       return failure();
 
     int64_t offset;
     SmallVector<int64_t, 4> strides;
     auto res = getStridesAndOffset(dstType, strides, offset);
     if (failed(res) || llvm::any_of(strides, [](int64_t val) {
           return ShapedType::isDynamicStrideOrOffset(val);
         }))
       return failure();
 
     edsc::ScopedContext context(rewriter, op->getLoc());
     ReshapeOpOperandAdaptor adaptor(operands);
-    BaseViewConversionHelper baseDesc(adaptor.view());
+    BaseViewConversionHelper baseDesc(adaptor.src());
     BaseViewConversionHelper desc(typeConverter.convertType(dstType));
     desc.setAllocatedPtr(baseDesc.allocatedPtr());
     desc.setAlignedPtr(baseDesc.alignedPtr());
     desc.setOffset(baseDesc.offset());
     for (auto en : llvm::enumerate(dstType.getShape()))
       desc.setConstantSize(en.index(), en.value());
     for (auto en : llvm::enumerate(strides))
       desc.setConstantStride(en.index(), en.value());

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

 //===- LinalgOps.cpp - Implementation of the linalg operations ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
else
       intStrides.push_back(ShapedType::kDynamicStrideOrOffset);
   }
   auto layout =
       makeStridedLinearLayoutMap(intStrides, intOffset, type.getContext());
   return canonicalizeStridedLayout(
       MemRefType::Builder(type).setShape(newSizes).setAffineMaps({layout}));
 }
 
+/// Compute the RankedTensorType obtained by applying `reassociation` to `type`.
+static RankedTensorType
+computeReshapeCollapsedType(RankedTensorType type,
+                            ArrayRef<AffineMap> reassociation) {
+  auto shape = type.getShape();
+  SmallVector<int64_t, 4> newShape;
+  newShape.reserve(reassociation.size());
+
+  // Use the fact that reassociation is valid to simplify the logic: only use
+  // each map's rank.
+  assert(isReassociationValid(reassociation) && "invalid reassociation");
+  unsigned currentDim = 0;
+  for (AffineMap m : reassociation) {
+    unsigned dim = m.getNumResults();
+    auto band = shape.drop_front(currentDim).take_front(dim);
+    int64_t size = 1;
+    if (llvm::any_of(band,
+                     [](int64_t v) { return v == ShapedType::kDynamicSize; }))
+      size = ShapedType::kDynamicSize;
+    else
+      for (unsigned d = 0; d < dim; ++d)
+        size *= shape[currentDim + d];
+    newShape.push_back(size);
+    currentDim += dim;
+  }
+
+  return RankedTensorType::get(newShape, type.getElementType());
+}
+
 /// Helper functions assert Attribute of the proper type in attr and returns the
 /// corresponding vector.
 /// TODO(rridle,ntv) this should be evolved into a generic
 /// `getRangeOfType<AffineMap>(ArrayAttr attrs)` that does not copy.
 static SmallVector<AffineMap, 4> getAffineMaps(ArrayAttr attrs) {
   return functional::map(
       [](Attribute a) { return a.cast<AffineMapAttr>().getValue(); }, attrs);
 }
getSymbolLessAffineMaps(ArrayRef<ArrayRef<AffineExpr>> reassociation) {
   SmallVector<AffineMap, 4> maps;
   maps.reserve(reassociation.size());
   for (auto exprs : reassociation)
     maps.push_back(AffineMap::get(maxDim + 1, 0, exprs));
   return maps;
 }
 
 void mlir::linalg::ReshapeOp::build(
-    Builder *b, OperationState &result, Value view,
+    Builder *b, OperationState &result, Value src,
     ArrayRef<ArrayRef<AffineExpr>> reassociation,
     ArrayRef<NamedAttribute> attrs) {
   auto maps = getSymbolLessAffineMaps(reassociation);
-  auto memRefType = view.getType().cast<MemRefType>();
-  auto resultType = computeReshapeCollapsedType(memRefType, maps);
-  build(b, result, resultType, view, attrs);
+  Type resultType;
+  if (auto memRefType = src.getType().dyn_cast<MemRefType>())
+    resultType = computeReshapeCollapsedType(memRefType, maps);
+  else
+    resultType = computeReshapeCollapsedType(
+        src.getType().cast<RankedTensorType>(), maps);
+  build(b, result, resultType, src, attrs);
   result.addAttribute(ReshapeOp::getReassociationAttrName(),
                       b->getAffineMapArrayAttr(maps));
 }
 
 void mlir::linalg::ReshapeOp::build(
-    Builder *b, OperationState &result, Type resultType, Value view,
+    Builder *b, OperationState &result, Type resultType, Value src,
     ArrayRef<ArrayRef<AffineExpr>> reassociation,
     ArrayRef<NamedAttribute> attrs) {
   auto maps = getSymbolLessAffineMaps(reassociation);
-  build(b, result, resultType, view, attrs);
+  build(b, result, resultType, src, attrs);
   result.addAttribute(ReshapeOp::getReassociationAttrName(),
                       b->getAffineMapArrayAttr(maps));
 }
 
 static LogicalResult verify(ReshapeOp op) {
-  MemRefType expandedType = op.getViewType();
-  MemRefType collapsedType = op.getResult().getType().cast<MemRefType>();
+  ShapedType expandedType = op.getSrcShapedType();
+  ShapedType collapsedType = op.getResult().getType().cast<ShapedType>();
   unsigned expandedRank = expandedType.getRank();
   unsigned collapsedRank = collapsedType.getRank();
   bool isCollapse = expandedRank > collapsedRank;
   if (!isCollapse) {
     std::swap(expandedRank, collapsedRank);
     std::swap(expandedType, collapsedType);
   }
   if (expandedRank == 0 || collapsedRank == 0)
     return op.emitOpError("expected non-zero memref ranks");
   if (expandedRank == collapsedRank)
     return op.emitOpError("expected to collapse or expand dims");
 
   if (collapsedRank != op.reassociation().size())
-    return op.emitOpError("expected rank of the collapsed view(")
+    return op.emitOpError("expected rank of the collapsed type(")
            << collapsedRank << ") to be the number of reassociation maps("
            << op.reassociation().size() << ")";
   auto maps = getAffineMaps(op.reassociation());
   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";
-  MemRefType expectedType = computeReshapeCollapsedType(expandedType, maps);
+  if ((expandedType.isa<MemRefType>() && !collapsedType.isa<MemRefType>()) ||
+      (expandedType.isa<RankedTensorType>() &&
+       !collapsedType.isa<RankedTensorType>()))
+    return op.emitOpError("expected source and result of same shaped type");
+
+  if (auto expandedMemRefType = expandedType.dyn_cast<MemRefType>()) {
+    MemRefType expectedType =
+        computeReshapeCollapsedType(expandedMemRefType, maps);
+    if (collapsedType != expectedType)
+      return op.emitOpError("expected collapsed type to be ")
+             << expectedType << ", but got " << collapsedType;
+  } else {
+    // TODO(ntv): expanding a ? with a non-constant is under-specified. Error
+    // out.
+    auto expandedTensorType = expandedType.dyn_cast<RankedTensorType>();
+    RankedTensorType expectedType =
+        computeReshapeCollapsedType(expandedTensorType, maps);
     if (collapsedType != expectedType)
       return op.emitOpError("expected collapsed type to be ")
              << expectedType << ", but got " << collapsedType;
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SliceOp
 //===----------------------------------------------------------------------===//
 void mlir::linalg::SliceOp::build(Builder *b, OperationState &result,
                                   Value base, ValueRange indexings) {

mlir/test/Dialect/Linalg/invalid.mlir

 func @reshape(%arg0: memref<?xf32>) {
   // expected-error @+1 {{expected to collapse or expand dims}}
   %0 = linalg.reshape %arg0 [affine_map<(i)->(i)>] : memref<?xf32> into memref<?xf32>
 }
 
 // -----
 
 func @reshape(%arg0: memref<?x?x?xf32>) {
-  // expected-error @+1 {{expected rank of the collapsed view(2) to be the number of reassociation maps(1)}}
+  // expected-error @+1 {{expected rank of the collapsed type(2) to be the number of reassociation maps(1)}}
   %0 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j)>] :
     memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
 }
 
 // -----
 
 func @reshape(%arg0: memref<?x?x?xf32>) {
   // expected-error @+1 {{expected reassociation map #0 of same rank as expanded memref(3), but got 1}}
   %0 = linalg.reshape %arg0 [affine_map<(i) -> (i)>, affine_map<(i, j, k) -> (k)>] :
     memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
 }
 
 // -----
 
 func @reshape(%arg0: memref<?x?x?xf32>) {
   // expected-error @+1 {{expected reassociation map #1 to be valid and contiguous}}
   %0 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j)>, affine_map<(i, j, k) -> (k, j)>] :
     memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
 }
 
 // -----
 
+func @reshape(%arg0: tensor<?x?x?xf32>) {
+  // expected-error @+1 {{expected source and result of same shaped type}}
+  %0 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j)>, affine_map<(i, j, k) -> (k)>] :
+    tensor<?x?x?xf32> into memref<?x?xf32>
+}
+
+// -----
+
 func @reshape(%arg0: memref<?x?x?xf32>) {
   // expected-error @+1 {{expected collapsed type to be 'memref<?x?xf32>', but got 'memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>'}}
   %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, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>
 }
 
 // -----
 
 func @pooling_rank_mismatch(%arg0: memref<?x?x?xf32>,
                             %arg1: memref<2x3xf32>,
                             %arg2: memref<?x?x?xf32>) {
   // expected-error @+1 {{expects memref ranks to match}}
   linalg.pooling_max(%arg0, %arg1, %arg2) {strides = [2, 1, 2]}:
     memref<?x?x?xf32>, memref<2x3xf32>, memref<?x?x?xf32>
   return
 }

mlir/test/Dialect/Linalg/roundtrip.mlir

 // CHECK-DAG: #[[reshapeD2:.*]] = affine_map<(d0, d1, d2) -> (d2)>
 // CHECK-DAG: #[[reshapeD0:.*]] = affine_map<(d0, d1, d2) -> (d0)>
 // CHECK-DAG: #[[reshapeD12:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>
 // CHECK-DAG: #[[reshapeD012:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-DAG: #[[reshape5D01:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>
 // CHECK-DAG: #[[reshape5D2:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2)>
 // CHECK-DAG: #[[reshape5D34:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
 
-func @reshape_static(%arg0: memref<3x4x5xf32>) {
-  // Reshapes that collapse and expand back a contiguous tensor.
+func @reshape_static(%arg0: memref<3x4x5xf32>, %arg1: tensor<3x4x5xf32>, %arg2: tensor<3x?x5xf32>) {
+  // Reshapes that collapse and expand back a contiguous buffer.
   %0 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j)>,
                              affine_map<(i, j, k) -> (k)>] :
     memref<3x4x5xf32> into memref<12x5xf32>
   %r0 = linalg.reshape %0 [affine_map<(i, j, k) -> (i, j)>,
                            affine_map<(i, j, k) -> (k)>] :
     memref<12x5xf32> into memref<3x4x5xf32>
   %1 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i)>,
                              affine_map<(i, j, k) -> (j, k)>] :
     memref<3x4x5xf32> into memref<3x20xf32>
   %r1 = linalg.reshape %1 [affine_map<(i, j, k) -> (i)>,
                            affine_map<(i, j, k) -> (j, k)>] :
     memref<3x20xf32> into memref<3x4x5xf32>
   %2 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j, k)>] :
     memref<3x4x5xf32> into memref<60xf32>
   %r2 = linalg.reshape %2 [affine_map<(i, j, k) -> (i, j, k)>] :
     memref<60xf32> into memref<3x4x5xf32>
-  // Reshapes that expand and collapse back a contiguous tensor with some 1's.
+  // Reshapes that expand and collapse back a contiguous buffer with some 1's.
   %3 = linalg.reshape %arg0 [affine_map<(i, j, k, l, m) -> (i, j)>,
                              affine_map<(i, j, k, l, m) -> (k)>,
                              affine_map<(i, j, k, l, m) -> (l, m)>] :
     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
   %r3 = linalg.reshape %3 [affine_map<(i, j, k, l, m) -> (i, j)>,
                            affine_map<(i, j, k, l, m) -> (k)>,
                            affine_map<(i, j, k, l, m) -> (l, m)>] :
     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+  // Reshapes on tensors.
+  %t0 = linalg.reshape %arg1 [affine_map<(i, j, k, l, m) -> (i, j)>,
+                              affine_map<(i, j, k, l, m) -> (k)>,
+                              affine_map<(i, j, k, l, m) -> (l, m)>] :
+    tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
+  %rt0 = linalg.reshape %t0 [affine_map<(i, j, k, l, m) -> (i, j)>,
+                            affine_map<(i, j, k, l, m) -> (k)>,
+                            affine_map<(i, j, k, l, m) -> (l, m)>] :
+    tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
+  %t1 = linalg.reshape %arg2 [affine_map<(i, j, k, l, m) -> (i, j)>,
+                              affine_map<(i, j, k, l, m) -> (k)>,
+                              affine_map<(i, j, k, l, m) -> (l, m)>] :
+    tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
+  %rt1 = linalg.reshape %t1 [affine_map<(i, j, k, l, m) -> (i)>,
+                             affine_map<(i, j, k, l, m) -> (j, k)>,
+                             affine_map<(i, j, k, l, m) -> (l, m)>] :
+    tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
   return
 }
 // CHECK-LABEL: func @reshape_static
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
 //  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD0]], #[[reshapeD12]]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD0]], #[[reshapeD12]]]
 //  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD012]]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD012]]]
 //  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshape5D01]], #[[reshape5D2]], #[[reshape5D34]]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshape5D01]], #[[reshape5D2]], #[[reshape5D34]]]
 //  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+//
+//       CHECK:   linalg.reshape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
+//       CHECK:   linalg.reshape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
+//       CHECK:   linalg.reshape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
+//       CHECK:   linalg.reshape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
 
 // -----
 
 // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[strided2DOFF0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>
 // CHECK-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 // CHECK-DAG: #[[strided3DOFF0:.*]] = affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s0 + d1 * s1 + d2)>