1// RUN: mlir-opt -buffer-loop-hoisting -split-input-file %s | FileCheck %s
2
3// This file checks the behavior of BufferLoopHoisting pass for moving Alloc
4// operations in their correct positions.
5
6// Test Case:
7//    bb0
8//   /   \
9//  bb1  bb2 <- Initial position of AllocOp
10//   \   /
11//    bb3
12// BufferLoopHoisting expected behavior: It should not move the AllocOp.
13
14// CHECK-LABEL: func @condBranch
15func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
16  cond_br %arg0, ^bb1, ^bb2
17^bb1:
18  br ^bb3(%arg1 : memref<2xf32>)
19^bb2:
20  %0 = memref.alloc() : memref<2xf32>
21  test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
22  br ^bb3(%0 : memref<2xf32>)
23^bb3(%1: memref<2xf32>):
24  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
25  return
26}
27
28// CHECK-NEXT: cond_br
29//      CHECK: %[[ALLOC:.*]] = memref.alloc()
30
31// -----
32
33// Test Case:
34//    bb0
35//   /   \
36//  bb1  bb2 <- Initial position of AllocOp
37//   \   /
38//    bb3
39// BufferLoopHoisting expected behavior: It should not move the existing AllocOp
40// to any other block since the alloc has a dynamic dependency to block argument
41// %0 in bb2.
42
43// CHECK-LABEL: func @condBranchDynamicType
44func @condBranchDynamicType(
45  %arg0: i1,
46  %arg1: memref<?xf32>,
47  %arg2: memref<?xf32>,
48  %arg3: index) {
49  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
50^bb1:
51  br ^bb3(%arg1 : memref<?xf32>)
52^bb2(%0: index):
53  %1 = memref.alloc(%0) : memref<?xf32>
54  test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
55  br ^bb3(%1 : memref<?xf32>)
56^bb3(%2: memref<?xf32>):
57  test.copy(%2, %arg2) : (memref<?xf32>, memref<?xf32>)
58  return
59}
60
61// CHECK-NEXT: cond_br
62//      CHECK: ^bb2
63//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
64// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])
65// CHECK-NEXT: test.buffer_based
66
67// -----
68
69// Test Case: Nested regions - This test defines a BufferBasedOp inside the
70// region of a RegionBufferBasedOp.
71// BufferLoopHoisting expected behavior: The AllocOp for the BufferBasedOp
72// should remain inside the region of the RegionBufferBasedOp. The AllocOp of
73// the RegionBufferBasedOp should not be moved during this pass.
74
75// CHECK-LABEL: func @nested_regions_and_cond_branch
76func @nested_regions_and_cond_branch(
77  %arg0: i1,
78  %arg1: memref<2xf32>,
79  %arg2: memref<2xf32>) {
80  cond_br %arg0, ^bb1, ^bb2
81^bb1:
82  br ^bb3(%arg1 : memref<2xf32>)
83^bb2:
84  %0 = memref.alloc() : memref<2xf32>
85  test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
86  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
87    %1 = memref.alloc() : memref<2xf32>
88    test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
89    %tmp1 = math.exp %gen1_arg0 : f32
90    test.region_yield %tmp1 : f32
91  }
92  br ^bb3(%0 : memref<2xf32>)
93^bb3(%1: memref<2xf32>):
94  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
95  return
96}
97// CHECK-NEXT:   cond_br
98//      CHECK:   %[[ALLOC0:.*]] = memref.alloc()
99//      CHECK:   test.region_buffer_based
100//      CHECK:     %[[ALLOC1:.*]] = memref.alloc()
101// CHECK-NEXT:     test.buffer_based
102
103// -----
104
105// Test Case: nested region control flow
106// The alloc position of %1 does not need to be changed and flows through
107// both if branches until it is finally returned.
108
109// CHECK-LABEL: func @nested_region_control_flow
110func @nested_region_control_flow(
111  %arg0 : index,
112  %arg1 : index) -> memref<?x?xf32> {
113  %0 = cmpi eq, %arg0, %arg1 : index
114  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
115  %2 = scf.if %0 -> (memref<?x?xf32>) {
116    scf.yield %1 : memref<?x?xf32>
117  } else {
118    %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
119    scf.yield %1 : memref<?x?xf32>
120  }
121  return %2 : memref<?x?xf32>
122}
123
124//      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
125// CHECK-NEXT: %{{.*}} = scf.if
126//      CHECK: else
127// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
128
129// -----
130
131// Test Case: structured control-flow loop using a nested alloc.
132// The alloc positions of %3 should not be changed.
133
134// CHECK-LABEL: func @loop_alloc
135func @loop_alloc(
136  %lb: index,
137  %ub: index,
138  %step: index,
139  %buf: memref<2xf32>,
140  %res: memref<2xf32>) {
141  %0 = memref.alloc() : memref<2xf32>
142  %1 = scf.for %i = %lb to %ub step %step
143    iter_args(%iterBuf = %buf) -> memref<2xf32> {
144    %2 = cmpi eq, %i, %ub : index
145    %3 = memref.alloc() : memref<2xf32>
146    scf.yield %3 : memref<2xf32>
147  }
148  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
149  return
150}
151
152//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
153// CHECK-NEXT: {{.*}} scf.for
154//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
155
156// -----
157
158// Test Case: structured control-flow loop with a nested if operation using
159// a deeply nested buffer allocation.
160// The allocation %4 should not be moved upwards due to a back-edge dependency.
161
162// CHECK-LABEL: func @loop_nested_if_alloc
163func @loop_nested_if_alloc(
164  %lb: index,
165  %ub: index,
166  %step: index,
167  %buf: memref<2xf32>) -> memref<2xf32> {
168  %0 = memref.alloc() : memref<2xf32>
169  %1 = scf.for %i = %lb to %ub step %step
170    iter_args(%iterBuf = %buf) -> memref<2xf32> {
171    %2 = cmpi eq, %i, %ub : index
172    %3 = scf.if %2 -> (memref<2xf32>) {
173      %4 = memref.alloc() : memref<2xf32>
174      scf.yield %4 : memref<2xf32>
175    } else {
176      scf.yield %0 : memref<2xf32>
177    }
178    scf.yield %3 : memref<2xf32>
179  }
180  return %1 : memref<2xf32>
181}
182
183//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
184// CHECK-NEXT: {{.*}} scf.for
185//      CHECK: %[[ALLOC1:.*]] = memref.alloc()
186
187// -----
188
189// Test Case: several nested structured control-flow loops with deeply nested
190// buffer allocations inside an if operation.
191// Behavior: The allocs %0, %4 and %9 are moved upwards, while %7 and %8 stay
192// in their positions.
193
194// CHECK-LABEL: func @loop_nested_alloc
195func @loop_nested_alloc(
196  %lb: index,
197  %ub: index,
198  %step: index,
199  %buf: memref<2xf32>,
200  %res: memref<2xf32>) {
201  %0 = memref.alloc() : memref<2xf32>
202  %1 = scf.for %i = %lb to %ub step %step
203    iter_args(%iterBuf = %buf) -> memref<2xf32> {
204    %2 = scf.for %i2 = %lb to %ub step %step
205      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
206      %3 = scf.for %i3 = %lb to %ub step %step
207        iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
208        %4 = memref.alloc() : memref<2xf32>
209        %5 = cmpi eq, %i, %ub : index
210        %6 = scf.if %5 -> (memref<2xf32>) {
211          %7 = memref.alloc() : memref<2xf32>
212          %8 = memref.alloc() : memref<2xf32>
213          scf.yield %8 : memref<2xf32>
214        } else {
215          scf.yield %iterBuf3 : memref<2xf32>
216        }
217        %9 = memref.alloc() : memref<2xf32>
218        scf.yield %6 : memref<2xf32>
219      }
220      scf.yield %3 : memref<2xf32>
221    }
222    scf.yield %2 : memref<2xf32>
223  }
224  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
225  return
226}
227
228//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
229// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc()
230// CHECK-NEXT: %[[ALLOC2:.*]] = memref.alloc()
231// CHECK-NEXT: {{.*}} = scf.for
232// CHECK-NEXT: {{.*}} = scf.for
233// CHECK-NEXT: {{.*}} = scf.for
234//      CHECK: {{.*}} = scf.if
235//      CHECK: %[[ALLOC3:.*]] = memref.alloc()
236//      CHECK: %[[ALLOC4:.*]] = memref.alloc()
237
238// -----
239
240// CHECK-LABEL: func @loop_nested_alloc_dyn_dependency
241func @loop_nested_alloc_dyn_dependency(
242  %lb: index,
243  %ub: index,
244  %step: index,
245  %arg0: index,
246  %buf: memref<?xf32>,
247  %res: memref<?xf32>) {
248  %0 = memref.alloc(%arg0) : memref<?xf32>
249  %1 = scf.for %i = %lb to %ub step %step
250    iter_args(%iterBuf = %buf) -> memref<?xf32> {
251    %2 = scf.for %i2 = %lb to %ub step %step
252      iter_args(%iterBuf2 = %iterBuf) -> memref<?xf32> {
253      %3 = scf.for %i3 = %lb to %ub step %step
254        iter_args(%iterBuf3 = %iterBuf2) -> memref<?xf32> {
255        %4 = memref.alloc(%i3) : memref<?xf32>
256        %5 = cmpi eq, %i, %ub : index
257        %6 = scf.if %5 -> (memref<?xf32>) {
258          %7 = memref.alloc(%i3) : memref<?xf32>
259          scf.yield %7 : memref<?xf32>
260        } else {
261          scf.yield %iterBuf3 : memref<?xf32>
262        }
263        %8 = memref.alloc(%i3) : memref<?xf32>
264        scf.yield %6 : memref<?xf32>
265      }
266      scf.yield %3 : memref<?xf32>
267    }
268    scf.yield %0 : memref<?xf32>
269  }
270  test.copy(%1, %res) : (memref<?xf32>, memref<?xf32>)
271  return
272}
273
274//      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
275// CHECK-NEXT: {{.*}} = scf.for
276// CHECK-NEXT: {{.*}} = scf.for
277// CHECK-NEXT: {{.*}} = scf.for
278//      CHECK: %[[ALLOC1:.*]] = memref.alloc({{.*}})
279//      CHECK: %[[ALLOC2:.*]] = memref.alloc({{.*}})
280
281// -----
282
283// CHECK-LABEL: func @hoist_one_loop
284func @hoist_one_loop(
285  %lb: index,
286  %ub: index,
287  %step: index,
288  %buf: memref<2xf32>,
289  %res: memref<2xf32>) {
290  %0 = memref.alloc() : memref<2xf32>
291  %1 = scf.for %i = %lb to %ub step %step
292    iter_args(%iterBuf = %buf) -> memref<2xf32> {
293      %2 = memref.alloc() : memref<2xf32>
294      scf.yield %0 : memref<2xf32>
295  }
296  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
297  return
298}
299
300//      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
301// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
302// CHECK-NEXT: {{.*}} = scf.for
303
304// -----
305
306// CHECK-LABEL: func @no_hoist_one_loop
307func @no_hoist_one_loop(
308  %lb: index,
309  %ub: index,
310  %step: index,
311  %buf: memref<2xf32>,
312  %res: memref<2xf32>) {
313  %0 = scf.for %i = %lb to %ub step %step
314    iter_args(%iterBuf = %buf) -> memref<2xf32> {
315      %1 = memref.alloc() : memref<2xf32>
316      scf.yield %1 : memref<2xf32>
317  }
318  test.copy(%0, %res) : (memref<2xf32>, memref<2xf32>)
319  return
320}
321
322//      CHECK: {{.*}} = scf.for
323// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})
324
325// -----
326
327// CHECK-LABEL: func @hoist_multiple_loop
328func @hoist_multiple_loop(
329  %lb: index,
330  %ub: index,
331  %step: index,
332  %buf: memref<2xf32>,
333  %res: memref<2xf32>) {
334  %0 = memref.alloc() : memref<2xf32>
335  %1 = scf.for %i = %lb to %ub step %step
336    iter_args(%iterBuf = %buf) -> memref<2xf32> {
337    %2 = scf.for %i2 = %lb to %ub step %step
338      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
339        %3 = memref.alloc() : memref<2xf32>
340        scf.yield %0 : memref<2xf32>
341    }
342    scf.yield %0 : memref<2xf32>
343  }
344  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
345  return
346}
347
348//      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
349// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
350// CHECK-NEXT: {{.*}} = scf.for
351
352// -----
353
354// CHECK-LABEL: func @no_hoist_one_loop_conditional
355func @no_hoist_one_loop_conditional(
356  %lb: index,
357  %ub: index,
358  %step: index,
359  %buf: memref<2xf32>,
360  %res: memref<2xf32>) {
361  %0 = scf.for %i = %lb to %ub step %step
362    iter_args(%iterBuf = %buf) -> memref<2xf32> {
363      %1 = cmpi eq, %i, %ub : index
364      %2 = scf.if %1 -> (memref<2xf32>) {
365        %3 = memref.alloc() : memref<2xf32>
366        scf.yield %3 : memref<2xf32>
367      } else {
368        scf.yield %iterBuf : memref<2xf32>
369      }
370    scf.yield %2 : memref<2xf32>
371  }
372  test.copy(%0, %res) : (memref<2xf32>, memref<2xf32>)
373  return
374}
375
376//      CHECK: {{.*}} = scf.for
377//      CHECK: {{.*}} = scf.if
378// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})
379
380// -----
381
382// CHECK-LABEL: func @hoist_one_loop_conditional
383func @hoist_one_loop_conditional(
384  %lb: index,
385  %ub: index,
386  %step: index,
387  %buf: memref<2xf32>,
388  %res: memref<2xf32>) {
389  %0 = memref.alloc() : memref<2xf32>
390  %1 = cmpi eq, %lb, %ub : index
391  %2 = scf.if %1 -> (memref<2xf32>) {
392    %3 = scf.for %i = %lb to %ub step %step
393    iter_args(%iterBuf = %buf) -> memref<2xf32> {
394      %4 = memref.alloc() : memref<2xf32>
395      scf.yield %0 : memref<2xf32>
396    }
397    scf.yield %0 : memref<2xf32>
398  }
399  else
400  {
401    scf.yield %0 : memref<2xf32>
402  }
403  test.copy(%2, %res) : (memref<2xf32>, memref<2xf32>)
404  return
405}
406
407//      CHECK: {{.*}} = scf.if
408// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})
409//      CHECK: {{.*}} = scf.for
410
411// -----
412
413// CHECK-LABEL: func @no_hoist_one_loop_dependency
414func @no_hoist_one_loop_dependency(
415  %lb: index,
416  %ub: index,
417  %step: index,
418  %buf: memref<2xf32>,
419  %res: memref<2xf32>) {
420  %0 = memref.alloc() : memref<2xf32>
421  %1 = scf.for %i = %lb to %ub step %step
422    iter_args(%iterBuf = %buf) -> memref<2xf32> {
423      %2 = memref.alloc(%i) : memref<?xf32>
424      scf.yield %0 : memref<2xf32>
425  }
426  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
427  return
428}
429
430//      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
431// CHECK-NEXT: {{.*}} = scf.for
432// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
433
434// -----
435
436// CHECK-LABEL: func @partial_hoist_multiple_loop_dependency
437func @partial_hoist_multiple_loop_dependency(
438  %lb: index,
439  %ub: index,
440  %step: index,
441  %buf: memref<2xf32>,
442  %res: memref<2xf32>) {
443  %0 = memref.alloc() : memref<2xf32>
444  %1 = scf.for %i = %lb to %ub step %step
445    iter_args(%iterBuf = %buf) -> memref<2xf32> {
446    %2 = scf.for %i2 = %lb to %ub step %step
447      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
448        %3 = memref.alloc(%i) : memref<?xf32>
449        scf.yield %0 : memref<2xf32>
450    }
451    scf.yield %0 : memref<2xf32>
452  }
453  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
454  return
455}
456
457//      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
458// CHECK-NEXT: {{.*}} = scf.for
459// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
460// CHECK-NEXT: {{.*}} = scf.for
461
462// -----
463
464// Test with allocas to ensure that op is also considered.
465
466// CHECK-LABEL: func @hoist_alloca
467func @hoist_alloca(
468  %lb: index,
469  %ub: index,
470  %step: index,
471  %buf: memref<2xf32>,
472  %res: memref<2xf32>) {
473  %0 = memref.alloca() : memref<2xf32>
474  %1 = scf.for %i = %lb to %ub step %step
475    iter_args(%iterBuf = %buf) -> memref<2xf32> {
476    %2 = scf.for %i2 = %lb to %ub step %step
477      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
478        %3 = memref.alloca() : memref<2xf32>
479        scf.yield %0 : memref<2xf32>
480    }
481    scf.yield %0 : memref<2xf32>
482  }
483  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
484  return
485}
486
487//      CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
488// CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
489// CHECK-NEXT: {{.*}} = scf.for
490