[mlir] Convert memref_reshape to memref_reinterpret_cast.

[pifon2a]

Summary

I had to remove some checks in MemRefReinterpretCastOp::verify. It might make
sense to convert to LLVM directly to avoid that.

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

[pifon2a]

Nit

[ekatz]

mlir/lib/Dialect/StandardOps/Transforms/ExpandMemRefReshape.cpp

27

This comment is incorrect

[pifon2a]

Fix the comment

[herhut]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2270

Missing space after comma.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2293

Why use an assign here and not construct directly?

2395

Can you explain why these checks needed to go?

[pifon2a]

Address the comments, add integration test.

[pifon2a]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2395

I put the checks back, but only when the result has an affine map specified.

What was happening is the following:

when

%result = memref_reshape %input(%shape)
              : (memref<*xf32>, memref<3xi32>) -> memref<?x?x?xf32>

is transformed into

%result = memref_reinterpret_cast %input
   to offset: [0], sizes: [%s0, %s1, %s2],
   strides: [%str0, %str1, %c1] : memref<*xf32> to memref<?x?x?xf32>

the verifier would fail because it would try to compare strides and offsets in the affine map of the result.

[herhut]

Please address comments, otherwise lgtm.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2434–2458

Remove tabs and the comment could state that this checks whether the offsets and strides match the map in the type, if any is present.

[pifon2a]

Address the comments.

[pifon2a]

Remove tabs

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2434–2458

those are not the tabs :) I think that's how diff is rendered here.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let arguments = (ins
     Variadic<Index>:$offsets,
     Variadic<Index>:$sizes,
     Variadic<Index>:$strides,
     I64ArrayAttr:$static_offsets,
     I64ArrayAttr:$static_sizes,
     I64ArrayAttr:$static_strides
   );
   let results = (outs AnyMemRef:$result);
+
+  let builders = [
+    // Build a ReinterpretCastOp with mixed static and dynamic entries.
+    OpBuilder<
+      "MemRefType resultType, Value source, int64_t staticOffset, "
+      "ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides, "
+      "ValueRange offset, ValueRange sizes, ValueRange strides, "
+      "ArrayRef<NamedAttribute> attrs = {}">,
+    // Build a ReinterpretCastOp with all dynamic entries.
+    OpBuilder<
+      "MemRefType resultType, Value source, Value offset, ValueRange sizes, "
+      "ValueRange strides, ArrayRef<NamedAttribute> attrs = {}">,
+  ];
+
   let extraClassDeclaration = extraBaseClassDeclaration # [{
     // The result of the op is always a ranked memref.
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
     Value getViewSource() { return source(); }
   }];
 }
 
 //===----------------------------------------------------------------------===//
