[mlir] [VectorOps] Introduce vector.transpose

[aartbik]

Summary

Introduced in order to introduce specialized lowering passes that implement transposition operations efficiently.

[aartbik]

fixed accidentally removed line

[andydavis1]

Thanks Aart. How hard would this be to generalize this op to make it N-D transpose op, instead of limit to 2D cases? For example, I may want to transpose the inner two dimensions of a N-D vector

[aartbik]

In D76460#1932449, @andydavis1 wrote:

Thanks Aart. How hard would this be to generalize this op to make it N-D transpose op, instead of limit to 2D cases? For example, I may want to transpose the inner two dimensions of a N-D vector

Not very hard for the reference implementation (so the direct lowering into a bunch of extra/insert operations), but a lot harder for efficient implementations.
Would you rather that I lift the restriction on 2-D right now already and provide a reference lowering for all ranks and sizes?
(happy to do that)

[andydavis1]

Thanks Aart. Yes, if its not too much work, I'd love to see this op take a permutation of dimensions of an N-D vector.

However, if you are concerned about having to support lowerings for all permutations, then how about removing the restriction on the operand being 2-D vectors, but limit the transpose op to picking a single pair of dimensions to transpose.

For example:

func @transpose(%arg0: vector<3x4x11x7xi32>) -> vector<3x4x7x11xi32> {

%0 = vector.transpose %arg0, 2, 3 : vector<3x4x11x7xi32>
return %0 : vector<3x4x7x11xi32>

}

Thanks!

[nicolasvasilache]

Thanks Aart!

@andydavis1 if you don't mind, I'd like to go DFS on the 2-d transpose and connect it all the way to the LLVM intrinsics from this state.
We have short-term use/benchmarking cases for it.

I think it is easy to extend the op semantics now and have the matcher fail on the non-2-D case.
How do we want to represent permutation? AffineMapAttr or I32ArrayAttr?

[nicolasvasilache]

limit the transpose op to picking a single pair of dimensions to transpose

This will still be tricky because we'll have a bunch of slices etc.
I'd just go with fail if not a 2-D op so that the lowering remains very simple towards LLVM matrix intrinsics (for now).

[aartbik]

How do we want to represent permutation? AffineMapAttr or I32ArrayAttr?

If we go the route of general semantics and dimensions, I think arrayattr would the the right syntax

%1 = vector.tranpose %0 [0,1,3,2] : vector<2x3x4x5xf32>

<----->
  rank-length, all members 0 <= i < rank, proper permutation

It will not be much more work to introduce this syntax, and still verify for proper 2-D transpose (or only lower those for the time being).

WDYT? Right syntax, shall I do it?

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

1534

Drop all this and use assemblyFormat instead please.
See: https://reviews.llvm.org/D75987 for the usage of TypesMatchWith for dependent types.

[nicolasvasilache]

WDYT? Right syntax, shall I do it?

Sure go for it, it is intuitive enough IMO.

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

1534

Yes, the sole purpose of having this as C++ was to compute the type dependence. But from this feedback, I am guessing we can express this in *.td now also. Cool! Let me look into that and make the change.

[aartbik]

generalized to n-D vectors, more invalid tests

[aartbik]

PTAL

[nicolasvasilache]

Thanks Aart!

[rriddle]

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

1284

Please make sure that you use mlir code blocks for descriptions.

mlir/lib/Dialect/Vector/VectorOps.cpp

1538

nit: Prefer using llvm::SmallBitVector or llvm::BitVector.

[aartbik]

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

1284

You mean

block

right? If so, we seem to have been a bit inconsistent with that anyway.
Follow up coming....

mlir/lib/Dialect/Vector/VectorOps.cpp

1538

noted, will follow up

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

def Vector_TupleOp :
 
   let extraClassDeclaration = [{
     TupleType getResultTupleType() {
       return getResult().getType().cast<TupleType>();
     }
   }];
 }
 
