Skip to content

Commit 64025b8

Browse files
authored
[mlir][SVE] Add an e2e test for vectorization of linalg.matmul (#69592)
--delete-branch
1 parent 658874e commit 64025b8

File tree

1 file changed

+71
-0
lines changed
  • mlir/test/Integration/Dialect/Linalg/CPU/ArmSVE

1 file changed

+71
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule \
2+
// RUN: -one-shot-bufferize -func-bufferize -cse -canonicalize -convert-vector-to-scf -arm-sve-legalize-vector-storage \
3+
// RUN: -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm | \
4+
// RUN: %mcr_aarch64_cmd -e=entry -entry-point-result=void --march=aarch64 --mattr="+sve" -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils | \
5+
// RUN: FileCheck %s
6+
7+
func.func @entry() {
8+
%c1 = arith.constant 1 : index
9+
%c2 = arith.constant 2 : index
10+
%c4 = arith.constant 4 : index
11+
%c0 = arith.constant 0 : index
12+
%step = arith.constant 1 : index
13+
%c0_f32 = arith.constant 0.0 : f32
14+
15+
%vscale = vector.vscale
16+
%vl_fp = arith.muli %c4, %vscale : index
17+
%A_alloc = bufferization.alloc_tensor(%c2, %c1) : tensor<?x?xf32>
18+
%B_alloc = bufferization.alloc_tensor(%c1, %vl_fp) : tensor<?x?xf32>
19+
%C_alloc = bufferization.alloc_tensor(%c2, %vl_fp) : tensor<?x?xf32>
20+
21+
%pi = arith.constant 3.14 : f32
22+
%A = linalg.fill ins(%pi : f32) outs(%A_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>
23+
%B = linalg.fill ins(%pi : f32) outs(%B_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>
24+
%C_in = linalg.fill ins(%c0_f32 : f32) outs(%C_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>
25+
26+
%C_out = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>) outs(%C_in: tensor<?x?xf32>) -> tensor<?x?xf32>
27+
28+
// CHECK-LABEL: SVE: START OF TEST OUTPUT
29+
vector.print str "SVE: START OF TEST OUTPUT"
30+
31+
// There are at least 4 x f32 elements in every SVE vector, i.e.
32+
// * %vscale >= 1.
33+
// Hence, when checking the outupt there will always be at least 4 elements
34+
// in every row. For implementations with wider vectors, you should see more
35+
// elements being printed.
36+
// CHECK-NEXT: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [2, 16] strides = [16, 1] data =
37+
// CHECK-NEXT: [9.8596, 9.8596, 9.8596, 9.8596
38+
// CHECK-NEXT: [9.8596, 9.8596, 9.8596, 9.8596
39+
40+
%xf = tensor.cast %C_out : tensor<?x?xf32> to tensor<*xf32>
41+
call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()
42+
43+
// CHECK-NEXT: SVE: END OF TEST OUTPUT
44+
vector.print str "SVE: END OF TEST OUTPUT"
45+
46+
return
47+
}
48+
49+
transform.sequence failures(propagate) {
50+
^bb1(%module_op: !transform.any_op):
51+
%0 = transform.structured.match ops{["linalg.matmul"]} in %module_op : (!transform.any_op) -> !transform.any_op
52+
%func_op = get_parent_op %0 : (!transform.any_op) -> !transform.op<"func.func">
53+
// The tile sizes match the output matrix sizes
54+
%1, %loops:3 = transform.structured.tile_using_for %0 [2, [4], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
55+
%2 = transform.structured.match ops{["linalg.matmul"]} in %module_op : (!transform.any_op) -> !transform.any_op
56+
// The vector sizes match the output matrix sizes
57+
// TOOD: Use variables to re-use "shared" sizes
58+
transform.structured.vectorize %2 vector_sizes [2, [4], 1] : !transform.any_op
59+
60+
transform.apply_patterns to %func_op {
61+
transform.apply_patterns.vector.reduction_to_contract
62+
transform.apply_patterns.vector.transfer_permutation_patterns
63+
transform.apply_patterns.vector.lower_masked_transfers
64+
} : !transform.op<"func.func">
65+
transform.apply_patterns to %func_op {
66+
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
67+
transform.apply_patterns.vector.lower_outerproduct
68+
} : !transform.op<"func.func">
69+
}
70+
71+
func.func private @printMemrefF32(%ptr : tensor<*xf32>)

0 commit comments

Comments
 (0)