[mlir][memref] Fix ExpandShapeOp verifier

[springerm]

Summary

• Complete rewrite of the verifier.
• CollapseShapeOp verifier will be updated in a subsequent commit.
• Update and expand op documentation.
• Add a new builder that infers the result type based on the source type, result shape and reassociation indices. In essence, only the result layout map is inferred.

[nicolasvasilache]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1817

I find the logic unwieldy, and I can't tell offhand whether correct or not.
You should be able to rewrite with something resembling:

unsigned shapeIndex = resultShape.size() - 1;
// 1-1 mapping between srcStrides and reassociation packs.
// Each srcStride starts with the given value and gets expanded according to the proper entries in resultShape.
for (auto it : llvm::zip(llvm::reverse(reassociation), llvm::reverse(srcStrides)))
  ReassociationIndices reassoc = std::get<0>(it);
  int64_t currentStrideToExpand = std::get<1>(it);
  // Note that a stride does not get expanded along the first entry of each shape pack.
  // Example:
  //   strides = [10000, 100, 1], 
  //   reassociations = [[0], [1, 2, 3], [4]], 
  //   resultSizes = [2, 5, 4, 3, 2] = [[2], [5, 4, 3], [2]]
  // For the purpose of stride calculation, the useful sizes are: [x, x, 4, 3, x] = [[x], [x, 4, 3], [x]].
  //   resultStrides = [10000, 1200, 300, 100, 1]
  for (unsigned idx = 0, e = reassoc.size() ; idx < e; ++idx) {
    using namespace saturated_arith;
    reverseResultStrides.push_back(currentStrideToExpand);
    currentStrideToExpand = Wrapper(currentStrideToExpand) * resultShape[shapeIndex--];
  }
} 
return makeStridedLinearLayoutMap(llvm::to_vector<8>(llvm::reverse(reverseResultStrides)), srcOffset, srcType.getContext())

1853

I forget how no-layout / identity layout are both represented, please double check you have both cases working.

1925

typo

1955

Please extract the code below in a helper function that you can also reuse in the builder.
If you need the intermediate state to return different errors in the verifier/builder, consider returning an error status enum.

[springerm]

update

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1817

The saturated_path thing didn't help much because I also have to deal with dynamic dim sizes (not just strides). There are different constants for dynamic strides and dynamic dim sizes.

But good idea with the zip. This made things already much simpler.

Also added an example. Is this easier to follow now?

1853

From my understanding, "no layout map specified" = "standard layout" = "isIdentity()". Both cases are covered by unit tests.

[springerm]

typo

[mravishankar]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

This is where things become interesting. The ordering of strides might not be same as the ordering of dimensions. Affine maps (unfortunately?) allows you to specify non-inner dimension to be the fastest varying. I am OK if the decision is to say that we arent going to allow that case. But this patch seems to be making an implicit assumption about striding (unless I am mistaken), and I'd like it to be explicit.

1908

This should be logic that is common to collapse and expand for tensor and memrefs. Could we put this in a utility ?

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

You mean for the expanded dims, right? The strides of the other dims are already specified in the source memref. But for the expanded dims within a reassociation group, yes, I currently make that assumption. (I didn't think about this much, but when you don't specify any custom layout, this is what you get.)

This commit adds a builder that infers the layout map. The function that we are looking at here is doing that. The strides of the expanded dims have to be decided during bufferization. Either as we do it here or in a different way.

When it comes to the verifier, we could make it more permissive, also allowing memref result types where the expanded strides have a different ordering. That would make the logic a bit more complex though, because we could not reuse computeExpandedLayoutMap for both the builder and the verifier that easily anymore.

[nicolasvasilache]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

@mravishankar could you please elaborate a bit?

The algorithm I described does not care about stride value ordering; i.e. we could (and we should) replace the doc and test by e.g. [1, ?, 1000].

1817

Yes the fact that dynamic dims use -1 as a sentinel is an unfortunate accident of history.
The use of the -1 magic constant that has crept in very deeply.

Still, this is easily solvable in the saturated_arith case with an explicit constructor for sizes and strides and a saturated bit.
Then we would also drop the binary operators that work with int to avoid surprises:

/// Helpers to write more idiomatic operations.
namespace saturated_arith {
struct Wrapper {
  static Wrapper stride(int64_t v) {
    if (ShapedType::isDynamicStrideOrOffset(v))
      return Wrapper{true, 0}; // the value does not matter, use 0 for more defensive programming
    return Wrapper{false, v};
  }
  static Wrapper offset(int64_t v) {...} 
  static Wrapper size(int64_t v) { ... }
  int64_t asOffset() { return saturated ? ShapedType::kDynamicStrideOrOffset : v; }
  int64_t asSize() { return saturated ? ShapedType::kDynamicSize : v; }
  int64_t asStride() { return saturated ? ShapedType::kDynamicStrideOrOffset : v; }
  bool saturated;
  int64_t v;
};
Wrapper operator+(Wrapper a) {
  if (saturated || a.saturated)
    return Wrapper{true, 0};
  return Wrapper{false, a.v + b};
}
Wrapper operator*(Wrapper a) {
  if (saturated || a.saturated)
    return Wrapper{true, 0};
  return Wrapper{false, a.v * b};
}
} // namespace saturated_arith

