435114f9fe
To keep the test filenames consistent, this patch:
* removes "test-" from file names (there used to be a mix of
"test-feature-1.mlir" and "feature-2.mlir"),
* replaces "_" with "-" (there used to be a mix of "feature-3.mlir"
and "feature_4.mlir").
Only files under test/Integration/Dialect/Vector/CPU are updated.
384 lines
12 KiB
MLIR
384 lines
12 KiB
MLIR
// RUN: mlir-opt %s -convert-vector-to-scf -convert-scf-to-cf -convert-vector-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \
|
|
// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
|
|
// RUN: -shared-libs=%mlir_c_runner_utils | \
|
|
// RUN: FileCheck %s
|
|
|
|
#dotp_accesses = [
|
|
affine_map<(i) -> (i)>,
|
|
affine_map<(i) -> (i)>,
|
|
affine_map<(i) -> ()>
|
|
]
|
|
#dotp_trait = {
|
|
indexing_maps = #dotp_accesses,
|
|
iterator_types = ["reduction"]
|
|
}
|
|
|
|
#matvec_accesses = [
|
|
affine_map<(i, j) -> (i, j)>,
|
|
affine_map<(i, j) -> (j)>,
|
|
affine_map<(i, j) -> (i)>
|
|
]
|
|
#matvec_trait = {
|
|
indexing_maps = #matvec_accesses,
|
|
iterator_types = ["parallel", "reduction"]
|
|
}
|
|
|
|
#mattransvec_accesses = [
|
|
affine_map<(i, j) -> (j, i)>,
|
|
affine_map<(i, j) -> (j)>,
|
|
affine_map<(i, j) -> (i)>
|
|
]
|
|
#mattransvec_trait = {
|
|
indexing_maps = #mattransvec_accesses,
|
|
iterator_types = ["parallel", "reduction"]
|
|
}
|
|
|
|
#matmat_accesses = [
|
|
affine_map<(i, j, k) -> (i, k)>,
|
|
affine_map<(i, j, k) -> (k, j)>,
|
|
affine_map<(i, j, k) -> (i, j)>
|
|
]
|
|
#matmat_trait = {
|
|
indexing_maps = #matmat_accesses,
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
|
|
#mattransmat_accesses = [
|
|
affine_map<(i, j, k) -> (k, i)>,
|
|
affine_map<(i, j, k) -> (k, j)>,
|
|
affine_map<(i, j, k) -> (i, j)>
|
|
]
|
|
#mattransmat_trait = {
|
|
indexing_maps = #mattransmat_accesses,
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
|
|
#matmattrans_accesses = [
|
|
affine_map<(i, j, k) -> (i, k)>,
|
|
affine_map<(i, j, k) -> (j, k)>,
|
|
affine_map<(i, j, k) -> (i, j)>
|
|
]
|
|
#matmattrans_trait = {
|
|
indexing_maps = #matmattrans_accesses,
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
|
|
#mattransmattrans_accesses = [
|
|
affine_map<(i, j, k) -> (k, i)>,
|
|
affine_map<(i, j, k) -> (j, k)>,
|
|
affine_map<(i, j, k) -> (i, j)>
|
|
]
|
|
#mattransmattrans_trait = {
|
|
indexing_maps = #mattransmattrans_accesses,
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
|
|
#matmat_then_trans_accesses = [
|
|
affine_map<(i, j, k) -> (i, k)>,
|
|
affine_map<(i, j, k) -> (k, j)>,
|
|
affine_map<(i, j, k) -> (j, i)>
|
|
]
|
|
#matmat_then_trans_trait = {
|
|
indexing_maps = #matmat_then_trans_accesses,
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
|
|
#contract2d_accesses = [
|
|
affine_map<(i, j) -> (i, j)>,
|
|
affine_map<(i, j) -> (i, j)>,
|
|
affine_map<(i, j) -> ()>
|
|
]
|
|
#contract2d_trait = {
|
|
indexing_maps = #contract2d_accesses,
|
|
iterator_types = ["reduction", "reduction"]
|
|
}
|
|
|
|
#contract2d_alt_accesses = [
|
|
affine_map<(i, j) -> (j, i)>,
|
|
affine_map<(i, j) -> (j, i)>,
|
|
affine_map<(i, j) -> ()>
|
|
]
|
|
#contract2d_alt_trait = {
|
|
indexing_maps = #contract2d_alt_accesses,
|
|
iterator_types = ["reduction", "reduction"]
|
|
}
|
|
|
|
#contract2d_trans_accesses = [
|
|
affine_map<(i, j) -> (i, j)>,
|
|
affine_map<(i, j) -> (j, i)>,
|
|
affine_map<(i, j) -> ()>
|
|
]
|
|
#contract2d_trans_trait = {
|
|
indexing_maps = #contract2d_trans_accesses,
|
|
iterator_types = ["reduction", "reduction"]
|
|
}
|
|
|
|
#contract2d_trans_alt_accesses = [
|
|
affine_map<(i, j) -> (j, i)>,
|
|
affine_map<(i, j) -> (i, j)>,
|
|
affine_map<(i, j) -> ()>
|
|
]
|
|
#contract2d_trans_alt_trait = {
|
|
indexing_maps = #contract2d_trans_alt_accesses,
|
|
iterator_types = ["reduction", "reduction"]
|
|
}
|
|
|
|
#column_major_matmat_accesses = [
|
|
affine_map<(i, j, k) -> (k, j)>,
|
|
affine_map<(i, j, k) -> (i, k)>,
|
|
affine_map<(i, j, k) -> (j, i)>
|
|
]
|
|
#column_major_matmat_trait = {
|
|
indexing_maps = #column_major_matmat_accesses,
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
|
|
func.func @entry() {
|
|
%f0 = arith.constant 0.0: f32
|
|
%f1 = arith.constant 1.0: f32
|
|
%f2 = arith.constant 2.0: f32
|
|
%f3 = arith.constant 3.0: f32
|
|
%f4 = arith.constant 4.0: f32
|
|
%f5 = arith.constant 5.0: f32
|
|
%f6 = arith.constant 6.0: f32
|
|
%f7 = arith.constant 7.0: f32
|
|
%f8 = arith.constant 8.0: f32
|
|
|
|
// Zero vectors.
|
|
%z1 = vector.broadcast %f0 : f32 to vector<2xf32>
|
|
%z2 = vector.broadcast %f0 : f32 to vector<2x2xf32>
|
|
%z3 = vector.broadcast %f0 : f32 to vector<3x4xf32>
|
|
|
|
// Construct test vectors.
|
|
%0 = vector.broadcast %f1 : f32 to vector<2xf32>
|
|
%a = vector.insert %f2, %0[1] : f32 into vector<2xf32>
|
|
%1 = vector.broadcast %f3 : f32 to vector<2xf32>
|
|
%b = vector.insert %f4, %1[1] : f32 into vector<2xf32>
|
|
%2 = vector.broadcast %f5 : f32 to vector<2xf32>
|
|
%c = vector.insert %f6, %2[1] : f32 into vector<2xf32>
|
|
%3 = vector.broadcast %f7 : f32 to vector<2xf32>
|
|
%d = vector.insert %f8, %3[1] : f32 into vector<2xf32>
|
|
|
|
vector.print %a : vector<2xf32>
|
|
vector.print %b : vector<2xf32>
|
|
vector.print %c : vector<2xf32>
|
|
vector.print %d : vector<2xf32>
|
|
//
|
|
// test vectors:
|
|
//
|
|
// CHECK: ( 1, 2 )
|
|
// CHECK: ( 3, 4 )
|
|
// CHECK: ( 5, 6 )
|
|
// CHECK: ( 7, 8 )
|
|
|
|
// Construct test matrices.
|
|
%4 = vector.broadcast %f0 : f32 to vector<2x2xf32>
|
|
%5 = vector.insert %a, %4[0] : vector<2xf32> into vector<2x2xf32>
|
|
%A = vector.insert %b, %5[1] : vector<2xf32> into vector<2x2xf32>
|
|
%6 = vector.broadcast %f0 : f32 to vector<2x2xf32>
|
|
%7 = vector.insert %c, %6[0] : vector<2xf32> into vector<2x2xf32>
|
|
%B = vector.insert %d, %7[1] : vector<2xf32> into vector<2x2xf32>
|
|
%8 = vector.broadcast %f0 : f32 to vector<3x2xf32>
|
|
%9 = vector.insert %a, %8[0] : vector<2xf32> into vector<3x2xf32>
|
|
%10 = vector.insert %b, %9[1] : vector<2xf32> into vector<3x2xf32>
|
|
%C = vector.insert %c, %10[2] : vector<2xf32> into vector<3x2xf32>
|
|
%cst = arith.constant dense<0.000000e+00> : vector<2x4xf32>
|
|
%11 = vector.insert_strided_slice %A, %cst {offsets = [0, 0], strides = [1, 1]} : vector<2x2xf32> into vector<2x4xf32>
|
|
%D = vector.insert_strided_slice %B, %11 {offsets = [0, 2], strides = [1, 1]} : vector<2x2xf32> into vector<2x4xf32>
|
|
|
|
vector.print %A : vector<2x2xf32>
|
|
vector.print %B : vector<2x2xf32>
|
|
vector.print %C : vector<3x2xf32>
|
|
vector.print %D : vector<2x4xf32>
|
|
//
|
|
// test matrices:
|
|
//
|
|
// CHECK: ( ( 1, 2 ), ( 3, 4 ) )
|
|
// CHECK: ( ( 5, 6 ), ( 7, 8 ) )
|
|
// CHECK: ( ( 1, 2 ), ( 3, 4 ), ( 5, 6 ) )
|
|
// CHECK: ( ( 1, 2, 5, 6 ), ( 3, 4, 7, 8 ) )
|
|
|
|
// Contraction: dot-product a x b
|
|
%dp1 = vector.contract #dotp_trait %a, %b, %f0
|
|
: vector<2xf32>, vector<2xf32> into f32
|
|
%dp2 = vector.contract #dotp_trait %a, %b, %f1
|
|
: vector<2xf32>, vector<2xf32> into f32
|
|
|
|
vector.print %dp1 : f32
|
|
vector.print %dp2 : f32
|
|
//
|
|
// dot products:
|
|
//
|
|
// CHECK: 11
|
|
// CHECK: 12
|
|
|
|
// Contraction: matrix-vector A x c
|
|
%mv1 = vector.contract #matvec_trait %A, %c, %z1
|
|
: vector<2x2xf32>, vector<2xf32> into vector<2xf32>
|
|
%mv2 = vector.contract #matvec_trait %A, %c, %a
|
|
: vector<2x2xf32>, vector<2xf32> into vector<2xf32>
|
|
|
|
vector.print %mv1 : vector<2xf32>
|
|
vector.print %mv2 : vector<2xf32>
|
|
//
|
|
// matrix x vector:
|
|
//
|
|
// CHECK: ( 17, 39 )
|
|
// CHECK: ( 18, 41 )
|
|
|
|
// Contraction: matrix-trans-vector A^T x c
|
|
%mv3 = vector.contract #mattransvec_trait %A, %c, %z1
|
|
: vector<2x2xf32>, vector<2xf32> into vector<2xf32>
|
|
%mv4 = vector.contract #mattransvec_trait %A, %c, %a
|
|
: vector<2x2xf32>, vector<2xf32> into vector<2xf32>
|
|
|
|
vector.print %mv3 : vector<2xf32>
|
|
vector.print %mv4 : vector<2xf32>
|
|
//
|
|
// matrix x vector:
|
|
//
|
|
// CHECK: ( 23, 34 )
|
|
// CHECK: ( 24, 36 )
|
|
|
|
// Contraction: matrix-matrix A x B
|
|
%mm1 = vector.contract #matmat_trait %A, %B, %z2
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
%mm2 = vector.contract #matmat_trait %A, %B, %A
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
|
|
vector.print %mm1 : vector<2x2xf32>
|
|
vector.print %mm2 : vector<2x2xf32>
|
|
//
|
|
// matrix x matrix:
|
|
//
|
|
// CHECK: ( ( 19, 22 ), ( 43, 50 ) )
|
|
// CHECK: ( ( 20, 24 ), ( 46, 54 ) )
|
|
|
|
// Contraction: matrix-matrix A x B where A, B, C have column-major layout.
|
|
// ( 1 * 5 + 3 * 6 = 23, 2 * 5 + 4 * 6 = 34)
|
|
// ( 1 * 7 + 3 * 8 = 31, 2 * 7 + 4 * 8 = 46)
|
|
// +
|
|
// ( ( 1, 2 ), ( 3, 4 ) )
|
|
%llvm_matrix_column_major_mm0 =
|
|
vector.contract #column_major_matmat_trait %A, %B, %z2
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
%llvm_matrix_column_major_mm1 =
|
|
vector.contract #column_major_matmat_trait %A, %B, %A
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
|
|
vector.print %llvm_matrix_column_major_mm0 : vector<2x2xf32>
|
|
vector.print %llvm_matrix_column_major_mm1 : vector<2x2xf32>
|
|
//
|
|
// matrix x matrix:
|
|
//
|
|
// CHECK: ( ( 23, 31 ), ( 34, 46 ) )
|
|
// CHECK: ( ( 24, 33 ), ( 37, 50 ) )
|
|
|
|
// Contraction: matrix-trans-matrix A^T x B
|
|
%mm3 = vector.contract #mattransmat_trait %A, %B, %z2
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
%mm4 = vector.contract #mattransmat_trait %A, %B, %A
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
|
|
vector.print %mm3 : vector<2x2xf32>
|
|
vector.print %mm4 : vector<2x2xf32>
|
|
//
|
|
// matrix x matrix:
|
|
//
|
|
// CHECK: ( ( 26, 30 ), ( 38, 44 ) )
|
|
// CHECK: ( ( 27, 32 ), ( 41, 48 ) )
|
|
|
|
// Contraction: matrix-matrix-trans A x B^T
|
|
%mm5 = vector.contract #matmattrans_trait %A, %B, %z2
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
%mm6 = vector.contract #matmattrans_trait %A, %B, %A
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
|
|
vector.print %mm5 : vector<2x2xf32>
|
|
vector.print %mm6 : vector<2x2xf32>
|
|
//
|
|
// matrix x matrix:
|
|
//
|
|
// CHECK: ( ( 17, 23 ), ( 39, 53 ) )
|
|
// CHECK: ( ( 18, 25 ), ( 42, 57 ) )
|
|
|
|
// Contraction: matrix-trans-matrix-trans A^T x B^T
|
|
%mm7 = vector.contract #mattransmattrans_trait %A, %B, %z2
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
%mm8 = vector.contract #mattransmattrans_trait %A, %B, %A
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
|
|
vector.print %mm7 : vector<2x2xf32>
|
|
vector.print %mm8 : vector<2x2xf32>
|
|
//
|
|
// matrix x matrix:
|
|
//
|
|
// CHECK: ( ( 23, 31 ), ( 34, 46 ) )
|
|
// CHECK: ( ( 24, 33 ), ( 37, 50 ) )
|
|
|
|
// Contraction: matrix-matrix-then-trans (A x B)^T
|
|
%mm9 = vector.contract #matmat_then_trans_trait %A, %B, %z2
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
%mm10 = vector.contract #matmat_then_trans_trait %A, %B, %A
|
|
: vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
|
|
|
|
vector.print %mm9 : vector<2x2xf32>
|
|
vector.print %mm10 : vector<2x2xf32>
|
|
//
|
|
// matrix x matrix:
|
|
//
|
|
// CHECK: ( ( 19, 43 ), ( 22, 50 ) )
|
|
// CHECK: ( ( 20, 45 ), ( 25, 54 ) )
|
|
|
|
// Contraction: matrix-matrix C x D
|
|
%mm11 = vector.contract #matmat_trait %C, %D, %z3
|
|
: vector<3x2xf32>, vector<2x4xf32> into vector<3x4xf32>
|
|
%mm12 = vector.contract #matmat_trait %C, %D, %mm11
|
|
: vector<3x2xf32>, vector<2x4xf32> into vector<3x4xf32>
|
|
|
|
vector.print %mm11 : vector<3x4xf32>
|
|
vector.print %mm12 : vector<3x4xf32>
|
|
// CHECK: ( ( 7, 10, 19, 22 ), ( 15, 22, 43, 50 ), ( 23, 34, 67, 78 ) )
|
|
// CHECK: ( ( 14, 20, 38, 44 ), ( 30, 44, 86, 100 ), ( 46, 68, 134, 156 ) )
|
|
|
|
// Contractions in 2D.
|
|
%c1 = vector.contract #contract2d_trait %A, %B, %f0
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c2 = vector.contract #contract2d_trait %A, %B, %f1
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c3 = vector.contract #contract2d_alt_trait %A, %B, %f0
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c4 = vector.contract #contract2d_alt_trait %A, %B, %f1
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c5 = vector.contract #contract2d_trans_trait %A, %B, %f0
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c6 = vector.contract #contract2d_trans_trait %A, %B, %f1
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c7 = vector.contract #contract2d_trans_alt_trait %A, %B, %f0
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
%c8 = vector.contract #contract2d_trans_alt_trait %A, %B, %f1
|
|
: vector<2x2xf32>, vector<2x2xf32> into f32
|
|
|
|
vector.print %c1 : f32
|
|
vector.print %c2 : f32
|
|
vector.print %c3 : f32
|
|
vector.print %c4 : f32
|
|
vector.print %c5 : f32
|
|
vector.print %c6 : f32
|
|
vector.print %c7 : f32
|
|
vector.print %c8 : f32
|
|
//
|
|
// 2D contractions:
|
|
//
|
|
// CHECK: 70
|
|
// CHECK: 71
|
|
// CHECK: 70
|
|
// CHECK: 71
|
|
// CHECK: 69
|
|
// CHECK: 70
|
|
// CHECK: 69
|
|
// CHECK: 70
|
|
|
|
return
|
|
}
|