1// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-simplify-structures | FileCheck %s
2
3// CHECK-DAG: #[[$SET_2D:.*]] = affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0)>
4// CHECK-DAG: #[[$SET_7_11:.*]] = affine_set<(d0, d1) : (d0 * 7 + d1 * 5 + 88 == 0, d0 * 5 - d1 * 11 + 60 == 0, d0 * 11 + d1 * 7 - 24 == 0, d0 * 7 + d1 * 5 + 88 == 0)>
5
6// An external function that we will use in bodies to avoid DCE.
7func.func private @external() -> ()
8
9// CHECK-LABEL: func @test_gaussian_elimination_empty_set0() {
10func.func @test_gaussian_elimination_empty_set0() {
11  affine.for %arg0 = 1 to 10 {
12    affine.for %arg1 = 1 to 100 {
13      // CHECK-NOT: affine.if
14      affine.if affine_set<(d0, d1) : (2 == 0)>(%arg0, %arg1) {
15        func.call @external() : () -> ()
16      }
17    }
18  }
19  return
20}
21
22// CHECK-LABEL: func @test_gaussian_elimination_empty_set1() {
23func.func @test_gaussian_elimination_empty_set1() {
24  affine.for %arg0 = 1 to 10 {
25    affine.for %arg1 = 1 to 100 {
26      // CHECK-NOT: affine.if
27      affine.if affine_set<(d0, d1) : (1 >= 0, -1 >= 0)> (%arg0, %arg1) {
28        func.call @external() : () -> ()
29      }
30    }
31  }
32  return
33}
34
35// CHECK-LABEL: func @test_gaussian_elimination_non_empty_set2() {
36func.func @test_gaussian_elimination_non_empty_set2() {
37  affine.for %arg0 = 1 to 10 {
38    affine.for %arg1 = 1 to 100 {
39      // CHECK: #[[$SET_2D]](%arg0, %arg1)
40      affine.if affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0, d1 + 101 >= 0)>(%arg0, %arg1) {
41        func.call @external() : () -> ()
42      }
43    }
44  }
45  return
46}
47
48// CHECK-LABEL: func @test_gaussian_elimination_empty_set3() {
49func.func @test_gaussian_elimination_empty_set3() {
50  %c7 = arith.constant 7 : index
51  %c11 = arith.constant 11 : index
52  affine.for %arg0 = 1 to 10 {
53    affine.for %arg1 = 1 to 100 {
54      // CHECK-NOT: affine.if
55      affine.if affine_set<(d0, d1)[s0, s1] : (d0 - s0 == 0, d0 + s0 == 0, s0 - 1 == 0)>(%arg0, %arg1)[%c7, %c11] {
56        func.call @external() : () -> ()
57      }
58    }
59  }
60  return
61}
62
63// Set for test case: test_gaussian_elimination_non_empty_set4
64#set_2d_non_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
65                                       d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0,
66                                       d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0,
67                                       d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0)>
68
69// CHECK-LABEL: func @test_gaussian_elimination_non_empty_set4() {
70func.func @test_gaussian_elimination_non_empty_set4() {
71  %c7 = arith.constant 7 : index
72  %c11 = arith.constant 11 : index
73  affine.for %arg0 = 1 to 10 {
74    affine.for %arg1 = 1 to 100 {
75      // CHECK: #[[$SET_7_11]](%arg0, %arg1)
76      affine.if #set_2d_non_empty(%arg0, %arg1)[%c7, %c11] {
77        func.call @external() : () -> ()
78      }
79    }
80  }
81  return
82}
83
84// Add invalid constraints to previous non-empty set to make it empty.
85#set_2d_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
86                                       d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0,
87                                       d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0,
88                                       d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
89                                       d0 - 1 == 0, d0 + 2 == 0)>
90
91// CHECK-LABEL: func @test_gaussian_elimination_empty_set5() {
92func.func @test_gaussian_elimination_empty_set5() {
93  %c7 = arith.constant 7 : index
94  %c11 = arith.constant 11 : index
95  affine.for %arg0 = 1 to 10 {
96    affine.for %arg1 = 1 to 100 {
97      // CHECK-NOT: affine.if
98      affine.if #set_2d_empty(%arg0, %arg1)[%c7, %c11] {
99        func.call @external() : () -> ()
100      }
101    }
102  }
103  return
104}
105
106// This is an artificially created system to exercise the worst case behavior of
107// FM elimination - as a safeguard against improperly constructed constraint
108// systems or fuzz input.
109#set_fuzz_virus = affine_set<(d0, d1, d2, d3, d4, d5) : (
110                            1089234*d0 + 203472*d1 + 82342 >= 0,
111                            -55*d0 + 24*d1 + 238*d2 - 234*d3 - 9743 >= 0,
112                            -5445*d0 - 284*d1 + 23*d2 + 34*d3 - 5943 >= 0,
113                            -5445*d0 + 284*d1 + 238*d2 - 34*d3 >= 0,
114                            445*d0 + 284*d1 + 238*d2 + 39*d3 >= 0,
115                            -545*d0 + 214*d1 + 218*d2 - 94*d3 >= 0,
116                            44*d0 - 184*d1 - 231*d2 + 14*d3 >= 0,
117                            -45*d0 + 284*d1 + 138*d2 - 39*d3 >= 0,
118                            154*d0 - 84*d1 + 238*d2 - 34*d3 >= 0,
119                            54*d0 - 284*d1 - 223*d2 + 384*d3 >= 0,
120                            -55*d0 + 284*d1 + 23*d2 + 34*d3 >= 0,
121                            54*d0 - 84*d1 + 28*d2 - 34*d3 >= 0,
122                            54*d0 - 24*d1 - 23*d2 + 34*d3 >= 0,
123                            -55*d0 + 24*d1 + 23*d2 + 4*d3 >= 0,
124                            15*d0 - 84*d1 + 238*d2 - 3*d3 >= 0,
125                            5*d0 - 24*d1 - 223*d2 + 84*d3 >= 0,
126                            -5*d0 + 284*d1 + 23*d2 - 4*d3 >= 0,
127                            14*d0 + 4*d2 + 7234 >= 0,
128                            -174*d0 - 534*d2 + 9834 >= 0,
129                            194*d0 - 954*d2 + 9234 >= 0,
130                            47*d0 - 534*d2 + 9734 >= 0,
131                            -194*d0 - 934*d2 + 984 >= 0,
132                            -947*d0 - 953*d2 + 234 >= 0,
133                            184*d0 - 884*d2 + 884 >= 0,
134                            -174*d0 + 834*d2 + 234 >= 0,
135                            844*d0 + 634*d2 + 9874 >= 0,
136                            -797*d2 - 79*d3 + 257 >= 0,
137                            2039*d0 + 793*d2 - 99*d3 - 24*d4 + 234*d5 >= 0,
138                            78*d2 - 788*d5 + 257 >= 0,
139                            d3 - (d5 + 97*d0) floordiv 423 >= 0,
140                            234* (d0 + d3 mod 5 floordiv 2342) mod 2309
141                            + (d0 + 2038*d3) floordiv 208 >= 0,
142                            239* (d0 + 2300 * d3) floordiv 2342
143                            mod 2309 mod 239423 == 0,
144                            d0 + d3 mod 2642 + (d3 + 2*d0) mod 1247
145                            mod 2038 mod 2390 mod 2039 floordiv 55 >= 0
146)>
147
148// CHECK-LABEL: func @test_fuzz_explosion
149func.func @test_fuzz_explosion(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) {
150  affine.for %arg4 = 1 to 10 {
151    affine.for %arg5 = 1 to 100 {
152      affine.if #set_fuzz_virus(%arg4, %arg5, %arg0, %arg1, %arg2, %arg3) {
153        func.call @external() : () -> ()
154      }
155    }
156  }
157  return
158}
159
160// CHECK-LABEL: func @test_empty_set(%arg0: index) {
161func.func @test_empty_set(%N : index) {
162  affine.for %i = 0 to 10 {
163    affine.for %j = 0 to 10 {
164      // CHECK-NOT: affine.if
165      affine.if affine_set<(d0, d1) : (d0 - d1 >= 0, d1 - d0 - 1 >= 0)>(%i, %j) {
166        "foo"() : () -> ()
167      }
168      // CHECK-NOT: affine.if
169      affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {
170        "bar"() : () -> ()
171      }
172      // CHECK-NOT: affine.if
173      affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {
174        "foo"() : () -> ()
175      }
176      // CHECK-NOT: affine.if
177      affine.if affine_set<(d0)[s0, s1] : (d0 >= 0, -d0 + s0 - 1 >= 0, -s0 >= 0)>(%i)[%N, %N] {
178        "bar"() : () -> ()
179      }
180      // CHECK-NOT: affine.if
181      // The set below implies d0 = d1; so d1 >= d0, but d0 >= d1 + 1.
182      affine.if affine_set<(d0, d1, d2) : (d0 - d1 == 0, d2 - d0 >= 0, d0 - d1 - 1 >= 0)>(%i, %j, %N) {
183        "foo"() : () -> ()
184      }
185      // CHECK-NOT: affine.if
186      // The set below has rational solutions but no integer solutions; GCD test catches it.
187      affine.if affine_set<(d0, d1) : (d0*2 -d1*2 - 1 == 0, d0 >= 0, -d0 + 100 >= 0, d1 >= 0, -d1 + 100 >= 0)>(%i, %j) {
188        "foo"() : () -> ()
189      }
190      // CHECK-NOT: affine.if
191      affine.if affine_set<(d0, d1) : (d1 == 0, d0 - 1 >= 0, - d0 - 1 >= 0)>(%i, %j) {
192        "foo"() : () -> ()
193      }
194    }
195  }
196  // The tests below test GCDTightenInequalities().
197  affine.for %k = 0 to 10 {
198    affine.for %l = 0 to 10 {
199      // Empty because no multiple of 8 lies between 4 and 7.
200      // CHECK-NOT: affine.if
201      affine.if affine_set<(d0) : (8*d0 - 4 >= 0, -8*d0 + 7 >= 0)>(%k) {
202        "foo"() : () -> ()
203      }
204      // Same as above but with equalities and inequalities.
205      // CHECK-NOT: affine.if
206      affine.if affine_set<(d0, d1) : (d0 - 4*d1 == 0, 4*d1 - 5 >= 0, -4*d1 + 7 >= 0)>(%k, %l) {
207        "foo"() : () -> ()
208      }
209      // Same as above but with a combination of multiple identifiers. 4*d0 +
210      // 8*d1 here is a multiple of 4, and so can't lie between 9 and 11. GCD
211      // tightening will tighten constraints to 4*d0 + 8*d1 >= 12 and 4*d0 +
212      // 8*d1 <= 8; hence infeasible.
213      // CHECK-NOT: affine.if
214      affine.if affine_set<(d0, d1) : (4*d0 + 8*d1 - 9 >= 0, -4*d0 - 8*d1 + 11 >= 0)>(%k, %l) {
215        "foo"() : () -> ()
216      }
217      // Same as above but with equalities added into the mix.
218      // CHECK-NOT: affine.if
219      affine.if affine_set<(d0, d1, d2) : (d0 - 4*d2 == 0, d0 + 8*d1 - 9 >= 0, -d0 - 8*d1 + 11 >= 0)>(%k, %k, %l) {
220        "foo"() : () -> ()
221      }
222    }
223  }
224
225  affine.for %m = 0 to 10 {
226    // CHECK-NOT: affine.if
227    affine.if affine_set<(d0) : (d0 mod 2 - 3 == 0)> (%m) {
228      "foo"() : () -> ()
229    }
230  }
231
232  return
233}
234
235// -----
236
237// An external function that we will use in bodies to avoid DCE.
238func.func private @external() -> ()
239
240// CHECK-DAG: #[[$SET:.*]] = affine_set<()[s0] : (s0 >= 0, -s0 + 50 >= 0)
241
242// CHECK-LABEL: func @simplify_set
243func.func @simplify_set(%a : index, %b : index) {
244  // CHECK: affine.if #[[$SET]]
245  affine.if affine_set<(d0, d1) : (d0 - d1 + d1 + d0 >= 0, 2 >= 0, d0 >= 0, -d0 + 50 >= 0, -d0 + 100 >= 0)>(%a, %b) {
246    func.call @external() : () -> ()
247  }
248  // CHECK-NOT: affine.if
249  affine.if affine_set<(d0, d1) : (d0 mod 2 - 1 == 0, d0 - 2 * (d0 floordiv 2) == 0)>(%a, %b) {
250    func.call @external() : () -> ()
251  }
252  // CHECK-NOT: affine.if
253  affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%a, %b) {
254    func.call @external() : () -> ()
255  }
256	return
257}
258
259// -----
260
261// CHECK-DAG: -> (s0 * 2 + 1)
262
263// Test "op local" simplification on affine.apply. DCE on arith.addi will not happen.
264func.func @affine.apply(%N : index) -> index {
265  %v = affine.apply affine_map<(d0, d1) -> (d0 + d1 + 1)>(%N, %N)
266  %res = arith.addi %v, %v : index
267  // CHECK: affine.apply #map{{.*}}()[%arg0]
268  // CHECK-NEXT: arith.addi
269  return %res: index
270}
271
272// -----
273
274// CHECK-LABEL: func @simplify_zero_dim_map
275func.func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {
276  %out = affine.load %in[] : memref<f32>
277  return %out : f32
278}
279
280// -----
281
282// Tests the simplification of a semi-affine expression in various cases.
283// CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 2)>
284// CHECK-DAG: #[[$map1:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 42)>
285
286// Tests the simplification of a semi-affine expression with a modulo operation on a floordiv and multiplication.
287// CHECK-LABEL: func @semiaffine_mod
288func.func @semiaffine_mod(%arg0: index, %arg1: index) -> index {
289  %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * s0) mod s0)> (%arg0)[%arg1]
290  // CHECK:       %[[CST:.*]] = arith.constant 0
291  return %a : index
292}
293
294// Tests the simplification of a semi-affine expression with a nested floordiv and a floordiv on modulo operation.
295// CHECK-LABEL: func @semiaffine_floordiv
296func.func @semiaffine_floordiv(%arg0: index, %arg1: index) -> index {
297  %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + ((2 * s0) mod (3 * s0))) floordiv s0)> (%arg0)[%arg1]
298  // CHECK: affine.apply #[[$map0]]()[%arg1, %arg0]
299  return %a : index
300}
301
302// Tests the simplification of a semi-affine expression with a ceildiv operation and a division of arith.constant 0 by a symbol.
303// CHECK-LABEL: func @semiaffine_ceildiv
304func.func @semiaffine_ceildiv(%arg0: index, %arg1: index) -> index {
305  %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * 42 + ((5-5) floordiv s0)) ceildiv  s0)> (%arg0)[%arg1]
306  // CHECK: affine.apply #[[$map1]]()[%arg1, %arg0]
307  return %a : index
308}
309
310// Tests the simplification of a semi-affine expression with a nested ceildiv operation and further simplifications after performing ceildiv.
311// CHECK-LABEL: func @semiaffine_composite_floor
312func.func @semiaffine_composite_floor(%arg0: index, %arg1: index) -> index {
313  %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) ceildiv 4) * 5) + s0 * 42) ceildiv s0)> (%arg0)[%arg1]
314  // CHECK:       %[[CST:.*]] = arith.constant 47
315  return %a : index
316}
317
318// Tests the simplification of a semi-affine expression with a modulo operation with a second operand that simplifies to symbol.
319// CHECK-LABEL: func @semiaffine_unsimplified_symbol
320func.func @semiaffine_unsimplified_symbol(%arg0: index, %arg1: index) -> index {
321  %a = affine.apply affine_map<(d0)[s0] ->(s0 mod (2 * s0 - s0))> (%arg0)[%arg1]
322  // CHECK:       %[[CST:.*]] = arith.constant 0
323  return %a : index
324}
325
326// -----
327
328// Two external functions that we will use in bodies to avoid DCE.
329func.func private @external() -> ()
330func.func private @external1() -> ()
331
332// CHECK-LABEL: func @test_always_true_if_elimination() {
333func.func @test_always_true_if_elimination() {
334  affine.for %arg0 = 1 to 10 {
335    affine.for %arg1 = 1 to 100 {
336      affine.if affine_set<(d0, d1) : (1 >= 0)> (%arg0, %arg1) {
337        func.call @external() : () -> ()
338      } else {
339        func.call @external1() : () -> ()
340      }
341    }
342  }
343  return
344}
345
346// CHECK:      affine.for
347// CHECK-NEXT:   affine.for
348// CHECK-NEXT:     call @external()
349// CHECK-NEXT:   }
350// CHECK-NEXT: }
351
352// CHECK-LABEL: func @test_always_false_if_elimination() {
353func.func @test_always_false_if_elimination() {
354  // CHECK: affine.for
355  affine.for %arg0 = 1 to 10 {
356    // CHECK: affine.for
357    affine.for %arg1 = 1 to 100 {
358      // CHECK: call @external1()
359      // CHECK-NOT: affine.if
360      affine.if affine_set<(d0, d1) : (-1 >= 0)> (%arg0, %arg1) {
361        func.call @external() : () -> ()
362      } else {
363        func.call @external1() : () -> ()
364      }
365    }
366  }
367  return
368}
369
370
371// Testing: affine.if is not trivially true or false, nothing happens.
372// CHECK-LABEL: func @test_dimensional_if_elimination() {
373func.func @test_dimensional_if_elimination() {
374  affine.for %arg0 = 1 to 10 {
375    affine.for %arg1 = 1 to 100 {
376      // CHECK: affine.if
377      // CHECK: } else {
378      affine.if affine_set<(d0, d1) : (d0-1 == 0)> (%arg0, %arg1) {
379        func.call @external() : () -> ()
380      } else {
381        func.call @external() : () -> ()
382      }
383    }
384  }
385  return
386}
387
388// Testing: affine.if gets removed.
389// CHECK-LABEL: func @test_num_results_if_elimination
390func.func @test_num_results_if_elimination() -> index {
391  // CHECK: %[[zero:.*]] = arith.constant 0 : index
392  %zero = arith.constant 0 : index
393  %0 = affine.if affine_set<() : ()> () -> index {
394    affine.yield %zero : index
395  } else {
396    affine.yield %zero : index
397  }
398  // CHECK-NEXT: return %[[zero]] : index
399  return %0 : index
400}
401
402
403// Three more test functions involving affine.if operations which are
404// returning results:
405
406// Testing: affine.if gets removed. `Else` block get promoted.
407// CHECK-LABEL: func @test_trivially_false_returning_two_results
408// CHECK-SAME: (%[[arg0:.*]]: index)
409func.func @test_trivially_false_returning_two_results(%arg0: index) -> (index, index) {
410  // CHECK: %[[c7:.*]] = arith.constant 7 : index
411  // CHECK: %[[c13:.*]] = arith.constant 13 : index
412  %c7 = arith.constant 7 : index
413  %c13 = arith.constant 13 : index
414  // CHECK: %[[c2:.*]] = arith.constant 2 : index
415  // CHECK: %[[c3:.*]] = arith.constant 3 : index
416  %res:2 = affine.if affine_set<(d0, d1) : (5 >= 0, -2 >= 0)> (%c7, %c13) -> (index, index) {
417    %c0 = arith.constant 0 : index
418    %c1 = arith.constant 1 : index
419    affine.yield %c0, %c1 : index, index
420  } else {
421    %c2 = arith.constant 2 : index
422    %c3 = arith.constant 3 : index
423    affine.yield %c7, %arg0 : index, index
424  }
425  // CHECK-NEXT: return %[[c7]], %[[arg0]] : index, index
426  return %res#0, %res#1 : index, index
427}
428
429// Testing: affine.if gets removed. `Then` block get promoted.
430// CHECK-LABEL: func @test_trivially_true_returning_five_results
431func.func @test_trivially_true_returning_five_results() -> (index, index, index, index, index) {
432  // CHECK: %[[c12:.*]] = arith.constant 12 : index
433  // CHECK: %[[c13:.*]] = arith.constant 13 : index
434  %c12 = arith.constant 12 : index
435  %c13 = arith.constant 13 : index
436  // CHECK: %[[c0:.*]] = arith.constant 0 : index
437  // CHECK: %[[c1:.*]] = arith.constant 1 : index
438  // CHECK: %[[c2:.*]] = arith.constant 2 : index
439  // CHECK: %[[c3:.*]] = arith.constant 3 : index
440  // CHECK: %[[c4:.*]] = arith.constant 4 : index
441  %res:5 = affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%c12, %c13) -> (index, index, index, index, index) {
442    %c0 = arith.constant 0 : index
443    %c1 = arith.constant 1 : index
444    %c2 = arith.constant 2 : index
445    %c3 = arith.constant 3 : index
446    %c4 = arith.constant 4 : index
447    affine.yield %c0, %c1, %c2, %c3, %c4 : index, index, index, index, index
448  } else {
449    %c5 = arith.constant 5 : index
450    %c6 = arith.constant 6 : index
451    %c7 = arith.constant 7 : index
452    %c8 = arith.constant 8 : index
453    %c9 = arith.constant 9 : index
454    affine.yield %c5, %c6, %c7, %c8, %c9 : index, index, index, index, index
455  }
456  // CHECK-NEXT: return %[[c0]], %[[c1]], %[[c2]], %[[c3]], %[[c4]] : index, index, index, index, index
457  return %res#0, %res#1, %res#2, %res#3, %res#4 : index, index, index, index, index
458}
459
460// Testing: affine.if doesn't get removed.
461// CHECK-LABEL: func @test_not_trivially_true_or_false_returning_three_results
462func.func @test_not_trivially_true_or_false_returning_three_results() -> (index, index, index) {
463  // CHECK: %[[c8:.*]] = arith.constant 8 : index
464  // CHECK: %[[c13:.*]] = arith.constant 13 : index
465  %c8 = arith.constant 8 : index
466  %c13 = arith.constant 13 : index
467  // CHECK: affine.if
468  %res:3 = affine.if affine_set<(d0, d1) : (d0 - 1 == 0)>(%c8, %c13) -> (index, index, index) {
469    %c0 = arith.constant 0 : index
470    %c1 = arith.constant 1 : index
471    %c2 = arith.constant 2 : index
472    affine.yield %c0, %c1, %c2 : index, index, index
473  // CHECK: } else {
474  } else {
475    %c3 = arith.constant 3 : index
476    %c4 = arith.constant 4 : index
477    %c5 = arith.constant 5 : index
478    affine.yield %c3, %c4, %c5 : index, index, index
479  }
480  return %res#0, %res#1, %res#2 : index, index, index
481}
482
483// -----
484
485// Test simplification of mod expressions.
486// CHECK-DAG:   #[[$MOD:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 * s1 + (s0 - s1) mod s2)>
487// CHECK-DAG:   #[[$SIMPLIFIED_MOD_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 mod (s2 - s0 * s1))>
488// CHECK-DAG:   #[[$MODULO_AND_PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s3 - (-s0 + s3) mod s2)>
489// CHECK-LABEL: func @semiaffine_simplification_mod
490// CHECK-SAME:  (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
491func.func @semiaffine_simplification_mod(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index, index) {
492  %a = affine.apply affine_map<(d0, d1)[s0, s1, s2, s3] -> ((-(d1 * s0 - (s0 - s1) mod s2) + s3) + (d0 * s1 + d1 * s0))>(%arg0, %arg1)[%arg2, %arg3, %arg4, %arg5]
493  %b = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 mod (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]
494  %c = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 + (d0 + s0) mod s2 + s0 * s1 - (d0 + s0) mod s2 - (d0 - s0) mod s2)>(%arg0)[%arg1, %arg2, %arg3]
495  return %a, %b, %c : index, index, index
496}
497// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$MOD]]()[%[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG0]]]
498// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SIMPLIFIED_MOD_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
499// CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$MODULO_AND_PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
500// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]]
501
502// -----
503
504// Test simplification of floordiv and ceildiv expressions.
505// CHECK-DAG:   #[[$SIMPLIFIED_FLOORDIV_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 floordiv (s2 - s0 * s1))>
506// CHECK-DAG:   #[[$FLOORDIV:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 + s3 + (s0 - s1) floordiv s2)>
507// CHECK-DAG:   #[[$SIMPLIFIED_CEILDIV_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 ceildiv (s2 - s0 * s1))>
508// CHECK-LABEL: func @semiaffine_simplification_floordiv_and_ceildiv
509// CHECK-SAME:  (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)
510func.func @semiaffine_simplification_floordiv_and_ceildiv(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
511  %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 floordiv (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]
512  %b = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(d0 * s1 - (s0 - s1) floordiv s2) + s3) + (d0 * s1 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4]
513  %c = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 ceildiv (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]
514  return %a, %b, %c : index, index, index
515}
516// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$SIMPLIFIED_FLOORDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
517// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$FLOORDIV]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]]
518// CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$SIMPLIFIED_CEILDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
519// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]]
520
521// -----
522
523// Test simplification of product expressions.
524// CHECK-DAG:   #[[$PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 + (s0 - s1) * s2)>
525// CHECK-DAG:   #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s2 * s0 + s2 + s3 * s0 + s3 * s1 + s3 + s4 * s1 + s4)>
526// CHECK-LABEL: func @semiaffine_simplification_product
527// CHECK-SAME:  (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
528func.func @semiaffine_simplification_product(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
529  %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(s0 - (s0 - s1) * s2) + s3) + (d0 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4]
530  %b = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + d1 * s1 + d1 + d0 * s0 + d1 * s0 + d2 * s1 + d2)>(%arg0, %arg1, %arg2)[%arg3, %arg4]
531  return %a, %b : index, index
532}
533// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG0]]]
534// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SUM_OF_PRODUCTS]]()[%[[ARG3]], %[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG2]]]
535// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]]
536
537// -----
538
539// CHECK-DAG: #[[$SIMPLIFIED_MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> ((-s0 + s2 + s3) mod (s0 + s1))>
540// CHECK-LABEL: func @semi_affine_simplification_euclidean_lemma
541// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
542func.func @semi_affine_simplification_euclidean_lemma(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
543  %a = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1) - ((d0 + d1) floordiv (s0 - s1)) * (s0 - s1) - (d0 + d1) mod (s0 - s1))>(%arg0, %arg1)[%arg2, %arg3]
544  %b = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1 - s0) - ((d0 + d1 - s0) floordiv (s0 + s1)) * (s0 + s1))>(%arg0, %arg1)[%arg2, %arg3]
545  return %a, %b : index, index
546}
547// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index
548// CHECK-NEXT: %[[RESULT:.*]] = affine.apply #[[$SIMPLIFIED_MAP]]()[%[[ARG2]], %[[ARG3]], %[[ARG0]], %[[ARG1]]]
549// CHECK-NEXT: return %[[ZERO]], %[[RESULT]]
550