[mlir][VectorOps] Add ShapeCastOp to the vector ops dialect.

[andydavis1]

Summary

Add ShapeCastOp to the vector ops dialect.

The shape_cast operation casts between an n-D source vector shape and a k-D result vector shape (the element type remains the same).

[merge_guards_bot]

Unit tests: fail. 62283 tests passed, 1 failed and 831 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_recursive/lock.pass.cpp

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1416

Use /// for top level comments

1476

.getOperation isn't necessary.

1492

Same here.

[aartbik]

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

975

we have a mixed bag of notation (variables with and without quotes) in our descriptions in this file and at some point we should unify this; in this case, I find n > k easier to read, and the n and k do no refer back to any parameters in real code anyway

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1418

technically this is not a verify since it returns a bool, not a LogicalResult.
how about hasMatchingShapeCastDims or something like that

[andydavis1]

addressing comments

[merge_guards_bot]

Unit tests: pass. 62322 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

addressing comments

[nicolasvasilache]

Thanks Andy, this should make a bunch of things easier!

[rriddle]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1397

nit: Can you use the declarative assembly form now that it has landed? There isn't any docs yet(sorry, some stuff is still in-flight), but this one should be something like:

let assemblyFormat = "$source attr-dict type($source) to type($result)";

[andydavis1]

addressing comments

[andydavis1]

Addressed River's comment.

[merge_guards_bot]

Unit tests: fail. 62357 tests passed, 1 failed and 839 were skipped.

failed: libc++.std/containers/sequences/array/array_creation/to_array.fail.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

Sync with github master

[merge_guards_bot]

Unit tests: fail. 62357 tests passed, 1 failed and 839 were skipped.

failed: libc++.std/containers/sequences/array/array_creation/to_array.fail.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1397

We should be able to remove these now, right?

[andydavis1]

rebasing with head

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

addressing comments

[andydavis1]

addressing comments

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

fixed format.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

fixing formating take 2

[andydavis1]

rebasing with head.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

let extraClassDeclaration = [{
     }
   }];
   let assemblyFormat = [{
     $vector `,` $memref `[` $indices `]` attr-dict `:` type($vector) `,`
     type($memref)
   }];
 }
 
