[mlir] Convert `memref_reshape` to LLVM.

[pifon2a]

Summary

https://llvm.discourse.group/t/rfc-standard-memref-cast-ops/1454/15

[pifon2a]

remove print

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2420

Nit: could we have a doc for this function?

2455–2456

Nit: would it be reasonable to have these things as additional accessors on UnrankedMemRefDescriptor?

2468–2471

I wonder if we cannot run into issues because of alignment properties on some weird architectures. The descriptor structure is not packed, so it is subject to alignment rules between elements as defined by LLVM's data layout. I don't have an example offhand, so it may be always fine, so I'd appreciate an argument why it is a safe thing to do in a comment. Otherwise, LLVMTypeConverter contains the llvm::DataLayout that we are targeting, which can be used to get proper offsets of elements in bytes and do all indexing arithmetic after bitcasting to i8*.

2572

Please fix

2609–2611

Nit: LLVM::LLVMType was fine here

2618–2619

Can't we use createIndexConstant instead here? Or at least not hardcode i32, I think there was some option controlling the bit size of the address arithmetic

2625–2630

Putting these addressing tricks (which are cool, I must admit!) into the UnrankedMemRefDescriptor sounds even more appealing to me. Since we seem to always need the triple allocated/aligned/store, we can have a function for those.

2651–2652

addArgument is invalid in conversion patterns, similarly to other in-place updates. rewriter.createBlock creates block with arguments, but you'd need to clone the remaining operations in there. Maybe splitBlock can be extended to also add arguments to the newly created block.

(I suspect the addArgument will most likely just work even if we rollback the change, but it may run into some bad use-def loop)

2675

(Beyond the scope): I wonder if we could later refactor this and the SCF-to-std lowering in a createLoop(function_ref bodyBuilder, function_ref conditionBuilder) that produces std control flow; the pattern infra will then lower the std ops into LLVM automatically.

[pifon2a]

Address all of the comments but one.

[pifon2a]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2468–2471

I am not sure how many architectures, not to mention weird ones, are using unranked code generation. I would leave it like that for now and use llvm::DataLayout later if needed.

2625–2630

I added setters/getters to UnrankedMemRefDescriptor. At first I intended to have a separate PR that adds them, but it looks like the current PR is the best way to actually test them.

2675

Yes, that would be much more readable.

[pifon2a]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2651–2652

extending splitBlock looks much harder than just using createBlock and cloning the remaining ops. I ll do that tomorrow.

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2468–2471

Let's keep a TODO comment then. If we ever run into a problem, will be easier to debug.

[pifon2a]

Address the remaining comments.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

if (srcType.isa<MemRefType>() && dstType.isa<UnrankedMemRefType>()) {
       auto loadOp = rewriter.create<LLVM::LoadOp>(loc, castPtr);
       rewriter.replaceOp(op, loadOp.getResult());
     } else {
       llvm_unreachable("Unsupported unranked memref to unranked memref cast");
     }
   }
 };
 
