uniform_divergent.mlir
1.93 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" -split-input-file | FileCheck %s
// Specific tests to check vectorization of uniform/divergent values.
// CHECK-LABEL: @uniform_arg
// CHECK-SAME: %[[in:.*]]: memref<512xf32>,
// CHECK-SAME: %[[uniform:.*]]: f32
func @uniform_arg(%in : memref<512xf32>, %uniform : f32) {
affine.for %i = 0 to 512 {
%ld = affine.load %in[%i] : memref<512xf32>
%add = addf %ld, %uniform : f32
}
return
}
// CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform]] : f32 to vector<128xf32>
// CHECK-NEXT: affine.for
// CHECK: addf %{{.*}}, %[[bcast]] : vector<128xf32>
// -----
// CHECK-LABEL: @multi_use_uniform_arg
// CHECK-SAME: %[[in:.*]]: memref<512xf32>
// CHECK-SAME: %[[uniform:.*]]: f32
func @multi_use_uniform_arg(%in : memref<512xf32>, %uniform : f32) {
affine.for %i = 0 to 512 {
%ld = affine.load %in[%i] : memref<512xf32>
%user0 = addf %ld, %uniform : f32
%user1 = addf %ld, %uniform : f32
}
return
}
// CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform]] : f32 to vector<128xf32>
// CHECK-NOT: vector.broadcast
// CHECK-NEXT: affine.for
// CHECK: addf %{{.*}}, %[[bcast]] : vector<128xf32>
// CHECK: addf %{{.*}}, %[[bcast]] : vector<128xf32>
// -----
// CHECK-LABEL: @uniform_load
func @uniform_load(%A : memref<?x?xf32>, %C : memref<?x?xf32>) {
%c0 = constant 0 : index
%N = dim %A, %c0 : memref<?x?xf32>
affine.for %i = 0 to %N {
%uniform_ld = affine.load %A[%i, %i] : memref<?x?xf32>
affine.for %j = 0 to %N {
%b = affine.load %A[%i, %j] : memref<?x?xf32>
%c = addf %uniform_ld, %b : f32
}
}
return
}
// CHECK: affine.for
// CHECK-NEXT: %[[uniform_ld:.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform_ld]] : f32 to vector<128xf32>
// CHECK-NEXT: affine.for
// CHECK: addf %[[bcast]], %{{.*}} : vector<128xf32>