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