1// RUN: mlir-opt %s -affine-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      affine.store %v0, %m[%arg0] : memref<10xf32>
12    }
13  }
14
15  // CHECK: %0 = alloc() : memref<10xf32>
16  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
17  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
18  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
19  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
20  // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
21
22  return
23}
24
25// The store-load forwarding can see through affine apply's since it relies on
26// dependence information.
27// CHECK-LABEL: func @store_affine_apply
28func @store_affine_apply() -> memref<10xf32> {
29  %cf7 = constant 7.0 : f32
30  %m = alloc() : memref<10xf32>
31  affine.for %arg0 = 0 to 10 {
32      %t0 = affine.apply (d1) -> (d1 + 1)(%arg0)
33      affine.store %cf7, %m[%t0] : memref<10xf32>
34  }
35  return %m : memref<10xf32>
36// CHECK:       %cst = constant 7.000000e+00 : f32
37// CHECK-NEXT:  %0 = alloc() : memref<10xf32>
38// CHECK-NEXT:  affine.for %arg0 = 0 to 10 {
39// CHECK-NEXT:      %1 = affine.apply #map3(%arg0)
40// CHECK-NEXT:      affine.store %cst, %0[%1] : memref<10xf32>
41// CHECK-NEXT:  }
42// CHECK-NEXT:  return %0 : memref<10xf32>
43}
44
45func @nested_loops_code_invariant_to_both() {
46  %m = alloc() : memref<10xf32>
47  %cf7 = constant 7.0 : f32
48  %cf8 = constant 8.0 : f32
49
50  affine.for %arg0 = 0 to 10 {
51    affine.for %arg1 = 0 to 10 {
52      %v0 = addf %cf7, %cf8 : f32
53    }
54  }
55
56  // CHECK: %0 = alloc() : memref<10xf32>
57  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
58  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
59  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
60
61  return
62}
63
64func @single_loop_nothing_invariant() {
65  %m1 = alloc() : memref<10xf32>
66  %m2 = alloc() : memref<10xf32>
67  affine.for %arg0 = 0 to 10 {
68    %v0 = affine.load %m1[%arg0] : memref<10xf32>
69    %v1 = affine.load %m2[%arg0] : memref<10xf32>
70    %v2 = addf %v0, %v1 : f32
71    affine.store %v2, %m1[%arg0] : memref<10xf32>
72  }
73
74  // CHECK: %0 = alloc() : memref<10xf32>
75  // CHECK-NEXT: %1 = alloc() : memref<10xf32>
76  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
77  // CHECK-NEXT: %2 = affine.load %0[%arg0] : memref<10xf32>
78  // CHECK-NEXT: %3 = affine.load %1[%arg0] : memref<10xf32>
79  // CHECK-NEXT: %4 = addf %2, %3 : f32
80  // CHECK-NEXT: affine.store %4, %0[%arg0] : memref<10xf32>
81
82  return
83}
84
85func @invariant_code_inside_affine_if() {
86  %m = alloc() : memref<10xf32>
87  %cf8 = constant 8.0 : f32
88
89  affine.for %arg0 = 0 to 10 {
90    %t0 = affine.apply (d1) -> (d1 + 1)(%arg0)
91    affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %t0) {
92        %cf9 = addf %cf8, %cf8 : f32
93        affine.store %cf9, %m[%arg0] : memref<10xf32>
94
95    }
96  }
97
98  // CHECK: %0 = alloc() : memref<10xf32>
99  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
100  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
101  // CHECK-NEXT: %1 = affine.apply #map3(%arg0)
102  // CHECK-NEXT: affine.if #set0(%arg0, %1) {
103  // CHECK-NEXT: %2 = addf %cst, %cst : f32
104  // CHECK-NEXT: affine.store %2, %0[%arg0] : memref<10xf32>
105  // CHECK-NEXT: }
106
107
108  return
109}
110
111
112func @dependent_stores() {
113  %m = alloc() : memref<10xf32>
114  %cf7 = constant 7.0 : f32
115  %cf8 = constant 8.0 : f32
116
117  affine.for %arg0 = 0 to 10 {
118    %v0 = addf %cf7, %cf8 : f32
119    affine.for %arg1 = 0 to 10 {
120      %v1 = addf %cf7, %cf7 : f32
121      affine.store %v1, %m[%arg1] : memref<10xf32>
122      affine.store %v0, %m[%arg0] : memref<10xf32>
123    }
124  }
125
126  // CHECK: %0 = alloc() : memref<10xf32>
127  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
128  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
129  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
130  // CHECK-NEXT: %2 = addf %cst, %cst : f32
131  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
132
133  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
134  // CHECK-NEXT:   affine.store %2, %0[%arg1] : memref<10xf32>
135  // CHECK-NEXT:   affine.store %1, %0[%arg0] : memref<10xf32>
136
137  return
138}
139
140func @independent_stores() {
141  %m = alloc() : memref<10xf32>
142  %cf7 = constant 7.0 : f32
143  %cf8 = constant 8.0 : f32
144
145  affine.for %arg0 = 0 to 10 {
146    %v0 = addf %cf7, %cf8 : f32
147    affine.for %arg1 = 0 to 10 {
148      %v1 = addf %cf7, %cf7 : f32
149      affine.store %v0, %m[%arg0] : memref<10xf32>
150      affine.store %v1, %m[%arg1] : memref<10xf32>
151    }
152  }
153
154  // CHECK: %0 = alloc() : memref<10xf32>
155  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
156  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
157  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
158  // CHECK-NEXT: %2 = addf %cst, %cst : f32
159  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
160  // CHECK-NEXT:   affine.for %arg1 = 0 to 10 {
161  // CHECK-NEXT:     affine.store %1, %0[%arg0] : memref<10xf32>
162  // CHECK-NEXT:     affine.store %2, %0[%arg1] : memref<10xf32>
163  // CHECK-NEXT:    }
164
165  return
166}
167
168func @load_dependent_store() {
169  %m = alloc() : memref<10xf32>
170  %cf7 = constant 7.0 : f32
171  %cf8 = constant 8.0 : f32
172
173  affine.for %arg0 = 0 to 10 {
174    %v0 = addf %cf7, %cf8 : f32
175    affine.for %arg1 = 0 to 10 {
176      %v1 = addf %cf7, %cf7 : f32
177      affine.store %v0, %m[%arg1] : memref<10xf32>
178      %v2 = affine.load %m[%arg0] : memref<10xf32>
179    }
180  }
181
182  // CHECK: %0 = alloc() : memref<10xf32>
183  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
184  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
185  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
186  // CHECK-NEXT: %2 = addf %cst, %cst : f32
187  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
188  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
189  // CHECK-NEXT:   affine.store %1, %0[%arg1] : memref<10xf32>
190  // CHECK-NEXT:   %3 = affine.load %0[%arg0] : memref<10xf32>
191
192  return
193}
194
195func @load_after_load() {
196  %m = alloc() : memref<10xf32>
197  %cf7 = constant 7.0 : f32
198  %cf8 = constant 8.0 : f32
199
200  affine.for %arg0 = 0 to 10 {
201    %v0 = addf %cf7, %cf8 : f32
202    affine.for %arg1 = 0 to 10 {
203      %v1 = addf %cf7, %cf7 : f32
204      %v3 = affine.load %m[%arg1] : memref<10xf32>
205      %v2 = affine.load %m[%arg0] : memref<10xf32>
206    }
207  }
208
209  // CHECK: %0 = alloc() : memref<10xf32>
210  // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
211  // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
212  // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
213  // CHECK-NEXT: %2 = addf %cst, %cst : f32
214  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
215  // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
216  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
217  // CHECK-NEXT: %4 = affine.load %0[%arg1] : memref<10xf32>
218
219  return
220}
221
222func @invariant_affine_if() {
223  %m = alloc() : memref<10xf32>
224  %cf8 = constant 8.0 : f32
225  affine.for %arg0 = 0 to 10 {
226    affine.for %arg1 = 0 to 10 {
227      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
228          %cf9 = addf %cf8, %cf8 : f32
229          affine.store %cf9, %m[%arg0] : memref<10xf32>
230
231      }
232    }
233  }
234
235  // CHECK: %0 = alloc() : memref<10xf32>
236  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
237  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
238  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
239  // CHECK-NEXT: %1 = addf %cst, %cst : f32
240  // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
241  // CHECK-NEXT: }
242
243
244  return
245}
246
247func @invariant_affine_if2() {
248  %m = alloc() : memref<10xf32>
249  %cf8 = constant 8.0 : f32
250  affine.for %arg0 = 0 to 10 {
251    affine.for %arg1 = 0 to 10 {
252      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
253          %cf9 = addf %cf8, %cf8 : f32
254          affine.store %cf9, %m[%arg1] : memref<10xf32>
255
256      }
257    }
258  }
259
260  // CHECK: %0 = alloc() : memref<10xf32>
261  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
262  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
263  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
264  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
265  // CHECK-NEXT: %1 = addf %cst, %cst : f32
266  // CHECK-NEXT: affine.store %1, %0[%arg1] : memref<10xf32>
267  // CHECK-NEXT: }
268  // CHECK-NEXT: }
269
270
271  return
272}
273
274func @invariant_affine_nested_if() {
275  %m = alloc() : memref<10xf32>
276  %cf8 = constant 8.0 : f32
277  affine.for %arg0 = 0 to 10 {
278    affine.for %arg1 = 0 to 10 {
279      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
280          %cf9 = addf %cf8, %cf8 : f32
281          affine.store %cf9, %m[%arg0] : memref<10xf32>
282          affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
283            affine.store %cf9, %m[%arg1] : memref<10xf32>
284          }
285      }
286    }
287  }
288
289  // CHECK: %0 = alloc() : memref<10xf32>
290  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
291  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
292  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
293  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
294  // CHECK-NEXT: %1 = addf %cst, %cst : f32
295  // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
296  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
297  // CHECK-NEXT: affine.store %1, %0[%arg1] : memref<10xf32>
298  // CHECK-NEXT: }
299  // CHECK-NEXT: }
300  // CHECK-NEXT: }
301
302
303  return
304}
305
306func @invariant_affine_nested_if_else() {
307  %m = alloc() : memref<10xf32>
308  %cf8 = constant 8.0 : f32
309  affine.for %arg0 = 0 to 10 {
310    affine.for %arg1 = 0 to 10 {
311      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
312          %cf9 = addf %cf8, %cf8 : f32
313          affine.store %cf9, %m[%arg0] : memref<10xf32>
314          affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
315            affine.store %cf9, %m[%arg0] : memref<10xf32>
316          } else {
317            affine.store %cf9, %m[%arg1] : memref<10xf32>
318          }
319      }
320    }
321  }
322
323  // CHECK: %0 = alloc() : memref<10xf32>
324  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
325  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
326  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
327  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
328  // CHECK-NEXT: %1 = addf %cst, %cst : f32
329  // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
330  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
331  // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
332  // CHECK-NEXT: } else {
333  // CHECK-NEXT: affine.store %1, %0[%arg1] : memref<10xf32>
334  // CHECK-NEXT: }
335  // CHECK-NEXT: }
336  // CHECK-NEXT: }
337
338
339  return
340}
341
342func @invariant_affine_nested_if_else2() {
343  %m = alloc() : memref<10xf32>
344  %m2 = alloc() : memref<10xf32>
345  %cf8 = constant 8.0 : f32
346  affine.for %arg0 = 0 to 10 {
347    affine.for %arg1 = 0 to 10 {
348      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
349          %cf9 = addf %cf8, %cf8 : f32
350          %tload1 = affine.load %m[%arg0] : memref<10xf32>
351          affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
352            affine.store %cf9, %m2[%arg0] : memref<10xf32>
353          } else {
354            %tload2 = affine.load %m[%arg0] : memref<10xf32>
355          }
356      }
357    }
358  }
359
360  // CHECK: %0 = alloc() : memref<10xf32>
361  // CHECK-NEXT: %1 = alloc() : memref<10xf32>
362  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
363  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
364  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
365  // CHECK-NEXT: %2 = addf %cst, %cst : f32
366  // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
367  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
368  // CHECK-NEXT: affine.store %2, %1[%arg0] : memref<10xf32>
369  // CHECK-NEXT: } else {
370  // CHECK-NEXT: %4 = affine.load %0[%arg0] : memref<10xf32>
371  // CHECK-NEXT: }
372  // CHECK-NEXT: }
373
374
375  return
376}
377
378
379func @invariant_affine_nested_if2() {
380  %m = alloc() : memref<10xf32>
381  %cf8 = constant 8.0 : f32
382  affine.for %arg0 = 0 to 10 {
383    affine.for %arg1 = 0 to 10 {
384      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
385          %cf9 = addf %cf8, %cf8 : f32
386          %v1 = affine.load %m[%arg0] : memref<10xf32>
387          affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
388            %v2 = affine.load %m[%arg0] : memref<10xf32>
389          }
390      }
391    }
392  }
393
394  // CHECK: %0 = alloc() : memref<10xf32>
395  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
396  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
397  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
398  // CHECK-NEXT: %1 = addf %cst, %cst : f32
399  // CHECK-NEXT: %2 = affine.load %0[%arg0] : memref<10xf32>
400  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
401  // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
402  // CHECK-NEXT: }
403  // CHECK-NEXT: }
404
405
406  return
407}
408
409func @invariant_affine_for_inside_affine_if() {
410  %m = alloc() : memref<10xf32>
411  %cf8 = constant 8.0 : f32
412  affine.for %arg0 = 0 to 10 {
413    affine.for %arg1 = 0 to 10 {
414      affine.if (d0, d1) : (d1 - d0 >= 0) (%arg0, %arg0) {
415          %cf9 = addf %cf8, %cf8 : f32
416          affine.store %cf9, %m[%arg0] : memref<10xf32>
417          affine.for %arg2 = 0 to 10 {
418            affine.store %cf9, %m[%arg2] : memref<10xf32>
419          }
420      }
421    }
422  }
423
424  // CHECK: %0 = alloc() : memref<10xf32>
425  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
426  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
427  // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
428  // CHECK-NEXT: affine.if #set0(%arg0, %arg0) {
429  // CHECK-NEXT: %1 = addf %cst, %cst : f32
430  // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
431  // CHECK-NEXT: affine.for %arg2 = 0 to 10 {
432  // CHECK-NEXT: affine.store %1, %0[%arg2] : memref<10xf32>
433  // CHECK-NEXT: }
434  // CHECK-NEXT: }
435  // CHECK-NEXT: }
436
437
438  return
439}
440
441
442func @invariant_constant_and_load() {
443  %m = alloc() : memref<100xf32>
444  %m2 = alloc() : memref<100xf32>
445  affine.for %arg0 = 0 to 5 {
446    %c0 = constant 0 : index
447    %v = affine.load %m2[%c0] : memref<100xf32>
448    affine.store %v, %m[%arg0] : memref<100xf32>
449  }
450
451  // CHECK: %0 = alloc() : memref<100xf32>
452  // CHECK-NEXT: %1 = alloc() : memref<100xf32>
453  // CHECK-NEXT: %c0 = constant 0 : index
454  // CHECK-NEXT: %2 = affine.load %1[%c0] : memref<100xf32>
455  // CHECK-NEXT: affine.for %arg0 = 0 to 5 {
456  // CHECK-NEXT:  affine.store %2, %0[%arg0] : memref<100xf32>
457
458
459  return
460}
461
462
463func @nested_load_store_same_memref() {
464  %m = alloc() : memref<10xf32>
465  %cst = constant 8.0 : f32
466  %c0 = constant 0 : index
467   affine.for %arg0 = 0 to 10 {
468    %v0 = affine.load %m[%c0] : memref<10xf32>
469    affine.for %arg1 = 0 to 10 {
470      affine.store %cst, %m[%arg1] : memref<10xf32>
471    }
472  }
473
474  // CHECK: %0 = alloc() : memref<10xf32>
475  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
476  // CHECK-NEXT: %c0 = constant 0 : index
477  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
478  // CHECK-NEXT:  %1 = affine.load %0[%c0] : memref<10xf32>
479  // CHECK-NEXT:   affine.for %arg1 = 0 to 10 {
480  // CHECK-NEXT:    affine.store %cst, %0[%arg1] : memref<10xf32>
481
482
483  return
484}
485
486
487func @nested_load_store_same_memref2() {
488  %m = alloc() : memref<10xf32>
489  %cst = constant 8.0 : f32
490  %c0 = constant 0 : index
491   affine.for %arg0 = 0 to 10 {
492     affine.store %cst, %m[%c0] : memref<10xf32>
493      affine.for %arg1 = 0 to 10 {
494        %v0 = affine.load %m[%arg0] : memref<10xf32>
495    }
496  }
497
498  // CHECK: %0 = alloc() : memref<10xf32>
499  // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
500  // CHECK-NEXT: %c0 = constant 0 : index
501  // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
502  // CHECK-NEXT:   affine.store %cst, %0[%c0] : memref<10xf32>
503  // CHECK-NEXT:   %1 = affine.load %0[%arg0] : memref<10xf32>
504
505
506  return
507}
508