[mlir] [VectorOps] Add 'vector.flat_transpose' operation

[aartbik]

Summary

Provides a representation of the linearized LLVM instrinsic.
With tests and lowering implementation to LLVM IR dialect.
Prepares better lowering for 2-D vector.transpose.

[aartbik]

fixed typo

[dcaballe]

Thanks! It looks good to me in general. I just added some comments about making it a bit more generic.
For my understanding, how do we plan to lower to these vector matrix multiply and transpose? From a vector contraction?

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

1557

Since the Vector dialect supports multi-dim vectors, I think we should be explicit about the flattened nature of this op in the name class and the mnemonic. Otherwise, it would be surprising that vector.matrix_transpose only accepts 1D vectors. What about something like FlatMatrixTransposeOp/flat_matrix_transpose?

1570

Since this is the Vector dialect and should be backend independent, I would suggest making the description more generic so that other backends can also use it. We can talk about the llvm.matrix.transpose at the end just as one potential lowering example.

1576

Could you also briefly describe how to represent the multi-dim version of this in the Vector dialect?

[aartbik]

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

1557

Note that we already have an instance of an op needed for progressive lowering and type change (see line 1467 that starts this section). The "matrix" in the name gives the special meaning. I a

regular                       mid-level

--------------------+-----------------
vector.contract -> vector.matrix_multiply

vector.transpose -> vector_matrix_transpose

So hopefully this is sufficient. If you insist, however, we should rename them both into flat_matrix..

1570

I moved the comment around a bit, see if that's better (note however that we have more instances of ops close to intrinsics)

1576

Added more comments

[aartbik]

clarified that certain ops ('matrix') are only there for progressive lowering

[aartbik]

fixed typo in new comments

[ftynse]

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

1557

I tend to agree with @dcaballe: while vector.matrix_multiply on flat vectors makes sense more or less, vector.matrix_transpose on flat vectors and vector.transpose on nD vectors sounds the like opposite of what one would naturally assume.

1595

You can omit = ? lines, these are the default anyway.

1596

Nit: do we really need parenthesis around type($matrix) ?

mlir/test/Dialect/Vector/invalid.mlir

1149

Nit: we don't have to verify auto-generated messages

[dcaballe]

Thanks, LGTM. You can rename both intrinsics in a separate commit.

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

1557

Yeah, I thing having vector.flat_transpose and vector.flat_matrix_multiply would be more clear. Feel free to change it in a separate commit, though. No problem.

[aartbik]

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

1557

Ok, since requested by two now, I renamed this flat_tranpose since renaming later is a bit of a pain always (but I still feel my naming was closer the llvm conventions). I leave the renaming of the existing matrix_multiply to Nicolas then, since he has some other code invested in that already.

1595

Removed, for matrix_multiply too

1596

I felt it was consistent with the matrix multiply case. As is, I feel we are getting a bit inconsistent with our "to", "into" and "->". I removed it here though, since you requested it.

mlir/test/Dialect/Vector/invalid.mlir

1149

Fair enough. I wanted something here. I think we should actually verify that rows x columns is valid, also for the matrix multiply.

[aartbik]

addressed comments

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

Type getPrintType() {
       return source().getType();
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source)";
 }
 
 //===----------------------------------------------------------------------===//
 // Ops used for supporting progressive lowering and conversion type changes.
