xref: /llvm-project/mlir/test/Dialect/Affine/canonicalize.mlir (revision 9f5cefebb482331796ceaebbfcebcd5aee1eb339)
1// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
2// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence top-down=0" | FileCheck %s --check-prefix=CHECK-BOTTOM-UP
3
4// -----
5
6// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 - 1)>
7// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0 + 1)>
8
9// CHECK-LABEL: func @compose_affine_maps_1dto2d_no_symbols() {
10func.func @compose_affine_maps_1dto2d_no_symbols() {
11  %0 = memref.alloc() : memref<4x4xf32>
12
13  affine.for %i0 = 0 to 15 {
14    // Test load[%x, %x]
15
16    %x0 = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
17    %x1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %x0)
18    %x1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %x0)
19
20    // CHECK: %[[I0A:.*]] = affine.apply #[[$MAP0]](%{{.*}})
21    // CHECK-NEXT: %[[V0:.*]] = memref.load %{{.*}}[%[[I0A]], %[[I0A]]]
22    %v0 = memref.load %0[%x1_0, %x1_1] : memref<4x4xf32>
23
24    // Test store[%y, %y]
25    %y0 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
26    %y1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %y0)
27    %y1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %y0)
28
29    // CHECK-NEXT: %[[I1A:.*]] = affine.apply #[[$MAP1]](%{{.*}})
30    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I1A]], %[[I1A]]]
31    memref.store %v0, %0[%y1_0, %y1_1] : memref<4x4xf32>
32
33    // Test store[%x, %y]
34    %xy_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %y0)
35    %xy_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %y0)
36
37    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I0A]], %[[I1A]]]
38    memref.store %v0, %0[%xy_0, %xy_1] : memref<4x4xf32>
39
40    // Test store[%y, %x]
41    %yx_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %x0)
42    %yx_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %x0)
43    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I1A]], %[[I0A]]]
44    memref.store %v0, %0[%yx_0, %yx_1] : memref<4x4xf32>
45  }
46  return
47}
48
49// -----
50
51// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> (d0 - 4)>
52// CHECK-DAG: #[[$MAP7:.*]] = affine_map<(d0) -> (d0 * 2 - 3)>
53// CHECK-DAG: #[[$MAP7a:.*]] = affine_map<(d0) -> (d0 * 2 + 1)>
54
55// CHECK-LABEL: func @compose_affine_maps_1dto2d_with_symbols() {
56func.func @compose_affine_maps_1dto2d_with_symbols() {
57  %0 = memref.alloc() : memref<4x4xf32>
58
59  affine.for %i0 = 0 to 15 {
60    // Test load[%x0, %x0] with symbol %c4
61    %c4 = arith.constant 4 : index
62    %x0 = affine.apply affine_map<(d0)[s0] -> (d0 - s0)> (%i0)[%c4]
63
64    // CHECK: %[[I0:.*]] = affine.apply #[[$MAP4]](%{{.*}})
65    // CHECK-NEXT: %[[V0:.*]] = memref.load %{{.*}}[%[[I0]], %[[I0]]]
66    %v0 = memref.load %0[%x0, %x0] : memref<4x4xf32>
67
68    // Test load[%x0, %x1] with symbol %c4 captured by '%x0' map.
69    %x1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
70    %y1 = affine.apply affine_map<(d0, d1) -> (d0+d1)> (%x0, %x1)
71    // CHECK-NEXT: %[[I1:.*]] = affine.apply #[[$MAP7]](%{{.*}})
72    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I1]], %[[I1]]]
73    memref.store %v0, %0[%y1, %y1] : memref<4x4xf32>
74
75    // Test store[%x1, %x0] with symbol %c4 captured by '%x0' map.
76    %y2 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x1, %x0)
77    // CHECK-NEXT: %[[I2:.*]] = affine.apply #[[$MAP7]](%{{.*}})
78    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I2]], %[[I2]]]
79    memref.store %v0, %0[%y2, %y2] : memref<4x4xf32>
80
81    // Test store[%x2, %x0] with symbol %c4 from '%x0' and %c5 from '%x2'
82    %c5 = arith.constant 5 : index
83    %x2 = affine.apply affine_map<(d0)[s0] -> (d0 + s0)> (%i0)[%c5]
84    %y3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x2, %x0)
85    // CHECK: %[[I3:.*]] = affine.apply #[[$MAP7a]](%{{.*}})
86    // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I3]], %[[I3]]]
87    memref.store %v0, %0[%y3, %y3] : memref<4x4xf32>
88  }
89  return
90}
91
92// -----
93
94// CHECK-DAG: #[[$MAP8:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
95// CHECK-DAG: #[[$MAP8a:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
96
97// CHECK-LABEL: func @compose_affine_maps_2d_tile
98func.func @compose_affine_maps_2d_tile(%0: memref<16x32xf32>, %1: memref<16x32xf32>) {
99  %c4 = arith.constant 4 : index
100  %c8 = arith.constant 8 : index
101
102  affine.for %i0 = 0 to 16 {
103    %x0 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i0)[%c4]
104    affine.for %i1 = 0 to 16 {
105      %x1 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i1)[%c8]
106      affine.for %i2 = 0 to 16 {
107        %x2 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i2)[%c4]
108        affine.for %i3 = 0 to 16 {
109          %x3 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i3)[%c8]
110
111          %x40 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
112            ((d0 * s0) + d2)> (%x0, %x1, %x2, %x3)[%c4, %c8]
113          %x41 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
114            ((d1 * s1) + d3)> (%x0, %x1, %x2, %x3)[%c4, %c8]
115          // CHECK: %[[I0:.*]] = affine.apply #[[$MAP8]](%{{.*}}, %{{.*}})
116          // CHECK: %[[I1:.*]] = affine.apply #[[$MAP8a]](%{{.*}}, %{{.*}})
117          // CHECK-NEXT: %[[L0:.*]] = memref.load %{{.*}}[%[[I0]], %[[I1]]]
118          %v0 = memref.load %0[%x40, %x41] : memref<16x32xf32>
119
120          // CHECK-NEXT: memref.store %[[L0]], %{{.*}}[%[[I0]], %[[I1]]]
121          memref.store %v0, %1[%x40, %x41] : memref<16x32xf32>
122        }
123      }
124    }
125  }
126  return
127}
128
129// -----
130
131// CHECK-DAG: #[[$MAP4b:.*]] = affine_map<(d0) -> (d0 - 7)>
132// CHECK-DAG: #[[$MAP9:.*]] = affine_map<(d0) -> (d0 + 3)>
133// CHECK-DAG: #[[$MAP10:.*]] = affine_map<(d0) -> (d0 * 3)>
134// CHECK-DAG: #[[$MAP11:.*]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)>
135// CHECK-DAG: #[[$MAP12:.*]] = affine_map<(d0) -> (d0 * 7 - 49)>
136
137// CHECK-LABEL: func @compose_affine_maps_dependent_loads() {
138func.func @compose_affine_maps_dependent_loads() {
139  %0 = memref.alloc() : memref<16x32xf32>
140  %1 = memref.alloc() : memref<16x32xf32>
141
142  affine.for %i0 = 0 to 3 {
143    affine.for %i1 = 0 to 3 {
144      affine.for %i2 = 0 to 3 {
145        %c3 = arith.constant 3 : index
146        %c7 = arith.constant 7 : index
147
148        %x00 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + s0)>
149            (%i0, %i1, %i2)[%c3, %c7]
150        %x01 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d1 - s1)>
151            (%i0, %i1, %i2)[%c3, %c7]
152        %x02 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d2 * s0)>
153            (%i0, %i1, %i2)[%c3, %c7]
154
155        // CHECK: %[[I0:.*]] = affine.apply #[[$MAP9]](%{{.*}})
156        // CHECK: %[[I1:.*]] = affine.apply #[[$MAP4b]](%{{.*}})
157        // CHECK: %[[I2:.*]] = affine.apply #[[$MAP10]](%{{.*}})
158        // CHECK-NEXT: %[[V0:.*]] = memref.load %{{.*}}[%[[I0]], %[[I1]]]
159        %v0 = memref.load %0[%x00, %x01] : memref<16x32xf32>
160
161        // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I0]], %[[I2]]]
162        memref.store %v0, %0[%x00, %x02] : memref<16x32xf32>
163
164        // Swizzle %i0, %i1
165        // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I1]], %[[I0]]]
166        memref.store %v0, %0[%x01, %x00] : memref<16x32xf32>
167
168        // Swizzle %x00, %x01 and %c3, %c7
169        %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)>
170           (%x01, %x00)[%c3, %c7]
171        %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)>
172           (%x01, %x00)[%c3, %c7]
173
174        // CHECK-NEXT: %[[I2A:.*]] = affine.apply #[[$MAP12]](%{{.*}})
175        // CHECK-NEXT: %[[I2B:.*]] = affine.apply #[[$MAP11]](%{{.*}})
176        // CHECK-NEXT: memref.store %[[V0]], %{{.*}}[%[[I2A]], %[[I2B]]]
177        memref.store %v0, %0[%x10, %x11] : memref<16x32xf32>
178      }
179    }
180  }
181  return
182}
183
184// -----
185
186// CHECK-DAG: #[[$MAP13A:.*]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
187// CHECK-DAG: #[[$MAP13B:.*]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
188
189// CHECK-LABEL: func @compose_affine_maps_diamond_dependency
190func.func @compose_affine_maps_diamond_dependency(%arg0: f32, %arg1: memref<4x4xf32>) {
191  affine.for %i0 = 0 to 15 {
192    %a = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
193    %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
194    %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
195    %d0 = affine.apply affine_map<(d0, d1) -> (d0 ceildiv 8)> (%b, %c)
196    %d1 = affine.apply affine_map<(d0, d1) -> (d1 floordiv 3)> (%b, %c)
197    // CHECK: %[[I0:.*]] = affine.apply #[[$MAP13A]](%{{.*}})
198    // CHECK: %[[I1:.*]] = affine.apply #[[$MAP13B]](%{{.*}})
199    // CHECK-NEXT: memref.store %arg0, %arg1[%[[I0]], %[[I1]]]
200    memref.store %arg0, %arg1[%d0, %d1] : memref<4x4xf32>
201  }
202
203  return
204}
205
206// -----
207
208// CHECK-DAG: #[[$MAP14:.*]] = affine_map<()[s0, s1] -> ((s0 * 4 + s1 * 4) floordiv s0)>
209
210// CHECK-LABEL: func @compose_affine_maps_multiple_symbols
211func.func @compose_affine_maps_multiple_symbols(%arg0: index, %arg1: index) -> index {
212  %a = affine.apply affine_map<(d0)[s0] -> (s0 + d0)> (%arg0)[%arg1]
213  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
214  %e = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%c)[%arg1]
215  // CHECK: [[I0:.*]] = affine.apply #[[$MAP14]]()[%{{.*}}, %{{.*}}]
216  return %e : index
217}
218
219// -----
220
221// CHECK-LABEL: func @arg_used_as_dim_and_symbol
222func.func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) -> (memref<100x100xf32, 1>, memref<1xi32>) {
223  %c9 = arith.constant 9 : index
224  %1 = memref.alloc() : memref<100x100xf32, 1>
225  %2 = memref.alloc() : memref<1xi32>
226  affine.for %i0 = 0 to 100 {
227    affine.for %i1 = 0 to 100 {
228      %3 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 + s0 + s1)>
229        (%i0, %i1)[%arg1, %c9]
230      %4 = affine.apply affine_map<(d0, d1, d3) -> (d3 - (d0 + d1))>
231        (%arg1, %c9, %3)
232      // CHECK: memref.store %arg2, %{{.*}}[%{{.*}}, %{{.*}}]
233      memref.store %arg2, %1[%4, %arg1] : memref<100x100xf32, 1>
234    }
235  }
236  return %1, %2 : memref<100x100xf32, 1>, memref<1xi32>
237}
238
239// -----
240
241// CHECK-LABEL: func @trivial_maps
242func.func @trivial_maps() {
243  // CHECK-NOT: affine.apply
244
245  %0 = memref.alloc() : memref<10xf32>
246  %c0 = arith.constant 0 : index
247  %cst = arith.constant 0.000000e+00 : f32
248  affine.for %i1 = 0 to 10 {
249    %1 = affine.apply affine_map<()[s0] -> (s0)>()[%c0]
250    memref.store %cst, %0[%1] : memref<10xf32>
251    %2 = memref.load %0[%c0] : memref<10xf32>
252
253    %3 = affine.apply affine_map<()[] -> (0)>()[]
254    memref.store %cst, %0[%3] : memref<10xf32>
255    memref.store %2, %0[%c0] : memref<10xf32>
256  }
257  return
258}
259
260// -----
261
262// CHECK-DAG: #[[$MAP15:.*]] = affine_map<()[s0] -> (s0 - 42)>
263
264// CHECK-LABEL: func @partial_fold_map
265func.func @partial_fold_map(%arg1: index, %arg2: index) -> index {
266  // TODO: Constant fold one index into affine.apply
267  %c42 = arith.constant 42 : index
268  %2 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%arg1, %c42)
269  // CHECK: [[X:.*]] = affine.apply #[[$MAP15]]()[%{{.*}}]
270  return %2 : index
271}
272
273// -----
274
275// CHECK-DAG: #[[$MAP_symbolic_composition_a:.*]] = affine_map<()[s0] -> (s0 * 512)>
276
277// CHECK-LABEL: func @symbolic_composition_a(%{{.*}}: index, %{{.*}}: index) -> index {
278func.func @symbolic_composition_a(%arg0: index, %arg1: index) -> index {
279  %0 = affine.apply affine_map<(d0) -> (d0 * 4)>(%arg0)
280  %1 = affine.apply affine_map<()[s0, s1] -> (8 * s0)>()[%0, %arg0]
281  %2 = affine.apply affine_map<()[s0, s1] -> (16 * s1)>()[%arg1, %1]
282  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_a]]()[%{{.*}}]
283  return %2 : index
284}
285
286// -----
287
288// CHECK-DAG: #[[$MAP_symbolic_composition_b:.*]] = affine_map<()[s0] -> (s0 * 4)>
289
290// CHECK-LABEL: func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
291func.func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
292  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
293  %1 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %0]
294  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
295  return %1 : index
296}
297
298// -----
299
300// CHECK-DAG: #[[$MAP_symbolic_composition_c:.*]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
301
302// CHECK-LABEL: func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
303func.func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
304  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
305  %1 = affine.apply affine_map<(d0) -> (d0)>(%arg1)
306  %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
307  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_c]]()[%{{.*}}, %{{.*}}]
308  return %2 : index
309}
310
311// -----
312
313// CHECK-DAG: #[[$MAP_symbolic_composition_d:.*]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
314
315// CHECK-LABEL: func @symbolic_composition_d(
316//  CHECK-SAME:   %[[ARG0:[0-9a-zA-Z]+]]: index
317//  CHECK-SAME:   %[[ARG1:[0-9a-zA-Z]+]]: index
318func.func @symbolic_composition_d(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
319  %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
320  %1 = affine.apply affine_map<()[s0] -> (s0)>()[%arg1]
321  %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
322  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_d]]()[%[[ARG0]], %[[ARG1]]]
323  return %2 : index
324}
325
326// -----
327
328// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_b:.*]] = affine_map<()[s0, s1] -> (s0 * 42 + s1 + 6)>
329
330// CHECK-LABEL: func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
331func.func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
332  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
333  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
334  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_b]]()[%{{.*}}, %{{.*}}]
335
336  return %b : index
337}
338
339// -----
340
341// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_c:.*]] = affine_map<()[s0, s1] -> (s0 * 168 + s1 * 4 - 4)>
342
343// CHECK-LABEL: func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
344func.func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
345  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
346  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
347  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
348  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_c]]()[%{{.*}}, %{{.*}}]
349  return %c : index
350}
351
352// -----
353
354// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_d:.*]] = affine_map<()[s0, s1] -> ((s0 * 42 + s1 + 6) ceildiv 8)>
355
356// CHECK-LABEL: func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
357func.func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
358  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
359  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
360  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
361  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
362  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_d]]()[%{{.*}}, %{{.*}}]
363  return %d : index
364}
365
366// -----
367
368// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_e:.*]] = affine_map<()[s0, s1] -> ((s0 * 168 + s1 * 4 - 4) floordiv 3)>
369
370// CHECK-LABEL: func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
371func.func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
372  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
373  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
374  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
375  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
376  %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
377  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_e]]()[%{{.*}}, %{{.*}}]
378  return %e : index
379}
380
381// -----
382
383// CHECK-LABEL: func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
384func.func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
385  %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
386  %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
387  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
388  %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
389  %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
390  %f = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 - s1 +  d1 - s0)> (%d, %e)[%e, %d]
391  // CHECK: {{.*}} = arith.constant 0 : index
392
393  return %f : index
394}
395
396// -----
397
398// CHECK-DAG: #[[$MAP_symbolic_composition_b:.*]] = affine_map<()[s0] -> (s0 * 4)>
399
400// CHECK-LABEL: func @mix_dims_and_symbols_g(%arg0: index, %arg1: index) -> (index, index, index) {
401func.func @mix_dims_and_symbols_g(%M: index, %N: index) -> (index, index, index) {
402  %K = affine.apply affine_map<(d0) -> (4*d0)> (%M)
403  %res1 = affine.apply affine_map<()[s0, s1] -> (4 * s0)>()[%N, %K]
404  %res2 = affine.apply affine_map<()[s0, s1] -> (s1)>()[%N, %K]
405  %res3 = affine.apply affine_map<()[s0, s1] -> (1024)>()[%N, %K]
406  // CHECK-DAG: {{.*}} = arith.constant 1024 : index
407  // CHECK-DAG: {{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
408  // CHECK-DAG: {{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
409  return %res1, %res2, %res3 : index, index, index
410}
411
412// -----
413
414// CHECK-DAG: #[[$symbolic_semi_affine:.*]] = affine_map<(d0)[s0] -> (d0 floordiv (s0 + 1))>
415
416// CHECK-LABEL: func @symbolic_semi_affine(%arg0: index, %arg1: index, %arg2: memref<?xf32>) {
417func.func @symbolic_semi_affine(%M: index, %N: index, %A: memref<?xf32>) {
418  %f1 = arith.constant 1.0 : f32
419  affine.for %i0 = 1 to 100 {
420    %1 = affine.apply affine_map<()[s0] -> (s0 + 1)> ()[%M]
421    %2 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%i0)[%1]
422    // CHECK-DAG: {{.*}} = affine.apply #[[$symbolic_semi_affine]](%{{.*}})[%{{.*}}]
423    memref.store %f1, %A[%2] : memref<?xf32>
424  }
425  return
426}
427
428// -----
429
430// CHECK: #[[$MAP0:.*]] = affine_map<()[s0] -> (0, s0)>
431// CHECK: #[[$MAP1:.*]] = affine_map<()[s0] -> (100, s0)>
432
433// CHECK-LABEL:  func @constant_fold_bounds(%arg0: index) {
434func.func @constant_fold_bounds(%N : index) {
435  // CHECK:      arith.constant 3 : index
436  // CHECK-NEXT: "foo"() : () -> index
437  %c9 = arith.constant 9 : index
438  %c1 = arith.constant 1 : index
439  %c2 = arith.constant 2 : index
440  %c3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%c1, %c2)
441  %l = "foo"() : () -> index
442
443  // CHECK:  affine.for %{{.*}} = 5 to 7 {
444  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) {
445    "foo"(%i, %c3) : (index, index) -> ()
446  }
447
448  // Bound takes a non-constant argument but can still be folded.
449  // CHECK:  affine.for %{{.*}} = 1 to 7 {
450  affine.for %j = max affine_map<(d0) -> (0, 1)> (%N) to min affine_map<(d0, d1) -> (7, 9)> (%N, %l) {
451    "foo"(%j, %c3) : (index, index) -> ()
452  }
453
454  // None of the bounds can be folded.
455  // CHECK: affine.for %{{.*}} = max #[[$MAP0]]()[%{{.*}}] to min #[[$MAP1]]()[%{{.*}}] {
456  affine.for %k = max affine_map<()[s0] -> (0, s0)> ()[%l] to min affine_map<()[s0] -> (100, s0)> ()[%N] {
457    "foo"(%k, %c3) : (index, index) -> ()
458  }
459  return
460}
461
462// -----
463
464// CHECK-LABEL:  func @fold_empty_loops()
465func.func @fold_empty_loops() -> index {
466  %c0 = arith.constant 0 : index
467  affine.for %i = 0 to 10 {
468  }
469  %res = affine.for %i = 0 to 10 iter_args(%arg = %c0) -> index {
470    affine.yield %arg : index
471  }
472  // CHECK-NEXT: %[[zero:.*]] = arith.constant 0
473  // CHECK-NEXT: return %[[zero]]
474  return %res : index
475}
476
477// -----
478
479// CHECK-LABEL:  func @fold_empty_loop()
480func.func @fold_empty_loop() -> (index, index) {
481  %c0 = arith.constant 0 : index
482  %c1 = arith.constant 1 : index
483  %c2 = arith.constant 2 : index
484  %res:2 = affine.for %i = 0 to 10 iter_args(%arg0 = %c0, %arg1 = %c1) -> (index, index) {
485    affine.yield %c2, %arg1 : index, index
486  }
487  // CHECK-DAG: %[[one:.*]] = arith.constant 1
488  // CHECK-DAG: %[[two:.*]] = arith.constant 2
489  // CHECK-NEXT: return %[[two]], %[[one]]
490  return %res#0, %res#1 : index, index
491}
492
493// -----
494
495// CHECK-LABEL:  func @fold_empty_loops_trip_count_1()
496func.func @fold_empty_loops_trip_count_1() -> (index, index, index, index) {
497  %c0 = arith.constant 0 : index
498  %c1 = arith.constant 1 : index
499  %c2 = arith.constant 2 : index
500  %res1:2 = affine.for %i = 0 to 1 iter_args(%arg0 = %c2, %arg1 = %c0) -> (index, index) {
501    affine.yield %c1, %arg0 : index, index
502  }
503  %res2:2 = affine.for %i = 0 to 2 step 3 iter_args(%arg0 = %c2, %arg1 = %c0) -> (index, index) {
504    affine.yield %arg1, %arg0 : index, index
505  }
506  // CHECK-DAG: %[[zero:.*]] = arith.constant 0
507  // CHECK-DAG: %[[one:.*]] = arith.constant 1
508  // CHECK-DAG: %[[two:.*]] = arith.constant 2
509  // CHECK-NEXT: return %[[one]], %[[two]], %[[zero]], %[[two]]
510  return %res1#0, %res1#1, %res2#0, %res2#1 : index, index, index, index
511}
512
513// -----
514
515// CHECK-LABEL:  func @fold_empty_loop_trip_count_0()
516func.func @fold_empty_loop_trip_count_0() -> (index, index) {
517  %c0 = arith.constant 0 : index
518  %c1 = arith.constant 1 : index
519  %c2 = arith.constant 2 : index
520  %res:2 = affine.for %i = 0 to 0 iter_args(%arg0 = %c2, %arg1 = %c0) -> (index, index) {
521    affine.yield %c1, %arg0 : index, index
522  }
523  // CHECK-DAG: %[[zero:.*]] = arith.constant 0
524  // CHECK-DAG: %[[two:.*]] = arith.constant 2
525  // CHECK-NEXT: return %[[two]], %[[zero]]
526  return %res#0, %res#1 : index, index
527}
528
529// -----
530
531// CHECK-LABEL:  func @fold_empty_loop_trip_count_unknown
532func.func @fold_empty_loop_trip_count_unknown(%in : index) -> (index, index) {
533  %c0 = arith.constant 0 : index
534  %c1 = arith.constant 1 : index
535  %res:2 = affine.for %i = 0 to %in iter_args(%arg0 = %c0, %arg1 = %c1) -> (index, index) {
536    affine.yield %arg0, %arg1 : index, index
537  }
538  // CHECK-DAG: %[[zero:.*]] = arith.constant 0
539  // CHECK-DAG: %[[one:.*]] = arith.constant 1
540  // CHECK-NEXT: return %[[zero]], %[[one]]
541  return %res#0, %res#1 : index, index
542}
543
544// -----
545
546// CHECK-LABEL:  func @empty_loops_not_folded_1
547func.func @empty_loops_not_folded_1(%in : index) -> index {
548  %c0 = arith.constant 0 : index
549  %c1 = arith.constant 1 : index
550  // CHECK: affine.for
551  %res = affine.for %i = 0 to %in iter_args(%arg = %c0) -> index {
552    affine.yield %c1 : index
553  }
554  return %res : index
555}
556
557// -----
558
559// CHECK-LABEL:  func @empty_loops_not_folded_2
560func.func @empty_loops_not_folded_2(%in : index) -> (index, index) {
561  %c0 = arith.constant 0 : index
562  %c1 = arith.constant 1 : index
563  // CHECK: affine.for
564  %res:2 = affine.for %i = 0 to %in iter_args(%arg0 = %c0, %arg1 = %c1) -> (index, index) {
565    affine.yield %arg1, %arg0 : index, index
566  }
567  return %res#0, %res#1 : index, index
568}
569
570// -----
571
572// CHECK-LABEL:  func @empty_loops_not_folded_3
573func.func @empty_loops_not_folded_3() -> (index, index) {
574  %c0 = arith.constant 0 : index
575  %c1 = arith.constant 1 : index
576  // CHECK: affine.for
577  %res:2 = affine.for %i = 0 to 10 iter_args(%arg0 = %c0, %arg1 = %c1) -> (index, index) {
578    affine.yield %arg1, %arg0 : index, index
579  }
580  return %res#0, %res#1 : index, index
581}
582
583// -----
584
585// CHECK-LABEL:  func @zero_iter_loop_not_folded
586func.func @zero_iter_loop_not_folded() {
587  %A = memref.alloc() : memref<4xf32>
588  affine.for %i = 0 to 0 {
589      %load = affine.load %A[%i] : memref<4xf32>
590      affine.store %load, %A[%i] : memref<4xf32>
591  }
592  // CHECK:   affine.for {{.*}} = 0 to 0 {
593  return
594}
595
596// -----
597
598// CHECK-LABEL:  func @fold_zero_iter_loops
599// CHECK-SAME: %[[ARG:.*]]: index
600func.func @fold_zero_iter_loops(%in : index) -> index {
601  %c1 = arith.constant 1 : index
602  %res = affine.for %i = 0 to 0 iter_args(%loop_arg = %in) -> index {
603    %yield = arith.addi %loop_arg, %c1 : index
604    affine.yield %yield : index
605  }
606  // CHECK-NEXT: return %[[ARG]]
607  return %res : index
608}
609
610// -----
611
612// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
613
614// CHECK-LABEL: func @canonicalize_affine_if
615//  CHECK-SAME:   %[[M:[0-9a-zA-Z]*]]: index,
616//  CHECK-SAME:   %[[N:[0-9a-zA-Z]*]]: index)
617func.func @canonicalize_affine_if(%M : index, %N : index) {
618  %c1022 = arith.constant 1022 : index
619  // Drop unused operand %M, propagate %c1022, and promote %N to symbolic.
620  affine.for %i = 0 to 1024 {
621    affine.for %j = 0 to %N {
622      // CHECK: affine.if #[[$SET]](%{{.*}}, %{{.*}})[%[[N]]]
623      affine.if affine_set<(d0, d1, d2, d3)[s0] : (d1 >= 0, d0 - d1 >= 0, d2 >= 0, d3 - d2 - 2 >= 0)>
624          (%c1022, %i, %j, %N)[%M] {
625        "foo"() : () -> ()
626      }
627      "bar"() : () -> ()
628    }
629  }
630  return
631}
632
633// -----
634
635// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 - 1 >= 0, d1 - 1 == 0, -d0 + s0 + 10 >= 0)>
636
637// CHECK-LABEL: func @canonicalize_affine_if_compose_apply
638// CHECK-SAME:   %[[N:.*]]: index
639func.func @canonicalize_affine_if_compose_apply(%N: index) {
640  %M = affine.apply affine_map<()[s0] -> (s0 + 10)> ()[%N]
641  // CHECK-NEXT: affine.for %[[I:.*]] =
642  affine.for %i = 0 to 1024 {
643    // CHECK-NEXT: affine.for %[[J:.*]] =
644    affine.for %j = 0 to 100 {
645      %j_ = affine.apply affine_map<(d0)[] -> (d0 + 1)> (%j)
646      // CHECK-NEXT: affine.if #[[$SET]](%[[I]], %[[J]])[%[[N]]]
647      affine.if affine_set<(d0, d1)[s0] : (d0 - 1 >= 0, d1 - 2 == 0, -d0 + s0 >= 0)>(%i, %j_)[%M] {
648        "test.foo"() : ()->()
649      }
650    }
651  }
652  return
653}
654
655// -----
656
657// CHECK-DAG: #[[$LBMAP:.*]] = affine_map<()[s0] -> (0, s0)>
658// CHECK-DAG: #[[$UBMAP:.*]] = affine_map<()[s0] -> (1024, s0 * 2)>
659
660// CHECK-LABEL: func @canonicalize_bounds
661// CHECK-SAME: %[[M:.*]]: index,
662// CHECK-SAME: %[[N:.*]]: index)
663func.func @canonicalize_bounds(%M : index, %N : index) {
664  %c0 = arith.constant 0 : index
665  %c1024 = arith.constant 1024 : index
666  // Drop unused operand %N, drop duplicate operand %M, propagate %c1024, and
667  // promote %M to a symbolic one.
668  // CHECK: affine.for %{{.*}} = 0 to min #[[$UBMAP]]()[%[[M]]]
669  affine.for %i = 0 to min affine_map<(d0, d1, d2, d3) -> (d0, d1 + d2)> (%c1024, %M, %M, %N) {
670    "foo"() : () -> ()
671  }
672  // Promote %M to symbolic position.
673  // CHECK: affine.for %{{.*}} = 0 to #{{.*}}()[%[[M]]]
674  affine.for %i = 0 to affine_map<(d0) -> (4 * d0)> (%M) {
675    "foo"() : () -> ()
676  }
677  // Lower bound canonicalize.
678  // CHECK: affine.for %{{.*}} = max #[[$LBMAP]]()[%[[N]]] to %[[M]]
679  affine.for %i = max affine_map<(d0, d1) -> (d0, d1)> (%c0, %N) to %M {
680    "foo"() : () -> ()
681  }
682  return
683}
684
685// -----
686
687// Compose maps into affine load and store ops.
688
689// CHECK-LABEL: @compose_into_affine_load_store
690func.func @compose_into_affine_load_store(%A : memref<1024xf32>, %u : index) {
691  // CHECK: affine.for %[[IV:.*]] = 0 to 1024
692  affine.for %i = 0 to 1024 {
693    // Make sure the unused operand (%u below) gets dropped as well.
694    %idx = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i, %u)
695    %0 = affine.load %A[%idx] : memref<1024xf32>
696    affine.store %0, %A[%idx] : memref<1024xf32>
697    // CHECK-NEXT: affine.load %{{.*}}[%[[IV]] + 1]
698    // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%[[IV]] + 1]
699
700    // Map remains the same, but operand changes on composition.
701    %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
702    %1 = affine.load %A[%copy] : memref<1024xf32>
703    "prevent.dce"(%1) : (f32) -> ()
704    // CHECK-NEXT: affine.load %{{.*}}[%[[IV]]]
705  }
706  return
707}
708
709// -----
710
711func.func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
712  %c511 = arith.constant 511 : index
713  %c1 = arith.constant 0 : index
714  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0 + 1)> (%c1)[%c511]
715  "op0"(%0) : (index) -> ()
716  // CHECK:       %[[CST:.*]] = arith.constant 512 : index
717  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
718  // CHECK-NEXT:  return
719  return
720}
721
722// -----
723
724func.func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
725  %c3 = arith.constant 3 : index
726  %c20 = arith.constant 20 : index
727  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
728  "op0"(%0) : (index) -> ()
729  // CHECK:       %[[CST:.*]] = arith.constant 4 : index
730  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
731  // CHECK-NEXT:  return
732  return
733}
734
735// -----
736
737func.func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
738  %c511 = arith.constant 511 : index
739  %c1 = arith.constant 0 : index
740  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0 + 1)> (%c1)[%c511]
741  "op0"(%0) : (index) -> ()
742  // CHECK:       %[[CST:.*]] = arith.constant 1000 : index
743  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
744  // CHECK-NEXT:  return
745  return
746}
747
748// -----
749
750func.func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
751  %c3 = arith.constant 3 : index
752  %c20 = arith.constant 20 : index
753  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
754  "op0"(%0) : (index) -> ()
755  // CHECK:       %[[CST:.*]] = arith.constant 1000 : index
756  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
757  // CHECK-NEXT:  return
758  return
759}
760
761// -----
762
763// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d1 - 2, d0)>
764
765func.func @affine_min(%arg0: index) {
766  affine.for %i = 0 to %arg0 {
767    affine.for %j = 0 to %arg0 {
768      %c2 = arith.constant 2 : index
769      // CHECK: affine.min #[[$MAP]]
770      %0 = affine.min affine_map<(d0,d1,d2)->(d0, d1 - d2)>(%i, %j, %c2)
771      "consumer"(%0) : (index) -> ()
772    }
773  }
774  return
775}
776
777// -----
778
779// Reproducer for PR45031. This used to fold into an incorrect map because
780// symbols were concatenated in the wrong order during map folding. Map
781// composition places the symbols of the original map before those of the map
782// it is composed with, e.g. A.compose(B) will first have all symbols of A,
783// then all symbols of B.
784
785#map1 = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
786#map2 = affine_map<(d0)[s0] -> (1024, -d0 + s0)>
787
788// CHECK: #[[$MAP:.*]] = affine_map<()[s0, s1] -> (1024, s0 - s1 * 1024)>
789
790// CHECK: func @rep(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
791func.func @rep(%arg0 : index, %arg1 : index) -> index {
792  // CHECK-NOT: arith.constant
793  %c0 = arith.constant 0 : index
794  %c1024 = arith.constant 1024 : index
795  // CHECK-NOT: affine.apply
796  %0 = affine.apply #map1(%arg0)[%c1024, %c0]
797
798  // CHECK: affine.min #[[$MAP]]()[%[[ARG1]], %[[ARG0]]]
799  %1 = affine.min #map2(%0)[%arg1]
800  return %1 : index
801}
802
803// -----
804
805// CHECK-DAG: #[[ub:.*]] = affine_map<()[s0] -> (s0 + 2)>
806
807func.func @drop_duplicate_bounds(%N : index) {
808  // affine.for %i = max #lb(%arg0) to min #ub(%arg0)
809  affine.for %i = max affine_map<(d0) -> (d0, d0)>(%N) to min affine_map<(d0) -> (d0 + 2, d0 + 2)>(%N) {
810    "foo"() : () -> ()
811  }
812  return
813}
814
815// -----
816
817// Ensure affine.parallel bounds expressions are canonicalized.
818
819#map3 = affine_map<(d0) -> (d0 * 5)>
820
821// CHECK-LABEL: func @affine_parallel_const_bounds
822func.func @affine_parallel_const_bounds() {
823  %cst = arith.constant 1.0 : f32
824  %c0 = arith.constant 0 : index
825  %c4 = arith.constant 4 : index
826  %0 = memref.alloc() : memref<4xf32>
827  // CHECK: affine.parallel (%{{.*}}) = (0) to (4)
828  affine.parallel (%i) = (%c0) to (%c0 + %c4) {
829    %1 = affine.apply #map3(%i)
830    // CHECK: affine.parallel (%{{.*}}) = (0) to (%{{.*}} * 5)
831    affine.parallel (%j) = (%c0) to (%1) {
832      affine.store %cst, %0[%j] : memref<4xf32>
833    }
834  }
835  return
836}
837
838// -----
839
840func.func @compose_affine_maps_div_symbol(%A : memref<i64>, %i0 : index, %i1 : index) {
841  %0 = affine.apply affine_map<()[s0] -> (2 * s0)> ()[%i0]
842  %1 = affine.apply affine_map<()[s0] -> (3 * s0)> ()[%i0]
843  %2 = affine.apply affine_map<(d0)[s0, s1] -> (d0 mod s1 + s0 * s1 + s0 * 4)> (%i1)[%0, %1]
844  %3 = arith.index_cast %2: index to i64
845  memref.store %3, %A[]: memref<i64>
846  affine.for %i2 = 0 to 3 {
847    %4 = affine.apply affine_map<(d0)[s0, s1] -> (d0 ceildiv s1 + s0 + s0 * 3)> (%i2)[%0, %1]
848    %5 = arith.index_cast %4: index to i64
849    memref.store %5, %A[]: memref<i64>
850  }
851  return
852}
853
854// -----
855
856// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
857
858// CHECK: func @deduplicate_affine_min_expressions
859// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
860func.func @deduplicate_affine_min_expressions(%i0: index, %i1: index) -> index {
861  // CHECK:  affine.min #[[MAP]]()[%[[I0]], %[[I1]]]
862  %0 = affine.min affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
863  return %0: index
864}
865
866// -----
867
868// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
869
870// CHECK: func @deduplicate_affine_max_expressions
871// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
872func.func @deduplicate_affine_max_expressions(%i0: index, %i1: index) -> index {
873  // CHECK:  affine.max #[[MAP]]()[%[[I0]], %[[I1]]]
874  %0 = affine.max affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
875  return %0: index
876}
877
878// -----
879
880// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1, s2] -> (-s1 + s2, 16, s0 * 3)>
881// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (-s0 + s1, -s2 + 5, 16)>
882
883// CHECK: func @merge_affine_min_ops
884// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index, %[[I3:.+]]: index)
885func.func @merge_affine_min_ops(%i0: index, %i1: index, %i2: index, %i3: index) -> (index, index) {
886  %0 = affine.min affine_map<(d0)[s0] -> (16, d0 - s0)> (%i0)[%i1]
887
888 // CHECK: affine.min #[[MAP0]]()[%[[I2]], %[[I1]], %[[I0]]]
889  %1 = affine.min affine_map<(d0)[s0] -> (3 * s0, d0)> (%0)[%i2] // Use as dim
890 // CHECK: affine.min #[[MAP1]]()[%[[I1]], %[[I0]], %[[I3]]]
891  %2 = affine.min affine_map<(d0)[s0] -> (s0, 5 - d0)> (%i3)[%0] // Use as symbol
892
893  return %1, %2: index, index
894}
895
896// -----
897
898// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s2 + 8, s2 * 4, s1 + 16, s1 * 8, s0 + 7)>
899
900// CHECK: func @merge_multiple_affine_min_ops
901// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index)
902func.func @merge_multiple_affine_min_ops(%i0: index, %i1: index, %i2: index) -> index {
903  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
904  %1 = affine.min affine_map<()[s0] -> (s0 + 8, s0 * 4)> ()[%i1]
905  // CHECK: affine.min #[[MAP]]()[%[[I2]], %[[I0]], %[[I1]]]
906  %2 = affine.min affine_map<()[s0, s1, s2] -> (s0, 7 + s1, s2)> ()[%0, %i2, %1]
907  return %2: index
908}
909
910// -----
911
912// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s1 + 16, s1 * 8, s0 * 2)>
913
914// CHECK: func @merge_multiple_uses_of_affine_min_ops
915// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
916func.func @merge_multiple_uses_of_affine_min_ops(%i0: index, %i1: index) -> index {
917  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
918  // CHECK: affine.min #[[MAP]]()[%[[I1]], %[[I0]]]
919  %2 = affine.min affine_map<()[s0, s1, s2] -> (s0, s1, s2 * 2)> ()[%0, %0, %i1]
920  return %2: index
921}
922
923// -----
924
925// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 16, s0 * 8)>
926// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (s2 + 16, s2 * 8, s1 * 2, s0 + 1)>
927
928// CHECK: func @merge_mixed_uses_of_affine_min_ops
929// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
930func.func @merge_mixed_uses_of_affine_min_ops(%i0: index, %i1: index) -> index {
931  // CHECK: %[[AFFINE:.+]] = affine.min #[[MAP0]]()[%[[I0]]]
932  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
933  // %0 is bound to a symbol that is both a standalone expression and a part
934  // of other expressions.
935  // CHECK: affine.min #[[MAP1]]()[%[[AFFINE]], %[[I1]], %[[I0]]]
936  %2 = affine.min affine_map<()[s0, s1, s2] -> (s0, s1 + 1, s2 * 2)> ()[%0, %0, %i1]
937  return %2: index
938}
939
940// -----
941
942// CHECK-LABEL: func @dont_merge_affine_min_if_not_single_dim
943func.func @dont_merge_affine_min_if_not_single_dim(%i0: index, %i1: index, %i2: index) -> index {
944  // CHECK-COUNT-2: affine.min
945  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
946  %1 = affine.min affine_map<(d0)[s0] -> (s0 + 4, 7 + d0)> (%0)[%i2]
947  return %1: index
948}
949
950// -----
951
952// CHECK-LABEL: func @dont_merge_affine_min_if_not_single_sym
953func.func @dont_merge_affine_min_if_not_single_sym(%i0: index, %i1: index, %i2: index) -> index {
954  // CHECK-COUNT-2: affine.min
955  %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
956  %1 = affine.min affine_map<()[s0, s1] -> (s0 + 4, 7 + s1)> ()[%0, %i2]
957  return %1: index
958}
959
960// -----
961
962// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1, s2] -> (-s1 + s2, 16, s0 * 3)>
963// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (-s0 + s1, -s2 + 5, 16)>
964
965// CHECK: func @merge_affine_max_ops
966// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index, %[[I3:.+]]: index)
967func.func @merge_affine_max_ops(%i0: index, %i1: index, %i2: index, %i3: index) -> (index, index) {
968  %0 = affine.max affine_map<(d0)[s0] -> (16, d0 - s0)> (%i0)[%i1]
969
970 // CHECK: affine.max #[[MAP0]]()[%[[I2]], %[[I1]], %[[I0]]]
971  %1 = affine.max affine_map<(d0)[s0] -> (3 * s0, d0)> (%0)[%i2] // Use as dim
972 // CHECK: affine.max #[[MAP1]]()[%[[I1]], %[[I0]], %[[I3]]]
973  %2 = affine.max affine_map<(d0)[s0] -> (s0, 5 - d0)> (%i3)[%0] // Use as symbol
974
975  return %1, %2: index, index
976}
977
978// -----
979
980// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s2 + 8, s2 * 4, s1 + 16, s1 * 8, s0 + 7)>
981
982// CHECK: func @merge_multiple_affine_max_ops
983// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index, %[[I2:.+]]: index)
984func.func @merge_multiple_affine_max_ops(%i0: index, %i1: index, %i2: index) -> index {
985  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
986  %1 = affine.max affine_map<()[s0] -> (s0 + 8, s0 * 4)> ()[%i1]
987  // CHECK: affine.max #[[MAP]]()[%[[I2]], %[[I0]], %[[I1]]]
988  %2 = affine.max affine_map<()[s0, s1, s2] -> (s0, 7 + s1, s2)> ()[%0, %i2, %1]
989  return %2: index
990}
991
992// -----
993
994// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s1 + 16, s1 * 8, s0 * 2)>
995
996// CHECK: func @merge_multiple_uses_of_affine_max_ops
997// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
998func.func @merge_multiple_uses_of_affine_max_ops(%i0: index, %i1: index) -> index {
999  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
1000  // CHECK: affine.max #[[MAP]]()[%[[I1]], %[[I0]]]
1001  %2 = affine.max affine_map<()[s0, s1, s2] -> (s0, s1, s2 * 2)> ()[%0, %0, %i1]
1002  return %2: index
1003}
1004
1005// -----
1006
1007// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 16, s0 * 8)>
1008// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (s2 + 16, s2 * 8, s1 * 2, s0 + 1)>
1009
1010// CHECK: func @merge_mixed_uses_of_affine_max_ops
1011// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
1012func.func @merge_mixed_uses_of_affine_max_ops(%i0: index, %i1: index) -> index {
1013  // CHECK: %[[AFFINE:.+]] = affine.max #[[MAP0]]()[%[[I0]]]
1014  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
1015  // %0 is bound to a symbol that is both a standalone expression and a part
1016  // of other expressions.
1017  // CHECK: affine.max #[[MAP1]]()[%[[AFFINE]], %[[I1]], %[[I0]]]
1018  %2 = affine.max affine_map<()[s0, s1, s2] -> (s0, s1 + 1, s2 * 2)> ()[%0, %0, %i1]
1019  return %2: index
1020}
1021
1022// -----
1023
1024// CHECK-LABEL: func @dont_merge_affine_max_if_not_single_dim
1025func.func @dont_merge_affine_max_if_not_single_dim(%i0: index, %i1: index, %i2: index) -> index {
1026  // CHECK-COUNT-2: affine.max
1027  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
1028  %1 = affine.max affine_map<(d0)[s0] -> (s0 + 4, 7 + d0)> (%0)[%i2]
1029  return %1: index
1030}
1031
1032// -----
1033
1034// CHECK-LABEL: func @dont_merge_affine_max_if_not_single_sym
1035func.func @dont_merge_affine_max_if_not_single_sym(%i0: index, %i1: index, %i2: index) -> index {
1036  // CHECK-COUNT-2: affine.max
1037  %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
1038  %1 = affine.max affine_map<()[s0, s1] -> (s0 + 4, 7 + s1)> ()[%0, %i2]
1039  return %1: index
1040}
1041
1042// -----
1043
1044// Ensure bounding maps of affine.for are composed.
1045
1046// CHECK-DAG: #[[$MAP0]] = affine_map<()[s0] -> (s0 - 2)>
1047// CHECK-DAG: #[[$MAP1]] = affine_map<()[s0] -> (s0 + 2)>
1048
1049// CHECK-LABEL: func @compose_affine_for_bounds
1050// CHECK-SAME:   %[[N:.*]]: index)
1051// CHECK: affine.for %{{.*}} = #[[$MAP0]]()[%[[N]]] to #[[$MAP1]]()[%[[N]]] {
1052
1053func.func @compose_affine_for_bounds(%N: index) {
1054  %u = affine.apply affine_map<(d0) -> (d0 + 2)>(%N)
1055  %l = affine.apply affine_map<(d0) -> (d0 - 2)>(%N)
1056  affine.for %i = %l to %u {
1057    "foo"() : () -> ()
1058  }
1059  return
1060}
1061
1062// -----
1063
1064// Compose maps into affine.vector_load / affine.vector_store
1065
1066// CHECK-LABEL: func @compose_into_affine_vector_load_vector_store
1067// CHECK: affine.for %[[IV:.*]] = 0 to 1024
1068// CHECK-NEXT: affine.vector_load %{{.*}}[%[[IV]] + 1]
1069// CHECK-NEXT: affine.vector_store %{{.*}}, %{{.*}}[%[[IV]] + 1]
1070// CHECK-NEXT: affine.vector_load %{{.*}}[%[[IV]]]
1071func.func @compose_into_affine_vector_load_vector_store(%A : memref<1024xf32>, %u : index) {
1072  affine.for %i = 0 to 1024 {
1073    // Make sure the unused operand (%u below) gets dropped as well.
1074    %idx = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i, %u)
1075    %0 = affine.vector_load %A[%idx] : memref<1024xf32>, vector<8xf32>
1076    affine.vector_store %0, %A[%idx] : memref<1024xf32>, vector<8xf32>
1077
1078    // Map remains the same, but operand changes on composition.
1079    %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
1080    %1 = affine.vector_load %A[%copy] : memref<1024xf32>, vector<8xf32>
1081    "prevent.dce"(%1) : (vector<8xf32>) -> ()
1082  }
1083  return
1084}
1085
1086// -----
1087
1088// CHECK-LABEL: func @no_fold_of_store
1089//  CHECK:   %[[cst:.+]] = memref.cast %arg
1090//  CHECK:   affine.store %[[cst]]
1091func.func @no_fold_of_store(%arg : memref<32xi8>, %holder: memref<memref<?xi8>>) {
1092  %0 = memref.cast %arg : memref<32xi8> to memref<?xi8>
1093  affine.store %0, %holder[] : memref<memref<?xi8>>
1094  return
1095}
1096
1097// -----
1098
1099// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 + 16)>
1100// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 * 4)>
1101
1102// CHECK: func @canonicalize_single_min_max
1103// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
1104func.func @canonicalize_single_min_max(%i0: index, %i1: index) -> (index, index) {
1105  // CHECK-NOT: affine.min
1106  // CHECK-NEXT: affine.apply #[[$MAP0]]()[%[[I0]]]
1107  %0 = affine.min affine_map<()[s0] -> (s0 + 16)> ()[%i0]
1108
1109  // CHECK-NOT: affine.max
1110  // CHECK-NEXT: affine.apply #[[$MAP1]]()[%[[I1]]]
1111  %1 = affine.min affine_map<()[s0] -> (s0 * 4)> ()[%i1]
1112
1113  return %0, %1: index, index
1114}
1115
1116// -----
1117
1118// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (32, s1 + 16, s0 + s1)>
1119
1120// CHECK-LABEL: func @canonicalize_multi_min_max
1121// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
1122func.func @canonicalize_multi_min_max(%i0: index, %i1: index) -> (index, index) {
1123  // CHECK-NEXT: affine.min #[[$MAP]]()[%[[I0]], %[[I1]]]
1124  %0 = affine.min affine_map<()[s0, s1] -> (s0 + s1, s1 + 16, 32)> ()[%i0, %i1]
1125
1126  // CHECK-NEXT: affine.max #[[$MAP]]()[%[[I0]], %[[I1]]]
1127  %1 = affine.max affine_map<()[s0, s1] -> (s0 + s1, 32, s1 + 16)> ()[%i0, %i1]
1128
1129  return %0, %1: index, index
1130}
1131
1132// -----
1133
1134module {
1135  memref.global "private" constant @__constant_1x5x1xf32 : memref<1x5x1xf32> = dense<[[[6.250000e-02], [2.500000e-01], [3.750000e-01], [2.500000e-01], [6.250000e-02]]]>
1136  memref.global "private" constant @__constant_32x64xf32 : memref<32x64xf32> = dense<0.000000e+00>
1137  // CHECK-LABEL: func @fold_const_init_global_memref
1138  func.func @fold_const_init_global_memref() -> (f32, f32) {
1139    %m = memref.get_global @__constant_1x5x1xf32 : memref<1x5x1xf32>
1140    %v0 = affine.load %m[0, 0, 0] : memref<1x5x1xf32>
1141    %v1 = affine.load %m[0, 1, 0] : memref<1x5x1xf32>
1142    return %v0, %v1 : f32, f32
1143    // CHECK-DAG: %[[C0:.*]] = arith.constant 6.250000e-02 : f32
1144    // CHECK-DAG: %[[C1:.*]] = arith.constant 2.500000e-01 : f32
1145    // CHECK-NEXT: return %[[C0]], %[[C1]]
1146  }
1147
1148  // CHECK-LABEL: func @fold_const_splat_global
1149  func.func @fold_const_splat_global() -> memref<32x64xf32> {
1150    // CHECK-NEXT: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
1151    %m = memref.get_global @__constant_32x64xf32 : memref<32x64xf32>
1152    %s = memref.alloc() : memref<32x64xf32>
1153    affine.for %i = 0 to 32 {
1154      affine.for %j = 0 to 64 {
1155        %v = affine.load %m[%i, %j] : memref<32x64xf32>
1156        affine.store %v, %s[%i, %j] : memref<32x64xf32>
1157        // CHECK: affine.store %[[CST]], %{{.*}}
1158      }
1159    }
1160    return %s: memref<32x64xf32>
1161  }
1162}
1163
1164// -----
1165
1166// Simplification of maps exploiting operand info.
1167
1168// CHECK: #[[$MAP_SIMPLER:.*]] = affine_map<(d0, d1) -> (((d0 + d1) mod 458313) floordiv 227)>
1169
1170// CHECK-LABEL: func @simplify_with_operands
1171func.func @simplify_with_operands(%N: index, %A: memref<?x32xf32>) {
1172  // CHECK-NEXT: affine.for %[[I:.*]] = 0 to %{{.*}}
1173  affine.for %i = 0 to %N step 32 {
1174    // CHECK-NEXT: affine.for %[[II:.*]] = 0 to 32
1175    affine.for %ii = 0 to 32 {
1176      // %ii is less than 32 and %i divides 32.
1177      // CHECK: affine.load %{{.*}}[0, 0]
1178      %x = affine.load %A[%ii floordiv 32, %i mod 32] : memref<?x32xf32>
1179      "test.foo"(%x) : (f32) -> ()
1180
1181      // %i is aligned at 32 boundary and %ii < 32.
1182      // CHECK: affine.load %{{.*}}[%[[I]] floordiv 32, %[[II]] mod 16]
1183      %a = affine.load %A[(%i + %ii) floordiv 32, (%i + %ii) mod 16] : memref<?x32xf32>
1184      "test.foo"(%a) : (f32) -> ()
1185      // CHECK: affine.load %{{.*}}[%[[I]] floordiv 64, (%[[I]] + %[[II]]) mod 64]
1186      %b = affine.load %A[(%i + %ii) floordiv 64, (%i + %ii) mod 64] : memref<?x32xf32>
1187      "test.foo"(%b) : (f32) -> ()
1188      // CHECK: affine.load %{{.*}}[(%[[I]] + %[[II]]) floordiv 16, %[[II]] mod 16]
1189      %c = affine.load %A[(%i + %ii) floordiv 16, (%i + %ii) mod 16] : memref<?x32xf32>
1190      "test.foo"(%c) : (f32) -> ()
1191    }
1192  }
1193
1194  // Should not simplify.
1195  affine.for %i = -1 to 32 {
1196    // CHECK: affine.load %{{.*}}[%{{.*}} floordiv {{.*}}, %{{.*}} mod {{.*}}] :
1197    %x = affine.load %A[%i floordiv 32, %i mod 32] : memref<?x32xf32>
1198    "test.foo"(%x) : (f32) -> ()
1199  }
1200
1201  affine.for %arg0 = 0 to %N step 128 {
1202    affine.for %arg4 = 0 to 32 step 32 {
1203      affine.for %arg5 = 0 to 128 {
1204        // CHECK: affine.apply #[[$MAP_SIMPLER]]
1205        %x = affine.apply affine_map<(d0, d1, d2) -> (((d0 + d2) mod 458313) floordiv 227 + d1 floordiv 256)>(%arg0, %arg4, %arg5)
1206        "test.foo"(%x) : (index) -> ()
1207      }
1208    }
1209  }
1210
1211  return
1212}
1213
1214// CHECK-LABEL: func @simplify_div_mod_with_operands
1215func.func @simplify_div_mod_with_operands(%N: index, %A: memref<64xf32>, %unknown: index) {
1216  // CHECK: affine.for %[[I:.*]] = 0 to 32
1217  %cst = arith.constant 1.0 : f32
1218  affine.for %i = 0 to 32 {
1219    // CHECK: affine.store %{{.*}}, %{{.*}}[0]
1220    affine.store %cst, %A[%i floordiv 32] : memref<64xf32>
1221    // CHECK: affine.store %{{.*}}, %{{.*}}[1]
1222    affine.store %cst, %A[(%i + 1) ceildiv 32] : memref<64xf32>
1223    // CHECK: affine.store %{{.*}}, %{{.*}}[%[[I]]]
1224    affine.store %cst, %A[%i mod 32] : memref<64xf32>
1225    // CHECK: affine.store %{{.*}}, %{{.*}}[0]
1226    affine.store %cst, %A[2 * %i floordiv 64] : memref<64xf32>
1227    // CHECK: affine.store %{{.*}}, %{{.*}}[0]
1228    affine.store %cst, %A[(%i mod 16) floordiv 16] : memref<64xf32>
1229
1230    // The ones below can't be simplified.
1231    affine.store %cst, %A[%i floordiv 16] : memref<64xf32>
1232    affine.store %cst, %A[%i mod 16] : memref<64xf32>
1233    affine.store %cst, %A[(%i mod 16) floordiv 15] : memref<64xf32>
1234    affine.store %cst, %A[%i mod 31] : memref<64xf32>
1235    // CHECK:      affine.store %{{.*}}, %{{.*}}[%{{.*}} floordiv 16] : memref<64xf32>
1236    // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 16] : memref<64xf32>
1237    // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[(%{{.*}} mod 16) floordiv 15] : memref<64xf32>
1238    // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 31] : memref<64xf32>
1239  }
1240
1241  affine.for %i = -8 to 32 {
1242    // Can't be simplified.
1243    // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} floordiv 32] : memref<64xf32>
1244    affine.store %cst, %A[%i floordiv 32] : memref<64xf32>
1245    // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 32] : memref<64xf32>
1246    affine.store %cst, %A[%i mod 32] : memref<64xf32>
1247    // floordiv rounds toward -inf; (%i - 96) floordiv 64 will be -2.
1248    // CHECK: affine.store %{{.*}}, %{{.*}}[0] : memref<64xf32>
1249    affine.store %cst, %A[2 + (%i - 96) floordiv 64] : memref<64xf32>
1250  }
1251
1252  // CHECK: affine.for %[[II:.*]] = 8 to 16
1253  affine.for %i = 8 to 16 {
1254    // CHECK: affine.store %{{.*}}, %{{.*}}[1] : memref<64xf32>
1255    affine.store %cst, %A[%i floordiv 8] : memref<64xf32>
1256    // CHECK: affine.store %{{.*}}, %{{.*}}[2] : memref<64xf32>
1257    affine.store %cst, %A[(%i + 1) ceildiv 8] : memref<64xf32>
1258    // CHECK: affine.store %{{.*}}, %{{.*}}[%[[II]] mod 8] : memref<64xf32>
1259    affine.store %cst, %A[%i mod 8] : memref<64xf32>
1260    // CHECK: affine.store %{{.*}}, %{{.*}}[%[[II]]] : memref<64xf32>
1261    affine.store %cst, %A[%i mod 32] : memref<64xf32>
1262    // Upper bound on the mod 32 expression will be 15.
1263    // CHECK: affine.store %{{.*}}, %{{.*}}[0] : memref<64xf32>
1264    affine.store %cst, %A[(%i mod 32) floordiv 16] : memref<64xf32>
1265    // Lower bound on the mod 16 expression will be 8.
1266    // CHECK: affine.store %{{.*}}, %{{.*}}[1] : memref<64xf32>
1267    affine.store %cst, %A[(%i mod 16) floordiv 8] : memref<64xf32>
1268    // CHECK: affine.store %{{.*}}, %{{.*}}[0] : memref<64xf32>
1269    affine.store %cst, %A[(%unknown mod 16) floordiv 16] : memref<64xf32>
1270  }
1271  return
1272}
1273
1274// -----
1275
1276#map0 = affine_map<(d0) -> (32, d0 * -32 + 32)>
1277#map1 = affine_map<(d0) -> (32, d0 * -32 + 64)>
1278#map3 = affine_map<(d0) -> (16, d0 * -16 + 32)>
1279
1280// CHECK-DAG: #[[$SIMPLE_MAP:.*]] = affine_map<()[s0] -> (3, s0)>
1281// CHECK-DAG: #[[$SIMPLE_MAP_MAX:.*]] = affine_map<()[s0] -> (5, s0)>
1282// CHECK-DAG: #[[$SIMPLIFIED_MAP:.*]] = affine_map<(d0, d1) -> (-9, d0 * 4 - d1 * 4)>
1283// CHECK-DAG: #[[$FLOORDIV:.*]] = affine_map<(d0) -> (d0 floordiv 2)>
1284
1285// CHECK-LABEL: func @simplify_min_max_bounds_simple
1286func.func @simplify_min_max_bounds_simple(%M: index) {
1287
1288  // CHECK-NEXT: affine.for %{{.*}} = 0 to min #[[$SIMPLE_MAP]]
1289  affine.for %i = 0 to min affine_map<(d0) -> (3, 5, d0)>(%M) {
1290    "test.foo"() : () -> ()
1291  }
1292
1293  // CHECK: affine.for %{{.*}} = 0 to min #[[$SIMPLE_MAP]]
1294  affine.for %i = 0 to min affine_map<(d0) -> (3, 3, d0)>(%M) {
1295    "test.foo"() : () -> ()
1296  }
1297
1298  // CHECK: affine.for %{{.*}} = max #[[$SIMPLE_MAP_MAX]]
1299  affine.for %i = max affine_map<(d0) -> (3, 5, d0)>(%M) to 10 {
1300    "test.foo"() : () -> ()
1301  }
1302
1303  // CHECK: affine.for %{{.*}} = max #[[$SIMPLE_MAP_MAX]]
1304  affine.for %i = max affine_map<(d0) -> (5, 5, d0)>(%M) to 10 {
1305    "test.foo"() : () -> ()
1306  }
1307
1308  return
1309}
1310
1311// CHECK-LABEL: func @simplify_bounds_tiled
1312func.func @simplify_bounds_tiled() {
1313  affine.for %arg5 = 0 to 1 {
1314    affine.for %arg6 = 0 to 2 {
1315      affine.for %arg8 = 0 to min #map0(%arg5) step 16 {
1316        affine.for %arg9 = 0 to min #map1(%arg6) step 16 {
1317          affine.for %arg10 = 0 to 2 {
1318            affine.for %arg12 = 0 to min #map3(%arg10) step 16 {
1319              "test.foo"() : () -> ()
1320            }
1321          }
1322        }
1323      }
1324    }
1325  }
1326  // CHECK:      affine.for
1327  // CHECK-NEXT:   affine.for
1328  // CHECK-NEXT:     affine.for %{{.*}} = 0 to 32 step 16
1329  // CHECK-NEXT:       affine.for %{{.*}} = 0 to 32 step 16
1330  // CHECK-NEXT:         affine.for %{{.*}} = 0 to 2
1331  // CHECK-NEXT:           affine.for %{{.*}} = 0 to 16 step 16
1332
1333  return
1334}
1335
1336// CHECK-LABEL: func @simplify_min_max_multi_expr
1337func.func @simplify_min_max_multi_expr() {
1338  // Lower bound max.
1339  // CHECK: affine.for
1340  affine.for %i = 0 to 2 {
1341    // CHECK: affine.for %{{.*}} = 5 to
1342    affine.for %j = max affine_map<(d0) -> (5, 4 * d0)> (%i) to affine_map<(d0) -> (4 * d0 + 3)>(%i) {
1343      "test.foo"() : () -> ()
1344    }
1345  }
1346
1347  // Expressions with multiple operands.
1348  // CHECK: affine.for
1349  affine.for %i = 0 to 2 {
1350    // CHECK: affine.for
1351    affine.for %j = 0 to 4 {
1352      // The first upper bound expression will not be lower than -9. So, it's redundant.
1353      // CHECK-NEXT: affine.for %{{.*}} = -10 to -9
1354      affine.for %k = -10 to min affine_map<(d0, d1) -> (4 * d0 - 3 * d1, -9)>(%i, %j) {
1355        "test.foo"() : () -> ()
1356      }
1357    }
1358  }
1359
1360  // One expression is redundant but not the others.
1361  // CHECK: affine.for
1362  affine.for %i = 0 to 2 {
1363    // CHECK: affine.for
1364    affine.for %j = 0 to 4 {
1365      // The first upper bound expression will not be lower than -9. So, it's redundant.
1366      // CHECK-NEXT: affine.for %{{.*}} = -10 to min #[[$SIMPLIFIED_MAP]]
1367      affine.for %k = -10 to min affine_map<(d0, d1) -> (4 * d0 - 3 * d1, -9, 4 * d0 - 4 * d1)>(%i, %j) {
1368        "test.foo"() : () -> ()
1369      }
1370    }
1371  }
1372
1373  // CHECK: affine.for %{{.*}} = 0 to 1
1374  affine.for %i = 0 to 2 {
1375    affine.for %j = max affine_map<(d0) -> (d0 floordiv 2, 0)>(%i) to 1 {
1376      "test.foo"() : () -> ()
1377    }
1378  }
1379
1380  // The constant bound is redundant here.
1381  // CHECK: affine.for %{{.*}} = #[[$FLOORDIV]](%{{.*}} to 10
1382  affine.for %i = 0 to 8 {
1383    affine.for %j = max affine_map<(d0) -> (d0 floordiv 2, 0)>(%i) to 10 {
1384      "test.foo"() : () -> ()
1385    }
1386  }
1387
1388  return
1389}
1390
1391// CHECK-LABEL: func @no_simplify_min_max
1392func.func @no_simplify_min_max(%M: index) {
1393  // Negative test cases.
1394  // CHECK: affine.for
1395  affine.for %i = 0 to 4 {
1396    // CHECK-NEXT: affine.for %{{.*}} = 0 to min
1397    affine.for %j = 0 to min affine_map<(d0) -> (2 * d0, 2)>(%i) {
1398      "test.foo"() : () -> ()
1399    }
1400    // CHECK:      affine.for %{{.*}} = 0 to min {{.*}}(%{{.*}})[%{{.*}}]
1401    affine.for %j = 0 to min affine_map<(d0)[s0] -> (d0, s0)>(%i)[%M] {
1402      "test.foo"() : () -> ()
1403    }
1404  }
1405
1406  return
1407}
1408
1409// -----
1410
1411//           CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 * ((-s0 + 40961) ceildiv 512))>
1412// CHECK-BOTTOM-UP: #[[$map:.*]] = affine_map<()[s0] -> (s0 * ((-s0 + 40961) ceildiv 512))>
1413//           CHECK-LABEL: func @regression_do_not_perform_invalid_replacements
1414// CHECK-BOTTOM-UP-LABEL: func @regression_do_not_perform_invalid_replacements
1415func.func @regression_do_not_perform_invalid_replacements(%arg0: index) {
1416  // Dim must be promoted to sym before combining both maps.
1417  //           CHECK: %[[apply:.*]] = affine.apply #[[$map]]()[%{{.*}}]
1418  // CHECK-BOTTOM-UP: %[[apply:.*]] = affine.apply #[[$map]]()[%{{.*}}]
1419  %0 = affine.apply affine_map<(d0) -> (-d0 + 40961)>(%arg0)
1420  %1 = affine.apply affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 512))>(%arg0)[%0]
1421  //           CHECK: "test.foo"(%[[apply]])
1422  // CHECK-BOTTOM-UP: "test.foo"(%[[apply]])
1423  "test.foo"(%1) : (index) -> ()
1424  return
1425}
1426
1427// -----
1428// CHECK-LABEL: func @min.oneval(%arg0: index)
1429func.func @min.oneval(%arg0: index) -> index {
1430  %min = affine.min affine_map<()[s0] -> (s0)> ()[%arg0]
1431  // CHECK: return %arg0 : index
1432  return %min: index
1433}
1434
1435// -----
1436// CHECK-LABEL: func @max.oneval(%arg0: index)
1437func.func @max.oneval(%arg0: index) -> index {
1438  %max = affine.max affine_map<()[s0] -> (s0)> ()[%arg0]
1439  // CHECK: return %arg0 : index
1440  return %max: index
1441}
1442
1443// -----
1444
1445// CHECK-LABEL: func @mod_of_mod(
1446//       CHECK:   %[[c0:.*]] = arith.constant 0
1447//       CHECK:   return %[[c0]], %[[c0]]
1448func.func @mod_of_mod(%lb: index, %ub: index, %step: index) -> (index, index) {
1449  // Simplify: (ub - ub % step) % step == 0
1450  %0 = affine.apply affine_map<()[s0, s1] -> ((s0 - (s0 mod s1)) mod s1)> ()[%ub, %step]
1451  // Simplify: (ub - (ub - lb) % step - lb) % step == 0
1452  %1 = affine.apply affine_map<()[s0, s1, s2] -> ((s0 - ((s0 - s2) mod s1) - s2) mod s1)> ()[%ub, %step, %lb]
1453  return %0, %1 : index, index
1454}
1455
1456// -----
1457
1458// CHECK-LABEL:  func.func @prefetch_canonicalize
1459// CHECK-SAME:   ([[PARAM_0_:%.+]]: memref<512xf32>) {
1460func.func @prefetch_canonicalize(%arg0: memref<512xf32>) -> () {
1461  // CHECK: affine.for [[I_0_:%.+]] = 0 to 8 {
1462  affine.for %arg3 = 0 to 8  {
1463    %1 = affine.apply affine_map<()[s0] -> (s0 * 64)>()[%arg3]
1464    // CHECK: affine.prefetch [[PARAM_0_]][symbol([[I_0_]]) * 64], read, locality<3>, data : memref<512xf32>
1465    affine.prefetch %arg0[%1], read, locality<3>, data : memref<512xf32>
1466  }
1467  return
1468}
1469
1470// -----
1471
1472// CHECK-LABEL: @delinearize_fold_constant
1473// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
1474// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
1475// CHECK-NOT: affine.delinearize_index
1476// CHECK: return %[[C1]], %[[C1]], %[[C2]]
1477func.func @delinearize_fold_constant() -> (index, index, index) {
1478  %c22 = arith.constant 22 : index
1479  %0:3 = affine.delinearize_index %c22 into (2, 3, 5) : index, index, index
1480  return %0#0, %0#1, %0#2 : index, index, index
1481}
1482
1483// -----
1484
1485// CHECK-LABEL: @delinearize_fold_negative_constant
1486// CHECK-DAG: %[[C_2:.+]] = arith.constant -2 : index
1487// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
1488// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
1489// CHECK-NOT: affine.delinearize_index
1490// CHECK: return %[[C_2]], %[[C1]], %[[C3]]
1491func.func @delinearize_fold_negative_constant() -> (index, index, index) {
1492  %c_22 = arith.constant -22 : index
1493  %0:3 = affine.delinearize_index %c_22 into (2, 3, 5) : index, index, index
1494  return %0#0, %0#1, %0#2 : index, index, index
1495}
1496
1497// -----
1498
1499// CHECK-LABEL: @delinearize_fold_negative_constant_no_outer_bound
1500// CHECK-DAG: %[[C_2:.+]] = arith.constant -2 : index
1501// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
1502// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
1503// CHECK-NOT: affine.delinearize_index
1504// CHECK: return %[[C_2]], %[[C1]], %[[C3]]
1505func.func @delinearize_fold_negative_constant_no_outer_bound() -> (index, index, index) {
1506  %c_22 = arith.constant -22 : index
1507  %0:3 = affine.delinearize_index %c_22 into (3, 5) : index, index, index
1508  return %0#0, %0#1, %0#2 : index, index, index
1509}
1510
1511// -----
1512
1513// CHECK-LABEL: @delinearize_dont_fold_constant_dynamic_basis
1514// CHECK-DAG: %[[C22:.+]] = arith.constant 22 : index
1515// CHECK: %[[RET:.+]]:3 = affine.delinearize_index %[[C22]]
1516// CHECK: return %[[RET]]#0, %[[RET]]#1, %[[RET]]#2
1517func.func @delinearize_dont_fold_constant_dynamic_basis(%arg0: index) -> (index, index, index) {
1518  %c22 = arith.constant 22 : index
1519  %0:3 = affine.delinearize_index %c22 into (2, %arg0, 5) : index, index, index
1520  return %0#0, %0#1, %0#2 : index, index, index
1521}
1522
1523// -----
1524
1525func.func @drop_unit_basis_in_delinearize(%arg0 : index, %arg1 : index, %arg2 : index) ->
1526    (index, index, index, index, index, index) {
1527  %c1 = arith.constant 1 : index
1528  %0:6 = affine.delinearize_index %arg0 into (1, %arg1, 1, 1, %arg2, %c1)
1529      : index, index, index, index, index, index
1530  return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 : index, index, index, index, index, index
1531}
1532// CHECK-LABEL: func @drop_unit_basis_in_delinearize(
1533//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1534//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1535//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1536//   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
1537//   CHECK-DAG:   %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], %[[ARG2]])
1538//       CHECK:   return %[[C0]], %[[DELINEARIZE]]#0, %[[C0]], %[[C0]], %[[DELINEARIZE]]#1, %[[C0]]
1539
1540// -----
1541
1542func.func @drop_unit_basis_in_delinearize_no_outer_bound(%arg0 : index, %arg1 : index, %arg2 : index) ->
1543    (index, index, index, index, index, index) {
1544  %c1 = arith.constant 1 : index
1545  %0:6 = affine.delinearize_index %arg0 into (%arg1, 1, 1, %arg2, %c1)
1546      : index, index, index, index, index, index
1547  return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 : index, index, index, index, index, index
1548}
1549// CHECK-LABEL: func @drop_unit_basis_in_delinearize_no_outer_bound(
1550//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1551//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1552//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1553//   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
1554//   CHECK-DAG:   %[[DELINEARIZE:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], %[[ARG2]])
1555//       CHECK:   return %[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[C0]], %[[C0]], %[[DELINEARIZE]]#2, %[[C0]]
1556
1557// -----
1558
1559func.func @drop_all_unit_bases(%arg0 : index) -> (index, index) {
1560  %0:2 = affine.delinearize_index %arg0 into (1, 1) : index, index
1561  return %0#0, %0#1 : index, index
1562}
1563// CHECK-LABEL: func @drop_all_unit_bases(
1564//  CHECK-SAME:     %[[ARG0:.+]]: index)
1565//   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
1566//   CHECK-NOT:   affine.delinearize_index
1567//       CHECK:   return %[[C0]], %[[C0]]
1568
1569// -----
1570
1571func.func @drop_all_unit_bases_no_outer_bound(%arg0 : index) -> (index, index, index) {
1572  %0:3 = affine.delinearize_index %arg0 into (1, 1) : index, index, index
1573  return %0#0, %0#1, %0#2 : index, index, index
1574}
1575// CHECK-LABEL: func @drop_all_unit_bases_no_outer_bound(
1576//  CHECK-SAME:     %[[ARG0:.+]]: index)
1577//   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
1578//   CHECK-NOT:   affine.delinearize_index
1579//       CHECK:   return %[[ARG0]], %[[C0]], %[[C0]]
1580
1581// -----
1582
1583func.func @drop_single_loop_delinearize(%arg0 : index, %arg1 : index) -> index {
1584  %c0 = arith.constant 0 : index
1585  %c1 = arith.constant 1 : index
1586  %2 = scf.for %iv = %c0 to %arg1 step %c1 iter_args(%arg2 = %c0) -> index {
1587    %0 = affine.delinearize_index %iv into (%arg1) : index
1588    %1 = "some_use"(%arg2, %0) : (index, index) -> (index)
1589    scf.yield %1 : index
1590  }
1591  return %2 : index
1592}
1593// CHECK-LABEL: func @drop_single_loop_delinearize(
1594//  CHECK-SAME:     %[[ARG0:.+]]: index)
1595//       CHECK:   scf.for %[[IV:[a-zA-Z0-9]+]] =
1596//   CHECK-NOT:     affine.delinearize_index
1597//       CHECK:     "some_use"(%{{.+}}, %[[IV]])
1598
1599// -----
1600
1601// CHECK-LABEL: func @delinearize_non_induction_variable
1602// CHECK-NOT: affine.delinearize
1603func.func @delinearize_non_induction_variable(%arg0: memref<?xi32>, %i : index, %t0 : index, %t1 : index, %t2 : index) -> index {
1604  %1 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 + s0 + s1 * 64 + s2 * 128)>(%i)[%t0, %t1, %t2]
1605  %2 = affine.delinearize_index %1 into (1024) : index
1606  return %2 : index
1607}
1608
1609// -----
1610
1611// CHECK-LABEL: func @delinearize_non_loop_like
1612// CHECK-NOT: affine.delinearize
1613func.func @delinearize_non_loop_like(%arg0: memref<?xi32>, %i : index) -> index {
1614  %2 = affine.delinearize_index %i into (1024) : index
1615  return %2 : index
1616}
1617
1618// -----
1619
1620// CHECK-LABEL: func @delinearize_empty_basis
1621// CHECK-SAME: (%[[ARG0:.+]]: index)
1622// CHECK-NOT: affine.delinearize
1623// CHECK: return %[[ARG0]]
1624func.func @delinearize_empty_basis(%arg0: index) -> index {
1625  %0 = affine.delinearize_index %arg0 into () : index
1626  return %0 : index
1627}
1628
1629// -----
1630
1631// CHECK-LABEL: @linearize_fold_constants
1632// CHECK-DAG: %[[C22:.+]] = arith.constant 22 : index
1633// CHECK-NOT: affine.linearize
1634// CHECK: return %[[C22]]
1635func.func @linearize_fold_constants() -> index {
1636  %c2 = arith.constant 2 : index
1637  %c1 = arith.constant 1 : index
1638
1639  %ret = affine.linearize_index [%c1, %c1, %c2] by (2, 3, 5) : index
1640  return %ret : index
1641}
1642
1643// -----
1644
1645// CHECK-LABEL: @linearize_fold_constants_no_outer_bound
1646// CHECK-DAG: %[[C22:.+]] = arith.constant 22 : index
1647// CHECK-NOT: affine.linearize
1648// CHECK: return %[[C22]]
1649func.func @linearize_fold_constants_no_outer_bound() -> index {
1650  %c2 = arith.constant 2 : index
1651  %c1 = arith.constant 1 : index
1652
1653  %ret = affine.linearize_index [%c1, %c1, %c2] by (3, 5) : index
1654  return %ret : index
1655}
1656
1657// -----
1658
1659// CHECK-LABEL: @linearize_fold_empty_basis
1660// CHECK-SAME: (%[[ARG0:.+]]: index)
1661// CHECK-NOT: affine.linearize
1662// CHECK: return %[[ARG0]]
1663func.func @linearize_fold_empty_basis(%arg0: index) -> index {
1664  %ret = affine.linearize_index [%arg0] by () : index
1665  return %ret : index
1666}
1667
1668// -----
1669
1670// CHECK-LABEL: @linearize_fold_only_outer_bound
1671// CHECK-SAME: (%[[ARG0:.+]]: index)
1672// CHECK-NOT: affine.linearize
1673// CHECK: return %[[ARG0]]
1674func.func @linearize_fold_only_outer_bound(%arg0: index) -> index {
1675  %ret = affine.linearize_index [%arg0] by (2) : index
1676  return %ret : index
1677}
1678
1679// -----
1680
1681// CHECK-LABEL: @linearize_dont_fold_dynamic_basis
1682// CHECK: %[[RET:.+]] = affine.linearize_index
1683// CHECK: return %[[RET]]
1684func.func @linearize_dont_fold_dynamic_basis(%arg0: index) -> index {
1685  %c2 = arith.constant 2 : index
1686  %c1 = arith.constant 1 : index
1687
1688  %ret = affine.linearize_index [%c1, %c1, %c2] by (2, %arg0, 5) : index
1689  return %ret : index
1690}
1691
1692// -----
1693
1694// CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_exact(
1695//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1696//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1697//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
1698//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
1699//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index)
1700//       CHECK:     return %[[ARG0]], %[[ARG1]], %[[ARG2]]
1701func.func @cancel_delinearize_linearize_disjoint_exact(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
1702  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (%arg3, 4, %arg4) : index
1703  %1:3 = affine.delinearize_index %0 into (%arg3, 4, %arg4)
1704      : index, index, index
1705  return %1#0, %1#1, %1#2 : index, index, index
1706}
1707
1708// -----
1709
1710// CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_linearize_extra_bound(
1711//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1712//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1713//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
1714//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
1715//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index)
1716//       CHECK:     return %[[ARG0]], %[[ARG1]], %[[ARG2]]
1717func.func @cancel_delinearize_linearize_disjoint_linearize_extra_bound(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
1718  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (4, %arg4) : index
1719  %1:3 = affine.delinearize_index %0 into (4, %arg4)
1720      : index, index, index
1721  return %1#0, %1#1, %1#2 : index, index, index
1722}
1723
1724// -----
1725
1726// CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_delinearize_extra_bound(
1727//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1728//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1729//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
1730//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
1731//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index)
1732//       CHECK:     return %[[ARG0]], %[[ARG1]], %[[ARG2]]
1733func.func @cancel_delinearize_linearize_disjoint_delinearize_extra_bound(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
1734  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (4, %arg4) : index
1735  %1:3 = affine.delinearize_index %0 into (%arg3, 4, %arg4)
1736      : index, index, index
1737  return %1#0, %1#1, %1#2 : index, index, index
1738}
1739
1740// -----
1741
1742// CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_partial(
1743//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1744//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1745//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
1746//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
1747//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index)
1748//       CHECK:     %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]]] by (%[[ARG3]], 4) : index
1749//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[LIN]] into (8) : index, index
1750//       CHECK:     return %[[DELIN]]#0, %[[DELIN]]#1, %[[ARG2]]
1751func.func @cancel_delinearize_linearize_disjoint_partial(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
1752  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (%arg3, 4, %arg4) : index
1753  %1:3 = affine.delinearize_index %0 into (8, %arg4)
1754      : index, index, index
1755  return %1#0, %1#1, %1#2 : index, index, index
1756}
1757
1758// -----
1759
1760// Without `disjoint`, the cancelation isn't guaranteed to be the identity.
1761// CHECK-LABEL: func @no_cancel_delinearize_linearize_exact(
1762//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1763//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1764//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
1765//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
1766//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index)
1767//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG0]], %[[ARG1]], %[[ARG2]]] by (%[[ARG3]], 4, %[[ARG4]])
1768//       CHECK:     %[[DELIN:.+]]:3 = affine.delinearize_index %[[LIN]] into (%[[ARG3]], 4, %[[ARG4]])
1769//       CHECK:     return %[[DELIN]]#0, %[[DELIN]]#1, %[[DELIN]]#2
1770func.func @no_cancel_delinearize_linearize_exact(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
1771  %0 = affine.linearize_index [%arg0, %arg1, %arg2] by (%arg3, 4, %arg4) : index
1772  %1:3 = affine.delinearize_index %0 into (%arg3, 4, %arg4)
1773      : index, index, index
1774  return %1#0, %1#1, %1#2 : index, index, index
1775}
1776
1777// -----
1778
1779// These don't cancel because the delinearize and linearize have a different basis.
1780// CHECK-LABEL: func @no_cancel_delinearize_linearize_different_basis(
1781//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1782//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1783//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
1784//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
1785//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index)
1786//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG0]], %[[ARG1]], %[[ARG2]]] by (%[[ARG3]], 4, %[[ARG4]])
1787//       CHECK:     %[[DELIN:.+]]:3 = affine.delinearize_index %[[LIN]] into (%[[ARG3]], 8, %[[ARG4]])
1788//       CHECK:     return %[[DELIN]]#0, %[[DELIN]]#1, %[[DELIN]]#2
1789func.func @no_cancel_delinearize_linearize_different_basis(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
1790  %0 = affine.linearize_index [%arg0, %arg1, %arg2] by (%arg3, 4, %arg4) : index
1791  %1:3 = affine.delinearize_index %0 into (%arg3, 8, %arg4)
1792      : index, index, index
1793  return %1#0, %1#1, %1#2 : index, index, index
1794}
1795
1796// -----
1797
1798// CHECK-LABEL: func @split_delinearize_spanning_final_part
1799//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1800//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1801//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1802//       CHECK:     %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]]] by (2, 4)
1803//       CHECK:     %[[DELIN1:.+]]:2 = affine.delinearize_index %[[LIN]] into (2)
1804//       CHECK:     %[[DELIN2:.+]]:2 = affine.delinearize_index %[[ARG2]] into (8, 8)
1805//       CHECK:     return %[[DELIN1]]#0, %[[DELIN1]]#1, %[[DELIN2]]#0, %[[DELIN2]]#1
1806func.func @split_delinearize_spanning_final_part(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index, index) {
1807  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (2, 4, 64) : index
1808  %1:4 = affine.delinearize_index %0 into (2, 8, 8)
1809      : index, index, index, index
1810  return %1#0, %1#1, %1#2, %1#3 : index, index, index, index
1811}
1812
1813// -----
1814
1815// CHECK-LABEL: func @split_delinearize_spanning_final_part_and_cancel
1816//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1817//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1818//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1819//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG2]] into (8, 8)
1820//       CHECK:     return %[[ARG0]], %[[ARG1]], %[[DELIN]]#0, %[[DELIN]]#1
1821func.func @split_delinearize_spanning_final_part_and_cancel(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index, index) {
1822  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (2, 4, 64) : index
1823  %1:4 = affine.delinearize_index %0 into (2, 4, 8, 8)
1824      : index, index, index, index
1825  return %1#0, %1#1, %1#2, %1#3 : index, index, index, index
1826}
1827
1828// -----
1829
1830// The delinearize basis doesn't match the last basis element before
1831// overshooting it, don't simplify.
1832// CHECK-LABEL: func @dont_split_delinearize_overshooting_target
1833//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1834//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1835//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1836//       CHECK:     %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]], %[[ARG2]]] by (2, 4, 64)
1837//       CHECK:     %[[DELIN:.+]]:4 = affine.delinearize_index %[[LIN]] into (2, 16, 8)
1838//       CHECK:     return %[[DELIN]]#0, %[[DELIN]]#1, %[[DELIN]]#2, %[[DELIN]]#3
1839func.func @dont_split_delinearize_overshooting_target(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index, index) {
1840  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (2, 4, 64) : index
1841  %1:4 = affine.delinearize_index %0 into (2, 16, 8)
1842      : index, index, index, index
1843  return %1#0, %1#1, %1#2, %1#3 : index, index, index, index
1844}
1845
1846// -----
1847
1848// The delinearize basis doesn't fully multiply to the final basis element.
1849// CHECK-LABEL: func @dont_split_delinearize_undershooting_target
1850//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1851//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
1852//       CHECK:     %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]]] by (2, 64)
1853//       CHECK:     %[[DELIN:.+]]:3 = affine.delinearize_index %[[LIN]] into (4, 8)
1854//       CHECK:     return %[[DELIN]]#0, %[[DELIN]]#1
1855func.func @dont_split_delinearize_undershooting_target(%arg0: index, %arg1: index) -> (index, index, index) {
1856  %0 = affine.linearize_index disjoint [%arg0, %arg1] by (2, 64) : index
1857  %1:3 = affine.delinearize_index %0 into (4, 8)
1858      : index, index, index
1859  return %1#0, %1#1, %1#2 : index, index, index
1860}
1861
1862// -----
1863
1864// CHECK-LABEL: @linearize_unit_basis_disjoint
1865// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index)
1866// CHECK: %[[ret:.+]] = affine.linearize_index disjoint [%[[arg0]], %[[arg2]]] by (3, %[[arg3]]) : index
1867// CHECK: return %[[ret]]
1868func.func @linearize_unit_basis_disjoint(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
1869  %ret = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (3, 1, %arg3) : index
1870  return %ret : index
1871}
1872
1873// -----
1874
1875// CHECK-LABEL: @linearize_unit_basis_disjoint_no_outer_bound
1876// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index)
1877// CHECK: %[[ret:.+]] = affine.linearize_index disjoint [%[[arg0]], %[[arg2]]] by (%[[arg3]]) : index
1878// CHECK: return %[[ret]]
1879func.func @linearize_unit_basis_disjoint_no_outer_bound(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
1880  %ret = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (1, %arg3) : index
1881  return %ret : index
1882}
1883
1884// -----
1885
1886// CHECK-LABEL: @linearize_unit_basis_zero
1887// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index)
1888// CHECK: %[[ret:.+]] = affine.linearize_index [%[[arg0]], %[[arg1]]] by (3, %[[arg2]]) : index
1889// CHECK: return %[[ret]]
1890func.func @linearize_unit_basis_zero(%arg0: index, %arg1: index, %arg2: index) -> index {
1891  %c0 = arith.constant 0 : index
1892  %ret = affine.linearize_index [%arg0, %c0, %arg1] by (3, 1, %arg2) : index
1893  return %ret : index
1894}
1895
1896// -----
1897
1898// CHECK-LABEL: @linearize_all_zero_unit_basis
1899// CHECK: arith.constant 0 : index
1900// CHECK-NOT: affine.linearize_index
1901func.func @linearize_all_zero_unit_basis() -> index {
1902  %c0 = arith.constant 0 : index
1903  %ret = affine.linearize_index [%c0, %c0] by (1, 1) : index
1904  return %ret : index
1905}
1906
1907// -----
1908
1909// CHECK-LABEL: @linearize_one_element_basis
1910// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index)
1911// CHECK-NOT: affine.linearize_index
1912// CHECK: return %[[arg0]]
1913func.func @linearize_one_element_basis(%arg0: index, %arg1: index) -> index {
1914  %ret = affine.linearize_index [%arg0] by (%arg1) : index
1915  return %ret : index
1916}
1917
1918// -----
1919
1920// CHECK-LABEL: func @cancel_linearize_delinearize_exact(
1921//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1922//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1923//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1924//       CHECK:     return %[[ARG0]]
1925func.func @cancel_linearize_delinearize_exact(%arg0: index, %arg1: index, %arg2: index) -> index {
1926  %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index
1927  %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (%arg1, 4, %arg2) : index
1928  return %1 : index
1929}
1930
1931// -----
1932
1933// CHECK-LABEL: func @cancel_linearize_delinearize_linearize_extra_bound(
1934//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1935//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1936//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1937//       CHECK:     return %[[ARG0]]
1938func.func @cancel_linearize_delinearize_linearize_extra_bound(%arg0: index, %arg1: index, %arg2: index) -> index {
1939  %0:3 = affine.delinearize_index %arg0 into (4, %arg2) : index, index, index
1940  %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (%arg1, 4, %arg2) : index
1941  return %1 : index
1942}
1943
1944// -----
1945
1946// CHECK-LABEL: func @cancel_linearize_delinearize_delinearize_extra_bound(
1947//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1948//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
1949//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
1950//       CHECK:     return %[[ARG0]]
1951func.func @cancel_linearize_delinearize_delinearize_extra_bound(%arg0: index, %arg1: index, %arg2: index) -> index {
1952  %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index
1953  %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (4, %arg2) : index
1954  return %1 : index
1955}
1956
1957// -----
1958
1959// CHECK-LABEL: func @cancel_linearize_delinearize_head(
1960//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1961//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
1962//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (12, 8)
1963//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (12, 16)
1964//       CHECK:     return %[[LIN]]
1965func.func @cancel_linearize_delinearize_head(%arg0: index, %arg1: index) -> index {
1966  %0:3 = affine.delinearize_index %arg0 into (3, 4, 8) : index, index, index
1967  %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (3, 4, 16) : index
1968  return %1 : index
1969}
1970
1971// -----
1972
1973// CHECK-LABEL: func @cancel_linearize_delinearize_head_delinearize_unbounded(
1974//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1975//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
1976//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (12, 8)
1977//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (12, 16)
1978//       CHECK:     return %[[LIN]]
1979func.func @cancel_linearize_delinearize_head_delinearize_unbounded(%arg0: index, %arg1: index) -> index {
1980  %0:3 = affine.delinearize_index %arg0 into (4, 8) : index, index, index
1981  %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (3, 4, 16) : index
1982  return %1 : index
1983}
1984
1985// -----
1986
1987// CHECK-LABEL: func @cancel_linearize_delinearize_head_linearize_unbounded(
1988//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
1989//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
1990//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8)
1991//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (16)
1992//       CHECK:     return %[[LIN]]
1993func.func @cancel_linearize_delinearize_head_linearize_unbounded(%arg0: index, %arg1: index) -> index {
1994  %0:3 = affine.delinearize_index %arg0 into (3, 4, 8) : index, index, index
1995  %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (4, 16) : index
1996  return %1 : index
1997}
1998
1999// -----
2000
2001// CHECK-LABEL: func @cancel_linearize_delinearize_head_both_unbounded(
2002//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2003//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
2004//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8)
2005//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (16)
2006//       CHECK:     return %[[LIN]]
2007func.func @cancel_linearize_delinearize_head_both_unbounded(%arg0: index, %arg1: index) -> index {
2008  %0:3 = affine.delinearize_index %arg0 into (4, 8) : index, index, index
2009  %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (4, 16) : index
2010  return %1 : index
2011}
2012
2013// -----
2014
2015// CHECK-LABEL: func @cancel_linearize_delinearize_tail(
2016//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2017//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
2018//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (3, 32)
2019//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[DELIN]]#1] by (5, 32)
2020//       CHECK:     return %[[LIN]]
2021func.func @cancel_linearize_delinearize_tail(%arg0: index, %arg1: index) -> index {
2022  %0:3 = affine.delinearize_index %arg0 into (3, 4, 8) : index, index, index
2023  %1 = affine.linearize_index [%arg1, %0#1, %0#2] by (5, 4, 8) : index
2024  return %1 : index
2025}
2026
2027// -----
2028
2029// CHECK-LABEL: func @cancel_linearize_delinearize_middle_exact(
2030//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2031//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2032//  CHECK-SAME:     %[[ARG2:[a-zA-z0-9]+]]: index)
2033//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[ARG0]], %[[ARG2]]] by (9, 30, 7)
2034//       CHECK:     return %[[LIN]]
2035func.func @cancel_linearize_delinearize_middle_exact(%arg0: index, %arg1: index, %arg2: index) -> index {
2036  %0:3 = affine.delinearize_index %arg0 into (2, 3, 5) : index, index, index
2037  %1 = affine.linearize_index [%arg1, %0#0, %0#1, %0#2, %arg2] by (9, 2, 3, 5, 7) : index
2038  return %1 : index
2039}
2040
2041// -----
2042
2043// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> ((s0 * s1) * 16)>
2044
2045// CHECK-LABEL: func @cancel_linearize_delinearize_middle_exact_dynamic_basis(
2046//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2047//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2048//  CHECK-SAME:     %[[ARG2:[a-zA-z0-9]+]]: index)
2049//       CHECK:     %[[C1:.+]] = arith.constant 1 : index
2050//       CHECK:     %[[SIZEPROD:.+]] = affine.apply #[[$MAP]]()[%[[ARG1]], %[[ARG2]]]
2051//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[C1]], %[[ARG0]], %[[C1]]] by (3, %[[SIZEPROD]], 4)
2052//       CHECK:     return %[[LIN]]
2053func.func @cancel_linearize_delinearize_middle_exact_dynamic_basis(%arg0: index, %arg1: index, %arg2: index) -> index {
2054  %c1 = arith.constant 1 : index
2055  %0:4 = affine.delinearize_index %arg0 into (2, %arg1, %arg2, 8) : index, index, index, index
2056  %1 = affine.linearize_index [%c1, %0#0, %0#1, %0#2, %0#3, %c1] by (3, 2, %arg1, %arg2, 8, 4) : index
2057  return %1 : index
2058}
2059
2060// -----
2061
2062// CHECK-LABEL: func @cancel_linearize_delinearize_middle_exact_delinearize_unbounded_disjoint(
2063//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2064//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2065//  CHECK-SAME:     %[[ARG2:[a-zA-z0-9]+]]: index)
2066//       CHECK:     %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG1]], %[[ARG0]], %[[ARG2]]] by (9, 30, 7)
2067//       CHECK:     return %[[LIN]]
2068func.func @cancel_linearize_delinearize_middle_exact_delinearize_unbounded_disjoint(%arg0: index, %arg1: index, %arg2: index) -> index {
2069  %0:3 = affine.delinearize_index %arg0 into (3, 5) : index, index, index
2070  %1 = affine.linearize_index disjoint [%arg1, %0#0, %0#1, %0#2, %arg2] by (9, 2, 3, 5, 7) : index
2071  return %1 : index
2072}
2073
2074// -----
2075
2076// Unlike in the test above, the linerize indices aren't asserted to be disjoint, so
2077// we can't know if the `2` from the basis is a correct bound.
2078// CHECK-LABEL: func @dont_cancel_linearize_delinearize_middle_exact_delinearize_unbounded(
2079//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2080//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2081//  CHECK-SAME:     %[[ARG2:[a-zA-z0-9]+]]: index)
2082//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (3)
2083//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[DELIN]]#0, %[[DELIN]]#1, %[[ARG2]]] by (9, 2, 3, 7)
2084//       CHECK:     return %[[LIN]]
2085
2086func.func @dont_cancel_linearize_delinearize_middle_exact_delinearize_unbounded(%arg0: index, %arg1: index, %arg2: index) -> index {
2087  %0:2 = affine.delinearize_index %arg0 into (3) : index, index
2088  %1 = affine.linearize_index [%arg1, %0#0, %0#1, %arg2] by (9, 2, 3, 7) : index
2089  return %1 : index
2090}
2091
2092// -----
2093
2094// The presence of a `disjoint` here tells us that the "unbounded" term on the
2095// delinearization can't have been above 2.
2096// CHECK-LABEL: func @cancel_linearize_delinearize_middle_delinearize_unbounded_disjoint_implied_bound(
2097//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2098//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2099//  CHECK-SAME:     %[[ARG2:[a-zA-z0-9]+]]: index)
2100//       CHECK:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (6, 5)
2101//       CHECK:     %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG1]], %[[DELIN]]#0, %[[ARG2]]] by (9, 6, 7)
2102//       CHECK:     return %[[LIN]]
2103func.func @cancel_linearize_delinearize_middle_delinearize_unbounded_disjoint_implied_bound(%arg0: index, %arg1: index, %arg2: index) -> index {
2104  %0:3 = affine.delinearize_index %arg0 into (3, 5) : index, index, index
2105  %1 = affine.linearize_index disjoint [%arg1, %0#0, %0#1, %arg2] by (9, 2, 3, 7) : index
2106  return %1 : index
2107}
2108
2109// -----
2110
2111// CHECK-LABEL: func @cancel_linearize_delinearize_multiple_matches(
2112//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2113//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
2114//       CHECK:     %[[C0:.+]] = arith.constant 0
2115//       CHECK:     %[[DELIN:.+]]:4 = affine.delinearize_index %[[ARG0]] into (4, 16, 4, 64)
2116//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[DELIN]]#1, %[[C0]], %[[DELIN]]#3] by (4, 16, 4, 64)
2117//       CHECK:     return %[[LIN]]
2118func.func @cancel_linearize_delinearize_multiple_matches(%arg0: index, %arg1: index) -> index {
2119  %c0 = arith.constant 0 : index
2120  %0:7 = affine.delinearize_index %arg0 into (4, 4, 4, 4, 4, 4, 4) : index, index, index, index, index, index, index
2121  %1 = affine.linearize_index [%arg1, %0#1, %0#2, %c0, %0#4, %0#5, %0#6] by (4, 4, 4, 4, 4, 4, 4) : index
2122  return %1 : index
2123}
2124
2125// -----
2126
2127// CHECK-LABEL: func @cancel_linearize_delinearize_multiple_delinearizes(
2128//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2129//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
2130//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[ARG0]], %[[ARG1]]] by (32, 32)
2131//       CHECK:     return %[[LIN]]
2132func.func @cancel_linearize_delinearize_multiple_delinearizes(%arg0: index, %arg1: index) -> index {
2133  %0:2 = affine.delinearize_index %arg0 into (4, 8) : index, index
2134  %1:2 = affine.delinearize_index %arg1 into (2, 16) : index, index
2135  %2 = affine.linearize_index [%0#0, %0#1, %1#0, %1#1] by (4, 8, 2, 16) : index
2136  return %2 : index
2137}
2138
2139// -----
2140
2141// Don't cancel because the values from the delinearize aren't used in order
2142// CHECK-LABEL: func @no_cancel_linearize_delinearize_permuted(
2143//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2144//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2145//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
2146//       CHECK:     %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[ARG2]])
2147//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#2, %[[DELIN]]#1] by (%[[ARG1]], %[[ARG2]], 4)
2148//       CHECK:     return %[[LIN]]
2149func.func @no_cancel_linearize_delinearize_permuted(%arg0: index, %arg1: index, %arg2: index) -> index {
2150  %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index
2151  %1 = affine.linearize_index [%0#0, %0#2, %0#1] by (%arg1, %arg2, 4) : index
2152  return %1 : index
2153}
2154
2155// -----
2156
2157// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 3)>
2158// But these cancel because they're a contiguous segment
2159// CHECK-LABEL: func @partial_cancel_linearize_delinearize_not_fully_permuted(
2160//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2161//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2162//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
2163//       CHECK:     %[[SIZEPROD:.+]] = affine.apply #[[$MAP]]()[%[[ARG2]]]
2164//       CHECK:     %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[SIZEPROD]])
2165//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#2, %[[DELIN]]#1] by (%[[ARG1]], %[[SIZEPROD]], 4)
2166//       CHECK:     return %[[LIN]]
2167func.func @partial_cancel_linearize_delinearize_not_fully_permuted(%arg0: index, %arg1: index, %arg2: index) -> index {
2168  %0:4 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2, 3) : index, index, index, index
2169  %1 = affine.linearize_index [%0#0, %0#2, %0#3, %0#1] by (%arg1, %arg2, 3, 4) : index
2170  return %1 : index
2171}
2172
2173// -----
2174
2175// Ensure we don't get SSA errors when creating new `affine.delinearize` operations.
2176// CHECK-LABEL: func @cancel_linearize_delinearize_placement
2177// CHECK-SAME: (%[[ARG0:.+]]: index)
2178// CHECK: %[[C0:.+]] = arith.constant 0 : index
2179// CHECK: %[[NEW_DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8, 32) : index, index
2180// CHECK-NEXT: %[[DELIN_PART:.+]]:2 = affine.delinearize_index %[[NEW_DELIN]]#1 into (8, 4) : index, index
2181// CHECK-NEXT: %[[L1:.+]] = affine.linearize_index disjoint [%[[DELIN_PART]]#1, %[[NEW_DELIN]]#0, %[[C0]], %[[C0]]] by (4, 8, 4, 8)
2182// CHECK-NEXT: %[[L2:.+]] = affine.linearize_index disjoint [%[[NEW_DELIN]]#1, %[[C0]], %[[C0]]] by (32, 8, 4)
2183// CHECK-NEXT: %[[L3:.+]] = affine.linearize_index disjoint [%[[DELIN_PART]]#0, %[[NEW_DELIN]]#0, %[[C0]], %[[C0]]] by (8, 8, 4, 4)
2184// CHECK-NEXT: return %[[L1]], %[[L2]], %[[L3]]
2185func.func @cancel_linearize_delinearize_placement(%arg0: index) -> (index, index, index) {
2186  %c0 = arith.constant 0 : index
2187  %0:3 = affine.delinearize_index %arg0 into (8, 8, 4) : index, index, index
2188  %1 = affine.linearize_index disjoint [%0#2, %0#0, %c0, %c0] by (4, 8, 4, 8) : index
2189  %2 = affine.linearize_index disjoint [%0#1, %0#2, %c0, %c0] by (8, 4, 8, 4) : index
2190  %3 = affine.linearize_index disjoint [%0#1, %0#0, %c0, %c0] by (8, 8, 4, 4) : index
2191  return %1, %2, %3 : index, index, index
2192}
2193
2194// -----
2195
2196// Won't cancel because the linearize and delinearize are using a different basis
2197// CHECK-LABEL: func @no_cancel_linearize_delinearize_different_basis(
2198//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2199//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
2200//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index)
2201//       CHECK:     %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[ARG2]])
2202//       CHECK:     %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#1, %[[DELIN]]#2] by (%[[ARG1]], 8, %[[ARG2]])
2203//       CHECK:     return %[[LIN]]
2204func.func @no_cancel_linearize_delinearize_different_basis(%arg0: index, %arg1: index, %arg2: index) -> index {
2205  %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index
2206  %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (%arg1, 8, %arg2) : index
2207  return %1 : index
2208}
2209
2210// -----
2211
2212// CHECK-LABEL: func @affine_leading_zero(
2213//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2214//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
2215//       CHECK:     %[[RET:.+]] = affine.linearize_index [%[[ARG0]], %[[ARG1]]] by (3, 5)
2216//       CHECK:     return %[[RET]]
2217func.func @affine_leading_zero(%arg0: index, %arg1: index) -> index {
2218  %c0 = arith.constant 0 : index
2219  %ret = affine.linearize_index [%c0, %arg0, %arg1] by (2, 3, 5) : index
2220  return %ret : index
2221}
2222
2223// -----
2224
2225// CHECK-LABEL: func @affine_leading_zero_no_outer_bound(
2226//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
2227//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
2228//       CHECK:     %[[RET:.+]] = affine.linearize_index [%[[ARG0]], %[[ARG1]]] by (3, 5)
2229//       CHECK:     return %[[RET]]
2230func.func @affine_leading_zero_no_outer_bound(%arg0: index, %arg1: index) -> index {
2231  %c0 = arith.constant 0 : index
2232  %ret = affine.linearize_index [%c0, %arg0, %arg1] by (3, 5) : index
2233  return %ret : index
2234}
2235
2236// -----
2237
2238// CHECK-LABEL: @cst_value_to_cst_attr_basis_delinearize_index
2239// CHECK-SAME:    (%[[ARG0:.*]]: index)
2240// CHECK:         %[[RET:.*]]:3 = affine.delinearize_index %[[ARG0]] into (3, 4, 2) : index, index
2241// CHECK:         return %[[RET]]#0, %[[RET]]#1, %[[RET]]#2 : index, index, index
2242func.func @cst_value_to_cst_attr_basis_delinearize_index(%arg0 : index) ->
2243    (index, index, index) {
2244  %c4 = arith.constant 4 : index
2245  %c3 = arith.constant 3 : index
2246  %c2 = arith.constant 2 : index
2247  %0:3 = affine.delinearize_index %arg0 into (%c3, %c4, %c2)
2248      : index, index, index
2249  return %0#0, %0#1, %0#2 : index, index, index
2250}
2251
2252// -----
2253
2254// CHECK-LABEL: @cst_value_to_cst_attr_basis_linearize_index
2255// CHECK-SAME:    (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
2256// CHECK:         %[[RET:.*]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]], %[[ARG2]]] by (2, 3, 4) : index
2257// CHECK:         return %[[RET]] : index
2258func.func @cst_value_to_cst_attr_basis_linearize_index(%arg0 : index, %arg1 : index, %arg2 : index) ->
2259    (index) {
2260  %c4 = arith.constant 4 : index
2261  %c2 = arith.constant 2 : index
2262  %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by  (%c2, 3, %c4) : index
2263  return %0 : index
2264}
2265