Confining the branchy logic to this helper greatly simplifies the logic and reader's ability to follow comments.
Even in the way it is currently written with the zip + reverse, it is unclear to me that 1728-1763 is correct.

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

I think what Mahesh said is that there are multiple ways to compute the result strides. E.g., the middle dim of memref<20x10x5xf32, offset: 0, strides: [50, 5, 1]> could be expanded into:

• memref<20x2x5x5xf32, offset: 0, strides: [50, 25, 5, 1]>
• memref<20x2x5x5xf32, offset: 0, strides: [50, 5, 25, 1]>

Is this correct? I currently choose the first one.

When calling the builder that infers the layout map (which we need for bufferization), we have to make decision. This decision is not currently documented. When it comes to the verifier, we could generalize it to allow both of the above options. At the moment only the first one is allowed.

[springerm]

address comments

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1817

Added a helper function computeNextStride that contains the branching logic and is reused during the CollapseShapeOp computation.

[springerm]

simplify code

[nicolasvasilache]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

I see, we can reduce it to just 10xf32 with stride = [1] and split that into 2x5xf32.
The data looks like this: 0, 1 ... 9.

Now if we reshape to 2x5 with stride = [5, 1], the data is:

0, 1, 2, 3, 4,
5, 6, 7, 8, 9

an i, j iterator with 0 <= i < 2, 0 <= j < 5 would access data in the expected order.
In particular, (i, j) = (1, 3) accesses 8 = [1, 3] * [5, 1].

However, if we reshape to 2x5 with stride = [2, 1], the data is:

0, 1, 
2, 3, 
4, 5, 
6, 7, 
8, 9

This is incorrect because only a subset of the data would be accessed and some elements would be accessed twice.
I.e. when size > stride_increment along a given dimension, the memref has internal aliasing.
In particular, (i, j) = (1, 3) accesses 5 = [1, 3] * [2, 1].

So this internal permutation of strides is generally not correct and one should just use memref.transpose which does the right thing.

[springerm]

update

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1908

Can't be used for the tensor op verification yet because it's still in MemRefOps.cpp, but extracted it into a separate function that will be used in the CollapseShapeOp verification in the next commit.

[mravishankar]

Removing my blocker cause I was mistaken. Please add tests for non-row major striding.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

Actually my bad. I missed that the code does not make assumptions about the stride layout ordering. For an expanded dimension the stride increasing from right to left of the expanded dims is fine. I was looking for tests and didnt find any. Please do add test as Nicolas suggested.

[nicolasvasilache]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1841

Ok, I see this is now the same logic as I've written down so we can go ahead.
I'll still want to clean up and unify the saturated arith path for consistency, I can do that myself as a followup.

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

It's not clear to me yet what additional test I should add. Expansions into memrefs such as memref<20x2x5x5xf32, offset: 0, strides: [50, 5, 25, 1]> will correctly fail verification. Do we need to extend this?

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

I just did another experiment. This implementation only works if the inner-dimensions are the fastest-varying ones when it comes to strides. Otherwise, we can get faulty stride computations like here:

%r6 = memref.expand_shape %arg6 [[0], [1], [2, 3], [4]] :
    memref<20x2x10x5xf32, offset: 0, strides: [100, 10, 50, 1]>
    into memref<20x2x2x5x5xf32, offset : 0, strides : [100, 10, 250, 50, 1]>

This op passes the verifier. (All other result types will be rejected.) I think we should just limit the verifier/op to source memrefs where the strides come in the same order as the dims. And then extend the op if we actually need to support other layout maps.

[springerm]

update

[springerm]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

Added a check and expanded the op docs to disallow such memrefs. Can be extended in a follow-up revision if necessary.

[springerm]

fix warning

[nicolasvasilache]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1806

I don't see any faulty stride computation but your source memref has internal aliasing. The restriction seems artificial and based on incorrect assumptions.

I addressed in https://reviews.llvm.org/D122845.

1806

Ordering strides looks adhoc and undesirable.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     MemRef_Op<mnemonic, !listconcat(traits,
       [NoSideEffect, ViewLikeOpInterface])>,
     Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyStridedMemRef:$result)>{
   
   code commonExtraClassDeclaration = [{
     SmallVector<AffineMap, 4> getReassociationMaps();
+
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
+
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : reassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
             llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
               return indexAttr.cast<IntegerAttr>().getInt();
             })));
       return reassociationIndices;
     };
+
     MemRefType getSrcType() { return src().getType().cast<MemRefType>(); }
+
     MemRefType getResultType() { return result().getType().cast<MemRefType>(); }