+// The Ops are typically not used directly by higher level dialects, but are
+// used by intra-dialect rewriting rules to bring vector operations closer
+// to the hardware ISA.
 //===----------------------------------------------------------------------===//
 
 /// Vector dialect matrix multiplication op that operates on flattened 1-D
 /// MLIR vectors. This is the counterpart of llvm.matrix.multiply in MLIR.
 /// This may seem redundant with vector.contract but it serves the purposes of
 /// more progressive lowering and localized type conversion on the path:
 ///   `vector<...x...xf32> -> vector<...xf32> -> !llvm<... x float>`.
 def Vector_MatmulOp : Vector_Op<"matrix_multiply", [NoSideEffect,
def Vector_MatmulOp : Vector_Op<"matrix_multiply", [NoSideEffect,
         outs VectorOfRankAndType<[1],
                [AnySignlessInteger, AnySignedInteger, AnyFloat]>:$res)>
 {
   let summary = "Vector matrix multiplication op that operates on flattened 1-D"
     " MLIR vectors";
   let description = [{
     This is the counterpart of llvm.matrix.multiply in MLIR. It serves the
     purposes of more progressive lowering and localized type conversion.
+    Higher levels typically lower matrix multiplications into 'vector.contract'
+    operations. Subsequent rewriting rule progressively lower these operations
+    into 'vector.matrix_multiply' operations to bring the operations closer
+    to the hardware ISA.
 
     The ‘vector.matrix_multiply’ op treats `lhs` as matrix with <lhs_rows> rows
     and <lhs_columns> columns, `rhs` as matrix with <lhs_columns> rows and
     <rhs_columns> and multiplies them. The result matrix is returned embedded in
     the result vector.
 
+    Also see:
+
+    http://llvm.org/docs/LangRef.html#llvm-matrix-multiply-intrinsic
+
     Example:
 
     ```mlir
     %C = vector.matrix_multiply %A, %B
       { lhs_rows = 4: i32, lhs_columns = 16: i32 , rhs_columns = 3: i32 } :
       (vector<64xf64>, vector<48xf64>) -> vector<12xf64>
     ```
   }];
result.addTypes(VectorType::get(lhsRows * rhsColumns,
        lhs.getType().cast<VectorType>().getElementType()));
    }]>,
   ];
   let verifier = ?;
   let assemblyFormat = "$lhs `,` $rhs attr-dict "
     "`:` `(` type($lhs) `,` type($rhs) `)` `->` type($res)";
 }
 
+/// Vector dialect matrix tranposition op that operates on flattened 1-D
+/// MLIR vectors. This is the counterpart of llvm.matrix.transpose in MLIR.
+/// This may seem redundant with vector.transpose but it serves the purposes of
+/// more progressive lowering and localized type conversion on the path:
+///   `vector<...x...xf32> -> vector<...xf32> -> !llvm<... x float>`.
+def Vector_FlatTransposeOp : Vector_Op<"flat_transpose", [NoSideEffect,
+  PredOpTrait<"source operand and result have same element type",
+                 TCresVTEtIsSameAsOpBase<0, 0>>]>,
+    Arguments<(
+      // TODO(ntv, fhahn, ajcbik): tighten vector element types that make sense.
+      ins VectorOfRankAndType<[1],
+            [AnySignlessInteger, AnySignedInteger, AnyFloat]>:$matrix,
+          I32Attr:$rows, I32Attr:$columns)>,
+    Results<(
+      outs VectorOfRankAndType<[1],
+             [AnySignlessInteger, AnySignedInteger, AnyFloat]>:$res)> {
+  let summary = "Vector matrix transposition on flattened 1-D MLIR vectors";
+  let description = [{
+    This is the counterpart of llvm.matrix.transpose in MLIR. It serves
+    the purposes of more progressive lowering and localized type conversion.
+    Higher levels typically lower matrix tranpositions into 'vector.transpose'
+    operations. Subsequent rewriting rule progressively lower these operations
+    into 'vector.flat_transpose' operations to bring the operations closer
+    to the hardware ISA.
+
+    The ‘vector.flat_transpose’ op treats the 1-D input `matrix` as
+    a 2-D matrix with <rows> rows and <columns> columns, and returns the
+    transposed matrix in flattened form in 'res'.
+
+    Also see:
+
+    http://llvm.org/docs/LangRef.html#llvm-matrix-transpose-intrinsic
+
+    Example:
+
+    ```mlir
+    %1 = vector.flat_transpose %0 { rows = 4: i32, columns = 4: i32 }
+       : (vector<16xf32>) -> vector<16xf32>
+    ```
+  }];
+  let verifier = ?;
+  let assemblyFormat = "$matrix attr-dict `:` type($matrix) `->` type($res)";
+}
+
 #endif // VECTOR_OPS

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<LLVM::MatrixMultiplyOp>(
         op, typeConverter.convertType(matmulOp.res().getType()), adaptor.lhs(),
         adaptor.rhs(), matmulOp.lhs_rows(), matmulOp.lhs_columns(),
         matmulOp.rhs_columns());
     return success();
   }
 };
 
