1// RUN: mlir-opt %s -split-input-file -test-affine-loop-unswitch | FileCheck %s
2
3// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0) : (d0 - 2 >= 0)>
4
5// CHECK-LABEL: func @if_else_imperfect
6func @if_else_imperfect(%A : memref<100xi32>, %B : memref<100xi32>, %v : i32) {
7// CHECK: %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref
8  affine.for %i = 0 to 100 {
9    affine.store %v, %A[%i] : memref<100xi32>
10    affine.for %j = 0 to 100 {
11      affine.store %v, %A[%j] : memref<100xi32>
12      affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {
13        affine.store %v, %B[%j] : memref<100xi32>
14      }
15      call @external() : () -> ()
16    }
17    affine.store %v, %A[%i] : memref<100xi32>
18  }
19  return
20}
21func private @external()
22
23// CHECK:       affine.for %[[I:.*]] = 0 to 100 {
24// CHECK-NEXT:    affine.store %{{.*}}, %[[A]][%[[I]]]
25// CHECK-NEXT:    affine.if #[[$SET]](%[[I]]) {
26// CHECK-NEXT:      affine.for %[[J:.*]] = 0 to 100 {
27// CHECK-NEXT:        affine.store %{{.*}}, %[[A]][%[[J]]]
28// CHECK-NEXT:        affine.store %{{.*}}, %[[B]][%[[J]]]
29// CHECK-NEXT:        call
30// CHECK-NEXT:      }
31// CHECK-NEXT:    } else {
32// CHECK-NEXT:      affine.for %[[JJ:.*]] = 0 to 100 {
33// CHECK-NEXT:        affine.store %{{.*}}, %[[A]][%[[J]]]
34// CHECK-NEXT:        call
35// CHECK-NEXT:      }
36// CHECK-NEXT:    }
37// CHECK-NEXT:    affine.store %{{.*}}, %[[A]][%[[I]]]
38// CHECK-NEXT:  }
39// CHECK-NEXT:  return
40
41// -----
42
43func private @foo()
44func private @bar()
45func private @abc()
46func private @xyz()
47
48// CHECK-LABEL: func @if_then_perfect
49func @if_then_perfect(%A : memref<100xi32>, %v : i32) {
50  affine.for %i = 0 to 100 {
51    affine.for %j = 0 to 100 {
52      affine.for %k = 0 to 100 {
53        affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {
54          affine.store %v, %A[%i] : memref<100xi32>
55        }
56      }
57    }
58  }
59  return
60}
61// CHECK:      affine.for
62// CHECK-NEXT:   affine.if
63// CHECK-NEXT:     affine.for
64// CHECK-NEXT:       affine.for
65// CHECK-NOT:    else
66
67
68// CHECK-LABEL: func @if_else_perfect
69func @if_else_perfect(%A : memref<100xi32>, %v : i32) {
70  affine.for %i = 0 to 99 {
71    affine.for %j = 0 to 100 {
72      affine.for %k = 0 to 100 {
73        call @foo() : () -> ()
74        affine.if affine_set<(d0, d1) : (d0 - 2 >= 0, -d1 + 80 >= 0)>(%i, %j) {
75          affine.store %v, %A[%i] : memref<100xi32>
76          call @abc() : () -> ()
77        } else {
78          affine.store %v, %A[%i + 1] : memref<100xi32>
79          call @xyz() : () -> ()
80        }
81        call @bar() : () -> ()
82      }
83    }
84  }
85  return
86}
87// CHECK:      affine.for
88// CHECK-NEXT:   affine.for
89// CHECK-NEXT:     affine.if
90// CHECK-NEXT:       affine.for
91// CHECK-NEXT:         call @foo
92// CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}]
93// CHECK-NEXT:         call @abc
94// CHECK-NEXT:         call @bar
95// CHECK-NEXT:       }
96// CHECK-NEXT:     else
97// CHECK-NEXT:       affine.for
98// CHECK-NEXT:         call @foo
99// CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}} + 1]
100// CHECK-NEXT:         call @xyz
101// CHECK-NEXT:         call @bar
102// CHECK-NEXT:       }
103// CHECK-NEXT:     }
104// CHECK-NEXT:   }
105// CHECK-NEXT: }
106
107// CHECK-LABEL: func @if_then_imperfect
108func @if_then_imperfect(%A : memref<100xi32>, %N : index, %v: i32) {
109  affine.for %i = 0 to 100 {
110    affine.store %v, %A[0] : memref<100xi32>
111    affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%N) {
112      affine.store %v, %A[%i] : memref<100xi32>
113    }
114  }
115  return
116}
117// CHECK:       affine.if
118// CHECK-NEXT:    affine.for
119// CHECK-NEXT:      affine.store
120// CHECK-NEXT:      affine.store
121// CHECK-NEXT:    }
122// CHECK-NEXT:  } else {
123// CHECK-NEXT:    affine.for
124// CHECK-NEXT:      affine.store
125// CHECK-NEXT:    }
126// CHECK-NEXT:  }
127// CHECK-NEXT:  return
128
129// Check if unused operands are dropped: hence, hoisting is possible.
130// CHECK-LABEL: func @hoist_after_canonicalize
131func @hoist_after_canonicalize() {
132  affine.for %i = 0 to 100 {
133    affine.for %j = 0 to 100 {
134      affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%j)  {
135        affine.if affine_set<(d0, d1) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%i, %j)  {
136          // The call to external is to avoid DCE on affine.if.
137          call @foo() : () -> ()
138        }
139      }
140    }
141  }
142  return
143}
144// CHECK:      affine.for
145// CHECK-NEXT:   affine.if
146// CHECK-NEXT:     affine.for
147// CHECK-NEXT:       affine.if
148// CHECK-NEXT:         call
149// CHECK-NEXT:       }
150// CHECK-NEXT:     }
151// CHECK-NEXT:   }
152// CHECK-NEXT: }
153// CHECK-NEXT: return
154
155// CHECK-LABEL: func @handle_dead_if
156func @handle_dead_if(%N : index) {
157  affine.for %i = 0 to 100 {
158      affine.if affine_set<(d0) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%N)  {
159      }
160  }
161  return
162}
163// CHECK-NEXT: affine.for
164// CHECK-NEXT: }
165// CHECK-NEXT: return
166
167// -----
168
169// A test case with affine.parallel.
170
171#flb1 = affine_map<(d0) -> (d0 * 3)>
172#fub1 = affine_map<(d0) -> (d0 * 3 + 3)>
173#flb0 = affine_map<(d0) -> (d0 * 16)>
174#fub0 = affine_map<(d0) -> (d0 * 16 + 16)>
175#pub1 = affine_map<(d0)[s0] -> (s0, d0 * 3 + 3)>
176#pub0 = affine_map<(d0)[s0] -> (s0, d0 * 16 + 16)>
177#lb1 = affine_map<(d0) -> (d0 * 480)>
178#ub1 = affine_map<(d0)[s0] -> (s0, d0 * 480 + 480)>
179#lb0 = affine_map<(d0) -> (d0 * 110)>
180#ub0 = affine_map<(d0)[s0] -> (d0 * 110 + 110, s0 floordiv 3)>
181
182#set0 = affine_set<(d0, d1)[s0, s1] : (d0 * -16 + s0 - 16 >= 0, d1 * -3 + s1 - 3 >= 0)>
183
184// CHECK-LABEL: func @perfect_if_else
185func @perfect_if_else(%arg0 : memref<?x?xf64>, %arg1 : memref<?x?xf64>, %v : f64,
186            %arg4 : index, %arg5 : index, %arg6 : index, %sym : index) {
187  affine.for %arg7 = #lb0(%arg5) to min #ub0(%arg5)[%sym] {
188    affine.parallel (%i0, %j0) = (0, 0) to (symbol(%sym), 100) step (10, 10) {
189      affine.for %arg8 = #lb1(%arg4) to min #ub1(%arg4)[%sym] {
190        affine.if #set0(%arg6, %arg7)[%sym, %sym] {
191          affine.for %arg9 = #flb0(%arg6) to #fub0(%arg6) {
192            affine.for %arg10 = #flb1(%arg7) to #fub1(%arg7) {
193              affine.store %v, %arg0[0, 0] : memref<?x?xf64>
194            }
195          }
196        } else {
197          affine.for %arg9 = #lb0(%arg6) to min #pub0(%arg6)[%sym] {
198            affine.for %arg10 = #lb1(%arg7) to min #pub1(%arg7)[%sym] {
199              affine.store %v, %arg0[0, 0] : memref<?x?xf64>
200            }
201          }
202        }
203      }
204    }
205  }
206  return
207}
208
209// CHECK:       affine.for
210// CHECK-NEXT:    affine.if
211// CHECK-NEXT:      affine.parallel
212// CHECK-NEXT:        affine.for
213// CHECK-NEXT:          affine.for
214// CHECK-NEXT:            affine.for
215// CHECK-NEXT:              affine.store
216// CHECK-NEXT:            }
217// CHECK-NEXT:          }
218// CHECK-NEXT:        }
219// CHECK-NEXT:      }
220// CHECK-NEXT:    } else {
221// CHECK-NEXT:      affine.parallel
222// CHECK-NEXT:        affine.for
223// CHECK-NEXT:          affine.for
224// CHECK-NEXT:            affine.for
225// CHECK-NEXT:              affine.store
226// CHECK-NEXT:            }
227// CHECK-NEXT:          }
228// CHECK-NEXT:        }
229// CHECK-NEXT:      }
230// CHECK-NEXT:    }
231// CHECK-NEXT:  }
232
233// With multiple if ops in a function, the test pass just looks for the first if
234// op that it is able to successfully hoist.
235
236// CHECK-LABEL: func @multiple_if
237func @multiple_if(%N : index) {
238  affine.if affine_set<() : (0 == 0)>() {
239    call @external() : () -> ()
240  }
241  affine.for %i = 0 to 100 {
242    affine.if affine_set<()[s0] : (s0 >= 0)>()[%N] {
243      call @external() : () -> ()
244    }
245  }
246  return
247}
248// CHECK:      call
249// CHECK-NEXT: affine.if
250// CHECK-NEXT:   affine.for
251// CHECK-NEXT:     call
252// CHECK-NEXT:   }
253// CHECK-NEXT: }
254// CHECK-NEXT: return
255
256func private @external()
257