Skip to content

Commit 42a82b1

Browse files
[mlir][linalg][NFC] Add test case for memref vectorization
Add test cases for vectorizing linalg.matmul and linalg.copy on tensors. Differential Revision: https://reviews.llvm.org/D153357
1 parent dae8c72 commit 42a82b1

File tree

1 file changed

+45
-0
lines changed

1 file changed

+45
-0
lines changed

mlir/test/Dialect/Linalg/transform-op-vectorize.mlir

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,51 @@ transform.sequence failures(propagate) {
2525

2626
// -----
2727

28+
// CHECK-LABEL: @vectorize_matmul_memref
29+
// CHECK-SAME: %[[A:.*]]: memref<24x12xf32>
30+
// CHECK-SAME: %[[B:.*]]: memref<12x25xf32>
31+
// CHECK-SAME: %[[C:.*]]: memref<24x25xf32>
32+
func.func @vectorize_matmul_memref(%arg0: memref<24x12xf32>,
33+
%arg1: memref<12x25xf32>,
34+
%arg2: memref<24x25xf32>) {
35+
// CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
36+
// CHECK: %[[vB:.+]] = vector.transfer_read %[[B]]
37+
// CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
38+
// CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
39+
// CHECK: vector.transfer_write %[[vR]], %[[C]]
40+
linalg.matmul ins(%arg0, %arg1 : memref<24x12xf32>, memref<12x25xf32>) outs(%arg2 : memref<24x25xf32>)
41+
return
42+
}
43+
44+
transform.sequence failures(propagate) {
45+
^bb1(%arg1: !transform.any_op):
46+
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
47+
%1 = get_closest_isolated_parent %0 : (!transform.any_op) -> !transform.any_op
48+
%2 = transform.structured.vectorize %1 : (!transform.any_op) -> !transform.any_op
49+
}
50+
51+
// -----
52+
53+
// CHECK-LABEL: @vectorize_copy_memref
54+
// CHECK-SAME: %[[A:.*]]: memref<100x100xf32>,
55+
// CHECK-SAME: %[[B:.*]]: memref<100x100xf32>
56+
func.func @vectorize_copy_memref(%arg0: memref<100x100xf32>,
57+
%arg1: memref<100x100xf32>) {
58+
// CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
59+
// CHECK: vector.transfer_write %[[vA]], %[[B]]
60+
linalg.copy ins(%arg0 : memref<100x100xf32>) outs(%arg1 : memref<100x100xf32>)
61+
return
62+
}
63+
64+
transform.sequence failures(propagate) {
65+
^bb1(%arg1: !transform.any_op):
66+
%0 = transform.structured.match ops{["linalg.copy"]} in %arg1 : (!transform.any_op) -> !transform.any_op
67+
%1 = get_closest_isolated_parent %0 : (!transform.any_op) -> !transform.any_op
68+
%2 = transform.structured.vectorize %1 : (!transform.any_op) -> !transform.any_op
69+
}
70+
71+
// -----
72+
2873
#map0 = affine_map<()[s0] -> (-s0 + 12, 7)>
2974
#map1 = affine_map<()[s0] -> (-s0 + 7)>
3075

0 commit comments

Comments
 (0)