[mlir] Add `tensor.reshape`.

[pifon2a]

Summary

This operation a counterpart of memref.reshape.

RFC Reshape Ops Restructuring

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

239

is this needed ? I don't see it declaring the whatever view interface would require this?

243

I don't think the result type matches your doc ?

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

453

Plz, verify the total number of elements when all is static.

[pifon2a]

Address the comments.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

205

Your IR examples and assembly format want

: (tensor<?xf32>, tensor<2xi32>) -> tensor<?x?xf32>

Either is fine with me but plz be consistent.

238

Nit: I probably would have called this getResultTensorType, your call.

[pifon2a]

More comments.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let builders = [
     OpBuilder<(ins "Type":$resultTy, "ValueRange":$dynamicExtents,
       "function_ref<void(OpBuilder &, Location, ValueRange)>")>,
   ];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// ReshapeOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_ReshapeOp: Tensor_Op<"reshape", [NoSideEffect]>  {
+  let summary = "tensor reshape operation";
+  let description = [{
+    The `reshape` operation converts a tensor from one type to an equivalent
+    type with a provided shape. The source and destination types are compatible
+    if both have the same element type, same number of elements. The following
+    combinations are possible:
+
+    a. Source type is ranked or unranked. Shape argument has static size.
+    Result type is ranked.
+
+    ```mlir
+    // Reshape statically-shaped tensor.
+    %dst = tensor.reshape %src(%shape)
+             : (tensor<4x1xf32>, tensor<1xi32>) -> tensor<4xf32>
+    %dst0 = tensor.reshape %src(%shape0)
+             : (tensor<4x1xf32>, tensor<2xi32>) -> tensor<2x2xf32>
+    // Flatten unranked tensor.
+    %dst = tensor.reshape %src(%shape)
+             : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xf32>
+    ```
+
+    b. Source type is ranked or unranked. Shape argument has dynamic size.
+    Result type is unranked.
+
+    ```mlir
+    // Reshape dynamically-shaped 1D tensor.
+    %dst = tensor.reshape %src(%shape)
+             : (tensor<?xf32>, tensor<?xi32>) -> tensor<*xf32>
+    // Reshape unranked tensor.
+    %dst = tensor.reshape %src(%shape)
+             : (tensor<*xf32>, tensor<?xi32>) -> tensor<*xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    AnyTensor:$source,
+    TensorRankOf<[AnySignlessInteger, Index], [1]>:$shape
+  );
+  let results = (outs AnyTensor:$result);
+
+  let builders = [OpBuilder<
+     (ins "TensorType":$resultType, "Value":$operand, "Value":$shape), [{
+       $_state.addOperands(operand);
+       $_state.addOperands(shape);
+       $_state.addTypes(resultType);
+     }]>];
+
+  let extraClassDeclaration = [{
+    TensorType getResultType() { return getResult().getType().cast<TensorType>(); }
+  }];
+
+  let assemblyFormat = [{
+    $source `(` $shape `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_YieldOp : Tensor_Op<"yield",
     [NoSideEffect, ReturnLike, Terminator,
      HasParent<"::mlir::tensor::GenerateOp">]> {
   let summary = "Yield a value from a region";
   let description = [{

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 void GenerateOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   // TODO: Move extract patterns to tensor::ExtractOp.
   results.add<ExtractFromTensorGenerate, ExtractFromTensorCast,
               StaticTensorGenerate>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// ReshapeOp
+//===----------------------------------------------------------------------===//
+
+static int64_t GetNumElements(ShapedType type) {
+  int64_t numElements = 1;
+  for (auto dim : type.getShape())
+    numElements *= dim;
+  return numElements;
+}
+
+static LogicalResult verify(ReshapeOp op) {
+  TensorType operandType = op.source().getType().cast<TensorType>();
+  TensorType resultType = op.result().getType().cast<TensorType>();
+
+  if (operandType.getElementType() != resultType.getElementType())
+    return op.emitOpError("element types of source and destination tensor "
+                          "types should be the same");
+
+  int64_t shapeSize =
+      op.shape().getType().cast<RankedTensorType>().getDimSize(0);
+  auto resultRankedType = resultType.dyn_cast<RankedTensorType>();
+  auto operandRankedType = operandType.dyn_cast<RankedTensorType>();
+
+  if (resultRankedType) {
+    if (operandRankedType && resultRankedType.hasStaticShape() &&
+        operandRankedType.hasStaticShape()) {
+      if (GetNumElements(operandRankedType) != GetNumElements(resultRankedType))
+        return op.emitOpError("source and destination tensor should have the "
+                              "same number of elements");
+    }
+    if (shapeSize == TensorType::kDynamicSize)
+      return op.emitOpError("cannot use shape operand with dynamic length to "
+                            "reshape to statically-ranked tensor type");
+    if (shapeSize != resultRankedType.getRank())
+      return op.emitOpError(
+          "length of shape operand differs from the result's tensor rank");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"

mlir/test/Dialect/Tensor/invalid.mlir

func @tensor.generate(%m : index, %n : index)
   // expected-error @+1 {{body must be terminated with a `yield` operation of the tensor element type}}
   %tnsr = tensor.generate %m, %n {
     ^bb0(%i : index, %j : index, %k : index):
       %elem = constant 8 : i32
       tensor.yield %elem : i32
   } : tensor<?x3x?xf32>
   return %tnsr : tensor<?x3x?xf32>
 }
+// -----
+
+func @tensor.reshape_element_type_mismatch(
+       %buf: tensor<*xf32>, %shape: tensor<1xi32>) {
+  // expected-error @+1 {{element types of source and destination tensor types should be the same}}
+  tensor.reshape %buf(%shape) : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xi32>
+}
+
+// -----
+
+func @tensor.reshape_dst_ranked_shape_unranked(
+       %buf: tensor<*xf32>, %shape: tensor<?xi32>) {
+  // expected-error @+1 {{cannot use shape operand with dynamic length to reshape to statically-ranked tensor type}}
+  tensor.reshape %buf(%shape) : (tensor<*xf32>, tensor<?xi32>) -> tensor<?xf32>
+}
+
+// -----
+
+func @tensor.reshape_dst_shape_rank_mismatch(
+       %buf: tensor<*xf32>, %shape: tensor<1xi32>) {
+  // expected-error @+1 {{length of shape operand differs from the result's tensor rank}}
+  tensor.reshape %buf(%shape)
+    : (tensor<*xf32>, tensor<1xi32>) -> tensor<?x?xf32>
+}
+
+// -----
+
+func @tensor.reshape_num_elements_mismatch(
+       %buf: tensor<1xf32>, %shape: tensor<1xi32>) {
+  // expected-error @+1 {{source and destination tensor should have the same number of elements}}
+  tensor.reshape %buf(%shape)
+    : (tensor<1xf32>, tensor<1xi32>) -> tensor<10xf32>
+}

mlir/test/Dialect/Tensor/ops.mlir

func @tensor.generate(%m : index, %n : index)
     -> tensor<?x3x?xf32> {
   %tnsr = tensor.generate %m, %n {
     ^bb0(%i : index, %j : index, %k : index):
       %elem = constant 8.0 : f32
       tensor.yield %elem : f32
   } : tensor<?x3x?xf32>
   return %tnsr : tensor<?x3x?xf32>
 }
+
+// CHECK-LABEL: func @tensor_reshape
+func @tensor_reshape(%unranked: tensor<*xf32>, %shape1: tensor<1xi32>,
+         %shape2: tensor<2xi32>, %shape3: tensor<?xi32>) -> tensor<*xf32> {
+  %dyn_vec = tensor.reshape %unranked(%shape1)
+               : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xf32>
+  %dyn_mat = tensor.reshape %dyn_vec(%shape2)
+               : (tensor<?xf32>, tensor<2xi32>) -> tensor<?x?xf32>
+  %new_unranked = tensor.reshape %dyn_mat(%shape3)
+               : (tensor<?x?xf32>, tensor<?xi32>) -> tensor<*xf32>
+  return %new_unranked : tensor<*xf32>
+}