+
     Value getViewSource() { return src(); }
   }];
 
   let assemblyFormat = [{
     $src $reassociation attr-dict `:` type($src) `into` type($result)
   }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape"> {
   let summary = "operation to produce a memref with a higher rank.";
   let description = [{
     The `memref.expand_shape` op produces a new view with a higher rank 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.
-
-    A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of I64ArrayAttr attribute.
-
-    For now, 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.
+    sizes are a reassociation of the original `view`. The operation is limited
+    to such reassociations, where a dimension is expanded into one or multiple
+    contiguous dimensions. Such reassociations never require additional allocs
+    or copies.
 
-    The operand memref type when dimensions can be zero-ranked if the result
-    memref type is statically shaped with all dimensions being unit extent. In
-    such case the reassociation map is empty.
-
-    The verification rule is that the reassociation maps are applied to the
-    result memref with the larger rank to obtain the operand memref with the
-    smaller rank.
+    A reassociation is defined as a grouping of dimensions and is represented
+    with an array of I64ArrayAttr attributes.
 
     Example:
 
     ```mlir
-    // Dimension expansion i -> (i', j') and (k) -> (k')
-    %1 = memref.expand_shape %0 [[0, 1], [2]] :
-      memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>
+    %r = memref.expand_shape %0 [[0, 1], [2]]
+        : memref<?x?xf32> into memref<?x5x?xf32>
     ```
+
+    At most one dimension of a reassociation group (e.g., [0, 1] above) may be
+    dynamic in the result type. Otherwise, the op would be ambiguous, as it
+    would not be clear how the source dimension is extended.
+
+    If an op can be statically proven to be invalid (e.g, an expansion from
+    `memref<10xf32>` to `memref<2x6xf32>`), it is rejected by the verifier. If
+    it cannot statically be proven invalid (e.g., the full example above; it is
+    unclear whether the first source dimension is divisible by 5), the op is
+    accepted by the verifier. However, if the op is in fact invalid at runtime,
+    the behavior is undefined.
+
+    The source memref can be zero-ranked. In that case, the reassociation
+    indices must be empty and the the result shape may only consist of unit
+    dimensions.
+
+    For simplicity, this op may not be used to cast dynamicity of dimension
+    sizes and/or strides. I.e., if and only if a source dimension is dynamic,
+    there must be a dynamic result dimension in the corresponding reassociation
+    group. Same for strides.
+
+    Note: This op currently assumes that the inner strides are of the
+    source/result layout map are the faster-varying ones.
   }];
+
   let builders = [
     // Builders using ReassociationIndices.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       build($_builder, $_state, resultType, src, attrs);
       $_state.addAttribute("reassociation",
                           getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
+
+    // Builder using ReassociationExprs.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
           convertReassociationMapsToIndices($_builder, reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
-    }]>
+    }]>,
+
+    // Builder that infers the result layout map. The result shape must be
+    // specified. Otherwise, the op may be ambiguous.
+    OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
+               "ArrayRef<ReassociationIndices>":$reassociation)>
   ];