let description = [{
     // Reshape unranked memref.
     %dst = memref_reshape %src(%shape)
              : (memref<*xf32>, memref<?xi32>) to memref<*xf32>
     ```
   }];
 
   let arguments = (ins
     AnyRankedOrUnrankedMemRef:$source,
-    MemRefRankOf<[AnySignlessInteger], [1]>:$shape
+    MemRefRankOf<[AnySignlessInteger, Index], [1]>:$shape
   );
   let results = (outs AnyRankedOrUnrankedMemRef:$result);
 
   let builders = [OpBuilderDAG<
      (ins "MemRefType":$resultType, "Value":$operand, "Value":$shape), [{
        $_state.addOperands(operand);
        $_state.addOperands(shape);
        $_state.addTypes(resultType);

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

 
 namespace mlir {
 
 class OwningRewritePatternList;
 
 /// Creates an instance of the ExpandAtomic pass.
 std::unique_ptr<Pass> createExpandAtomicPass();
 
+void populateExpandMemRefReshapePattern(OwningRewritePatternList &patterns,
+                                        MLIRContext *ctx);
+
 void populateExpandTanhPattern(OwningRewritePatternList &patterns,
                                MLIRContext *ctx);
 
 void populateStdBufferizePatterns(MLIRContext *context,
                                   BufferizeTypeConverter &typeConverter,
                                   OwningRewritePatternList &patterns);
 
 /// Creates an instance of std bufferization pass.

mlir/lib/Conversion/StandardToLLVM/CMakeLists.txt

 add_mlir_conversion_library(MLIRStandardToLLVM
   StandardToLLVM.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/StandardToLLVM
 
   DEPENDS
   MLIRConversionPassIncGen
   intrinsics_gen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRLLVMIR
+  MLIRStandardOpsTransforms
   MLIRTransforms
   )

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LogicalResult.h"
else
     patterns.insert<FuncOpConversion>(converter);
 }
 
 void mlir::populateStdToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   populateStdToLLVMFuncOpConversionPattern(converter, patterns);
   populateStdToLLVMNonMemoryConversionPatterns(converter, patterns);
   populateStdToLLVMMemoryConversionPatterns(converter, patterns);
+  populateExpandMemRefReshapePattern(patterns, &converter.getContext());
 }
 
 /// Convert a non-empty list of types to be returned from a function into a
 /// supported LLVM IR type.  In particular, if more than one value is returned,
 /// create an LLVM IR structure type with elements that correspond to each of
 /// the MLIR types converted with `convertType`.
 Type LLVMTypeConverter::packFunctionResults(ArrayRef<Type> types) {
   assert(!types.empty() && "expected non-empty list of type");

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 OpFoldResult MemRefCastOp::fold(ArrayRef<Attribute> operands) {
   return impl::foldCastOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefReinterpretCastOp
 //===----------------------------------------------------------------------===//
 
+void mlir::MemRefReinterpretCastOp::build(
+    OpBuilder &b, OperationState &result, MemRefType resultType, Value source,
+    int64_t staticOffset, ArrayRef<int64_t> staticSizes,
+    ArrayRef<int64_t> staticStrides, ValueRange offset, ValueRange sizes,
+    ValueRange strides, ArrayRef<NamedAttribute> attrs) {
+  build(b, result, resultType, source, offset, sizes, strides,
+        b.getI64ArrayAttr(staticOffset), b.getI64ArrayAttr(staticSizes),
+        b.getI64ArrayAttr(staticStrides));
+  result.addAttributes(attrs);
+}
+
+/// Build a MemRefReinterpretCastOp with all dynamic entries: `staticOffsets`,
+/// `staticSizes` and `staticStrides` are  automatically filled with
+/// source-memref-rank sentinel values that encode dynamic entries.
+void mlir::MemRefReinterpretCastOp::build(OpBuilder &b, OperationState &result,
+                                          MemRefType resultType, Value source,
+                                          Value offset, ValueRange sizes,
+                                          ValueRange strides,
+                                          ArrayRef<NamedAttribute> attrs) {
+  unsigned rank = resultType.getRank();
+  SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamicSize);
+  SmallVector<int64_t, 4> staticStridesVector(
+      rank, ShapedType::kDynamicStrideOrOffset);
+  build(b, result, resultType, source,
+        /*staticOffset=*/ShapedType::kDynamicStrideOrOffset, staticSizesVector,
+        staticStridesVector, offset, sizes, strides, attrs);
+}
+
 /// Print of the form:
 /// ```
 ///   `name` ssa-name to
 ///       offset: `[` offset `]`
 ///       sizes: `[` size-list `]`
 ///       strides:`[` stride-list `]`
 ///   `:` any-memref-type to strided-memref-type
 /// ```
if (failed(verifyOpWithOffsetSizesAndStridesPart(
           op.static_sizes(), ShapedType::isDynamic, op.sizes())))
     return failure();
   if (failed(verifyOpWithOffsetSizesAndStridesPart(
           op, "stride", resultRank, op.getStaticStridesAttrName(),
           op.static_strides(), ShapedType::isDynamicStrideOrOffset,
           op.strides())))
     return failure();
 
-  // Extract source offset and strides.
-  int64_t resultOffset;
-  SmallVector<int64_t, 4> resultStrides;
-  if (failed(getStridesAndOffset(resultType, resultStrides, resultOffset)))
-    return failure();
-
-  // Match offset in result memref type and in static_offsets attribute.
-  int64_t expectedOffset = extractFromI64ArrayAttr(op.static_offsets()).front();
-  if (resultOffset != expectedOffset)
-    return op.emitError("expected result type with offset = ")
-           << resultOffset << " instead of " << expectedOffset;
-
   // Match sizes in result memref type and in static_sizes attribute.
   for (auto &en :
        llvm::enumerate(llvm::zip(resultType.getShape(),
                                  extractFromI64ArrayAttr(op.static_sizes())))) {
     int64_t resultSize = std::get<0>(en.value());
     int64_t expectedSize = std::get<1>(en.value());
     if (resultSize != expectedSize)
       return op.emitError("expected result type with size = ")
              << expectedSize << " instead of " << resultSize
              << " in dim = " << en.index();
   }
 
+  // Match offset and strides in static_offset and static_strides attributes if
+  // result memref type has an affine map specified.
+  if (!resultType.getAffineMaps().empty()) {
+    int64_t resultOffset;
+    SmallVector<int64_t, 4> resultStrides;
+    if (failed(getStridesAndOffset(resultType, resultStrides, resultOffset)))
+      return failure();
+
+    // Match offset in result memref type and in static_offsets attribute.
+    int64_t expectedOffset =
+        extractFromI64ArrayAttr(op.static_offsets()).front();
+    if (resultOffset != expectedOffset)
+      return op.emitError("expected result type with offset = ")
+             << resultOffset << " instead of " << expectedOffset;
+
     // Match strides in result memref type and in static_strides attribute.
     for (auto &en : llvm::enumerate(llvm::zip(
              resultStrides, extractFromI64ArrayAttr(op.static_strides())))) {
       int64_t resultStride = std::get<0>(en.value());
       int64_t expectedStride = std::get<1>(en.value());
       if (resultStride != expectedStride)
         return op.emitError("expected result type with stride = ")
                << expectedStride << " instead of " << resultStride
                << " in dim = " << en.index();
     }
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefReshapeOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(MemRefReshapeOp op) {

mlir/lib/Dialect/StandardOps/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
   ExpandAtomic.cpp
+  ExpandMemRefReshape.cpp
   ExpandTanh.cpp
   FuncConversions.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS
   MLIRStandardTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRSCF
   MLIRStandard
   MLIRTransforms
   )

mlir/lib/Dialect/StandardOps/Transforms/ExpandMemRefReshape.cpp

+//===- ExpandMemRefReshape.cpp - Code to perform expanding memref_reshape -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements expansion of MemRefReshapeOp into
+// MemRefReinterpretCastOp.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+
+/// Converts `memref_reshape` that has a target shape of a statically-known
+/// size to `memref_reinterpret_cast`.
+struct MemRefReshapeOpConverter : public OpRewritePattern<MemRefReshapeOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MemRefReshapeOp op,
+                                PatternRewriter &rewriter) const final {
+    auto shapeType = op.shape().getType().cast<MemRefType>();
+    if (!shapeType.hasStaticShape())
+      return failure();
+
+    int64_t rank = shapeType.cast<MemRefType>().getDimSize(0);
+    SmallVector<Value, 4> sizes, strides;
+    sizes.resize(rank);
+    strides.resize(rank);
+
+    Location loc = op.getLoc();
+    Value stride = rewriter.create<ConstantIndexOp>(loc, 1);
+    for (int i = rank - 1; i >= 0; --i) {
+      Value index = rewriter.create<ConstantIndexOp>(loc, i);
+      Value size = rewriter.create<LoadOp>(loc, op.shape(), index);
+      if (!size.getType().isa<IndexType>())
+        size = rewriter.create<IndexCastOp>(loc, size, rewriter.getIndexType());
+      sizes[i] = size;
+      strides[i] = stride;
+      if (i > 0)
+        stride = rewriter.create<MulIOp>(loc, stride, size);
+    }
+    SmallVector<int64_t, 2> staticSizes(rank, ShapedType::kDynamicSize);
+    SmallVector<int64_t, 2> staticStrides(rank,
+                                          ShapedType::kDynamicStrideOrOffset);
+    rewriter.replaceOpWithNewOp<MemRefReinterpretCastOp>(
+        op, op.getType(), op.source(), /*staticOffset = */ 0, staticSizes,
+        staticStrides, /*offset=*/llvm::None, sizes, strides);
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::populateExpandMemRefReshapePattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  patterns.insert<MemRefReshapeOpConverter>(ctx);
+}

mlir/test/Dialect/Standard/expand-memref-reshape.mlir

+// RUN: mlir-opt %s -test-expand-memref-reshape | FileCheck %s
+
+// CHECK-LABEL: func @memref_reshape(
+func @memref_reshape(%input: memref<*xf32>,
+                     %shape: memref<3xi32>) -> memref<?x?x?xf32> {
+  %result = memref_reshape %input(%shape)
+               : (memref<*xf32>, memref<3xi32>) -> memref<?x?x?xf32>
+  return %result : memref<?x?x?xf32>
+}
+// CHECK-SAME: [[SRC:%.*]]: memref<*xf32>,
+// CHECK-SAME: [[SHAPE:%.*]]: memref<3xi32>) -> memref<?x?x?xf32> {
+// CHECK: [[C2:%.*]] = constant 2 : index
+// CHECK: [[C1:%.*]] = constant 1 : index
+// CHECK: [[C0:%.*]] = constant 0 : index
+// CHECK: [[DIM_2:%.*]] = load [[SHAPE]]{{\[}}[[C2]]] : memref<3xi32>
+// CHECK: [[SIZE_2:%.*]] = index_cast [[DIM_2]] : i32 to index
+// CHECK: [[DIM_1:%.*]] = load [[SHAPE]]{{\[}}[[C1]]] : memref<3xi32>
+// CHECK: [[SIZE_1:%.*]] = index_cast [[DIM_1]] : i32 to index
+// CHECK: [[STRIDE_0:%.*]] = muli [[SIZE_2]], [[SIZE_1]] : index
+// CHECK: [[DIM_0:%.*]] = load [[SHAPE]]{{\[}}[[C0]]] : memref<3xi32>
+// CHECK: [[SIZE_0:%.*]] = index_cast [[DIM_0]] : i32 to index
+
+// CHECK: [[RESULT:%.*]] = memref_reinterpret_cast [[SRC]]
+// CHECK-SAME: to offset: [0], sizes: {{\[}}[[SIZE_0]], [[SIZE_1]], [[SIZE_2]]],
+// CHECK-SAME: strides: {{\[}}[[STRIDE_0]], [[SIZE_2]], [[C1]]]
+// CHECK-SAME: : memref<*xf32> to memref<?x?x?xf32>

mlir/test/Dialect/Standard/invalid.mlir

%out = memref_reinterpret_cast %in to
          : memref<*xf32> to memref<10xi32, offset: 0, strides: [1], 2>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_reinterpret_cast_offset_mismatch
 func @memref_reinterpret_cast_offset_mismatch(%in: memref<?xf32>) {
-  // expected-error @+1 {{expected result type with offset = 0 instead of 1}}
+  // expected-error @+1 {{expected result type with offset = 2 instead of 1}}
   %out = memref_reinterpret_cast %in to
            offset: [1], sizes: [10], strides: [1]
-         : memref<?xf32> to memref<10xf32>
+         : memref<?xf32> to memref<10xf32, offset: 2, strides: [1]>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_reinterpret_cast_size_mismatch
 func @memref_reinterpret_cast_size_mismatch(%in: memref<*xf32>) {
   // expected-error @+1 {{expected result type with size = 10 instead of 1 in dim = 0}}
   %out = memref_reinterpret_cast %in to
            offset: [0], sizes: [10], strides: [1]
          : memref<*xf32> to memref<1xf32, offset: 0, strides: [1]>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_reinterpret_cast_stride_mismatch
 func @memref_reinterpret_cast_offset_mismatch(%in: memref<?xf32>) {
   // expected-error @+1 {{expected result type with stride = 2 instead of 1 in dim = 0}}
   %out = memref_reinterpret_cast %in to
-           offset: [0], sizes: [10], strides: [2]
-         : memref<?xf32> to memref<10xf32>
+           offset: [2], sizes: [10], strides: [2]
+         : memref<?xf32> to memref<10xf32, offset: 2, strides: [1]>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_reinterpret_cast_dynamic_size_mismatch
 func @memref_reinterpret_cast_offset_mismatch(%in: memref<?xf32>) {
   %c0 = constant 0 : index

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAffineLoopParametricTiling.cpp
   TestBufferPlacement.cpp
+  TestExpandMemRefReshape.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvVectorization.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDominance.cpp

mlir/test/lib/Transforms/TestExpandMemRefReshape.cpp

+//===- TestExpandMemRefReshape.cpp - Test expansion of memref_reshape -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for expanding memref reshape.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestExpandMemRefReshapePass
+    : public PassWrapper<TestExpandMemRefReshapePass, FunctionPass> {
+  void runOnFunction() override;
+};
+} // end anonymous namespace
+
+void TestExpandMemRefReshapePass::runOnFunction() {
+  OwningRewritePatternList patterns;
+  populateExpandMemRefReshapePattern(patterns, &getContext());
+  applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+
+namespace mlir {
+void registerTestExpandMemRefReshapePass() {
+  PassRegistration<TestExpandMemRefReshapePass> pass(
+      "test-expand-memref-reshape", "Test expanding memref reshape");
+}
+} // namespace mlir

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

+// RUN: mlir-opt %s -convert-scf-to-std -convert-std-to-llvm --print-ir-after-all \
+// RUN: | mlir-cpu-runner -e main -entry-point-result=void \
+// RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
+// RUN: | FileCheck %s
+
+
+func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
+
+func @main() -> () {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+
+  // Initialize input.
+  %input = alloc() : memref<2x3xf32>
+  %dim_x = dim %input, %c0 : memref<2x3xf32>
+  %dim_y = dim %input, %c1 : memref<2x3xf32>
+  scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) {
+    %prod = muli %i,  %dim_y : index
+    %val = addi %prod, %j : index
+    %val_i64 = index_cast %val : index to i64
+    %val_f32 = sitofp %val_i64 : i64 to f32
+    store %val_f32, %input[%i, %j] : memref<2x3xf32>
+  }
+  %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32>
+  call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
+  // CHECK-NEXT: [0,   1,   2]
+  // CHECK-NEXT: [3,   4,   5]
+
+  // Initialize shape.
+  %shape = alloc() : memref<2xindex>
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  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>) -> ()
+  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]  
+  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]  
+  return
+}

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvVectorization();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestDialect(DialectRegistry &);
 void registerTestDynamicPipelinePass();
+void registerTestExpandMemRefReshapePass();
 void registerTestExpandTanhPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgHoisting();
#if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
   registerTestAffineLoopParametricTilingPass();
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
   registerTestDynamicPipelinePass();
   registerTestFunc();
   registerTestExpandTanhPass();
+  registerTestExpandMemRefReshapePass();
   registerTestGpuMemoryPromotionPass();
   registerTestInterfaces();
   registerTestLinalgCodegenStrategy();
   registerTestLinalgFusionTransforms();
   registerTestLinalgHoisting();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();