Support non identity layout map for reshape ops in MemRefToLLVM lowering

[cathyzhyi]

Summary

This change borrows the ideas from computeExpanded/CollapsedLayoutMap
and computes the dynamic strides at runtime for the memref descriptors.

[cathyzhyi]

• Fix cases where all collapsed dims in an association are size 1.
• Fix tests.

[cathyzhyi]

Clean ups

[cathyzhyi]

Fix testcases format

[springerm]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1352

calculated

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

1825–1836 ↗

(On Diff #423738)

Can you factor this part of the CL out into a separate one? Along with the bufferization test case.

1826–1827 ↗

(On Diff #423738)

Good catch, there is definitely a problem here in the old implementation. The old implementation assumes that strides are sorted according to size. This is not always the case, however. E.g., you can use memref.transpose (which only touches the memref descriptor, not the data) to produce a new view on the data that is iterated over in a transposed fashion.

The check for 1 only solves this issue partial. Instead, we should be using the minimum stride in the current reassociation group. (The old code always used the last one.)

1828 ↗

(On Diff #423738)

This is problematic because we would be failing silently at runtime with a wrong result. We cannot just return failure() unfortunately. We should do the same computation as described above, choosing the minimum stride in the reassociation group. We don't need a loop for this, just a sequence of min operations.

[cathyzhyi]

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

1826–1827 ↗

(On Diff #423738)

The checking for 1 here is need for non transposed cases when the strides are in order. When the innermost dim is of size 1 the result stride needs to take the smallest stride on the src dim that's not size 1. This is the case for the test case collapse_shape_of_slice3 I added. I realize the input of collpase_shape for this case would not be contiguous but it can still be represented by the strided form and is needed for our testcase. I can't think of a good way to detect this at runtime unless inserting runtime assertion to assume dynamic dims are not size 1. Do you have any suggestions on that?

[springerm]

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

1826–1827 ↗

(On Diff #423738)

We had a discussion and concluded that this handling of dims of size 1 is correct. More work is needed to handle cases where strides are not sorted. This can be done in separate change.

[cathyzhyi]

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

1825–1836 ↗

(On Diff #423738)

Moved this part to a new patch https://reviews.llvm.org/D124137.

[cathyzhyi]

Factored out the part fixing size 1 cases into https://reviews.llvm.org/D124137.
The remaining memref to llvm lowering part needs to be updated to
take into consideration where the dynamic size is 1 as well.

[cathyzhyi]

Calculate with the new logic at runtime. I got an error about dominance here https://gist.github.com/949dafb96801866cf067b5baba4ac614. It's complaining the last mem ref descriptor doesn't dominate the use tensor.collapse_shape but it seems to me it does dominate tensor.collapse_shape.

[cathyzhyi]

Fixed the previous error about dominance.

[springerm]

This looks good. Can you add another test case for collapse and expand, so that we have one for the static case (taking the stride directly from the src type) and one for the dynamic case (also checking that the blocks+branches are generated correctly).

mlir/include/mlir/Conversion/LLVMCommon/MemRefBuilder.h

82

Returns

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1369

calculated

1382

Add a TODO that we should take the minimum stride in the reassociation group instead of just the first where the dimension is not 1.

1517–1523

Add a comment that there could be mixed static/dynamic strides. For simplicity, we recompute all strides if there is at least one dynamic stride.

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

852

Why did this test case change? This commit is just adding new functionality for a previously not handled case, so existing test cases should not change.

[springerm]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1412

should this be called "index type" maybe?

1419

I think there could be another small optimization here:

if (srcShape[srcIndex] == 1 && srcIndex != ref.front())
  continue;

[cathyzhyi]

Address comments.

[cathyzhyi]

Thanks for the review!

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

852

yea, I should have kept the implementation for the identity layout case instead of just using the new logic. Changed back the test cases and added the identity layout logic.

[cathyzhyi]

Move the comment about dynamic stride to a more propriate place

[springerm]

thanks!

mlir/include/mlir/Conversion/LLVMCommon/MemRefBuilder.h

public:
   /// Builds IR extracting the pos-th size from the descriptor.
   Value stride(OpBuilder &builder, Location loc, unsigned pos);
 
   /// Builds IR inserting the pos-th stride into the descriptor
   void setStride(OpBuilder &builder, Location loc, unsigned pos, Value stride);
   void setConstantStride(OpBuilder &builder, Location loc, unsigned pos,
                          uint64_t stride);
 
+  /// Returns the type of array element in this descriptor.
+  Type getIndexType() { return indexType; };
+
   /// Returns the (LLVM) pointer type this descriptor contains.
   LLVM::LLVMPointerType getElementPtrType();
 
   /// Builds IR populating a MemRef descriptor structure from a list of
   /// individual values composing that descriptor, in the following order:
   /// - allocated pointer;
   /// - aligned pointer;
   /// - offset;

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

for (auto inDimIndex : reassocation[outDimIndex]) {
     outDimSizeDynamic =
         b.create<LLVM::MulOp>(loc, outDimSizeDynamic, inDimSizeDynamic);
   }
   return outDimSizeDynamic;
 }
 
 static SmallVector<OpFoldResult, 4>
 getCollapsedOutputShape(OpBuilder &b, Location loc, Type &llvmIndexType,
-                        ArrayRef<ReassociationIndices> reassocation,
+                        ArrayRef<ReassociationIndices> reassociation,
                         ArrayRef<int64_t> inStaticShape,
                         MemRefDescriptor &inDesc,
                         ArrayRef<int64_t> outStaticShape) {
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, outStaticShape.size()), [&](int64_t outDimIndex) {
         return getCollapsedOutputDimSize(b, loc, llvmIndexType, outDimIndex,
                                          outStaticShape[outDimIndex],
-                                         inStaticShape, inDesc, reassocation);
+                                         inStaticShape, inDesc, reassociation);
       }));
 }
 
 static SmallVector<OpFoldResult, 4>
 getExpandedOutputShape(OpBuilder &b, Location loc, Type &llvmIndexType,
-                       ArrayRef<ReassociationIndices> reassocation,
+                       ArrayRef<ReassociationIndices> reassociation,
                        ArrayRef<int64_t> inStaticShape,
                        MemRefDescriptor &inDesc,
                        ArrayRef<int64_t> outStaticShape) {
   DenseMap<int64_t, int64_t> outDimToInDimMap =
-      getExpandedDimToCollapsedDimMap(reassocation);
+      getExpandedDimToCollapsedDimMap(reassociation);
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, outStaticShape.size()), [&](int64_t outDimIndex) {
         return getExpandedOutputDimSize(b, loc, llvmIndexType, outDimIndex,
                                         outStaticShape, inDesc, inStaticShape,
-                                        reassocation, outDimToInDimMap);
+                                        reassociation, outDimToInDimMap);
       }));
 }
 
 static SmallVector<Value>
 getDynamicOutputShape(OpBuilder &b, Location loc, Type &llvmIndexType,
-                      ArrayRef<ReassociationIndices> reassocation,
+                      ArrayRef<ReassociationIndices> reassociation,
                       ArrayRef<int64_t> inStaticShape, MemRefDescriptor &inDesc,
                       ArrayRef<int64_t> outStaticShape) {
   return outStaticShape.size() < inStaticShape.size()
              ? getAsValues(b, loc, llvmIndexType,
                            getCollapsedOutputShape(b, loc, llvmIndexType,
-                                                   reassocation, inStaticShape,
+                                                   reassociation, inStaticShape,
                                                    inDesc, outStaticShape))
              : getAsValues(b, loc, llvmIndexType,
                            getExpandedOutputShape(b, loc, llvmIndexType,
-                                                  reassocation, inStaticShape,
+                                                  reassociation, inStaticShape,
                                                   inDesc, outStaticShape));
 }
 