+
   let extraClassDeclaration = commonExtraClassDeclaration;
   let hasVerifier = 1;
 }
 
 def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape"> {
   let summary = "operation to produce a memref with a smaller rank.";
   let description = [{
     The `memref.collapse_shape` op produces a new view with a smaller rank

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

OpFoldResult ReinterpretCastOp::fold(ArrayRef<Attribute> /*operands*/) {
 
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // Reassociative reshape ops
 //===----------------------------------------------------------------------===//
 
+/// Helper function that computes a stride based on the size/stride of the
+/// previous dimension.
+///
+/// E.g., memref<20x10x5xf32, offset: 0, strides: [50, 5, 1]>
+///                                                ^^
+///                                        compute this one
+///   prevStride = 5, prevDimSize = 10
+///   nextStride = 5 * 10 = 50
+static int64_t computeNextStride(int64_t prevStride, int64_t prevDimSize) {
+  if (ShapedType::isDynamicStrideOrOffset(prevStride))
+    return ShapedType::kDynamicStrideOrOffset;
+
+  if (ShapedType::isDynamic(prevDimSize))
+    return ShapedType::kDynamicStrideOrOffset;
+
+  return prevStride * prevDimSize;
+}
+
+/// Helper function for verifying the shape of ExpandShapeOp and ResultShapeOp
+/// result and operand. Layout maps are verified separately.
+///
+/// If `allowMultipleDynamicDimsPerGroup`, multiple dynamic dimensions are
+/// allowed in a reassocation group.
+static LogicalResult
+verifyCollapsedShape(Operation *op, ArrayRef<int64_t> collapsedShape,
+                     ArrayRef<int64_t> expandedShape,
+                     ArrayRef<ReassociationIndices> reassociation,
+                     bool allowMultipleDynamicDimsPerGroup) {
+  // There must be one reassociation group per collapsed dimension.
+  if (collapsedShape.size() != reassociation.size())
+    return op->emitOpError("invalid number of reassociation groups: found ")
+           << reassociation.size() << ", expected " << collapsedShape.size();
+
+  // The next expected expanded dimension index (while iterating over
+  // reassociation indices).
+  int64_t nextDim = 0;
+  for (const auto &it : llvm::enumerate(reassociation)) {
+    ReassociationIndices group = it.value();
+    int64_t collapsedDim = it.index();
+
+    bool foundDynamic = false;
+    for (int64_t expandedDim : group) {
+      if (expandedDim != nextDim++)
+        return op->emitOpError("reassociation indices must be contiguous");
+
+      if (expandedDim >= static_cast<int64_t>(expandedShape.size()))
+        return op->emitOpError("reassociation index ")
+               << expandedDim << " is out of bounds";
+
+      // Check if there are multiple dynamic dims in a reassociation group.
+      if (ShapedType::isDynamic(expandedShape[expandedDim])) {
+        if (foundDynamic && !allowMultipleDynamicDimsPerGroup)
+          return op->emitOpError(
+              "at most one dimension in a reassociation group may be dynamic");
+        foundDynamic = true;
+      }
+    }
+
+    // ExpandShapeOp/CollapseShapeOp may not be used to cast dynamicity.
+    if (ShapedType::isDynamic(collapsedShape[collapsedDim]) != foundDynamic)
+      return op->emitOpError("collapsed dim (")
+             << collapsedDim
+             << ") must be dynamic if and only if reassociation group is "
+                "dynamic";
+
+    // If all dims in the reassociation group are static, the size of the
+    // collapsed dim can be verified.
+    if (!foundDynamic) {
+      int64_t groupSize = 1;
+      for (int64_t expandedDim : group)
+        groupSize *= expandedShape[expandedDim];
+      if (groupSize != collapsedShape[collapsedDim])
+        return op->emitOpError("collapsed dim size (")
+               << collapsedShape[collapsedDim]
+               << ") must equal reassociation group size (" << groupSize << ")";
+    }
+  }
+
+  if (collapsedShape.empty()) {
+    // Rank 0: All expanded dimensions must be 1.
+    for (int64_t d : expandedShape)
+      if (d != 1)
+        return op->emitOpError(
+            "rank 0 memrefs can only be extended/collapsed with/from ones");
+  } else if (nextDim != static_cast<int64_t>(expandedShape.size())) {
+    // Rank >= 1: Number of dimensions among all reassociation groups must match
+    // the result memref rank.
+    return op->emitOpError("expanded rank (")
+           << expandedShape.size()
+           << ") inconsistent with number of reassociation indices (" << nextDim
+           << ")";
+  }
+
+  return success();
+}
+
 SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {
   return getSymbolLessAffineMaps(getReassociationExprs());
 }
+
 SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() {
   return convertReassociationIndicesToExprs(getContext(),
                                             getReassociationIndices());
 }
 
 SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() {
   return getSymbolLessAffineMaps(getReassociationExprs());
 }
+
 SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() {
   return convertReassociationIndicesToExprs(getContext(),
                                             getReassociationIndices());
 }
 
 /// Detect whether memref dims [dim, dim + extent) can be reshaped without
 /// copies.
 static bool isReshapableDimBand(unsigned dim, unsigned extent,
static LogicalResult verifyReshapeOp(ReshapeOp op, MemRefType expandedType,
   auto maps = op.getReassociationMaps();
   MemRefType expectedType = computeReshapeCollapsedType(expandedType, maps);
   if (collapsedType != expectedType)
     return op.emitOpError("expected collapsed type to be ")
            << expectedType << ", but got " << collapsedType;
   return success();
 }
 
+/// Compute the layout map after expanding a given source MemRef type with the
+/// specified reassociation indices.
+static FailureOr<AffineMap>
+computeExpandedLayoutMap(MemRefType srcType, ArrayRef<int64_t> resultShape,
+                         ArrayRef<ReassociationIndices> reassociation) {
+  SmallVector<int64_t> srcStrides, resultStrides(resultShape.size(), 0);
+  int64_t srcOffset;
+  if (failed(getStridesAndOffset(srcType, srcStrides, srcOffset)))
+    return failure();
+  assert(srcStrides.size() == reassociation.size() && "invalid reassociation");
+
+  // Ensure that inner strides are the fastest-varying ones. Other source layout
+  // maps are currently not supported.
+  int64_t lastStride = 0;
+  for (int64_t s : llvm::reverse(srcStrides)) {
+    if (!ShapedType::isDynamicStrideOrOffset(s)) {
+      if (s < lastStride)
+        return failure();
+      lastStride = s;
+    }
+  }
+
+  // Iterate over all reassociation groups from the back. Example:
+  // strides       = [1000, ?, 2]
+  // source shape  = [20,  10, 5]
+  // result shape  = [ 2, 10,   2, 5,   5]
+  // reassociation = [[0,  1], [2, 3], [4]]
+  for (const auto &it : llvm::reverse(llvm::zip(reassociation, srcStrides))) {
+    ReassociationIndices indices = std::get<0>(it);
+    int64_t srcGroupStride = std::get<1>(it);
+
+    // The first result dimension (least significant one) in each reassociation
+    // group has the same stride as the corresponding source dimension. E.g.:
+    // reassociation = [[0, 1], [2, 3], [4]]
+    //                      |       |    |
+    //                      v       v    v
+    //                    1000      ?    2
+    resultStrides[indices.pop_back_val()] = srcGroupStride;
+
+    // Compute the strides for the remaining dims in the reassociation group.
+    for (int64_t resultDim : llvm::reverse(indices)) {
+      // E.g.:
+      // reassociation = [[0, 1], [2, 3], [4]]
+      //                   |
+      //                   v
+      //               1000 * 10 = 10000
+      //
+      // If the previous stride or the previous dimension was dynamic, then this
+      // stride will also be dynamic.
+      resultStrides[resultDim] = computeNextStride(resultStrides[resultDim + 1],
+                                                   resultShape[resultDim + 1]);
+    }
+  }
+
+  return makeStridedLinearLayoutMap(resultStrides, srcOffset,
+                                    srcType.getContext());
+}
+
+static FailureOr<MemRefType>
+computeExpandedType(MemRefType srcType, ArrayRef<int64_t> resultShape,
+                    ArrayRef<ReassociationIndices> reassociation) {
+  if (srcType.getLayout().isIdentity()) {
+    // If the source is contiguous (i.e., no layout map specified), so is the
+    // result.
+    MemRefLayoutAttrInterface layout;
+    return MemRefType::get(resultShape, srcType.getElementType(), layout,
+                           srcType.getMemorySpace());
+  }
+
+  // Source may not be contiguous. Compute the layout map.
+  FailureOr<AffineMap> computedLayout =
+      computeExpandedLayoutMap(srcType, resultShape, reassociation);
+  if (failed(computedLayout))
+    return failure();
+  auto computedType =
+      MemRefType::get(resultShape, srcType.getElementType(), *computedLayout,
+                      srcType.getMemorySpaceAsInt());
+  return canonicalizeStridedLayout(computedType);
+}
+
+void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
+                          ArrayRef<int64_t> resultShape, Value src,
+                          ArrayRef<ReassociationIndices> reassociation) {
+  // Only ranked memref source values are supported.
+  auto srcType = src.getType().cast<MemRefType>();
+  FailureOr<MemRefType> resultType =
+      computeExpandedType(srcType, resultShape, reassociation);
+  // Failure of this assertion usually indicates a problem with the source
+  // type, e.g., could not get strides/offset.
+  assert(succeeded(resultType) && "could not compute layout");
+  build(builder, result, *resultType, src, reassociation);
+}
+
 LogicalResult ExpandShapeOp::verify() {
-  return verifyReshapeOp(*this, getResultType(), getSrcType());
+  MemRefType srcType = getSrcType();
+  MemRefType resultType = getResultType();
+
+  // Verify result shape.
+  if (failed(verifyCollapsedShape(getOperation(), srcType.getShape(),
+                                  resultType.getShape(),
+                                  getReassociationIndices(),
+                                  /*allowMultipleDynamicDimsPerGroup=*/false)))
+    return failure();
+
+  // Compute expected result type (including layout map).
+  FailureOr<MemRefType> expectedResultType = computeExpandedType(
+      srcType, resultType.getShape(), getReassociationIndices());
+  if (failed(expectedResultType))
+    return emitOpError("invalid source layout map");
+
+  // Check actual result type.
+  auto canonicalizedResultType = canonicalizeStridedLayout(resultType);
+  if (*expectedResultType != canonicalizedResultType)
+    return emitOpError("expected expanded type to be ")
+           << *expectedResultType << " but found " << canonicalizedResultType;
+
+  return success();
 }
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<CollapseReshapeOps<ExpandShapeOp>,
               CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>>(context);
 }
 
 LogicalResult CollapseShapeOp::verify() {
   return verifyReshapeOp(*this, getSrcType(), getResultType());
 }
 
 struct CollapseShapeOpMemRefCastFolder
     : public OpRewritePattern<CollapseShapeOp> {
Type newResultType = computeReshapeCollapsedType(
         cast.getOperand().getType().cast<MemRefType>(),
         op.getReassociationMaps());
 
     if (newResultType == op.getResultType()) {
       rewriter.updateRootInPlace(
           op, [&]() { op.srcMutable().assign(cast.source()); });
     } else {
       Value newOp = rewriter.create<CollapseShapeOp>(
           op->getLoc(), cast.source(), op.getReassociationIndices());
       rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(), newOp);
     }
     return success();
   }
 };
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {

mlir/test/Dialect/MemRef/invalid.mlir

func @copy_different_eltype(%arg0: memref<2xf32>, %arg1: memref<2xf16>) {
   // expected-error @+1 {{op requires the same element type for all operands}}
   memref.copy %arg0, %arg1 : memref<2xf32> to memref<2xf16>
   return
 }
 
 // -----
 
 func @expand_shape(%arg0: memref<f32>) {
-  // expected-error @+1 {{expected non-zero memref ranks}}
+  // expected-error @+1 {{invalid number of reassociation groups: found 1, expected 0}}
   %0 = memref.expand_shape %arg0 [[0]] : memref<f32> into memref<f32>
+  return
+}
+
+// -----
+
+func @expand_shape(%arg0: memref<f32>) {
+  // expected-error @+1 {{rank 0 memrefs can only be extended/collapsed with/from ones}}
+  %0 = memref.expand_shape %arg0 [] : memref<f32> into memref<1x2xf32>
+  return
 }
 
 // -----
 
 func @collapse_shape_to_higher_rank(%arg0: memref<f32>) {
   // expected-error @+1 {{expected the type 'memref<f32>' to have higher rank than the type = 'memref<1xf32>'}}
   %0 = memref.collapse_shape %arg0 [[0]] : memref<f32> into memref<1xf32>
 }
 
 // -----
 
 func @expand_shape_to_smaller_rank(%arg0: memref<1xf32>) {
-  // expected-error @+1 {{expected the type 'memref<f32>' to have higher rank than the type = 'memref<1xf32>'}}
+  // expected-error @+1 {{op reassociation index 0 is out of bounds}}
   %0 = memref.expand_shape %arg0 [[0]] : memref<1xf32> into memref<f32>
 }
 
 // -----
 
+func @expand_shape_invalid_result_layout(
+    %arg0: memref<30x20xf32, offset : 100, strides : [4000, 2]>) {
+  // expected-error @+1 {{expected expanded type to be 'memref<2x15x20xf32, affine_map<(d0, d1, d2) -> (d0 * 60000 + d1 * 4000 + d2 * 2 + 100)>>' but found 'memref<2x15x20xf32, affine_map<(d0, d1, d2) -> (d0 * 5000 + d1 * 4000 + d2 * 2 + 100)>>'}}
+  %0 = memref.expand_shape %arg0 [[0, 1], [2]] :
+      memref<30x20xf32, offset : 100, strides : [4000, 2]>
+      into memref<2x15x20xf32, offset : 100, strides : [5000, 4000, 2]>
+}
+
+// -----
+
 func @collapse_shape(%arg0: memref<?xf32>) {
   // expected-error @+1 {{expected to collapse or expand dims}}
   %0 = memref.collapse_shape %arg0 [[0]] : memref<?xf32> into memref<?xf32>
 }
 
 // -----
 
 func @collapse_shape_mismatch_indices_num(%arg0: memref<?x?x?xf32>) {
func @collapse_shape_wrong_collapsed_type(%arg0: memref<?x?x?xf32>) {
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
     memref<?x?x?xf32> into memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>
 }
 
 // -----
 
 func @expand_shape_illegal_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)}}
