[mlir] Add MemRefReshapeOp definition to Standard.

[pifon2a]

Summary

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

[herhut]

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

2035

There is also a restriction on the number of elements, if known, right?

2038

nit: ranked or ranked?

2052

nit: here too.

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

2152

You have extraClassDeclaration so why not move this there, too?

2173

Is this tested somewhere?

[ftynse]

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

2038

Nit: s/length/size

2043–2044

I'm not sure why we need this version. The shape of the result is known from the type, so this is essentially asserting that the shape argument contains the same values as the result type.

2072

We don't do OpBuilder, OperationState here anymore.

2084–2085

Nit: would functional-type(operands, results) work here?

[pifon2a]

Address the comments.

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

2035

Right. In the design doc @ftynse had the idea of having a separate pass that would add std.asserts to ensure the same number of elements.

2043–2044

I am not sure what you mean here. Do you mean that there is no need to print the shape if the result is statically-shaped and %shape is const? We need a case when the shape arg is non-constant. I'll add some examples.

2072

that's much nicer

[herhut]

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

2043–2044

I guess the idea would be to drop the shape operand if it is static anyway. That would give us a static and dynamic version of this operation. Could be worthwhile, especially as we otherwise have to allocate a memref with static contents.
One could also envision ranked/unranked versions that take scalars vs. memrefs.

I'd say we start with the most general form and can add refinements later if they turn out to be useful.

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