+static void fillInStridesForExpandedMemDescriptor(
+    OpBuilder &b, Location loc, MemRefType srcType, MemRefDescriptor &srcDesc,
+    MemRefDescriptor &dstDesc, ArrayRef<ReassociationIndices> reassociation) {
+  // See comments for computeExpandedLayoutMap for details on how the strides
+  // are calculated.
+  for (auto &en : llvm::enumerate(reassociation)) {
+    auto currentStrideToExpand = srcDesc.stride(b, loc, en.index());
+    for (auto dstIndex : llvm::reverse(en.value())) {
+      dstDesc.setStride(b, loc, dstIndex, currentStrideToExpand);
+      Value size = dstDesc.size(b, loc, dstIndex);
+      currentStrideToExpand =
+          b.create<LLVM::MulOp>(loc, size, currentStrideToExpand);
+    }
+  }
+}
+
+static void fillInStridesForCollapsedMemDescriptor(
+    ConversionPatternRewriter &rewriter, Location loc, Operation *op,
+    TypeConverter *typeConverter, MemRefType srcType, MemRefDescriptor &srcDesc,
+    MemRefDescriptor &dstDesc, ArrayRef<ReassociationIndices> reassociation) {
+  // See comments for computeCollapsedLayoutMap for details on how the strides
+  // are calculated.
+  auto srcShape = srcType.getShape();
+  for (auto &en : llvm::enumerate(reassociation)) {
+    rewriter.setInsertionPoint(op);
+    auto dstIndex = en.index();
+    ArrayRef<int64_t> ref = llvm::makeArrayRef(en.value());
+    while (srcShape[ref.back()] == 1 && ref.size() > 1)
+      ref = ref.drop_back();
+    if (!ShapedType::isDynamic(srcShape[ref.back()]) || ref.size() == 1) {
+      dstDesc.setStride(rewriter, loc, dstIndex,
+                        srcDesc.stride(rewriter, loc, ref.back()));
+    } else {
+      // Iterate over the source strides in reverse order. Skip over the
+      // dimensions whose size is 1.
+      // TODO: we should take the minimum stride in the reassociation group
+      // instead of just the first where the dimension is not 1.
+      //
+      // +------------------------------------------------------+
+      // | curEntry:                                            |
+      // |   %srcStride = strides[srcIndex]                     |
+      // |   %neOne = cmp sizes[srcIndex],1                     +--+
+      // |   cf.cond_br %neOne, continue(%srcStride), nextEntry |  |
+      // +-------------------------+----------------------------+  |
+      //                           |                               |
+      //                           v                               |
+      //            +-----------------------------+                |
+      //            | nextEntry:                  |                |
+      //            |   ...                       +---+            |
+      //            +--------------+--------------+   |            |
+      //                           |                  |            |
+      //                           v                  |            |
+      //            +-----------------------------+   |            |
+      //            | nextEntry:                  |   |            |
+      //            |   ...                       |   |            |
+      //            +--------------+--------------+   |   +--------+
+      //                           |                  |   |
+      //                           v                  v   v
+      //   +--------------------------------------------------+
+      //   | continue(%newStride):                            |
+      //   |   %newMemRefDes = setStride(%newStride,dstIndex) |
+      //   +--------------------------------------------------+
+      OpBuilder::InsertionGuard guard(rewriter);
+      Block *initBlock = rewriter.getInsertionBlock();
+      Block *continueBlock =
+          rewriter.splitBlock(initBlock, rewriter.getInsertionPoint());
+      continueBlock->insertArgument(unsigned(0), srcDesc.getIndexType(), loc);
+      rewriter.setInsertionPointToStart(continueBlock);
+      dstDesc.setStride(rewriter, loc, dstIndex, continueBlock->getArgument(0));
+
+      Block *curEntryBlock = initBlock;
+      Block *nextEntryBlock;
+      for (auto srcIndex : llvm::reverse(ref)) {
+        if (srcShape[srcIndex] == 1 && srcIndex != ref.front())
+          continue;
+        rewriter.setInsertionPointToEnd(curEntryBlock);
+        Value srcStride = srcDesc.stride(rewriter, loc, srcIndex);
+        if (srcIndex == ref.front()) {
+          rewriter.create<LLVM::BrOp>(loc, srcStride, continueBlock);
+          break;
+        }
+        Value one = rewriter.create<LLVM::ConstantOp>(
+            loc, typeConverter->convertType(rewriter.getI64Type()),
+            rewriter.getI32IntegerAttr(1));
+        Value predNeOne = rewriter.create<LLVM::ICmpOp>(
+            loc, LLVM::ICmpPredicate::ne, srcDesc.size(rewriter, loc, srcIndex),
+            one);
+        {
+          OpBuilder::InsertionGuard guard(rewriter);
+          nextEntryBlock = rewriter.createBlock(
+              initBlock->getParent(), Region::iterator(continueBlock), {});
+        }
+        rewriter.create<LLVM::CondBrOp>(loc, predNeOne, continueBlock,
+                                        srcStride, nextEntryBlock, llvm::None);
+        curEntryBlock = nextEntryBlock;
+      }
+    }
+  }
+}
+
+static void fillInDynamicStridesForMemDescriptor(
+    ConversionPatternRewriter &b, Location loc, Operation *op,
+    TypeConverter *typeConverter, MemRefType srcType, MemRefType dstType,
+    MemRefDescriptor &srcDesc, MemRefDescriptor &dstDesc,
+    ArrayRef<ReassociationIndices> reassociation) {
+  if (srcType.getRank() > dstType.getRank())
+    fillInStridesForCollapsedMemDescriptor(b, loc, op, typeConverter, srcType,
+                                           srcDesc, dstDesc, reassociation);
+  else
+    fillInStridesForExpandedMemDescriptor(b, loc, srcType, srcDesc, dstDesc,
+                                          reassociation);
+}
+
 // ReshapeOp creates a new view descriptor of the proper rank.
 // For now, the only conversion supported is for target MemRef with static sizes
 // and strides.
 template <typename ReshapeOp>
 class ReassociatingReshapeOpConversion
     : public ConvertOpToLLVMPattern<ReshapeOp> {
 public:
   using ConvertOpToLLVMPattern<ReshapeOp>::ConvertOpToLLVMPattern;
   using ReshapeOpAdaptor = typename ReshapeOp::Adaptor;
 
   LogicalResult
   matchAndRewrite(ReshapeOp reshapeOp, typename ReshapeOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     MemRefType dstType = reshapeOp.getResultType();
     MemRefType srcType = reshapeOp.getSrcType();
 
-    // The condition on the layouts can be ignored when all shapes are static.
-    if (!srcType.hasStaticShape() || !dstType.hasStaticShape()) {
-      if (!srcType.getLayout().isIdentity() ||
-          !dstType.getLayout().isIdentity()) {
-        return rewriter.notifyMatchFailure(
-            reshapeOp, "only empty layout map is supported");
-      }
-    }
-
     int64_t offset;
     SmallVector<int64_t, 4> strides;
     if (failed(getStridesAndOffset(dstType, strides, offset))) {
       return rewriter.notifyMatchFailure(
           reshapeOp, "failed to get stride and offset exprs");
     }
 
     MemRefDescriptor srcDesc(adaptor.src());
for (auto &en : llvm::enumerate(dstShape))
       dstDesc.setSize(rewriter, loc, en.index(), en.value());
 
     auto isStaticStride = [](int64_t stride) {
       return !ShapedType::isDynamicStrideOrOffset(stride);
     };
     if (llvm::all_of(strides, isStaticStride)) {
       for (auto &en : llvm::enumerate(strides))
         dstDesc.setConstantStride(rewriter, loc, en.index(), en.value());
-    } else {
+    } else if (srcType.getLayout().isIdentity() &&
+               dstType.getLayout().isIdentity()) {
       Value c1 = rewriter.create<LLVM::ConstantOp>(loc, llvmIndexType,
                                                    rewriter.getIndexAttr(1));
       Value stride = c1;
       for (auto dimIndex :
            llvm::reverse(llvm::seq<int64_t>(0, dstShape.size()))) {
         dstDesc.setStride(rewriter, loc, dimIndex, stride);
         stride = rewriter.create<LLVM::MulOp>(loc, dstShape[dimIndex], stride);
       }
+    } else {
+      // There could be mixed static/dynamic strides. For simplicity, we
+      // recompute all strides if there is at least one dynamic stride.
+      fillInDynamicStridesForMemDescriptor(
+          rewriter, loc, reshapeOp, this->typeConverter, srcType, dstType,
+          srcDesc, dstDesc, reshapeOp.getReassociationIndices());
     }
     rewriter.replaceOp(reshapeOp, {dstDesc});
     return success();
   }
 };
 
 /// Conversion pattern that transforms a subview op into:
 ///   1. An `llvm.mlir.undef` operation to create a memref descriptor

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

 //       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 //       CHECK:    llvm.mlir.constant(5 : index) : i64
 //       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 //       CHECK:    llvm.mlir.constant(1 : index) : i64
 //       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 
 // -----
 