+def Vector_TransposeOp :
+  Vector_Op<"transpose", [NoSideEffect,
+    PredOpTrait<"operand and result have same element type",
+                 TCresVTEtIsSameAsOpBase<0, 0>>]>,
+    Arguments<(ins AnyVector:$vector, I64ArrayAttr:$transp)>,
+    Results<(outs AnyVector:$result)> {
+  let summary = "vector transpose operation";
+  let description = [{
+    Takes a n-D vector and returns the transposed n-D vector defined by
+    the permutation of ranks in the n-sized integer array attribute.
+    In the operation
+
+    %1 = vector.tranpose %0, [i_1, .., i_n]
+      : vector<d_1 x .. x d_n x f32>
+      to vector<d_trans[0] x .. x d_trans[n-1] x f32>
+
+    the transp array [i_1, .., i_n] must be a permutation of [0, .., n-1].
+
+    Example:
+
+    %1 = vector.tranpose %0, [1, 0] : vector<2x3xf32> to vector<3x2xf32>
+
+     [ [a, b, c],       [ [a, d],
+       [d, e, f] ]  ->    [b, e],
+                          [c, f] ]
+  }];
+  let extraClassDeclaration = [{
+    VectorType getVectorType() {
+      return vector().getType().cast<VectorType>();
+    }
+    VectorType getResultType() {
+      return result().getType().cast<VectorType>();
+    }
+  }];
+  let assemblyFormat = [{
+    $vector `,` $transp attr-dict `:` type($vector) `to` type($result)
+  }];
+}
+
 def Vector_TupleGetOp :
   Vector_Op<"tuple_get", [NoSideEffect]>,
     Arguments<(ins TupleOf<[AnyVector]>:$vectors, APIntAttr:$index)>,
     Results<(outs AnyVector)> {
   let summary = "vector tuple get operation";
   let description = [{
     Returns the tuple element of 'vectors' at 'index'.
 

mlir/lib/Dialect/Vector/VectorOps.cpp

static void print(OpAsmPrinter &p, TupleOp op) {
   p.printOptionalAttrDict(op.getAttrs());
   p << " : ";
   interleaveComma(op.getOperation()->getOperandTypes(), p);
 }
 
 static LogicalResult verify(TupleOp op) { return success(); }
 
 //===----------------------------------------------------------------------===//
+// TransposeOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(TransposeOp op) {
+  VectorType vectorType = op.getVectorType();
+  VectorType resultType = op.getResultType();
+  int64_t rank = resultType.getRank();
+  if (vectorType.getRank() != rank)
+    return op.emitOpError("vector result rank mismatch: ") << rank;
+  // Verify transposition array.
+  auto transpAttr = op.transp().getValue();
+  int64_t size = transpAttr.size();
+  if (rank != size)
+    return op.emitOpError("transposition length mismatch: ") << size;
+  SmallVector<bool, 8> seen(rank, false);
+  for (auto ta : llvm::enumerate(transpAttr)) {
+    int64_t i = ta.value().cast<IntegerAttr>().getInt();
+    if (i < 0 || i >= rank)
+      return op.emitOpError("transposition index out of range: ") << i;
+    if (seen[i])
+      return op.emitOpError("duplicate position index: ") << i;
+    seen[i] = true;
+    if (resultType.getDimSize(ta.index()) != vectorType.getDimSize(i))
+      return op.emitOpError("dimension size mismatch at: ") << i;
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TupleGetOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseTupleGetOp(OpAsmParser &parser,
                                    OperationState &result) {
   OpAsmParser::OperandType operandInfo;
   IntegerAttr indexAttr;
   StringRef indexAttrName = TupleGetOp::getIndexAttrName();

mlir/test/Dialect/Vector/invalid.mlir

 
 func @reduce_unsupported_rank(%arg0: vector<4x16xf32>) -> f32 {
   // expected-error@+1 {{'vector.reduction' op unsupported reduction rank: 2}}
   %0 = vector.reduction "add", %arg0 : vector<4x16xf32> into f32
 }
 
 // -----
 
+func @transpose_rank_mismatch(%arg0: vector<4x16x11xf32>) {
+  // expected-error@+1 {{'vector.transpose' op vector result rank mismatch: 1}}
+  %0 = vector.transpose %arg0, [2, 1, 0] : vector<4x16x11xf32> to vector<100xf32>
+}
+
+// -----
+
+func @transpose_length_mismatch(%arg0: vector<4x4xf32>) {
+  // expected-error@+1 {{'vector.transpose' op transposition length mismatch: 3}}
+  %0 = vector.transpose %arg0, [2, 0, 1] : vector<4x4xf32> to vector<4x4xf32>
+}
+
+// -----
+
+func @transpose_index_oob(%arg0: vector<4x4xf32>) {
+  // expected-error@+1 {{'vector.transpose' op transposition index out of range: 2}}
+  %0 = vector.transpose %arg0, [2, 0] : vector<4x4xf32> to vector<4x4xf32>
+}
+
+// -----
+
+func @transpose_index_dup(%arg0: vector<4x4xf32>) {
+  // expected-error@+1 {{'vector.transpose' op duplicate position index: 0}}
+  %0 = vector.transpose %arg0, [0, 0] : vector<4x4xf32> to vector<4x4xf32>
+}
+
+// -----
+
+func @transpose_dim_size_mismatch(%arg0: vector<11x7x3x2xi32>) {
+  // expected-error@+1 {{'vector.transpose' op dimension size mismatch at: 0}}
+  %0 = vector.transpose %arg0, [3, 0, 1, 2] : vector<11x7x3x2xi32> to vector<2x3x7x11xi32>
+}
+
+// -----
+
 func @type_cast_layout(%arg0: memref<4x3xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2)>>) {
   // expected-error@+1 {{expects operand to be a memref with no layout}}
   %0 = vector.type_cast %arg0: memref<4x3xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2)>> to memref<vector<4x3xf32>>
 }

mlir/test/Dialect/Vector/ops.mlir

func @reduce_int(%arg0: vector<16xi32>) -> i32 {
   vector.reduction "and", %arg0 : vector<16xi32> into i32
   // CHECK:    vector.reduction "or", %{{.*}} : vector<16xi32> into i32
   vector.reduction "or", %arg0 : vector<16xi32> into i32
   // CHECK:    %[[X:.*]] = vector.reduction "xor", %{{.*}} : vector<16xi32> into i32
   %0 = vector.reduction "xor", %arg0 : vector<16xi32> into i32
   // CHECK:    return %[[X]] : i32
   return %0 : i32
 }
+
+// CHECK-LABEL: transpose_fp
+func @transpose_fp(%arg0: vector<3x7xf32>) -> vector<7x3xf32> {
+  %0 = vector.transpose %arg0, [1, 0] : vector<3x7xf32> to vector<7x3xf32>
+  return %0 : vector<7x3xf32>
+}
+
+// CHECK-LABEL: transpose_int
+func @transpose_int(%arg0: vector<11x7x3x2xi32>) -> vector<2x11x7x3xi32> {
+  %0 = vector.transpose %arg0, [3, 0, 1, 2] : vector<11x7x3x2xi32> to vector<2x11x7x3xi32>
+  return %0 : vector<2x11x7x3xi32>
+}