let extraClassDeclaration = [{
     static bool areCastCompatible(Type a, Type b);
 
     /// The result of a memref_cast is always a memref.
     Type getType() { return getResult().getType(); }
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// MemRefReshapeOp
+//===----------------------------------------------------------------------===//
+
+def MemRefReshapeOp: Std_Op<"memref_reshape", [
+    DeclareOpInterfaceMethods<ViewLikeOpInterface>,
+    NoSideEffect]>  {
+  let summary = "memref reshape operation";
+  let description = [{
+    The `memref_reshape` operation converts a memref from one type to an
+    equivalent type with a provided shape. The data is never copied or
+    modified.  The source and destination types are compatible if both have the
+    same element type, address space and identity layout map. The following
+    combinations are possible:
+
+    a. Source type is ranked or ranked. Shape argument has static length.
+    Result type is ranked.
+
+    ```mlir
+    // Reshape statically-shaped memref.
+    %dst = memref_reshape %src(%shape)
+             : (memref<4x1xf32>, memref<1xi32>) to memref<4xf32>
+    %dst0 = memref_reshape %src(%shape0)
+             : (memref<4x1xf32>, memref<2xi32>) to memref<2x2xf32>
+    // Flatten unranked memref.
+    %dst = memref_reshape %src(%shape)
+             : (memref<*xf32>, memref<1xi32>) to memref<?xf32>
+    ```
+
+    a. Source type is ranked or ranked. Shape argument has dynamic length.
+    Result type is unranked.
+
+    ```mlir
+    // Reshape dynamically-shaped 1D memref.
+    %dst = memref_reshape %src(%shape)
+             : (memref<?xf32>, memref<?xi32>) to memref<*xf32>
+    // Reshape unranked memref.
+    %dst = memref_reshape %src(%shape)
+             : (memref<*xf32>, memref<?xi32>) to memref<*xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    AnyRankedOrUnrankedMemRef:$source,
+    MemRefRankOf<[AnySignlessInteger], [1]>:$shape
+  );
+  let results = (outs AnyRankedOrUnrankedMemRef:$result);
+
+  let builders = [OpBuilder<
+    "OpBuilder &builder, OperationState &result, MemRefType resultType, " #
+    "Value operand, Value shape", [{
+       result.addOperands(operand);
+       result.addOperands(shape);
+       result.types.push_back(resultType);
+     }]>];
+
+  let extraClassDeclaration = [{
+    MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
+  }];
+
+  let assemblyFormat = [{
+    $source `(` $shape `)` attr-dict `:` `(` type($source) `,` type($shape)
+    `)` `->` type($result)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 def MulFOp : FloatArithmeticOp<"mulf"> {
   let summary = "floating point multiplication operation";
   let description = [{
     Syntax:
 

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

bool MemRefCastOp::areCastCompatible(Type a, Type b) {
   return false;
 }
 
 OpFoldResult MemRefCastOp::fold(ArrayRef<Attribute> operands) {
   return impl::foldCastOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// MemRefReshapeOp
+//===----------------------------------------------------------------------===//
+
+Value MemRefReshapeOp::getViewSource() { return source(); }
+
+static LogicalResult verify(MemRefReshapeOp op) {
+  Type operandType = op.source().getType();
+  Type resultType = op.result().getType();
+
+  Type operandElementType = operandType.cast<ShapedType>().getElementType();
+  Type resultElementType = resultType.cast<ShapedType>().getElementType();
+  if (operandElementType != resultElementType)
+    return op.emitOpError("element types of source and destination memref "
+                          "types should be the same");
+
+  if (auto operandMemRefType = operandType.dyn_cast<MemRefType>())
+    if (!operandMemRefType.getAffineMaps().empty())
+      return op.emitOpError(
+          "source memref type should have identity affine map");
+
+  int64_t shapeSize = op.shape().getType().cast<MemRefType>().getDimSize(0);
+  auto resultMemRefType = resultType.dyn_cast<MemRefType>();
+  if (resultMemRefType) {
+    if (!resultMemRefType.getAffineMaps().empty())
+      return op.emitOpError(
+          "result memref type should have identity affine map");
+    if (shapeSize == ShapedType::kDynamicSize)
+      return op.emitOpError("cannot use shape operand with dynamic length to "
+                            "reshape to statically-ranked memref type");
+    if (shapeSize != resultMemRefType.getRank())
+      return op.emitOpError(
+          "length of shape operand differs from the result's memref rank");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MulFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOp<FloatAttr>(
       operands, [](APFloat a, APFloat b) { return a * b; });
 }
 

mlir/test/Dialect/Standard/invalid.mlir

 }
 
 // -----
 
 func @transpose_wrong_type(%v : memref<?x?xf32, affine_map<(i, j)[off, M]->(off + M * i + j)>>) {
   // expected-error @+1 {{output type 'memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>>' does not match transposed input type 'memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>>'}}
   transpose %v (i, j) -> (j, i) : memref<?x?xf32, affine_map<(i, j)[off, M]->(off + M * i + j)>> to memref<?x?xf32, affine_map<(i, j)[off, M]->(off + M * i + j)>>
 }
+
+// -----
+
+// CHECK-LABEL: memref_reshape_element_type_mismatch
+func @memref_reshape_element_type_mismatch(
+       %buf: memref<*xf32>, %shape: memref<1xi32>) {
+  // expected-error @+1 {{element types of source and destination memref types should be the same}}
+  memref_reshape %buf(%shape) : (memref<*xf32>, memref<1xi32>) -> memref<?xi32>
+}
+
+// -----
+
+// CHECK-LABEL: memref_reshape_dst_ranked_shape_unranked
+func @memref_reshape_dst_ranked_shape_unranked(
+       %buf: memref<*xf32>, %shape: memref<?xi32>) {
+  // expected-error @+1 {{cannot use shape operand with dynamic length to reshape to statically-ranked memref type}}
+  memref_reshape %buf(%shape) : (memref<*xf32>, memref<?xi32>) -> memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: memref_reshape_dst_shape_rank_mismatch
+func @memref_reshape_dst_shape_rank_mismatch(
+       %buf: memref<*xf32>, %shape: memref<1xi32>) {
+  // expected-error @+1 {{length of shape operand differs from the result's memref rank}}
+  memref_reshape %buf(%shape)
+    : (memref<*xf32>, memref<1xi32>) -> memref<?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: memref_reshape_affine_map_is_not_identity
+func @memref_reshape_affine_map_is_not_identity(
+        %buf: memref<4x4xf32, offset: 0, strides: [3, 2]>,
+        %shape: memref<1xi32>) {
+  // expected-error @+1 {{source memref type should have identity affine map}}
+  memref_reshape %buf(%shape)
+    : (memref<4x4xf32, offset: 0, strides: [3, 2]>, memref<1xi32>)
+    -> memref<8xf32>
+}

mlir/test/Dialect/Standard/ops.mlir

func @atan(%arg : f32) -> f32 {
   return %result : f32
 }
 
 // CHECK-LABEL: @atan2
 func @atan2(%arg0 : f32, %arg1 : f32) -> f32 {
   %result = atan2 %arg0, %arg1 : f32
   return %result : f32
 }
+
+// CHECK-LABEL: func @memref_reshape(
+func @memref_reshape(%unranked: memref<*xf32>, %shape1: memref<1xi32>,
+         %shape2: memref<2xi32>, %shape3: memref<?xi32>) -> memref<*xf32> {
+  %dyn_vec = memref_reshape %unranked(%shape1)
+               : (memref<*xf32>, memref<1xi32>) -> memref<?xf32>
+  %dyn_mat = memref_reshape %dyn_vec(%shape2)
+               : (memref<?xf32>, memref<2xi32>) -> memref<?x?xf32>
+  %new_unranked = memref_reshape %dyn_mat(%shape3)
+               : (memref<?x?xf32>, memref<?xi32>) -> memref<*xf32>
+  return %new_unranked : memref<*xf32>
+}