+  // expected-error @+1 {{at most one dimension in a reassociation group may be dynamic}}
   %0 = memref.expand_shape %arg0 [[0], [1], [2, 3, 4]]
       : memref<?x?x?xf32> into memref<?x?x?x4x?xf32>
   return %0 : memref<?x?x?x4x?xf32>
 }
 
 // -----
 
 func @expand_shape_illegal_static_memref
   (%arg0: memref<2x3x20xf32>) -> memref<2x3x2x4x5xf32> {
-  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  // expected-error @+1 {{collapsed dim size (20) must equal reassociation group size (40)}}
   %0 = memref.expand_shape %arg0 [[0], [1], [2, 3, 4]]
       : memref<2x3x20xf32> into memref<2x3x2x4x5xf32>
   return %0 : memref<2x3x2x4x5xf32>
 }
 
 // -----
 
 func @collapse_shape_illegal_static_memref
   (%arg0: memref<2x3x2x4x5xf32>) -> memref<2x3x20xf32> {
   // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
   %0 = memref.collapse_shape %arg0 [[0], [1], [2, 3, 4]]
       : memref<2x3x2x4x5xf32> into memref<2x3x20xf32>
   return %0 : memref<2x3x20xf32>
 }
 
 // -----
 
 func @expand_shape_illegal_mixed_memref(%arg0 : memref<?x?xf32>)
     -> memref<?x4x5xf32> {
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  // expected-error @+1 {{collapsed dim (1) must be dynamic if and only if reassociation group is dynamic}}
   %0 = memref.expand_shape %arg0 [[0, 1], [2]]
       : memref<?x?xf32> into memref<?x4x5xf32>
   return %0 : memref<?x4x5xf32>
 }
 
 // -----
 
 func @expand_shape_illegal_mixed_memref_2(%arg0 : memref<?x?xf32>)
     -> memref<?x4x5xf32> {
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  // expected-error @+1 {{collapsed dim (1) must be dynamic if and only if reassociation group is dynamic}}
   %0 = memref.expand_shape %arg0 [[0], [1, 2]]
       : memref<?x?xf32> into memref<?x4x5xf32>
   return %0 : memref<?x4x5xf32>
 }
 
 // -----
 
