c1bccea4c9f5f669f75cc3dbd8af18293761b831
[lldb.git] / mlir / test / Transforms / loop-fusion.mlir
1 // RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-fusion -split-input-file | FileCheck %s
2 // RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-fusion="fusion-maximal" -split-input-file | FileCheck %s --check-prefix=MAXIMAL
3
4 // TODO: Add more tests:
5 // *) Add nested fusion test cases when non-constant loop bound support is
6 //    added to iteration domain in dependence check.
7 // *) Add a test w/ floordiv/ceildiv/mod when supported in dependence check.
8 // *) Add tests which check fused computation slice indexing and loop bounds.
9 // TODO: Test clean up: move memref allocs to func args.
10
11 // -----
12
13 // CHECK-LABEL: func @should_fuse_raw_dep_for_locality() {
14 func @should_fuse_raw_dep_for_locality() {
15   %m = alloc() : memref<10xf32>
16   %cf7 = constant 7.0 : f32
17
18   affine.for %i0 = 0 to 10 {
19     affine.store %cf7, %m[%i0] : memref<10xf32>
20   }
21   affine.for %i1 = 0 to 10 {
22     %v0 = affine.load %m[%i1] : memref<10xf32>
23   }
24   // CHECK:      affine.for %{{.*}} = 0 to 10 {
25   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
26   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
27   // CHECK-NEXT: }
28   // CHECK-NEXT: return
29   return
30 }
31
32 // -----
33
34 // CHECK-LABEL: func @should_fuse_reduction_to_pointwise() {
35 func @should_fuse_reduction_to_pointwise() {
36   %a = alloc() : memref<10x10xf32>
37   %b = alloc() : memref<10xf32>
38   %c = alloc() : memref<10xf32>
39
40   %cf7 = constant 7.0 : f32
41
42   affine.for %i0 = 0 to 10 {
43     affine.for %i1 = 0 to 10 {
44       %v0 = affine.load %b[%i0] : memref<10xf32>
45       %v1 = affine.load %a[%i0, %i1] : memref<10x10xf32>
46       %v3 = addf %v0, %v1 : f32
47       affine.store %v3, %b[%i0] : memref<10xf32>
48     }
49   }
50   affine.for %i2 = 0 to 10 {
51     %v4 = affine.load %b[%i2] : memref<10xf32>
52     affine.store %v4, %c[%i2] : memref<10xf32>
53   }
54
55   // Should fuse in entire inner loop on %i1 from source loop nest, as %i1
56   // is not used in the access function of the store/load on %b.
57   // CHECK:       affine.for %{{.*}} = 0 to 10 {
58   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
59   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
60   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
61   // CHECK-NEXT:      addf %{{.*}}, %{{.*}} : f32
62   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
63   // CHECK-NEXT:    }
64   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
65   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
66   // CHECK-NEXT:  }
67   // CHECK-NEXT:  return
68   return
69 }
70
71 // -----
72
73 // CHECK-DAG: [[$MAP_SHIFT_MINUS_ONE_R1:#map[0-9]+]] = affine_map<(d0) -> (d0 - 1)>
74 // CHECK-DAG: [[$MAP_SHIFT_D0_BY_ONE:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 + 1)>
75 // CHECK-DAG: [[$MAP_SHIFT_D1_BY_ONE:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + 1)>
76
77 // CHECK-LABEL: func @should_fuse_loop_nests_with_shifts() {
78 func @should_fuse_loop_nests_with_shifts() {
79   %a = alloc() : memref<10x10xf32>
80   %cf7 = constant 7.0 : f32
81
82   affine.for %i0 = 0 to 9 {
83     affine.for %i1 = 0 to 9 {
84       affine.store %cf7, %a[%i0 + 1, %i1 + 1] : memref<10x10xf32>
85     }
86   }
87   affine.for %i2 = 1 to 10 {
88     affine.for %i3 = 1 to 10 {
89       %v0 = affine.load %a[%i2, %i3] : memref<10x10xf32>
90     }
91   }
92
93   // Source slice affine apply sequence:
94   // *) First two affine apply's map from the dst to src iteration space.
95   // *) Third affine apply is access function around src store.
96   // *) Fourth affine apply shifts the stores access function by '-1', because
97   //    of the offset induced by reducing the memref shape from 10x10 to 9x9.
98   // *) Fifth affine apply shifts the loads access function by '-1', because
99   //    of the offset induced by reducing the memref shape from 10x10 to 9x9.
100   // NOTE: Should create a private memref with reduced shape 9x9xf32.
101   // CHECK:      affine.for %{{.*}} = 1 to 10 {
102   // CHECK-NEXT:   affine.for %{{.*}} = 1 to 10 {
103   // CHECK-NEXT:     %[[I:.*]] = affine.apply [[$MAP_SHIFT_MINUS_ONE_R1]](%{{.*}})
104   // CHECK-NEXT:     %[[J:.*]] = affine.apply [[$MAP_SHIFT_MINUS_ONE_R1]](%{{.*}})
105   // CHECK-NEXT:     affine.apply [[$MAP_SHIFT_D0_BY_ONE]](%[[I]], %[[J]])
106   // CHECK-NEXT:     affine.apply [[$MAP_SHIFT_D1_BY_ONE]](%[[I]], %[[J]])
107   // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
108   // CHECK-NEXT:     affine.load %{{.*}}[0, 0] : memref<1x1xf32>
109   // CHECK-NEXT:   }
110   // CHECK-NEXT: }
111   // CHECK-NEXT: return
112   return
113 }
114
115 // -----
116
117 // CHECK-LABEL: func @should_fuse_loop_nest() {
118 func @should_fuse_loop_nest() {
119   %a = alloc() : memref<10x10xf32>
120   %b = alloc() : memref<10x10xf32>
121   %cf7 = constant 7.0 : f32
122
123   affine.for %i0 = 0 to 10 {
124     affine.for %i1 = 0 to 10 {
125       affine.store %cf7, %a[%i0, %i1] : memref<10x10xf32>
126     }
127   }
128   affine.for %i2 = 0 to 10 {
129     affine.for %i3 = 0 to 10 {
130       %v0 = affine.load %a[%i3, %i2] : memref<10x10xf32>
131       affine.store %v0, %b[%i2, %i3] : memref<10x10xf32>
132     }
133   }
134   affine.for %i4 = 0 to 10 {
135     affine.for %i5 = 0 to 10 {
136       %v1 = affine.load %b[%i4, %i5] : memref<10x10xf32>
137     }
138   }
139   // Expecting private memref for '%a' first, then private memref for '%b'.
140   // CHECK-DAG:  [[NEWA:%[0-9]+]] = alloc() : memref<1x1xf32>
141   // CHECK-DAG:  [[NEWB:%[0-9]+]] = alloc() : memref<1x1xf32>
142   // CHECK:      affine.for %{{.*}} = 0 to 10 {
143   // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10 {
144   // CHECK-NEXT:     affine.store %{{.*}}, [[NEWA]][0, 0] : memref<1x1xf32>
145   // CHECK-NEXT:     affine.load [[NEWA]][0, 0] : memref<1x1xf32>
146   // CHECK-NEXT:     affine.store %{{.*}}, [[NEWB]][0, 0] : memref<1x1xf32>
147   // CHECK-NEXT:     affine.load [[NEWB]][0, 0] : memref<1x1xf32>
148   // CHECK-NEXT:   }
149   // CHECK-NEXT: }
150   // CHECK-NEXT: return
151   return
152 }
153
154 // -----
155
156 // CHECK-LABEL: func @should_fuse_across_intermediate_loop_with_no_deps() {
157 func @should_fuse_across_intermediate_loop_with_no_deps() {
158   %a = alloc() : memref<10xf32>
159   %b = alloc() : memref<10xf32>
160   %c = alloc() : memref<10xf32>
161
162   %cf7 = constant 7.0 : f32
163
164   affine.for %i0 = 0 to 10 {
165     %v0 = affine.load %a[%i0] : memref<10xf32>
166     affine.store %v0, %b[%i0] : memref<10xf32>
167   }
168   affine.for %i1 = 0 to 10 {
169     affine.store %cf7, %c[%i1] : memref<10xf32>
170   }
171   affine.for %i2 = 0 to 10 {
172     %v1 = affine.load %b[%i2] : memref<10xf32>
173   }
174
175   // Should fuse first loop (past second loop with no dependences) into third.
176   // Note that fusion creates a private memref '%2' for the fused loop nest.
177   // CHECK:      affine.for %{{.*}} = 0 to 10 {
178   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
179   // CHECK-NEXT: }
180   // CHECK:      affine.for %{{.*}} = 0 to 10 {
181   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
182   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
183   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
184   // CHECK-NEXT: }
185   // CHECK-NEXT: return
186   return
187 }
188
189 // -----
190
191 // CHECK-LABEL: func @should_fuse_all_loops() {
192 func @should_fuse_all_loops() {
193   %a = alloc() : memref<10xf32>
194   %b = alloc() : memref<10xf32>
195   %cf7 = constant 7.0 : f32
196
197   // Set up flow dependences from first and second loops to third.
198   affine.for %i0 = 0 to 10 {
199     affine.store %cf7, %a[%i0] : memref<10xf32>
200   }
201   affine.for %i1 = 0 to 10 {
202     affine.store %cf7, %b[%i1] : memref<10xf32>
203   }
204   affine.for %i2 = 0 to 10 {
205     %v0 = affine.load %a[%i2] : memref<10xf32>
206     %v1 = affine.load %b[%i2] : memref<10xf32>
207   }
208
209   // Should fuse first and second loops into third.
210   // Expecting private memref for '%a' first, then private memref for '%b'.
211   // CHECK-DAG: [[NEWA:%[0-9]+]] = alloc() : memref<1xf32>
212   // CHECK-DAG: [[NEWB:%[0-9]+]] = alloc() : memref<1xf32>
213   // CHECK:      affine.for %{{.*}} = 0 to 10 {
214   // CHECK-NEXT:   affine.store %{{.*}}, [[NEWA]][0] : memref<1xf32>
215   // CHECK-NEXT:   affine.store %{{.*}}, [[NEWB]][0] : memref<1xf32>
216   // CHECK-NEXT:   affine.load [[NEWA]][0] : memref<1xf32>
217   // CHECK-NEXT:   affine.load [[NEWB]][0] : memref<1xf32>
218   // CHECK-NEXT: }
219   // CHECK-NEXT: return
220   return
221 }
222
223 // -----
224
225 // CHECK-LABEL: func @should_fuse_first_and_second_loops() {
226 func @should_fuse_first_and_second_loops() {
227   %a = alloc() : memref<10xf32>
228   %b = alloc() : memref<10xf32>
229   %c = alloc() : memref<10xf32>
230
231   %cf7 = constant 7.0 : f32
232
233   affine.for %i0 = 0 to 10 {
234     affine.store %cf7, %a[%i0] : memref<10xf32>
235   }
236   affine.for %i1 = 0 to 10 {
237     %v0 = affine.load %a[%i1] : memref<10xf32>
238     affine.store %cf7, %b[%i1] : memref<10xf32>
239   }
240   affine.for %i2 = 0 to 10 {
241     %v1 = affine.load %c[%i2] : memref<10xf32>
242   }
243
244   // Should fuse first loop into the second (last loop should not be fused).
245   // Should create private memref '%2' for fused scf.
246   // CHECK:      affine.for %{{.*}} = 0 to 10 {
247   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
248   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
249   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
250   // CHECK-NEXT: }
251   // CHECK:      affine.for %{{.*}} = 0 to 10 {
252   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
253   // CHECK-NEXT: }
254   // CHECK-NEXT: return
255
256   return
257 }
258
259 // -----
260
261 // CHECK-LABEL: func @should_not_fuse_would_create_cycle() {
262 func @should_not_fuse_would_create_cycle() {
263   %a = alloc() : memref<10xf32>
264   %b = alloc() : memref<10xf32>
265   %c = alloc() : memref<10xf32>
266
267   %cf7 = constant 7.0 : f32
268
269   // Set up the following dependences:
270   // 1) loop0 -> loop1 on memref '%{{.*}}'
271   // 2) loop0 -> loop2 on memref '%{{.*}}'
272   // 3) loop1 -> loop2 on memref '%{{.*}}'
273   affine.for %i0 = 0 to 10 {
274     %v0 = affine.load %a[%i0] : memref<10xf32>
275     affine.store %cf7, %b[%i0] : memref<10xf32>
276   }
277   affine.for %i1 = 0 to 10 {
278     affine.store %cf7, %a[%i1] : memref<10xf32>
279     %v1 = affine.load %c[%i1] : memref<10xf32>
280   }
281   affine.for %i2 = 0 to 10 {
282     %v2 = affine.load %b[%i2] : memref<10xf32>
283     affine.store %cf7, %c[%i2] : memref<10xf32>
284   }
285   // Should not fuse: fusing loop first loop into last would create a cycle.
286   // CHECK:      affine.for %{{.*}} = 0 to 10 {
287   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
288   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
289   // CHECK-NEXT: }
290   // CHECK:      affine.for %{{.*}} = 0 to 10 {
291   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
292   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
293   // CHECK-NEXT: }
294   // CHECK:      affine.for %{{.*}} = 0 to 10 {
295   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
296   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
297   // CHECK-NEXT: }
298   // CHECK-NEXT: return
299   return
300 }
301
302 // -----
303
304 // CHECK-LABEL: func @should_fuse_producer_consumer() {
305 func @should_fuse_producer_consumer() {
306   %m = alloc() : memref<10xf32>
307   %cf7 = constant 7.0 : f32
308
309   affine.for %i0 = 0 to 10 {
310     affine.store %cf7, %m[%i0] : memref<10xf32>
311   }
312   affine.for %i1 = 0 to 10 {
313     affine.store %cf7, %m[%i1] : memref<10xf32>
314   }
315   affine.for %i2 = 0 to 10 {
316     %v1 = affine.load %m[%i2] : memref<10xf32>
317   }
318   // Fusing loop %i0 to %i2 would violate the WAW dependence between %i0 and
319   // %i1, but OK to fuse %i1 into %i2.
320   // TODO: When the fusion pass is run to a fixed-point, it should
321   // fuse all three of these loop nests.
322   // CHECK:      alloc() : memref<1xf32>
323   // CHECK:      affine.for %{{.*}} = 0 to 10 {
324   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
325   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
326   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
327   // CHECK-NEXT: }
328   // CHECK-NEXT: return
329   return
330 }
331
332 // -----
333
334 // CHECK-LABEL: func @should_fuse_and_move_to_preserve_war_dep() {
335 func @should_fuse_and_move_to_preserve_war_dep() {
336   %a = alloc() : memref<10xf32>
337   %b = alloc() : memref<10xf32>
338   %cf7 = constant 7.0 : f32
339
340   affine.for %i0 = 0 to 10 {
341     %v0 = affine.load %a[%i0] : memref<10xf32>
342     affine.store %v0, %b[%i0] : memref<10xf32>
343   }
344   affine.for %i1 = 0 to 10 {
345     affine.store %cf7, %a[%i1] : memref<10xf32>
346   }
347   affine.for %i2 = 0 to 10 {
348     %v1 = affine.load %b[%i2] : memref<10xf32>
349   }
350   // Loops '%i1' and '%i2' have no dependences. We can fuse a slice of '%i0'
351   // into '%i2' if we move the fused loop nest before '%i1', which preserves
352   // the WAR dependence from load '%a' in '%i0' to the store '%a' in loop '%i1'.
353   // CHECK:       affine.for %{{.*}} = 0 to 10 {
354   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
355   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
356   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
357   // CHECK-NEXT:  }
358   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
359   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
360   // CHECK-NEXT:  }
361   // CHECK-NEXT:  return
362   return
363 }
364
365 // -----
366
367 // CHECK-LABEL: func @should_fuse_if_top_level_access() {
368 func @should_fuse_if_top_level_access() {
369   %m = alloc() : memref<10xf32>
370   %cf7 = constant 7.0 : f32
371
372   affine.for %i0 = 0 to 10 {
373     affine.store %cf7, %m[%i0] : memref<10xf32>
374   }
375   affine.for %i1 = 0 to 10 {
376     %v0 = affine.load %m[%i1] : memref<10xf32>
377   }
378
379   %c0 = constant 4 : index
380   %v1 = affine.load %m[%c0] : memref<10xf32>
381   // Top-level load to '%m' should prevent creating a private memref but
382   // loop nests should be fused and '%i0' should be removed.
383   // CHECK:      %[[m:.*]] = alloc() : memref<10xf32>
384   // CHECK-NOT:  alloc
385
386   // CHECK:      affine.for %[[i1:.*]] = 0 to 10 {
387   // CHECK-NEXT:   affine.store %{{.*}}, %[[m]][%[[i1]]] : memref<10xf32>
388   // CHECK-NEXT:   affine.load %[[m]][%[[i1]]] : memref<10xf32>
389   // CHECK-NEXT: }
390   // CHECK:      affine.load %[[m]][%{{.*}}] : memref<10xf32>
391   return
392 }
393
394 // -----
395
396 // CHECK-LABEL: func @should_fuse_but_not_remove_src() {
397 func @should_fuse_but_not_remove_src() {
398   %m = alloc() : memref<100xf32>
399   %cf7 = constant 7.0 : f32
400
401   affine.for %i0 = 0 to 100 {
402     affine.store %cf7, %m[%i0] : memref<100xf32>
403   }
404   affine.for %i1 = 0 to 17 {
405     %v0 = affine.load %m[%i1] : memref<100xf32>
406   }
407   %v1 = affine.load %m[99] : memref<100xf32>
408
409   // Loop '%i0' and '%i1' should be fused but '%i0' shouldn't be removed to
410   // preserve the dependence with the top-level access.
411   // CHECK:      affine.for %{{.*}} = 0 to 100 {
412   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<100xf32>
413   // CHECK-NEXT: }
414   // CHECK-NEXT: affine.for %{{.*}} = 0 to 17 {
415   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
416   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
417   // CHECK-NEXT: }
418   // CHECK-NEXT: affine.load %{{.*}}[99] : memref<100xf32>
419   // CHECK-NEXT: return
420   return
421 }
422
423 // -----
424
425 // CHECK-LABEL: func @should_fuse_no_top_level_access() {
426 func @should_fuse_no_top_level_access() {
427   %m = alloc() : memref<10xf32>
428   %cf7 = constant 7.0 : f32
429
430   affine.for %i0 = 0 to 10 {
431     affine.store %cf7, %m[%i0] : memref<10xf32>
432   }
433   affine.for %i1 = 0 to 10 {
434     %v0 = affine.load %m[%i1] : memref<10xf32>
435   }
436   // CHECK:      affine.for %{{.*}} = 0 to 10 {
437   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
438   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
439   // CHECK-NEXT: }
440   // CHECK-NEXT: return
441   return
442 }
443
444 // -----
445
446 #set0 = affine_set<(d0) : (1 == 0)>
447
448 // CHECK-LABEL: func @should_not_fuse_if_inst_at_top_level() {
449 func @should_not_fuse_if_inst_at_top_level() {
450   %m = alloc() : memref<10xf32>
451   %cf7 = constant 7.0 : f32
452
453   affine.for %i0 = 0 to 10 {
454     affine.store %cf7, %m[%i0] : memref<10xf32>
455   }
456   affine.for %i1 = 0 to 10 {
457     %v0 = affine.load %m[%i1] : memref<10xf32>
458   }
459   %c0 = constant 4 : index
460   affine.if #set0(%c0) {
461   }
462   // Top-level IfOp should prevent fusion.
463   // CHECK:      affine.for %{{.*}} = 0 to 10 {
464   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
465   // CHECK-NEXT: }
466   // CHECK:      affine.for %{{.*}} = 0 to 10 {
467   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
468   // CHECK-NEXT: }
469   return
470 }
471
472 // -----
473
474 #set0 = affine_set<(d0) : (1 == 0)>
475
476 // CHECK-LABEL: func @should_not_fuse_if_inst_in_loop_nest() {
477 func @should_not_fuse_if_inst_in_loop_nest() {
478   %m = alloc() : memref<10xf32>
479   %cf7 = constant 7.0 : f32
480   %c4 = constant 4 : index
481
482   affine.for %i0 = 0 to 10 {
483     affine.store %cf7, %m[%i0] : memref<10xf32>
484   }
485   affine.for %i1 = 0 to 10 {
486     affine.if #set0(%c4) {
487     }
488     %v0 = affine.load %m[%i1] : memref<10xf32>
489   }
490
491   // IfOp in ForInst should prevent fusion.
492   // CHECK:      affine.for %{{.*}} = 0 to 10 {
493   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
494   // CHECK-NEXT: }
495   // CHECK:      affine.for %{{.*}} = 0 to 10 {
496   // CHECK-NEXT:   affine.if #set(%{{.*}}) {
497   // CHECK-NEXT:   }
498   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
499   // CHECK-NEXT: }
500   return
501 }
502
503 // -----
504
505 // CHECK-LABEL: func @permute_and_fuse() {
506 func @permute_and_fuse() {
507   %m = alloc() : memref<10x20x30xf32>
508
509   %cf7 = constant 7.0 : f32
510   affine.for %i0 = 0 to 10 {
511     affine.for %i1 = 0 to 20 {
512       affine.for %i2 = 0 to 30 {
513         affine.store %cf7, %m[%i0, %i1, %i2] : memref<10x20x30xf32>
514       }
515     }
516   }
517   affine.for %i3 = 0 to 30 {
518     affine.for %i4 = 0 to 10 {
519       affine.for %i5 = 0 to 20 {
520         %v0 = affine.load %m[%i4, %i5, %i3] : memref<10x20x30xf32>
521         "foo"(%v0) : (f32) -> ()
522       }
523     }
524   }
525 // CHECK:       affine.for %{{.*}} = 0 to 30 {
526 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
527 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 20 {
528 // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
529 // CHECK-NEXT:        affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
530 // CHECK-NEXT:        "foo"(%{{.*}}) : (f32) -> ()
531 // CHECK-NEXT:      }
532 // CHECK-NEXT:    }
533 // CHECK-NEXT:  }
534 // CHECK-NEXT:  return
535
536   return
537 }
538
539 // -----
540
541 // CHECK-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 4 + d1)>
542 // CHECK-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 floordiv 4)>
543 // CHECK-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0) -> (d0 mod 4)>
544
545 // Reshape from a 64 x f32 to 16 x 4 x f32.
546 // CHECK-LABEL: func @fuse_reshape_64_16_4
547 func @fuse_reshape_64_16_4(%in : memref<64xf32>) {
548   %out = alloc() : memref<16x4xf32>
549
550   affine.for %i0 = 0 to 64 {
551     %v = affine.load %in[%i0] : memref<64xf32>
552     affine.store %v, %out[%i0 floordiv 4, %i0 mod 4] : memref<16x4xf32>
553   }
554
555   affine.for %i1 = 0 to 16 {
556     affine.for %i2 = 0 to 4 {
557       %w = affine.load %out[%i1, %i2] : memref<16x4xf32>
558       "foo"(%w) : (f32) -> ()
559     }
560   }
561   return
562   // CHECK:      affine.for %{{.*}} =
563   // CHECK-NEXT:   affine.for %{{.*}} =
564   // CHECK-NOT:    for
565   // CHECK:        }
566   // CHECK-NEXT: }
567   // CHECK-NEXT: return
568 }
569
570 // -----
571 // CHECK-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 floordiv 4)>
572 // CHECK-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 mod 4)>
573 // CHECK-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 4 + d1)>
574
575 // Reshape a 16x4xf32 to 64xf32.
576 // CHECK-LABEL: func @fuse_reshape_16_4_64
577 func @fuse_reshape_16_4_64() {
578   %in = alloc() : memref<16x4xf32>
579   %out = alloc() : memref<64xf32>
580
581   affine.for %i0 = 0 to 16 {
582     affine.for %i1 = 0 to 4 {
583       %v = affine.load %in[%i0, %i1] : memref<16x4xf32>
584       affine.store %v, %out[4*%i0 + %i1] : memref<64xf32>
585     }
586   }
587
588   affine.for %i2 = 0 to 64 {
589     %w = affine.load %out[%i2] : memref<64xf32>
590     "foo"(%w) : (f32) -> ()
591   }
592 // CHECK:       affine.for %{{.*}} = 0 to 64 {
593 // CHECK-NEXT:    affine.apply [[$MAP0]](%{{.*}})
594 // CHECK-NEXT:    affine.apply [[$MAP1]](%{{.*}})
595 // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<16x4xf32>
596 // CHECK-NEXT:    affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
597 // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
598 // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
599 // CHECK-NEXT:    "foo"(%{{.*}}) : (f32) -> ()
600 // CHECK-NEXT:  }
601 // CHECK-NEXT:  return
602   return
603 }
604
605
606 // -----
607
608 // All three loop nests below (6-d one, 2-d one, 2-d one is fused into a single
609 // 2-d loop nest).
610 func @R6_to_R2_reshape_square() -> memref<64x9xi32> {
611   %in = alloc() : memref<2x2x3x3x16x1xi32>
612   %out = alloc() : memref<64x9xi32>
613   %live_out = alloc() : memref<64x9xi32>
614
615   // Initialize input.
616   affine.for %i0 = 0 to 2 {
617     affine.for %i1 = 0 to 2 {
618       affine.for %i2 = 0 to 3 {
619         affine.for %i3 = 0 to 3 {
620           affine.for %i4 = 0 to 16 {
621             affine.for %i5 = 0 to 1 {
622               %val = "foo"(%i0, %i1, %i2, %i3, %i4, %i5) : (index, index, index, index, index, index) -> i32
623               affine.store %val, %in[%i0, %i1, %i2, %i3, %i4, %i5] : memref<2x2x3x3x16x1xi32>
624             }
625           }
626         }
627       }
628     }
629   }
630
631   affine.for %ii = 0 to 64 {
632     affine.for %jj = 0 to 9 {
633       // Convert output coordinates to linear index.
634       %a0 = affine.apply affine_map<(d0, d1) -> (d0 * 9 + d1)> (%ii, %jj)
635       %0 = affine.apply affine_map<(d0) -> (d0 floordiv (2 * 3 * 3 * 16 * 1))>(%a0)
636       %1 = affine.apply affine_map<(d0) -> ((d0 mod 288) floordiv (3 * 3 * 16 * 1))>(%a0)
637       %2 = affine.apply affine_map<(d0) -> (((d0 mod 288) mod 144) floordiv (3 * 16 * 1))>(%a0)
638       %3 = affine.apply affine_map<(d0) -> ((((d0 mod 288) mod 144) mod 48) floordiv (16 * 1))>(%a0)
639       %4 = affine.apply affine_map<(d0) -> ((((d0 mod 288) mod 144) mod 48) mod 16)>(%a0)
640       %5 = affine.apply affine_map<(d0) -> (((((d0 mod 144) mod 144) mod 48) mod 16) mod 1)>(%a0)
641       %v = affine.load %in[%0, %1, %2, %3, %4, %5] : memref<2x2x3x3x16x1xi32>
642       affine.store %v, %out[%ii, %jj] : memref<64x9xi32>
643     }
644   }
645
646   affine.for %i = 0 to 64 {
647     affine.for %j = 0 to 9 {
648       %a = affine.load %out[%i, %j] : memref<64x9xi32>
649       %b = muli %a, %a : i32
650       affine.store %b, %live_out[%i, %j] : memref<64x9xi32>
651     }
652   }
653   return %live_out : memref<64x9xi32>
654 }
655 // Everything above is fused to a single 2-d loop nest, and the 6-d tensor %in
656 // is eliminated if -memref-dataflow-opt is also supplied.
657 //
658 // CHECK-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> ((d0 * 9 + d1) floordiv 288)>
659 // CHECK-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0, d1) -> (((d0 * 9 + d1) mod 288) floordiv 144)>
660 // CHECK-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0, d1) -> ((((d0 * 9 + d1) mod 288) mod 144) floordiv 48)>
661 // CHECK-DAG: [[$MAP3:#map[0-9]+]] = affine_map<(d0, d1) -> (((((d0 * 9 + d1) mod 288) mod 144) mod 48) floordiv 16)>
662 // CHECK-DAG: [[$MAP4:#map[0-9]+]] = affine_map<(d0, d1) -> (((((d0 * 9 + d1) mod 288) mod 144) mod 48) mod 16)>
663 // CHECK-DAG: [[$MAP11:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 9 + d1)>
664 // CHECK-DAG: [[$MAP12:#map[0-9]+]] = affine_map<(d0) -> (d0 floordiv 288)>
665 // CHECK-DAG: [[$MAP13:#map[0-9]+]] = affine_map<(d0) -> ((d0 mod 288) floordiv 144)>
666 // CHECK-DAG: [[$MAP14:#map[0-9]+]] = affine_map<(d0) -> (((d0 mod 288) mod 144) floordiv 48)>
667 // CHECK-DAG: [[$MAP15:#map[0-9]+]] = affine_map<(d0) -> ((((d0 mod 288) mod 144) mod 48) floordiv 16)>
668 // CHECK-DAG: [[$MAP16:#map[0-9]+]] = affine_map<(d0) -> ((((d0 mod 288) mod 144) mod 48) mod 16)>
669 // CHECK-DAG: [[$MAP17:#map[0-9]+]] = affine_map<(d0) -> (0)>
670
671 //
672 // CHECK-LABEL: func @R6_to_R2_reshape
673 // CHECK:       alloc() : memref<1x2x3x3x16x1xi32>
674 // CHECK:       alloc() : memref<1x1xi32>
675 // CHECK:       alloc() : memref<64x9xi32>
676 // CHECK-NEXT:  affine.for %{{.*}} = 0 to 64 {
677 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 9 {
678 // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
679 // CHECK-NEXT:      affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
680 // CHECK-NEXT:      affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
681 // CHECK-NEXT:      affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
682 // CHECK-NEXT:      affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
683 // CHECK-NEXT:      "foo"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index, index, index, index) -> i32
684 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, ((%{{.*}} * 9 + %{{.*}}) mod 288) floordiv 144, (((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) floordiv 48, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) floordiv 16, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) mod 16, 0] : memref<1x2x3x3x16x1xi32>
685 // CHECK-NEXT:      affine.apply [[$MAP11]](%{{.*}}, %{{.*}})
686 // CHECK-NEXT:      affine.apply [[$MAP12]](%{{.*}})
687 // CHECK-NEXT:      affine.apply [[$MAP13]](%{{.*}})
688 // CHECK-NEXT:      affine.apply [[$MAP14]](%{{.*}})
689 // CHECK-NEXT:      affine.apply [[$MAP15]](%{{.*}})
690 // CHECK-NEXT:      affine.apply [[$MAP16]](%{{.*}})
691 // CHECK-NEXT:      affine.apply [[$MAP17]](%{{.*}})
692 // CHECK-NEXT:      affine.load %{{.*}}[0, ((%{{.*}} * 9 + %{{.*}}) mod 288) floordiv 144, (((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) floordiv 48, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) floordiv 16, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) mod 16, 0] : memref<1x2x3x3x16x1xi32>
693 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xi32>
694 // CHECK-NEXT:      affine.load %{{.*}}[0, 0] : memref<1x1xi32>
695 // CHECK-NEXT:      muli %{{.*}}, %{{.*}} : i32
696 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x9xi32>
697 // CHECK-NEXT:    }
698 // CHECK-NEXT:  }
699 // CHECK-NEXT:  return %{{.*}} : memref<64x9xi32>
700
701 // -----
702
703 // CHECK-LABEL: func @fuse_symbolic_bounds
704 func @fuse_symbolic_bounds(%M : index, %N : index) {
705   %N_plus_5 = affine.apply affine_map<(d0) -> (d0 + 5)>(%N)
706   %m = alloc(%M, %N_plus_5) : memref<? x ? x f32>
707
708   %c0 = constant 0.0 : f32
709   %s = constant 5 : index
710
711   affine.for %i0 = 0 to %M {
712     affine.for %i1 = 0 to affine_map<(d0) -> (d0 + 5)> (%N) {
713       affine.store %c0, %m[%i0, %i1] : memref<? x ? x f32>
714     }
715   }
716
717   affine.for %i2 = 0 to %M {
718     affine.for %i3 = 0 to %N {
719       %v = affine.load %m[%i2, %i3 + symbol(%s)] : memref<? x ? x f32>
720     }
721   }
722
723   return
724 }
725
726 // -----
727
728 // CHECK-LABEL: func @should_fuse_reduction_at_depth_of_one
729 func @should_fuse_reduction_at_depth_of_one() {
730   %a = alloc() : memref<10x100xf32>
731   %b = alloc() : memref<10xf32>
732
733   affine.for %i0 = 0 to 10 {
734     affine.for %i1 = 0 to 100 {
735       %v0 = affine.load %b[%i0] : memref<10xf32>
736       %v1 = affine.load %a[%i0, %i1] : memref<10x100xf32>
737       %v2 = "maxf"(%v0, %v1) : (f32, f32) -> f32
738       affine.store %v2, %b[%i0] : memref<10xf32>
739     }
740   }
741   affine.for %i2 = 0 to 10 {
742     affine.for %i3 = 0 to 100 {
743       %v3 = affine.load %b[%i2] : memref<10xf32>
744       %v4 = affine.load %a[%i2, %i3] : memref<10x100xf32>
745       %v5 = subf %v4, %v3 : f32
746       affine.store %v5, %b[%i2] : memref<10xf32>
747     }
748   }
749   // This test should fuse the src reduction loop at depth 1 in the destination
750   // loop nest, which improves locality and enables subsequence passes to
751   // decrease the reduction memref size and possibly place it in a faster
752   // memory space.
753   // CHECK:       affine.for %{{.*}} = 0 to 10 {
754   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 100 {
755   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
756   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x100xf32>
757   // CHECK-NEXT:      "maxf"(%{{.*}}, %{{.*}}) : (f32, f32) -> f32
758   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
759   // CHECK-NEXT:    }
760   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 100 {
761   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
762   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x100xf32>
763   // CHECK-NEXT:      subf %{{.*}}, %{{.*}} : f32
764   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
765   // CHECK-NEXT:    }
766   // CHECK-NEXT:  }
767   // CHECK-NEXT:  return
768   return
769 }
770
771 // -----
772
773 // CHECK-LABEL: func @should_fuse_at_src_depth1_and_dst_depth1
774 func @should_fuse_at_src_depth1_and_dst_depth1() {
775   %a = alloc() : memref<100x16xf32>
776   %b = alloc() : memref<100x16xf32>
777
778   affine.for %i0 = 0 to 100 {
779     affine.for %i1 = 0 to 16 {
780       %v0 = affine.load %a[%i0, %i1] : memref<100x16xf32>
781       "op0"(%v0) : (f32) -> ()
782     }
783     affine.for %i2 = 0 to 16 {
784       %v1 = "op1"() : () -> (f32)
785       affine.store %v1, %b[%i0, %i2] : memref<100x16xf32>
786     }
787   }
788
789   affine.for %i3 = 0 to 100 {
790     affine.for %i4 = 0 to 16 {
791       %v2 = affine.load %b[%i3, %i4] : memref<100x16xf32>
792       "op2"(%v2) : (f32) -> ()
793     }
794   }
795   // We can slice iterations of the '%i0' and '%i1' loops in the source
796   // loop nest, but slicing at depth 2 and inserting the slice in the
797   // destination loop nest at depth2 causes extra computation. Instead,
798   // the fusion algorithm should detect that the source loop should be sliced
799   // at depth 1 and the slice should be inserted at depth 1.
800   // CHECK:       affine.for %{{.*}} = 0 to 100 {
801   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
802   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<100x16xf32>
803   // CHECK-NEXT:      "op0"(%{{.*}}) : (f32) -> ()
804   // CHECK-NEXT:    }
805   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
806   // CHECK-NEXT:      %{{.*}} = "op1"() : () -> f32
807   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
808   // CHECK-NEXT:    }
809   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
810   // CHECK-NEXT:      affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
811   // CHECK-NEXT:      "op2"(%{{.*}}) : (f32) -> ()
812   // CHECK-NEXT:    }
813   // CHECK-NEXT:  }
814   // CHECK-NEXT:  return
815   return
816 }
817
818 // -----
819 // CHECK: [[$MAP0:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 10 + d1)>
820
821 // CHECK-LABEL: func @should_fuse_src_depth1_at_dst_depth2
822 func @should_fuse_src_depth1_at_dst_depth2() {
823   %a = alloc() : memref<100xf32>
824   %c0 = constant 0.0 : f32
825
826   affine.for %i0 = 0 to 100 {
827     affine.store %c0, %a[%i0] : memref<100xf32>
828   }
829
830   affine.for %i1 = 0 to 10 {
831     affine.for %i2 = 0 to 10 {
832       %a0 = affine.apply affine_map<(d0, d1) -> (d0 * 10 + d1)> (%i1, %i2)
833       %v0 = affine.load %a[%a0] : memref<100xf32>
834     }
835   }
836   // The source loop nest slice loop bound is a function of both destination
837   // loop IVs, so we should slice at depth 1 and insert the slice at depth 2.
838   // CHECK:       affine.for %{{.*}} = 0 to 10 {
839   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
840   // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
841   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
842   // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
843   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
844   // CHECK-NEXT:    }
845   // CHECK-NEXT:  }
846   // CHECK-NEXT:  return
847   return
848 }
849
850 // -----
851
852 // CHECK-LABEL: func @fusion_at_depth0_not_currently_supported
853 func @fusion_at_depth0_not_currently_supported() {
854   %0 = alloc() : memref<10xf32>
855   %c0 = constant 0 : index
856   %cst = constant 0.000000e+00 : f32
857   affine.for %i0 = 0 to 10 {
858     affine.store %cst, %0[%i0] : memref<10xf32>
859   }
860   affine.for %i1 = 0 to 10 {
861     %1 = affine.load %0[%c0] : memref<10xf32>
862   }
863   // NOTE: Should shrink memref size to 1 element access by load in dst loop
864   // nest, and make the store in the slice store to the same element.
865   // CHECK-DAG:   alloc() : memref<1xf32>
866   // CHECK:       affine.for %{{.*}} = 0 to 10 {
867   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
868   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
869   // CHECK-NEXT:  }
870   // CHECK-NEXT:  return
871   return
872 }
873
874 // -----
875
876 // CHECK-LABEL: func @should_fuse_deep_loop_nests
877 func @should_fuse_deep_loop_nests() {
878   %0 = alloc() : memref<2x2x3x3x16x10xf32, 2>
879   %1 = alloc() : memref<2x2x3x3x16x10xf32, 2>
880   %2 = alloc() : memref<3x3x3x3x16x10xf32, 2>
881   %c0 = constant 0 : index
882   %c1 = constant 1 : index
883   %c1_0 = constant 1 : index
884   %cst = constant 0.000000e+00 : f32
885   affine.for %i0 = 0 to 2 {
886     affine.for %i1 = 0 to 2 {
887       affine.for %i2 = 0 to 3 {
888         affine.for %i3 = 0 to 3 {
889           affine.for %i4 = 0 to 16 {
890             affine.for %i5 = 0 to 10 {
891               %3 = affine.load %0[%i0, %i1, %i2, %i3, %i4, %i5]
892                 : memref<2x2x3x3x16x10xf32, 2>
893             }
894           }
895           affine.for %i6 = 0 to 16 {
896             affine.for %i7 = 0 to 10 {
897               affine.store %cst, %1[%i0, %i1, %i2, %i3, %i6, %i7]
898                 : memref<2x2x3x3x16x10xf32, 2>
899             }
900           }
901         }
902       }
903     }
904   }
905   affine.for %i8 = 0 to 3 {
906     affine.for %i9 = 0 to 3 {
907       affine.for %i10 = 0 to 2 {
908         affine.for %i11 = 0 to 2 {
909           affine.for %i12 = 0 to 3 {
910             affine.for %i13 = 0 to 3 {
911               affine.for %i14 = 0 to 2 {
912                 affine.for %i15 = 0 to 2 {
913                   affine.for %i16 = 0 to 16 {
914                     affine.for %i17 = 0 to 10 {
915                       %5 = affine.load %0[%i14, %i15, %i12, %i13, %i16, %i17]
916                         : memref<2x2x3x3x16x10xf32, 2>
917                     }
918                   }
919                   affine.for %i18 = 0 to 16 {
920                     affine.for %i19 = 0 to 10 {
921                       %6 = affine.load %1[%i10, %i11, %i8, %i9, %i18, %i19]
922                         : memref<2x2x3x3x16x10xf32, 2>
923                     }
924                   }
925                 }
926               }
927             }
928           }
929         }
930       }
931     }
932   }
933 // The first four loops of the source loop nest can be sliced with iteration
934 // bounds which are a function of the first four loops of destination loop nest,
935 // where the destination loops nests have been interchanged.
936
937 // CHECK-DAG:   alloc() : memref<1x1x1x1x16x10xf32, 2>
938 // CHECK:       affine.for %{{.*}} = 0 to 3 {
939 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 3 {
940 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 2 {
941 // CHECK-NEXT:        affine.for %{{.*}} = 0 to 2 {
942 // CHECK-NEXT:          affine.for %{{.*}} = 0 to 3 {
943 // CHECK-NEXT:            affine.for %{{.*}} = 0 to 3 {
944 // CHECK-NEXT:              affine.for %{{.*}} = 0 to 16 {
945 // CHECK-NEXT:                affine.for %{{.*}} = 0 to 10 {
946 // CHECK-NEXT:                  affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x10xf32, 2>
947 // CHECK-NEXT:                }
948 // CHECK-NEXT:              }
949 // CHECK-NEXT:              affine.for %{{.*}} = 0 to 16 {
950 // CHECK-NEXT:                affine.for %{{.*}} = 0 to 10 {
951 // CHECK-NEXT:                  affine.store %{{.*}}, %{{.*}}[0, 0, 0, 0, %{{.*}}, %{{.*}}] : memref<1x1x1x1x16x10xf32, 2>
952 // CHECK-NEXT:                }
953 // CHECK-NEXT:              }
954 // CHECK-NEXT:              affine.for %{{.*}} = 0 to 2 {
955 // CHECK-NEXT:                affine.for %{{.*}} = 0 to 2 {
956 // CHECK-NEXT:                  affine.for %{{.*}} = 0 to 16 {
957 // CHECK-NEXT:                    affine.for %{{.*}} = 0 to 10 {
958 // CHECK-NEXT:                      affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x10xf32, 2>
959 // CHECK-NEXT:                    }
960 // CHECK-NEXT:                  }
961 // CHECK-NEXT:                  affine.for %{{.*}} = 0 to 16 {
962 // CHECK-NEXT:                    affine.for %{{.*}} = 0 to 10 {
963 // CHECK-NEXT:                      affine.load %{{.*}}[0, 0, 0, 0, %{{.*}}, %{{.*}}] : memref<1x1x1x1x16x10xf32, 2>
964 // CHECK-NEXT:                    }
965 // CHECK-NEXT:                  }
966 // CHECK-NEXT:                }
967 // CHECK-NEXT:              }
968 // CHECK-NEXT:            }
969 // CHECK-NEXT:          }
970 // CHECK-NEXT:        }
971 // CHECK-NEXT:      }
972 // CHECK-NEXT:    }
973 // CHECK-NEXT:  }
974 // CHECK-NEXT:  return
975   return
976 }
977
978 // -----
979
980 // CHECK-LABEL: func @should_fuse_at_depth1_and_reduce_slice_trip_count
981 func @should_fuse_at_depth1_and_reduce_slice_trip_count() {
982   %a = alloc() : memref<4x256xf32>
983   %b = alloc() : memref<4x256xf32>
984
985   %c0 = constant 0 : index
986   %cf0 = constant 0.0 : f32
987
988   affine.for %i0 = 0 to 4 {
989     affine.for %i1 = 0 to 256 {
990       %v0 = affine.load %b[%i0, %i1] : memref<4x256xf32>
991     }
992     affine.for %i2 = 0 to 256 {
993       affine.store %cf0, %a[%i0, %i2] : memref<4x256xf32>
994     }
995   }
996
997   affine.for %d0 = 0 to 4 {
998     affine.for %d1 = 0 to 16 {
999       %v1 = affine.load %a[%d0, %d1] : memref<4x256xf32>
1000     }
1001   }
1002   // The cost of fusing at depth 2 is greater than the cost of fusing at depth 1
1003   // for two reasons:
1004   // 1) Inserting the unsliceable src loop %i1 to a higher depth removes
1005   //    redundant computation and reduces costs.
1006   // 2) Inserting the sliceable src loop %i2 at depth 1, we can still reduce
1007   //    its trip count to 16 (from 256) reducing costs.
1008   // NOTE: the size of the private memref created for the fused loop nest
1009   // is reduced from the original shape from 4x256 to 4x16 because of the
1010   // data accessed by the load.
1011   // CHECK-DAG:   alloc() : memref<1x16xf32>
1012   // CHECK:       affine.for %{{.*}} = 0 to 4 {
1013   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 256 {
1014   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4x256xf32>
1015   // CHECK-NEXT:    }
1016   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
1017   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
1018   // CHECK-NEXT:    }
1019   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
1020   // CHECK-NEXT:      affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
1021   // CHECK-NEXT:    }
1022   // CHECK-NEXT:  }
1023   // CHECK-NEXT:  return
1024   return
1025 }
1026
1027 // -----
1028
1029 // CHECK-LABEL: func @should_fuse_at_depth1_with_trip_count_20
1030 func @should_fuse_at_depth1_with_trip_count_20() {
1031   %a = alloc() : memref<100xf32>
1032   %c0 = constant 0 : index
1033   %cf0 = constant 0.0 : f32
1034
1035   affine.for %i0 = 0 to 100 {
1036     affine.store %cf0, %a[%i0]: memref<100xf32>
1037   }
1038
1039   affine.for %i1 = 0 to 5 {
1040     affine.for %i2 = 0 to 10 {
1041       %v0 = affine.load %a[%i2]: memref<100xf32>
1042     }
1043     affine.for %i3 = 0 to 10 {
1044       affine.for %i4 = 0 to 20 {
1045         %v1 = affine.load %a[%i4]: memref<100xf32>
1046       }
1047     }
1048   }
1049   // NOTE: The size of the private memref created for fusion is shrunk to 20xf32
1050   // CHECK-DAG:   alloc() : memref<20xf32>
1051   // CHECK:       affine.for %{{.*}} = 0 to 5 {
1052   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 20 {
1053   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<20xf32>
1054   // CHECK-NEXT:    }
1055   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
1056   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}] : memref<20xf32>
1057   // CHECK-NEXT:    }
1058   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
1059   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 20 {
1060   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}] : memref<20xf32>
1061   // CHECK-NEXT:      }
1062   // CHECK-NEXT:    }
1063   // CHECK-NEXT:  }
1064   // CHECK-NEXT:  return
1065   return
1066 }
1067
1068 // -----
1069
1070 // CHECK-LABEL: func @should_fuse_at_depth1_with_trip_count_19
1071 func @should_fuse_at_depth1_with_trip_count_19() {
1072   %a = alloc() : memref<100xf32>
1073   %c0 = constant 0 : index
1074   %cf0 = constant 0.0 : f32
1075
1076   affine.for %i0 = 0 to 100 {
1077     affine.store %cf0, %a[%i0]: memref<100xf32>
1078   }
1079
1080   affine.for %i1 = 0 to 5 {
1081     affine.for %i2 = 0 to 19 {
1082       %v0 = affine.load %a[%i2]: memref<100xf32>
1083     }
1084     affine.for %i3 = 0 to 10 {
1085       affine.for %i4 = 0 to 10 {
1086         %v1 = affine.load %a[%i4]: memref<100xf32>
1087       }
1088     }
1089   }
1090   // NOTE: The size of the private memref created for fusion is shrunk to 19xf32
1091   // CHECK-DAG:   alloc() : memref<19xf32>
1092   // CHECK:       affine.for %{{.*}} = 0 to 5 {
1093   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 19 {
1094   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<19xf32>
1095   // CHECK-NEXT:    }
1096   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 19 {
1097   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}] : memref<19xf32>
1098   // CHECK-NEXT:    }
1099   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
1100   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 10 {
1101   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}] : memref<19xf32>
1102   // CHECK-NEXT:      }
1103   // CHECK-NEXT:    }
1104   // CHECK-NEXT:  }
1105   // CHECK-NEXT:  return
1106   return
1107 }
1108
1109
1110 // -----
1111
1112 // CHECK-LABEL: func @should_fuse_with_private_memrefs_with_diff_shapes() {
1113 func @should_fuse_with_private_memrefs_with_diff_shapes() {
1114   %m = alloc() : memref<100xf32>
1115   %cf7 = constant 7.0 : f32
1116
1117   affine.for %i0 = 0 to 100 {
1118     affine.store %cf7, %m[%i0] : memref<100xf32>
1119   }
1120   affine.for %i1 = 0 to 17 {
1121     %v0 = affine.load %m[%i1] : memref<100xf32>
1122   }
1123   affine.for %i2 = 0 to 82 {
1124     %v1 = affine.load %m[%i2] : memref<100xf32>
1125   }
1126   // Should create two new private memrefs customized to the shapes accessed
1127   // by loops %{{.*}} and %{{.*}}.
1128   // CHECK-DAG:  alloc() : memref<1xf32>
1129   // CHECK-DAG:  alloc() : memref<1xf32>
1130   // CHECK:      affine.for %{{.*}} = 0 to 17 {
1131   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1132   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
1133   // CHECK-NEXT: }
1134   // CHECK-NEXT: affine.for %{{.*}} = 0 to 82 {
1135   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1136   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
1137   // CHECK-NEXT: }
1138   // CHECK-NEXT: return
1139   return
1140 }
1141
1142 // -----
1143
1144 // CHECK-LABEL: func @should_fuse_live_out_arg_but_preserve_src_loop(%{{.*}}: memref<10xf32>) {
1145 func @should_fuse_live_out_arg_but_preserve_src_loop(%arg0: memref<10xf32>) {
1146   %cf7 = constant 7.0 : f32
1147
1148   affine.for %i0 = 0 to 10 {
1149     affine.store %cf7, %arg0[%i0] : memref<10xf32>
1150   }
1151   affine.for %i1 = 0 to 9 {
1152     %v0 = affine.load %arg0[%i1] : memref<10xf32>
1153   }
1154   // This tests that the loop nest '%i0' should not be removed after fusion
1155   // because it writes to memref argument '%arg0', and its read region
1156   // does not cover its write region (so fusion would shrink the write region
1157   // in the fused loop nest, so complete live out data region would not
1158   // be written).
1159   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1160   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1161   // CHECK-NEXT:  }
1162   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 9 {
1163   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1164   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1165   // CHECK-NEXT:  }
1166   // CHECK-NEXT:  return
1167   return
1168 }
1169
1170 // -----
1171
1172 // CHECK-LABEL: func @should_fuse_live_out_arg(%{{.*}}: memref<10xf32>) {
1173 func @should_fuse_live_out_arg(%arg0: memref<10xf32>) {
1174   %cf7 = constant 7.0 : f32
1175
1176   affine.for %i0 = 0 to 10 {
1177     affine.store %cf7, %arg0[%i0] : memref<10xf32>
1178   }
1179   affine.for %i1 = 0 to 10 {
1180     %v0 = affine.load %arg0[%i1] : memref<10xf32>
1181   }
1182   // The read/write regions for memref '%{{.*}}' are the same for both
1183   // loops, so they should fuse.
1184
1185   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1186   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1187   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1188   // CHECK-NEXT:  }
1189   // CHECK-NEXT:  return
1190   return
1191 }
1192
1193 // -----
1194
1195 // CHECK-LABEL: func @should_fuse_escaping_memref_but_preserve_src_loop() -> memref<10xf32>
1196 func @should_fuse_escaping_memref_but_preserve_src_loop() -> memref<10xf32> {
1197   %cf7 = constant 7.0 : f32
1198   %m = alloc() : memref<10xf32>
1199   affine.for %i0 = 0 to 10 {
1200     affine.store %cf7, %m[%i0] : memref<10xf32>
1201   }
1202   affine.for %i1 = 0 to 9 {
1203     %v0 = affine.load %m[%i1] : memref<10xf32>
1204   }
1205   // This tests that the loop nest '%i0' should not be removed after fusion
1206   // because it writes to memref '%m', which is returned by the function, and
1207   // the '%i1' memory region does not cover '%i0' memory region.
1208
1209   // CHECK-DAG:   alloc() : memref<10xf32>
1210   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1211   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1212   // CHECK-NEXT:  }
1213   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 9 {
1214   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1215   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1216   // CHECK-NEXT:  }
1217   // CHECK-NEXT:  return %{{.*}} : memref<10xf32>
1218   return %m : memref<10xf32>
1219 }
1220 // -----
1221
1222 // This should fuse with the %in becoming a 1x1x1.
1223 func @R3_to_R2_reshape() {
1224   %in = alloc() : memref<2x3x16xi32>
1225
1226   %c0 = constant 0 : index
1227
1228   affine.for %i0 = 0 to 2 {
1229     affine.for %i1 = 0 to 3 {
1230       affine.for %i2 = 0 to 16 {
1231         %val = "foo"(%i0, %i1, %i2) : (index, index, index) -> i32
1232         affine.store %val, %in[%i0, %i1, %i2] : memref<2x3x16xi32>
1233       }
1234     }
1235   }
1236
1237   affine.for %ii = 0 to 32 {
1238     affine.for %jj = 0 to 3 {
1239       %a0 = affine.apply affine_map<(d0, d1) -> (d0 * 3 + d1)> (%ii, %jj)
1240       %idx = affine.apply affine_map<(d0) -> (d0 floordiv (3 * 16))> (%a0)
1241       %v = affine.load %in[%idx, %jj, %c0]
1242         : memref<2x3x16xi32>
1243     }
1244   }
1245   return
1246 }
1247 // CHECK-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> ((d0 * 3 + d1) floordiv 48)>
1248 // CHECK-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 3 + d1)>
1249 // CHECK-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0) -> (d0 floordiv 48)>
1250
1251 // CHECK-LABEL: func @R3_to_R2_reshape()
1252 // CHECK-DAG:    alloc() : memref<1x1x1xi32>
1253 // CHECK:        affine.for %{{.*}} = 0 to 32 {
1254 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 3 {
1255 // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
1256 // CHECK-NEXT:      "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> i32
1257 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
1258 // CHECK-NEXT:      affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
1259 // CHECK-NEXT:      affine.apply [[$MAP2]](%{{.*}})
1260 // CHECK-NEXT:      affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
1261 // CHECK-NEXT:    }
1262 // CHECK-NEXT:  }
1263 // CHECK-NEXT:  return
1264
1265 // -----
1266
1267 func @should_fuse_multi_output_producer() {
1268   %a = alloc() : memref<10xf32>
1269   %b = alloc() : memref<10xf32>
1270
1271   %cf7 = constant 7.0 : f32
1272
1273   affine.for %i0 = 0 to 10 {
1274     affine.store %cf7, %a[%i0] : memref<10xf32>
1275     affine.store %cf7, %b[%i0] : memref<10xf32>
1276   }
1277   affine.for %i1 = 0 to 10 {
1278     %v0 = affine.load %a[%i1] : memref<10xf32>
1279     %v1 = affine.load %b[%i1] : memref<10xf32>
1280   }
1281
1282   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1283   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1284   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1285   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
1286   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
1287   // CHECK-NEXT:  }
1288   // CHECK-NEXT:  return
1289   return
1290 }
1291
1292 // -----
1293
1294 // CHECK-LABEL: func @fusion_preventing_deps_on_middle_loop() {
1295 func @fusion_preventing_deps_on_middle_loop() {
1296   %a = alloc() : memref<10xf32>
1297   %b = alloc() : memref<10xf32>
1298   %c = alloc() : memref<10xf32>
1299
1300   %cf7 = constant 7.0 : f32
1301
1302   affine.for %i0 = 0 to 10 {
1303     %v0 = affine.load %a[%i0] : memref<10xf32>
1304     affine.store %v0, %b[%i0] : memref<10xf32>
1305   }
1306   affine.for %i1 = 0 to 10 {
1307     affine.store %cf7, %a[%i1] : memref<10xf32>
1308     %v1 = affine.load %c[%i1] : memref<10xf32>
1309   }
1310   affine.for %i2 = 0 to 10 {
1311     %v2 = affine.load %b[%i2] : memref<10xf32>
1312     affine.store %v2, %c[%i2] : memref<10xf32>
1313   }
1314   // Loops '%i0' and '%i2' cannot fuse along producer/consumer edge on memref
1315   // '%b', because of the WAR dep from '%i0' to '%i1' on memref '%a' and
1316   // because of the WAR dep from '%i1' to '%i2' on memref '%c'.
1317   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1318   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1319   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1320   // CHECK-NEXT:  }
1321   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1322   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1323   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1324   // CHECK-NEXT:  }
1325   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1326   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1327   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1328   // CHECK-NEXT:  }
1329   // CHECK-NEXT:  return
1330   return
1331 }
1332
1333 // -----
1334
1335 // CHECK-LABEL: func @should_fuse_and_move_to_preserve_war_dep() {
1336 func @should_fuse_and_move_to_preserve_war_dep() {
1337   %a = alloc() : memref<10xf32>
1338   %b = alloc() : memref<10xf32>
1339   %c = alloc() : memref<10xf32>
1340
1341   %cf7 = constant 7.0 : f32
1342
1343   affine.for %i0 = 0 to 10 {
1344     %v0 = affine.load %b[%i0] : memref<10xf32>
1345     affine.store %v0, %a[%i0] : memref<10xf32>
1346   }
1347   affine.for %i1 = 0 to 3 {
1348     %v2 = affine.load %c[%i1] : memref<10xf32>
1349   }
1350   affine.for %i2 = 0 to 5 {
1351     affine.store %cf7, %b[%i2] : memref<10xf32>
1352   }
1353   affine.for %i3 = 0 to 10 {
1354     %v1 = affine.load %a[%i3] : memref<10xf32>
1355     affine.store %cf7, %c[%i3] : memref<10xf32>
1356   }
1357
1358   // Dependence graph:
1359   //
1360   //         %i0 ---------
1361   //               |     |
1362   //     --- %i1   | %b  | %a
1363   //     |         |     |
1364   //  %c |   %i2 <--     |
1365   //     |               |
1366   //     --> %i3 <--------
1367   //
1368   // It is possible to fuse loop '%i0' into '%i3' and preserve dependences
1369   // if the fused loop nest is inserted between loops '%i1' and '%i2'.
1370
1371   // CHECK-DAG:   alloc() : memref<1xf32>
1372   // CHECK:       affine.for %{{.*}} = 0 to 3 {
1373   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1374   // CHECK-NEXT:  }
1375   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1376   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1377   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1378   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
1379   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1380   // CHECK-NEXT:  }
1381   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 5 {
1382   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1383   // CHECK-NEXT:  }
1384   // CHECK-NEXT:  return
1385   return
1386 }
1387
1388 // -----
1389
1390 // CHECK-LABEL: func @fusion_preventing_dep_on_constant() {
1391 func @fusion_preventing_dep_on_constant() {
1392   %a = alloc() : memref<10xf32>
1393   %b = alloc() : memref<10xf32>
1394   %c = alloc() : memref<10xf32>
1395
1396   %cf7 = constant 7.0 : f32
1397
1398   affine.for %i0 = 0 to 10 {
1399     %v0 = affine.load %b[%i0] : memref<10xf32>
1400     affine.store %cf7, %a[%i0] : memref<10xf32>
1401   }
1402   affine.for %i1 = 0 to 10 {
1403     affine.store %cf7, %b[%i1] : memref<10xf32>
1404   }
1405   %cf11 = constant 11.0 : f32
1406   affine.for %i2 = 0 to 10 {
1407     %v2 = affine.load %a[%i2] : memref<10xf32>
1408     affine.store %cf11, %c[%i2] : memref<10xf32>
1409   }
1410   // Loops '%i0' and '%i2' cannot fuse along producer/consumer edge on memref
1411   // '%a', because of the WAR dep from '%i0' to '%i1' on memref '%b' and
1412   // because of the SSA value dep from '%cf11' def to use in '%i2'.
1413   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1414   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1415   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1416   // CHECK-NEXT:  }
1417   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1418   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1419   // CHECK-NEXT:  }
1420   // CHECK-NEXT:  %{{.*}} = constant 1.100000e+01 : f32
1421   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1422   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1423   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1424   // CHECK-NEXT:  }
1425   // CHECK-NEXT:  return
1426   return
1427 }
1428
1429 // -----
1430
1431 // CHECK-LABEL: func @should_fuse_and_preserve_dep_on_constant() {
1432 func @should_fuse_and_preserve_dep_on_constant() {
1433   %a = alloc() : memref<10xf32>
1434   %b = alloc() : memref<10xf32>
1435   %c = alloc() : memref<10xf32>
1436
1437   %cf7 = constant 7.0 : f32
1438   %cf11 = constant 11.0 : f32
1439   affine.for %i0 = 0 to 10 {
1440     %v0 = affine.load %b[%i0] : memref<10xf32>
1441     affine.store %cf7, %a[%i0] : memref<10xf32>
1442   }
1443   affine.for %i1 = 0 to 10 {
1444     affine.store %cf7, %b[%i1] : memref<10xf32>
1445   }
1446   affine.for %i2 = 0 to 10 {
1447     %v2 = affine.load %a[%i2] : memref<10xf32>
1448     affine.store %cf11, %c[%i2] : memref<10xf32>
1449   }
1450
1451   // Loops '%i0' and '%i2' can fuse along producer/consumer edge on memref
1452   // '%a', and preserve the WAR dep from '%i0' to '%i1' on memref '%b', and
1453   // the SSA value dep from '%cf11' def to use in '%i2'.
1454
1455   // CHECK:       constant 1.100000e+01 : f32
1456   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1457   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1458   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1459   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
1460   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1461   // CHECK-NEXT:  }
1462   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1463   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1464   // CHECK-NEXT:  }
1465   // CHECK-NEXT:  return
1466   return
1467 }
1468
1469 // -----
1470
1471 // CHECK-LABEL: func @should_fuse_at_depth_above_loop_carried_dependence(%{{.*}}: memref<64x4xf32>, %{{.*}}: memref<64x4xf32>) {
1472 func @should_fuse_at_depth_above_loop_carried_dependence(%arg0: memref<64x4xf32>, %arg1: memref<64x4xf32>) {
1473   %out = alloc() : memref<64x4xf32>
1474   %0 = constant 0.0 : f32
1475   affine.for %i0 = 0 to 64 {
1476     affine.for %i1 = 0 to 4 {
1477       affine.store %0, %out[%i0, %i1] : memref<64x4xf32>
1478     }
1479   }
1480   affine.for %i2 = 0 to 4 {
1481     affine.for %i3 = 0 to 4 {
1482       affine.for %i4 = 0 to 16 {
1483         %v = affine.load %arg1[16 * %i3 - %i4 + 15, %i2] : memref<64x4xf32>
1484         "op0"(%v) : (f32) -> ()
1485       }
1486       affine.for %i5 = 0 to 4 {
1487         affine.for %i6 = 0 to 16 {
1488           %v = affine.load %arg0[16 * %i5 - %i6 + 15, %i3] : memref<64x4xf32>
1489           "op1"(%v) : (f32) -> ()
1490         }
1491         affine.for %i7 = 0 to 16 {
1492           %r = "op2"() : () -> (f32)
1493           %v = affine.load %out[16 * %i5 + %i7, %i2] : memref<64x4xf32>
1494           %s = addf %v, %r : f32
1495           affine.store %s, %out[16 * %i5 + %i7, %i2] : memref<64x4xf32>
1496         }
1497       }
1498     }
1499   }
1500
1501   // We can fuse source loop nest '%i0' into dst loop nest '%i2', but the
1502   // depth at which we can insert the src loop nest slice into the dst loop
1503   // lest must be decreased because of a loop carried dependence on loop '%i3'.
1504   // As a result, the source loop nest is inserted at dst loop nest depth 1,
1505   // just above the loop with the carried dependence. In addition, the source
1506   // loop nest iteration bounds on its loop '%i1' are reduced to 1, so the
1507   // memref size can be reduced to 128x1xf32.
1508
1509   // CHECK:       alloc() : memref<64x1xf32>
1510   // CHECK:       affine.for %{{.*}} = 0 to 4 {
1511   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 64 {
1512   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<64x1xf32>
1513   // CHECK-NEXT:    }
1514   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 4 {
1515   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 16 {
1516   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, %{{.*}}] : memref<64x4xf32>
1517   // CHECK-NEXT:        "op0"(%{{.*}}) : (f32) -> ()
1518   // CHECK-NEXT:      }
1519   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 4 {
1520   // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
1521   // CHECK-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, %{{.*}}] : memref<64x4xf32>
1522   // CHECK-NEXT:          "op1"(%{{.*}}) : (f32) -> ()
1523   // CHECK-NEXT:        }
1524   // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
1525   // CHECK-NEXT:          %{{.*}} = "op2"() : () -> f32
1526   // CHECK:               affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
1527   // CHECK-NEXT:          addf %{{.*}}, %{{.*}} : f32
1528   // CHECK:               affine.store %{{.*}}, %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
1529   // CHECK-NEXT:        }
1530   // CHECK-NEXT:      }
1531   // CHECK-NEXT:    }
1532   // CHECK-NEXT:  }
1533   // CHECK-NEXT:  return
1534   return
1535 }
1536
1537 // -----
1538
1539 // CHECK-LABEL: func @should_fuse_only_two_loops_and_remove_producer() {
1540 func @should_fuse_only_two_loops_and_remove_producer() {
1541   %a = alloc() : memref<10xf32>
1542   %b = alloc() : memref<10xf32>
1543
1544   %cf7 = constant 7.0 : f32
1545
1546   affine.for %i0 = 0 to 10 {
1547     affine.store %cf7, %a[%i0] : memref<10xf32>
1548   }
1549   affine.for %i1 = 0 to 10 {
1550     %v0 = affine.load %a[%i1] : memref<10xf32>
1551     affine.store %v0, %b[%i1] : memref<10xf32>
1552   }
1553   affine.for %i2 = 0 to 10 {
1554     %v1 = affine.load %a[%i2] : memref<10xf32>
1555     affine.store %v1, %b[%i2] : memref<10xf32>
1556   }
1557
1558   // On the first visit to '%i2', the fusion algorithm can not fuse loop nest
1559   // '%i0' into '%i2' because of the dependences '%i0' and '%i2' each have on
1560   // '%i1'. Then, '%i0' is fused into '%i1' and no private memref is created for
1561   // memref '%a' to be able to remove '%i0' and still preserve the depencence on
1562   // '%a' with '%i2'.
1563   // TODO: Alternatively, we could fuse '%i0' into '%i1' with a private memref,
1564   // the dependence between '%i0' and '%i1' on memref '%a' would no longer exist,
1565   // and '%i0' could be fused into '%i2' as well. Note that this approach would
1566   // duplicate the computation in loop nest '%i0' to loop nests '%i1' and '%i2',
1567   // which would limit its profitability.
1568   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1569   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1570   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1571   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1572   // CHECK-NEXT:  }
1573   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1574   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1575   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1576   // CHECK-NEXT:  }
1577   // CHECK-NEXT:   return
1578   return
1579 }
1580
1581 // -----
1582
1583 // CHECK-LABEL: func @should_fuse_after_one_loop_interchange() {
1584 func @should_fuse_after_one_loop_interchange() {
1585   %a = alloc() : memref<10xf32>
1586
1587   %cf0 = constant 0.0 : f32
1588   affine.for %i0 = 0 to 10 {
1589     affine.store %cf0, %a[%i0] : memref<10xf32>
1590   }
1591
1592   affine.for %i1 = 0 to 5 {
1593     affine.for %i2 = 0 to 10 {
1594       %v0 = affine.load %a[%i2] : memref<10xf32>
1595       affine.store %v0, %a[%i2] : memref<10xf32>
1596     }
1597   }
1598
1599   // The dependence between the load and affine.store is carried on loop '%i1', and
1600   // cannot be fused with loop '%i0' without violating this dependence.
1601   // Once loops '%i1' and %i2' are interchanged, loop '%i0' can be fused
1602   // at loop depth 1, because the loop carrying the dependence has been
1603   // interchanged and is now at depth 2.
1604
1605   // CHECK:       affine.for %{{.*}} = 0 to 10 {
1606   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1607   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 5 {
1608   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
1609   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
1610   // CHECK-NEXT:    }
1611   // CHECK-NEXT:  }
1612   // CHECK-NEXT:  return
1613   return
1614 }
1615
1616 // -----
1617
1618 // CHECK-LABEL: func @should_fuse_after_two_loop_interchanges() {
1619 func @should_fuse_after_two_loop_interchanges() {
1620   %a = alloc() : memref<6x8xf32>
1621
1622   %cf0 = constant 0.0 : f32
1623   affine.for %i0 = 0 to 6 {
1624     affine.for %i1 = 0 to 8 {
1625       affine.store %cf0, %a[%i0, %i1] : memref<6x8xf32>
1626     }
1627   }
1628
1629   affine.for %i2 = 0 to 4 {
1630     affine.for %i3 = 0 to 6 {
1631       affine.for %i4 = 0 to 2 {
1632         affine.for %i5 = 0 to 8 {
1633           %v0 = affine.load %a[%i3, %i5] : memref<6x8xf32>
1634           %v1 = addf %v0, %v0 : f32
1635           affine.store %v1, %a[%i3, %i5] : memref<6x8xf32>
1636         }
1637       }
1638     }
1639   }
1640
1641   // The dependence between the load and affine.store is carried on loops '%i2' and
1642   // '%i4', and cannot be fused with loop '%i0' without violating this
1643   // dependence.
1644   // Once loop '%i2' is interchanged with loop '%i3', and again with loop
1645   // '%i5', then loop '%i0' can be fused at loop depth 2, because the loop
1646   // carrying the dependences have been interchanged with loops at depth > 2.
1647
1648   // CHECK:       affine.for %{{.*}} = 0 to 6 {
1649   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8 {
1650   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
1651   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 4 {
1652   // CHECK-NEXT:        affine.for %{{.*}} = 0 to 2 {
1653   // CHECK-NEXT:          affine.load %{{.*}}[0, 0] : memref<1x1xf32>
1654   // CHECK-NEXT:          addf %{{.*}}, %{{.*}} : f32
1655   // CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
1656   // CHECK-NEXT:        }
1657   // CHECK-NEXT:      }
1658   // CHECK-NEXT:    }
1659   // CHECK-NEXT:  }
1660   // CHECK-NEXT:  return
1661   return
1662 }
1663
1664 // -----
1665
1666 func @should_fuse_live_out_writer(%arg0 : memref<10xf32>) -> memref<10xf32> {
1667   %cst = constant 0.000000e+00 : f32
1668   affine.for %i0 = 0 to 10 {
1669     affine.store %cst, %arg0[%i0] : memref<10xf32>
1670   }
1671   affine.for %i1 = 0 to 10 {
1672     %1 = affine.load %arg0[%i1] : memref<10xf32>
1673     affine.store %1, %arg0[%i1] : memref<10xf32>
1674   }
1675   return %arg0 : memref<10xf32>
1676
1677   // CHECK:       %{{.*}} = constant 0.000000e+00 : f32
1678   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
1679   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1680   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1681   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1682   // CHECK-NEXT:  }
1683   // CHECK-NEXT:  return %{{.*}} : memref<10xf32>
1684 }
1685
1686 // -----
1687
1688 // The fused slice has 16 iterations from along %i0.
1689
1690 // CHECK-DAG: [[$MAP_LB:#map[0-9]+]] = affine_map<(d0) -> (d0 * 16)>
1691 // CHECK-DAG: [[$MAP_UB:#map[0-9]+]] = affine_map<(d0) -> (d0 * 16 + 16)>
1692
1693 // CHECK-LABEL: slice_tile
1694 func @slice_tile(%arg0: memref<128x8xf32>, %arg1: memref<32x8xf32>, %0 : f32) -> memref<32x8xf32> {
1695   affine.for %i0 = 0 to 32 {
1696     affine.for %i1 = 0 to 8 {
1697       affine.store %0, %arg1[%i0, %i1] : memref<32x8xf32>
1698     }
1699   }
1700   affine.for %i = 0 to 2 {
1701     affine.for %j = 0 to 8 {
1702       affine.for %k = 0 to 8 {
1703         affine.for %kk = 0 to 16 {
1704           %v = affine.load %arg0[16 * %k + %kk, %j] : memref<128x8xf32>
1705           %r = "foo"(%v) : (f32) -> f32
1706         }
1707         affine.for %ii = 0 to 16 {
1708           %v = affine.load %arg1[16 * %i + %ii, %j] : memref<32x8xf32>
1709           %s = addf %v, %v : f32
1710           affine.store %s, %arg1[16 * %i + %ii, %j] : memref<32x8xf32>
1711         }
1712       }
1713     }
1714   }
1715   return %arg1 : memref<32x8xf32>
1716 }
1717 // CHECK:       affine.for %{{.*}} = 0 to 2 {
1718 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8 {
1719 // CHECK-NEXT:      affine.for %{{.*}} = [[$MAP_LB]](%{{.*}}) to [[$MAP_UB]](%{{.*}}) {
1720 // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<32x8xf32>
1721 // CHECK-NEXT:      }
1722 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 8 {
1723 // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
1724 // CHECK-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<128x8xf32>
1725 // CHECK-NEXT:          "foo"(%{{.*}}) : (f32) -> f32
1726 // CHECK-NEXT:        }
1727 // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
1728 // CHECK-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<32x8xf32>
1729 // CHECK-NEXT:          addf %{{.*}}, %{{.*}} : f32
1730 // CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<32x8xf32>
1731 // CHECK-NEXT:        }
1732 // CHECK-NEXT:      }
1733 // CHECK-NEXT:    }
1734 // CHECK-NEXT:  }
1735 // CHECK-NEXT:  return %{{.*}} : memref<32x8xf32>
1736 // CHECK-NEXT:}
1737
1738 // -----
1739
1740 // Test case which illustrates fix for b/126454413
1741 func @test_add_slice_bounds() {
1742   %a = alloc() : memref<10xf32>
1743   %b = alloc() : memref<10xf32>
1744   %cf7 = constant 7.0 : f32
1745   %c0 = constant 0 : index
1746
1747   affine.for %i0 = 0 to 10 {
1748     affine.for %i1 = 0 to 10 {
1749       affine.for %i2 = 0 to 10 {
1750         %a0 = affine.apply affine_map<(d0) -> (d0)> (%i0)
1751         %a1 = affine.apply affine_map<(d0) -> (d0)> (%i0)
1752         %a2 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%a0, %a1)
1753         affine.store %cf7, %a[%a2] : memref<10xf32>
1754       }
1755     }
1756   }
1757   affine.for %i3 = 0 to 10 {
1758     affine.for %i4 = 0 to 10 {
1759       affine.for %i5 = 0 to 10 {
1760         %v0 = affine.load %a[%c0] : memref<10xf32>
1761       }
1762     }
1763   }
1764
1765 // CHECK:        affine.for %{{.*}} = 0 to 10 {
1766 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
1767 // CHECK-NEXT:       affine.for %{{.*}} = 0 to 10 {
1768 // CHECK-NEXT:         affine.apply #map0(%{{.*}})
1769 // CHECK-NEXT:         affine.apply #map0(%{{.*}})
1770 // CHECK-NEXT:         affine.apply #map1(%{{.*}}, %{{.*}})
1771 // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1772 // CHECK-NEXT:       }
1773 // CHECK-NEXT:     }
1774 // CHECK-NEXT:   }
1775 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10 {
1776 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
1777 // CHECK-NEXT:       affine.for %{{.*}} = 0 to 10 {
1778 // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1779 // CHECK-NEXT:       }
1780 // CHECK-NEXT:     }
1781 // CHECK-NEXT:   }
1782   return
1783 }
1784
1785 // -----
1786
1787 func @should_fuse_init_loops_siblings_then_shared_producer(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>) {
1788   %0 = alloc() : memref<10x10xf32>
1789   %cst = constant 0.000000e+00 : f32
1790   %cst_0 = constant 1.000000e+00 : f32
1791   %cst_1 = constant 7.000000e+00 : f32
1792   affine.for %i0 = 0 to 10 {
1793     affine.for %i1 = 0 to 10 {
1794       affine.store %cst_1, %0[%i0, %i1] : memref<10x10xf32>
1795     }
1796   }
1797   affine.for %i2 = 0 to 3 {
1798     affine.for %i3 = 0 to 3 {
1799       affine.store %cst, %arg0[%i2, %i3] : memref<10x10xf32>
1800     }
1801   }
1802   affine.for %i4 = 0 to 3 {
1803     affine.for %i5 = 0 to 3 {
1804       %1 = affine.load %0[%i4, %i5] : memref<10x10xf32>
1805       %2 = affine.load %arg0[%i4, %i5] : memref<10x10xf32>
1806       %3 = mulf %1, %2 : f32
1807       affine.store %3, %arg0[%i4, %i5] : memref<10x10xf32>
1808     }
1809   }
1810   affine.for %i6 = 0 to 3 {
1811     affine.for %i7 = 0 to 3 {
1812       affine.store %cst_0, %arg1[%i6, %i7] : memref<10x10xf32>
1813     }
1814   }
1815   affine.for %i8 = 0 to 3 {
1816     affine.for %i9 = 0 to 3 {
1817       %4 = affine.load %0[%i8, %i9] : memref<10x10xf32>
1818       %5 = affine.load %arg1[%i8, %i9] : memref<10x10xf32>
1819       %6 = addf %4, %5 : f32
1820       affine.store %6, %arg1[%i8, %i9] : memref<10x10xf32>
1821     }
1822   }
1823
1824   // Pass 1: should fuse single-use producer loop nests into their unique user,
1825   //         so '%i2' will fuse into '%i4' and '%i6' will fuse into '%i8'.
1826   // Pass 2: should fuse sibling loop nests which share no dependence edges,
1827   //         so should fuse '%i4' into '%i8'.
1828   // Pass 3: should fuse single-use producer loop nest '%i0' into '%i8'. Note
1829   //         that loop nest '%i0' now has a single user after Pass 2 fused its
1830   //         two users together).
1831
1832 // CHECK:        affine.for %{{.*}} = 0 to 3 {
1833 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 3 {
1834 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
1835 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
1836 // CHECK-NEXT:       affine.load %{{.*}}[0, 0] : memref<1x1xf32>
1837 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
1838 // CHECK-NEXT:       mulf %{{.*}}, %{{.*}} : f32
1839 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
1840 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
1841 // CHECK-NEXT:       affine.load %{{.*}}[0, 0] : memref<1x1xf32>
1842 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
1843 // CHECK-NEXT:       addf %{{.*}}, %{{.*}} : f32
1844 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
1845 // CHECK-NEXT:     }
1846 // CHECK-NEXT:   }
1847 // CHECK-NEXT:   return
1848
1849   return
1850 }
1851
1852 // -----
1853
1854 func @two_matrix_vector_products() {
1855   %in_matrix = alloc() : memref<10x10xf32>
1856   %in_vec0 = alloc() : memref<10xf32>
1857   %in_vec1 = alloc() : memref<10xf32>
1858   %out_vec0 = alloc() : memref<10xf32>
1859   %out_vec1 = alloc() : memref<10xf32>
1860   %cf7 = constant 7.0 : f32
1861
1862   // Populate input matrix.
1863   affine.for %i0 = 0 to 10 {
1864     affine.for %i1 = 0 to 10 {
1865       affine.store %cf7, %in_matrix[%i0, %i1] : memref<10x10xf32>
1866     }
1867   }
1868   // out_vec0 = in_matrix x in_vec0
1869   affine.for %i2 = 0 to 10 {
1870     affine.for %i3 = 0 to 10 {
1871       %v0 = affine.load %in_matrix[%i2, %i3] : memref<10x10xf32>
1872       %v1 = affine.load %in_vec0[%i3] : memref<10xf32>
1873       %v2 = mulf %v0, %v1 : f32
1874       %v3 = affine.load %out_vec0[%i3] : memref<10xf32>
1875       %v4 = addf %v2, %v3 : f32
1876       affine.store %v4, %out_vec0[%i3] : memref<10xf32>
1877     }
1878   }
1879   // out_vec1 = in_matrix x in_vec1
1880   affine.for %i4 = 0 to 10 {
1881     affine.for %i5 = 0 to 10 {
1882       %v5 = affine.load %in_matrix[%i4, %i5] : memref<10x10xf32>
1883       %v6 = affine.load %in_vec1[%i5] : memref<10xf32>
1884       %v7 = mulf %v5, %v6 : f32
1885       %v8 = affine.load %out_vec1[%i5] : memref<10xf32>
1886       %v9 = addf %v7, %v8 : f32
1887       affine.store %v9, %out_vec1[%i5] : memref<10xf32>
1888     }
1889   }
1890
1891 // CHECK:        affine.for %{{.*}} = 0 to 10 {
1892 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
1893 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<10x1xf32>
1894 // CHECK-NEXT:     }
1895 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
1896 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, 0] : memref<10x1xf32>
1897 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1898 // CHECK-NEXT:       mulf %{{.*}}, %{{.*}} : f32
1899 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1900 // CHECK-NEXT:       addf %{{.*}}, %{{.*}} : f32
1901 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1902 // CHECK-NEXT:     }
1903 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
1904 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, 0] : memref<10x1xf32>
1905 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1906 // CHECK-NEXT:       mulf %{{.*}}, %{{.*}} : f32
1907 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
1908 // CHECK-NEXT:       addf %{{.*}}, %{{.*}} : f32
1909 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
1910 // CHECK-NEXT:     }
1911 // CHECK-NEXT:   }
1912 // CHECK-NEXT:   return
1913   return
1914 }
1915
1916 // -----
1917
1918 func @should_not_slice_past_slice_barrier() {
1919   %0 = alloc() : memref<100x16xf32>
1920   affine.for %i0 = 0 to 100 {
1921     affine.for %i1 = 0 to 16 {
1922       %1 = "op1"() : () -> f32
1923       affine.store %1, %0[%i0, %i1] : memref<100x16xf32>
1924     } {slice_fusion_barrier = true}
1925   }
1926   affine.for %i2 = 0 to 100 {
1927     affine.for %i3 = 0 to 16 {
1928       %2 = affine.load %0[%i2, %i3] : memref<100x16xf32>
1929       "op2"(%2) : (f32) -> ()
1930     }
1931   }
1932   // The 'slice_fusion_barrier' attribute on '%i1' prevents slicing the
1933   // iteration space of '%i1' and any enclosing loop nests.
1934 // CHECK:        affine.for %{{.*}} = 0 to 100 {
1935 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 16 {
1936 // CHECK-NEXT:       %{{.*}} = "op1"() : () -> f32
1937 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
1938 // CHECK-NEXT:     } {slice_fusion_barrier = true}
1939 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 16 {
1940 // CHECK-NEXT:       affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
1941 // CHECK-NEXT:       "op2"(%{{.*}}) : (f32) -> ()
1942 // CHECK-NEXT:     }
1943 // CHECK-NEXT:   }
1944   return
1945 }
1946
1947 // -----
1948
1949 #map0 = affine_map<(d0, d1) -> (d0 * 16 + d1)>
1950 func @fuse_across_dim_mismatch(%arg0: memref<4x4x16x1xf32>, %arg1: memref<144x9xf32>, %arg2: memref<9xf32>) {
1951   %1 = alloc() : memref<144x4xf32>
1952   %2 = constant 0.0 : f32
1953   affine.for %i2 = 0 to 9 {
1954     affine.for %i3 = 0 to 4 {
1955       affine.for %i5 = 0 to 16 {
1956         %7 = affine.apply #map0(%i2, %i5)
1957         affine.store %2, %1[%7, %i3] : memref<144x4xf32>
1958       }
1959     }
1960   }
1961   affine.for %i6 = 0 to 9 {
1962     affine.for %i7 = 0 to 9 {
1963       affine.for %i8 = 0 to 4 {
1964         affine.for %i10 = 0 to 16 {
1965           %10 = affine.apply #map0(%i6, %i10)
1966           %11 = affine.load %1[%10, %i8] : memref<144x4xf32>
1967         }
1968       }
1969     }
1970   }
1971   return
1972 }
1973 // MAXIMAL:      #map = affine_map<(d0, d1) -> (d0 * 16 + d1)>
1974 // MAXIMAL-LABEL: func @fuse_across_dim_mismatch
1975 // MAXIMAL:        alloc() : memref<1x1xf32>
1976 // MAXIMAL:        affine.for %{{.*}} = 0 to 9 {
1977 // MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 9 {
1978 // MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 4 {
1979 // MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 16 {
1980 // MAXIMAL-NEXT:          affine.apply #map(%{{.*}}, %{{.*}})
1981 // MAXIMAL-NEXT:          affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
1982 // MAXIMAL-NEXT:          affine.apply #map(%{{.*}}, %{{.*}})
1983 // MAXIMAL-NEXT:          affine.load %{{.*}}[0, 0] : memref<1x1xf32>
1984 // MAXIMAL-NEXT:        }
1985 // MAXIMAL-NEXT:      }
1986 // MAXIMAL-NEXT:    }
1987 // MAXIMAL-NEXT:  }
1988
1989 // -----
1990
1991 #map3 = affine_map<(d0, d1) -> ((d0 * 72 + d1) floordiv 2304)>
1992 #map4 = affine_map<(d0, d1) -> (((d0 * 72 + d1) mod 2304) floordiv 1152)>
1993 #map5 = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) floordiv 9) floordiv 8)>
1994 #map6 = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) mod 9) floordiv 3)>
1995 #map7 = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) mod 9) mod 3)>
1996 #map10 = affine_map<(d0, d1) -> (d0 * 16 + d1)>
1997 #map11 = affine_map<(d0, d1) -> (d0 * 16 + d1)>
1998 #map12 = affine_map<(d0, d1) -> (d0 * 16 - d1 + 15)>
1999 func @fuse_across_varying_dims_complex(%arg0: f32) {
2000   %c0 = constant 0 : index
2001   %0 = alloc() : memref<2x2x3x3x16x1xf32>
2002   %1 = alloc() : memref<64x9xf32>
2003   %2 = alloc() : memref<144x4xf32>
2004   affine.for %i0 = 0 to 64 {
2005     affine.for %i1 = 0 to 9 {
2006       %4 = affine.apply #map3(%i0, %i1)
2007       %5 = affine.apply #map4(%i0, %i1)
2008       %6 = affine.apply #map5(%i0, %i1)
2009       %7 = affine.apply #map6(%i0, %i1)
2010       %8 = affine.apply #map7(%i0, %i1)
2011       %9 = affine.load %0[%4, %5, %7, %8, %6, %c0] : memref<2x2x3x3x16x1xf32>
2012       affine.store %9, %1[%i0, %i1] : memref<64x9xf32>
2013     }
2014   }
2015   affine.for %i2 = 0 to 9 {
2016     affine.for %i3 = 0 to 4 {
2017       affine.for %i4 = 0 to 16 {
2018         %10 = affine.apply #map10(%i3, %i4)
2019         %11 = affine.load %1[%10, %i2] : memref<64x9xf32>
2020       }
2021       affine.for %i5 = 0 to 16 {
2022         %14 = affine.apply #map11(%i2, %i5)
2023         affine.store %arg0, %2[%14, %i3] : memref<144x4xf32>
2024       }
2025     }
2026   }
2027   affine.for %i6 = 0 to 9 {
2028     affine.for %i7 = 0 to 9 {
2029       affine.for %i8 = 0 to 4 {
2030         affine.for %i9 = 0 to 16 {
2031           %15 = affine.apply #map12(%i8, %i9)
2032           %16 = affine.load %1[%15, %i7] : memref<64x9xf32>
2033         }
2034       }
2035     }
2036   }
2037   return
2038 }
2039 // MAXIMAL-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> ((d0 * 72 + d1) floordiv 2304)>
2040 // MAXIMAL-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0, d1) -> (((d0 * 72 + d1) mod 2304) floordiv 1152)>
2041 // MAXIMAL-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) floordiv 9) floordiv 8)>
2042 // MAXIMAL-DAG: [[$MAP3:#map[0-9]+]] = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) mod 9) floordiv 3)>
2043 // MAXIMAL-DAG: [[$MAP4:#map[0-9]+]] = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) mod 9) mod 3)>
2044 // MAXIMAL-DAG: [[$MAP7:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
2045 // MAXIMAL-DAG: [[$MAP8:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 16 - d1 + 15)>
2046 // MAXIMAL-LABEL: func @fuse_across_varying_dims_complex
2047 // MAXIMAL-NEXT:  alloc() : memref<64x1xf32>
2048 // MAXIMAL-NEXT:  constant 0 : index
2049 // MAXIMAL-NEXT:  alloc() : memref<2x2x3x3x16x1xf32>
2050 // MAXIMAL-NEXT:  alloc() : memref<144x4xf32>
2051 // MAXIMAL-NEXT:  affine.for %{{.*}} = 0 to 9 {
2052 // MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 9 {
2053 // MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 4 {
2054 // MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 16 {
2055 // MAXIMAL-NEXT:          affine.for %{{.*}} = 0 to 64 {
2056 // MAXIMAL-NEXT:            affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
2057 // MAXIMAL-NEXT:            affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
2058 // MAXIMAL-NEXT:            affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
2059 // MAXIMAL-NEXT:            affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
2060 // MAXIMAL-NEXT:            affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
2061 // MAXIMAL-NEXT:            affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x1xf32>
2062 // MAXIMAL-NEXT:            affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<64x1xf32>
2063 // MAXIMAL-NEXT:          }
2064 // MAXIMAL-NEXT:          affine.for %{{.*}} = 0 to 4 {
2065 // MAXIMAL-NEXT:            affine.for %{{.*}} = 0 to 16 {
2066 // MAXIMAL-NEXT:              affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
2067 // MAXIMAL-NEXT:              affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
2068 // MAXIMAL-NEXT:            }
2069 // MAXIMAL-NEXT:            affine.for %{{.*}} = 0 to 16 {
2070 // MAXIMAL-NEXT:              affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
2071 // MAXIMAL-NEXT:              affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<144x4xf32>
2072 // MAXIMAL-NEXT:            }
2073 // MAXIMAL-NEXT:          }
2074 // MAXIMAL-NEXT:          affine.apply [[$MAP8]](%{{.*}}, %{{.*}})
2075 // MAXIMAL-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, 0] : memref<64x1xf32>
2076 // MAXIMAL-NEXT:        }
2077 // MAXIMAL-NEXT:      }
2078 // MAXIMAL-NEXT:    }
2079 // MAXIMAL-NEXT:  }
2080
2081 // -----
2082
2083 func @should_fuse_with_slice_union() {
2084   %a = alloc() : memref<100xf32>
2085   %c0 = constant 0 : index
2086   %cf0 = constant 0.0 : f32
2087
2088   affine.for %i0 = 0 to 100 {
2089     affine.store %cf0, %a[%i0]: memref<100xf32>
2090   }
2091
2092   affine.for %i1 = 10 to 20 {
2093     %v0 = affine.load %a[%i1]: memref<100xf32>
2094     affine.for %i2 = 15 to 25 {
2095       %v1 = affine.load %a[%i2]: memref<100xf32>
2096     }
2097   }
2098   // The union of two slice bounds (calculated between the store and each of
2099   // the loads) is computed and used in the fusion cost calculation, index
2100   // remapping, and private memref size. The result is that the temporary
2101   // memref is reduced from 100xf32 to 15xf32 and properly indexed by
2102   // the fused loops based on the union calculation.
2103 // CHECK:      affine.for %{{.*}} = 10 to 20 {
2104 // CHECK-NEXT:   affine.for %{{.*}} = 10 to 25 {
2105 // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[%{{.*}} - 10] : memref<15xf32>
2106 // CHECK-NEXT:   }
2107 // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}} - 10] : memref<15xf32>
2108 // CHECK-NEXT:   affine.for %{{.*}} = 15 to 25 {
2109 // CHECK-NEXT:     affine.load %{{.*}}[%{{.*}} - 10] : memref<15xf32>
2110 // CHECK-NEXT:   }
2111 // CHECK-NEXT: }
2112 // CHECK-NEXT: return
2113   return
2114 }
2115
2116 // -----
2117
2118 func @affine_add_mm_fused(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>, %arg3: memref<1024x1024xf32>) {
2119   affine.for %i2 = 0 to 1024 {
2120     affine.for %i3 = 0 to 1024 {
2121       %0 = affine.load %arg3[%i2, %i3] : memref<1024x1024xf32>
2122       %1 = affine.load %arg2[%i2, %i3] : memref<1024x1024xf32>
2123       %2 = addf %1, %0 : f32
2124       affine.store %2, %arg2[%i2, %i3] : memref<1024x1024xf32>
2125     }
2126   }
2127   affine.for %i4 = 0 to 1024 {
2128     affine.for %i5 = 0 to 1024 {
2129       affine.for %i6 = 0 to 1024 {
2130         %3 = affine.load %arg1[%i6, %i5] : memref<1024x1024xf32>
2131         %4 = affine.load %arg0[%i4, %i6] : memref<1024x1024xf32>
2132         %5 = mulf %4, %3 : f32
2133         %6 = affine.load %arg2[%i4, %i5] : memref<1024x1024xf32>
2134         %7 = addf %6, %5 : f32
2135         affine.store %7, %arg2[%i4, %i5] : memref<1024x1024xf32>
2136       }
2137     }
2138   }
2139   // Should fuse elementwise add loop at loop depth 2, above loop-carried
2140   // dependence between load/store on '%arg2', carried on reduction loop %i6.
2141   // CHECK:       affine.for %{{.*}} = 0 to 1024 {
2142   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 1024 {
2143   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2144   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2145   // CHECK-NEXT:      addf %{{.*}}, %{{.*}} : f32
2146   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2147   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 1024 {
2148   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2149   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2150   // CHECK-NEXT:        mulf %{{.*}}, %{{.*}} : f32
2151   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2152   // CHECK-NEXT:        addf %{{.*}}, %{{.*}} : f32
2153   // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2154   // CHECK-NEXT:      }
2155   // CHECK-NEXT:    }
2156   // CHECK-NEXT:  }
2157   return
2158 }
2159
2160 // -----
2161
2162 func @affine_2mm_fused(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>, %arg3: memref<1024x1024xf32>, %arg4: memref<1024x1024xf32>) {
2163   %cst = constant 0.000000e+00 : f32
2164   affine.for %i0 = 0 to 1024 {
2165     affine.for %i1 = 0 to 1024 {
2166       affine.store %cst, %arg2[%i0, %i1] : memref<1024x1024xf32>
2167     }
2168   }
2169   affine.for %i2 = 0 to 1024 {
2170     affine.for %i3 = 0 to 1024 {
2171       affine.store %cst, %arg4[%i2, %i3] : memref<1024x1024xf32>
2172     }
2173   }
2174   affine.for %i4 = 0 to 1024 {
2175     affine.for %i5 = 0 to 1024 {
2176       affine.for %i6 = 0 to 1024 {
2177         %0 = affine.load %arg1[%i6, %i5] : memref<1024x1024xf32>
2178         %1 = affine.load %arg0[%i4, %i6] : memref<1024x1024xf32>
2179         %2 = mulf %1, %0 : f32
2180         %3 = affine.load %arg2[%i4, %i5] : memref<1024x1024xf32>
2181         %4 = addf %3, %2 : f32
2182         affine.store %4, %arg2[%i4, %i5] : memref<1024x1024xf32>
2183       }
2184     }
2185   }
2186   affine.for %i7 = 0 to 1024 {
2187     affine.for %i8 = 0 to 1024 {
2188       affine.for %i9 = 0 to 1024 {
2189         %5 = affine.load %arg1[%i9, %i8] : memref<1024x1024xf32>
2190         %6 = affine.load %arg0[%i7, %i9] : memref<1024x1024xf32>
2191         %7 = mulf %6, %5 : f32
2192         %8 = affine.load %arg4[%i7, %i8] : memref<1024x1024xf32>
2193         %9 = addf %8, %7 : f32
2194         affine.store %9, %arg4[%i7, %i8] : memref<1024x1024xf32>
2195       }
2196     }
2197   }
2198
2199   // Should fuse MM initialization loops into their consumers, then fuse the
2200   // two matmul loops together for input reuse on '%arg0/%arg1'.
2201
2202   // CHECK:        affine.for %{{.*}} = 0 to 1024 {
2203   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
2204   // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2205   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
2206   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2207   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2208   // CHECK-NEXT:         mulf %{{.*}}, %{{.*}} : f32
2209   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2210   // CHECK-NEXT:         addf %{{.*}}, %{{.*}} : f32
2211   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2212   // CHECK-NEXT:       }
2213   // CHECK-NEXT:     }
2214   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
2215   // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2216   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
2217   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2218   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2219   // CHECK-NEXT:         mulf %{{.*}}, %{{.*}} : f32
2220   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2221   // CHECK-NEXT:         addf %{{.*}}, %{{.*}} : f32
2222   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2223   // CHECK-NEXT:       }
2224   // CHECK-NEXT:     }
2225   // CHECK-NEXT:   }
2226
2227   return
2228 }
2229
2230 // -----
2231
2232 func @affine_2_dependent_mm_fused(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>, %arg3: memref<1024x1024xf32>, %arg4: memref<1024x1024xf32>) {
2233   affine.for %i0 = 0 to 1024 {
2234     affine.for %i1 = 0 to 1024 {
2235       affine.for %i2 = 0 to 1024 {
2236         %0 = affine.load %arg1[%i2, %i1] : memref<1024x1024xf32>
2237         %1 = affine.load %arg0[%i0, %i2] : memref<1024x1024xf32>
2238         %2 = mulf %1, %0 : f32
2239         %3 = affine.load %arg2[%i0, %i1] : memref<1024x1024xf32>
2240         %4 = addf %3, %2 : f32
2241         affine.store %4, %arg2[%i0, %i1] : memref<1024x1024xf32>
2242       }
2243     }
2244   }
2245   affine.for %i3 = 0 to 1024 {
2246     affine.for %i4 = 0 to 1024 {
2247       affine.for %i5 = 0 to 1024 {
2248         %5 = affine.load %arg3[%i5, %i4] : memref<1024x1024xf32>
2249         %6 = affine.load %arg2[%i3, %i5] : memref<1024x1024xf32>
2250         %7 = mulf %6, %5 : f32
2251         %8 = affine.load %arg4[%i3, %i4] : memref<1024x1024xf32>
2252         %9 = addf %8, %7 : f32
2253         affine.store %9, %arg4[%i3, %i4] : memref<1024x1024xf32>
2254       }
2255     }
2256   }
2257
2258   // CHECK:        affine.for %{{.*}} = 0 to 1024 {
2259   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
2260   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
2261   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2262   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2263   // CHECK-NEXT:         mulf %{{.*}}, %{{.*}} : f32
2264   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2265   // CHECK-NEXT:         addf %{{.*}}, %{{.*}} : f32
2266   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2267   // CHECK-NEXT:       }
2268   // CHECK-NEXT:     }
2269   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
2270   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
2271   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2272   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2273   // CHECK-NEXT:         mulf %{{.*}}, %{{.*}} : f32
2274   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2275   // CHECK-NEXT:         addf %{{.*}}, %{{.*}} : f32
2276   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
2277   // CHECK-NEXT:       }
2278   // CHECK-NEXT:     }
2279   // CHECK-NEXT:   }
2280   return
2281 }
2282
2283 // -----
2284
2285 // CHECK-LABEL: func @should_fuse_self_dependence_multi_store_producer() {
2286 func @should_fuse_self_dependence_multi_store_producer() {
2287   %m = alloc() : memref<10xf32>
2288   %local_m = alloc() : memref<10xf32>
2289   %cf7 = constant 7.0 : f32
2290
2291   affine.for %i0 = 0 to 10 {
2292     affine.store %cf7, %local_m[%i0] : memref<10xf32>
2293     %v0 = affine.load %local_m[%i0] : memref<10xf32>
2294     affine.store %v0, %m[%i0] : memref<10xf32>
2295   }
2296   affine.for %i1 = 0 to 10 {
2297     %v1 = affine.load %m[%i1] : memref<10xf32>
2298   }
2299   // CHECK:      affine.for %[[i0:.*]] = 0 to 10 {
2300   // CHECK-NEXT:   affine.store %{{.*}}, [[LOCAL_M:%.*]][%[[i0]]] : memref<10xf32>
2301   // CHECK-NEXT:   [[v0:%.*]] = affine.load [[LOCAL_M]][%[[i0]]] : memref<10xf32>
2302   // CHECK-NEXT:   affine.store [[v0]], %{{.*}}[0] : memref<1xf32>
2303   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
2304   // CHECK-NEXT: }
2305   // CHECK-NEXT: return
2306   return
2307 }
2308
2309 // -----
2310
2311 // CHECK-LABEL: func @should_fuse_dead_multi_store_producer() {
2312 func @should_fuse_dead_multi_store_producer() {
2313   %m = alloc() : memref<10xf32>
2314   %dead_m = alloc() : memref<10xf32>
2315   %cf7 = constant 7.0 : f32
2316
2317   affine.for %i0 = 0 to 10 {
2318     affine.store %cf7, %dead_m[%i0] : memref<10xf32>
2319     affine.store %cf7, %m[%i0] : memref<10xf32>
2320   }
2321   affine.for %i1 = 0 to 10 {
2322     %v0 = affine.load %m[%i1] : memref<10xf32>
2323   }
2324   // CHECK:      affine.for %[[i0:.*]] = 0 to 10 {
2325   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%[[i0]]] : memref<10xf32>
2326   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
2327   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
2328   // CHECK-NEXT: }
2329   // CHECK-NEXT: return
2330   return
2331 }
2332
2333 // -----
2334
2335 // CHECK-LABEL: func @should_fuse_function_live_out_multi_store_producer
2336 func @should_fuse_function_live_out_multi_store_producer(%live_in_out_m : memref<10xf32>) {
2337   %m = alloc() : memref<10xf32>
2338   %cf7 = constant 7.0 : f32
2339
2340   affine.for %i0 = 0 to 10 {
2341     affine.store %cf7, %live_in_out_m[%i0] : memref<10xf32>
2342     affine.store %cf7, %m[%i0] : memref<10xf32>
2343   }
2344   affine.for %i1 = 0 to 10 {
2345     %v0 = affine.load %m[%i1] : memref<10xf32>
2346   }
2347   // CHECK:      affine.for %[[i0:.*]] = 0 to 10 {
2348   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%[[i0]]] : memref<10xf32>
2349   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
2350   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
2351   // CHECK-NEXT: }
2352   // CHECK-NEXT: return
2353   return
2354 }
2355
2356 // -----
2357
2358 // Test case from github bug 777.
2359 // CHECK-LABEL: func @mul_add_0
2360 func @mul_add_0(%arg0: memref<3x4xf32>, %arg1: memref<4x3xf32>, %arg2: memref<3x3xf32>, %arg3: memref<3x3xf32>) {
2361   %cst = constant 0.000000e+00 : f32
2362   %0 = alloc() : memref<3x3xf32>
2363   affine.for %arg4 = 0 to 3 {
2364     affine.for %arg5 = 0 to 3 {
2365       affine.store %cst, %0[%arg4, %arg5] : memref<3x3xf32>
2366     }
2367   }
2368   affine.for %arg4 = 0 to 3 {
2369     affine.for %arg5 = 0 to 3 {
2370       affine.for %arg6 = 0 to 4 {
2371         %1 = affine.load %arg1[%arg6, %arg5] : memref<4x3xf32>
2372         %2 = affine.load %arg0[%arg4, %arg6] : memref<3x4xf32>
2373         %3 = mulf %2, %1 : f32
2374         %4 = affine.load %0[%arg4, %arg5] : memref<3x3xf32>
2375         %5 = addf %4, %3 : f32
2376         affine.store %5, %0[%arg4, %arg5] : memref<3x3xf32>
2377       }
2378     }
2379   }
2380   affine.for %arg4 = 0 to 3 {
2381     affine.for %arg5 = 0 to 3 {
2382       %6 = affine.load %arg2[%arg4, %arg5] : memref<3x3xf32>
2383       %7 = affine.load %0[%arg4, %arg5] : memref<3x3xf32>
2384       %8 = addf %7, %6 : f32
2385       affine.store %8, %arg3[%arg4, %arg5] : memref<3x3xf32>
2386     }
2387   }
2388   // CHECK:      affine.for %[[i0:.*]] = 0 to 3 {
2389   // CHECK-NEXT:   affine.for %[[i1:.*]] = 0 to 3 {
2390   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
2391   // CHECK-NEXT:     affine.for %[[i2:.*]] = 0 to 4 {
2392   // CHECK-NEXT:       affine.load %{{.*}}[%[[i2]], %[[i1]]] : memref<4x3xf32>
2393   // CHECK-NEXT:       affine.load %{{.*}}[%[[i0]], %[[i2]]] : memref<3x4xf32>
2394   // CHECK-NEXT:       mulf %{{.*}}, %{{.*}} : f32
2395   // CHECK-NEXT:       affine.load %{{.*}}[0, 0] : memref<1x1xf32>
2396   // CHECK-NEXT:       addf %{{.*}}, %{{.*}} : f32
2397   // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
2398   // CHECK-NEXT:     }
2399   // CHECK-NEXT:     affine.load %{{.*}}[%[[i0]], %[[i1]]] : memref<3x3xf32>
2400   // CHECK-NEXT:     affine.load %{{.*}}[0, 0] : memref<1x1xf32>
2401   // CHECK-NEXT:     addf %{{.*}}, %{{.*}} : f32
2402   // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[%[[i0]], %[[i1]]] : memref<3x3xf32>
2403   // CHECK-NEXT:   }
2404   // CHECK-NEXT: }
2405   // CHECK-NEXT: return
2406   return
2407 }
2408
2409 // -----
2410
2411 // Verify that 'fuseProducerConsumerNodes' fuse a producer loop with a store
2412 // that has multiple outgoing edges.
2413
2414 // CHECK-LABEL: func @should_fuse_multi_outgoing_edge_store_producer
2415 func @should_fuse_multi_outgoing_edge_store_producer(%a : memref<1xf32>) {
2416   %cst = constant 0.000000e+00 : f32
2417   affine.for %arg0 = 0 to 1 {
2418     affine.store %cst, %a[%arg0] : memref<1xf32>
2419   }
2420
2421   affine.for %arg0 = 0 to 1 {
2422     %0 = affine.load %a[%arg0] : memref<1xf32>
2423   }
2424
2425   affine.for %arg0 = 0 to 1 {
2426     %0 = affine.load %a[%arg0] : memref<1xf32>
2427   }
2428   // CHECK:      affine.for %{{.*}} = 0 to 1 {
2429   // CHECK-NEXT:   affine.store
2430   // CHECK-NEXT:   affine.load
2431   // CHECK-NEXT:   affine.load
2432   // CHECK-NEXT: }
2433
2434   return
2435 }
2436
2437 // -----
2438
2439 // Verify that 'fuseProducerConsumerNodes' fuses a producer loop that: 1) has
2440 // multiple outgoing edges, 2) producer store has a single outgoing edge.
2441 // Sibling loop fusion should not fuse any of these loops due to
2442 // dependencies on external memrefs '%a' and '%b'.
2443
2444 // CHECK-LABEL: func @should_fuse_producer_with_multi_outgoing_edges
2445 func @should_fuse_producer_with_multi_outgoing_edges(%a : memref<1xf32>, %b : memref<1xf32>) {
2446   %cst = constant 0.000000e+00 : f32
2447   affine.for %arg0 = 0 to 1 {
2448     %0 = affine.load %a[%arg0] : memref<1xf32>
2449     affine.store %cst, %b[%arg0] : memref<1xf32>
2450   }
2451
2452   affine.for %arg0 = 0 to 1 {
2453     affine.store %cst, %a[%arg0] : memref<1xf32>
2454     %1 = affine.load %b[%arg0] : memref<1xf32>
2455   }
2456   // CHECK: affine.for %{{.*}} = 0 to 1
2457   // CHECK-NEXT: affine.load %[[A:.*]][{{.*}}]
2458   // CHECK-NEXT: affine.store %{{.*}}, %[[B:.*]][{{.*}}]
2459   // CHECK-NEXT: affine.store %{{.*}}, %[[A]]
2460   // CHECK-NEXT: affine.load %[[B]]
2461   // CHECK-NOT: affine.for %{{.*}}
2462   // CHECK: return
2463   return
2464 }
2465
2466 // -----
2467
2468 // MAXIMAL-LABEL: func @reshape_into_matmul
2469 func @reshape_into_matmul(%lhs : memref<1024x1024xf32>,
2470               %R: memref<16x64x1024xf32>, %out: memref<1024x1024xf32>) {
2471   %rhs = alloc() :  memref<1024x1024xf32>
2472
2473   // Reshape from 3-d to 2-d.
2474   affine.for %i0 = 0 to 16 {
2475     affine.for %i1 = 0 to 64 {
2476       affine.for %k = 0 to 1024 {
2477         %v = affine.load %R[%i0, %i1, %k] : memref<16x64x1024xf32>
2478         affine.store %v, %rhs[64*%i0 + %i1, %k] : memref<1024x1024xf32>
2479       }
2480     }
2481   }
2482
2483   // Matmul.
2484   affine.for %i = 0 to 1024 {
2485     affine.for %j = 0 to 1024 {
2486       affine.for %k = 0 to 1024 {
2487         %0 = affine.load %rhs[%k, %j] : memref<1024x1024xf32>
2488         %1 = affine.load %lhs[%i, %k] : memref<1024x1024xf32>
2489         %2 = mulf %1, %0 : f32
2490         %3 = affine.load %out[%i, %j] : memref<1024x1024xf32>
2491         %4 = addf %3, %2 : f32
2492         affine.store %4, %out[%i, %j] : memref<1024x1024xf32>
2493       }
2494     }
2495   }
2496   return
2497 }
2498 // MAXIMAL-NEXT: alloc
2499 // MAXIMAL-NEXT: affine.for
2500 // MAXIMAL-NEXT:   affine.for
2501 // MAXIMAL-NEXT:     affine.for
2502 // MAXIMAL-NOT:      affine.for
2503 // MAXIMAL:      return
2504
2505 // -----
2506
2507 // CHECK-LABEL: func @vector_loop
2508 func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>,
2509                   %c : memref<10x20xf32>) {
2510   affine.for %j = 0 to 10 {
2511     affine.for %i = 0 to 5 {
2512       %ld0 = affine.vector_load %a[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
2513       affine.vector_store %ld0, %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
2514     }
2515   }
2516
2517   affine.for %j = 0 to 10 {
2518     affine.for %i = 0 to 5 {
2519       %ld0 = affine.vector_load %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
2520       affine.vector_store %ld0, %c[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
2521     }
2522   }
2523
2524   return
2525 }
2526 // CHECK:      affine.for
2527 // CHECK-NEXT:   affine.for
2528 // CHECK-NEXT:     affine.vector_load
2529 // CHECK-NEXT:     affine.vector_store
2530 // CHECK-NEXT:     affine.vector_load
2531 // CHECK-NEXT:     affine.vector_store
2532 // CHECK-NOT:  affine.for
2533
2534 // -----
2535
2536 // CHECK-LABEL: func @multi_outgoing_edges
2537 func @multi_outgoing_edges(%in0 : memref<32xf32>,
2538                       %in1 : memref<32xf32>) {
2539   affine.for %d = 0 to 32 {
2540     %lhs = affine.load %in0[%d] : memref<32xf32>
2541     %rhs = affine.load %in1[%d] : memref<32xf32>
2542     %add = addf %lhs, %rhs : f32
2543     affine.store %add, %in0[%d] : memref<32xf32>
2544   }
2545   affine.for %d = 0 to 32 {
2546     %lhs = affine.load %in0[%d] : memref<32xf32>
2547     %rhs = affine.load %in1[%d] : memref<32xf32>
2548     %add = subf %lhs, %rhs : f32
2549     affine.store %add, %in0[%d] : memref<32xf32>
2550   }
2551   affine.for %d = 0 to 32 {
2552     %lhs = affine.load %in0[%d] : memref<32xf32>
2553     %rhs = affine.load %in1[%d] : memref<32xf32>
2554     %add = mulf %lhs, %rhs : f32
2555     affine.store %add, %in0[%d] : memref<32xf32>
2556   }
2557   affine.for %d = 0 to 32 {
2558     %lhs = affine.load %in0[%d] : memref<32xf32>
2559     %rhs = affine.load %in1[%d] : memref<32xf32>
2560     %add = divf %lhs, %rhs : f32
2561     affine.store %add, %in0[%d] : memref<32xf32>
2562   }
2563   return
2564 }
2565
2566 // CHECK:      affine.for
2567 // CHECK-NOT:  affine.for
2568 // CHECK:        addf
2569 // CHECK-NOT:  affine.for
2570 // CHECK:        subf
2571 // CHECK-NOT:  affine.for
2572 // CHECK:        mulf
2573 // CHECK-NOT:  affine.for
2574 // CHECK:        divf
2575
2576 // -----
2577
2578 // Test fusion when dynamically shaped memrefs are used with constant trip count loops.
2579
2580 // CHECK-LABEL: func @calc
2581 func @calc(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %len: index) {
2582   %c1 = constant 1 : index
2583   %1 = alloc(%len) : memref<?xf32>
2584   affine.for %arg4 = 1 to 10 {
2585     %7 = affine.load %arg0[%arg4] : memref<?xf32>
2586     %8 = affine.load %arg1[%arg4] : memref<?xf32>
2587     %9 = addf %7, %8 : f32
2588     affine.store %9, %1[%arg4] : memref<?xf32>
2589   }
2590   affine.for %arg4 = 1 to 10 {
2591     %7 = affine.load %1[%arg4] : memref<?xf32>
2592     %8 = affine.load %arg1[%arg4] : memref<?xf32>
2593     %9 = mulf %7, %8 : f32
2594     affine.store %9, %arg2[%arg4] : memref<?xf32>
2595   }
2596   return
2597 }
2598 // CHECK:       alloc() : memref<1xf32>
2599 // CHECK:       affine.for %arg{{.*}} = 1 to 10 {
2600 // CHECK-NEXT:    affine.load %arg{{.*}}
2601 // CHECK-NEXT:    affine.load %arg{{.*}}
2602 // CHECK-NEXT:    addf
2603 // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
2604 // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
2605 // CHECK-NEXT:    affine.load %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
2606 // CHECK-NEXT:    mulf
2607 // CHECK-NEXT:    affine.store %{{.*}}, %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
2608 // CHECK-NEXT:  }
2609 // CHECK-NEXT:  return
2610
2611 // -----
2612
2613 // CHECK-LABEL: func @should_not_fuse_since_non_affine_users
2614 func @should_not_fuse_since_non_affine_users(%in0 : memref<32xf32>,
2615                       %in1 : memref<32xf32>) {
2616   affine.for %d = 0 to 32 {
2617     %lhs = affine.load %in0[%d] : memref<32xf32>
2618     %rhs = affine.load %in1[%d] : memref<32xf32>
2619     %add = addf %lhs, %rhs : f32
2620     affine.store %add, %in0[%d] : memref<32xf32>
2621   }
2622   affine.for %d = 0 to 32 {
2623     %lhs = load %in0[%d] : memref<32xf32>
2624     %rhs = load %in1[%d] : memref<32xf32>
2625     %add = subf %lhs, %rhs : f32
2626     store %add, %in0[%d] : memref<32xf32>
2627   }
2628   affine.for %d = 0 to 32 {
2629     %lhs = affine.load %in0[%d] : memref<32xf32>
2630     %rhs = affine.load %in1[%d] : memref<32xf32>
2631     %add = mulf %lhs, %rhs : f32
2632     affine.store %add, %in0[%d] : memref<32xf32>
2633   }
2634   return
2635 }
2636
2637 // CHECK:  affine.for
2638 // CHECK:    addf
2639 // CHECK:  affine.for
2640 // CHECK:    subf
2641 // CHECK:  affine.for
2642 // CHECK:    mulf
2643
2644 // -----
2645
2646 // CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_users
2647 func @should_not_fuse_since_top_level_non_affine_users(%in0 : memref<32xf32>,
2648                       %in1 : memref<32xf32>) {
2649   %sum = alloc() : memref<f32>
2650   affine.for %d = 0 to 32 {
2651     %lhs = affine.load %in0[%d] : memref<32xf32>
2652     %rhs = affine.load %in1[%d] : memref<32xf32>
2653     %add = addf %lhs, %rhs : f32
2654     store %add, %sum[] : memref<f32>
2655     affine.store %add, %in0[%d] : memref<32xf32>
2656   }
2657   %load_sum = load %sum[] : memref<f32>
2658   affine.for %d = 0 to 32 {
2659     %lhs = affine.load %in0[%d] : memref<32xf32>
2660     %rhs = affine.load %in1[%d] : memref<32xf32>
2661     %add = mulf %lhs, %rhs : f32
2662     %sub = subf %add, %load_sum: f32
2663     affine.store %sub, %in0[%d] : memref<32xf32>
2664   }
2665   dealloc %sum : memref<f32>
2666   return
2667 }
2668
2669 // CHECK:  affine.for
2670 // CHECK:    addf
2671 // CHECK:  affine.for
2672 // CHECK:    mulf
2673 // CHECK:    subf
2674
2675 // -----
2676
2677 // MAXIMAL-LABEL: func @fuse_minor_affine_map
2678 func @fuse_minor_affine_map(%in: memref<128xf32>, %out: memref<20x512xf32>) {
2679   %tmp = alloc() : memref<128xf32>
2680
2681   affine.for %arg4 = 0 to 128 {
2682     %ld = affine.load %in[%arg4] : memref<128xf32>
2683     affine.store %ld, %tmp[%arg4] : memref<128xf32>
2684   }
2685
2686   affine.for %arg3 = 0 to 20 {
2687     affine.for %arg4 = 0 to 512 {
2688       %ld = affine.load %tmp[%arg4 mod 128] : memref<128xf32>
2689       affine.store %ld, %out[%arg3, %arg4] : memref<20x512xf32>
2690     }
2691   }
2692
2693   return
2694 }
2695
2696 // TODO: The size of the private memref is not properly computed in the presence
2697 // of the 'mod' operation. It should be memref<1xf32> instead of
2698 // memref<128xf32>: https://bugs.llvm.org/show_bug.cgi?id=46973
2699 // MAXIMAL:       alloc() : memref<128xf32>
2700 // MAXIMAL:       affine.for
2701 // MAXIMAL-NEXT:    affine.for
2702 // MAXIMAL-NOT:   affine.for
2703 // MAXIMAL:       return
2704
2705 // -----
2706
2707 // CHECK-LABEL: func @should_fuse_multi_store_producer_and_privatize_memfefs
2708 func @should_fuse_multi_store_producer_and_privatize_memfefs() {
2709   %a = alloc() : memref<10xf32>
2710   %b = alloc() : memref<10xf32>
2711   %c = alloc() : memref<10xf32>
2712   %cst = constant 0.000000e+00 : f32
2713   affine.for %arg0 = 0 to 10 {
2714     affine.store %cst, %a[%arg0] : memref<10xf32>
2715     affine.store %cst, %b[%arg0] : memref<10xf32>
2716     affine.store %cst, %c[%arg0] : memref<10xf32>
2717     %0 = affine.load %c[%arg0] : memref<10xf32>
2718   }
2719
2720   affine.for %arg0 = 0 to 10 {
2721     %0 = affine.load %a[%arg0] : memref<10xf32>
2722   }
2723
2724   affine.for %arg0 = 0 to 10 {
2725     %0 = affine.load %b[%arg0] : memref<10xf32>
2726   }
2727
2728         // All the memrefs should be privatized except '%c', which is not involved in
2729   // the producer-consumer fusion.
2730   // CHECK:      affine.for %{{.*}} = 0 to 10 {
2731   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
2732   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
2733   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2734   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
2735   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
2736   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
2737   // CHECK-NEXT: }
2738
2739   return
2740 }
2741
2742 // -----
2743
2744 func @should_fuse_multi_store_producer_with_scaping_memrefs_and_remove_src(
2745     %a : memref<10xf32>, %b : memref<10xf32>) {
2746   %cst = constant 0.000000e+00 : f32
2747   affine.for %i0 = 0 to 10 {
2748     affine.store %cst, %a[%i0] : memref<10xf32>
2749     affine.store %cst, %b[%i0] : memref<10xf32>
2750   }
2751
2752   affine.for %i1 = 0 to 10 {
2753     %0 = affine.load %a[%i1] : memref<10xf32>
2754   }
2755
2756   affine.for %i2 = 0 to 10 {
2757     %0 = affine.load %b[%i2] : memref<10xf32>
2758   }
2759
2760         // Producer loop '%i0' should be removed after fusion since fusion is maximal.
2761   // No memref should be privatized since they escape the function.
2762   // CHECK:       affine.for %{{.*}} = 0 to 10 {
2763   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2764   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2765   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
2766   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
2767   // CHECK-NEXT:  }
2768   // CHECK-NOT:   affine.for
2769
2770   return
2771 }
2772
2773 // -----
2774
2775 func @should_fuse_multi_store_producer_with_scaping_memrefs_and_preserve_src(
2776     %a : memref<10xf32>, %b : memref<10xf32>) {
2777   %cst = constant 0.000000e+00 : f32
2778   affine.for %i0 = 0 to 10 {
2779     affine.store %cst, %a[%i0] : memref<10xf32>
2780     affine.store %cst, %b[%i0] : memref<10xf32>
2781   }
2782
2783   affine.for %i1 = 0 to 5 {
2784     %0 = affine.load %a[%i1] : memref<10xf32>
2785   }
2786
2787   affine.for %i2 = 0 to 10 {
2788     %0 = affine.load %b[%i2] : memref<10xf32>
2789   }
2790
2791         // Loops '%i0' and '%i2' should be fused first and '%i0' should be removed
2792   // since fusion is maximal. Then the fused loop and '%i1' should be fused
2793   // and the fused loop shouldn't be removed since fusion is not maximal.
2794   // CHECK:       affine.for %{{.*}} = 0 to 10 {
2795   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2796   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2797   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
2798   // CHECK-NEXT:  }
2799   // CHECK:       affine.for %{{.*}} = 0 to 5 {
2800   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2801   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
2802   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
2803   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
2804   // CHECK-NEXT:  }
2805   // CHECK-NOT:   affine.for
2806
2807   return
2808 }