1// RUN: mlir-opt -buffer-hoisting -split-input-file %s | FileCheck %s
2
3// This file checks the behaviour of BufferHoisting pass for moving Alloc
4// operations to their correct positions.
5
6// Test Case:
7//    bb0
8//   /   \
9//  bb1  bb2 <- Initial position of AllocOp
10//   \   /
11//    bb3
12// BufferHoisting expected behavior: It should move the existing AllocOp to
13// the entry block.
14
15// CHECK-LABEL: func @condBranch
16func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
17  cond_br %arg0, ^bb1, ^bb2
18^bb1:
19  br ^bb3(%arg1 : memref<2xf32>)
20^bb2:
21  %0 = memref.alloc() : memref<2xf32>
22  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
23  br ^bb3(%0 : memref<2xf32>)
24^bb3(%1: memref<2xf32>):
25  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
26  return
27}
28
29// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
30// CHECK-NEXT: cond_br
31
32// -----
33
34// Test Case:
35//    bb0
36//   /   \
37//  bb1  bb2 <- Initial position of AllocOp
38//   \   /
39//    bb3
40// BufferHoisting expected behavior: It should not move the existing AllocOp
41// to any other block since the alloc has a dynamic dependency to block argument
42// %0 in bb2.
43
44// CHECK-LABEL: func @condBranchDynamicType
45func @condBranchDynamicType(
46  %arg0: i1,
47  %arg1: memref<?xf32>,
48  %arg2: memref<?xf32>,
49  %arg3: index) {
50  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
51^bb1:
52  br ^bb3(%arg1 : memref<?xf32>)
53^bb2(%0: index):
54  %1 = memref.alloc(%0) : memref<?xf32>
55  test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
56  br ^bb3(%1 : memref<?xf32>)
57^bb3(%2: memref<?xf32>):
58  test.copy(%2, %arg2) : (memref<?xf32>, memref<?xf32>)
59  return
60}
61
62// CHECK-NEXT: cond_br
63//      CHECK: ^bb2
64//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
65// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])
66// CHECK-NEXT: test.buffer_based
67
68// -----
69
70// Test Case:
71//      bb0
72//     /    \
73//   bb1    bb2 <- Initial position of AllocOp
74//    |     /  \
75//    |   bb3  bb4
76//    |     \  /
77//    \     bb5
78//     \    /
79//       bb6
80//        |
81//       bb7
82// BufferHoisting expected behavior: It should not move the existing AllocOp
83// to any other block since the alloc has a dynamic dependency to block argument
84// %0 in bb2.
85
86// CHECK-LABEL: func @condBranchDynamicTypeNested
87func @condBranchDynamicTypeNested(
88  %arg0: i1,
89  %arg1: memref<?xf32>,
90  %arg2: memref<?xf32>,
91  %arg3: index) {
92  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
93^bb1:
94  br ^bb6(%arg1 : memref<?xf32>)
95^bb2(%0: index):
96  %1 = memref.alloc(%0) : memref<?xf32>
97  test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
98  cond_br %arg0, ^bb3, ^bb4
99^bb3:
100  br ^bb5(%1 : memref<?xf32>)
101^bb4:
102  br ^bb5(%1 : memref<?xf32>)
103^bb5(%2: memref<?xf32>):
104  br ^bb6(%2 : memref<?xf32>)
105^bb6(%3: memref<?xf32>):
106  br ^bb7(%3 : memref<?xf32>)
107^bb7(%4: memref<?xf32>):
108  test.copy(%4, %arg2) : (memref<?xf32>, memref<?xf32>)
109  return
110}
111
112// CHECK-NEXT: cond_br
113//      CHECK: ^bb2
114//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
115// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])
116// CHECK-NEXT: test.buffer_based
117
118// -----
119
120// Test Case:
121//    bb0
122//   /   \
123//  |    bb1 <- Initial position of AllocOp
124//   \   /
125//    bb2
126// BufferHoisting expected behavior: It should move the existing AllocOp to
127// the entry block.
128
129// CHECK-LABEL: func @criticalEdge
130func @criticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
131  cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
132^bb1:
133  %0 = memref.alloc() : memref<2xf32>
134  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
135  br ^bb2(%0 : memref<2xf32>)
136^bb2(%1: memref<2xf32>):
137  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
138  return
139}
140
141// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
142// CHECK-NEXT: cond_br
143
144// -----
145
146// Test Case:
147//    bb0 <- Initial position of the first AllocOp
148//   /   \
149//  bb1  bb2
150//   \   /
151//    bb3 <- Initial position of the second AllocOp
152// BufferHoisting expected behavior: It shouldn't move the AllocOps.
153
154// CHECK-LABEL: func @ifElse
155func @ifElse(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
156  %0 = memref.alloc() : memref<2xf32>
157  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
158  cond_br %arg0,
159    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
160    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
161^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
162  br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
163^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
164  br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
165^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
166  %7 = memref.alloc() : memref<2xf32>
167  test.buffer_based in(%7: memref<2xf32>) out(%7: memref<2xf32>)
168  test.copy(%7, %arg2) : (memref<2xf32>, memref<2xf32>)
169  return
170}
171
172// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
173// CHECK-NEXT: test.buffer_based
174//      CHECK: br ^bb3
175//      CHECK: br ^bb3
176// CHECK-NEXT: ^bb3
177//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
178// CHECK-NEXT: test.buffer_based
179//      CHECK: test.copy(%[[ALLOC1]]
180// CHECK-NEXT: return
181
182// -----
183
184// Test Case: No users for buffer in if-else CFG
185//    bb0 <- Initial position of AllocOp
186//   /   \
187//  bb1  bb2
188//   \   /
189//    bb3
190// BufferHoisting expected behavior: It shouldn't move the AllocOp.
191
192// CHECK-LABEL: func @ifElseNoUsers
193func @ifElseNoUsers(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
194  %0 = memref.alloc() : memref<2xf32>
195  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
196  cond_br %arg0,
197    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
198    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
199^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
200  br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
201^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
202  br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
203^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
204  test.copy(%arg1, %arg2) : (memref<2xf32>, memref<2xf32>)
205  return
206}
207
208// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
209// CHECK-NEXT: test.buffer_based
210
211// -----
212
213// Test Case:
214//      bb0 <- Initial position of the first AllocOp
215//     /    \
216//   bb1    bb2
217//    |     /  \
218//    |   bb3  bb4
219//    \     \  /
220//     \     /
221//       bb5 <- Initial position of the second AllocOp
222// BufferHoisting expected behavior: AllocOps shouldn't be moved.
223
224// CHECK-LABEL: func @ifElseNested
225func @ifElseNested(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
226  %0 = memref.alloc() : memref<2xf32>
227  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
228  cond_br %arg0,
229    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
230    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
231^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
232  br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
233^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
234  cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
235^bb3(%5: memref<2xf32>):
236  br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
237^bb4(%6: memref<2xf32>):
238  br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
239^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
240  %9 = memref.alloc() : memref<2xf32>
241  test.buffer_based in(%7: memref<2xf32>) out(%9: memref<2xf32>)
242  test.copy(%9, %arg2) : (memref<2xf32>, memref<2xf32>)
243  return
244}
245
246// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
247// CHECK-NEXT: test.buffer_based
248//      CHECK: br ^bb5
249//      CHECK: br ^bb5
250//      CHECK: br ^bb5
251// CHECK-NEXT: ^bb5
252//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
253// CHECK-NEXT: test.buffer_based
254
255// -----
256
257// Test Case: Dead operations in a single block.
258// BufferHoisting expected behavior: It shouldn't move the AllocOps.
259
260// CHECK-LABEL: func @redundantOperations
261func @redundantOperations(%arg0: memref<2xf32>) {
262  %0 = memref.alloc() : memref<2xf32>
263  test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
264  %1 = memref.alloc() : memref<2xf32>
265  test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
266  return
267}
268
269// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
270// CHECK-NEXT: test.buffer_based
271//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
272// CHECK-NEXT: test.buffer_based
273
274// -----
275
276// Test Case:
277//                                     bb0
278//                                    /   \
279// Initial pos of the 1st AllocOp -> bb1  bb2 <- Initial pos of the 2nd AllocOp
280//                                    \   /
281//                                     bb3
282// BufferHoisting expected behavior: Both AllocOps should be moved to the
283// entry block.
284
285// CHECK-LABEL: func @moving_alloc_and_inserting_missing_dealloc
286func @moving_alloc_and_inserting_missing_dealloc(
287  %cond: i1,
288    %arg0: memref<2xf32>,
289    %arg1: memref<2xf32>) {
290  cond_br %cond, ^bb1, ^bb2
291^bb1:
292  %0 = memref.alloc() : memref<2xf32>
293  test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
294  br ^exit(%0 : memref<2xf32>)
295^bb2:
296  %1 = memref.alloc() : memref<2xf32>
297  test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
298  br ^exit(%1 : memref<2xf32>)
299^exit(%arg2: memref<2xf32>):
300  test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
301  return
302}
303
304// CHECK-NEXT: %{{.*}} = memref.alloc()
305// CHECK-NEXT: %{{.*}} = memref.alloc()
306// CHECK-NEXT: cond_br
307
308// -----
309
310// Test Case: Invalid position of the DeallocOp. There is a user after
311// deallocation.
312//   bb0
313//  /   \
314// bb1  bb2 <- Initial position of AllocOp
315//  \   /
316//   bb3
317// BufferHoisting expected behavior: It should move the AllocOp to the entry
318// block.
319
320// CHECK-LABEL: func @moving_invalid_dealloc_op_complex
321func @moving_invalid_dealloc_op_complex(
322  %cond: i1,
323    %arg0: memref<2xf32>,
324    %arg1: memref<2xf32>) {
325  cond_br %cond, ^bb1, ^bb2
326^bb1:
327  br ^exit(%arg0 : memref<2xf32>)
328^bb2:
329  %1 = memref.alloc() : memref<2xf32>
330  test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
331  memref.dealloc %1 : memref<2xf32>
332  br ^exit(%1 : memref<2xf32>)
333^exit(%arg2: memref<2xf32>):
334  test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
335  return
336}
337
338// CHECK-NEXT: %{{.*}} = memref.alloc()
339// CHECK-NEXT: cond_br
340
341// -----
342
343// Test Case: Nested regions - This test defines a BufferBasedOp inside the
344// region of a RegionBufferBasedOp.
345// BufferHoisting expected behavior: The AllocOp for the BufferBasedOp should
346// remain inside the region of the RegiobBufferBasedOp. The AllocOp of the
347// RegionBufferBasedOp should be moved to the entry block.
348
349// CHECK-LABEL: func @nested_regions_and_cond_branch
350func @nested_regions_and_cond_branch(
351  %arg0: i1,
352  %arg1: memref<2xf32>,
353  %arg2: memref<2xf32>) {
354  cond_br %arg0, ^bb1, ^bb2
355^bb1:
356  br ^bb3(%arg1 : memref<2xf32>)
357^bb2:
358  %0 = memref.alloc() : memref<2xf32>
359  test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
360  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
361    %1 = memref.alloc() : memref<2xf32>
362    test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
363    %tmp1 = math.exp %gen1_arg0 : f32
364    test.region_yield %tmp1 : f32
365  }
366  br ^bb3(%0 : memref<2xf32>)
367^bb3(%1: memref<2xf32>):
368  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
369  return
370}
371// CHECK-NEXT:   %[[ALLOC0:.*]] = memref.alloc()
372// CHECK-NEXT:   cond_br
373//      CHECK:   test.region_buffer_based
374//      CHECK:     %[[ALLOC1:.*]] = memref.alloc()
375// CHECK-NEXT:     test.buffer_based
376
377// -----
378
379// Test Case: nested region control flow
380// The alloc position of %1 does not need to be changed and flows through
381// both if branches until it is finally returned.
382
383// CHECK-LABEL: func @nested_region_control_flow
384func @nested_region_control_flow(
385  %arg0 : index,
386  %arg1 : index) -> memref<?x?xf32> {
387  %0 = cmpi eq, %arg0, %arg1 : index
388  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
389  %2 = scf.if %0 -> (memref<?x?xf32>) {
390    scf.yield %1 : memref<?x?xf32>
391  } else {
392    %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
393    scf.yield %1 : memref<?x?xf32>
394  }
395  return %2 : memref<?x?xf32>
396}
397
398//      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
399// CHECK-NEXT: %{{.*}} = scf.if
400//      CHECK: else
401// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
402
403// -----
404
405// Test Case: nested region control flow with a nested buffer allocation in a
406// divergent branch.
407// The alloc positions of %1 does not need to be changed. %3 is moved upwards.
408
409// CHECK-LABEL: func @nested_region_control_flow_div
410func @nested_region_control_flow_div(
411  %arg0 : index,
412  %arg1 : index) -> memref<?x?xf32> {
413  %0 = cmpi eq, %arg0, %arg1 : index
414  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
415  %2 = scf.if %0 -> (memref<?x?xf32>) {
416    scf.yield %1 : memref<?x?xf32>
417  } else {
418    %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
419    scf.yield %3 : memref<?x?xf32>
420  }
421  return %2 : memref<?x?xf32>
422}
423
424//      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
425// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
426// CHECK-NEXT: %{{.*}} = scf.if
427
428// -----
429
430// Test Case: deeply nested region control flow with a nested buffer allocation
431// in a divergent branch.
432// The alloc position of %1 does not need to be changed. Allocs %4 and %5 are
433// moved upwards.
434
435// CHECK-LABEL: func @nested_region_control_flow_div_nested
436func @nested_region_control_flow_div_nested(
437  %arg0 : index,
438  %arg1 : index) -> memref<?x?xf32> {
439  %0 = cmpi eq, %arg0, %arg1 : index
440  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
441  %2 = scf.if %0 -> (memref<?x?xf32>) {
442    %3 = scf.if %0 -> (memref<?x?xf32>) {
443      scf.yield %1 : memref<?x?xf32>
444    } else {
445      %4 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
446      scf.yield %4 : memref<?x?xf32>
447    }
448    scf.yield %3 : memref<?x?xf32>
449  } else {
450    %5 = memref.alloc(%arg1, %arg1) : memref<?x?xf32>
451    scf.yield %5 : memref<?x?xf32>
452  }
453  return %2 : memref<?x?xf32>
454}
455//      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
456// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
457// CHECK-NEXT: %[[ALLOC2:.*]] = memref.alloc(%arg1, %arg1)
458// CHECK-NEXT: %{{.*}} = scf.if
459
460// -----
461
462// Test Case: deeply nested region control flow with a nested buffer allocation
463// that has dependency within a nested region should not be moved outside of
464// this region.
465
466// CHECK-LABEL: func @nested_region_control_flow_div_nested_dependencies
467func @nested_region_control_flow_div_nested_dependencies(
468  %arg0: i32,
469  %arg1: i1,
470  %arg2: index) -> memref<?x?xf32> {
471  %0 = scf.if %arg1 -> (memref<?x?xf32>) {
472    %1 = constant 1 : i32
473    %2 = addi %arg0, %1 : i32
474    %3 = index_cast %2 : i32 to index
475    %4 = memref.alloc(%arg2, %3) : memref<?x?xf32>
476    scf.yield %4 : memref<?x?xf32>
477  } else {
478    %1 = constant 2 : i32
479    %2 = addi %arg0, %1 : i32
480    %3 = index_cast %2 : i32 to index
481    %4 = memref.alloc(%arg2, %3) : memref<?x?xf32>
482    scf.yield %4 : memref<?x?xf32>
483  }
484  return %0 : memref<?x?xf32>
485}
486
487//      CHECK: (%[[ARG0:.*]]: {{.*}}
488// CHECK-NEXT: %{{.*}} = scf.if
489// CHECK-NEXT: %{{.*}} = constant
490// CHECK-NEXT: %{{.*}} = addi
491// CHECK-NEXT: %[[FUNC:.*]] = index_cast
492// CHECK-NEXT: alloc(%arg2, %[[FUNC]])
493// CHECK-NEXT: scf.yield
494// CHECK-NEXT: } else {
495// CHECK-NEXT: %{{.*}} = constant
496// CHECK-NEXT: %{{.*}} = addi
497// CHECK-NEXT: %[[FUNC:.*]] = index_cast
498// CHECK-NEXT: alloc(%arg2, %[[FUNC]])
499
500// -----
501
502// Test Case: nested region control flow within a region interface.
503// The alloc positions of %0 does not need to be changed.
504
505// CHECK-LABEL: func @inner_region_control_flow
506func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
507  %0 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
508  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
509    ^bb0(%arg1 : memref<?x?xf32>):
510      test.region_if_yield %arg1 : memref<?x?xf32>
511  } else {
512    ^bb0(%arg1 : memref<?x?xf32>):
513      test.region_if_yield %arg1 : memref<?x?xf32>
514  } join {
515    ^bb0(%arg1 : memref<?x?xf32>):
516      test.region_if_yield %arg1 : memref<?x?xf32>
517  }
518  return %1 : memref<?x?xf32>
519}
520
521//      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
522// CHECK-NEXT: {{.*}} test.region_if
523
524// -----
525
526// Test Case: nested region control flow within a region interface including an
527// allocation in a divergent branch.
528// The alloc positions of %0 does not need to be changed. %2 is moved upwards.
529
530// CHECK-LABEL: func @inner_region_control_flow_div
531func @inner_region_control_flow_div(
532  %arg0 : index,
533  %arg1 : index) -> memref<?x?xf32> {
534  %0 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
535  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
536    ^bb0(%arg2 : memref<?x?xf32>):
537      test.region_if_yield %arg2 : memref<?x?xf32>
538  } else {
539    ^bb0(%arg2 : memref<?x?xf32>):
540      %2 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
541      test.region_if_yield %2 : memref<?x?xf32>
542  } join {
543    ^bb0(%arg2 : memref<?x?xf32>):
544      test.region_if_yield %arg2 : memref<?x?xf32>
545  }
546  return %1 : memref<?x?xf32>
547}
548
549//      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
550// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
551// CHECK-NEXT: {{.*}} test.region_if
552
553// -----
554
555// Test Case: Alloca operations shouldn't be moved.
556
557// CHECK-LABEL: func @condBranchAlloca
558func @condBranchAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
559  cond_br %arg0, ^bb1, ^bb2
560^bb1:
561  br ^bb3(%arg1 : memref<2xf32>)
562^bb2:
563  %0 = memref.alloca() : memref<2xf32>
564  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
565  br ^bb3(%0 : memref<2xf32>)
566^bb3(%1: memref<2xf32>):
567  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
568  return
569}
570
571// CHECK-NEXT: cond_br
572//      CHECK: ^bb2
573//      CHECK: ^bb2
574// CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca()
575// CHECK-NEXT: test.buffer_based
576
577// -----
578
579// Test Case: Alloca operations shouldn't be moved. The alloc operation also
580// shouldn't be moved analogously to the ifElseNested test.
581
582// CHECK-LABEL: func @ifElseNestedAlloca
583func @ifElseNestedAlloca(
584  %arg0: i1,
585  %arg1: memref<2xf32>,
586  %arg2: memref<2xf32>) {
587  %0 = memref.alloca() : memref<2xf32>
588  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
589  cond_br %arg0,
590    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
591    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
592^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
593  br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
594^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
595  cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
596^bb3(%5: memref<2xf32>):
597  br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
598^bb4(%6: memref<2xf32>):
599  br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
600^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
601  %9 = memref.alloc() : memref<2xf32>
602  test.buffer_based in(%7: memref<2xf32>) out(%9: memref<2xf32>)
603  test.copy(%9, %arg2) : (memref<2xf32>, memref<2xf32>)
604  return
605}
606
607// CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca()
608// CHECK-NEXT: test.buffer_based
609//      CHECK: ^bb5
610//      CHECK: ^bb5
611//      CHECK: ^bb5
612// CHECK-NEXT: ^bb5
613// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
614// CHECK-NEXT: test.buffer_based
615
616// -----
617
618// Test Case: Alloca operations shouldn't be moved. The alloc operation should
619// be moved in the beginning analogous to the nestedRegionsAndCondBranch test.
620
621// CHECK-LABEL: func @nestedRegionsAndCondBranchAlloca
622func @nestedRegionsAndCondBranchAlloca(
623  %arg0: i1,
624  %arg1: memref<2xf32>,
625  %arg2: memref<2xf32>) {
626  cond_br %arg0, ^bb1, ^bb2
627^bb1:
628  br ^bb3(%arg1 : memref<2xf32>)
629^bb2:
630  %0 = memref.alloc() : memref<2xf32>
631  test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
632  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
633    %1 = memref.alloca() : memref<2xf32>
634    test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
635    %tmp1 = math.exp %gen1_arg0 : f32
636    test.region_yield %tmp1 : f32
637  }
638  br ^bb3(%0 : memref<2xf32>)
639^bb3(%1: memref<2xf32>):
640  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
641  return
642}
643// CHECK-NEXT:   %[[ALLOC:.*]] = memref.alloc()
644// CHECK-NEXT:   cond_br
645//      CHECK:   test.region_buffer_based
646//      CHECK:     %[[ALLOCA:.*]] = memref.alloca()
647// CHECK-NEXT:     test.buffer_based
648
649// -----
650
651// Test Case: structured control-flow loop using a nested alloc.
652// The alloc positions of %3 will be moved upwards.
653
654// CHECK-LABEL: func @loop_alloc
655func @loop_alloc(
656  %lb: index,
657  %ub: index,
658  %step: index,
659  %buf: memref<2xf32>,
660  %res: memref<2xf32>) {
661  %0 = memref.alloc() : memref<2xf32>
662  %1 = scf.for %i = %lb to %ub step %step
663    iter_args(%iterBuf = %buf) -> memref<2xf32> {
664    %2 = cmpi eq, %i, %ub : index
665    %3 = memref.alloc() : memref<2xf32>
666    scf.yield %3 : memref<2xf32>
667  }
668  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
669  return
670}
671
672//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
673// CHECK-NEXT: {{.*}} scf.for
674//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
675
676// -----
677
678// Test Case: structured control-flow loop with a nested if operation using
679// a deeply nested buffer allocation.
680// The allocation %4 is not moved upwards.
681
682// CHECK-LABEL: func @loop_nested_if_alloc
683func @loop_nested_if_alloc(
684  %lb: index,
685  %ub: index,
686  %step: index,
687  %buf: memref<2xf32>) -> memref<2xf32> {
688  %0 = memref.alloc() : memref<2xf32>
689  %1 = scf.for %i = %lb to %ub step %step
690    iter_args(%iterBuf = %buf) -> memref<2xf32> {
691    %2 = cmpi eq, %i, %ub : index
692    %3 = scf.if %2 -> (memref<2xf32>) {
693      %4 = memref.alloc() : memref<2xf32>
694      scf.yield %4 : memref<2xf32>
695    } else {
696      scf.yield %0 : memref<2xf32>
697    }
698    scf.yield %3 : memref<2xf32>
699  }
700  return %1 : memref<2xf32>
701}
702
703//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
704// CHECK-NEXT: {{.*}} scf.for
705//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
706
707// -----
708
709// Test Case: several nested structured control-flow loops with a deeply nested
710// buffer allocation inside an if operation.
711// Same behavior is an loop_nested_if_alloc: The allocs are not moved upwards.
712
713// CHECK-LABEL: func @loop_nested_alloc
714func @loop_nested_alloc(
715  %lb: index,
716  %ub: index,
717  %step: index,
718  %buf: memref<2xf32>,
719  %res: memref<2xf32>) {
720  %0 = memref.alloc() : memref<2xf32>
721  %1 = scf.for %i = %lb to %ub step %step
722    iter_args(%iterBuf = %buf) -> memref<2xf32> {
723    %2 = scf.for %i2 = %lb to %ub step %step
724      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
725      %3 = scf.for %i3 = %lb to %ub step %step
726        iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
727        %4 = memref.alloc() : memref<2xf32>
728        %5 = cmpi eq, %i, %ub : index
729        %6 = scf.if %5 -> (memref<2xf32>) {
730          %7 = memref.alloc() : memref<2xf32>
731          scf.yield %7 : memref<2xf32>
732        } else {
733          scf.yield %iterBuf3 : memref<2xf32>
734        }
735        scf.yield %6 : memref<2xf32>
736      }
737      scf.yield %3 : memref<2xf32>
738    }
739    scf.yield %2 : memref<2xf32>
740  }
741  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
742  return
743}
744
745//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
746// CHECK-NEXT: {{.*}} = scf.for
747// CHECK-NEXT: {{.*}} = scf.for
748// CHECK-NEXT: {{.*}} = scf.for
749// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc()
750//      CHECK: %[[ALLOC2:.*]] = memref.alloc()
751
752// -----
753
754// CHECK-LABEL: func @loop_nested_alloc_dyn_dependency
755func @loop_nested_alloc_dyn_dependency(
756  %lb: index,
757  %ub: index,
758  %step: index,
759  %arg0: index,
760  %buf: memref<?xf32>,
761  %res: memref<?xf32>) {
762  %0 = memref.alloc(%arg0) : memref<?xf32>
763  %1 = scf.for %i = %lb to %ub step %step
764    iter_args(%iterBuf = %buf) -> memref<?xf32> {
765    %2 = scf.for %i2 = %lb to %ub step %step
766      iter_args(%iterBuf2 = %iterBuf) -> memref<?xf32> {
767      %3 = scf.for %i3 = %lb to %ub step %step
768        iter_args(%iterBuf3 = %iterBuf2) -> memref<?xf32> {
769        %5 = cmpi eq, %i, %ub : index
770        %6 = scf.if %5 -> (memref<?xf32>) {
771          %7 = memref.alloc(%i3) : memref<?xf32>
772          scf.yield %7 : memref<?xf32>
773        } else {
774          scf.yield %iterBuf3 : memref<?xf32>
775        }
776        scf.yield %6 : memref<?xf32>
777      }
778      scf.yield %3 : memref<?xf32>
779    }
780    scf.yield %0 : memref<?xf32>
781  }
782  test.copy(%1, %res) : (memref<?xf32>, memref<?xf32>)
783  return
784}
785
786
787//      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
788// CHECK-NEXT: {{.*}} = scf.for
789// CHECK-NEXT: {{.*}} = scf.for
790// CHECK-NEXT: {{.*}} = scf.for
791//      CHECK: %[[ALLOC1:.*]] = memref.alloc({{.*}})
792