1// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -pass-pipeline='func(canonicalize)' | FileCheck %s
2
3// -----
4
5// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 - 1)>
6// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0 + 1)>
7
8// CHECK-LABEL: func @compose_affine_maps_1dto2d_no_symbols() {
9func @compose_affine_maps_1dto2d_no_symbols() {
10  %0 = memref.alloc() : memref<4x4xf32>
11
12  affine.for %i0 = 0 to 15 {
13    // Test load[%x, %x]
14
15    %x0 = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
16    %x1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %x0)
17    %x1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %x0)
18
19    // CHECK: %[[I0A:.*]] = affine.apply #[[$MAP0]](%{{.*}})
20    // CHECK-NEXT: %[[V0:.*]] = memref.load %0[%[[I0A]], %[[I0A]]]
21    %v0 = memref.load %0[%x1_0, %x1_1] : memref<4x4xf32>
22
23    // Test store[%y, %y]
24    %y0 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
25    %y1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %y0)
26    %y1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %y0)
27
28    // CHECK-NEXT: %[[I1A:.*]] = affine.apply #[[$MAP1]](%{{.*}})
29    // CHECK-NEXT: memref.store %[[V0]], %0[%[[I1A]], %[[I1A]]]
30    memref.store %v0, %0[%y1_0, %y1_1] : memref<4x4xf32>
31
32    // Test store[%x, %y]
33    %xy_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %y0)
34    %xy_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %y0)
35
36    // CHECK-NEXT: memref.store %[[V0]], %0[%[[I0A]], %[[I1A]]]
37    memref.store %v0, %0[%xy_0, %xy_1] : memref<4x4xf32>
38
39    // Test store[%y, %x]
40    %yx_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %x0)
41    %yx_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %x0)
42    // CHECK-NEXT: memref.store %[[V0]], %0[%[[I1A]], %[[I0A]]]
43    memref.store %v0, %0[%yx_0, %yx_1] : memref<4x4xf32>
44  }
45  return
46}
47
48// -----
49
50// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> (d0 - 4)>
51// CHECK-DAG: #[[$MAP7:.*]] = affine_map<(d0) -> (d0 * 2 - 3)>
52// CHECK-DAG: #[[$MAP7a:.*]] = affine_map<(d0) -> (d0 * 2 + 1)>
53
54// CHECK-LABEL: func @compose_affine_maps_1dto2d_with_symbols() {
55func @compose_affine_maps_1dto2d_with_symbols() {
56  %0 = memref.alloc() : memref<4x4xf32>
57
58  affine.for %i0 = 0 to 15 {
59    // Test load[%x0, %x0] with symbol %c4
60    %c4 = constant 4 : index
61    %x0 = affine.apply affine_map<(d0)[s0] -> (d0 - s0)> (%i0)[%c4]
62
63    // CHECK: %[[I0:.*]] = affine.apply #[[$MAP4]](%{{.*}})
64    // CHECK-NEXT: %[[V0:.*]] = memref.load %{{.*}}[%[[I0]], %[[I0]]]
65    %v0 = memref.load %0[%x0, %x0] : memref<4x4xf32>
66
67    // Test load[%x0, %x1] with symbol %c4 captured by '%x0' map.
68    %x1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
69    %y1 = affine.apply affine_map<(d0, d1) -> (d0+d1)> (%x0, %x1)
70    // CHECK-NEXT: %[[I1:.*]] = affine.apply #[[$MAP7]](%{{.*}})
71    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I1]], %[[I1]]]
72    memref.store %v0, %0[%y1, %y1] : memref<4x4xf32>
73
74    // Test store[%x1, %x0] with symbol %c4 captured by '%x0' map.
75    %y2 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x1, %x0)
76    // CHECK-NEXT: %[[I2:.*]] = affine.apply #[[$MAP7]](%{{.*}})
77    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I2]], %[[I2]]]
78    memref.store %v0, %0[%y2, %y2] : memref<4x4xf32>
79
80    // Test store[%x2, %x0] with symbol %c4 from '%x0' and %c5 from '%x2'
81    %c5 = constant 5 : index
82    %x2 = affine.apply affine_map<(d0)[s0] -> (d0 + s0)> (%i0)[%c5]
83    %y3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x2, %x0)
84    // CHECK: %[[I3:.*]] = affine.apply #[[$MAP7a]](%{{.*}})
85    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I3]], %[[I3]]]
86    memref.store %v0, %0[%y3, %y3] : memref<4x4xf32>
87  }
88  return
89}
90
91// -----
92
93// CHECK-DAG: #[[$MAP8:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
94// CHECK-DAG: #[[$MAP8a:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
95
96// CHECK-LABEL: func @compose_affine_maps_2d_tile
97func @compose_affine_maps_2d_tile(%0: memref<16x32xf32>, %1: memref<16x32xf32>) {
98  %c4 = constant 4 : index
99  %c8 = constant 8 : index
100
101  affine.for %i0 = 0 to 3 {
102    %x0 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i0)[%c4]
103    affine.for %i1 = 0 to 3 {
104      %x1 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i1)[%c8]
105      affine.for %i2 = 0 to 3 {
106        %x2 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i2)[%c4]
107        affine.for %i3 = 0 to 3 {
108          %x3 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i3)[%c8]
109
110          %x40 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
111            ((d0 * s0) + d2)> (%x0, %x1, %x2, %x3)[%c4, %c8]
112          %x41 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
113            ((d1 * s1) + d3)> (%x0, %x1, %x2, %x3)[%c4, %c8]
114          // CHECK: %[[I0:.*]] = affine.apply #[[$MAP8]](%{{.*}}, %{{.*}})
115          // CHECK: %[[I1:.*]] = affine.apply #[[$MAP8a]](%{{.*}}, %{{.*}})
116          // CHECK-NEXT: %[[L0:.*]] = memref.load %{{.*}}[%[[I0]], %[[I1]]]
117          %v0 = memref.load %0[%x40, %x41] : memref<16x32xf32>
118
119          // CHECK-NEXT: memref.store %[[L0]], %{{.*}}[%[[I0]], %[[I1]]]
120          memref.store %v0, %1[%x40, %x41] : memref<16x32xf32>
121        }
122      }
123    }
124  }
125  return
126}
127
128// -----
129
130// CHECK-DAG: #[[$MAP4b:.*]] = affine_map<(d0) -> (d0 - 7)>
131// CHECK-DAG: #[[$MAP9:.*]] = affine_map<(d0) -> (d0 + 3)>
132// CHECK-DAG: #[[$MAP10:.*]] = affine_map<(d0) -> (d0 * 3)>
133// CHECK-DAG: #[[$MAP11:.*]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)>
134// CHECK-DAG: #[[$MAP12:.*]] = affine_map<(d0) -> (d0 * 7 - 49)>
135
136// CHECK-LABEL: func @compose_affine_maps_dependent_loads() {
137func @compose_affine_maps_dependent_loads() {
138  %0 = memref.alloc() : memref<16x32xf32>
139  %1 = memref.alloc() : memref<16x32xf32>
140
141  affine.for %i0 = 0 to 3 {
142    affine.for %i1 = 0 to 3 {
143      affine.for %i2 = 0 to 3 {
144        %c3 = constant 3 : index
145        %c7 = constant 7 : index
146
147        %x00 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + s0)>
148            (%i0, %i1, %i2)[%c3, %c7]
149        %x01 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d1 - s1)>
150            (%i0, %i1, %i2)[%c3, %c7]
151        %x02 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d2 * s0)>
152            (%i0, %i1, %i2)[%c3, %c7]
153
154        // CHECK: %[[I0:.*]] = affine.apply #[[$MAP9]](%{{.*}})
155        // CHECK: %[[I1:.*]] = affine.apply #[[$MAP4b]](%{{.*}})
156        // CHECK: %[[I2:.*]] = affine.apply #[[$MAP10]](%{{.*}})
157        // CHECK-NEXT: %[[V0:.*]] = memref.load %{{.*}}[%[[I0]], %[[I1]]]
158        %v0 = memref.load %0[%x00, %x01] : memref<16x32xf32>
159
160        // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I0]], %[[I2]]]
161        memref.store %v0, %0[%x00, %x02] : memref<16x32xf32>
162
163        // Swizzle %i0, %i1
164        // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I1]], %[[I0]]]
165        memref.store %v0, %0[%x01, %x00] : memref<16x32xf32>
166
167        // Swizzle %x00, %x01 and %c3, %c7
168        %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)>
169           (%x01, %x00)[%c3, %c7]
170        %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)>
171           (%x01, %x00)[%c3, %c7]
172
173        // CHECK-NEXT: %[[I2A:.*]] = affine.apply #[[$MAP12]](%{{.*}})
174        // CHECK-NEXT: %[[I2B:.*]] = affine.apply #[[$MAP11]](%{{.*}})
175        // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I2A]], %[[I2B]]]
176        memref.store %v0, %0[%x10, %x11] : memref<16x32xf32>
177      }
178    }
179  }
180  return
181}
182
183// -----
184
185// CHECK-DAG: #[[$MAP13A:.*]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
186// CHECK-DAG: #[[$MAP13B:.*]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
187
188// CHECK-LABEL: func @compose_affine_maps_diamond_dependency
189func @compose_affine_maps_diamond_dependency(%arg0: f32, %arg1: memref<4x4xf32>) {
190  affine.for %i0 = 0 to 15 {
191    %a = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
192    %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
193    %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
194    %d0 = affine.apply affine_map<(d0, d1) -> (d0 ceildiv 8)> (%b, %c)
195    %d1 = affine.apply affine_map<(d0, d1) -> (d1 floordiv 3)> (%b, %c)
196    // CHECK: %[[I0:.*]] = affine.apply #[[$MAP13A]](%{{.*}})
197    // CHECK: %[[I1:.*]] = affine.apply #[[$MAP13B]](%{{.*}})
198    // CHECK-NEXT: memref.store %arg0, %arg1[%[[I0]], %[[I1]]]
199    memref.store %arg0, %arg1[%d0, %d1] : memref<4x4xf32>
200  }
201
202  return
203}
204
205// -----
206
207// CHECK-DAG: #[[$MAP14:.*]] = affine_map<()[s0, s1] -> ((s0 * 4 + s1 * 4) floordiv s0)>
208
209// CHECK-LABEL: func @compose_affine_maps_multiple_symbols
210func @compose_affine_maps_multiple_symbols(%arg0: index, %arg1: index) -> index {
211  %a = affine.apply affine_map<(d0)[s0] -> (s0 + d0)> (%arg0)[%arg1]
212  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
213  %e = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%c)[%arg1]
214  // CHECK: [[I0:.*]] = affine.apply #[[$MAP14]]()[%{{.*}}, %{{.*}}]
215  return %e : index
216}
217
218// -----
219
220// CHECK-LABEL: func @arg_used_as_dim_and_symbol
221func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) -> (memref<100x100xf32, 1>, memref<1xi32>) {
222  %c9 = constant 9 : index
223  %1 = memref.alloc() : memref<100x100xf32, 1>
224  %2 = memref.alloc() : memref<1xi32>
225  affine.for %i0 = 0 to 100 {
226    affine.for %i1 = 0 to 100 {
227      %3 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 + s0 + s1)>
228        (%i0, %i1)[%arg1, %c9]
229      %4 = affine.apply affine_map<(d0, d1, d3) -> (d3 - (d0 + d1))>
230        (%arg1, %c9, %3)
231      // CHECK: memref.store %arg2, %{{.*}}[%{{.*}}, %{{.*}}]
232      memref.store %arg2, %1[%4, %arg1] : memref<100x100xf32, 1>
233    }
234  }
235  return %1, %2 : memref<100x100xf32, 1>, memref<1xi32>
236}
237
238// -----
239
240// CHECK-LABEL: func @trivial_maps
241func @trivial_maps() {
242  // CHECK-NOT: affine.apply
243
244  %0 = memref.alloc() : memref<10xf32>
245  %c0 = constant 0 : index
246  %cst = constant 0.000000e+00 : f32
247  affine.for %i1 = 0 to 10 {
248    %1 = affine.apply affine_map<()[s0] -> (s0)>()[%c0]
249    memref.store %cst, %0[%1] : memref<10xf32>
250    %2 = memref.load %0[%c0] : memref<10xf32>
251
252    %3 = affine.apply affine_map<()[] -> (0)>()[]
253    memref.store %cst, %0[%3] : memref<10xf32>
254    memref.store %2, %0[%c0] : memref<10xf32>
255  }
256  return
257}
258
259// -----
260
261// CHECK-DAG: #[[$MAP15:.*]] = affine_map<()[s0] -> (s0 - 42)>
262
263// CHECK-LABEL: func @partial_fold_map
264func @partial_fold_map(%arg1: index, %arg2: index) -> index {
265  // TODO: Constant fold one index into affine.apply
266  %c42 = constant 42 : index
267  %2 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%arg1, %c42)
268  // CHECK: [[X:.*]] = affine.apply #[[$MAP15]]()[%{{.*}}]
269  return %2 : index
270}
271
272// -----
273
274// CHECK-DAG: #[[$MAP_symbolic_composition_a:.*]] = affine_map<()[s0] -> (s0 * 512)>
275
276// CHECK-LABEL: func @symbolic_composition_a(%{{.*}}: index, %{{.*}}: index) -> index {
277func @symbolic_composition_a(%arg0: index, %arg1: index) -> index {
278  %0 = affine.apply affine_map<(d0) -> (d0 * 4)>(%arg0)
279  %1 = affine.apply affine_map<()[s0, s1] -> (8 * s0)>()[%0, %arg0]
280  %2 = affine.apply affine_map<()[s0, s1] -> (16 * s1)>()[%arg1, %1]
281  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_a]]()[%{{.*}}]
282  return %2 : index
283}
284
285// -----
286
287// CHECK-DAG: #[[$MAP_symbolic_composition_b:.*]] = affine_map<()[s0] -> (s0 * 4)>
288
289// CHECK-LABEL: func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
290func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
291  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
292  %1 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %0]
293  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
294  return %1 : index
295}
296
297// -----
298
299// CHECK-DAG: #[[$MAP_symbolic_composition_c:.*]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
300
301// CHECK-LABEL: func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
302func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
303  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
304  %1 = affine.apply affine_map<(d0) -> (d0)>(%arg1)
305  %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
306  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_c]]()[%{{.*}}, %{{.*}}]
307  return %2 : index
308}
309
310// -----
311
312// CHECK-DAG: #[[$MAP_symbolic_composition_d:.*]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
313
314// CHECK-LABEL: func @symbolic_composition_d(
315//  CHECK-SAME:   %[[ARG0:[0-9a-zA-Z]+]]: index
316//  CHECK-SAME:   %[[ARG1:[0-9a-zA-Z]+]]: index
317func @symbolic_composition_d(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
318  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
319  %1 = affine.apply affine_map<()[s0] -> (s0)>()[%arg1]
320  %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
321  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_d]]()[%[[ARG0]], %[[ARG1]]]
322  return %2 : index
323}
324
325// -----
326
327// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_b:.*]] = affine_map<()[s0, s1] -> (s0 * 42 + s1 + 6)>
328
329// CHECK-LABEL: func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
330func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
331  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
332  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
333  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_b]]()[%{{.*}}, %{{.*}}]
334
335  return %b : index
336}
337
338// -----
339
340// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_c:.*]] = affine_map<()[s0, s1] -> (s0 * 168 + s1 * 4 - 4)>
341
342// CHECK-LABEL: func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
343func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
344  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
345  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
346  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
347  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_c]]()[%{{.*}}, %{{.*}}]
348  return %c : index
349}
350
351// -----
352
353// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_d:.*]] = affine_map<()[s0, s1] -> ((s0 * 42 + s1 + 6) ceildiv 8)>
354
355// CHECK-LABEL: func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
356func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
357  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
358  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
359  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
360  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
361  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_d]]()[%{{.*}}, %{{.*}}]
362  return %d : index
363}
364
365// -----
366
367// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_e:.*]] = affine_map<()[s0, s1] -> ((s0 * 168 + s1 * 4 - 4) floordiv 3)>
368
369// CHECK-LABEL: func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
370func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
371  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
372  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
373  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
374  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
375  %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
376  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_e]]()[%{{.*}}, %{{.*}}]
377  return %e : index
378}
379
380// -----
381
382// CHECK-LABEL: func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
383func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
384  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
385  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
386  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
387  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
388  %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
389  %f = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 - s1 +  d1 - s0)> (%d, %e)[%e, %d]
390  // CHECK: {{.*}} = constant 0 : index
391
392  return %f : index
393}
394
395// -----
396
397// CHECK-DAG: #[[$MAP_symbolic_composition_b:.*]] = affine_map<()[s0] -> (s0 * 4)>
398
399// CHECK-LABEL: func @mix_dims_and_symbols_g(%arg0: index, %arg1: index) -> (index, index, index) {
400func @mix_dims_and_symbols_g(%M: index, %N: index) -> (index, index, index) {
401  %K = affine.apply affine_map<(d0) -> (4*d0)> (%M)
402  %res1 = affine.apply affine_map<()[s0, s1] -> (4 * s0)>()[%N, %K]
403  %res2 = affine.apply affine_map<()[s0, s1] -> (s1)>()[%N, %K]
404  %res3 = affine.apply affine_map<()[s0, s1] -> (1024)>()[%N, %K]
405  // CHECK-DAG: {{.*}} = constant 1024 : index
406  // CHECK-DAG: {{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
407  // CHECK-DAG: {{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
408  return %res1, %res2, %res3 : index, index, index
409}
410
411// -----
412
413// CHECK-DAG: #[[$symbolic_semi_affine:.*]] = affine_map<(d0)[s0] -> (d0 floordiv (s0 + 1))>
414
415// CHECK-LABEL: func @symbolic_semi_affine(%arg0: index, %arg1: index, %arg2: memref<?xf32>) {
416func @symbolic_semi_affine(%M: index, %N: index, %A: memref<?xf32>) {
417  %f1 = constant 1.0 : f32
418  affine.for %i0 = 1 to 100 {
419    %1 = affine.apply affine_map<()[s0] -> (s0 + 1)> ()[%M]
420    %2 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%i0)[%1]
421    // CHECK-DAG: {{.*}} = affine.apply #[[$symbolic_semi_affine]](%{{.*}})[%{{.*}}]
422    memref.store %f1, %A[%2] : memref<?xf32>
423  }
424  return
425}
426
427// -----
428
429// CHECK: #[[$MAP0:.*]] = affine_map<()[s0] -> (0, s0)>
430// CHECK: #[[$MAP1:.*]] = affine_map<()[s0] -> (100, s0)>
431
432// CHECK-LABEL:  func @constant_fold_bounds(%arg0: index) {
433func @constant_fold_bounds(%N : index) {
434  // CHECK:      constant 3 : index
435  // CHECK-NEXT: "foo"() : () -> index
436  %c9 = constant 9 : index
437  %c1 = constant 1 : index
438  %c2 = constant 2 : index
439  %c3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%c1, %c2)
440  %l = "foo"() : () -> index
441
442  // CHECK:  affine.for %{{.*}} = 5 to 7 {
443  affine.for %i = max affine_map<(d0, d1) -> (0, d0 + d1)> (%c2, %c3) to min affine_map<(d0, d1) -> (d0 - 2, 32*d1)> (%c9, %c1) {
444    "foo"(%i, %c3) : (index, index) -> ()
445  }
446
447  // Bound takes a non-constant argument but can still be folded.
448  // CHECK:  affine.for %{{.*}} = 1 to 7 {
449  affine.for %j = max affine_map<(d0) -> (0, 1)> (%N) to min affine_map<(d0, d1) -> (7, 9)> (%N, %l) {
450    "foo"(%j, %c3) : (index, index) -> ()
451  }
452
453  // None of the bounds can be folded.
454  // CHECK: affine.for %{{.*}} = max #[[$MAP0]]()[%{{.*}}] to min #[[$MAP1]]()[%{{.*}}] {
455  affine.for %k = max affine_map<()[s0] -> (0, s0)> ()[%l] to min affine_map<()[s0] -> (100, s0)> ()[%N] {
456    "foo"(%k, %c3) : (index, index) -> ()
457  }
458  return
459}
460
461// -----
462
463// CHECK-LABEL:  func @fold_empty_loops()
464func @fold_empty_loops() -> index {
465  %c0 = constant 0 : index
466  affine.for %i = 0 to 10 {
467  }
468  %res = affine.for %i = 0 to 10 iter_args(%arg = %c0) -> index {
469    affine.yield %arg : index
470  }
471  // CHECK-NEXT: %[[zero:.*]] = constant 0
472  // CHECK-NEXT: return %[[zero]]
473  return %res : index
474}
475
476// -----
477
478// CHECK-LABEL:  func @fold_zero_iter_loops
479// CHECK-SAME: %[[ARG:.*]]: index
480func @fold_zero_iter_loops(%in : index) -> index {
481  %c1 = constant 1 : index
482  affine.for %i = 0 to 0 {
483    affine.for %j = 0 to -1 {
484    }
485  }
486  %res = affine.for %i = 0 to 0 iter_args(%loop_arg = %in) -> index {
487    %yield = addi %loop_arg, %c1 : index
488    affine.yield %yield : index
489  }
490  // CHECK-NEXT: return %[[ARG]]
491  return %res : index
492}
493
494// -----
495
496// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
497
498// CHECK-LABEL: func @canonicalize_affine_if
499//  CHECK-SAME:   %[[M:[0-9a-zA-Z]*]]: index,
500//  CHECK-SAME:   %[[N:[0-9a-zA-Z]*]]: index)
501func @canonicalize_affine_if(%M : index, %N : index) {
502  %c1022 = constant 1022 : index
503  // Drop unused operand %M, propagate %c1022, and promote %N to symbolic.
504  affine.for %i = 0 to 1024 {
505    affine.for %j = 0 to %N {
506      // CHECK: affine.if #[[$SET]](%{{.*}}, %{{.*}})[%[[N]]]
507      affine.if affine_set<(d0, d1, d2, d3)[s0] : (d1 >= 0, d0 - d1 >= 0, d2 >= 0, d3 - d2 - 2 >= 0)>
508          (%c1022, %i, %j, %N)[%M] {
509        "foo"() : () -> ()
510      }
511      "bar"() : () -> ()
512    }
513  }
514  return
515}
516
517// -----
518
519// CHECK-DAG: #[[$LBMAP:.*]] = affine_map<()[s0] -> (0, s0)>
520// CHECK-DAG: #[[$UBMAP:.*]] = affine_map<()[s0] -> (1024, s0 * 2)>
521
522// CHECK-LABEL: func @canonicalize_bounds
523// CHECK-SAME: %[[M:.*]]: index,
524// CHECK-SAME: %[[N:.*]]: index)
525func @canonicalize_bounds(%M : index, %N : index) {
526  %c0 = constant 0 : index
527  %c1024 = constant 1024 : index
528  // Drop unused operand %N, drop duplicate operand %M, propagate %c1024, and
529  // promote %M to a symbolic one.
530  // CHECK: affine.for %{{.*}} = 0 to min #[[$UBMAP]]()[%[[M]]]
531  affine.for %i = 0 to min affine_map<(d0, d1, d2, d3) -> (d0, d1 + d2)> (%c1024, %M, %M, %N) {
532    "foo"() : () -> ()
533  }
534  // Promote %M to symbolic position.
535  // CHECK: affine.for %{{.*}} = 0 to #{{.*}}()[%[[M]]]
536  affine.for %i = 0 to affine_map<(d0) -> (4 * d0)> (%M) {
537    "foo"() : () -> ()
538  }
539  // Lower bound canonicalize.
540  // CHECK: affine.for %{{.*}} = max #[[$LBMAP]]()[%[[N]]] to %[[M]]
541  affine.for %i = max affine_map<(d0, d1) -> (d0, d1)> (%c0, %N) to %M {
542    "foo"() : () -> ()
543  }
544  return
545}
546
547// -----
548
549// Compose maps into affine load and store ops.
550
551// CHECK-LABEL: @compose_into_affine_load_store
552func @compose_into_affine_load_store(%A : memref<1024xf32>, %u : index) {
553  // CHECK: affine.for %[[IV:.*]] = 0 to 1024
554  affine.for %i = 0 to 1024 {
555    // Make sure the unused operand (%u below) gets dropped as well.
556    %idx = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i, %u)
557    %0 = affine.load %A[%idx] : memref<1024xf32>
558    affine.store %0, %A[%idx] : memref<1024xf32>
559    // CHECK-NEXT: affine.load %{{.*}}[%[[IV]] + 1]
560    // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%[[IV]] + 1]
561
562    // Map remains the same, but operand changes on composition.
563    %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
564    %1 = affine.load %A[%copy] : memref<1024xf32>
565    "prevent.dce"(%1) : (f32) -> ()
566    // CHECK-NEXT: affine.load %{{.*}}[%[[IV]]]
567  }
568  return
569}
570
571// -----
572
573func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
574  %c511 = constant 511 : index
575  %c1 = constant 0 : index
576  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0 + 1)> (%c1)[%c511]
577  "op0"(%0) : (index) -> ()
578  // CHECK:       %[[CST:.*]] = constant 512 : index
579  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
580  // CHECK-NEXT:  return
581  return
582}
583
584// -----
585
586func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
587  %c3 = constant 3 : index
588  %c20 = constant 20 : index
589  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
590  "op0"(%0) : (index) -> ()
591  // CHECK:       %[[CST:.*]] = constant 4 : index
592  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
593  // CHECK-NEXT:  return
594  return
595}
596
597// -----
598
599func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
600  %c511 = constant 511 : index
601  %c1 = constant 0 : index
602  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0 + 1)> (%c1)[%c511]
603  "op0"(%0) : (index) -> ()
604  // CHECK:       %[[CST:.*]] = constant 1000 : index
605  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
606  // CHECK-NEXT:  return
607  return
608}
609
610// -----
611
612func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
613  %c3 = constant 3 : index
614  %c20 = constant 20 : index
615  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
616  "op0"(%0) : (index) -> ()
617  // CHECK:       %[[CST:.*]] = constant 1000 : index
618  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
619  // CHECK-NEXT:  return
620  return
621}
622
623// -----
624
625// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d0, d1 - 2)>
626
627func @affine_min(%arg0: index) {
628  affine.for %i = 0 to %arg0 {
629    affine.for %j = 0 to %arg0 {
630      %c2 = constant 2 : index
631      // CHECK: affine.min #[[$MAP]]
632      %0 = affine.min affine_map<(d0,d1,d2)->(d0, d1 - d2)>(%i, %j, %c2)
633      "consumer"(%0) : (index) -> ()
634    }
635  }
636  return
637}
638
639// -----
640
641// Reproducer for PR45031. This used to fold into an incorrect map because
642// symbols were concatenated in the wrong order during map folding. Map
643// composition places the symbols of the original map before those of the map
644// it is composed with, e.g. A.compose(B) will first have all symbols of A,
645// then all symbols of B.
646
647#map1 = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
648#map2 = affine_map<(d0)[s0] -> (1024, -d0 + s0)>
649
650// CHECK: #[[$MAP:.*]] = affine_map<()[s0, s1] -> (1024, s0 - s1 * 1024)>
651
652// CHECK: func @rep(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
653func @rep(%arg0 : index, %arg1 : index) -> index {
654  // CHECK-NOT: constant
655  %c0 = constant 0 : index
656  %c1024 = constant 1024 : index
657  // CHECK-NOT: affine.apply
658  %0 = affine.apply #map1(%arg0)[%c1024, %c0]
659
660  // CHECK: affine.min #[[$MAP]]()[%[[ARG1]], %[[ARG0]]]
661  %1 = affine.min #map2(%0)[%arg1]
662  return %1 : index
663}
664
665// -----
666
667// CHECK-DAG: #[[ub:.*]] = affine_map<()[s0] -> (s0 + 2)>
668
669func @drop_duplicate_bounds(%N : index) {
670  // affine.for %i = max #lb(%arg0) to min #ub(%arg0)
671  affine.for %i = max affine_map<(d0) -> (d0, d0)>(%N) to min affine_map<(d0) -> (d0 + 2, d0 + 2)>(%N) {
672    "foo"() : () -> ()
673  }
674  return
675}
676
677// -----
678
679// Ensure affine.parallel bounds expressions are canonicalized.
680
681#map3 = affine_map<(d0) -> (d0 * 5)>
682
683// CHECK-LABEL: func @affine_parallel_const_bounds
684func @affine_parallel_const_bounds() {
685  %cst = constant 1.0 : f32
686  %c0 = constant 0 : index
687  %c4 = constant 4 : index
688  %0 = memref.alloc() : memref<4xf32>
689  // CHECK: affine.parallel (%{{.*}}) = (0) to (4)
690  affine.parallel (%i) = (%c0) to (%c0 + %c4) {
691    %1 = affine.apply #map3(%i)
692    // CHECK: affine.parallel (%{{.*}}) = (0) to (%{{.*}} * 5)
693    affine.parallel (%j) = (%c0) to (%1) {
694      affine.store %cst, %0[%j] : memref<4xf32>
695    }
696  }
697  return
698}
699
700// -----
701
702func @compose_affine_maps_div_symbol(%A : memref<i64>, %i0 : index, %i1 : index) {
703  %0 = affine.apply affine_map<()[s0] -> (2 * s0)> ()[%i0]
704  %1 = affine.apply affine_map<()[s0] -> (3 * s0)> ()[%i0]
705  %2 = affine.apply affine_map<(d0)[s0, s1] -> (d0 mod s1 + s0 * s1 + s0 * 4)> (%i1)[%0, %1]
706  %3 = index_cast %2: index to i64
707  memref.store %3, %A[]: memref<i64>
708  affine.for %i2 = 0 to 3 {
709    %4 = affine.apply affine_map<(d0)[s0, s1] -> (d0 ceildiv s1 + s0 + s0 * 3)> (%i2)[%0, %1]
710    %5 = index_cast %4: index to i64
711    memref.store %5, %A[]: memref<i64>
712  }
713  return
714}
715
716// -----
717
718// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
719
720// CHECK: func @deduplicate_affine_min_expressions
721// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
722func @deduplicate_affine_min_expressions(%i0: index, %i1: index) -> index {
723  // CHECK:  affine.min #[[MAP]]()[%[[I0]], %[[I1]]]
724  %0 = affine.min affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
725  return %0: index
726}
727
728// -----
729
730// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
731
732// CHECK: func @deduplicate_affine_max_expressions
733// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
734func @deduplicate_affine_max_expressions(%i0: index, %i1: index) -> index {
735  // CHECK:  affine.max #[[MAP]]()[%[[I0]], %[[I1]]]
736  %0 = affine.max affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
737  return %0: index
738}
739
740// -----
741
742// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1, s2] -> (s0 * 3, 16, -s1 + s2)>
743// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (-s2 + 5, 16, -s0 + s1)>
744
745// CHECK: func @merge_affine_min_ops
746// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index, %[[I3:.+]]: index)
747func @merge_affine_min_ops(%i0: index, %i1: index, %i2: index, %i3: index) -> (index, index) {
748  %0 = affine.min affine_map<(d0)[s0] -> (16, d0 - s0)> (%i0)[%i1]
749
750 // CHECK: affine.min #[[MAP0]]()[%[[I2]], %[[I1]], %[[I0]]]
751  %1 = affine.min affine_map<(d0)[s0] -> (3 * s0, d0)> (%0)[%i2] // Use as dim
752 // CHECK: affine.min #[[MAP1]]()[%[[I1]], %[[I0]], %[[I3]]]
753  %2 = affine.min affine_map<(d0)[s0] -> (s0, 5 - d0)> (%i3)[%0] // Use as symbol
754
755  return %1, %2: index, index
756}
757
758// -----
759
760// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + 7, s1 + 16, s1 * 8, s2 + 8, s2 * 4)>
761
762// CHECK: func @merge_multiple_affine_min_ops
763// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index)
764func @merge_multiple_affine_min_ops(%i0: index, %i1: index, %i2: index) -> index {
765  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
766  %1 = affine.min affine_map<()[s0] -> (s0 + 8, s0 * 4)> ()[%i1]
767  // CHECK: affine.min #[[MAP]]()[%[[I2]], %[[I0]], %[[I1]]]
768  %2 = affine.min affine_map<()[s0, s1, s2] -> (s0, 7 + s1, s2)> ()[%0, %i2, %1]
769  return %2: index
770}
771
772// -----
773
774// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 * 2, s1 + 16, s1 * 8)>
775
776// CHECK: func @merge_multiple_uses_of_affine_min_ops
777// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
778func @merge_multiple_uses_of_affine_min_ops(%i0: index, %i1: index) -> index {
779  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
780  // CHECK: affine.min #[[MAP]]()[%[[I1]], %[[I0]]]
781  %2 = affine.min affine_map<()[s0, s1, s2] -> (s0, s1, s2 * 2)> ()[%0, %0, %i1]
782  return %2: index
783}
784
785// -----
786
787// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 16, s0 * 8)>
788// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (s0 + 1, s1 * 2, s2 + 16, s2 * 8)>
789
790// CHECK: func @merge_mixed_uses_of_affine_min_ops
791// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
792func @merge_mixed_uses_of_affine_min_ops(%i0: index, %i1: index) -> index {
793  // CHECK: %[[AFFINE:.+]] = affine.min #[[MAP0]]()[%[[I0]]]
794  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
795  // %0 is bound to a symbol that is both a standalone expression and a part
796  // of other expressions.
797  // CHECK: affine.min #[[MAP1]]()[%[[AFFINE]], %[[I1]], %[[I0]]]
798  %2 = affine.min affine_map<()[s0, s1, s2] -> (s0, s1 + 1, s2 * 2)> ()[%0, %0, %i1]
799  return %2: index
800}
801
802// -----
803
804// CHECK-LABEL: func @dont_merge_affine_min_if_not_single_dim
805func @dont_merge_affine_min_if_not_single_dim(%i0: index, %i1: index, %i2: index) -> index {
806  // CHECK-COUNT-2: affine.min
807  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
808  %1 = affine.min affine_map<(d0)[s0] -> (s0 + 4, 7 + d0)> (%0)[%i2]
809  return %1: index
810}
811
812// -----
813
814// CHECK-LABEL: func @dont_merge_affine_min_if_not_single_sym
815func @dont_merge_affine_min_if_not_single_sym(%i0: index, %i1: index, %i2: index) -> index {
816  // CHECK-COUNT-2: affine.min
817  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
818  %1 = affine.min affine_map<()[s0, s1] -> (s0 + 4, 7 + s1)> ()[%0, %i2]
819  return %1: index
820}
821
822// -----
823
824// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1, s2] -> (s0 * 3, 16, -s1 + s2)>
825// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (-s2 + 5, 16, -s0 + s1)>
826
827// CHECK: func @merge_affine_max_ops
828// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index, %[[I3:.+]]: index)
829func @merge_affine_max_ops(%i0: index, %i1: index, %i2: index, %i3: index) -> (index, index) {
830  %0 = affine.max affine_map<(d0)[s0] -> (16, d0 - s0)> (%i0)[%i1]
831
832 // CHECK: affine.max #[[MAP0]]()[%[[I2]], %[[I1]], %[[I0]]]
833  %1 = affine.max affine_map<(d0)[s0] -> (3 * s0, d0)> (%0)[%i2] // Use as dim
834 // CHECK: affine.max #[[MAP1]]()[%[[I1]], %[[I0]], %[[I3]]]
835  %2 = affine.max affine_map<(d0)[s0] -> (s0, 5 - d0)> (%i3)[%0] // Use as symbol
836
837  return %1, %2: index, index
838}
839
840// -----
841
842// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + 7, s1 + 16, s1 * 8, s2 + 8, s2 * 4)>
843
844// CHECK: func @merge_multiple_affine_max_ops
845// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index)
846func @merge_multiple_affine_max_ops(%i0: index, %i1: index, %i2: index) -> index {
847  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
848  %1 = affine.max affine_map<()[s0] -> (s0 + 8, s0 * 4)> ()[%i1]
849  // CHECK: affine.max #[[MAP]]()[%[[I2]], %[[I0]], %[[I1]]]
850  %2 = affine.max affine_map<()[s0, s1, s2] -> (s0, 7 + s1, s2)> ()[%0, %i2, %1]
851  return %2: index
852}
853
854// -----
855
856// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 * 2, s1 + 16, s1 * 8)>
857
858// CHECK: func @merge_multiple_uses_of_affine_max_ops
859// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
860func @merge_multiple_uses_of_affine_max_ops(%i0: index, %i1: index) -> index {
861  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
862  // CHECK: affine.max #[[MAP]]()[%[[I1]], %[[I0]]]
863  %2 = affine.max affine_map<()[s0, s1, s2] -> (s0, s1, s2 * 2)> ()[%0, %0, %i1]
864  return %2: index
865}
866
867// -----
868
869// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 16, s0 * 8)>
870// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (s0 + 1, s1 * 2, s2 + 16, s2 * 8)>
871
872// CHECK: func @merge_mixed_uses_of_affine_max_ops
873// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
874func @merge_mixed_uses_of_affine_max_ops(%i0: index, %i1: index) -> index {
875  // CHECK: %[[AFFINE:.+]] = affine.max #[[MAP0]]()[%[[I0]]]
876  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
877  // %0 is bound to a symbol that is both a standalone expression and a part
878  // of other expressions.
879  // CHECK: affine.max #[[MAP1]]()[%[[AFFINE]], %[[I1]], %[[I0]]]
880  %2 = affine.max affine_map<()[s0, s1, s2] -> (s0, s1 + 1, s2 * 2)> ()[%0, %0, %i1]
881  return %2: index
882}
883
884// -----
885
886// CHECK-LABEL: func @dont_merge_affine_max_if_not_single_dim
887func @dont_merge_affine_max_if_not_single_dim(%i0: index, %i1: index, %i2: index) -> index {
888  // CHECK-COUNT-2: affine.max
889  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
890  %1 = affine.max affine_map<(d0)[s0] -> (s0 + 4, 7 + d0)> (%0)[%i2]
891  return %1: index
892}
893
894// -----
895
896// CHECK-LABEL: func @dont_merge_affine_max_if_not_single_sym
897func @dont_merge_affine_max_if_not_single_sym(%i0: index, %i1: index, %i2: index) -> index {
898  // CHECK-COUNT-2: affine.max
899  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
900  %1 = affine.max affine_map<()[s0, s1] -> (s0 + 4, 7 + s1)> ()[%0, %i2]
901  return %1: index
902}
903
904// -----
905
906// Ensure bounding maps of affine.for are composed.
907
908// CHECK-DAG: #[[$MAP0]] = affine_map<()[s0] -> (s0 - 2)>
909// CHECK-DAG: #[[$MAP1]] = affine_map<()[s0] -> (s0 + 2)>
910
911// CHECK-LABEL: func @compose_affine_for_bounds
912// CHECK-SAME:   %[[N:.*]]: index)
913// CHECK: affine.for %{{.*}} = #[[$MAP0]]()[%[[N]]] to #[[$MAP1]]()[%[[N]]] {
914
915func @compose_affine_for_bounds(%N: index) {
916  %u = affine.apply affine_map<(d0) -> (d0 + 2)>(%N)
917  %l = affine.apply affine_map<(d0) -> (d0 - 2)>(%N)
918  affine.for %i = %l to %u {
919    "foo"() : () -> ()
920  }
921  return
922}
923
924// -----
925
926// Compose maps into affine.vector_load / affine.vector_store
927
928// CHECK-LABEL: func @compose_into_affine_vector_load_vector_store
929// CHECK: affine.for %[[IV:.*]] = 0 to 1024
930// CHECK-NEXT: affine.vector_load %{{.*}}[%[[IV]] + 1]
931// CHECK-NEXT: affine.vector_store %{{.*}}, %{{.*}}[%[[IV]] + 1]
932// CHECK-NEXT: affine.vector_load %{{.*}}[%[[IV]]]
933func @compose_into_affine_vector_load_vector_store(%A : memref<1024xf32>, %u : index) {
934  affine.for %i = 0 to 1024 {
935    // Make sure the unused operand (%u below) gets dropped as well.
936    %idx = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i, %u)
937    %0 = affine.vector_load %A[%idx] : memref<1024xf32>, vector<8xf32>
938    affine.vector_store %0, %A[%idx] : memref<1024xf32>, vector<8xf32>
939
940    // Map remains the same, but operand changes on composition.
941    %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
942    %1 = affine.vector_load %A[%copy] : memref<1024xf32>, vector<8xf32>
943    "prevent.dce"(%1) : (vector<8xf32>) -> ()
944  }
945  return
946}
947
948// -----
949
950// CHECK-LABEL: func @no_fold_of_store
951//  CHECK:   %[[cst:.+]] = memref.cast %arg
952//  CHECK:   affine.store %[[cst]]
953func @no_fold_of_store(%arg : memref<32xi8>, %holder: memref<memref<?xi8>>) {
954  %0 = memref.cast %arg : memref<32xi8> to memref<?xi8>
955  affine.store %0, %holder[] : memref<memref<?xi8>>
956  return
957}
958
959// -----
960
961// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 + 16)>
962// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 * 4)>
963
964// CHECK: func @canonicalize_single_min_max
965// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
966func @canonicalize_single_min_max(%i0: index, %i1: index) -> (index, index) {
967  // CHECK-NOT: affine.min
968  // CHECK-NEXT: affine.apply #[[$MAP0]]()[%[[I0]]]
969  %0 = affine.min affine_map<()[s0] -> (s0 + 16)> ()[%i0]
970
971  // CHECK-NOT: affine.max
972  // CHECK-NEXT: affine.apply #[[$MAP1]]()[%[[I1]]]
973  %1 = affine.min affine_map<()[s0] -> (s0 * 4)> ()[%i1]
974
975  return %0, %1: index, index
976}
977