Skip to content

Commit 19119dd

Browse files
committed
[mlir][vector] Add integration test for vector distribute transformation
Differential Revision: https://reviews.llvm.org/D89062
1 parent cf402a1 commit 19119dd

File tree

1 file changed

+60
-0
lines changed

1 file changed

+60
-0
lines changed
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 \
2+
// RUN: -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm | \
3+
// RUN: mlir-cpu-runner -e main -entry-point-result=void \
4+
// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext | \
5+
// RUN: FileCheck %s
6+
7+
func @print_memref_f32(memref<*xf32>)
8+
9+
func @alloc_1d_filled_inc_f32(%arg0: index, %arg1: f32) -> memref<?xf32> {
10+
%c0 = constant 0 : index
11+
%c1 = constant 1 : index
12+
%0 = alloc(%arg0) : memref<?xf32>
13+
scf.for %arg2 = %c0 to %arg0 step %c1 {
14+
%tmp = index_cast %arg2 : index to i32
15+
%tmp1 = sitofp %tmp : i32 to f32
16+
%tmp2 = addf %tmp1, %arg1 : f32
17+
store %tmp2, %0[%arg2] : memref<?xf32>
18+
}
19+
return %0 : memref<?xf32>
20+
}
21+
22+
func @vector_add_cycle(%id : index, %A: memref<?xf32>, %B: memref<?xf32>, %C: memref<?xf32>) {
23+
%c0 = constant 0 : index
24+
%cf0 = constant 0.0 : f32
25+
%a = vector.transfer_read %A[%c0], %cf0: memref<?xf32>, vector<64xf32>
26+
%b = vector.transfer_read %B[%c0], %cf0: memref<?xf32>, vector<64xf32>
27+
%acc = addf %a, %b: vector<64xf32>
28+
vector.transfer_write %acc, %C[%c0]: vector<64xf32>, memref<?xf32>
29+
return
30+
}
31+
32+
// Loop over a function containinng a large add vector and distribute it so that
33+
// each iteration of the loop process part of the vector operation.
34+
func @main() {
35+
%cf1 = constant 1.0 : f32
36+
%cf2 = constant 2.0 : f32
37+
%c0 = constant 0 : index
38+
%c1 = constant 1 : index
39+
%c64 = constant 64 : index
40+
%out = alloc(%c64) : memref<?xf32>
41+
%in1 = call @alloc_1d_filled_inc_f32(%c64, %cf1) : (index, f32) -> memref<?xf32>
42+
%in2 = call @alloc_1d_filled_inc_f32(%c64, %cf2) : (index, f32) -> memref<?xf32>
43+
scf.for %arg5 = %c0 to %c64 step %c1 {
44+
call @vector_add_cycle(%arg5, %in1, %in2, %out) : (index, memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
45+
}
46+
%converted = memref_cast %out : memref<?xf32> to memref<*xf32>
47+
call @print_memref_f32(%converted): (memref<*xf32>) -> ()
48+
// CHECK: Unranked{{.*}}data =
49+
// CHECK: [
50+
// CHECK-SAME: 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27,
51+
// CHECK-SAME: 29, 31, 33, 35, 37, 39, 41, 43, 45, 47, 49, 51,
52+
// CHECK-SAME: 53, 55, 57, 59, 61, 63, 65, 67, 69, 71, 73, 75,
53+
// CHECK-SAME: 77, 79, 81, 83, 85, 87, 89, 91, 93, 95, 97, 99,
54+
// CHECK-SAME: 101, 103, 105, 107, 109, 111, 113, 115, 117, 119,
55+
// CHECK-SAME: 121, 123, 125, 127, 129]
56+
dealloc %out : memref<?xf32>
57+
dealloc %in1 : memref<?xf32>
58+
dealloc %in2 : memref<?xf32>
59+
return
60+
}

0 commit comments

Comments
 (0)