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