+/// Conversion pattern for a vector.flat_transpose.
+/// This is lowered directly to the proper llvm.intr.matrix.transpose.
+class VectorFlatTransposeOpConversion : public ConvertToLLVMPattern {
+public:
+  explicit VectorFlatTransposeOpConversion(MLIRContext *context,
+                                           LLVMTypeConverter &typeConverter)
+      : ConvertToLLVMPattern(vector::FlatTransposeOp::getOperationName(),
+                             context, typeConverter) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto transOp = cast<vector::FlatTransposeOp>(op);
+    auto adaptor = vector::FlatTransposeOpOperandAdaptor(operands);
+    rewriter.replaceOpWithNewOp<LLVM::MatrixTransposeOp>(
+        transOp, typeConverter.convertType(transOp.res().getType()),
+        adaptor.matrix(), transOp.rows(), transOp.columns());
+    return success();
+  }
+};
+
 class VectorReductionOpConversion : public ConvertToLLVMPattern {
 public:
   explicit VectorReductionOpConversion(MLIRContext *context,
                                        LLVMTypeConverter &typeConverter)
       : ConvertToLLVMPattern(vector::ReductionOp::getOperationName(), context,
                              typeConverter) {}
 
   LogicalResult
patterns
               VectorTypeCastOpConversion>(ctx, converter);
   // clang-format on
 }
 
 void mlir::populateVectorToLLVMMatrixConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   MLIRContext *ctx = converter.getDialect()->getContext();
   patterns.insert<VectorMatmulOpConversion>(ctx, converter);
+  patterns.insert<VectorFlatTransposeOpConversion>(ctx, converter);
 }
 
 namespace {
 struct LowerVectorToLLVMPass
     : public ConvertVectorToLLVMBase<LowerVectorToLLVMPass> {
   void runOnOperation() override;
 };
 } // namespace

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

 // CHECK: %[[T3:.*]] = llvm.insertelement %[[T0]], %[[T1]][%[[T2]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: %[[T4:.*]] = llvm.mlir.constant(1 : i64) : !llvm.i64
 // CHECK: %[[T5:.*]] = llvm.insertelement %[[T0]], %[[T3]][%[[T4]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: %[[T6:.*]] = llvm.mlir.constant(2 : i64) : !llvm.i64
 // CHECK: %[[T7:.*]] = llvm.insertelement %[[T0]], %[[T5]][%[[T6]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: %[[T8:.*]] = llvm.mlir.constant(3 : i64) : !llvm.i64
 // CHECK: %[[T9:.*]] = llvm.insertelement %[[T0]], %[[T7]][%[[T8]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: llvm.return %9 : !llvm<"<8 x i1>">
+
+// CHECK-LABEL: func @flat_transpose
+// CHECK-SAME:  %[[A:.*]]: !llvm<"<16 x float>">
+// CHECK:       %[[T:.*]] = llvm.intr.matrix.transpose %[[A]]
+// CHECK-SAME:      {columns = 4 : i32, rows = 4 : i32} :
+// CHECK-SAME:      !llvm<"<16 x float>"> into !llvm<"<16 x float>">
+// CHECK:       llvm.return %[[T]] : !llvm<"<16 x float>">
+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>
+}

mlir/test/Dialect/Vector/invalid.mlir

 
 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 @flat_transpose_type_mismatch(%arg0: vector<16xf32>) {
+  // expected-error@+1 {{'vector.flat_transpose' op failed to verify that source operand and result have same element type}}
+  %0 = vector.flat_transpose %arg0 { rows = 4: i32, columns = 4: i32 } : vector<16xf32> -> vector<16xf64>
+}
+
+// -----
+
 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

#contraction_to_scalar_accesses = [
   affine_map<(i) -> (i)>,
   affine_map<(i) -> (i)>,
   affine_map<(i) -> ()>
 ]
 #contraction_to_scalar_trait = {
   indexing_maps = #contraction_to_scalar_accesses,
   iterator_types = ["reduction"]
 }
-// CHECK-LABEL: contraction_to_scalar
+// CHECK-LABEL: @contraction_to_scalar
 func @contraction_to_scalar(%arg0: vector<10xf32>, %arg1: vector<10xf32>) -> f32 {
   // CHECK:      %[[C0:.*]] = constant 0.000000e+00 : f32
   %f0 = constant 0.0: f32
   // CHECK:      %[[X:.*]] = vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["reduction"]} %{{.*}}, %{{.*}}, %[[C0]] : vector<10xf32>, vector<10xf32> into f32
   %0 = vector.contract #contraction_to_scalar_trait %arg0, %arg1, %f0
     : vector<10xf32>, vector<10xf32> into f32
   // CHECK:      return %[[X]] : f32
   return %0 : f32
affine_map<(f0, f1, f2, f3, c0, c1) -> (f1, c1, c0, f3)>,
                                         // 8,  8, 15,  5
   affine_map<(f0, f1, f2, f3, c0, c1) -> (f0, f1, f2, f3)>
 ]
 #contraction_trait1 = {
   indexing_maps = #contraction_accesses1,
   iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction",
                     "reduction"]
 }