+func @expand_shape_unsupported_src_layout(
+    %arg0 : memref<20x2x10x5xf32, offset: 0, strides: [100, 10, 50, 1]>)
+    -> memref<20x2x2x5x5xf32, offset : 0, strides : [100, 10, 250, 50, 1]> {
+  // expected-error @+1 {{invalid source layout map}}
+  %0 = memref.expand_shape %arg0 [[0], [1], [2, 3], [4]] :
+      memref<20x2x10x5xf32, offset: 0, strides: [100, 10, 50, 1]>
+      into memref<20x2x2x5x5xf32, offset : 0, strides : [100, 10, 250, 50, 1]>
+  return %0 : memref<20x2x2x5x5xf32, offset : 0, strides : [100, 10, 250, 50, 1]>
+}
+
+// -----
+
+func @expand_shape_invalid_static_dim_size(%arg0 : memref<?x21xf32>)
+    -> memref<?x4x5xf32> {
+  // expected-error @+1 {{collapsed dim size (21) must equal reassociation group size (20)}}
+  %0 = memref.expand_shape %arg0 [[0], [1, 2]]
+      : memref<?x21xf32> into memref<?x4x5xf32>
+  return %0 : memref<?x4x5xf32>
+}
+
+// -----
+
 func @collapse_shape_illegal_mixed_memref(%arg0 : memref<?x4x5xf32>)
     -> memref<?x?xf32> {
   // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]]
       : memref<?x4x5xf32> into memref<?x?xf32>
   return %0 : memref<?x?xf32>
 }
 

