mlir/test/Transforms/buffer-deallocation.mlir

   1 // RUN: mlir-opt -buffer-deallocation -split-input-file %s | FileCheck %s
   2
   3 // This file checks the behaviour of BufferDeallocation pass for moving and
   4 // inserting missing DeallocOps in their correct positions. Furthermore,
   5 // copies and their corresponding AllocOps are inserted.
   6
   7 // Test Case:
   8 //    bb0
   9 //   /   \
  10 //  bb1  bb2 <- Initial position of AllocOp
  11 //   \   /
  12 //    bb3
  13 // BufferDeallocation expected behavior: bb2 contains an AllocOp which is
  14 // passed to bb3. In the latter block, there should be an deallocation.
  15 // Since bb1 does not contain an adequate alloc and the alloc in bb2 is not
  16 // moved to bb0, we need to insert allocs and copies.
  17
  18 // CHECK-LABEL: func @condBranch
  19 func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  20   cond_br %arg0, ^bb1, ^bb2
  21 ^bb1:
  22   br ^bb3(%arg1 : memref<2xf32>)
  23 ^bb2:
  24   %0 = alloc() : memref<2xf32>
  25   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
  26   br ^bb3(%0 : memref<2xf32>)
  27 ^bb3(%1: memref<2xf32>):
  28   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
  29   return
  30 }
  31
  32 // CHECK-NEXT: cond_br
  33 //      CHECK: %[[ALLOC0:.*]] = alloc()
  34 // CHECK-NEXT: linalg.copy
  35 // CHECK-NEXT: br ^bb3(%[[ALLOC0]]
  36 //      CHECK: %[[ALLOC1:.*]] = alloc()
  37 // CHECK-NEXT: test.buffer_based
  38 //      CHECK: %[[ALLOC2:.*]] = alloc()
  39 // CHECK-NEXT: linalg.copy
  40 // CHECK-NEXT: dealloc %[[ALLOC1]]
  41 // CHECK-NEXT: br ^bb3(%[[ALLOC2]]
  42 //      CHECK: test.copy
  43 // CHECK-NEXT: dealloc
  44 // CHECK-NEXT: return
  45
  46 // -----
  47
  48 // Test Case:
  49 //    bb0
  50 //   /   \
  51 //  bb1  bb2 <- Initial position of AllocOp
  52 //   \   /
  53 //    bb3
  54 // BufferDeallocation expected behavior: The existing AllocOp has a dynamic
  55 // dependency to block argument %0 in bb2. Since the dynamic type is passed
  56 // to bb3 via the block argument %2, it is currently required to allocate a
  57 // temporary buffer for %2 that gets copies of %arg0 and %1 with their
  58 // appropriate shape dimensions. The copy buffer deallocation will be applied
  59 // to %2 in block bb3.
  60
  61 // CHECK-LABEL: func @condBranchDynamicType
  62 func @condBranchDynamicType(
  63   %arg0: i1,
  64   %arg1: memref<?xf32>,
  65   %arg2: memref<?xf32>,
  66   %arg3: index) {
  67   cond_br %arg0, ^bb1, ^bb2(%arg3: index)
  68 ^bb1:
  69   br ^bb3(%arg1 : memref<?xf32>)
  70 ^bb2(%0: index):
  71   %1 = alloc(%0) : memref<?xf32>
  72   test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
  73   br ^bb3(%1 : memref<?xf32>)
  74 ^bb3(%2: memref<?xf32>):
  75   test.copy(%2, %arg2) : (memref<?xf32>, memref<?xf32>)
  76   return
  77 }
  78
  79 // CHECK-NEXT: cond_br
  80 //      CHECK: %[[DIM0:.*]] = dim
  81 // CHECK-NEXT: %[[ALLOC0:.*]] = alloc(%[[DIM0]])
  82 // CHECK-NEXT: linalg.copy(%{{.*}}, %[[ALLOC0]])
  83 // CHECK-NEXT: br ^bb3(%[[ALLOC0]]
  84 //      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
  85 // CHECK-NEXT: %[[ALLOC1:.*]] = alloc(%[[IDX]])
  86 // CHECK-NEXT: test.buffer_based
  87 //      CHECK: %[[DIM1:.*]] = dim %[[ALLOC1]]
  88 // CHECK-NEXT: %[[ALLOC2:.*]] = alloc(%[[DIM1]])
  89 // CHECK-NEXT: linalg.copy(%[[ALLOC1]], %[[ALLOC2]])
  90 // CHECK-NEXT: dealloc %[[ALLOC1]]
  91 // CHECK-NEXT: br ^bb3
  92 // CHECK-NEXT: ^bb3(%[[ALLOC3:.*]]:{{.*}})
  93 //      CHECK: test.copy(%[[ALLOC3]],
  94 // CHECK-NEXT: dealloc %[[ALLOC3]]
  95 // CHECK-NEXT: return
  96
  97 // -----
  98
  99 // Test Case:
 100 //      bb0
 101 //     /    \
 102 //   bb1    bb2 <- Initial position of AllocOp
 103 //    |     /  \
 104 //    |   bb3  bb4
 105 //    |     \  /
 106 //    \     bb5
 107 //     \    /
 108 //       bb6
 109 //        |
 110 //       bb7
 111 // BufferDeallocation expected behavior: The existing AllocOp has a dynamic
 112 // dependency to block argument %0 in bb2. Since the dynamic type is passed to
 113 // bb5 via the block argument %2 and to bb6 via block argument %3, it is
 114 // currently required to allocate temporary buffers for %2 and %3 that gets
 115 // copies of %1 and %arg0 1 with their appropriate shape dimensions. The copy
 116 // buffer deallocations will be applied to %2 in block bb5 and to %3 in block
 117 // bb6. Furthermore, there should be no copy inserted for %4.
 118
 119 // CHECK-LABEL: func @condBranchDynamicTypeNested
 120 func @condBranchDynamicTypeNested(
 121   %arg0: i1,
 122   %arg1: memref<?xf32>,
 123   %arg2: memref<?xf32>,
 124   %arg3: index) {
 125   cond_br %arg0, ^bb1, ^bb2(%arg3: index)
 126 ^bb1:
 127   br ^bb6(%arg1 : memref<?xf32>)
 128 ^bb2(%0: index):
 129   %1 = alloc(%0) : memref<?xf32>
 130   test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
 131   cond_br %arg0, ^bb3, ^bb4
 132 ^bb3:
 133   br ^bb5(%1 : memref<?xf32>)
 134 ^bb4:
 135   br ^bb5(%1 : memref<?xf32>)
 136 ^bb5(%2: memref<?xf32>):
 137   br ^bb6(%2 : memref<?xf32>)
 138 ^bb6(%3: memref<?xf32>):
 139   br ^bb7(%3 : memref<?xf32>)
 140 ^bb7(%4: memref<?xf32>):
 141   test.copy(%4, %arg2) : (memref<?xf32>, memref<?xf32>)
 142   return
 143 }
 144
 145 // CHECK-NEXT: cond_br
 146 //      CHECK: ^bb1
 147 //      CHECK: %[[DIM0:.*]] = dim
 148 // CHECK-NEXT: %[[ALLOC0:.*]] = alloc(%[[DIM0]])
 149 // CHECK-NEXT: linalg.copy(%{{.*}}, %[[ALLOC0]])
 150 // CHECK-NEXT: br ^bb6
 151 //      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
 152 // CHECK-NEXT: %[[ALLOC1:.*]] = alloc(%[[IDX]])
 153 // CHECK-NEXT: test.buffer_based
 154 //      CHECK: cond_br
 155 //      CHECK: ^bb3:
 156 // CHECK-NEXT: br ^bb5(%[[ALLOC1]]{{.*}})
 157 //      CHECK: ^bb4:
 158 // CHECK-NEXT: br ^bb5(%[[ALLOC1]]{{.*}})
 159 // CHECK-NEXT: ^bb5(%[[ALLOC2:.*]]:{{.*}})
 160 //      CHECK: %[[DIM2:.*]] = dim %[[ALLOC2]]
 161 // CHECK-NEXT: %[[ALLOC3:.*]] = alloc(%[[DIM2]])
 162 // CHECK-NEXT: linalg.copy(%[[ALLOC2]], %[[ALLOC3]])
 163 // CHECK-NEXT: dealloc %[[ALLOC1]]
 164 // CHECK-NEXT: br ^bb6(%[[ALLOC3]]{{.*}})
 165 // CHECK-NEXT: ^bb6(%[[ALLOC4:.*]]:{{.*}})
 166 // CHECK-NEXT: br ^bb7(%[[ALLOC4]]{{.*}})
 167 // CHECK-NEXT: ^bb7(%[[ALLOC5:.*]]:{{.*}})
 168 //      CHECK: test.copy(%[[ALLOC5]],
 169 // CHECK-NEXT: dealloc %[[ALLOC4]]
 170 // CHECK-NEXT: return
 171
 172 // -----
 173
 174 // Test Case: Existing AllocOp with no users.
 175 // BufferDeallocation expected behavior: It should insert a DeallocOp right
 176 // before ReturnOp.
 177
 178 // CHECK-LABEL: func @emptyUsesValue
 179 func @emptyUsesValue(%arg0: memref<4xf32>) {
 180   %0 = alloc() : memref<4xf32>
 181   return
 182 }
 183 // CHECK-NEXT: %[[ALLOC:.*]] = alloc()
 184 // CHECK-NEXT: dealloc %[[ALLOC]]
 185 // CHECK-NEXT: return
 186
 187 // -----
 188
 189 // Test Case:
 190 //    bb0
 191 //   /   \
 192 //  |    bb1 <- Initial position of AllocOp
 193 //   \   /
 194 //    bb2
 195 // BufferDeallocation expected behavior: It should insert a DeallocOp at the
 196 // exit block after CopyOp since %1 is an alias for %0 and %arg1. Furthermore,
 197 // we have to insert a copy and an alloc in the beginning of the function.
 198
 199 // CHECK-LABEL: func @criticalEdge
 200 func @criticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 201   cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
 202 ^bb1:
 203   %0 = alloc() : memref<2xf32>
 204   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 205   br ^bb2(%0 : memref<2xf32>)
 206 ^bb2(%1: memref<2xf32>):
 207   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
 208   return
 209 }
 210
 211 // CHECK-NEXT: %[[ALLOC0:.*]] = alloc()
 212 // CHECK-NEXT: linalg.copy
 213 // CHECK-NEXT: cond_br
 214 //      CHECK: %[[ALLOC1:.*]] = alloc()
 215 // CHECK-NEXT: test.buffer_based
 216 //      CHECK: %[[ALLOC2:.*]] = alloc()
 217 // CHECK-NEXT: linalg.copy
 218 // CHECK-NEXT: dealloc %[[ALLOC1]]
 219 //      CHECK: test.copy
 220 // CHECK-NEXT: dealloc
 221 // CHECK-NEXT: return
 222
 223 // -----
 224
 225 // Test Case:
 226 //    bb0 <- Initial position of AllocOp
 227 //   /   \
 228 //  |    bb1
 229 //   \   /
 230 //    bb2
 231 // BufferDeallocation expected behavior: It only inserts a DeallocOp at the
 232 // exit block after CopyOp since %1 is an alias for %0 and %arg1.
 233
 234 // CHECK-LABEL: func @invCriticalEdge
 235 func @invCriticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 236   %0 = alloc() : memref<2xf32>
 237   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 238   cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
 239 ^bb1:
 240   br ^bb2(%0 : memref<2xf32>)
 241 ^bb2(%1: memref<2xf32>):
 242   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
 243   return
 244 }
 245
 246 //      CHECK: dealloc
 247 // CHECK-NEXT: return
 248
 249 // -----
 250
 251 // Test Case:
 252 //    bb0 <- Initial position of the first AllocOp
 253 //   /   \
 254 //  bb1  bb2
 255 //   \   /
 256 //    bb3 <- Initial position of the second AllocOp
 257 // BufferDeallocation expected behavior: It only inserts two missing
 258 // DeallocOps in the exit block. %5 is an alias for %0. Therefore, the
 259 // DeallocOp for %0 should occur after the last BufferBasedOp. The Dealloc for
 260 // %7 should happen after CopyOp.
 261
 262 // CHECK-LABEL: func @ifElse
 263 func @ifElse(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 264   %0 = alloc() : memref<2xf32>
 265   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 266   cond_br %arg0,
 267     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
 268     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 269 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
 270   br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
 271 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
 272   br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
 273 ^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
 274   %7 = alloc() : memref<2xf32>
 275   test.buffer_based in(%5: memref<2xf32>) out(%7: memref<2xf32>)
 276   test.copy(%7, %arg2) : (memref<2xf32>, memref<2xf32>)
 277   return
 278 }
 279
 280 // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
 281 // CHECK-NEXT: test.buffer_based
 282 //      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
 283 // CHECK-NEXT: test.buffer_based
 284 //      CHECK: dealloc %[[FIRST_ALLOC]]
 285 //      CHECK: test.copy
 286 // CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
 287 // CHECK-NEXT: return
 288
 289 // -----
 290
 291 // Test Case: No users for buffer in if-else CFG
 292 //    bb0 <- Initial position of AllocOp
 293 //   /   \
 294 //  bb1  bb2
 295 //   \   /
 296 //    bb3
 297 // BufferDeallocation expected behavior: It only inserts a missing DeallocOp
 298 // in the exit block since %5 or %6 are the latest aliases of %0.
 299
 300 // CHECK-LABEL: func @ifElseNoUsers
 301 func @ifElseNoUsers(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 302   %0 = alloc() : memref<2xf32>
 303   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 304   cond_br %arg0,
 305     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
 306     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 307 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
 308   br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
 309 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
 310   br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
 311 ^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
 312   test.copy(%arg1, %arg2) : (memref<2xf32>, memref<2xf32>)
 313   return
 314 }
 315
 316 // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
 317 //      CHECK: test.copy
 318 // CHECK-NEXT: dealloc %[[FIRST_ALLOC]]
 319 // CHECK-NEXT: return
 320
 321 // -----
 322
 323 // Test Case:
 324 //      bb0 <- Initial position of the first AllocOp
 325 //     /    \
 326 //   bb1    bb2
 327 //    |     /  \
 328 //    |   bb3  bb4
 329 //    \     \  /
 330 //     \     /
 331 //       bb5 <- Initial position of the second AllocOp
 332 // BufferDeallocation expected behavior: Two missing DeallocOps should be
 333 // inserted in the exit block.
 334
 335 // CHECK-LABEL: func @ifElseNested
 336 func @ifElseNested(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 337   %0 = alloc() : memref<2xf32>
 338   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 339   cond_br %arg0,
 340     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
 341     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 342 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
 343   br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
 344 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
 345   cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
 346 ^bb3(%5: memref<2xf32>):
 347   br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
 348 ^bb4(%6: memref<2xf32>):
 349   br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
 350 ^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
 351   %9 = alloc() : memref<2xf32>
 352   test.buffer_based in(%7: memref<2xf32>) out(%9: memref<2xf32>)
 353   test.copy(%9, %arg2) : (memref<2xf32>, memref<2xf32>)
 354   return
 355 }
 356
 357 // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
 358 // CHECK-NEXT: test.buffer_based
 359 //      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
 360 // CHECK-NEXT: test.buffer_based
 361 //      CHECK: dealloc %[[FIRST_ALLOC]]
 362 //      CHECK: test.copy
 363 // CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
 364 // CHECK-NEXT: return
 365
 366 // -----
 367
 368 // Test Case: Dead operations in a single block.
 369 // BufferDeallocation expected behavior: It only inserts the two missing
 370 // DeallocOps after the last BufferBasedOp.
 371
 372 // CHECK-LABEL: func @redundantOperations
 373 func @redundantOperations(%arg0: memref<2xf32>) {
 374   %0 = alloc() : memref<2xf32>
 375   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
 376   %1 = alloc() : memref<2xf32>
 377   test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
 378   return
 379 }
 380
 381 //      CHECK: (%[[ARG0:.*]]: {{.*}})
 382 // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
 383 // CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}}out(%[[FIRST_ALLOC]]
 384 //      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
 385 // CHECK-NEXT: test.buffer_based in(%[[FIRST_ALLOC]]{{.*}}out(%[[SECOND_ALLOC]]
 386 //      CHECK: dealloc
 387 // CHECK-NEXT: dealloc
 388 // CHECK-NEXT: return
 389
 390 // -----
 391
 392 // Test Case:
 393 //                                     bb0
 394 //                                    /   \
 395 // Initial pos of the 1st AllocOp -> bb1  bb2 <- Initial pos of the 2nd AllocOp
 396 //                                    \   /
 397 //                                     bb3
 398 // BufferDeallocation expected behavior: We need to introduce a copy for each
 399 // buffer since the buffers are passed to bb3. The both missing DeallocOps are
 400 // inserted in the respective block of the allocs. The copy is freed in the exit
 401 // block.
 402
 403 // CHECK-LABEL: func @moving_alloc_and_inserting_missing_dealloc
 404 func @moving_alloc_and_inserting_missing_dealloc(
 405   %cond: i1,
 406     %arg0: memref<2xf32>,
 407     %arg1: memref<2xf32>) {
 408   cond_br %cond, ^bb1, ^bb2
 409 ^bb1:
 410   %0 = alloc() : memref<2xf32>
 411   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
 412   br ^exit(%0 : memref<2xf32>)
 413 ^bb2:
 414   %1 = alloc() : memref<2xf32>
 415   test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
 416   br ^exit(%1 : memref<2xf32>)
 417 ^exit(%arg2: memref<2xf32>):
 418   test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
 419   return
 420 }
 421
 422 // CHECK-NEXT: cond_br
 423 //      CHECK: ^bb1
 424 //      CHECK: ^bb1
 425 //      CHECK: %[[ALLOC0:.*]] = alloc()
 426 // CHECK-NEXT: test.buffer_based
 427 //      CHECK: %[[ALLOC1:.*]] = alloc()
 428 // CHECK-NEXT: linalg.copy
 429 // CHECK-NEXT: dealloc %[[ALLOC0]]
 430 // CHECK-NEXT: br ^bb3(%[[ALLOC1]]
 431 // CHECK-NEXT: ^bb2
 432 // CHECK-NEXT: %[[ALLOC2:.*]] = alloc()
 433 // CHECK-NEXT: test.buffer_based
 434 //      CHECK: %[[ALLOC3:.*]] = alloc()
 435 // CHECK-NEXT: linalg.copy
 436 // CHECK-NEXT: dealloc %[[ALLOC2]]
 437 // CHECK-NEXT: br ^bb3(%[[ALLOC3]]
 438 // CHECK-NEXT: ^bb3(%[[ALLOC4:.*]]:{{.*}})
 439 //      CHECK: test.copy
 440 // CHECK-NEXT: dealloc %[[ALLOC4]]
 441 // CHECK-NEXT: return
 442
 443 // -----
 444
 445 // Test Case: Invalid position of the DeallocOp. There is a user after
 446 // deallocation.
 447 //   bb0
 448 //  /   \
 449 // bb1  bb2 <- Initial position of AllocOp
 450 //  \   /
 451 //   bb3
 452 // BufferDeallocation expected behavior: The existing DeallocOp should be
 453 // moved to exit block.
 454
 455 // CHECK-LABEL: func @moving_invalid_dealloc_op_complex
 456 func @moving_invalid_dealloc_op_complex(
 457   %cond: i1,
 458     %arg0: memref<2xf32>,
 459     %arg1: memref<2xf32>) {
 460   %1 = alloc() : memref<2xf32>
 461   cond_br %cond, ^bb1, ^bb2
 462 ^bb1:
 463   br ^exit(%arg0 : memref<2xf32>)
 464 ^bb2:
 465   test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
 466   dealloc %1 : memref<2xf32>
 467   br ^exit(%1 : memref<2xf32>)
 468 ^exit(%arg2: memref<2xf32>):
 469   test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
 470   return
 471 }
 472
 473 // CHECK-NEXT: %[[ALLOC0:.*]] = alloc()
 474 // CHECK-NEXT: cond_br
 475 //      CHECK: test.copy
 476 // CHECK-NEXT: dealloc %[[ALLOC0]]
 477 // CHECK-NEXT: return
 478
 479 // -----
 480
 481 // Test Case: Inserting missing DeallocOp in a single block.
 482
 483 // CHECK-LABEL: func @inserting_missing_dealloc_simple
 484 func @inserting_missing_dealloc_simple(
 485   %arg0 : memref<2xf32>,
 486   %arg1: memref<2xf32>) {
 487   %0 = alloc() : memref<2xf32>
 488   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
 489   test.copy(%0, %arg1) : (memref<2xf32>, memref<2xf32>)
 490   return
 491 }
 492
 493 // CHECK-NEXT: %[[ALLOC0:.*]] = alloc()
 494 //      CHECK: test.copy
 495 // CHECK-NEXT: dealloc %[[ALLOC0]]
 496
 497 // -----
 498
 499 // Test Case: Moving invalid DeallocOp (there is a user after deallocation) in a
 500 // single block.
 501
 502 // CHECK-LABEL: func @moving_invalid_dealloc_op
 503 func @moving_invalid_dealloc_op(%arg0 : memref<2xf32>, %arg1: memref<2xf32>) {
 504   %0 = alloc() : memref<2xf32>
 505   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
 506   dealloc %0 : memref<2xf32>
 507   test.copy(%0, %arg1) : (memref<2xf32>, memref<2xf32>)
 508   return
 509 }
 510
 511 // CHECK-NEXT: %[[ALLOC0:.*]] = alloc()
 512 //      CHECK: test.copy
 513 // CHECK-NEXT: dealloc %[[ALLOC0]]
 514
 515 // -----
 516
 517 // Test Case: Nested regions - This test defines a BufferBasedOp inside the
 518 // region of a RegionBufferBasedOp.
 519 // BufferDeallocation expected behavior: The AllocOp for the BufferBasedOp
 520 // should remain inside the region of the RegionBufferBasedOp and it should insert
 521 // the missing DeallocOp in the same region. The missing DeallocOp should be
 522 // inserted after CopyOp.
 523
 524 // CHECK-LABEL: func @nested_regions_and_cond_branch
 525 func @nested_regions_and_cond_branch(
 526   %arg0: i1,
 527   %arg1: memref<2xf32>,
 528   %arg2: memref<2xf32>) {
 529   cond_br %arg0, ^bb1, ^bb2
 530 ^bb1:
 531   br ^bb3(%arg1 : memref<2xf32>)
 532 ^bb2:
 533   %0 = alloc() : memref<2xf32>
 534   test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
 535   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
 536     %1 = alloc() : memref<2xf32>
 537     test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
 538     %tmp1 = exp %gen1_arg0 : f32
 539     test.region_yield %tmp1 : f32
 540   }
 541   br ^bb3(%0 : memref<2xf32>)
 542 ^bb3(%1: memref<2xf32>):
 543   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
 544   return
 545 }
 546 //      CHECK: (%[[cond:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %{{.*}}: {{.*}})
 547 // CHECK-NEXT:   cond_br %[[cond]], ^[[BB1:.*]], ^[[BB2:.*]]
 548 //      CHECK:   %[[ALLOC0:.*]] = alloc()
 549 // CHECK-NEXT:   linalg.copy(%[[ARG1]], %[[ALLOC0]])
 550 //      CHECK: ^[[BB2]]:
 551 //      CHECK:   %[[ALLOC1:.*]] = alloc()
 552 // CHECK-NEXT:   test.region_buffer_based in(%[[ARG1]]{{.*}}out(%[[ALLOC1]]
 553 //      CHECK:     %[[ALLOC2:.*]] = alloc()
 554 // CHECK-NEXT:     test.buffer_based in(%[[ARG1]]{{.*}}out(%[[ALLOC2]]
 555 //      CHECK:     dealloc %[[ALLOC2]]
 556 // CHECK-NEXT:     %{{.*}} = exp
 557 //      CHECK:   %[[ALLOC3:.*]] = alloc()
 558 // CHECK-NEXT:   linalg.copy(%[[ALLOC1]], %[[ALLOC3]])
 559 // CHECK-NEXT:   dealloc %[[ALLOC1]]
 560 //      CHECK:  ^[[BB3:.*]]({{.*}}):
 561 //      CHECK:  test.copy
 562 // CHECK-NEXT:  dealloc
 563
 564 // -----
 565
 566 // Test Case: buffer deallocation escaping
 567 // BufferDeallocation expected behavior: It must not dealloc %arg1 and %x
 568 // since they are operands of return operation and should escape from
 569 // deallocating. It should dealloc %y after CopyOp.
 570
 571 // CHECK-LABEL: func @memref_in_function_results
 572 func @memref_in_function_results(
 573   %arg0: memref<5xf32>,
 574   %arg1: memref<10xf32>,
 575   %arg2: memref<5xf32>) -> (memref<10xf32>, memref<15xf32>) {
 576   %x = alloc() : memref<15xf32>
 577   %y = alloc() : memref<5xf32>
 578   test.buffer_based in(%arg0: memref<5xf32>) out(%y: memref<5xf32>)
 579   test.copy(%y, %arg2) : (memref<5xf32>, memref<5xf32>)
 580   return %arg1, %x : memref<10xf32>, memref<15xf32>
 581 }
 582 //      CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>,
 583 // CHECK-SAME: %[[RESULT:.*]]: memref<5xf32>)
 584 //      CHECK: %[[X:.*]] = alloc()
 585 //      CHECK: %[[Y:.*]] = alloc()
 586 //      CHECK: test.copy
 587 //      CHECK: dealloc %[[Y]]
 588 //      CHECK: return %[[ARG1]], %[[X]]
 589
 590 // -----
 591
 592 // Test Case: nested region control flow
 593 // The alloc %1 flows through both if branches until it is finally returned.
 594 // Hence, it does not require a specific dealloc operation. However, %3
 595 // requires a dealloc.
 596
 597 // CHECK-LABEL: func @nested_region_control_flow
 598 func @nested_region_control_flow(
 599   %arg0 : index,
 600   %arg1 : index) -> memref<?x?xf32> {
 601   %0 = cmpi "eq", %arg0, %arg1 : index
 602   %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
 603   %2 = scf.if %0 -> (memref<?x?xf32>) {
 604     scf.yield %1 : memref<?x?xf32>
 605   } else {
 606     %3 = alloc(%arg0, %arg1) : memref<?x?xf32>
 607     scf.yield %1 : memref<?x?xf32>
 608   }
 609   return %2 : memref<?x?xf32>
 610 }
 611
 612 //      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
 613 // CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
 614 //      CHECK: scf.yield %[[ALLOC0]]
 615 //      CHECK: %[[ALLOC2:.*]] = alloc(%arg0, %arg1)
 616 // CHECK-NEXT: dealloc %[[ALLOC2]]
 617 // CHECK-NEXT: scf.yield %[[ALLOC0]]
 618 //      CHECK: return %[[ALLOC1]]
 619
 620 // -----
 621
 622 // Test Case: nested region control flow with a nested buffer allocation in a
 623 // divergent branch.
 624 // Buffer deallocation places a copy for both  %1 and %3, since they are
 625 // returned in the end.
 626
 627 // CHECK-LABEL: func @nested_region_control_flow_div
 628 func @nested_region_control_flow_div(
 629   %arg0 : index,
 630   %arg1 : index) -> memref<?x?xf32> {
 631   %0 = cmpi "eq", %arg0, %arg1 : index
 632   %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
 633   %2 = scf.if %0 -> (memref<?x?xf32>) {
 634     scf.yield %1 : memref<?x?xf32>
 635   } else {
 636     %3 = alloc(%arg0, %arg1) : memref<?x?xf32>
 637     scf.yield %3 : memref<?x?xf32>
 638   }
 639   return %2 : memref<?x?xf32>
 640 }
 641
 642 //      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
 643 // CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
 644 //      CHECK: %[[ALLOC2:.*]] = alloc
 645 // CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC2]])
 646 //      CHECK: scf.yield %[[ALLOC2]]
 647 //      CHECK: %[[ALLOC3:.*]] = alloc(%arg0, %arg1)
 648 //      CHECK: %[[ALLOC4:.*]] = alloc
 649 // CHECK-NEXT: linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
 650 //      CHECK: dealloc %[[ALLOC3]]
 651 //      CHECK: scf.yield %[[ALLOC4]]
 652 //      CHECK: dealloc %[[ALLOC0]]
 653 // CHECK-NEXT: return %[[ALLOC1]]
 654
 655 // -----
 656
 657 // Test Case: nested region control flow within a region interface.
 658 // No copies are required in this case since the allocation finally escapes
 659 // the method.
 660
 661 // CHECK-LABEL: func @inner_region_control_flow
 662 func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
 663   %0 = alloc(%arg0, %arg0) : memref<?x?xf32>
 664   %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
 665     ^bb0(%arg1 : memref<?x?xf32>):
 666       test.region_if_yield %arg1 : memref<?x?xf32>
 667   } else {
 668     ^bb0(%arg1 : memref<?x?xf32>):
 669       test.region_if_yield %arg1 : memref<?x?xf32>
 670   } join {
 671     ^bb0(%arg1 : memref<?x?xf32>):
 672       test.region_if_yield %arg1 : memref<?x?xf32>
 673   }
 674   return %1 : memref<?x?xf32>
 675 }
 676
 677 //      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
 678 // CHECK-NEXT: %[[ALLOC1:.*]] = test.region_if
 679 // CHECK-NEXT: ^bb0(%[[ALLOC2:.*]]:{{.*}}):
 680 // CHECK-NEXT: test.region_if_yield %[[ALLOC2]]
 681 //      CHECK: ^bb0(%[[ALLOC3:.*]]:{{.*}}):
 682 // CHECK-NEXT: test.region_if_yield %[[ALLOC3]]
 683 //      CHECK: ^bb0(%[[ALLOC4:.*]]:{{.*}}):
 684 // CHECK-NEXT: test.region_if_yield %[[ALLOC4]]
 685 //      CHECK: return %[[ALLOC1]]
 686
 687 // -----
 688
 689 // CHECK-LABEL: func @subview
 690 func @subview(%arg0 : index, %arg1 : index, %arg2 : memref<?x?xf32>) {
 691   %0 = alloc() : memref<64x4xf32, offset: 0, strides: [4, 1]>
 692   %1 = subview %0[%arg0, %arg1][%arg0, %arg1][%arg0, %arg1] :
 693     memref<64x4xf32, offset: 0, strides: [4, 1]>
 694   to memref<?x?xf32, offset: ?, strides: [?, ?]>
 695   test.copy(%1, %arg2) :
 696     (memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32>)
 697   return
 698 }
 699
 700 // CHECK-NEXT: %[[ALLOC:.*]] = alloc()
 701 // CHECK-NEXT: subview
 702 // CHECK-NEXT: test.copy
 703 // CHECK-NEXT: dealloc %[[ALLOC]]
 704 // CHECK-NEXT: return
 705
 706 // -----
 707
 708 // Test Case: In the presence of AllocaOps only the AllocOps has top be freed.
 709 // Therefore, all allocas are not handled.
 710
 711 // CHECK-LABEL: func @condBranchAlloca
 712 func @condBranchAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 713   cond_br %arg0, ^bb1, ^bb2
 714 ^bb1:
 715   br ^bb3(%arg1 : memref<2xf32>)
 716 ^bb2:
 717   %0 = alloca() : memref<2xf32>
 718   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 719   br ^bb3(%0 : memref<2xf32>)
 720 ^bb3(%1: memref<2xf32>):
 721   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
 722   return
 723 }
 724
 725 // CHECK-NEXT: cond_br
 726 //      CHECK: %[[ALLOCA:.*]] = alloca()
 727 //      CHECK: br ^bb3(%[[ALLOCA:.*]])
 728 // CHECK-NEXT: ^bb3
 729 // CHECK-NEXT: test.copy
 730 // CHECK-NEXT: return
 731
 732 // -----
 733
 734 // Test Case: In the presence of AllocaOps only the AllocOps has top be freed.
 735 // Therefore, all allocas are not handled. In this case, only alloc %0 has a
 736 // dealloc.
 737
 738 // CHECK-LABEL: func @ifElseAlloca
 739 func @ifElseAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 740   %0 = alloc() : memref<2xf32>
 741   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 742   cond_br %arg0,
 743     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
 744     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 745 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
 746   br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
 747 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
 748   br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
 749 ^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
 750   %7 = alloca() : memref<2xf32>
 751   test.buffer_based in(%5: memref<2xf32>) out(%7: memref<2xf32>)
 752   test.copy(%7, %arg2) : (memref<2xf32>, memref<2xf32>)
 753   return
 754 }
 755
 756 // CHECK-NEXT: %[[ALLOC:.*]] = alloc()
 757 // CHECK-NEXT: test.buffer_based
 758 //      CHECK: %[[ALLOCA:.*]] = alloca()
 759 // CHECK-NEXT: test.buffer_based
 760 //      CHECK: dealloc %[[ALLOC]]
 761 //      CHECK: test.copy
 762 // CHECK-NEXT: return
 763
 764 // -----
 765
 766 // CHECK-LABEL: func @ifElseNestedAlloca
 767 func @ifElseNestedAlloca(
 768   %arg0: i1,
 769   %arg1: memref<2xf32>,
 770   %arg2: memref<2xf32>) {
 771   %0 = alloca() : memref<2xf32>
 772   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
 773   cond_br %arg0,
 774     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
 775     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 776 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
 777   br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
 778 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
 779   cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
 780 ^bb3(%5: memref<2xf32>):
 781   br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
 782 ^bb4(%6: memref<2xf32>):
 783   br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
 784 ^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
 785   %9 = alloc() : memref<2xf32>
 786   test.buffer_based in(%7: memref<2xf32>) out(%9: memref<2xf32>)
 787   test.copy(%9, %arg2) : (memref<2xf32>, memref<2xf32>)
 788   return
 789 }
 790
 791 // CHECK-NEXT: %[[ALLOCA:.*]] = alloca()
 792 // CHECK-NEXT: test.buffer_based
 793 //      CHECK: %[[ALLOC:.*]] = alloc()
 794 // CHECK-NEXT: test.buffer_based
 795 //      CHECK: test.copy
 796 // CHECK-NEXT: dealloc %[[ALLOC]]
 797 // CHECK-NEXT: return
 798
 799 // -----
 800
 801 // CHECK-LABEL: func @nestedRegionsAndCondBranchAlloca
 802 func @nestedRegionsAndCondBranchAlloca(
 803   %arg0: i1,
 804   %arg1: memref<2xf32>,
 805   %arg2: memref<2xf32>) {
 806   cond_br %arg0, ^bb1, ^bb2
 807 ^bb1:
 808   br ^bb3(%arg1 : memref<2xf32>)
 809 ^bb2:
 810   %0 = alloc() : memref<2xf32>
 811   test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
 812   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
 813     %1 = alloca() : memref<2xf32>
 814     test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
 815     %tmp1 = exp %gen1_arg0 : f32
 816     test.region_yield %tmp1 : f32
 817   }
 818   br ^bb3(%0 : memref<2xf32>)
 819 ^bb3(%1: memref<2xf32>):
 820   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
 821   return
 822 }
 823 //      CHECK: (%[[cond:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %{{.*}}: {{.*}})
 824 // CHECK-NEXT:   cond_br %[[cond]], ^[[BB1:.*]], ^[[BB2:.*]]
 825 //      CHECK: ^[[BB1]]:
 826 //      CHECK: %[[ALLOC0:.*]] = alloc()
 827 // CHECK-NEXT: linalg.copy
 828 //      CHECK: ^[[BB2]]:
 829 //      CHECK:   %[[ALLOC1:.*]] = alloc()
 830 // CHECK-NEXT:   test.region_buffer_based in(%[[ARG1]]{{.*}}out(%[[ALLOC1]]
 831 //      CHECK:     %[[ALLOCA:.*]] = alloca()
 832 // CHECK-NEXT:     test.buffer_based in(%[[ARG1]]{{.*}}out(%[[ALLOCA]]
 833 //      CHECK:     %{{.*}} = exp
 834 //      CHECK:  %[[ALLOC2:.*]] = alloc()
 835 // CHECK-NEXT:  linalg.copy
 836 // CHECK-NEXT:  dealloc %[[ALLOC1]]
 837 //      CHECK:  ^[[BB3:.*]]({{.*}}):
 838 //      CHECK:  test.copy
 839 // CHECK-NEXT:  dealloc
 840
 841 // -----
 842
 843 // CHECK-LABEL: func @nestedRegionControlFlowAlloca
 844 func @nestedRegionControlFlowAlloca(
 845   %arg0 : index,
 846   %arg1 : index) -> memref<?x?xf32> {
 847   %0 = cmpi "eq", %arg0, %arg1 : index
 848   %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
 849   %2 = scf.if %0 -> (memref<?x?xf32>) {
 850     scf.yield %1 : memref<?x?xf32>
 851   } else {
 852     %3 = alloca(%arg0, %arg1) : memref<?x?xf32>
 853     scf.yield %1 : memref<?x?xf32>
 854   }
 855   return %2 : memref<?x?xf32>
 856 }
 857
 858 //      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
 859 // CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
 860 //      CHECK: scf.yield %[[ALLOC0]]
 861 //      CHECK: %[[ALLOCA:.*]] = alloca(%arg0, %arg1)
 862 // CHECK-NEXT: scf.yield %[[ALLOC0]]
 863 //      CHECK: return %[[ALLOC1]]
 864
 865 // -----
 866
 867 // Test Case: structured control-flow loop using a nested alloc.
 868 // The iteration argument %iterBuf has to be freed before yielding %3 to avoid
 869 // memory leaks.
 870
 871 // CHECK-LABEL: func @loop_alloc
 872 func @loop_alloc(
 873   %lb: index,
 874   %ub: index,
 875   %step: index,
 876   %buf: memref<2xf32>,
 877   %res: memref<2xf32>) {
 878   %0 = alloc() : memref<2xf32>
 879   %1 = scf.for %i = %lb to %ub step %step
 880     iter_args(%iterBuf = %buf) -> memref<2xf32> {
 881     %2 = cmpi "eq", %i, %ub : index
 882     %3 = alloc() : memref<2xf32>
 883     scf.yield %3 : memref<2xf32>
 884   }
 885   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
 886   return
 887 }
 888
 889 //      CHECK: %[[ALLOC0:.*]] = alloc()
 890 // CHECK-NEXT: dealloc %[[ALLOC0]]
 891 // CHECK-NEXT: %[[ALLOC1:.*]] = alloc()
 892 //      CHECK: linalg.copy(%arg3, %[[ALLOC1]])
 893 //      CHECK: %[[ALLOC2:.*]] = scf.for {{.*}} iter_args
 894 // CHECK-SAME: (%[[IALLOC:.*]] = %[[ALLOC1]]
 895 //      CHECK:    cmpi
 896 //      CHECK:    dealloc %[[IALLOC]]
 897 //      CHECK:    %[[ALLOC3:.*]] = alloc()
 898 //      CHECK:    %[[ALLOC4:.*]] = alloc()
 899 //      CHECK:    linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
 900 //      CHECK:    dealloc %[[ALLOC3]]
 901 //      CHECK:    scf.yield %[[ALLOC4]]
 902 //      CHECK: }
 903 //      CHECK: test.copy(%[[ALLOC2]], %arg4)
 904 // CHECK-NEXT: dealloc %[[ALLOC2]]
 905
 906 // -----
 907
 908 // Test Case: structured control-flow loop with a nested if operation.
 909 // The loop yields buffers that have been defined outside of the loop and the
 910 // backeges only use the iteration arguments (or one of its aliases).
 911 // Therefore, we do not have to (and are not allowed to) free any buffers
 912 // that are passed via the backedges.
 913
 914 // CHECK-LABEL: func @loop_nested_if_no_alloc
 915 func @loop_nested_if_no_alloc(
 916   %lb: index,
 917   %ub: index,
 918   %step: index,
 919   %buf: memref<2xf32>,
 920   %res: memref<2xf32>) {
 921   %0 = alloc() : memref<2xf32>
 922   %1 = scf.for %i = %lb to %ub step %step
 923     iter_args(%iterBuf = %buf) -> memref<2xf32> {
 924     %2 = cmpi "eq", %i, %ub : index
 925     %3 = scf.if %2 -> (memref<2xf32>) {
 926       scf.yield %0 : memref<2xf32>
 927     } else {
 928       scf.yield %iterBuf : memref<2xf32>
 929     }
 930     scf.yield %3 : memref<2xf32>
 931   }
 932   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
 933   return
 934 }
 935
 936 //      CHECK: %[[ALLOC0:.*]] = alloc()
 937 // CHECK-NEXT: %[[ALLOC1:.*]] = scf.for {{.*}} iter_args(%[[IALLOC:.*]] =
 938 //      CHECK: %[[ALLOC2:.*]] = scf.if
 939 //      CHECK: scf.yield %[[ALLOC0]]
 940 //      CHECK: scf.yield %[[IALLOC]]
 941 //      CHECK: scf.yield %[[ALLOC2]]
 942 //      CHECK: test.copy(%[[ALLOC1]], %arg4)
 943 //      CHECK: dealloc %[[ALLOC0]]
 944
 945 // -----
 946
 947 // Test Case: structured control-flow loop with a nested if operation using
 948 // a deeply nested buffer allocation.
 949 // Since the innermost allocation happens in a divergent branch, we have to
 950 // introduce additional copies for the nested if operation. Since the loop's
 951 // yield operation "returns" %3, it will return a newly allocated buffer.
 952 // Therefore, we have to free the iteration argument %iterBuf before
 953 // "returning" %3.
 954
 955 // CHECK-LABEL: func @loop_nested_if_alloc
 956 func @loop_nested_if_alloc(
 957   %lb: index,
 958   %ub: index,
 959   %step: index,
 960   %buf: memref<2xf32>) -> memref<2xf32> {
 961   %0 = alloc() : memref<2xf32>
 962   %1 = scf.for %i = %lb to %ub step %step
 963     iter_args(%iterBuf = %buf) -> memref<2xf32> {
 964     %2 = cmpi "eq", %i, %ub : index
 965     %3 = scf.if %2 -> (memref<2xf32>) {
 966       %4 = alloc() : memref<2xf32>
 967       scf.yield %4 : memref<2xf32>
 968     } else {
 969       scf.yield %0 : memref<2xf32>
 970     }
 971     scf.yield %3 : memref<2xf32>
 972   }
 973   return %1 : memref<2xf32>
 974 }
 975
 976 //      CHECK: %[[ALLOC0:.*]] = alloc()
 977 //      CHECK: %[[ALLOC1:.*]] = alloc()
 978 // CHECK-NEXT: linalg.copy(%arg3, %[[ALLOC1]])
 979 // CHECK-NEXT: %[[ALLOC2:.*]] = scf.for {{.*}} iter_args
 980 // CHECK-SAME: (%[[IALLOC:.*]] = %[[ALLOC1]]
 981 //      CHECK: dealloc %[[IALLOC]]
 982 //      CHECK: %[[ALLOC3:.*]] = scf.if
 983
 984 //      CHECK: %[[ALLOC4:.*]] = alloc()
 985 // CHECK-NEXT: %[[ALLOC5:.*]] = alloc()
 986 // CHECK-NEXT: linalg.copy(%[[ALLOC4]], %[[ALLOC5]])
 987 // CHECK-NEXT: dealloc %[[ALLOC4]]
 988 // CHECK-NEXT: scf.yield %[[ALLOC5]]
 989
 990 //      CHECK: %[[ALLOC6:.*]] = alloc()
 991 // CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC6]])
 992 // CHECK-NEXT: scf.yield %[[ALLOC6]]
 993
 994 //      CHECK: %[[ALLOC7:.*]] = alloc()
 995 // CHECK-NEXT: linalg.copy(%[[ALLOC3:.*]], %[[ALLOC7]])
 996 // CHECK-NEXT: dealloc %[[ALLOC3]]
 997 // CHECK-NEXT: scf.yield %[[ALLOC7]]
 998
 999 //      CHECK: dealloc %[[ALLOC0]]
