Skip to content

Commit f24c443

Browse files
banach-spacebanach-space
authored
[mlir][SVE] Add an e2e test for vector.contract (#69845)
Adds an end-to-end test for `vector.contract` that targets SVE (i.e. scalable vectors). Note that this requires lifting the restriction on `vector.outerproduct` (to which `vector.contract` is lowered to) that would deem the following as invalid by the Op verifier (*): ``` vector.outerproduct %27, %28, %26 {kind = #vector.kind<add>} : vector<3xf32>, vector<[2]xf32> ``` This is indeed valid as the end-to-end test demonstrates (at least when compiling for SVE).
1 parent b13e9ef commit f24c443

File tree

4 files changed

+157
-19
lines changed

4 files changed

+157
-19
lines changed

‎mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+6-3
Original file line numberDiff line numberDiff line change
@@ -3067,9 +3067,12 @@ LogicalResult OuterProductOp::verify() {
30673067
return emitOpError("expected #1 operand dim to match result dim #1");
30683068
if (vRHS.getDimSize(0) != vRES.getDimSize(1))
30693069
return emitOpError("expected #2 operand dim to match result dim #2");
3070-
if (vRHS.isScalable() != vLHS.isScalable())
3071-
return emitOpError("expected either all or none of vector operands #1 "
3072-
"and #2 to be scalable");
3070+
if (vLHS.isScalable() && !vRHS.isScalable()) {
3071+
// This restriction reflects what's currently supported in terms of
3072+
// scalable vectors. However, we could relax this if there's a use case.
3073+
return emitOpError(
3074+
"expected either both or only #2 operand dim to be scalable");
3075+
}
30733076
} else {
30743077
// An AXPY operation.
30753078
if (vRES.getRank() != 1)

‎mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir

+44-15
Original file line numberDiff line numberDiff line change
@@ -79,21 +79,21 @@ func.func @masked_extract_contract2_scalable_parallel_dim(%arg0: vector<[2]x3xf3
7979
}
8080

8181
// CHECK-LABEL: func.func @masked_extract_contract4(
82-
// CHECK-SAME: %[[VAL_0:.*]]: vector<3x5xf32>,
83-
// CHECK-SAME: %[[VAL_1:.*]]: vector<5x7xf32>,
84-
// CHECK-SAME: %[[VAL_2:.*]]: vector<3x7xf32>,
85-
// CHECK-SAME: %[[VAL_3:.*]]: vector<3x7x5xi1>) -> vector<3x7xf32> {
86-
// CHECK: %[[VAL_5:.*]] = vector.transpose %[[VAL_3]], [2, 0, 1] : vector<3x7x5xi1> to vector<5x3x7xi1>
87-
// CHECK: %[[VAL_8:.*]] = vector.extract %[[VAL_5]][0] : vector<3x7xi1> from vector<5x3x7xi1>
88-
// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
89-
// CHECK: %[[VAL_12:.*]] = vector.extract %[[VAL_5]][1] : vector<3x7xi1> from vector<5x3x7xi1>
90-
// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_12]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
91-
// CHECK: %[[VAL_16:.*]] = vector.extract %[[VAL_5]][2] : vector<3x7xi1> from vector<5x3x7xi1>
92-
// CHECK: %[[VAL_17:.*]] = vector.mask %[[VAL_16]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
93-
// CHECK: %[[VAL_20:.*]] = vector.extract %[[VAL_5]][3] : vector<3x7xi1> from vector<5x3x7xi1>
94-
// CHECK: %[[VAL_21:.*]] = vector.mask %[[VAL_20]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
95-
// CHECK: %[[VAL_24:.*]] = vector.extract %[[VAL_5]][4] : vector<3x7xi1> from vector<5x3x7xi1>
96-
// CHECK: %[[VAL_25:.*]] = vector.mask %[[VAL_24]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
82+
// CHECK-SAME: %{{.*}}: vector<3x5xf32>,
83+
// CHECK-SAME: %{{.*}}: vector<5x7xf32>,
84+
// CHECK-SAME: %{{.*}}: vector<3x7xf32>,
85+
// CHECK-SAME: %[[IN_MASK:.*]]: vector<3x7x5xi1>) -> vector<3x7xf32> {
86+
// CHECK: %[[T_MASK:.*]] = vector.transpose %[[IN_MASK]], [2, 0, 1] : vector<3x7x5xi1> to vector<5x3x7xi1>
87+
// CHECK: %[[T_MASK_R0:.*]] = vector.extract %[[T_MASK]][0] : vector<3x7xi1> from vector<5x3x7xi1>
88+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R0]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
89+
// CHECK: %[[T_MASK_R1:.*]] = vector.extract %[[T_MASK]][1] : vector<3x7xi1> from vector<5x3x7xi1>
90+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R1]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
91+
// CHECK: %[[T_MASK_R2:.*]] = vector.extract %[[T_MASK]][2] : vector<3x7xi1> from vector<5x3x7xi1>
92+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R2]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
93+
// CHECK: %[[T_MASK_R3:.*]] = vector.extract %[[T_MASK]][3] : vector<3x7xi1> from vector<5x3x7xi1>
94+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R3]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
95+
// CHECK: %[[T_MASK_R4:.*]] = vector.extract %[[T_MASK]][4] : vector<3x7xi1> from vector<5x3x7xi1>
96+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R4]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
9797