-// CHECK-LABEL: contraction
+// CHECK-LABEL: @contraction
 func @contraction(%arg0 : vector<7x8x16x15xf32>, %arg1 : vector<8x16x7x5xf32>,
                   %arg2 : vector<8x15x5xf32>, %arg3 : vector<8x8x15x5xf32>,
                   %arg4 : index) {
   // Test contraction with batch and contracting dims.
   // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
   %0 = vector.contract #contraction_trait0 %arg0, %arg1, %arg2
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
   // Test contraction with only contracting dims. In this case the lhs/rhs
   // dimension of size 8 will be considered a parallel dim for lhs/rhs and will
   // appear twice in the output.
   // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   %1 = vector.contract #contraction_trait1 %arg0, %arg1, %arg3
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   // Test contraction with optional vector mask arguments.
   %lhs_mask = vector.constant_mask [7, 8, 16, 15] : vector<7x8x16x15xi1>
   %rhs_mask = vector.constant_mask [8, 16, 7, 5] : vector<8x16x7x5xi1>
   // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   %2 = vector.contract #contraction_trait1 %arg0, %arg1, %arg3, %lhs_mask,
                                            %rhs_mask
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   return
 }
 
-// CHECK-LABEL: create_vector_mask
+// CHECK-LABEL: @create_vector_mask
 func @create_vector_mask() {
   // CHECK:      %[[C2:.*]] = constant 2 : index
   %c2 = constant 2 : index
   // CHECK-NEXT: %[[C3:.*]] = constant 3 : index
   %c3 = constant 3 : index
   // CHECK-NEXT: vector.create_mask %[[C3]], %[[C2]] : vector<4x3xi1>
   %0 = vector.create_mask %c3, %c2 : vector<4x3xi1>
 
   return
 }
 
-// CHECK-LABEL: constant_vector_mask
+// CHECK-LABEL: @constant_vector_mask
 func @constant_vector_mask() {
   // CHECK: vector.constant_mask [3, 2] : vector<4x3xi1>
   %0 = vector.constant_mask [3, 2] : vector<4x3xi1>
   return
 }
 
-// CHECK-LABEL: extract_slices
+// CHECK-LABEL: @extract_slices
 func @extract_slices(%arg0 : vector<4x2xf32>)
   -> (tuple<vector<2x2xf32>, vector<2x2xf32>>) {
   // CHECK: vector.extract_slices %{{.*}}, [2, 2], [1, 1] : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
   %0 = vector.extract_slices %arg0, [2, 2], [1, 1]
     : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
   %1 = vector.tuple_get %0, 0 : tuple<vector<2x2xf32>, vector<2x2xf32>>
   %2 = vector.tuple_get %0, 1 : tuple<vector<2x2xf32>, vector<2x2xf32>>
   %3 = vector.tuple %1, %2 : vector<2x2xf32>, vector<2x2xf32>
   return %3 : tuple<vector<2x2xf32>, vector<2x2xf32>>
 }
 
-// CHECK-LABEL: insert_slices
+// CHECK-LABEL: @insert_slices
 func @insert_slices(%arg0 : tuple<vector<2x2xf32>, vector<2x2xf32>>)
   -> (vector<4x2xf32>) {
   // CHECK: vector.insert_slices %{{.*}}, [2, 2], [1, 1] : tuple<vector<2x2xf32>, vector<2x2xf32>> into vector<4x2xf32>
   %0 = vector.insert_slices %arg0, [2, 2], [1, 1]
     : tuple<vector<2x2xf32>, vector<2x2xf32>> into vector<4x2xf32>
   return %0 : vector<4x2xf32>
 }
 
 // CHECK-LABEL: @vector_print
 func @vector_print(%arg0: vector<8x4xf32>) {
   // CHECK: vector.print %{{.*}} : vector<8x4xf32>
   vector.print %arg0 : vector<8x4xf32>
   return
 }
 
