1// RUN: mlir-opt %s -loop-invariant-code-motion -split-input-file | FileCheck %s
2
3func @nested_loops_both_having_invariant_code() {
4  %m = alloc() : memref<10xf32>
5  %cf7 = constant 7.0 : f32
6  %cf8 = constant 8.0 : f32
7
8  affine.for %arg0 = 0 to 10 {
9    %v0 = addf %cf7, %cf8 : f32
10    affine.for %arg1 = 0 to 10 {
11      %v1 = addf %v0, %cf8 : f32
12      affine.store %v0, %m[%arg0] : memref<10xf32>
13    }
14  }
15
16  // CHECK: %0 = alloc() : memref<10xf32>
17  // CHECK-NEXT: %[[CST0:.*]] = constant 7.000000e+00 : f32
18  // CHECK-NEXT: %[[CST1:.*]] = constant 8.000000e+00 : f32
19  // CHECK-NEXT: %[[ADD0:.*]] = addf %[[CST0]], %[[CST1]] : f32
20  // CHECK-NEXT: addf %[[ADD0]], %[[CST1]] : f32
21  // CHECK-NEXT: affine.for
22  // CHECK-NEXT: affine.for
23  // CHECK-NEXT: affine.store
24
25  return
26}
27
28func @nested_loops_code_invariant_to_both() {
29  %m = alloc() : memref<10xf32>
30  %cf7 = constant 7.0 : f32
31  %cf8 = constant 8.0 : f32
32
33  affine.for %arg0 = 0 to 10 {
34    affine.for %arg1 = 0 to 10 {
35      %v0 = addf %cf7, %cf8 : f32
36    }
37  }
38
39  // CHECK: %0 = alloc() : memref<10xf32>
40  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
41  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
42  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
43
44  return
45}
46
47func @single_loop_nothing_invariant() {
48  %m1 = alloc() : memref<10xf32>
49  %m2 = alloc() : memref<10xf32>
50  affine.for %arg0 = 0 to 10 {
51    %v0 = affine.load %m1[%arg0] : memref<10xf32>
52    %v1 = affine.load %m2[%arg0] : memref<10xf32>
53    %v2 = addf %v0, %v1 : f32
54    affine.store %v2, %m1[%arg0] : memref<10xf32>
55  }
56
57  // CHECK: %0 = alloc() : memref<10xf32>
58  // CHECK-NEXT: %1 = alloc() : memref<10xf32>
59  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
60  // CHECK-NEXT: %2 = affine.load %0[%arg0] : memref<10xf32>
61  // CHECK-NEXT: %3 = affine.load %1[%arg0] : memref<10xf32>
62  // CHECK-NEXT: %4 = addf %2, %3 : f32
63  // CHECK-NEXT: affine.store %4, %0[%arg0] : memref<10xf32>
64
65  return
66}
67
68func @invariant_code_inside_affine_if() {
69  %m = alloc() : memref<10xf32>
70  %cf8 = constant 8.0 : f32
71
72  affine.for %arg0 = 0 to 10 {
73    %t0 = affine.apply affine_map<(d1) -> (d1 + 1)>(%arg0)
74    affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %t0) {
75        %cf9 = addf %cf8, %cf8 : f32
76        affine.store %cf9, %m[%arg0] : memref<10xf32>
77
78    }
79  }
80
81  // CHECK: %0 = alloc() : memref<10xf32>
82  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
83  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
84  // CHECK-NEXT: %1 = affine.apply #map3(%arg0)
85  // CHECK-NEXT: affine.if #set0(%arg0, %1) {
86  // CHECK-NEXT: %2 = addf %cst, %cst : f32
87  // CHECK-NEXT: affine.store %2, %0[%arg0] : memref<10xf32>
88  // CHECK-NEXT: }
89
90
91  return
92}
93
94func @invariant_affine_if() {
95  %m = alloc() : memref<10xf32>
96  %cf8 = constant 8.0 : f32
97  affine.for %arg0 = 0 to 10 {
98    affine.for %arg1 = 0 to 10 {
99      affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
100          %cf9 = addf %cf8, %cf8 : f32
101      }
102    }
103  }
104
105  // CHECK: %0 = alloc() : memref<10xf32>
106  // CHECK-NEXT: %[[CST:.*]] = constant 8.000000e+00 : f32
107  // CHECK-NEXT: affine.for %[[ARG:.*]] = 0 to 10 {
108  // CHECK-NEXT: affine.if #set0(%[[ARG]], %[[ARG]]) {
109  // CHECK-NEXT: addf %[[CST]], %[[CST]] : f32
110  // CHECK-NEXT: }
111
112  return
113}
114
115func @invariant_affine_if2() {
116  %m = alloc() : memref<10xf32>
117  %cf8 = constant 8.0 : f32
118  affine.for %arg0 = 0 to 10 {
119    affine.for %arg1 = 0 to 10 {
120      affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
121          %cf9 = addf %cf8, %cf8 : f32
122          affine.store %cf9, %m[%arg1] : memref<10xf32>
123      }
124    }
125  }
126
127  // CHECK: alloc
128  // CHECK-NEXT: constant
129  // CHECK-NEXT: affine.for
130  // CHECK-NEXT: affine.for
131  // CHECK-NEXT: affine.if
132  // CHECK-NEXT: addf
133  // CHECK-NEXT: affine.store
134  // CHECK-NEXT: }
135  // CHECK-NEXT: }
136
137  return
138}
139
140func @invariant_affine_nested_if() {
141  %m = alloc() : memref<10xf32>
142  %cf8 = constant 8.0 : f32
143  affine.for %arg0 = 0 to 10 {
144    affine.for %arg1 = 0 to 10 {
145      affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
146          %cf9 = addf %cf8, %cf8 : f32
147          affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
148            %cf10 = addf %cf9, %cf9 : f32
149          }
150      }
151    }
152  }
153
154  // CHECK: alloc
155  // CHECK-NEXT: constant
156  // CHECK-NEXT: affine.for
157  // CHECK-NEXT: affine.for
158  // CHECK-NEXT: affine.if
159  // CHECK-NEXT: addf
160  // CHECK-NEXT: affine.if
161  // CHECK-NEXT: addf
162  // CHECK-NEXT: }
163  // CHECK-NEXT: }
164  // CHECK-NEXT: }
165
166
167  return
168}
169
170func @invariant_affine_nested_if_else() {
171  %m = alloc() : memref<10xf32>
172  %cf8 = constant 8.0 : f32
173  affine.for %arg0 = 0 to 10 {
174    affine.for %arg1 = 0 to 10 {
175      affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
176          %cf9 = addf %cf8, %cf8 : f32
177          affine.store %cf9, %m[%arg0] : memref<10xf32>
178          affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
179            %cf10 = addf %cf9, %cf9 : f32
180          } else {
181            affine.store %cf9, %m[%arg1] : memref<10xf32>
182          }
183      }
184    }
185  }
186
187  // CHECK: alloc
188  // CHECK-NEXT: constant
189  // CHECK-NEXT: affine.for
190  // CHECK-NEXT: affine.for
191  // CHECK-NEXT: affine.if
192  // CHECK-NEXT: addf
193  // CHECK-NEXT: affine.store
194  // CHECK-NEXT: affine.if
195  // CHECK-NEXT: addf
196  // CHECK-NEXT: } else {
197  // CHECK-NEXT: affine.store
198  // CHECK-NEXT: }
199  // CHECK-NEXT: }
200  // CHECK-NEXT: }
201
202
203  return
204}
205
206func @invariant_loop_dialect() {
207  %ci0 = constant 0 : index
208  %ci10 = constant 10 : index
209  %ci1 = constant 1 : index
210  %m = alloc() : memref<10xf32>
211  %cf7 = constant 7.0 : f32
212  %cf8 = constant 8.0 : f32
213  loop.for %arg0 = %ci0 to %ci10 step %ci1 {
214    loop.for %arg1 = %ci0 to %ci10 step %ci1 {
215      %v0 = addf %cf7, %cf8 : f32
216    }
217  }
218
219  // CHECK: %0 = alloc() : memref<10xf32>
220  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
221  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
222  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
223
224  return
225}
226
227func @variant_loop_dialect() {
228  %ci0 = constant 0 : index
229  %ci10 = constant 10 : index
230  %ci1 = constant 1 : index
231  %m = alloc() : memref<10xf32>
232  loop.for %arg0 = %ci0 to %ci10 step %ci1 {
233    loop.for %arg1 = %ci0 to %ci10 step %ci1 {
234      %v0 = addi %arg0, %arg1 : index
235    }
236  }
237
238  // CHECK: %0 = alloc() : memref<10xf32>
239  // CHECK-NEXT: loop.for
240  // CHECK-NEXT: loop.for
241  // CHECK-NEXT: addi
242
243  return
244}
245