9898
func.func @masked_extract_contract4(%arg0: vector<3x5xf32>,
9999
%arg1: vector<5x7xf32>,
@@ -104,6 +104,35 @@ func.func @masked_extract_contract4(%arg0: vector<3x5xf32>,
104104
return %0 : vector<3x7xf32>
105105
}
106106

107+
// CHECK-LABEL: func.func @masked_extract_contract4_scalable_J_dim(
108+
// CHECK-SAME: %{{.*}}: vector<3x5xf32>,
109+
// CHECK-SAME: %{{.*}}: vector<5x[7]xf32>,
110+
// CHECK-SAME: %{{.*}}: vector<3x[7]xf32>,
111+
// CHECK-SAME: %[[IN_MASK:.*]]: vector<3x[7]x5xi1>) -> vector<3x[7]xf32> {
112+
// CHECK: %[[T_MASK:.*]] = vector.transpose %[[IN_MASK]], [2, 0, 1] : vector<3x[7]x5xi1> to vector<5x3x[7]xi1>
113+
// CHECK: %[[T_MASK_R0:.*]] = vector.extract %[[T_MASK]][0] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
114+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R0]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
115+
// CHECK: %[[T_MASK_R1:.*]] = vector.extract %[[T_MASK]][1] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
116+
// CHECK: %[[VAL_13:.*]] = vector.mask %[[T_MASK_R1]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
117+
// CHECK: %[[T_MASK_R2:.*]] = vector.extract %[[T_MASK]][2] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
118+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R2]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
119+
// CHECK: %[[T_MASK_R3:.*]] = vector.extract %[[T_MASK]][3] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
120+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R3]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
121+
// CHECK: %[[T_MASK_R4:.*]] = vector.extract %[[T_MASK]][4] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
122+
// CHECK: %{{.*}} = vector.mask %[[T_MASK_R4]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
123+
124+
// Note that only the J dimension is scalable in this example. In theory, all
125+
// dimensions could be be scalable, but there is no target yet for which this
126+
// would make sense.
127+
func.func @masked_extract_contract4_scalable_J_dim(%arg0: vector<3x5xf32>,
128+
%arg1: vector<5x[7]xf32>,
129+
%arg2: vector<3x[7]xf32>,
130+
%m : vector<3x[7]x5xi1>) -> vector<3x[7]xf32> {
131+
%0 = vector.mask %m { vector.contract #matmat_trait %arg0, %arg1, %arg2
132+
: vector<3x5xf32>, vector<5x[7]xf32> into vector<3x[7]xf32> } : vector<3x[7]x5xi1> -> vector<3x[7]xf32>
133+
return %0 : vector<3x[7]xf32>
134+
}
135+
107136
// CHECK-LABEL: func @matmul
108137
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x4xf32>,
109138
// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: vector<4x3xf32>,

‎mlir/test/Dialect/Vector/vector-scalable-outerproduct.mlir

+4-1
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,12 @@ func.func @invalid_outerproduct(%src : memref<?xf32>) {
2121
%0 = vector.load %src[%idx] : memref<?xf32>, vector<[4]xf32>
2222
%1 = vector.load %src[%idx] : memref<?xf32>, vector<4xf32>
2323

24-
// expected-error @+1 {{expected either all or none of vector operands #1 and #2 to be scalable}}
24+
// expected-error @+1 {{expected either both or only #2 operand dim to be scalable}}
2525
%op = vector.outerproduct %0, %1 : vector<[4]xf32>, vector<4xf32>
26+
27+
return
2628
}
29+
2730
// -----
2831