mlir/test/Dialect/MemRef/ops.mlir

 // CHECK-LABEL: func @memref_alloca_scope
 func @memref_alloca_scope() {
   memref.alloca_scope {
     memref.alloca_scope.return
   }
   return
 }
 
-func @expand_collapse_shape_static(%arg0: memref<3x4x5xf32>,
-                                   %arg1: tensor<3x4x5xf32>,
-                                   %arg2: tensor<3x?x5xf32>) {
+// CHECK-LABEL: func @expand_collapse_shape_static
+func @expand_collapse_shape_static(
+    %arg0: memref<3x4x5xf32>,
+    %arg2: tensor<3x?x5xf32>,
+    %arg3: memref<30x20xf32, offset : 100, strides : [4000, 2]>,
+    %arg4: memref<1x5xf32, affine_map<(d0, d1)[s0] -> (d0 * 5 + s0 + d1)>>,
+    %arg5: memref<f32>) {
   // Reshapes that collapse and expand back a contiguous buffer.
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
     memref<3x4x5xf32> into memref<12x5xf32>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
   %r0 = memref.expand_shape %0 [[0, 1], [2]] :
     memref<12x5xf32> into memref<3x4x5xf32>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0], [1, 2]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
   %1 = memref.collapse_shape %arg0 [[0], [1, 2]] :
     memref<3x4x5xf32> into memref<3x20xf32>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0], [1, 2]]
+//  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
   %r1 = memref.expand_shape %1 [[0], [1, 2]] :
     memref<3x20xf32> into memref<3x4x5xf32>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1, 2]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
   %2 = memref.collapse_shape %arg0 [[0, 1, 2]] :
     memref<3x4x5xf32> into memref<60xf32>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1, 2]]
+//  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
   %r2 = memref.expand_shape %2 [[0, 1, 2]] :
       memref<60xf32> into memref<3x4x5xf32>
+
+//       CHECK:   memref.expand_shape {{.*}} []
+//  CHECK-SAME:     memref<f32> into memref<1x1xf32>
+  %r5 = memref.expand_shape %arg5 [] :
+      memref<f32> into memref<1x1xf32>
+
+// Reshapes with a custom layout map.
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0], [1, 2]]
+  %l0 = memref.expand_shape %arg3 [[0], [1, 2]] :
+      memref<30x20xf32, offset : 100, strides : [4000, 2]>
+      into memref<30x4x5xf32, offset : 100, strides : [4000, 10, 2]>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+  %l1 = memref.expand_shape %arg3 [[0, 1], [2]] :
+      memref<30x20xf32, offset : 100, strides : [4000, 2]>
+      into memref<2x15x20xf32, offset : 100, strides : [60000, 4000, 2]>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0], [1, 2]]
+  %r4 = memref.expand_shape %arg4 [[0], [1, 2]] :
+      memref<1x5xf32, affine_map<(d0, d1)[s0] -> (d0 * 5 + s0 + d1)>> into
+      memref<1x1x5xf32, affine_map<(d0, d1, d2)[s0] -> (d0 * 5 + s0 + d2 + d1 * 5)>>
+
   // Reshapes that expand and collapse back a contiguous buffer with some 1's.
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
   %3 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]] :
     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