+static void extractPointersAndOffset(Location loc,
+                                     ConversionPatternRewriter &rewriter,
+                                     LLVMTypeConverter &typeConverter,
+                                     Value originalOperand,
+                                     Value convertedOperand,
+                                     Value *allocatedPtr, Value *alignedPtr,
+                                     Value *offset = nullptr) {
+  Type operandType = originalOperand.getType();
+  if (operandType.isa<MemRefType>()) {
+    MemRefDescriptor desc(convertedOperand);
+    *allocatedPtr = desc.allocatedPtr(rewriter, loc);
+    *alignedPtr = desc.alignedPtr(rewriter, loc);
+    if (offset != nullptr)
+      *offset = desc.offset(rewriter, loc);
+    return;
+  }
+
+  unsigned memorySpace =
+      operandType.cast<UnrankedMemRefType>().getMemorySpace();
+  Type elementType = operandType.cast<UnrankedMemRefType>().getElementType();
+  LLVM::LLVMType llvmElementType =
+      typeConverter.convertType(elementType).cast<LLVM::LLVMType>();
+  LLVM::LLVMType elementPtrPtrType =
+      llvmElementType.getPointerTo(memorySpace).getPointerTo();
+
+  // Extract pointer to the underlying ranked memref descriptor and cast it to
+  // ElemType**.
+  UnrankedMemRefDescriptor unrankedDesc(convertedOperand);
+  Value underlyingDescPtr = unrankedDesc.memRefDescPtr(rewriter, loc);
+  Value elementPtrPtr = rewriter.create<LLVM::BitcastOp>(loc, elementPtrPtrType,
+                                                         underlyingDescPtr);
+
+  LLVM::LLVMType int32Type =
+      typeConverter.convertType(rewriter.getI32Type()).cast<LLVM::LLVMType>();
+
+  // Extract and set allocated pointer.
+  *allocatedPtr = rewriter.create<LLVM::LoadOp>(loc, elementPtrPtr);
+
+  // Extract and set aligned pointer.
+  Value one = rewriter.create<LLVM::ConstantOp>(loc, int32Type,
+                                                rewriter.getI32IntegerAttr(1));
+  Value alignedGep = rewriter.create<LLVM::GEPOp>(
+      loc, elementPtrPtrType, elementPtrPtr, ValueRange({one}));
+  *alignedPtr = rewriter.create<LLVM::LoadOp>(loc, alignedGep);
+
+  // Extract and set offset.
+  if (offset != nullptr) {
+    auto llvmIndexPtrType = typeConverter.getIndexType().getPointerTo();
+    Value offsetGep = rewriter.create<LLVM::GEPOp>(
+        loc, elementPtrPtrType, alignedGep, ValueRange({one}));
+    offsetGep =
+        rewriter.create<LLVM::BitcastOp>(loc, llvmIndexPtrType, offsetGep);
+    *offset = rewriter.create<LLVM::LoadOp>(loc, offsetGep);
+  }
+}
+
 struct MemRefReinterpretCastOpLowering
     : public ConvertOpToLLVMPattern<MemRefReinterpretCastOp> {
   using ConvertOpToLLVMPattern<MemRefReinterpretCastOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto castOp = cast<MemRefReinterpretCastOp>(op);
if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy())
       return failure();
 
     // Create descriptor.
     Location loc = castOp.getLoc();
     auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
 
     // Set allocated and aligned pointers.
     Value allocatedPtr, alignedPtr;
-    extractPointers(loc, rewriter, castOp.source(), adaptor.source(),
-                    &allocatedPtr, &alignedPtr);
+    extractPointersAndOffset(loc, rewriter, typeConverter, castOp.source(),
+                             adaptor.source(), &allocatedPtr, &alignedPtr);
     desc.setAllocatedPtr(rewriter, loc, allocatedPtr);
     desc.setAlignedPtr(rewriter, loc, alignedPtr);
 
     // Set offset.
     if (castOp.isDynamicOffset(0))
       desc.setOffset(rewriter, loc, adaptor.offsets()[0]);
     else
       desc.setConstantOffset(rewriter, loc, castOp.getStaticOffset(0));
for (unsigned i = 0, e = targetMemRefType.getRank(); i < e; ++i) {
       if (castOp.isDynamicStride(i))
         desc.setStride(rewriter, loc, i, adaptor.strides()[dynStrideId++]);
       else
         desc.setConstantStride(rewriter, loc, i, castOp.getStaticStride(i));
     }
     *descriptor = desc;
     return success();
   }
+};
 
