[mlir][Vector] Support 0-D vectors in TransposeOp

[nicolasvasilache]

Summary

Co-authored-by: Michal Terepeta <[e-mail]>

Reviewed-by: ftynse

[ftynse]

mlir/test/Dialect/Vector/invalid.mlir

1149

Nit: I find this error message quite cryptic -- what is "1"? Expected rank? Actual rank? Of what? I see that it is not introduced by this patch, but if you can fix that in a separate follow-up, it would be appreciated.

[nicolasvasilache]

Rebase

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

def Vector_CreateMaskOp :
   let assemblyFormat = "$operands attr-dict `:` type(results)";
 }
 
 def Vector_TransposeOp :
   Vector_Op<"transpose", [NoSideEffect,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
     PredOpTrait<"operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>]>,
-    Arguments<(ins AnyVector:$vector, I64ArrayAttr:$transp)>,
-    Results<(outs AnyVector:$result)> {
+    Arguments<(ins AnyVectorOfAnyRank:$vector, I64ArrayAttr:$transp)>,
+    Results<(outs AnyVectorOfAnyRank:$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.
+    the permutation of ranks in the n-sized integer array attribute (in case
+    of 0-D vectors the array attribute must be empty).
     In the operation
 
     ```mlir
     %1 = vector.transpose %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>
     ```
 

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 // CHECK-SAME: %[[arg:.*]]: index
 // CHECK:  %[[indices:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xi32>
 // CHECK:  %[[arg_i32:.*]] = arith.index_cast %[[arg]] : index to i32
 // CHECK:  %[[boundsInsert:.*]] = llvm.insertelement %[[arg_i32]]
 // CHECK:  %[[bounds:.*]] = llvm.shufflevector %[[boundsInsert]]
 // CHECK:  %[[result:.*]] = arith.cmpi slt, %[[indices]], %[[bounds]] : vector<4xi32>
 // CHECK:  return %[[result]] : vector<4xi1>
 
+// -----
+
 func.func @create_mask_1d_scalable(%a : index) -> vector<[4]xi1> {
   %v = vector.create_mask %a : vector<[4]xi1>
   return %v: vector<[4]xi1>
 }
 
 // CHECK-LABEL: func @create_mask_1d_scalable
 // CHECK-SAME: %[[arg:.*]]: index
 // CHECK:  %[[indices:.*]] = llvm.intr.experimental.stepvector : vector<[4]xi32>
 // CHECK:  %[[arg_i32:.*]] = arith.index_cast %[[arg]] : index to i32
 // CHECK:  %[[boundsInsert:.*]] = llvm.insertelement %[[arg_i32]], {{.*}} : vector<[4]xi32>
 // CHECK:  %[[bounds:.*]] = llvm.shufflevector %[[boundsInsert]], {{.*}} : vector<[4]xi32>
 // CHECK:  %[[result:.*]] = arith.cmpi slt, %[[indices]], %[[bounds]] : vector<[4]xi32>
 // CHECK: return %[[result]] : vector<[4]xi1>
 
 // -----
 
+func.func @transpose_0d(%arg0: vector<f32>) -> vector<f32> {
+  %0 = vector.transpose %arg0, [] : vector<f32> to vector<f32>
+  return %0 : vector<f32>
+}
+
+// CHECK-LABEL: func @transpose_0d
+// CHECK-SAME:  %[[A:.*]]: vector<f32>
+// CHECK:       return %[[A]] : vector<f32>
+
+// -----
+
 func.func @flat_transpose(%arg0: vector<16xf32>) -> vector<16xf32> {
   %0 = vector.flat_transpose %arg0 { rows = 4: i32, columns = 4: i32 }
      : vector<16xf32> -> vector<16xf32>
   return %0 : vector<16xf32>
 }
 
 // CHECK-LABEL: func @flat_transpose
 // CHECK-SAME:  %[[A:.*]]: vector<16xf32>

mlir/test/Dialect/Vector/invalid.mlir

 
 func.func @multi_reduce_invalid_type(%arg0: vector<4x16xf32>, %acc: vector<16xf32>) -> f32 {
   // expected-error@+1 {{'vector.multi_reduction' op destination type 'vector<16xf32>' is incompatible with source type 'vector<4x16xf32>'}}
   %0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<4x16xf32> to vector<16xf32>
 }
 
 // -----
 
+func.func @transpose_rank_mismatch_0d(%arg0: vector<f32>) {
+  // expected-error@+1 {{'vector.transpose' op vector result rank mismatch: 1}}
+  %0 = vector.transpose %arg0, [] : vector<f32> to vector<100xf32>
+}
+
+// -----
+
 func.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.func @transpose_length_mismatch_0d(%arg0: vector<f32>) {
+  // expected-error@+1 {{'vector.transpose' op transposition length mismatch: 1}}
+  %0 = vector.transpose %arg0, [1] : vector<f32> to vector<f32>
+}
+
+// -----
+
 func.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.func @transpose_index_oob(%arg0: vector<4x4xf32>) {

mlir/test/Dialect/Vector/ops.mlir

 // CHECK-LABEL: @transpose_int
 func.func @transpose_int(%arg0: vector<11x7x3x2xi32>) -> vector<2x11x7x3xi32> {
   // CHECK: %[[X:.*]] = vector.transpose %{{.*}}, [3, 0, 1, 2] : vector<11x7x3x2xi32> to vector<2x11x7x3xi32>
   %0 = vector.transpose %arg0, [3, 0, 1, 2] : vector<11x7x3x2xi32> to vector<2x11x7x3xi32>
   // CHECK: return %[[X]] : vector<2x11x7x3xi32>
   return %0 : vector<2x11x7x3xi32>
 }
 
+// CHECK-LABEL: @transpose_fp_0d
+func.func @transpose_fp_0d(%arg0: vector<f32>) -> vector<f32> {
+  // CHECK: %[[X:.*]] = vector.transpose %{{.*}}, [] : vector<f32> to vector<f32>
+  %0 = vector.transpose %arg0, [] : vector<f32> to vector<f32>
+  // CHECK: return %[[X]] : vector<f32>
+  return %0 : vector<f32>
+}
+
+// CHECK-LABEL: @transpose_int_0d
+func.func @transpose_int_0d(%arg0: vector<i32>) -> vector<i32> {
+  // CHECK: %[[X:.*]] = vector.transpose %{{.*}}, [] : vector<i32> to vector<i32>
+  %0 = vector.transpose %arg0, [] : vector<i32> to vector<i32>
+  // CHECK: return %[[X]] : vector<i32>
+  return %0 : vector<i32>
+}
+
 // CHECK-LABEL: @flat_transpose_fp
 func.func @flat_transpose_fp(%arg0: vector<16xf32>) -> vector<16xf32> {
   // CHECK: %[[X:.*]] = vector.flat_transpose %{{.*}} {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32>
   %0 = vector.flat_transpose %arg0 { rows = 4: i32, columns = 4: i32 } : vector<16xf32> -> vector<16xf32>
   // CHECK: return %[[X]] : vector<16xf32>
   return %0 : vector<16xf32>
 }
 

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

 func.func @fma_0d(%four: vector<f32>) {
   %0 = vector.fma %four, %four, %four : vector<f32>
   // 4 * 4 + 4 = 20
   // CHECK: ( 20 )
   vector.print %0: vector<f32>
   return
 }
 
+func.func @transpose_0d(%arg: vector<i32>) {
+  %1 = vector.transpose %arg, [] : vector<i32> to vector<i32>
+  // CHECK: ( 42 )
+  vector.print %1: vector<i32>
+  return
+}
+
 func.func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
 
   %3 = arith.constant dense<42.0> : vector<f32>
   call  @print_vector_0d(%3) : (vector<f32>) -> ()
func.func @entry() {
   %one_idx = arith.constant 1 : index
   call  @create_mask_0d(%zero_idx, %one_idx) : (index, index) -> ()
 
   %red_array = arith.constant dense<5.0> : vector<f32>
   call  @reduce_add(%red_array) : (vector<f32>) -> ()
 
   %5 = arith.constant dense<4.0> : vector<f32>
   call  @fma_0d(%5) : (vector<f32>) -> ()
+  %6 = arith.constant dense<42> : vector<i32>
+  call @transpose_0d(%6) : (vector<i32>) -> ()
 
   return
 }