+func.func @collapse_shape_dynamic_with_non_identity_layout(
+        %arg0 : memref<4x?x?xf32, affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s1 + s0 + d1 * 4 + d2)>>) ->
+        memref<4x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>> {
+  %0 = memref.collapse_shape %arg0 [[0], [1, 2]]:
+    memref<4x?x?xf32, affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s1 + s0 + d1 * 4 + d2)>> into
+    memref<4x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>>
+  return %0 : memref<4x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>>
+}
+// CHECK-LABEL:   func @collapse_shape_dynamic_with_non_identity_layout(
+//       CHECK:      llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.mlir.constant(1 : index) : i64
+//       CHECK:      llvm.extractvalue %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.mul %{{.*}}, %{{.*}}  : i64
+//       CHECK:      llvm.extractvalue %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.mul %{{.*}}, %{{.*}}  : i64
+//       CHECK:      llvm.mlir.constant(4 : index) : i64
+//       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.extractvalue %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:      llvm.extractvalue %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.mlir.constant(1 : i32) : i64
+//       CHECK:      llvm.extractvalue %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.icmp "ne" %{{.*}}, %{{.*}} : i64
+//       CHECK:      llvm.cond_br %{{.*}}, ^bb2(%{{.*}} : i64), ^bb1
+//       CHECK:      ^bb1:
+//       CHECK:      llvm.extractvalue %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:      llvm.br ^bb2(%{{.*}} : i64)
+//       CHECK:      ^bb2(%[[STRIDE:.*]]: i64):
+//       CHECK:      llvm.insertvalue %[[STRIDE]], %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+
+// -----
+
 func.func @expand_shape_static(%arg0: memref<3x4x5xf32>) -> memref<1x3x4x1x5xf32> {
   // Reshapes that expand a contiguous tensor with some 1's.
   %0 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]]
       : memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
   return %0 : memref<1x3x4x1x5xf32>
 }
 // CHECK-LABEL: func @expand_shape_static
 //       CHECK:    llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>

 // CHECK:           llvm.mlir.constant(1 : index) : i64
 // CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
 // CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
 
 // -----
 
 func.func @expand_shape_dynamic(%arg0 : memref<1x?xf32>) -> memref<1x2x?xf32> {
   %0 = memref.expand_shape %arg0 [[0], [1, 2]]: memref<1x?xf32> into memref<1x2x?xf32>
   return %0 : memref<1x2x?xf32>
 }
 // CHECK-LABEL:   func @expand_shape_dynamic(
 // CHECK:           llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 // CHECK:           llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 // CHECK:           llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>

 // CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
 // CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 // CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
 // CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 // CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
 
 // -----
 