-  void extractPointers(Location loc, ConversionPatternRewriter &rewriter,
-                       Value originalOperand, Value convertedOperand,
-                       Value *allocatedPtr, Value *alignedPtr) const {
-    Type operandType = originalOperand.getType();
-    if (operandType.isa<MemRefType>()) {
-      MemRefDescriptor desc(convertedOperand);
-      *allocatedPtr = desc.allocatedPtr(rewriter, loc);
-      *alignedPtr = desc.alignedPtr(rewriter, loc);
-      return;
-    }
+struct MemRefReshapeOpLowering
+    : public ConvertOpToLLVMPattern<MemRefReshapeOp> {
+  using ConvertOpToLLVMPattern<MemRefReshapeOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto reshapeOp = cast<MemRefReshapeOp>(op);
+
+    MemRefReshapeOp::Adaptor adaptor(operands, op->getAttrDictionary());
+    Type srcType = reshapeOp.source().getType();
+
+    Value descriptor;
+    if (failed(convertSourceMemRefToDescriptor(rewriter, srcType, reshapeOp,
+                                               adaptor, &descriptor)))
+      return failure();
+    rewriter.replaceOp(op, {descriptor});
+    return success();
 
-    unsigned memorySpace =
-        operandType.cast<UnrankedMemRefType>().getMemorySpace();
-    Type elementType = operandType.cast<UnrankedMemRefType>().getElementType();
+private:
+  LogicalResult
+  convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
+                                  Type srcType, MemRefReshapeOp reshapeOp,
+                                  MemRefReshapeOp::Adaptor adaptor,
+                                  Value *descriptor) const {
+    // Conversion for statically-known shape args is performed via
+    // `memref_reinterpret_cast`.
+    auto shape_memref_type = reshapeOp.shape().getType().cast<MemRefType>();
+    if (shape_memref_type.hasStaticShape())
+      return failure();
+
+    // The shape is a rank-1 tensor with unknown length.
+    Location loc = reshapeOp.getLoc();
+    MemRefDescriptor shapeDesc(adaptor.shape());
+    Value resultRank = shapeDesc.size(rewriter, loc, 0);
+
+    // Extract address space and element type.
+    auto targetType =
+        reshapeOp.getResult().getType().cast<UnrankedMemRefType>();
+    unsigned addressSpace = targetType.getMemorySpace();
+    Type elementType = targetType.getElementType();
+
+    // Create the unranked memref descriptor that holds the ranked one. The
+    // inner descriptor is allocated on stack.
+    auto targetDesc = UnrankedMemRefDescriptor::undef(
+        rewriter, loc,
+        typeConverter.convertType(targetType).cast<LLVM::LLVMType>());
+    targetDesc.setRank(rewriter, loc, resultRank);
+    SmallVector<Value, 4> sizes;
+    UnrankedMemRefDescriptor::computeSizes(rewriter, loc, typeConverter,
+                                           targetDesc, sizes);
+    Value underlyingDescPtr = rewriter.create<LLVM::AllocaOp>(
+        loc, getVoidPtrType(), sizes.front(), llvm::None);
+    targetDesc.setMemRefDescPtr(rewriter, loc, underlyingDescPtr);
+
+    // Extract pointers and offset from the source memref.
+    Value allocatedPtr, alignedPtr, offset;
+    extractPointersAndOffset(loc, rewriter, typeConverter, reshapeOp.source(),
+                             adaptor.source(), &allocatedPtr, &alignedPtr,
+                             &offset);
+
+    // Set allocated pointer.
     LLVM::LLVMType llvmElementType =
         typeConverter.convertType(elementType).cast<LLVM::LLVMType>();
-    LLVM::LLVMType elementPtrPtrType =
-        llvmElementType.getPointerTo(memorySpace).getPointerTo();
-
-    // Extract pointer to the underlying ranked memref descriptor and cast it to
-    // ElemType**.
-    UnrankedMemRefDescriptor unrankedDesc(convertedOperand);
-    Value underlyingDescPtr = unrankedDesc.memRefDescPtr(rewriter, loc);
-    Value elementPtrPtr = rewriter.create<LLVM::BitcastOp>(
+    auto elementPtrPtrType =
+        llvmElementType.getPointerTo(addressSpace).getPointerTo();
+    Value allocatedGep = rewriter.create<LLVM::BitcastOp>(
         loc, elementPtrPtrType, underlyingDescPtr);
+    rewriter.create<LLVM::StoreOp>(loc, allocatedPtr, allocatedGep);
 
+    // Set aligned pointer.
     LLVM::LLVMType int32Type =
         typeConverter.convertType(rewriter.getI32Type()).cast<LLVM::LLVMType>();
-
-    // Extract and set allocated pointer.
-    *allocatedPtr = rewriter.create<LLVM::LoadOp>(loc, elementPtrPtr);
-
-    // Extract and set aligned pointer.
     Value one = rewriter.create<LLVM::ConstantOp>(
         loc, int32Type, rewriter.getI32IntegerAttr(1));
     Value alignedGep = rewriter.create<LLVM::GEPOp>(
-        loc, elementPtrPtrType, elementPtrPtr, ValueRange({one}));
-    *alignedPtr = rewriter.create<LLVM::LoadOp>(loc, alignedGep);
+        loc, elementPtrPtrType, allocatedGep, ValueRange({one}));
+    rewriter.create<LLVM::StoreOp>(loc, alignedPtr, alignedGep);
+
+    // Set offset.
+    auto llvmIndexPtrType = typeConverter.getIndexType().getPointerTo();
+    Value offsetGep = rewriter.create<LLVM::GEPOp>(
+        loc, llvmIndexPtrType, alignedGep, ValueRange({one}));
+    offsetGep =
+        rewriter.create<LLVM::BitcastOp>(loc, llvmIndexPtrType, offsetGep);
+    rewriter.create<LLVM::StoreOp>(loc, offset, offsetGep);
+
+    // Use the offset pointer as base for further addressing. Copy over the new
+    // shape and compute strides. For this, we create a loop from rank-1 to 0.
+    Value zeroIndex = createIndexConstant(rewriter, loc, 0);
+    Value oneIndex = createIndexConstant(rewriter, loc, 1);
+    Value targetSizesBase = rewriter.create<LLVM::GEPOp>(
+        loc, llvmIndexPtrType, offsetGep, ValueRange({one}));
+    Value targetStridesBase = rewriter.create<LLVM::GEPOp>(
+        loc, llvmIndexPtrType, targetSizesBase, ValueRange({resultRank}));
+    Value shapeOperandPtr = shapeDesc.alignedPtr(rewriter, loc);
+    Value resultRankMinusOne =
+        rewriter.create<LLVM::SubOp>(loc, resultRank, oneIndex);
+
+    Block *initBlock = rewriter.getInsertionBlock();
+    Block *condBlock =
+        rewriter.splitBlock(initBlock, rewriter.getInsertionPoint());
+    rewriter.setInsertionPointToEnd(initBlock);
+    rewriter.create<LLVM::BrOp>(loc, ValueRange({resultRankMinusOne, oneIndex}),
+                                condBlock);
+    rewriter.setInsertionPointToStart(condBlock);
+    auto indexArg = condBlock->addArgument(typeConverter.getIndexType());
+    auto strideArg = condBlock->addArgument(typeConverter.getIndexType());
+    auto pred = rewriter.create<LLVM::ICmpOp>(
+        loc, LLVM::LLVMType::getInt1Ty(rewriter.getContext()),
+        LLVM::ICmpPredicate::sge, indexArg, zeroIndex);
+
+    Block *bodyBlock =
+        rewriter.splitBlock(condBlock, rewriter.getInsertionPoint());
+    rewriter.setInsertionPointToStart(bodyBlock);
+
+    // Copy size from shape to descriptor.
+    Value sizeLoadGep = rewriter.create<LLVM::GEPOp>(
+        loc, llvmIndexPtrType, shapeOperandPtr, ValueRange{indexArg});
+    Value size = rewriter.create<LLVM::LoadOp>(loc, sizeLoadGep);
+    Value sizeStoreGep = rewriter.create<LLVM::GEPOp>(
+        loc, llvmIndexPtrType, targetSizesBase, ValueRange({indexArg}));
+    rewriter.create<LLVM::StoreOp>(loc, size, sizeStoreGep);
+
+    // Write stride value and compute next one.
+    Value strideStoreGep = rewriter.create<LLVM::GEPOp>(
+        loc, llvmIndexPtrType, targetStridesBase, ValueRange({indexArg}));
+    rewriter.create<LLVM::StoreOp>(loc, strideArg, strideStoreGep);
+    Value nextStride = rewriter.create<LLVM::MulOp>(loc, strideArg, size);
+
+    // Decrement loop counter and branch back.
+    Value decrement = rewriter.create<LLVM::SubOp>(loc, indexArg, oneIndex);
+    rewriter.create<LLVM::BrOp>(loc, ValueRange({decrement, nextStride}),
+                                condBlock);
+
+    Block *remainder =
+        rewriter.splitBlock(bodyBlock, rewriter.getInsertionPoint());
+
+    // Hook up the cond exit to the remainder.
+    rewriter.setInsertionPointToEnd(condBlock);
+    rewriter.create<LLVM::CondBrOp>(loc, pred, bodyBlock, llvm::None, remainder,
+                                    llvm::None);
+
+    // Reset position to beginning of new remainder block.
+    rewriter.setInsertionPointToStart(remainder);
+
+    *descriptor = targetDesc;
+    return success();
   }
 };
 
 struct DialectCastOpLowering
     : public ConvertOpToLLVMPattern<LLVM::DialectCastOp> {
   using ConvertOpToLLVMPattern<LLVM::DialectCastOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
void mlir::populateStdToLLVMMemoryConversionPatterns(
   // clang-format off
   patterns.insert<
       AssumeAlignmentOpLowering,
       DeallocOpLowering,
       DimOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       MemRefReinterpretCastOpLowering,
+      MemRefReshapeOpLowering,
       RankOpLowering,
       StoreOpLowering,
       SubViewOpLowering,
       TransposeOpLowering,
       ViewOpLowering>(converter);
   // clang-format on
   if (converter.getOptions().useAlignedAlloc)
     patterns.insert<AlignedAllocOpLowering>(converter);

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

 // CHECK: [[ALIGNED_PTR:%.*]] = llvm.load [[ALIGNED_PTR_PTR]] : !llvm.ptr<ptr<float>>
 // CHECK: [[OUT_1:%.*]] = llvm.insertvalue [[BASE_PTR]], [[OUT_0]][0] : [[TY]]
 // CHECK: [[OUT_2:%.*]] = llvm.insertvalue [[ALIGNED_PTR]], [[OUT_1]][1] : [[TY]]
 // CHECK: [[OUT_3:%.*]] = llvm.insertvalue [[OFFSET]], [[OUT_2]][2] : [[TY]]
 // CHECK: [[OUT_4:%.*]] = llvm.insertvalue [[SIZE_0]], [[OUT_3]][3, 0] : [[TY]]
 // CHECK: [[OUT_5:%.*]] = llvm.insertvalue [[STRIDE_0]], [[OUT_4]][4, 0] : [[TY]]
 // CHECK: [[OUT_6:%.*]] = llvm.insertvalue [[SIZE_1]], [[OUT_5]][3, 1] : [[TY]]
 // CHECK: [[OUT_7:%.*]] = llvm.insertvalue [[STRIDE_1]], [[OUT_6]][4, 1] : [[TY]]
+
+// CHECK-LABEL: @memref_reshape
+func @memref_reshape(%input : memref<2x3xf32>, %shape : memref<?xindex>) {
+  %output = memref_reshape %input(%shape)
+                : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32>
+  return
+}
+// CHECK: [[INPUT:%.*]] = llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : [[INPUT_TY:!.*]]
+// CHECK: [[SHAPE:%.*]] = llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : [[SHAPE_TY:!.*]]
+// CHECK: [[RANK:%.*]] = llvm.extractvalue [[SHAPE]][3, 0] : [[SHAPE_TY]]
+// CHECK: [[UNRANKED_OUT_O:%.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
+// CHECK: [[UNRANKED_OUT_1:%.*]] = llvm.insertvalue [[RANK]], [[UNRANKED_OUT_O]][0] : !llvm.struct<(i64, ptr<i8>)>
+
+// Compute size in bytes to allocate result ranked descriptor
+// CHECK: [[C1:%.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+// CHECK: [[C2:%.*]] = llvm.mlir.constant(2 : index) : !llvm.i64
+// CHECK: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : index) : !llvm.i64
+// CHECK: [[INDEX_SIZE:%.*]] = llvm.mlir.constant(8 : index) : !llvm.i64
+// CHECK: [[DOUBLE_PTR_SIZE:%.*]] = llvm.mul [[C2]], [[PTR_SIZE]] : !llvm.i64
+// CHECK: [[DESC_ALLOC_SIZE:%.*]] = llvm.add [[DOUBLE_PTR_SIZE]], %{{.*}}
+// CHECK: [[UNDERLYING_DESC:%.*]] = llvm.alloca [[DESC_ALLOC_SIZE]] x !llvm.i8
+// CHECK: llvm.insertvalue [[UNDERLYING_DESC]], [[UNRANKED_OUT_1]][1]
+
+// Set allocated, aligned pointers and offset.
+// CHECK: [[ALLOC_PTR:%.*]] = llvm.extractvalue [[INPUT]][0] : [[INPUT_TY]]
+// CHECK: [[ALIGN_PTR:%.*]] = llvm.extractvalue [[INPUT]][1] : [[INPUT_TY]]
+// CHECK: [[OFFSET:%.*]] = llvm.extractvalue [[INPUT]][2] : [[INPUT_TY]]
+// CHECK: [[BASE_PTR_PTR:%.*]] = llvm.bitcast [[UNDERLYING_DESC]]
+// CHECK-SAME:                     !llvm.ptr<i8> to !llvm.ptr<ptr<float>>
+// CHECK: llvm.store [[ALLOC_PTR]], [[BASE_PTR_PTR]] : !llvm.ptr<ptr<float>>
+// CHECK: [[C1_I32:%.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+// CHECK: [[ALIGNED_PTR_PTR:%.*]] = llvm.getelementptr [[BASE_PTR_PTR]]{{\[}}[[C1_I32]]]
+// CHECK: llvm.store [[ALIGN_PTR]], [[ALIGNED_PTR_PTR]] : !llvm.ptr<ptr<float>>
+// CHECK: [[OFFSET_PTR_:%.*]] = llvm.getelementptr [[ALIGNED_PTR_PTR]]{{\[}}[[C1_I32]]]
+// CHECK: [[OFFSET_PTR:%.*]] = llvm.bitcast [[OFFSET_PTR_]]
+// CHECK: llvm.store [[OFFSET]], [[OFFSET_PTR]] : !llvm.ptr<i64>
+
+// Iterate over shape operand in reverse order and set sizes and strides.
+// CHECK: [[C0:%.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+// CHECK: [[C1:%.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+// CHECK: [[SIZES_PTR:%.*]] = llvm.getelementptr [[OFFSET_PTR]]{{\[}}[[C1_I32]]]
+// CHECK: [[STRIDES_PTR:%.*]] = llvm.getelementptr [[SIZES_PTR]]{{\[}}[[RANK]]]
+// CHECK: [[SHAPE_IN_PTR:%.*]] = llvm.extractvalue [[SHAPE]][1] : [[SHAPE_TY]]
+// CHECK: [[RANK_MIN_1:%.*]] = llvm.sub [[RANK]], [[C1]] : !llvm.i64
+// CHECK: llvm.br ^bb1([[RANK_MIN_1]], [[C1]] : !llvm.i64, !llvm.i64)
+
+// CHECK: ^bb1([[DIM:%.*]]: !llvm.i64, [[CUR_STRIDE:%.*]]: !llvm.i64):
+// CHECK:   [[COND:%.*]] = llvm.icmp "sge" [[DIM]], [[C0]] : !llvm.i64
+// CHECK:   llvm.cond_br [[COND]], ^bb2, ^bb3
+
+// CHECK: ^bb2:
+// CHECK:   [[SIZE_PTR:%.*]] = llvm.getelementptr [[SHAPE_IN_PTR]]{{\[}}[[DIM]]]
+// CHECK:   [[SIZE:%.*]] = llvm.load [[SIZE_PTR]] : !llvm.ptr<i64>
+// CHECK:   [[TARGET_SIZE_PTR:%.*]] = llvm.getelementptr [[SIZES_PTR]]{{\[}}[[DIM]]]
+// CHECK:   llvm.store [[SIZE]], [[TARGET_SIZE_PTR]] : !llvm.ptr<i64>
+// CHECK:   [[TARGET_STRIDE_PTR:%.*]] = llvm.getelementptr [[STRIDES_PTR]]{{\[}}[[DIM]]]
+// CHECK:   llvm.store [[CUR_STRIDE]], [[TARGET_STRIDE_PTR]] : !llvm.ptr<i64>
+// CHECK:   [[UPDATE_STRIDE:%.*]] = llvm.mul [[CUR_STRIDE]], [[SIZE]] : !llvm.i64
+// CHECK:   [[STRIDE_COND:%.*]] = llvm.sub [[DIM]], [[C1]] : !llvm.i64
+// CHECK:   llvm.br ^bb1([[STRIDE_COND]], [[UPDATE_STRIDE]] : !llvm.i64, !llvm.i64)
+
+// CHECK: ^bb3:
+// CHECK:   llvm.return

mlir/test/mlir-cpu-runner/memref_reshape.mlir

func @main() -> () {
   store %c3, %shape[%c0] : memref<2xindex>
   store %c2, %shape[%c1] : memref<2xindex>
 
   // Test cases.
   call @reshape_ranked_memref_to_ranked(%input, %shape)
     : (memref<2x3xf32>, memref<2xindex>) -> ()
   call @reshape_unranked_memref_to_ranked(%input, %shape)
     : (memref<2x3xf32>, memref<2xindex>) -> ()
+  call @reshape_ranked_memref_to_unranked(%input, %shape)
+    : (memref<2x3xf32>, memref<2xindex>) -> ()
+  call @reshape_unranked_memref_to_unranked(%input, %shape)
+    : (memref<2x3xf32>, memref<2xindex>) -> ()
   return
 }
 
 func @reshape_ranked_memref_to_ranked(%input : memref<2x3xf32>,
                                       %shape : memref<2xindex>) {
   %output = memref_reshape %input(%shape)
                 : (memref<2x3xf32>, memref<2xindex>) -> memref<?x?xf32>
 
   %unranked_output = memref_cast %output : memref<?x?xf32> to memref<*xf32>
   call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
   // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
-  // CHECK: [0,   1],                          
-  // CHECK: [2,   3],                          
-  // CHECK: [4,   5]  
+  // CHECK: [0,   1],
+  // CHECK: [2,   3],
+  // CHECK: [4,   5]
   return
 }
 
 func @reshape_unranked_memref_to_ranked(%input : memref<2x3xf32>,
                                         %shape : memref<2xindex>) {
   %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32>
   %output = memref_reshape %input(%shape)
                 : (memref<2x3xf32>, memref<2xindex>) -> memref<?x?xf32>
 
   %unranked_output = memref_cast %output : memref<?x?xf32> to memref<*xf32>
   call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
   // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
-  // CHECK: [0,   1],                          
-  // CHECK: [2,   3],                          
-  // CHECK: [4,   5]  
+  // CHECK: [0,   1],
+  // CHECK: [2,   3],
+  // CHECK: [4,   5]
   return
 }
+
+func @reshape_ranked_memref_to_unranked(%input : memref<2x3xf32>,
+                                        %shape : memref<2xindex>) {
+  %dyn_size_shape = memref_cast %shape : memref<2xindex> to memref<?xindex>
+  %output = memref_reshape %input(%dyn_size_shape)
+                : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32>
+
+  call @print_memref_f32(%output) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
+  // CHECK: [0,   1],
+  // CHECK: [2,   3],
+  // CHECK: [4,   5]
+  return
+}
+
+func @reshape_unranked_memref_to_unranked(%input : memref<2x3xf32>,
+                                          %shape : memref<2xindex>) {
+  %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32>
+  %dyn_size_shape = memref_cast %shape : memref<2xindex> to memref<?xindex>
+  %output = memref_reshape %input(%dyn_size_shape)
+                : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32>
+
+  call @print_memref_f32(%output) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
+  // CHECK: [0,   1],
+  // CHECK: [2,   3],
+  // CHECK: [4,   5]
+  return
+}