1000 // CHECK-NEXT: return %[[ALLOC2]]
1001
1002 // -----
1003
1004 // Test Case: several nested structured control-flow loops with a deeply nested
1005 // buffer allocation inside an if operation.
1006 // Same behavior is an loop_nested_if_alloc: we have to insert deallocations
1007 // before each yield in all loops recursively.
1008
1009 // CHECK-LABEL: func @loop_nested_alloc
1010 func @loop_nested_alloc(
1011   %lb: index,
1012   %ub: index,
1013   %step: index,
1014   %buf: memref<2xf32>,
1015   %res: memref<2xf32>) {
1016   %0 = alloc() : memref<2xf32>
1017   %1 = scf.for %i = %lb to %ub step %step
1018     iter_args(%iterBuf = %buf) -> memref<2xf32> {
1019     %2 = scf.for %i2 = %lb to %ub step %step
1020       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
1021       %3 = scf.for %i3 = %lb to %ub step %step
1022         iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
1023         %4 = alloc() : memref<2xf32>
1024         %5 = cmpi "eq", %i, %ub : index
1025         %6 = scf.if %5 -> (memref<2xf32>) {
1026           %7 = alloc() : memref<2xf32>
1027           scf.yield %7 : memref<2xf32>
1028         } else {
1029           scf.yield %iterBuf3 : memref<2xf32>
1030         }
1031         scf.yield %6 : memref<2xf32>
1032       }
1033       scf.yield %3 : memref<2xf32>
1034     }
1035     scf.yield %2 : memref<2xf32>
1036   }
1037   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
1038   return
1039 }
1040
1041 //      CHECK: %[[ALLOC0:.*]] = alloc()
1042 // CHECK-NEXT: dealloc %[[ALLOC0]]
1043 // CHECK-NEXT: %[[ALLOC1:.*]] = alloc()
1044 // CHECK-NEXT: linalg.copy(%arg3, %[[ALLOC1]])
1045 // CHECK-NEXT: %[[VAL_7:.*]] = scf.for {{.*}} iter_args
1046 // CHECK-SAME: (%[[IALLOC0:.*]] = %[[ALLOC1]])
1047 //      CHECK: %[[ALLOC2:.*]] = alloc()
1048 // CHECK-NEXT: linalg.copy(%[[IALLOC0]], %[[ALLOC2]])
1049 // CHECK-NEXT: dealloc %[[IALLOC0]]
1050 // CHECK-NEXT: %[[ALLOC3:.*]] = scf.for {{.*}} iter_args
1051 // CHECK-SAME: (%[[IALLOC1:.*]] = %[[ALLOC2]])
1052 //      CHECK: %[[ALLOC5:.*]] = alloc()
1053 // CHECK-NEXT: linalg.copy(%[[IALLOC1]], %[[ALLOC5]])
1054 // CHECK-NEXT: dealloc %[[IALLOC1]]
1055
1056 //      CHECK: %[[ALLOC6:.*]] = scf.for {{.*}} iter_args
1057 // CHECK-SAME: (%[[IALLOC2:.*]] = %[[ALLOC5]])
1058 //      CHECK: %[[ALLOC8:.*]] = alloc()
1059 // CHECK-NEXT: dealloc %[[ALLOC8]]
1060 //      CHECK: %[[ALLOC9:.*]] = scf.if
1061
1062 //      CHECK: %[[ALLOC11:.*]] = alloc()
1063 // CHECK-NEXT: %[[ALLOC12:.*]] = alloc()
1064 // CHECK-NEXT: linalg.copy(%[[ALLOC11]], %[[ALLOC12]])
1065 // CHECK-NEXT: dealloc %[[ALLOC11]]
1066 // CHECK-NEXT: scf.yield %[[ALLOC12]]
1067
1068 //      CHECK: %[[ALLOC13:.*]] = alloc()
1069 // CHECK-NEXT: linalg.copy(%[[IALLOC2]], %[[ALLOC13]])
1070 // CHECK-NEXT: scf.yield %[[ALLOC13]]
1071
1072 //      CHECK: dealloc %[[IALLOC2]]
1073 // CHECK-NEXT: %[[ALLOC10:.*]] = alloc()
1074 // CHECK-NEXT: linalg.copy(%[[ALLOC9]], %[[ALLOC10]])
1075 // CHECK-NEXT: dealloc %[[ALLOC9]]
1076 // CHECK-NEXT: scf.yield %[[ALLOC10]]
1077
1078 //      CHECK: %[[ALLOC7:.*]] = alloc()
1079 // CHECK-NEXT: linalg.copy(%[[ALLOC6]], %[[ALLOC7]])
1080 // CHECK-NEXT: dealloc %[[ALLOC6]]
1081 // CHECK-NEXT: scf.yield %[[ALLOC7]]
1082
1083 //      CHECK: %[[ALLOC4:.*]] = alloc()
1084 // CHECK-NEXT: linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
1085 // CHECK-NEXT: dealloc %[[ALLOC3]]
1086 // CHECK-NEXT: scf.yield %[[ALLOC4]]
1087
1088 //      CHECK: test.copy(%[[VAL_7]], %arg4)
1089 // CHECK-NEXT: dealloc %[[VAL_7]]
1090
1091 // -----
1092
1093 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
1094 // The BufferDeallocation transformation should fail on this explicit
1095 // control-flow loop since they are not supported.
1096
1097 // CHECK-LABEL: func @loop_dynalloc
1098 func @loop_dynalloc(
1099   %arg0 : i32,
1100   %arg1 : i32,
1101   %arg2: memref<?xf32>,
1102   %arg3: memref<?xf32>) {
1103   %const0 = constant 0 : i32
1104   br ^loopHeader(%const0, %arg2 : i32, memref<?xf32>)
1105
1106 ^loopHeader(%i : i32, %buff : memref<?xf32>):
1107   %lessThan = cmpi "slt", %i, %arg1 : i32
1108   cond_br %lessThan,
1109     ^loopBody(%i, %buff : i32, memref<?xf32>),
1110     ^exit(%buff : memref<?xf32>)
1111
1112 ^loopBody(%val : i32, %buff2: memref<?xf32>):
1113   %const1 = constant 1 : i32
1114   %inc = addi %val, %const1 : i32
1115   %size = std.index_cast %inc : i32 to index
1116   %alloc1 = alloc(%size) : memref<?xf32>
1117   br ^loopHeader(%inc, %alloc1 : i32, memref<?xf32>)
1118
1119 ^exit(%buff3 : memref<?xf32>):
1120   test.copy(%buff3, %arg3) : (memref<?xf32>, memref<?xf32>)
1121   return
1122 }
1123
1124 // expected-error@+1 {{Structured control-flow loops are supported only}}
1125
1126 // -----
1127
1128 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
1129 // The BufferDeallocation transformation should fail on this explicit
1130 // control-flow loop since they are not supported.
1131
1132 // CHECK-LABEL: func @do_loop_alloc
1133 func @do_loop_alloc(
1134   %arg0 : i32,
1135   %arg1 : i32,
1136   %arg2: memref<2xf32>,
1137   %arg3: memref<2xf32>) {
1138   %const0 = constant 0 : i32
1139   br ^loopBody(%const0, %arg2 : i32, memref<2xf32>)
1140
1141 ^loopBody(%val : i32, %buff2: memref<2xf32>):
1142   %const1 = constant 1 : i32
1143   %inc = addi %val, %const1 : i32
1144   %alloc1 = alloc() : memref<2xf32>
1145   br ^loopHeader(%inc, %alloc1 : i32, memref<2xf32>)
1146
1147 ^loopHeader(%i : i32, %buff : memref<2xf32>):
1148   %lessThan = cmpi "slt", %i, %arg1 : i32
1149   cond_br %lessThan,
1150     ^loopBody(%i, %buff : i32, memref<2xf32>),
1151     ^exit(%buff : memref<2xf32>)
1152
1153 ^exit(%buff3 : memref<2xf32>):
1154   test.copy(%buff3, %arg3) : (memref<2xf32>, memref<2xf32>)
1155   return
1156 }
1157
1158 // expected-error@+1 {{Structured control-flow loops are supported only}}
1159
1160 // -----
1161
1162 // CHECK-LABEL: func @assumingOp(
1163 func @assumingOp(
1164   %arg0: !shape.witness,
1165   %arg2: memref<2xf32>,
1166   %arg3: memref<2xf32>) {
1167   // Confirm the alloc will be dealloc'ed in the block.
1168   %1 = shape.assuming %arg0 -> memref<2xf32> {
1169      %0 = alloc() : memref<2xf32>
1170     shape.assuming_yield %arg2 : memref<2xf32>
1171   }
1172   // Confirm the alloc will be returned and dealloc'ed after its use.
1173   %3 = shape.assuming %arg0 -> memref<2xf32> {
1174     %2 = alloc() : memref<2xf32>
1175     shape.assuming_yield %2 : memref<2xf32>
1176   }
1177   test.copy(%3, %arg3) : (memref<2xf32>, memref<2xf32>)
1178   return
1179 }
1180
1181 // CHECK-SAME: %[[ARG0:.*]]: !shape.witness,
1182 // CHECK-SAME: %[[ARG1:.*]]: {{.*}},
1183 // CHECK-SAME: %[[ARG2:.*]]: {{.*}}
1184 //      CHECK: %[[UNUSED_RESULT:.*]] = shape.assuming %[[ARG0]]
1185 // CHECK-NEXT:    %[[ALLOC0:.*]] = alloc()
1186 // CHECK-NEXT:    dealloc %[[ALLOC0]]
1187 // CHECK-NEXT:    shape.assuming_yield %[[ARG1]]
1188 //      CHECK: %[[ASSUMING_RESULT:.*]] = shape.assuming %[[ARG0]]
1189 // CHECK-NEXT:    %[[TMP_ALLOC:.*]] = alloc()
1190 // CHECK-NEXT:    %[[RETURNING_ALLOC:.*]] = alloc()
1191 // CHECK-NEXT:    linalg.copy(%[[TMP_ALLOC]], %[[RETURNING_ALLOC]])
1192 // CHECK-NEXT:    dealloc %[[TMP_ALLOC]]
1193 // CHECK-NEXT:    shape.assuming_yield %[[RETURNING_ALLOC]]
1194 //      CHECK: test.copy(%[[ASSUMING_RESULT:.*]], %[[ARG2]])
1195 // CHECK-NEXT: dealloc %[[ASSUMING_RESULT]]