[mlir] Add `ReshapeMemrefCastOp` to Standard.

[pifon2a]

Summary

The op allows to reshape memrefs using a 1D memref of extents. The shape memref
can also be dynamic.

[mehdi_amini]

Can you make sure there is a thread on discourse before adding ops to the standard dialect?

[herhut]

Thanks, some comments...

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

1839

Please keep the file sorted.

1875

Why is unranked to unranked not supported?

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

1911

Nit: You could also use isa<UnrankedMemRefType> here as you do not need the value.

1915

Nit: This is not the rank but the length of the shape operand.

1919

Maybe cannot use shape operand with dynamic length to cast ...?

1922

Maybe length of shape operand differs from the result's memref rank?

[pifon2a]

In D83068#2129809, @mehdi_amini wrote:

Can you make sure there is a thread on discourse before adding ops to the standard dialect?

https://llvm.discourse.group/t/standardops-add-reshapememrefcastop-to-standard/1350/2

[pifon2a]

Thank you!

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

1875

I really need to see a use case for that. We can always relax this requirement later.

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

1915

which is the rank of the resulting memref.

[pifon2a]

Addressed the comments.

[pifon2a]

Remove verification for unranked -> unranked memrefs.

[pifon2a]

Address the comments in https://llvm.discourse.group/t/standardops-add-reshapememrefcastop-to-standard/1350/2

[pifon2a]