+def Vector_ShapeCastOp :
+  Vector_Op<"shape_cast", [NoSideEffect]>,
+    Arguments<(ins AnyTypeOf<[AnyVector, TupleOf<[AnyVector]>]>:$source)>,
+    Results<(outs AnyTypeOf<[AnyVector, TupleOf<[AnyVector]>]>:$result)> {
+  let summary = "shape_cast casts between vector shapes";
+  let description = [{
+    The shape_cast operation casts between an n-D source vector shape and
+    a k-D result vector shape (the element type remains the same).
+
+    If reducing rank (n > k), result dimension sizes must be a product
+    of contiguous source dimension sizes.
+    If expanding rank (n < k), source dimensions must factor into a
+    contiguous sequence of destination dimension sizes.
+    Each source dim is expanded (or contiguous sequence of source dims combined)
+    in source dimension list order (i.e. 0 <= i < n), to produce a contiguous
+    sequence of result dims (or a single result dim), in result dimension list
+    order (i.e. 0 <= j < k). The product of all source dimension sizes and all
+    result dimension sizes must match.
+
+    If the source/result types are a tuple of vectors, the casting operation
+    described above is applied to each source/result tuple element pair.
+
+    It is currently assumed that this operation does not require moving data,
+    and that it will be canonicalized away before lowering vector operations.
+
+    Examples:
+
+    ```mlir
+    // Example casting to a lower vector rank.
+    %1 = vector.shape_cast %0 : vector<5x1x4x3xf32> to vector<20x3xf32>
+
+    // Example casting to a higher vector rank.
+    %3 = vector.shape_cast %2 : vector<10x12x8xf32> to vector<5x2x3x4x8xf32>
+
+    // Example casting a tuple of vectors of same rank, where tuple elements
+    // may have different shapes.
+    %5 = vector.shape_cast %4 : tuple<vector<3x4x2xf32>, vector<3x3x2xf32>> to
+                                tuple<vector<12x2xf32>, vector<9x2xf32>>
+    ```
+  }];
+  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
+}
+
 def Vector_TypeCastOp :
   Vector_Op<"type_cast", [NoSideEffect]>,
     Arguments<(ins StaticShapeMemRefOf<[AnyType]>:$memref)>,
     Results<(outs AnyMemRef)> {
   let summary = "type_cast op converts a scalar memref to a vector memref";
   let description = [{
     Performs a conversion from a memref with scalar element to a memref with a
     *single* vector element, copying the shape of the memref to the vector. This

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/Functional.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/StringSet.h"
+#include <numeric>
 
 using namespace mlir;
 using namespace mlir::vector;
 
 //===----------------------------------------------------------------------===//
 // VectorOpsDialect
 //===----------------------------------------------------------------------===//
 
if (failed(verifyTransferOp(op.getOperation(), memrefType, vectorType,
                               permutationMap)))
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 //===----------------------------------------------------------------------===//
+// ShapeCastOp
+//===----------------------------------------------------------------------===//
+
+/// Returns true if each element of 'a' is equal to the product of a contiguous
+/// sequence of the elements of 'b'. Returns false otherwise.
+static bool isValidShapeCast(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
+  unsigned rankA = a.size();
+  unsigned rankB = b.size();
+  assert(rankA < rankB);
+
+  unsigned i = 0;
+  unsigned j = 0;
+  while (i < rankA && j < rankB) {
+    int64_t dimA = a[i];
+    int64_t dimB = 1;
+    while (dimB < dimA && j < rankB)
+      dimB *= b[j++];
+    if (dimA != dimB)
+      break;
+    ++i;
+  }
+
+  return i == rankA && j == rankB;
+}
+
+static LogicalResult verifyVectorShapeCast(Operation *op,
+                                           VectorType sourceVectorType,
+                                           VectorType resultVectorType) {
+  // Check that element type is the same.
+  if (sourceVectorType.getElementType() != resultVectorType.getElementType())
+    return op->emitOpError("source/result vectors must have same element type");
+  auto sourceShape = sourceVectorType.getShape();
+  auto resultShape = resultVectorType.getShape();
+
+  // Check that product of source dim sizes matches product of result dim sizes.
+  int64_t sourceDimProduct = std::accumulate(
+      sourceShape.begin(), sourceShape.end(), 1LL, std::multiplies<int64_t>{});
+  int64_t resultDimProduct = std::accumulate(
+      resultShape.begin(), resultShape.end(), 1LL, std::multiplies<int64_t>{});
+  if (sourceDimProduct != resultDimProduct)
+    return op->emitOpError("source/result number of elements must match");
+
+  // Check that expanding/contracting rank cases.
+  unsigned sourceRank = sourceVectorType.getRank();
+  unsigned resultRank = resultVectorType.getRank();
+  if (sourceRank < resultRank) {
+    if (!isValidShapeCast(sourceShape, resultShape))
+      return op->emitOpError("invalid shape cast");
+  } else if (sourceRank > resultRank) {
+    if (!isValidShapeCast(resultShape, sourceShape))
+      return op->emitOpError("invalid shape cast");
+  }
+  return success();
+}
+
+static LogicalResult verify(ShapeCastOp op) {
+  auto sourceVectorType = op.source().getType().dyn_cast_or_null<VectorType>();
+  auto resultVectorType = op.result().getType().dyn_cast_or_null<VectorType>();
+
+  // Check if source/result are of vector type.
+  if (sourceVectorType && resultVectorType)
+    return verifyVectorShapeCast(op, sourceVectorType, resultVectorType);
+
+  // Check if source/result are "tuple of vectors" type.
+  auto sourceTupleType = op.source().getType().dyn_cast_or_null<TupleType>();
+  auto resultTupleType = op.result().getType().dyn_cast_or_null<TupleType>();
+  if (!sourceTupleType || !resultTupleType)
+    return op.emitOpError("source/result must be of same type");
+
+  // Check that source/result tuple sizes are the same.
+  if (sourceTupleType.size() != resultTupleType.size())
+    return op.emitOpError("source/result tuples must be the same size");
+
+  // Check each source/result tuple element pair.
+  for (unsigned i = 0, e = sourceTupleType.size(); i < e; ++i)
+    if (failed(verifyVectorShapeCast(
+            op, sourceTupleType.getType(i).cast<VectorType>(),
+            resultTupleType.getType(i).cast<VectorType>())))
+      return failure();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TypeCastOp
 //===----------------------------------------------------------------------===//
 
 static MemRefType inferVectorTypeCastResultType(MemRefType t) {
   return MemRefType::get({}, VectorType::get(t.getShape(), t.getElementType()));
 }
 
 void TypeCastOp::build(Builder *builder, OperationState &result, Value source) {
   result.addOperands(source);
   result.addTypes(
       inferVectorTypeCastResultType(source.getType().cast<MemRefType>()));
 }
 
 static void print(OpAsmPrinter &p, TypeCastOp op) {
   auto type = op.getOperand().getType().cast<MemRefType>();
   p << op.getOperationName() << ' ' << op.memref() << " : " << type << " to "
     << inferVectorTypeCastResultType(type);
 }
 
 static LogicalResult verify(TypeCastOp op) {
   auto resultType = inferVectorTypeCastResultType(op.getMemRefType());
   if (op.getResultMemRefType() != resultType)
     return op.emitOpError("expects result type to be: ") << resultType;

mlir/test/Dialect/VectorOps/invalid.mlir

func @reshape_bad_output_fixed_size(%arg0 : vector<3x2x4xf32>) {
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c6 = constant 6 : index
   %c9 = constant 9 : index
   // expected-error@+1 {{fixed vector size must match output vector for dim 0}}
   %1 = vector.reshape %arg0, [%c3, %c6], [%c2, %c9], [4]
     : vector<3x2x4xf32> to vector<2x3x5xf32>
 }
+
+// -----
+
+func @shape_cast_wrong_element_type(%arg0 : vector<5x1x3x2xf32>) {
+  // expected-error@+1 {{op source/result vectors must have same element type}}
+  %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<15x2xi32>
+}
+
+// -----
+
+func @shape_cast_wrong_element_type_tuple(%arg0 : tuple<vector<5x4x2xf32>,
+                                                        vector<3x4x2xf32>>) {
+  // expected-error@+1 {{op source/result vectors must have same element type}}
+  %0 = vector.shape_cast %arg0 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
+                                 tuple<vector<20x2xi32>, vector<12x2xi32>>
+}
+
+// -----
+
+func @shape_cast_wrong_num_elements(%arg0 : vector<5x1x3x2xf32>) {
+  // expected-error@+1 {{op source/result number of elements must match}}
+  %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<10x2xf32>
+}
+
+// -----
+
+func @shape_cast_wrong_num_elements_tuple(%arg0 : tuple<vector<5x4x2xf32>,
+                                                        vector<3x4x2xf32>>) {
+  // expected-error@+1 {{op source/result number of elements must match}}
+  %0 = vector.shape_cast %arg0 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
+                                 tuple<vector<21x2xf32>, vector<13x2xf32>>
+}
+
+// -----
+
+func @shape_cast_invalid_rank_reduction(%arg0 : vector<5x1x3x2xf32>) {
+  // expected-error@+1 {{invalid shape cast}}
+  %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<2x15xf32>
+}
+
+// -----
+
+func @shape_cast_invalid_rank_reduction_tuple(%arg0
+  : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>>) {
+  // expected-error@+1 {{invalid shape cast}}
+  %0 = vector.shape_cast %arg0: tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
+                                tuple<vector<10x4xf32>, vector<6x4xf32>>
+}
+
+// -----
+
+func @shape_cast_invalid_rank_expansion(%arg0 : vector<15x2xf32>) {
+  // expected-error@+1 {{invalid shape cast}}
+  %0 = vector.shape_cast %arg0 : vector<15x2xf32> to vector<5x2x3x1xf32>
+}
+
+// -----
+
+func @shape_cast_invalid_rank_expansion_tuple(%arg0 : tuple<vector<20x2xf32>,
+                                                            vector<12x2xf32>>) {
+  // expected-error@+1 {{invalid shape cast}}
+  %0 = vector.shape_cast %arg0 : tuple<vector<20x2xf32>, vector<12x2xf32>> to
+                                 tuple<vector<5x2x4xf32>, vector<4x3x2xf32>>
+}
+
+// -----
+
+func @shape_cast_source_result_different_types(
+  %arg1 : tuple<vector<20x2xf32>, vector<12x2xf32>>) {
+  // expected-error@+1 {{source/result must be of same type}}
+  %1 = vector.shape_cast %arg1 : tuple<vector<20x2xf32>, vector<12x2xf32>> to
+                                 vector<5x2x4xf32>
+}
+
+// -----
+
+func @shape_cast_different_tuple_sizes(
+  %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>>) {
+  // expected-error@+1 {{op source/result tuples must be the same size}}
+  %1 = vector.shape_cast %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
+                                 tuple<vector<20x2xf32>>
+}

mlir/test/Dialect/VectorOps/ops.mlir

func @reshape(%arg0 : vector<3x2x4xf32>) -> (vector<2x3x4xf32>) {
   // CHECK:      %[[C9:.*]] = constant 9 : index
   %c9 = constant 9 : index
   // CHECK: vector.reshape %{{.*}}, [%[[C3]], %[[C6]]], [%[[C2]], %[[C9]]], [4] : vector<3x2x4xf32> to vector<2x3x4xf32>
   %1 = vector.reshape %arg0, [%c3, %c6], [%c2, %c9], [4]
     : vector<3x2x4xf32> to vector<2x3x4xf32>
 
   return %1 : vector<2x3x4xf32>
 }
+
+// CHECK-LABEL: shape_cast
+func @shape_cast(%arg0 : vector<5x1x3x2xf32>,
+                 %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>>)
+  -> (vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>) {
+
+  // CHECK: vector.shape_cast %{{.*}} : vector<5x1x3x2xf32> to vector<15x2xf32>
+  %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<15x2xf32>
+
+  // CHECK-NEXT: vector.shape_cast %{{.*}} : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to tuple<vector<20x2xf32>, vector<12x2xf32>>
+  %1 = vector.shape_cast %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
+                                 tuple<vector<20x2xf32>, vector<12x2xf32>>
+
+  return %0, %1 : vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>
+}