[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

1739

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())

1775

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

1839

typo

1869

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

1739

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?

1775

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

1728

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.

1822

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

1728

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

1728

@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].

1739

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

1728

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

1739

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

1728

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

1822

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

1728

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

1763

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

1728

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

1728

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

1728

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

1728

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.

1728

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.
+    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.
 
-    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.
-
-    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.
   }];
+
   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;
+}
+
 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");
+
+  // Iterate over all reassociations group 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]]
+  int64_t resultDim = resultShape.size();
+  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[--resultDim] = srcGroupStride;
+
+    // Compute the strides for the remaining dims in the reassociation group.
+    for (int64_t i = indices.size() - 2; i >= 0; --i) {
+      --resultDim;
+
+      // 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();
+
+  // There must be one reassociation group per source dimension.
+  if (srcType.getRank() != getReassociationIndices().size())
+    return emitOpError("invalid number of reassociation groups: found ")
+           << getReassociationIndices().size() << ", expected "
+           << srcType.getRank();
+
+  // The next expected dimension (while iterating over reassociation indices)
+  int64_t nextDim = 0;
+  for (const auto &it : llvm::enumerate(getReassociationIndices())) {
+    ReassociationIndices group = it.value();
+    int64_t srcDim = it.index();
+
+    bool foundDynamic = false;
+    for (int64_t resultDim : group) {
+      if (resultDim != nextDim++)
+        return emitOpError("reassociation indices must be contiguous");
+
+      if (resultDim >= resultType.getRank())
+        return emitOpError("reassociation index ")
+               << resultDim << " is out of bounds";
+
+      // There may be at most dynamic result dim in a reassociation group.
+      if (resultType.isDynamicDim(resultDim)) {
+        if (foundDynamic)
+          return emitOpError(
+              "at most one dimension in a reassociation group may be dynamic");
+        foundDynamic = true;
+      }
+    }
+
+    if (srcType.isDynamicDim(srcDim) != foundDynamic)
+      return emitOpError(
+          "result dim must be dynamic if and only if source dim is dynamic");
+
+    if (!foundDynamic) {
+      int64_t groupSize = 1;
+      for (int64_t resultDim : group)
+        groupSize *= resultType.getDimSize(resultDim);
+      if (groupSize != srcType.getDimSize(srcDim))
+        return emitOpError("source dim must be a product of expanded dims");
+    }
+  }
+
+  if (srcType.getRank() == 0) {
+    // Rank 0: All result dimensions must be 1.
+    for (int64_t d : resultType.getShape())
+      if (d != 1)
+        return emitOpError("rank 0 memrefs can only be extended with ones");
+  } else if (nextDim != resultType.getRank()) {
+    // Rank >= 1: Number of dimensions among all reassociation groups must match
+    // the result memref rank.
+    return emitOpError("result type inconsistent with reassociation");
+  }
+
+  FailureOr<MemRefType> expectedResultType = computeExpandedType(
+      srcType, resultType.getShape(), getReassociationIndices());
+  if (failed(expectedResultType))
+    return emitOpError("invalid source layout map");
+
+  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);
 }
 

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 with 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 {{source dim must be a product of expanded dims}}
   %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 {{result dim must be dynamic if and only if source dim 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 {{result dim must be dynamic if and only if source dim is dynamic}}
   %0 = memref.expand_shape %arg0 [[0], [1, 2]]
       : memref<?x?xf32> into memref<?x4x5xf32>
   return %0 : memref<?x4x5xf32>
 }
 
 // -----
 
+func @expand_shape_invalid_static_dim_size(%arg0 : memref<?x21xf32>)
+    -> memref<?x4x5xf32> {
+  // expected-error @+1 {{source dim must be a product of expanded dims}}
+  %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