+//  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
   %r3 = memref.collapse_shape %3 [[0, 1], [2], [3, 4]] :
     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+
   // Reshapes on tensors.
+//       CHECK:   tensor.expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
   %t0 = tensor.expand_shape %arg1 [[0, 1], [2], [3, 4]] :
     tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
+
+//       CHECK:   tensor.collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
   %rt0 = tensor.collapse_shape %t0 [[0, 1], [2], [3, 4]] :
     tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
+
+//       CHECK:   tensor.expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
   %t1 = tensor.expand_shape %arg2 [[0, 1], [2], [3, 4]] :
     tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
+
+//       CHECK:   tensor.collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
   %rt1 = tensor.collapse_shape %t1 [[0], [1, 2], [3, 4]] :
     tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
   return
 }
-// CHECK-LABEL: func @expand_collapse_shape_static
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0], [1, 2]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0], [1, 2]]
-//  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1, 2]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1, 2]]
-//  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
-//  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
-//
-//       CHECK:   tensor.expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
-//       CHECK:   tensor.collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
-//       CHECK:   tensor.expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
-//       CHECK:   tensor.collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
-
 
+// CHECK-LABEL: func @expand_collapse_shape_dynamic
 func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
          %arg1: memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>,
          %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>,
          %arg3: memref<?x42xf32, offset : 0, strides : [42, 1]>) {
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
     memref<?x?x?xf32> into memref<?x?xf32>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?xf32> into memref<?x4x?xf32>
   %r0 = memref.expand_shape %0 [[0, 1], [2]] :
     memref<?x?xf32> into memref<?x4x?xf32>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3DOFF0]]> into memref<?x?xf32, #[[$strided2DOFF0]]>
   %1 = memref.collapse_shape %arg1 [[0, 1], [2]] :
     memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]> into
     memref<?x?xf32, offset : 0, strides : [?, 1]>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?xf32, #[[$strided2DOFF0]]> into memref<?x4x?xf32, #[[$strided3DOFF0]]>
   %r1 = memref.expand_shape %1 [[0, 1], [2]] :
     memref<?x?xf32, offset : 0, strides : [?, 1]> into
     memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3D]]> into memref<?x?xf32, #[[$strided2D]]>
   %2 = memref.collapse_shape %arg2 [[0, 1], [2]] :
     memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]> into
     memref<?x?xf32, offset : ?, strides : [?, 1]>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?xf32, #[[$strided2D]]> into memref<?x4x?xf32, #[[$strided3D]]>
   %r2 = memref.expand_shape %2 [[0, 1], [2]] :
     memref<?x?xf32, offset : ?, strides : [?, 1]> into
     memref<?x4x?xf32, offset : ?, strides : [?, ?, 1]>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1]]
+//  CHECK-SAME:     memref<?x42xf32, #[[$strided2D42]]> into memref<?xf32>
   %3 = memref.collapse_shape %arg3 [[0, 1]] :
     memref<?x42xf32, offset : 0, strides : [42, 1]> into
     memref<?xf32, offset : 0, strides : [1]>
+
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1]]
+//  CHECK-SAME:     memref<?xf32> into memref<?x42xf32>
   %r3 = memref.expand_shape %3 [[0, 1]] :
-    memref<?xf32, offset : 0, strides : [1]> into
-    memref<?x42xf32, offset : 0, strides : [42, 1]>
+    memref<?xf32, offset : 0, strides : [1]> into memref<?x42xf32>
   return
 }
-// CHECK-LABEL: func @expand_collapse_shape_dynamic
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?xf32> into memref<?x4x?xf32>
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3DOFF0]]> into memref<?x?xf32, #[[$strided2DOFF0]]>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?xf32, #[[$strided2DOFF0]]> into memref<?x4x?xf32, #[[$strided3DOFF0]]>
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3D]]> into memref<?x?xf32, #[[$strided2D]]>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?xf32, #[[$strided2D]]> into memref<?x4x?xf32, #[[$strided3D]]>
-//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1]]
-//  CHECK-SAME:     memref<?x42xf32, #[[$strided2D42]]> into memref<?xf32>
-//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1]]
-//  CHECK-SAME:     memref<?xf32> into memref<?x42xf32, #[[$strided2D42]]>
 
 func @expand_collapse_shape_zero_dim(%arg0 : memref<1x1xf32>, %arg1 : memref<f32>)
     -> (memref<f32>, memref<1x1xf32>) {
   %0 = memref.collapse_shape %arg0 [] : memref<1x1xf32> into memref<f32>
   %1 = memref.expand_shape %0 [] : memref<f32> into memref<1x1xf32>
   return %0, %1 : memref<f32>, memref<1x1xf32>
 }
 // CHECK-LABEL: func @expand_collapse_shape_zero_dim