2932
func.func @invalid_outerproduct1(%src : memref<?xf32>) {

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir

+103
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
// DEFINE: %{compile} = mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule\
2+
// DEFINE: -cse -canonicalize -convert-vector-to-scf -arm-sve-legalize-vector-storage\
3+
// DEFINE: -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm -o %t
4+
// DEFINE: %{entry} =
5+
// DEFINE: %{run} = %mcr_aarch64_cmd %t -e=%{entry} -entry-point-result=void --march=aarch64 --mattr="+sve" -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext
6+
7+
// This check whether the files compiles and generates a temporary that will be executed further down.
8+
// RUN: %{compile}
9+
10+
// REDEFINE: %{entry} = matmul_i32
11+
// RUN: %{run} | FileCheck %s --check-prefix=I32
12+
13+
// REDEFINE: %{entry} = matmul_f32
14+
// RUN: %{run} | FileCheck %s --check-prefix=F32
15+
16+
// NOTE: These tests are meant to complement the integration tests from:
17+
// * ../test-contraction.mlir
18+
// (tests with fixed width vectors). Rather than duplicating those tests, this
19+
// file focuses on excercissing scalable vectors in a few most common cases.
20+
21+
// TODO: Masks + matvec + dot product
22+
23+
#matmat_accesses = [
24+
affine_map<(i, j, k) -> (i, k)>,
25+
affine_map<(i, j, k) -> (k, j)>,
26+
affine_map<(i, j, k) -> (i, j)>
27+
]
28+
#matmat_trait = {
29+
indexing_maps = #matmat_accesses,
30+
iterator_types = ["parallel", "parallel", "reduction"]
31+
}
32+
33+
func.func @matmul_i32() {
34+
// Setup vector A:
35+
%vector_a = arith.constant dense<123> : vector<3x5xi32>
36+
37+
// Setup vector B:
38+
%vector_b = arith.constant dense<123> : vector<5x[2]xi32>
39+
40+
// Setup vector C:
41+
%vector_c = arith.constant dense<314> : vector<3x[2]xi32>
42+
43+
// Matmul
44+
%0 = vector.contract #matmat_trait %vector_a, %vector_b, %vector_c
45+
: vector<3x5xi32>, vector<5x[2]xi32> into vector<3x[2]xi32>
46+
47+
// Print the output
48+
%slice1 = vector.extract %0[0] : vector<[2]xi32> from vector<3x[2]xi32>
49+
// I32: ( 75959, 75959
50+
vector.print %slice1 : vector<[2]xi32>
51+
%slice2 = vector.extract %0[1] : vector<[2]xi32> from vector<3x[2]xi32>
52+
// I32-NEXT: ( 75959, 75959
53+
vector.print %slice2 : vector<[2]xi32>
54+
%slice3 = vector.extract %0[2] : vector<[2]xi32> from vector<3x[2]xi32>
55+
// I32-NEXT: ( 75959, 75959
56+
vector.print %slice3 : vector<[2]xi32>
57+
58+
// CHECK: SVE: END OF TEST OUTPUT
59+
vector.print str "SVE: END OF TEST OUTPUT"
60+
61+
return
62+
}
63+
64+
func.func @matmul_f32() {
65+
// Setup vector A:
66+
%vector_a = arith.constant dense<1.23> : vector<3x5xf32>
67+
68+
// Setup vector B:
69+
%vector_b = arith.constant dense<1.23> : vector<5x[2]xf32>
70+
71+
// Setup vector C:
72+
%vector_c = arith.constant dense<3.14> : vector<3x[2]xf32>
73+
74+
// Matmul
75+
%0 = vector.contract #matmat_trait %vector_a, %vector_b, %vector_c
76+
: vector<3x5xf32>, vector<5x[2]xf32> into vector<3x[2]xf32>
77+
78+
// Print the output
79+
%slice1 = vector.extract %0[0] : vector<[2]xf32> from vector<3x[2]xf32>
80+
// F32: ( 10.7045, 10.7045
81+
vector.print %slice1 : vector<[2]xf32>
82+
%slice2 = vector.extract %0[1] : vector<[2]xf32> from vector<3x[2]xf32>
83+
// F32-NEXT: ( 10.7045, 10.7045
84+
vector.print %slice2 : vector<[2]xf32>
85+
%slice3 = vector.extract %0[2] : vector<[2]xf32> from vector<3x[2]xf32>
86+
// F32-NEXT: ( 10.7045, 10.7045
87+
vector.print %slice3 : vector<[2]xf32>
88+
89+
// CHECK: SVE: END OF TEST OUTPUT
90+
vector.print str "SVE: END OF TEST OUTPUT"
91+
92+
return
93+
}
94+
95+
transform.sequence failures(propagate) {
96+
^bb1(%module_op: !transform.any_op):
97+
%f = transform.structured.match ops{["func.func"]} in %module_op
98+
: (!transform.any_op) -> !transform.any_op
99+
100+
transform.apply_patterns to %f {
101+
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
102+
} : !transform.any_op
103+
}

0 commit comments

Comments
 (0)