+func.func @expand_shape_dynamic_with_non_identity_layout(
+            %arg0 : memref<1x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>>) ->
+            memref<1x2x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>> {
+  %0 = memref.expand_shape %arg0 [[0], [1, 2]]:
+    memref<1x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>> into
+    memref<1x2x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>>
+  return %0 : memref<1x2x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>>
+}
+// CHECK-LABEL:   func @expand_shape_dynamic_with_non_identity_layout(
+// CHECK:           llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK:           llvm.mlir.constant(2 : index) : i64
+// CHECK:           llvm.sdiv %{{.*}}, %{{.*}}  : i64
+// CHECK:           llvm.mlir.constant(1 : index) : i64
+// CHECK:           llvm.mlir.constant(2 : index) : i64
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
+// CHECK:           llvm.extractvalue %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
+// CHECK:           llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.extractvalue %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+// CHECK:           llvm.mul %{{.*}}, %{{.*}}  : i64
+
+// -----
+
 // CHECK-LABEL: func @rank_of_unranked
 // CHECK32-LABEL: func @rank_of_unranked
 func.func @rank_of_unranked(%unranked: memref<*xi32>) {
   %rank = memref.rank %unranked : memref<*xi32>
   return
 }
 // CHECK: %[[UNRANKED_DESC:.*]] = builtin.unrealized_conversion_cast
 // CHECK-NEXT: llvm.extractvalue %[[UNRANKED_DESC]][0] : !llvm.struct<(i64, ptr<i8>)>