x index -> x i32

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(); }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 def MulFOp : FloatArithmeticOp<"mulf"> {
   let summary = "floating point multiplication operation";
   let description = [{
     Syntax:
 
     ```
let description = [{
     ```
 
     TODO: In the distant future, this will accept optional attributes for fast
     math, contraction, rounding mode, and other controls.
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MulIOp
 //===----------------------------------------------------------------------===//
 
 def MulIOp : IntArithmeticOp<"muli", [Commutative]> {
   let summary = "integer multiplication operation";
   let hasFolder = 1;
 }
 

 // RemFOp
 //===----------------------------------------------------------------------===//
 
 def RemFOp : FloatArithmeticOp<"remf"> {
   let summary = "floating point division remainder operation";
 }
 
 //===----------------------------------------------------------------------===//
+// ReshapeMemRefCastOp
+//===----------------------------------------------------------------------===//
+
+def ReshapeMemRefCastOp: Std_Op<"reshape_memref_cast", [
+    DeclareOpInterfaceMethods<ViewLikeOpInterface>,
+    NoSideEffect]>  {
+  let summary = "reshape memref cast operation";
+  let description = [{
+    The `reshape_memref_cast` operation converts a memref from one type to an
+    equivalent type with a provided shape. The data is never copied or moded.
+    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. Both are ranked memref types.
+
+    ```mlir
+    // Reshape statically-shaped memref.
+    %dst = reshape_memref_cast %src(%shape)
+             : (memref<4x1xf32>, memref<1xi32>) to memref<4xf32>
+    %dst0 = reshape_memref_cast %src(%shape0)
+             : (memref<4x1xf32>, memref<2xi32>) to memref<2x2xf32>
+    ```
+
+    b. Source type is ranked, destination type is unranked.
+
+    ```mlir
+    // Reshape dynamically-shaped 1D memref.
+    %dst = reshape_memref_cast %src(%shape)
+             : (memref<?xf32>, memref<?xi32>) to memref<*xf32>
+    ```
+
+    c. Source type is unranked, destination type is ranked.
+
+    ```mlir
+    // Flatten unranked memref.
+    %dst = reshape_memref_cast %src(%shape)
+             : (memref<*xf32>, memref<1xi32>) to memref<?xf32>
+    ```
+
+    d. Both are unranked memref types.
+
+    ```mlir
+    // Reshape unranked memref.
+    %dst = reshape_memref_cast %src(%shape)
+             : (memref<*xf32>, memref<?xi32>) to memref<*xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    AnyRankedOrUnrankedMemRef:$operand,
+    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 = [{
+    $operand `(` $shape `)` attr-dict `:` `(` type($operand) `,` type($shape) 
+    `)` `->` type($result)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">, ReturnLike,
                                  Terminator]> {
   let summary = "return operation";
   let description = [{
     The `return` operation represents a return operation within a function.

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

 OpFoldResult MulFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOp<FloatAttr>(
       operands, [](APFloat a, APFloat b) { return a * b; });
 }
 
 //===----------------------------------------------------------------------===//
 // MulIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MulIOp::fold(ArrayRef<Attribute> operands) {
   /// muli(x, 0) -> 0
   if (matchPattern(rhs(), m_Zero()))
     return rhs();
   /// muli(x, 1) -> x
   if (matchPattern(rhs(), m_One()))
     return getOperand(0);
 
   // TODO: Handle the overflow case.
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a * b; });
 }
 
 //===----------------------------------------------------------------------===//
 // OrOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult OrOp::fold(ArrayRef<Attribute> operands) {
   /// or(x, 0) -> x
OpFoldResult RankOp::fold(ArrayRef<Attribute> operands) {
   // Constant fold rank when the rank of the tensor is known.
   auto type = getOperand().getType();
   if (auto tensorType = type.dyn_cast<RankedTensorType>())
     return IntegerAttr::get(IndexType::get(getContext()), tensorType.getRank());
   return IntegerAttr();
 }
 
 //===----------------------------------------------------------------------===//
+// ReshapeMemrefCastOp
+//===----------------------------------------------------------------------===//
+
+Value ReshapeMemRefCastOp::getViewSource() { return operand(); }
+
+static LogicalResult verify(ReshapeMemRefCastOp op) {
+  Type operandType = op.operand().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(
+          "operand 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 (shapeSize == ShapedType::kDynamicSize)
+      return op.emitOpError("cannot use shape operand with dynamic length to "
+                            "cast statically-ranked memref type");
+    if (shapeSize != resultMemRefType.getRank())
+      return op.emitOpError(
+          "length of shape operand differs from the result's memref rank");
+
+    if (!resultMemRefType.getAffineMaps().empty())
+      return op.emitOpError(
+          "result memref type should have identity affine map");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReturnOp op) {
   auto function = cast<FuncOp>(op.getParentOp());
 
   // The operand number and types must match the function signature.
   const auto &results = function.getType().getResults();

mlir/test/Dialect/Standard/invalid.mlir

 // -----
 
 // CHECK-LABEL: test_index_cast_tensor_error
 func @test_index_cast_tensor_error(%arg0 : tensor<index>) -> i64 {
   // expected-error @+1 {{requires the same shape for all operands and results}}
   %0 = index_cast %arg0 : tensor<index> to i64
   return %0 : i64
 }
+
+// -----
+
+// CHECK-LABEL: reshape_memref_cast_element_type_mismatch
+func @reshape_memref_cast_element_type_mismatch(
+       %buf: memref<*xf32>, %shape: memref<1xi32>) {
+  // expected-error @+1 {{element types of source and destination memref types should be the same}}
+  reshape_memref_cast %buf(%shape)
+    : (memref<*xf32>, memref<1xi32>) -> memref<?xi32>
+}
+
+// -----
+
+// CHECK-LABEL: reshape_memref_cast_dst_ranked_shape_unranked
+func @reshape_memref_cast_dst_ranked_shape_unranked(
+       %buf: memref<*xf32>, %shape: memref<?xi32>) {
+  // expected-error @+1 {{cannot use shape operand with dynamic length to cast statically-ranked memref type}}
+  reshape_memref_cast %buf(%shape)
+    : (memref<*xf32>, memref<?xi32>) -> memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: reshape_memref_cast_dst_shape_rank_mismatch
+func @reshape_memref_cast_dst_shape_rank_mismatch(
+       %buf: memref<*xf32>, %shape: memref<1xi32>) {
+  // expected-error @+1 {{length of shape operand differs from the result's memref rank}}
+  reshape_memref_cast %buf(%shape)
+    : (memref<*xf32>, memref<1xi32>) -> memref<?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: reshape_memref_cast_affine_map_is_not_identity
+func @reshape_memref_cast_affine_map_is_not_identity(
+        %buf: memref<4x4xf32, offset: 0, strides: [3, 2]>,
+        %shape: memref<1xi32>) {
+  // expected-error @+1 {{operand memref type should have identity affine map}}
+  reshape_memref_cast %buf(%shape)
+    : (memref<4x4xf32, offset: 0, strides: [3, 2]>, memref<1xi32>)
+    -> memref<8xf32>
+}

mlir/test/IR/core-ops.mlir

func @memref_cast(%arg0: memref<4xf32>, %arg1 : memref<?xf32>, %arg2 : memref<64x16x4xf32, offset: 0, strides: [64, 4, 1]>) {
   // CHECK: memref_cast %{{.*}} : memref<4xf32> to memref<*xf32>
   %4 = memref_cast %1 : memref<4xf32> to memref<*xf32>
 
   // CHECK: memref_cast %{{.*}} : memref<*xf32> to memref<4xf32>
   %5 = memref_cast %4 : memref<*xf32> to memref<4xf32>
   return
 }
 
+// CHECK-LABEL: func @reshape_memref_cast(
+func @reshape_memref_cast(%unranked: memref<*xf32>, %shape1: memref<1xi32>,
+         %shape2: memref<2xi32>, %shape3: memref<?xi32>) {
+  // CHECK-SAME: [[UNRANKED:%.*]]: memref<*xf32>, [[SHAPE_1:%.*]]: memref<1xi32>,
+  // CHECK-SAME: [[SHAPE_2:%.*]]: memref<2xi32>, [[SHAPE_3:%.*]]: memref<?xi32>
+
+  // CHECK-NEXT: [[DYN_VEC:%.*]] = reshape_memref_cast [[UNRANKED]]
+  // CHECK-SAME:     : (memref<*xf32>, memref<1xi32>) -> memref<?xf32>
+  %dyn_vec = reshape_memref_cast %unranked(%shape1)
+               : (memref<*xf32>, memref<1xi32>) -> memref<?xf32>
+
+  // CHECK-NEXT: [[DYN_MAT:%.*]] = reshape_memref_cast [[DYN_VEC]]
+  // CHECK-SAME:     : (memref<?xf32>, memref<2xi32>) -> memref<?x?xf32>
+  %dyn_mat = reshape_memref_cast %dyn_vec(%shape2)
+               : (memref<?xf32>, memref<2xi32>) -> memref<?x?xf32>
+
+  // CHECK-NEXT: {{%.*}} = reshape_memref_cast [[DYN_MAT]]
+  // CHECK-SAME:     : (memref<?x?xf32>, memref<?xi32>) -> memref<*xf32>
+  %new_unranked = reshape_memref_cast %dyn_mat(%shape3)
+               : (memref<?x?xf32>, memref<?xi32>) -> memref<*xf32>
+  return
+}
+
 // CHECK-LABEL: func @memref_view(%arg0
 func @memref_view(%arg0 : index, %arg1 : index, %arg2 : index) {
   %0 = alloc() : memref<2048xi8>
   // Test two dynamic sizes and dynamic offset.
   // CHECK: %{{.*}} = std.view %0[%arg2][%arg0, %arg1] : memref<2048xi8> to memref<?x?xf32>
   %1 = view %0[%arg2][%arg0, %arg1] : memref<2048xi8> to memref<?x?xf32>
 
   // Test one dynamic size and dynamic offset.