Adds CastSliceOp to the vector ops dialect.

[andydavis1]

Summary

The cast_slices vector operation type casts slices from a 'source' vector to
slices of a 'result' vector, where there is a one-to-one mapping between source
vector and result vector slices (specified by the 'source_sizes' and 'result_sizes' arguments).

[aartbik]

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

951

I would like to see a bit more documentation of this operation in terms of ranks and sizes of the arrays, like we do for the other sliced operations (so, like "n-D" with restrictions on n array sizes for the input arguments and output).

[rriddle]

Please fix the description, BEGIN_PUBLIC/END_PUBLIC don't make sense. The summary should also have a [mlir] tag.

[rriddle]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1651

Cache the end iterator here and above.

mlir/test/Dialect/VectorOps/ops.mlir

248

Seems to be missing a newline here?

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

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-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[andydavis1]

Addressing comments.

[andydavis1]

addressing comments

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

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

[nicolasvasilache]

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

964

I find the semantics harder to grasp than it could, I think.
I think the issue is that the usage you have for "cast_slices" leaks into the semantics of the op and I don't think it is warranted.

%1 = vector.cast_slices %0, [1, 1, 3, 12], [3, 12]
      : vector<4x1x3x12xf32> to vector<12x12xf32>

I can think of different ways of specifying this that I would find more natural but first I'd like to understand why you need these attributes in this form and where ambiguity appears that makes it impossible to derive them automatically.

In particular I view this as conflating 2 unrelated things: (1) a reshape that does not move and (2) a multiplicity.
See how (1) is represented in linalg.reshape with affine maps.

It would be great if we could uniformize things (if is makes sense).
Note that this may be an indication that you need 2 ops to express what you want.

Marking as "request changes" until the discussion is resolved.

[andydavis1]

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

964

This operation specifies a cast from one vector type to another. I intentionally did not use affine maps (or even the less expressive permutation of dimensions to reshape), so that this operation was even more restricted to its purpose. Its purpose is express a trivial cast between slices of the source/result vectors. Because these aggegate vectors will eventually need to be decomposed to hardware vectors, whats important is that the slices of these vectors have trivial casts that can be canonicalized away.

As an analogy. Think of a reshape operation which starts with the linear address, and de-linearizes this into the source (resp result) vector shape, by using the source (resp result) strides. Instead of source/result permutation of dims, the cast_slices operation specifies the source/result decomposition of slices and further restricts the operation in the following simple ways:

1. source/result must generate the same number slices so there is a 1-to-1 relationship between a source slice and its associated result slice.
2. each source/result slice pair must be trivial castable to each other

[nicolasvasilache]

We sync'ed offline and the intuition is that there is a more idiomatic way involving just a "natural" vector.cast that works on both vector and tuple<vector<>,...>

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

VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
   }];
 }
 
+def Vector_CastSlicesOp :
+  Vector_Op<"cast_slices", [NoSideEffect]>,
+    Arguments<(ins AnyVector:$vector, I64ArrayAttr:$source_sizes,
+               I64ArrayAttr:$result_sizes)>,
+    Results<(outs AnyVector)> {
+  let summary = "The cast_slices op casts slices of a vector to another type";
+  let description = [{
+    The cast_slices vector operation type casts slices from a 'source' n-D
+    vector to slices of a 'result' k-D vector. The vector ranks 'n' and 'k'
+    can be arbitrary as long as the following conditions are met:
+    *) The number of slices generated by applying the target slice sizes
+       'source_sizes' to the source vector shape, must equal the number of
+       slices generated by applying the target slice sizes 'result_sizes' to
+       the operation's result vector shape.
+    *) The cast between source/result slice shapes must be trivial (i.e. it
+       must not move or reshape data).
+
+    Example:
+
+    ```mlir
+     // Cast slices of 'vector<4x1x3x12xf32>' from shape 1x1x3x12 to 3x12
+     %1 = vector.cast_slices %0, [1, 1, 3, 12], [3, 12]
+       : vector<4x1x3x12xf32> to vector<12x12xf32>
+
+     // Cast slices of 'vector<12x12xf32>' from shape 3x12 to 1x1x3x12
+     %2 = vector.cast_slices %1, [3, 12], [1, 1, 3, 12]
+       : vector<12x12xf32> to vector<4x1x3x12xf32>
+    ```
+  }];
+
+  let extraClassDeclaration = [{
+    VectorType getSourceVectorType() {
+      return vector().getType().cast<VectorType>();
+    }
+    VectorType getResultVectorType() {
+      return getResult().getType().cast<VectorType>();
+    }
+    void getSourceSizes(SmallVectorImpl<int64_t> &results);
+    void getResultSizes(SmallVectorImpl<int64_t> &results);
+    static StringRef getSourceSizesAttrName() { return "source_sizes"; }
+    static StringRef getResultSizesAttrName() { return "result_sizes"; }
+  }];
+}
+
 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

