loop-tiling.mlir
7.25 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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
// RUN: mlir-opt %s -split-input-file -affine-loop-tile -tile-size=32 | FileCheck %s
// RUN: mlir-opt %s -split-input-file -affine-loop-tile -tile-cache-size=512 | FileCheck %s --check-prefix=MODEL
// -----
// CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32)>
// CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32, 50)>
// CHECK-DAG: [[IDENTITY:#map[0-9]+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @loop_tiling()
// CHECK-NEXT: affine.for %{{.*}} = 0 to 256 step 32 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 512 step 32 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 1024 step 32 {
// CHECK-NEXT: affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to [[MAP0]](%{{.*}}) {
// CHECK-NEXT: affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to [[MAP0]](%{{.*}}) {
// CHECK-NEXT: affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to [[MAP0]](%{{.*}}) {
// CHECK-NEXT: "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> ()
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 50 step 32 {
// CHECK-NEXT: affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to min [[MAP1]](%{{.*}}) {
// CHECK-NEXT: "bar"(%{{.*}}, %{{.*}}) : (index, index) -> ()
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 21 step 32 {
// CHECK-NEXT: affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to 21 {
// CHECK-NEXT: "foobar"(%{{.*}}) : (index) -> ()
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: return
func @loop_tiling() {
affine.for %i = 0 to 256 {
affine.for %j = 0 to 512 {
affine.for %k = 0 to 1024 {
"foo"(%i, %j, %k) : (index, index, index) -> ()
}
}
}
affine.for %x = 0 to 50 {
"bar"(%x, %x) : (index, index) -> ()
}
// Intra-tile loop won't need a min expression.
affine.for %y = 0 to 21 {
"foobar"(%y) : (index) -> ()
}
return
}
// -----
// CHECK-DAG: [[IDENTITY:#map[0-9]+]] = affine_map<(d0) -> (d0)>
// CHECK-DAG: [[LB:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
// CHECK-DAG: [[UB:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
// CHECK-DAG: [[UB_INTRA_TILE:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0, 4096 floordiv s1)>
#lb = affine_map<()[s0] -> (0, s0)>
#ub = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
// CHECK-LABEL: func @loop_max_min_bound(%{{.*}}: memref<?xi32>, %{{.*}}: index, %{{.*}}: index) {
func @loop_max_min_bound(%A : memref<? x i32>, %L : index, %U : index) {
%M = dim %A, 0 : memref<? x i32>
affine.for %iTT = max #lb()[%L] to min #ub()[%M, %U] {
%out = affine.apply affine_map<(d0) -> (d0)> (%iTT)
}
return
// CHECK: affine.for %{{.*}} = max [[LB]]()[%{{.*}}] to min [[UB]]()[%{{.*}}, %{{.*}}] step 32 {
// CHECK-NEXT: affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to min [[UB_INTRA_TILE]](%{{.*}})[%{{.*}}, %{{.*}}] {
// CHECK-NEXT: %{{.*}} = affine.apply [[IDENTITY]](%{{.*}})
// CHECK-NEXT: }
// CHECK-NEXT: }
}
// -----
// Cache size is set to 512 KiB. This loop nest accesses about 49 MiB, and the
// tile sizes chosen would be 6 x 6 x 6. However, to avoid min/max, which is
// possible here, they are adjusted to 4 x 4 x 5.
// MODEL-LABEL: func @simple_matmul
func @simple_matmul(%arg0: memref<256x256xvector<64xf32>>, %arg1: memref<256x256xvector<64xf32>>, %arg2: memref<256x256xvector<64xf32>>) -> memref<256x256xvector<64xf32>> {
affine.for %i = 0 to 256 {
affine.for %j = 0 to 256 {
affine.for %k = 0 to 250 {
%l = affine.load %arg0[%i, %k] : memref<256x256xvector<64xf32>>
%r = affine.load %arg1[%k, %j] : memref<256x256xvector<64xf32>>
%o = affine.load %arg2[%i, %j] : memref<256x256xvector<64xf32>>
%m = mulf %l, %r : vector<64xf32>
%a = addf %o, %m : vector<64xf32>
affine.store %a, %arg2[%i, %j] : memref<256x256xvector<64xf32>>
}
}
}
return %arg2 : memref<256x256xvector<64xf32>>
}
// MODEL: affine.for %{{.*}} = 0 to 256 step 4 {
// MODEL-NEXT: affine.for %{{.*}} = 0 to 256 step 4 {
// MODEL-NEXT: affine.for %{{.*}} = 0 to 250 step 5 {
// -----
// CHECK-DAG: [[UBMAP:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>
func @tile_with_symbolic_loop_upper_bounds(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
%cst = constant 0.000000e+00 : f32
%0 = dim %arg0, 0 : memref<?x?xf32>
affine.for %i0 = 0 to %0 {
affine.for %i1 = 0 to %0 {
affine.store %cst, %arg2[%i0, %i1] : memref<?x?xf32>
affine.for %i2 = 0 to %0 {
%1 = affine.load %arg0[%i0, %i2] : memref<?x?xf32>
%2 = affine.load %arg1[%i2, %i1] : memref<?x?xf32>
%3 = mulf %1, %2 : f32
%4 = affine.load %arg2[%i0, %i1] : memref<?x?xf32>
%5 = addf %4, %3 : f32
affine.store %5, %arg2[%i0, %i1] : memref<?x?xf32>
}
}
}
return
}
// CHECK: %{{.*}} = dim %{{.*}}, 0 : memref<?x?xf32>
// CHECK-NEXT: affine.for %{{.*}} = 0 to %{{.*}} step 32 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to %{{.*}} step 32 {
// CHECK-NEXT: affine.for %{{.*}} = #map3(%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}] {
// CHECK-NEXT: affine.for %{{.*}} = #map3(%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}] {
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: affine.for %{{.*}} = 0 to %{{.*}} {
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: %{{.*}} = mulf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: %{{.*}} = addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: return
// -----
// CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0)>
// CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 + s1)>
// CHECK-DAG: [[UBMAP:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0 + s1)>
func @tile_with_loop_upper_bounds_in_two_symbols(%arg0: memref<?xf32>, %limit: index) {
%dim0 = dim %arg0, 0 : memref<?xf32>
affine.for %i0 = 0 to affine_map<()[s0, s1] -> (s0 + s1)> ()[%dim0, %limit] {
%v0 = affine.load %arg0[%i0] : memref<?xf32>
}
return
}
// CHECK: %{{.*}} = dim %{{.*}}, 0 : memref<?xf32>
// CHECK-NEXT: affine.for %{{.*}} = 0 to [[MAP1]]()[%{{.*}}, %{{.*}}] step 32 {
// CHECK-NEXT: affine.for %{{.*}} = [[MAP0]](%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}, %{{.*}}] {
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}] : memref<?xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
// -----
func @trip_count_1(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>)
-> memref<196608x1xf32> {
affine.for %i1 = 0 to 196608 {
affine.for %i3 = 0 to 1 {
%4 = affine.load %arg0[%i1, %i3] : memref<196608x1xf32>
affine.store %4, %arg1[%i1, %i3] : memref<196608x1xf32>
}
}
return %arg1 : memref<196608x1xf32>
}
// CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<196608x1xf32>