522 lines · plain
1// RUN: mlir-opt -buffer-loop-hoisting -split-input-file %s | FileCheck %s2 3// This file checks the behavior of BufferLoopHoisting pass for moving Alloc4// operations in their correct positions.5 6// Test Case:7// bb08// / \9// bb1 bb2 <- Initial position of AllocOp10// \ /11// bb312// BufferLoopHoisting expected behavior: It should not move the AllocOp.13 14// CHECK-LABEL: func @condBranch15func.func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {16 cf.cond_br %arg0, ^bb1, ^bb217^bb1:18 cf.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 cf.br ^bb3(%0 : memref<2xf32>)23^bb3(%1: memref<2xf32>):24 test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)25 return26}27 28// CHECK-NEXT: cf.cond_br29// CHECK: %[[ALLOC:.*]] = memref.alloc()30 31// -----32 33// Test Case:34// bb035// / \36// bb1 bb2 <- Initial position of AllocOp37// \ /38// bb339// BufferLoopHoisting expected behavior: It should not move the existing AllocOp40// to any other block since the alloc has a dynamic dependency to block argument41// %0 in bb2.42 43// CHECK-LABEL: func @condBranchDynamicType44func.func @condBranchDynamicType(45 %arg0: i1,46 %arg1: memref<?xf32>,47 %arg2: memref<?xf32>,48 %arg3: index) {49 cf.cond_br %arg0, ^bb1, ^bb2(%arg3: index)50^bb1:51 cf.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 cf.br ^bb3(%1 : memref<?xf32>)56^bb3(%2: memref<?xf32>):57 test.copy(%2, %arg2) : (memref<?xf32>, memref<?xf32>)58 return59}60 61// CHECK-NEXT: cf.cond_br62// CHECK: ^bb263// CHECK: ^bb2(%[[IDX:.*]]:{{.*}})64// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])65// CHECK-NEXT: test.buffer_based66 67// -----68 69// Test Case: Nested regions - This test defines a BufferBasedOp inside the70// region of a RegionBufferBasedOp.71// BufferLoopHoisting expected behavior: The AllocOp for the BufferBasedOp72// should remain inside the region of the RegionBufferBasedOp. The AllocOp of73// the RegionBufferBasedOp should not be moved during this pass.74 75// CHECK-LABEL: func @nested_regions_and_cond_branch76func.func @nested_regions_and_cond_branch(77 %arg0: i1,78 %arg1: memref<2xf32>,79 %arg2: memref<2xf32>) {80 cf.cond_br %arg0, ^bb1, ^bb281^bb1:82 cf.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 : f3290 test.region_yield %tmp1 : f3291 }92 cf.br ^bb3(%0 : memref<2xf32>)93^bb3(%1: memref<2xf32>):94 test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)95 return96}97// CHECK-NEXT: cf.cond_br98// CHECK: %[[ALLOC0:.*]] = memref.alloc()99// CHECK: test.region_buffer_based100// CHECK: %[[ALLOC1:.*]] = memref.alloc()101// CHECK-NEXT: test.buffer_based102 103// -----104 105// Test Case: nested region control flow106// The alloc position of %1 does not need to be changed and flows through107// both if branches until it is finally returned.108 109// CHECK-LABEL: func @nested_region_control_flow110func.func @nested_region_control_flow(111 %arg0 : index,112 %arg1 : index) -> memref<?x?xf32> {113 %0 = arith.cmpi eq, %arg0, %arg1 : index114 %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.if126// CHECK: else127// 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_alloc135func.func @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 %step143 iter_args(%iterBuf = %buf) -> memref<2xf32> {144 %2 = arith.cmpi eq, %i, %ub : index145 %3 = memref.alloc() : memref<2xf32>146 scf.yield %3 : memref<2xf32>147 }148 test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)149 return150}151 152// CHECK: %[[ALLOC0:.*]] = memref.alloc()153// CHECK-NEXT: {{.*}} scf.for154// CHECK: %[[ALLOC1:.*]] = memref.alloc()155 156// -----157 158// Test Case: structured control-flow loop with a nested if operation using159// 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_alloc163func.func @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 %step170 iter_args(%iterBuf = %buf) -> memref<2xf32> {171 %2 = arith.cmpi eq, %i, %ub : index172 %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.for185// CHECK: %[[ALLOC1:.*]] = memref.alloc()186 187// -----188 189// Test Case: several nested structured control-flow loops with deeply nested190// buffer allocations inside an if operation.191// Behavior: The allocs %0, %4 and %9 are moved upwards, while %7 and %8 stay192// in their positions.193 194// CHECK-LABEL: func @loop_nested_alloc195func.func @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 %step203 iter_args(%iterBuf = %buf) -> memref<2xf32> {204 %2 = scf.for %i2 = %lb to %ub step %step205 iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {206 %3 = scf.for %i3 = %lb to %ub step %step207 iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {208 %4 = memref.alloc() : memref<2xf32>209 %5 = arith.cmpi eq, %i, %ub : index210 %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 return226}227 228// CHECK: %[[ALLOC0:.*]] = memref.alloc()229// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc()230// CHECK-NEXT: %[[ALLOC2:.*]] = memref.alloc()231// CHECK-NEXT: {{.*}} = scf.for232// CHECK-NEXT: {{.*}} = scf.for233// CHECK-NEXT: {{.*}} = scf.for234// CHECK: {{.*}} = scf.if235// CHECK: %[[ALLOC3:.*]] = memref.alloc()236// CHECK: %[[ALLOC4:.*]] = memref.alloc()237 238// -----239 240// CHECK-LABEL: func @loop_nested_alloc_dyn_dependency241func.func @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 %step250 iter_args(%iterBuf = %buf) -> memref<?xf32> {251 %2 = scf.for %i2 = %lb to %ub step %step252 iter_args(%iterBuf2 = %iterBuf) -> memref<?xf32> {253 %3 = scf.for %i3 = %lb to %ub step %step254 iter_args(%iterBuf3 = %iterBuf2) -> memref<?xf32> {255 %4 = memref.alloc(%i3) : memref<?xf32>256 %5 = arith.cmpi eq, %i, %ub : index257 %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 return272}273 274// CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})275// CHECK-NEXT: {{.*}} = scf.for276// CHECK-NEXT: {{.*}} = scf.for277// CHECK-NEXT: {{.*}} = scf.for278// CHECK: %[[ALLOC1:.*]] = memref.alloc({{.*}})279// CHECK: %[[ALLOC2:.*]] = memref.alloc({{.*}})280 281// -----282 283// CHECK-LABEL: func @hoist_one_loop284func.func @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 %step292 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 return298}299 300// CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})301// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})302// CHECK-NEXT: {{.*}} = scf.for303 304// -----305 306// CHECK-LABEL: func @no_hoist_one_loop307func.func @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 %step314 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 return320}321 322// CHECK: {{.*}} = scf.for323// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})324 325// -----326 327// CHECK-LABEL: func @hoist_multiple_loop328func.func @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 %step336 iter_args(%iterBuf = %buf) -> memref<2xf32> {337 %2 = scf.for %i2 = %lb to %ub step %step338 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 return346}347 348// CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})349// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})350// CHECK-NEXT: {{.*}} = scf.for351 352// -----353 354// CHECK-LABEL: func @no_hoist_one_loop_conditional355func.func @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 %step362 iter_args(%iterBuf = %buf) -> memref<2xf32> {363 %1 = arith.cmpi eq, %i, %ub : index364 %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 return374}375 376// CHECK: {{.*}} = scf.for377// CHECK: {{.*}} = scf.if378// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})379 380// -----381 382// CHECK-LABEL: func @hoist_one_loop_conditional383func.func @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 = arith.cmpi eq, %lb, %ub : index391 %2 = scf.if %1 -> (memref<2xf32>) {392 %3 = scf.for %i = %lb to %ub step %step393 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 else400 {401 scf.yield %0 : memref<2xf32>402 }403 test.copy(%2, %res) : (memref<2xf32>, memref<2xf32>)404 return405}406 407// CHECK: {{.*}} = scf.if408// CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})409// CHECK: {{.*}} = scf.for410 411// -----412 413// CHECK-LABEL: func @no_hoist_one_loop_dependency414func.func @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 %step422 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 return428}429 430// CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})431// CHECK-NEXT: {{.*}} = scf.for432// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})433 434// -----435 436// CHECK-LABEL: func @partial_hoist_multiple_loop_dependency437func.func @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 %step445 iter_args(%iterBuf = %buf) -> memref<2xf32> {446 %2 = scf.for %i2 = %lb to %ub step %step447 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 return455}456 457// CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})458// CHECK-NEXT: {{.*}} = scf.for459// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})460// CHECK-NEXT: {{.*}} = scf.for461 462// -----463 464// CHECK-LABEL: func @no_hoist_parallel465func.func @no_hoist_parallel(466 %lb: index,467 %ub: index,468 %step: index) {469 scf.parallel (%i) = (%lb) to (%ub) step (%step) {470 %0 = memref.alloc() : memref<2xf32>471 scf.reduce472 }473 return474}475 476// CHECK: memref.alloc477// CHECK-NEXT: scf.reduce478 479// -----480 481func.func @no_hoist_forall(482 %lb: index,483 %ub: index,484 %step: index) {485 scf.forall (%i) = (%lb) to (%ub) step (%step) {486 %1 = memref.alloc() : memref<2xf32>487 }488 return489}490 491// CHECK: scf.forall492// CHECK-NEXT: memref.alloc493 494// -----495 496// Test with allocas to ensure that op is also considered.497 498// CHECK-LABEL: func @hoist_alloca499func.func @hoist_alloca(500 %lb: index,501 %ub: index,502 %step: index,503 %buf: memref<2xf32>,504 %res: memref<2xf32>) {505 %0 = memref.alloca() : memref<2xf32>506 %1 = scf.for %i = %lb to %ub step %step507 iter_args(%iterBuf = %buf) -> memref<2xf32> {508 %2 = scf.for %i2 = %lb to %ub step %step509 iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {510 %3 = memref.alloca() : memref<2xf32>511 scf.yield %0 : memref<2xf32>512 }513 scf.yield %0 : memref<2xf32>514 }515 test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)516 return517}518 519// CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})520// CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})521// CHECK-NEXT: {{.*}} = scf.for522