-// CHECK-LABEL: reshape
+// CHECK-LABEL: @reshape
 func @reshape(%arg0 : vector<3x2x4xf32>) -> (vector<2x3x4xf32>) {
   // CHECK:      %[[C2:.*]] = constant 2 : index
   %c2 = constant 2 : index
   // CHECK:      %[[C3:.*]] = constant 3 : index
   %c3 = constant 3 : index
   // CHECK:      %[[C6:.*]] = constant 6 : index
   %c6 = constant 6 : index
   // 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
+// 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>>
 }
 
 // CHECK-LABEL: @vector_fma
 func @vector_fma(%a: vector<8xf32>, %b: vector<8x4xf32>) {
   // CHECK: vector.fma %{{.*}} : vector<8xf32>
   vector.fma %a, %a, %a : vector<8xf32>
   // CHECK: vector.fma %{{.*}} : vector<8x4xf32>
   vector.fma %b, %b, %b : vector<8x4xf32>
   return
 }
 
-// CHECK-LABEL: reduce_fp
+// CHECK-LABEL: @reduce_fp
 func @reduce_fp(%arg0: vector<16xf32>, %arg1: f32) -> f32 {
   // CHECK:    vector.reduction "add", %{{.*}} : vector<16xf32> into f32
   vector.reduction "add", %arg0 : vector<16xf32> into f32
   // CHECK:    vector.reduction "add", %{{.*}}, %{{.*}} : vector<16xf32> into f32
   vector.reduction "add", %arg0, %arg1 : vector<16xf32> into f32
   // CHECK:    vector.reduction "mul", %{{.*}} : vector<16xf32> into f32
   vector.reduction "mul", %arg0 : vector<16xf32> into f32
   // CHECK:    vector.reduction "mul", %{{.*}}, %{{.*}} : vector<16xf32> into f32
   vector.reduction "mul", %arg0, %arg1 : vector<16xf32> into f32
   // CHECK:    vector.reduction "min", %{{.*}} : vector<16xf32> into f32
   vector.reduction "min", %arg0 : vector<16xf32> into f32
   // CHECK:    %[[X:.*]] = vector.reduction "max", %{{.*}} : vector<16xf32> into f32
   %0 = vector.reduction "max", %arg0 : vector<16xf32> into f32
   // CHECK:    return %[[X]] : f32
   return %0 : f32
 }
 
-// CHECK-LABEL: reduce_int
+// CHECK-LABEL: @reduce_int
 func @reduce_int(%arg0: vector<16xi32>) -> i32 {
   // CHECK:    vector.reduction "add", %{{.*}} : vector<16xi32> into i32
   vector.reduction "add", %arg0 : vector<16xi32> into i32
   // CHECK:    vector.reduction "mul", %{{.*}} : vector<16xi32> into i32
   vector.reduction "mul", %arg0 : vector<16xi32> into i32
   // CHECK:    vector.reduction "min", %{{.*}} : vector<16xi32> into i32
   vector.reduction "min", %arg0 : vector<16xi32> into i32
   // CHECK:    vector.reduction "max", %{{.*}} : vector<16xi32> into i32
   vector.reduction "max", %arg0 : vector<16xi32> into i32
   // CHECK:    vector.reduction "and", %{{.*}} : vector<16xi32> into 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
+// CHECK-LABEL: @transpose_fp
 func @transpose_fp(%arg0: vector<3x7xf32>) -> vector<7x3xf32> {
+  // CHECK: %[[X:.*]] = vector.transpose %{{.*}}, [1, 0] : vector<3x7xf32> to vector<7x3xf32>
   %0 = vector.transpose %arg0, [1, 0] : vector<3x7xf32> to vector<7x3xf32>
+  // CHECK: return %[[X]] : vector<7x3xf32>
   return %0 : vector<7x3xf32>
 }
 
-// CHECK-LABEL: transpose_int
+// CHECK-LABEL: @transpose_int
 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: @flat_transpose_fp
+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>
+}
+
+// CHECK-LABEL: @flat_transpose_int
+func @flat_transpose_int(%arg0: vector<16xi32>) -> vector<16xi32> {
+  // CHECK: %[[X:.*]] = vector.flat_transpose %{{.*}} {columns = 8 : i32, rows = 2 : i32} : vector<16xi32> -> vector<16xi32>
+  %0 = vector.flat_transpose %arg0 { rows = 2: i32, columns = 8: i32 } : vector<16xi32> -> vector<16xi32>
+  // CHECK: return %[[X]] : vector<16xi32>
+  return %0 : vector<16xi32>
+}