if (failed(verifyTransferOp(op.getOperation(), memrefType, vectorType,
                               permutationMap)))
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 //===----------------------------------------------------------------------===//
+// CastSlicesOp
+//===----------------------------------------------------------------------===//
+
+ParseResult parseCastSlicesOp(OpAsmParser &parser, OperationState &result) {
+  ArrayAttr sourceSizesAttr;
+  StringRef sourceSizesAttrName = CastSlicesOp::getSourceSizesAttrName();
+  ArrayAttr resultSizesAttr;
+  StringRef resultSizesAttrName = CastSlicesOp::getResultSizesAttrName();
+
+  OpAsmParser::OperandType source;
+  Type sourceType;
+  VectorType vectorType;
+  return failure(parser.parseOperand(source) || parser.parseComma() ||
+                 parser.parseAttribute(sourceSizesAttr, sourceSizesAttrName,
+                                       result.attributes) ||
+                 parser.parseComma() ||
+                 parser.parseAttribute(resultSizesAttr, resultSizesAttrName,
+                                       result.attributes) ||
+                 parser.parseOptionalAttrDict(result.attributes) ||
+                 parser.parseColonType(sourceType) ||
+                 parser.parseKeywordType("to", vectorType) ||
+                 parser.resolveOperand(source, sourceType, result.operands) ||
+                 parser.addTypeToList(vectorType, result.types));
+}
+
+static void print(OpAsmPrinter &p, CastSlicesOp op) {
+  p << op.getOperationName() << ' ' << op.vector() << ", ";
+  p << op.source_sizes() << ", " << op.result_sizes();
+  p.printOptionalAttrDict(
+      op.getAttrs(),
+      /*elidedAttrs=*/{CastSlicesOp::getSourceSizesAttrName(),
+                       CastSlicesOp::getResultSizesAttrName()});
+  p << " : " << op.getSourceVectorType() << " to " << op.getResultVectorType();
+}
+
+static SmallVector<int64_t, 4> canonicalizeShape(ArrayRef<int64_t> shape) {
+  SmallVector<int64_t, 4> result;
+  for (int64_t dim : shape)
+    if (dim != 1)
+      result.push_back(dim);
+  return result;
+}
+
+static LogicalResult verify(CastSlicesOp op) {
+  // Verify that slicing source vector shape by 'source_sizes' produces the
+  // same number of slices as slicing result vector shape by 'result_sizes'.
+  SmallVector<int64_t, 4> sourceSizes;
+  op.getSourceSizes(sourceSizes);
+  auto sourceVectorType = op.getSourceVectorType();
+  auto sourceVectorShape = sourceVectorType.getShape();
+  int64_t sourceSliceCount = 1;
+  for (unsigned i = 0, e = sourceVectorType.getRank(); i < e; ++i)
+    sourceSliceCount *= ceilDiv(sourceVectorShape[i], sourceSizes[i]);
+
+  SmallVector<int64_t, 4> resultSizes;
+  op.getResultSizes(resultSizes);
+  auto resultVectorType = op.getResultVectorType();
+  auto resultVectorShape = resultVectorType.getShape();
+  int64_t resultSliceCount = 1;
+  for (unsigned i = 0, e = resultVectorType.getRank(); i < e; ++i)
+    resultSliceCount *= ceilDiv(resultVectorShape[i], resultSizes[i]);
+
+  if (sourceSliceCount != resultSliceCount)
+    return op.emitOpError(
+        "requires the same number of source/result vector slices");
+
+  // Verify that the target slice shapes are trivially castable from one
+  // to another.
+  if (canonicalizeShape(sourceSizes) != canonicalizeShape(resultSizes))
+    return op.emitOpError("requires source/result with equivalent shapes");
+  return success();
+}
+
+void CastSlicesOp::getSourceSizes(SmallVectorImpl<int64_t> &results) {
+  populateFromInt64AttrArray(source_sizes(), results);
+}
+
+void CastSlicesOp::getResultSizes(SmallVectorImpl<int64_t> &results) {
+  populateFromInt64AttrArray(result_sizes(), results);
+}
+
+//===----------------------------------------------------------------------===//
 // TypeCastOp
 //===----------------------------------------------------------------------===//
 
 static MemRefType inferVectorTypeCastResultType(MemRefType t) {
   return MemRefType::get({}, VectorType::get(t.getShape(), t.getElementType()));
 }
 
 void TypeCastOp::build(Builder *builder, OperationState &result, Value source) {

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 @cast_slices_wrong_number_of_slices(%arg0 : vector<5x1x3x12xf32>) {
+  // expected-error@+1 {{requires the same number of source/result vector slices}}
+  %0 = vector.cast_slices %arg0, [1, 1, 3, 12], [3, 12]
+    : vector<5x1x3x12xf32> to vector<12x12xf32>
+}
+
+// -----
+
+func @cast_slices_cant_cast_slice_sizes(%arg0 : vector<4x1x3x12xf32>) {
+  // expected-error@+1 {{requires source/result with equivalent shapes}}
+  %0 = vector.cast_slices %arg0, [1, 2, 3, 12], [3, 12]
+    : vector<4x1x3x12xf32> to vector<12x12xf32>
+}

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: cast_slices
+func @cast_slices(%arg0 : vector<4x1x3x12xf32>) -> (vector<4x1x3x12xf32>) {
+  // CHECK: %[[CS0:.*]] = vector.cast_slices %{{.*}}, [1, 1, 3, 12], [3, 12] : vector<4x1x3x12xf32> to vector<12x12xf32>
+  %0 = vector.cast_slices %arg0, [1, 1, 3, 12], [3, 12]
+    : vector<4x1x3x12xf32> to vector<12x12xf32>
+  // CHECK: vector.cast_slices %[[CS0]], [3, 12], [1, 1, 3, 12] : vector<12x12xf32> to vector<4x1x3x12xf32>
+  %1 = vector.cast_slices %0, [3, 12], [1, 1, 3, 12]
+    : vector<12x12xf32> to vector<4x1x3x12xf32>
+
+  return %1 : vector<4x1x3x12xf32>
+}