1// RUN: mlir-opt %s -split-input-file -pass-pipeline='func(canonicalize)' | FileCheck %s
2
3// Affine maps for test case: compose_affine_maps_1dto2d_no_symbols
4// CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 - 1)>
5// CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 1)>
6
7// Affine maps for test case: compose_affine_maps_1dto2d_with_symbols
8// CHECK-DAG: [[MAP4:#map[0-9]+]] = affine_map<(d0) -> (d0 - 4)>
9// CHECK-DAG: [[MAP4b:#map[0-9]+]] = affine_map<(d0) -> (d0 - 7)>
10// CHECK-DAG: [[MAP7:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2 - 3)>
11// CHECK-DAG: [[MAP7a:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2 + 1)>
12
13// Affine map for test case: compose_affine_maps_d2_tile
14// CHECK-DAG: [[MAP8:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
15// CHECK-DAG: [[MAP8a:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
16
17// Affine maps for test case: compose_affine_maps_dependent_loads
18// CHECK-DAG: [[MAP9:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
19// CHECK-DAG: [[MAP10:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)>
20// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 7) ceildiv 3)>
21// CHECK-DAG: [[MAP12:#map[0-9]+]] = affine_map<(d0) -> (d0 * 7 - 49)>
22
23// Affine maps for test case: compose_affine_maps_diamond_dependency
24// CHECK-DAG: [[MAP13A:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
25// CHECK-DAG: [[MAP13B:#map[0-9]+]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
26
27// Affine maps for test case: partial_fold_map
28// CHECK-DAG: [[MAP15:#map[0-9]+]] = affine_map<()[s0] -> (s0 - 42)>
29
30// Affine maps for test cases: symbolic_composition_*
31// CHECK-DAG: [[map_symbolic_composition_a:#map[0-9]+]] = affine_map<()[s0] -> (s0 * 512)>
32// CHECK-DAG: [[map_symbolic_composition_b:#map[0-9]+]] = affine_map<()[s0] -> (s0 * 4)>
33// CHECK-DAG: [[map_symbolic_composition_c:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
34// CHECK-DAG: [[map_symbolic_composition_d:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 * 3 + s0)>
35
36// Affine maps for test cases: map_mix_dims_and_symbols_*
37// CHECK-DAG: [[map_mix_dims_and_symbols_b:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 + s0 * 42 + 6)>
38// CHECK-DAG: [[map_mix_dims_and_symbols_c:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 * 4 + s0 * 168 - 4)>
39// CHECK-DAG: [[map_mix_dims_and_symbols_d:#map[0-9]+]] = affine_map<()[s0, s1] -> ((s1 + s0 * 42 + 6) ceildiv 8)>
40// CHECK-DAG: [[map_mix_dims_and_symbols_e:#map[0-9]+]] = affine_map<()[s0, s1] -> ((s1 * 4 + s0 * 168 - 4) floordiv 3)>
41
42// Affine maps for test case: symbolic_semi_affine
43// CHECK-DAG: [[symbolic_semi_affine:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 floordiv (s0 + 1))>
44
45// CHECK-LABEL: func @compose_affine_maps_1dto2d_no_symbols() {
46func @compose_affine_maps_1dto2d_no_symbols() {
47  %0 = alloc() : memref<4x4xf32>
48
49  affine.for %i0 = 0 to 15 {
50    // Test load[%x, %x]
51
52    %x0 = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
53    %x1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %x0)
54    %x1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %x0)
55
56    // CHECK: [[I0A:%[0-9]+]] = affine.apply [[MAP0]](%{{.*}})
57    // CHECK-NEXT: load %0{{\[}}[[I0A]], [[I0A]]{{\]}}
58    %v0 = load %0[%x1_0, %x1_1] : memref<4x4xf32>
59
60    // Test load[%y, %y]
61    %y0 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
62    %y1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %y0)
63    %y1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %y0)
64
65    // CHECK-NEXT: [[I1A:%[0-9]+]] = affine.apply [[MAP1]](%{{.*}})
66    // CHECK-NEXT: load %0{{\[}}[[I1A]], [[I1A]]{{\]}}
67    %v1 = load %0[%y1_0, %y1_1] : memref<4x4xf32>
68
69    // Test load[%x, %y]
70    %xy_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %y0)
71    %xy_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %y0)
72
73    // CHECK-NEXT: load %0{{\[}}[[I0A]], [[I1A]]{{\]}}
74    %v2 = load %0[%xy_0, %xy_1] : memref<4x4xf32>
75
76    // Test load[%y, %x]
77    %yx_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %x0)
78    %yx_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %x0)
79    // CHECK-NEXT: load %0{{\[}}[[I1A]], [[I0A]]{{\]}}
80    %v3 = load %0[%yx_0, %yx_1] : memref<4x4xf32>
81  }
82  return
83}
84
85// CHECK-LABEL: func @compose_affine_maps_1dto2d_with_symbols() {
86func @compose_affine_maps_1dto2d_with_symbols() {
87  %0 = alloc() : memref<4x4xf32>
88
89  affine.for %i0 = 0 to 15 {
90    // Test load[%x0, %x0] with symbol %c4
91    %c4 = constant 4 : index
92    %x0 = affine.apply affine_map<(d0)[s0] -> (d0 - s0)> (%i0)[%c4]
93
94    // CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP4]](%{{.*}})
95    // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I0]], [[I0]]{{\]}}
96    %v0 = load %0[%x0, %x0] : memref<4x4xf32>
97
98    // Test load[%x0, %x1] with symbol %c4 captured by '%x0' map.
99    %x1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
100    %y1 = affine.apply affine_map<(d0, d1) -> (d0+d1)> (%x0, %x1)
101    // CHECK-NEXT: [[I1:%[0-9]+]] = affine.apply [[MAP7]](%{{.*}})
102    // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I1]], [[I1]]{{\]}}
103    %v1 = load %0[%y1, %y1] : memref<4x4xf32>
104
105    // Test load[%x1, %x0] with symbol %c4 captured by '%x0' map.
106    %y2 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x1, %x0)
107    // CHECK-NEXT: [[I2:%[0-9]+]] = affine.apply [[MAP7]](%{{.*}})
108    // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I2]], [[I2]]{{\]}}
109    %v2 = load %0[%y2, %y2] : memref<4x4xf32>
110
111    // Test load[%x2, %x0] with symbol %c4 from '%x0' and %c5 from '%x2'
112    %c5 = constant 5 : index
113    %x2 = affine.apply affine_map<(d0)[s0] -> (d0 + s0)> (%i0)[%c5]
114    %y3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x2, %x0)
115    // CHECK: [[I3:%[0-9]+]] = affine.apply [[MAP7a]](%{{.*}})
116    // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I3]], [[I3]]{{\]}}
117    %v3 = load %0[%y3, %y3] : memref<4x4xf32>
118  }
119  return
120}
121
122// CHECK-LABEL: func @compose_affine_maps_2d_tile() {
123func @compose_affine_maps_2d_tile() {
124  %0 = alloc() : memref<16x32xf32>
125  %1 = alloc() : memref<16x32xf32>
126
127  %c4 = constant 4 : index
128  %c8 = constant 8 : index
129
130  affine.for %i0 = 0 to 3 {
131    %x0 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i0)[%c4]
132    affine.for %i1 = 0 to 3 {
133      %x1 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i1)[%c8]
134      affine.for %i2 = 0 to 3 {
135        %x2 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i2)[%c4]
136        affine.for %i3 = 0 to 3 {
137          %x3 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i3)[%c8]
138
139          %x40 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
140            ((d0 * s0) + d2)> (%x0, %x1, %x2, %x3)[%c4, %c8]
141          %x41 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
142            ((d1 * s1) + d3)> (%x0, %x1, %x2, %x3)[%c4, %c8]
143          // CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP8]](%{{.*}}, %{{.*}})
144          // CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP8a]](%{{.*}}, %{{.*}})
145          // CHECK-NEXT: [[L0:%[0-9]+]] = load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
146          %v0 = load %0[%x40, %x41] : memref<16x32xf32>
147
148          // CHECK-NEXT: store [[L0]], %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
149          store %v0, %1[%x40, %x41] : memref<16x32xf32>
150        }
151      }
152    }
153  }
154  return
155}
156
157// CHECK-LABEL: func @compose_affine_maps_dependent_loads() {
158func @compose_affine_maps_dependent_loads() {
159  %0 = alloc() : memref<16x32xf32>
160  %1 = alloc() : memref<16x32xf32>
161
162  affine.for %i0 = 0 to 3 {
163    affine.for %i1 = 0 to 3 {
164      affine.for %i2 = 0 to 3 {
165        %c3 = constant 3 : index
166        %c7 = constant 7 : index
167
168        %x00 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + s0)>
169            (%i0, %i1, %i2)[%c3, %c7]
170        %x01 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d1 - s1)>
171            (%i0, %i1, %i2)[%c3, %c7]
172        %x02 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d2 * s0)>
173            (%i0, %i1, %i2)[%c3, %c7]
174
175        // CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP9]](%{{.*}})
176        // CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP4b]](%{{.*}})
177        // CHECK: [[I2:%[0-9]+]] = affine.apply [[MAP10]](%{{.*}})
178        // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
179        %v0 = load %0[%x00, %x01] : memref<16x32xf32>
180
181        // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I0]], [[I2]]{{\]}}
182        %v1 = load %0[%x00, %x02] : memref<16x32xf32>
183
184        // Swizzle %i0, %i1
185        // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I1]], [[I0]]{{\]}}
186        %v2 = load %0[%x01, %x00] : memref<16x32xf32>
187
188        // Swizzle %x00, %x01 and %c3, %c7
189        %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)>
190           (%x01, %x00)[%c7, %c3]
191        %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)>
192           (%x01, %x00)[%c7, %c3]
193
194        // CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[MAP12]](%{{.*}})
195        // CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[MAP11]](%{{.*}})
196        // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I2A]], [[I2B]]{{\]}}
197        %v3 = load %0[%x10, %x11] : memref<16x32xf32>
198      }
199    }
200  }
201  return
202}
203
204// CHECK-LABEL: func @compose_affine_maps_diamond_dependency() {
205func @compose_affine_maps_diamond_dependency() {
206  %0 = alloc() : memref<4x4xf32>
207
208  affine.for %i0 = 0 to 15 {
209    %a = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
210    %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
211    %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
212    %d0 = affine.apply affine_map<(d0, d1) -> (d0 ceildiv 8)> (%b, %c)
213    %d1 = affine.apply affine_map<(d0, d1) -> (d1 floordiv 3)> (%b, %c)
214    // CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP13A]](%{{.*}})
215    // CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP13B]](%{{.*}})
216    // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
217    %v = load %0[%d0, %d1] : memref<4x4xf32>
218  }
219
220  return
221}
222
223// CHECK-LABEL: func @arg_used_as_dim_and_symbol
224func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index) {
225  %c9 = constant 9 : index
226  %1 = alloc() : memref<100x100xf32, 1>
227  %2 = alloc() : memref<1xi32>
228  affine.for %i0 = 0 to 100 {
229    affine.for %i1 = 0 to 100 {
230      %3 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 + s0 + s1)>
231        (%i0, %i1)[%arg1, %c9]
232      %4 = affine.apply affine_map<(d0, d1, d3) -> (d3 - (d0 + d1))>
233        (%arg1, %c9, %3)
234      // CHECK: load %{{[0-9]+}}{{\[}}%{{.*}}, %{{.*}}{{\]}}
235      %5 = load %1[%4, %arg1] : memref<100x100xf32, 1>
236    }
237  }
238  return
239}
240
241// CHECK-LABEL: func @trivial_maps
242func @trivial_maps() {
243  // CHECK-NOT: affine.apply
244
245  %0 = alloc() : memref<10xf32>
246  %c0 = constant 0 : index
247  %cst = constant 0.000000e+00 : f32
248  affine.for %i1 = 0 to 10 {
249    %1 = affine.apply affine_map<()[s0] -> (s0)>()[%c0]
250    store %cst, %0[%1] : memref<10xf32>
251    %2 = load %0[%c0] : memref<10xf32>
252
253    %3 = affine.apply affine_map<()[] -> (0)>()[]
254    store %cst, %0[%3] : memref<10xf32>
255    %4 = load %0[%c0] : memref<10xf32>
256  }
257  return
258}
259
260// CHECK-LABEL: func @partial_fold_map
261func @partial_fold_map(%arg1: index, %arg2: index) -> index {
262  // TODO: Constant fold one index into affine.apply
263  %c42 = constant 42 : index
264  %2 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%arg1, %c42)
265  // CHECK: [[X:%[0-9]+]] = affine.apply [[MAP15]]()[%{{.*}}]
266  return %2 : index
267}
268
269// CHECK-LABEL: func @symbolic_composition_a(%{{.*}}: index, %{{.*}}: index) -> index {
270func @symbolic_composition_a(%arg0: index, %arg1: index) -> index {
271  %0 = affine.apply affine_map<(d0) -> (d0 * 4)>(%arg0)
272  %1 = affine.apply affine_map<()[s0, s1] -> (8 * s0)>()[%0, %arg0]
273  %2 = affine.apply affine_map<()[s0, s1] -> (16 * s1)>()[%arg1, %1]
274  // CHECK: %{{.*}} = affine.apply [[map_symbolic_composition_a]]()[%{{.*}}]
275  return %2 : index
276}
277
278// CHECK-LABEL: func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
279func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
280  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
281  %1 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %0]
282  // CHECK: %{{.*}} = affine.apply [[map_symbolic_composition_b]]()[%{{.*}}]
283  return %1 : index
284}
285
286// CHECK-LABEL: func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
287func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
288  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
289  %1 = affine.apply affine_map<(d0) -> (d0)>(%arg1)
290  %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
291  // CHECK: %{{.*}} = affine.apply [[map_symbolic_composition_c]]()[%{{.*}}, %{{.*}}]
292  return %2 : index
293}
294
295// CHECK-LABEL: func @symbolic_composition_d(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
296func @symbolic_composition_d(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
297  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
298  %1 = affine.apply affine_map<()[s0] -> (s0)>()[%arg1]
299  %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
300  // CHECK: %{{.*}} = affine.apply [[map_symbolic_composition_d]]()[%{{.*}}, %{{.*}}]
301  return %2 : index
302}
303
304
305// CHECK-LABEL: func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
306func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
307  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
308  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
309  // CHECK: {{.*}} = affine.apply [[map_mix_dims_and_symbols_b]]()[%{{.*}}, %{{.*}}]
310
311  return %b : index
312}
313
314// CHECK-LABEL: func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
315func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
316  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
317  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
318  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
319  // CHECK: {{.*}} = affine.apply [[map_mix_dims_and_symbols_c]]()[%{{.*}}, %{{.*}}]
320  return %c : index
321}
322
323// CHECK-LABEL: func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
324func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
325  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
326  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
327  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
328  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
329  // CHECK: {{.*}} = affine.apply [[map_mix_dims_and_symbols_d]]()[%{{.*}}, %{{.*}}]
330  return %d : index
331}
332
333// CHECK-LABEL: func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
334func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
335  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
336  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
337  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
338  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
339  %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
340  // CHECK: {{.*}} = affine.apply [[map_mix_dims_and_symbols_e]]()[%{{.*}}, %{{.*}}]
341  return %e : index
342}
343
344// CHECK-LABEL: func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
345func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
346  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
347  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
348  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
349  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
350  %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
351  %f = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 - s1 +  d1 - s0)> (%d, %e)[%e, %d]
352  // CHECK: {{.*}} = constant 0 : index
353
354  return %f : index
355}
356
357// CHECK-LABEL: func @mix_dims_and_symbols_g(%arg0: index, %arg1: index) -> (index, index, index) {
358func @mix_dims_and_symbols_g(%M: index, %N: index) -> (index, index, index) {
359  %K = affine.apply affine_map<(d0) -> (4*d0)> (%M)
360  %res1 = affine.apply affine_map<()[s0, s1] -> (4 * s0)>()[%N, %K]
361  %res2 = affine.apply affine_map<()[s0, s1] -> (s1)>()[%N, %K]
362  %res3 = affine.apply affine_map<()[s0, s1] -> (1024)>()[%N, %K]
363  // CHECK-DAG: {{.*}} = constant 1024 : index
364  // CHECK-DAG: {{.*}} = affine.apply [[map_symbolic_composition_b]]()[%{{.*}}]
365  // CHECK-DAG: {{.*}} = affine.apply [[map_symbolic_composition_b]]()[%{{.*}}]
366  return %res1, %res2, %res3 : index, index, index
367}
368
369// CHECK-LABEL: func @symbolic_semi_affine(%arg0: index, %arg1: index, %arg2: memref<?xf32>) {
370func @symbolic_semi_affine(%M: index, %N: index, %A: memref<?xf32>) {
371  %f1 = constant 1.0 : f32
372  affine.for %i0 = 1 to 100 {
373    %1 = affine.apply affine_map<()[s0] -> (s0 + 1)> ()[%M]
374    %2 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%i0)[%1]
375    // CHECK-DAG: {{.*}} = affine.apply [[symbolic_semi_affine]](%{{.*}})[%{{.*}}]
376    store %f1, %A[%2] : memref<?xf32>
377  }
378  return
379}
380
381// -----
382
383// CHECK: [[MAP0:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
384// CHECK: [[MAP1:#map[0-9]+]] = affine_map<()[s0] -> (100, s0)>
385
386// CHECK-LABEL:  func @constant_fold_bounds(%arg0: index) {
387func @constant_fold_bounds(%N : index) {
388  // CHECK:      constant 3 : index
389  // CHECK-NEXT: "foo"() : () -> index
390  %c9 = constant 9 : index
391  %c1 = constant 1 : index
392  %c2 = constant 2 : index
393  %c3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%c1, %c2)
394  %l = "foo"() : () -> index
395
396  // CHECK:  affine.for %{{.*}} = 5 to 7 {
397  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) {
398    "foo"(%i, %c3) : (index, index) -> ()
399  }
400
401  // Bound takes a non-constant argument but can still be folded.
402  // CHECK:  affine.for %{{.*}} = 1 to 7 {
403  affine.for %j = max affine_map<(d0) -> (0, 1)> (%N) to min affine_map<(d0, d1) -> (7, 9)> (%N, %l) {
404    "foo"(%j, %c3) : (index, index) -> ()
405  }
406
407  // None of the bounds can be folded.
408  // CHECK: affine.for %{{.*}} = max [[MAP0]]()[%{{.*}}] to min [[MAP1]]()[%{{.*}}] {
409  affine.for %k = max affine_map<()[s0] -> (0, s0)> ()[%l] to min affine_map<()[s0] -> (100, s0)> ()[%N] {
410    "foo"(%k, %c3) : (index, index) -> ()
411  }
412  return
413}
414
415// -----
416
417// CHECK-LABEL:  func @fold_empty_loop() {
418func @fold_empty_loop() {
419  // CHECK-NOT: affine.for
420  affine.for %i = 0 to 10 {
421  }
422  return
423}
424// CHECK: return
425
426// -----
427
428// CHECK-DAG: [[SET:#set[0-9]+]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
429
430// CHECK-LABEL: func @canonicalize_affine_if
431// CHECK-SAME: [[M:%.*]]: index,
432// CHECK-SAME: [[N:%.*]]: index)
433func @canonicalize_affine_if(%M : index, %N : index) {
434  %c1022 = constant 1022 : index
435  // Drop unused operand %M, propagate %c1022, and promote %N to symbolic.
436  affine.for %i = 0 to 1024 {
437    affine.for %j = 0 to %N {
438      // CHECK: affine.if [[SET]](%{{.*}}, %{{.*}}){{\[}}[[N]]{{\]}}
439      affine.if affine_set<(d0, d1, d2, d3)[s0] : (d1 >= 0, d0 - d1 >= 0, d2 >= 0, d3 - d2 - 2 >= 0)> (%c1022, %i, %j, %N)[%M] {
440        "foo"() : () -> ()
441      }
442      "bar"() : () -> ()
443    }
444  }
445  return
446}
447
448// -----
449
450// CHECK-DAG: [[LBMAP:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
451// CHECK-DAG: [[UBMAP:#map[0-9]+]] = affine_map<()[s0] -> (1024, s0 + s0)>
452
453// CHECK-LABEL: func @canonicalize_bounds
454// CHECK-SAME: [[M:%.*]]: index,
455// CHECK-SAME: [[N:%.*]]: index)
456func @canonicalize_bounds(%M : index, %N : index) {
457  %c0 = constant 0 : index
458  %c1024 = constant 1024 : index
459  // Drop unused operand %N, drop duplicate operand %M, propagate %c1024, and
460  // promote %M to a symbolic one.
461  // CHECK: affine.for %{{.*}} = 0 to min [[UBMAP]](){{\[}}[[M]]{{\]}}
462  affine.for %i = 0 to min affine_map<(d0, d1, d2, d3) -> (d0, d1 + d2)> (%c1024, %M, %M, %N) {
463    "foo"() : () -> ()
464  }
465  // Promote %M to symbolic position.
466  // CHECK: affine.for %{{.*}} = 0 to #map{{[0-9]+}}(){{\[}}[[M]]{{\]}}
467  affine.for %i = 0 to affine_map<(d0) -> (4 * d0)> (%M) {
468    "foo"() : () -> ()
469  }
470  // Lower bound canonicalize.
471  // CHECK: affine.for %{{.*}} = max [[LBMAP]](){{\[}}[[N]]{{\]}} to [[M]]
472  affine.for %i = max affine_map<(d0, d1) -> (d0, d1)> (%c0, %N) to %M {
473    "foo"() : () -> ()
474  }
475  return
476}
477
478// -----
479
480// Compose maps into affine load and store ops.
481
482// CHECK-DAG: #map{{[0-9]+}} = affine_map<(d0) -> (d0 + 1)>
483
484// CHECK-LABEL: @compose_into_affine_load_store
485func @compose_into_affine_load_store(%A : memref<1024xf32>, %u : index) {
486  %cf1 = constant 1.0 : f32
487  // CHECK: affine.for %[[IV:.*]] = 0 to 1024
488  affine.for %i = 0 to 1024 {
489    // Make sure the unused operand (%u below) gets dropped as well.
490    %idx = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i, %u)
491    affine.load %A[%idx] : memref<1024xf32>
492    affine.store %cf1, %A[%idx] : memref<1024xf32>
493    // CHECK-NEXT: affine.load %{{.*}}[%[[IV]] + 1]
494    // CHECK-NEXT: affine.store %cst, %{{.*}}[%[[IV]] + 1]
495
496    // Map remains the same, but operand changes on composition.
497    %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
498    affine.load %A[%copy] : memref<1024xf32>
499    // CHECK-NEXT: affine.load %{{.*}}[%[[IV]]]
500  }
501  return
502}
503
504// -----
505
506func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
507  %c511 = constant 511 : index
508  %c1 = constant 0 : index
509  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0 + 1)> (%c1)[%c511]
510  "op0"(%0) : (index) -> ()
511  // CHECK:       %[[CST:.*]] = constant 512 : index
512  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
513  // CHECK-NEXT:  return
514  return
515}
516
517// -----
518
519func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
520  %c3 = constant 3 : index
521  %c20 = constant 20 : index
522  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
523  "op0"(%0) : (index) -> ()
524  // CHECK:       %[[CST:.*]] = constant 4 : index
525  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
526  // CHECK-NEXT:  return
527  return
528}
529