// RUN: mlir-opt -buffer-placement -split-input-file %s | FileCheck %s

// This file checks the behaviour of BufferPlacement pass for moving Alloc and
// Dealloc operations and inserting the missing the DeallocOps in their correct
// positions.

// Test Case:
//    bb0
//   /   \
//  bb1  bb2 <- Initial position of AllocOp
//   \   /
//    bb3
// BufferPlacement Expected Behaviour: It should move the existing AllocOp to
// the entry block, and insert a DeallocOp at the exit block after CopyOp since
// %1 is an alias for %0 and %arg1.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @condBranch
func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  cond_br %arg0, ^bb1, ^bb2
^bb1:
  br ^bb3(%arg1 : memref<2xf32>)
^bb2:
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
      ins(%arg1: memref<2xf32>)
     outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[ALLOC:.*]] = alloc()
// CHECK-NEXT: cond_br
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc %[[ALLOC]]
// CHECK-NEXT: return

// -----

// Test Case:
//    bb0
//   /   \
//  bb1  bb2 <- Initial position of AllocOp
//   \   /
//    bb3
// BufferPlacement Expected Behaviour: It should not move the existing AllocOp
// to any other block since the alloc has a dynamic dependency to block argument
// %0 in bb2. Since the dynamic type is passed to bb3 via the block argument %2,
// it is currently required to allocate a temporary buffer for %2 that gets
// copies of %arg0 and %1 with their appropriate shape dimensions. The copy
// buffer deallocation will be applied to %2 in block bb3.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @condBranchDynamicType
func @condBranchDynamicType(
  %arg0: i1,
  %arg1: memref<?xf32>,
  %arg2: memref<?xf32>,
  %arg3: index) {
  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
^bb1:
  br ^bb3(%arg1 : memref<?xf32>)
^bb2(%0: index):
  %1 = alloc(%0) : memref<?xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<?xf32>)
   outs(%1: memref<?xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^bb3(%1 : memref<?xf32>)
^bb3(%2: memref<?xf32>):
  "linalg.copy"(%2, %arg2) : (memref<?xf32>, memref<?xf32>) -> ()
  return
}

// CHECK-NEXT: cond_br
//      CHECK: %[[DIM0:.*]] = dim
// CHECK-NEXT: %[[ALLOC0:.*]] = alloc(%[[DIM0]])
// CHECK-NEXT: linalg.copy(%{{.*}}, %[[ALLOC0]])
//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
// CHECK-NEXT: %[[ALLOC1:.*]] = alloc(%[[IDX]])
// CHECK-NEXT: linalg.generic
//      CHECK: %[[DIM1:.*]] = dim %[[ALLOC1]]
// CHECK-NEXT: %[[ALLOC2:.*]] = alloc(%[[DIM1]])
// CHECK-NEXT: linalg.copy(%[[ALLOC1]], %[[ALLOC2]])
// CHECK-NEXT: dealloc %[[ALLOC1]]
// CHECK-NEXT: br ^bb3
// CHECK-NEXT: ^bb3(%[[ALLOC3:.*]]:{{.*}})
//      CHECK: linalg.copy(%[[ALLOC3]],
// CHECK-NEXT: dealloc %[[ALLOC3]]
// CHECK-NEXT: return

// -----

// Test Case:
//      bb0
//     /    \
//   bb1    bb2 <- Initial position of AllocOp
//    |     /  \
//    |   bb3  bb4
//    |     \  /
//    \     bb5
//     \    /
//       bb6
//        |
//       bb7
// BufferPlacement Expected Behaviour: It should not move the existing AllocOp
// to any other block since the alloc has a dynamic dependency to block argument
// %0 in bb2. Since the dynamic type is passed to bb5 via the block argument %2
// and to bb6 via block argument %3, it is currently required to allocate
// temporary buffers for %2 and %3 that gets copies of %1 and %arg0 1 with their
// appropriate shape dimensions. The copy buffer deallocations will be applied
// to %2 in block bb5 and to %3 in block bb6. Furthermore, there should be no
// copy inserted for %4.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @condBranchDynamicType
func @condBranchDynamicTypeNested(
  %arg0: i1,
  %arg1: memref<?xf32>,
  %arg2: memref<?xf32>,
  %arg3: index) {
  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
^bb1:
  br ^bb6(%arg1 : memref<?xf32>)
^bb2(%0: index):
  %1 = alloc(%0) : memref<?xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<?xf32>)
   outs(%1: memref<?xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0, ^bb3, ^bb4
^bb3:
  br ^bb5(%1 : memref<?xf32>)
^bb4:
  br ^bb5(%1 : memref<?xf32>)
^bb5(%2: memref<?xf32>):
  br ^bb6(%2 : memref<?xf32>)
^bb6(%3: memref<?xf32>):
  br ^bb7(%3 : memref<?xf32>)
^bb7(%4: memref<?xf32>):
  "linalg.copy"(%4, %arg2) : (memref<?xf32>, memref<?xf32>) -> ()
  return
}

// CHECK-NEXT: cond_br
//      CHECK: ^bb1
//      CHECK: %[[DIM0:.*]] = dim
// CHECK-NEXT: %[[ALLOC0:.*]] = alloc(%[[DIM0]])
// CHECK-NEXT: linalg.copy(%{{.*}}, %[[ALLOC0]])
//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
// CHECK-NEXT: %[[ALLOC1:.*]] = alloc(%[[IDX]])
// CHECK-NEXT: linalg.generic
//      CHECK: cond_br
//      CHECK: ^bb3:
// CHECK-NEXT: br ^bb5(%[[ALLOC1]]{{.*}})
//      CHECK: ^bb4:
// CHECK-NEXT: br ^bb5(%[[ALLOC1]]{{.*}})
// CHECK-NEXT: ^bb5(%[[ALLOC2:.*]]:{{.*}})
//      CHECK: %[[DIM2:.*]] = dim %[[ALLOC2]]
// CHECK-NEXT: %[[ALLOC3:.*]] = alloc(%[[DIM2]])
// CHECK-NEXT: linalg.copy(%[[ALLOC2]], %[[ALLOC3]])
// CHECK-NEXT: dealloc %[[ALLOC1]]
// CHECK-NEXT: br ^bb6(%[[ALLOC3]]{{.*}})
// CHECK-NEXT: ^bb6(%[[ALLOC4:.*]]:{{.*}})
// CHECK-NEXT: br ^bb7(%[[ALLOC4]]{{.*}})
// CHECK-NEXT: ^bb7(%[[ALLOC5:.*]]:{{.*}})
//      CHECK: linalg.copy(%[[ALLOC5]],
// CHECK-NEXT: dealloc %[[ALLOC4]]
// CHECK-NEXT: return

// -----

// Test Case: Existing AllocOp with no users.
// BufferPlacement Expected Behaviour: It should insert a DeallocOp right before
// ReturnOp.

// CHECK-LABEL: func @emptyUsesValue
func @emptyUsesValue(%arg0: memref<4xf32>) {
  %0 = alloc() : memref<4xf32>
  return
}
// CHECK-NEXT: %[[ALLOC:.*]] = alloc()
// CHECK-NEXT: dealloc %[[ALLOC]]
// CHECK-NEXT: return

// -----

// Test Case:
//    bb0
//   /   \
//  |    bb1 <- Initial position of AllocOp
//   \   /
//    bb2
// BufferPlacement Expected Behaviour: It should move the existing AllocOp to
// the entry block and insert a DeallocOp at the exit block after CopyOp since
// %1 is an alias for %0 and %arg1.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @criticalEdge
func @criticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
^bb1:
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^bb2(%0 : memref<2xf32>)
^bb2(%1: memref<2xf32>):
  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[ALLOC:.*]] = alloc()
// CHECK-NEXT: cond_br
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc %[[ALLOC]]
// CHECK-NEXT: return

// -----

// Test Case:
//    bb0 <- Initial position of AllocOp
//   /   \
//  |    bb1
//   \   /
//    bb2
// BufferPlacement Expected Behaviour: It shouldn't move the alloc position. It
// only inserts a DeallocOp at the exit block after CopyOp since %1 is an alias
// for %0 and %arg1.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @invCriticalEdge
func @invCriticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
^bb1:
  br ^bb2(%0 : memref<2xf32>)
^bb2(%1: memref<2xf32>):
  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

//      CHECK: dealloc
// CHECK-NEXT: return

// -----

// Test Case:
//    bb0 <- Initial position of the first AllocOp
//   /   \
//  bb1  bb2
//   \   /
//    bb3 <- Initial position of the second AllocOp
// BufferPlacement Expected Behaviour: It shouldn't move the AllocOps. It only
// inserts two missing DeallocOps in the exit block. %5 is an alias for %0.
// Therefore, the DeallocOp for %0 should occur after the last GenericOp. The
// Dealloc for %7 should happen after the CopyOp.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @ifElse
func @ifElse(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0,
    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
  br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
  br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
  %7 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%5: memref<2xf32>)
   outs(%7: memref<2xf32>) {
  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
    %tmp2 = exp %gen2_arg0 : f32
    linalg.yield %tmp2 : f32
  }
  "linalg.copy"(%7, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic
//      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic
//      CHECK: dealloc %[[FIRST_ALLOC]]
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
// CHECK-NEXT: return

// -----

// Test Case: No users for buffer in if-else CFG
//    bb0 <- Initial position of AllocOp
//   /   \
//  bb1  bb2
//   \   /
//    bb3
// BufferPlacement Expected Behaviour: It shouldn't move the AllocOp. It only
// inserts a missing DeallocOp in the exit block since %5 or %6 are the latest
// aliases of %0.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @ifElseNoUsers
func @ifElseNoUsers(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0,
    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
  br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
  br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
  "linalg.copy"(%arg1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
//      CHECK: dealloc %[[FIRST_ALLOC]]
// CHECK-NEXT: return

// -----

// Test Case:
//      bb0 <- Initial position of the first AllocOp
//     /    \
//   bb1    bb2
//    |     /  \
//    |   bb3  bb4
//    \     \  /
//     \     /
//       bb5 <- Initial position of the second AllocOp
// BufferPlacement Expected Behaviour: AllocOps shouldn't be moved.
// Two missing DeallocOps should be inserted in the exit block.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @ifElseNested
func @ifElseNested(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0,
    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
  br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
  cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
^bb3(%5: memref<2xf32>):
  br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
^bb4(%6: memref<2xf32>):
  br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
  %9 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%7: memref<2xf32>)
   outs(%9: memref<2xf32>) {
  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
    %tmp2 = exp %gen2_arg0 : f32
    linalg.yield %tmp2 : f32
  }
  "linalg.copy"(%9, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic
//      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic
//      CHECK: dealloc %[[FIRST_ALLOC]]
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
// CHECK-NEXT: return

// -----

// Test Case: Dead operations in a single block.
// BufferPlacement Expected Behaviour: It shouldn't move the AllocOps. It only
// inserts the two missing DeallocOps after the last GenericOp.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @redundantOperations
func @redundantOperations(%arg0: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg0: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  %1 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%0: memref<2xf32>)
   outs(%1: memref<2xf32>) {
  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
    %tmp2 = exp %gen2_arg0 : f32
    linalg.yield %tmp2 : f32
  }
  return
}

//      CHECK: (%[[ARG0:.*]]: {{.*}})
// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic {{.*}} ins(%[[ARG0]]{{.*}}outs(%[[FIRST_ALLOC]]
//      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic {{.*}} ins(%[[FIRST_ALLOC]]{{.*}}outs(%[[SECOND_ALLOC]]
//      CHECK: dealloc
// CHECK-NEXT: dealloc
// CHECK-NEXT: return

// -----

// Test Case:
//                                     bb0
//                                    /   \
// Initial pos of the 1st AllocOp -> bb1  bb2 <- Initial pos of the 2nd AllocOp
//                                    \   /
//                                     bb3
// BufferPlacement Expected Behaviour: Both AllocOps should be moved to the
// entry block. Both missing DeallocOps should be moved to the exit block after
// CopyOp since %arg2 is an alias for %0 and %1.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @moving_alloc_and_inserting_missing_dealloc
func @moving_alloc_and_inserting_missing_dealloc(
  %cond: i1,
    %arg0: memref<2xf32>,
    %arg1: memref<2xf32>) {
  cond_br %cond, ^bb1, ^bb2
^bb1:
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg0: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^exit(%0 : memref<2xf32>)
^bb2:
  %1 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg0: memref<2xf32>)
   outs(%1: memref<2xf32>) {
  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
    %tmp2 = exp %gen2_arg0 : f32
    linalg.yield %tmp2 : f32
  }
  br ^exit(%1 : memref<2xf32>)
^exit(%arg2: memref<2xf32>):
  "linalg.copy"(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %{{.*}} = alloc()
// CHECK-NEXT: %{{.*}} = alloc()
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc
// CHECK-NEXT: dealloc
// CHECK-NEXT: return

// -----

// Test Case: Invalid position of the DeallocOp. There is a user after
// deallocation.
//   bb0
//  /   \
// bb1  bb2 <- Initial position of AllocOp
//  \   /
//   bb3
// BufferPlacement Expected Behaviour: It should move the AllocOp to the entry
// block. The existing DeallocOp should be moved to exit block.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @moving_invalid_dealloc_op_complex
func @moving_invalid_dealloc_op_complex(
  %cond: i1,
    %arg0: memref<2xf32>,
    %arg1: memref<2xf32>) {
  cond_br %cond, ^bb1, ^bb2
^bb1:
  br ^exit(%arg0 : memref<2xf32>)
^bb2:
  %1 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg0: memref<2xf32>)
   outs(%1: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  dealloc %1 : memref<2xf32>
  br ^exit(%1 : memref<2xf32>)
^exit(%arg2: memref<2xf32>):
  "linalg.copy"(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %{{.*}} = alloc()
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc
// CHECK-NEXT: return

// -----

// Test Case: Inserting missing DeallocOp in a single block.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @inserting_missing_dealloc_simple
func @inserting_missing_dealloc_simple(
  %arg0 : memref<2xf32>,
  %arg1: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg0: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  "linalg.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

//      CHECK: linalg.copy
// CHECK-NEXT: dealloc

// -----

// Test Case: Moving invalid DeallocOp (there is a user after deallocation) in a
// single block.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @moving_invalid_dealloc_op
func @moving_invalid_dealloc_op(%arg0 : memref<2xf32>, %arg1: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg0: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  dealloc %0 : memref<2xf32>
  "linalg.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

//      CHECK: linalg.copy
// CHECK-NEXT: dealloc

// -----

// Test Case: Nested regions - This test defines a GenericOp inside the region of
// another GenericOp.
// BufferPlacement Expected Behaviour: The AllocOp of inner GenericOp should remain
// inside the region of outer GenericOp and it should insert the missing DeallocOp
// in the same region. The AllocOp of the outer GenericOp should be moved to the
// entry block and its missing DeallocOp should be inserted after Linalg.Copy.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @nested_regions_and_cond_branch
func @nested_regions_and_cond_branch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  cond_br %arg0, ^bb1, ^bb2
^bb1:
  br ^bb3(%arg1 : memref<2xf32>)
^bb2:
  %0 = alloc() : memref<2xf32>
  linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %1 = alloc() : memref<2xf32>
    linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
      ins(%arg1: memref<2xf32>)
     outs(%1: memref<2xf32>) {
    ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
      %tmp2 = exp %gen2_arg0 : f32
      linalg.yield %tmp2 : f32
    }
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}
//      CHECK: (%[[cond:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %{{.*}}: {{.*}})
// CHECK-NEXT:   %[[GENERIC1_ALLOC:.*]] = alloc()
// CHECK-NEXT:   cond_br %[[cond]], ^[[BB1:.*]], ^[[BB2:.*]]
//      CHECK: ^[[BB2]]:
// CHECK-NEXT:   linalg.generic {{{.*}}} ins(%[[ARG1]]{{.*}}outs(%[[GENERIC1_ALLOC]]
//      CHECK:     %[[GENERIC2_ALLOC:.*]] = alloc()
// CHECK-NEXT:     linalg.generic {{{.*}}} ins(%[[ARG1]]{{.*}}outs(%[[GENERIC2_ALLOC]]
//      CHECK:     dealloc %[[GENERIC2_ALLOC]]
// CHECK-NEXT:     %{{.*}} = exp
//      CHECK:  ^[[BB3:.*]]({{.*}}):
//      CHECK:  linalg.copy
// CHECK-NEXT:  dealloc %[[GENERIC1_ALLOC]]

// -----

// Test Case: buffer deallocation escaping
// BufferPlacement Expected Behaviour: It must not dealloc %arg1 and %x
// since they are operands of return operation and should escape from
// deallocating. It should dealloc %y after linalg.copy.

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @memref_in_function_results
func @memref_in_function_results(%arg0: memref<5xf32>, %arg1: memref<10xf32>, %arg2: memref<5xf32>) -> (memref<10xf32>, memref<15xf32>) {
  %x = alloc() : memref<15xf32>
  %y = alloc() : memref<5xf32>
  linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
      ins(%arg0: memref<5xf32>)
     outs(%y: memref<5xf32>) {
  ^bb0(%arg3: f32, %arg4: f32):
    %2 = exp %arg3 : f32
    linalg.yield %2 : f32
  }
  linalg.copy(%y, %arg2) : memref<5xf32>, memref<5xf32>
  return %arg1, %x : memref<10xf32>, memref<15xf32>
}
// CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, %[[RESULT:.*]]: memref<5xf32>)
// CHECK: %[[X:.*]] = alloc()
// CHECK: %[[Y:.*]] = alloc()
// CHECK: linalg.copy
// CHECK: dealloc %[[Y]]
// CHECK: return %[[ARG1]], %[[X]]

// -----

// Test Case: nested region control flow
// The alloc position of %1 does not need to be changed and flows through
// both if branches until it is finally returned. Hence, it does not
// require a specific dealloc operation. However, %3 requires a dealloc.

// CHECK-LABEL: func @nested_region_control_flow
func @nested_region_control_flow(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = cmpi "eq", %arg0, %arg1 : index
  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    scf.yield %1 : memref<?x?xf32>
  } else {
    %3 = alloc(%arg0, %arg1) : memref<?x?xf32>
    scf.yield %1 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}

//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
//      CHECK: scf.yield %[[ALLOC0]]
//      CHECK: %[[ALLOC2:.*]] = alloc(%arg0, %arg1)
// CHECK-NEXT: dealloc %[[ALLOC2]]
// CHECK-NEXT: scf.yield %[[ALLOC0]]
//      CHECK: return %[[ALLOC1]]

// -----

// Test Case: nested region control flow with a nested buffer allocation in a
// divergent branch.
// The alloc positions of %1, %3 does not need to be changed since
// BufferPlacement does not move allocs out of nested regions at the moment.
// However, since %3 is allocated and "returned" in a divergent branch, we have
// to allocate a temporary buffer (like in condBranchDynamicTypeNested).

// CHECK-LABEL: func @nested_region_control_flow_div
func @nested_region_control_flow_div(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = cmpi "eq", %arg0, %arg1 : index
  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    scf.yield %1 : memref<?x?xf32>
  } else {
    %3 = alloc(%arg0, %arg1) : memref<?x?xf32>
    scf.yield %3 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}

//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
//      CHECK: %[[ALLOC2:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC2]])
//      CHECK: scf.yield %[[ALLOC2]]
//      CHECK: %[[ALLOC3:.*]] = alloc(%arg0, %arg1)
//      CHECK: %[[ALLOC4:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
//      CHECK: dealloc %[[ALLOC3]]
//      CHECK: scf.yield %[[ALLOC4]]
//      CHECK: dealloc %[[ALLOC0]]
// CHECK-NEXT: return %[[ALLOC1]]

// -----

// Test Case: deeply nested region control flow with a nested buffer allocation
// in a divergent branch.
// The alloc positions of %1, %4 and %5 does not need to be changed since
// BufferPlacement does not move allocs out of nested regions at the moment.
// However, since %4 is allocated and "returned" in a divergent branch, we have
// to allocate several temporary buffers (like in condBranchDynamicTypeNested).

// CHECK-LABEL: func @nested_region_control_flow_div_nested
func @nested_region_control_flow_div_nested(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = cmpi "eq", %arg0, %arg1 : index
  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    %3 = scf.if %0 -> (memref<?x?xf32>) {
      scf.yield %1 : memref<?x?xf32>
    } else {
      %4 = alloc(%arg0, %arg1) : memref<?x?xf32>
      scf.yield %4 : memref<?x?xf32>
    }
    scf.yield %3 : memref<?x?xf32>
  } else {
    %5 = alloc(%arg1, %arg1) : memref<?x?xf32>
    scf.yield %5 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}
//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
// CHECK-NEXT: %[[ALLOC2:.*]] = scf.if
//      CHECK: %[[ALLOC3:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC3]])
//      CHECK: scf.yield %[[ALLOC3]]
//      CHECK: %[[ALLOC4:.*]] = alloc(%arg0, %arg1)
//      CHECK: %[[ALLOC5:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC4]], %[[ALLOC5]])
//      CHECK: dealloc %[[ALLOC4]]
//      CHECK: scf.yield %[[ALLOC5]]
//      CHECK: %[[ALLOC6:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC2]], %[[ALLOC6]])
//      CHECK: dealloc %[[ALLOC2]]
//      CHECK: scf.yield %[[ALLOC6]]
//      CHECK: %[[ALLOC7:.*]] = alloc(%arg1, %arg1)
//      CHECK: %[[ALLOC8:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC7]], %[[ALLOC8]])
//      CHECK: dealloc %[[ALLOC7]]
//      CHECK: scf.yield %[[ALLOC8]]
//      CHECK: dealloc %[[ALLOC0]]
// CHECK-NEXT: return %[[ALLOC1]]

// -----

// Test Case: nested region control flow within a region interface.
// The alloc positions of %0 does not need to be changed and no copies are
// required in this case since the allocation finally escapes the method.

// CHECK-LABEL: func @inner_region_control_flow
func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
  %0 = alloc(%arg0, %arg0) : memref<?x?xf32>
  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
    ^bb0(%arg1 : memref<?x?xf32>):
      test.region_if_yield %arg1 : memref<?x?xf32>
  } else {
    ^bb0(%arg1 : memref<?x?xf32>):
      test.region_if_yield %arg1 : memref<?x?xf32>
  } join {
    ^bb0(%arg1 : memref<?x?xf32>):
      test.region_if_yield %arg1 : memref<?x?xf32>
  }
  return %1 : memref<?x?xf32>
}

//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
// CHECK-NEXT: %[[ALLOC1:.*]] = test.region_if
// CHECK-NEXT: ^bb0(%[[ALLOC2:.*]]:{{.*}}):
// CHECK-NEXT: test.region_if_yield %[[ALLOC2]]
//      CHECK: ^bb0(%[[ALLOC3:.*]]:{{.*}}):
// CHECK-NEXT: test.region_if_yield %[[ALLOC3]]
//      CHECK: ^bb0(%[[ALLOC4:.*]]:{{.*}}):
// CHECK-NEXT: test.region_if_yield %[[ALLOC4]]
//      CHECK: return %[[ALLOC1]]

// -----

// Test Case: nested region control flow within a region interface including an
// allocation in a divergent branch.
// The alloc positions of %1 and %2 does not need to be changed since
// BufferPlacement does not move allocs out of nested regions at the moment.
// However, since %2 is allocated and yielded in a divergent branch, we have
// to allocate several temporary buffers (like in condBranchDynamicTypeNested).

// CHECK-LABEL: func @inner_region_control_flow_div
func @inner_region_control_flow_div(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = alloc(%arg0, %arg0) : memref<?x?xf32>
  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
    ^bb0(%arg2 : memref<?x?xf32>):
      test.region_if_yield %arg2 : memref<?x?xf32>
  } else {
    ^bb0(%arg2 : memref<?x?xf32>):
      %2 = alloc(%arg0, %arg1) : memref<?x?xf32>
      test.region_if_yield %2 : memref<?x?xf32>
  } join {
    ^bb0(%arg2 : memref<?x?xf32>):
      test.region_if_yield %arg2 : memref<?x?xf32>
  }
  return %1 : memref<?x?xf32>
}

//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
// CHECK-NEXT: %[[ALLOC1:.*]] = test.region_if
// CHECK-NEXT: ^bb0(%[[ALLOC2:.*]]:{{.*}}):
//      CHECK: %[[ALLOC3:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC2]], %[[ALLOC3]])
// CHECK-NEXT: test.region_if_yield %[[ALLOC3]]
//      CHECK: ^bb0(%[[ALLOC4:.*]]:{{.*}}):
//      CHECK: %[[ALLOC5:.*]] = alloc
//      CHECK: %[[ALLOC6:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC5]], %[[ALLOC6]])
// CHECK-NEXT: dealloc %[[ALLOC5]]
// CHECK-NEXT: test.region_if_yield %[[ALLOC6]]
//      CHECK: ^bb0(%[[ALLOC7:.*]]:{{.*}}):
//      CHECK: %[[ALLOC8:.*]] = alloc
// CHECK-NEXT: linalg.copy(%[[ALLOC7]], %[[ALLOC8]])
// CHECK-NEXT: dealloc %[[ALLOC7]]
// CHECK-NEXT: test.region_if_yield %[[ALLOC8]]
//      CHECK: dealloc %[[ALLOC0]]
// CHECK-NEXT: return %[[ALLOC1]]

// -----

// CHECK-LABEL: func @subview
func @subview(%arg0 : index, %arg1 : index, %arg2 : memref<?x?xf32>) {
  %0 = alloc() : memref<64x4xf32, offset: 0, strides: [4, 1]>
  %1 = subview %0[%arg0, %arg1][%arg0, %arg1][%arg0, %arg1] :
    memref<64x4xf32, offset: 0, strides: [4, 1]>
  to memref<?x?xf32, offset: ?, strides: [?, ?]>
  "linalg.copy"(%1, %arg2) :
    (memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32>) -> ()
  return
}

// CHECK-NEXT: %[[ALLOC:.*]] = alloc()
// CHECK-NEXT: subview
// CHECK-NEXT: linalg.copy
// CHECK-NEXT: dealloc %[[ALLOC]]
// CHECK-NEXT: return

// -----

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @condBranchAlloca
func @condBranchAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  cond_br %arg0, ^bb1, ^bb2
^bb1:
  br ^bb3(%arg1 : memref<2xf32>)
^bb2:
  %0 = alloca() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: cond_br
//      CHECK: %[[ALLOCA:.*]] = alloca()
//      CHECK: br ^bb3(%[[ALLOCA:.*]])
// CHECK-NEXT: ^bb3
// CHECK-NEXT: linalg.copy
// CHECK-NEXT: return

// -----

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @ifElseAlloca
func @ifElseAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0,
    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
  br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
  br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
  %7 = alloca() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%5: memref<2xf32>)
   outs(%7: memref<2xf32>) {
  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
    %tmp2 = exp %gen2_arg0 : f32
    linalg.yield %tmp2 : f32
  }
  "linalg.copy"(%7, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic
//      CHECK: %[[ALLOCA:.*]] = alloca()
// CHECK-NEXT: linalg.generic
//      CHECK: dealloc %[[ALLOC]]
//      CHECK: linalg.copy
// CHECK-NEXT: return

// -----

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @ifElseNestedAlloca
func @ifElseNestedAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  %0 = alloca() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  cond_br %arg0,
    ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
    ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
  br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
  cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
^bb3(%5: memref<2xf32>):
  br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
^bb4(%6: memref<2xf32>):
  br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
  %9 = alloc() : memref<2xf32>
  linalg.generic {
    indexing_maps = [#map0, #map0],
    iterator_types = ["parallel"]}
    ins(%7: memref<2xf32>)
   outs(%9: memref<2xf32>) {
  ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
    %tmp2 = exp %gen2_arg0 : f32
    linalg.yield %tmp2 : f32
  }
  "linalg.copy"(%9, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-NEXT: %[[ALLOCA:.*]] = alloca()
// CHECK-NEXT: linalg.generic
//      CHECK: %[[ALLOC:.*]] = alloc()
// CHECK-NEXT: linalg.generic
//      CHECK: linalg.copy
// CHECK-NEXT: dealloc %[[ALLOC]]
// CHECK-NEXT: return

// -----

#map0 = affine_map<(d0) -> (d0)>

// CHECK-LABEL: func @nestedRegionsAndCondBranchAlloca
func @nestedRegionsAndCondBranchAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  cond_br %arg0, ^bb1, ^bb2
^bb1:
  br ^bb3(%arg1 : memref<2xf32>)
^bb2:
  %0 = alloc() : memref<2xf32>
  linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %1 = alloca() : memref<2xf32>
    linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
    ins(%arg1: memref<2xf32>)
   outs(%1: memref<2xf32>) {
    ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
      %tmp2 = exp %gen2_arg0 : f32
      linalg.yield %tmp2 : f32
    }
    %tmp1 = exp %gen1_arg0 : f32
    linalg.yield %tmp1 : f32
  }
  br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}
//      CHECK: (%[[cond:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %{{.*}}: {{.*}})
// CHECK-NEXT:   %[[ALLOC:.*]] = alloc()
// CHECK-NEXT:   cond_br %[[cond]], ^[[BB1:.*]], ^[[BB2:.*]]
//      CHECK: ^[[BB2]]:
// CHECK-NEXT:   linalg.generic {{{.*}}} ins(%[[ARG1]]{{.*}}outs(%[[ALLOC]]
//      CHECK:     %[[ALLOCA:.*]] = alloca()
// CHECK-NEXT:     linalg.generic {{{.*}}} ins(%[[ARG1]]{{.*}}outs(%[[ALLOCA]]
//      CHECK:     %{{.*}} = exp
//      CHECK:  ^[[BB3:.*]]({{.*}}):
//      CHECK:  linalg.copy
// CHECK-NEXT:  dealloc %[[ALLOC]]

// -----

// CHECK-LABEL: func @nestedRegionControlFlowAlloca
func @nestedRegionControlFlowAlloca(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = cmpi "eq", %arg0, %arg1 : index
  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    scf.yield %1 : memref<?x?xf32>
  } else {
    %3 = alloca(%arg0, %arg1) : memref<?x?xf32>
    scf.yield %1 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}

//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
//      CHECK: scf.yield %[[ALLOC0]]
//      CHECK: %[[ALLOCA:.*]] = alloca(%arg0, %arg1)
// CHECK-NEXT: scf.yield %[[ALLOC0]]
//      CHECK: return %[[ALLOC1]]

// -----

// Test Case: structured control-flow loop using a nested alloc.
// The alloc positions of %3 will not be changed, but the iteration argument
// %iterBuf has to be freed before yielding %3 to avoid memory leaks.

// -----

// CHECK-LABEL: func @loop_alloc
func @loop_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>,
  %res: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = cmpi "eq", %i, %ub : index
    %3 = alloc() : memref<2xf32>
    scf.yield %3 : memref<2xf32>
  }
  "linalg.copy"(%1, %res) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

//      CHECK: %[[ALLOC0:.*]] = alloc()
// CHECK-NEXT: dealloc %[[ALLOC0]]
// CHECK-NEXT: %[[ALLOC1:.*]] = alloc()
//      CHECK: linalg.copy(%arg3, %[[ALLOC1]])
//      CHECK: %[[ALLOC2:.*]] = scf.for {{.*}} iter_args(%[[IALLOC:.*]] = %[[ALLOC1]]
//      CHECK:    cmpi
//      CHECK:    dealloc %[[IALLOC]]
//      CHECK:    %[[ALLOC3:.*]] = alloc()
//      CHECK:    %[[ALLOC4:.*]] = alloc()
//      CHECK:    linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
//      CHECK:    dealloc %[[ALLOC3]]
//      CHECK:    scf.yield %[[ALLOC4]]
//      CHECK: }
//      CHECK: linalg.copy(%[[ALLOC2]], %arg4)
// CHECK-NEXT: dealloc %[[ALLOC2]]

// -----

// Test Case: structured control-flow loop with a nested if operation.
// The loop yields buffers that have been defined outside of the loop and the
// backeges only use the iteration arguments (or one of its aliases).
// Therefore, we do not have to (and are not allowed to) free any buffers
// that are passed via the backedges.

// CHECK-LABEL: func @loop_nested_if_no_alloc
func @loop_nested_if_no_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>,
  %res: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = cmpi "eq", %i, %ub : index
    %3 = scf.if %2 -> (memref<2xf32>) {
      scf.yield %0 : memref<2xf32>
    } else {
      scf.yield %iterBuf : memref<2xf32>
    }
    scf.yield %3 : memref<2xf32>
  }
  "linalg.copy"(%1, %res) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

//      CHECK: %[[ALLOC0:.*]] = alloc()
// CHECK-NEXT: %[[ALLOC1:.*]] = scf.for {{.*}} iter_args(%[[IALLOC:.*]] =
//      CHECK: %[[ALLOC2:.*]] = scf.if
//      CHECK: scf.yield %[[ALLOC0]]
//      CHECK: scf.yield %[[IALLOC]]
//      CHECK: scf.yield %[[ALLOC2]]
//      CHECK: linalg.copy(%[[ALLOC1]], %arg4)
//      CHECK: dealloc %[[ALLOC0]]

// -----

// Test Case: structured control-flow loop with a nested if operation using
// a deeply nested buffer allocation.
// Since the innermost allocation happens in a divergent branch, we have to
// introduce additional copies for the nested if operation. Since the loop's
// yield operation "returns" %3, it will return a newly allocated buffer.
// Therefore, we have to free the iteration argument %iterBuf before
// "returning" %3.

// CHECK-LABEL: func @loop_nested_if_alloc
func @loop_nested_if_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>) -> memref<2xf32> {
  %0 = alloc() : memref<2xf32>
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = cmpi "eq", %i, %ub : index
    %3 = scf.if %2 -> (memref<2xf32>) {
      %4 = alloc() : memref<2xf32>
      scf.yield %4 : memref<2xf32>
    } else {
      scf.yield %0 : memref<2xf32>
    }
    scf.yield %3 : memref<2xf32>
  }
  return %1 : memref<2xf32>
}

//      CHECK: %[[ALLOC0:.*]] = alloc()
//      CHECK: %[[ALLOC1:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%arg3, %[[ALLOC1]])
// CHECK-NEXT: %[[ALLOC2:.*]] = scf.for {{.*}} iter_args(%[[IALLOC:.*]] = %[[ALLOC1]]
//      CHECK: dealloc %[[IALLOC]]
//      CHECK: %[[ALLOC3:.*]] = scf.if

//      CHECK: %[[ALLOC4:.*]] = alloc()
// CHECK-NEXT: %[[ALLOC5:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC4]], %[[ALLOC5]])
// CHECK-NEXT: dealloc %[[ALLOC4]]
// CHECK-NEXT: scf.yield %[[ALLOC5]]

//      CHECK: %[[ALLOC6:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC6]])
// CHECK-NEXT: scf.yield %[[ALLOC6]]

//      CHECK: %[[ALLOC7:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC3:.*]], %[[ALLOC7]])
// CHECK-NEXT: dealloc %[[ALLOC3]]
// CHECK-NEXT: scf.yield %[[ALLOC7]]

//      CHECK: dealloc %[[ALLOC0]]
// CHECK-NEXT: return %[[ALLOC2]]

// -----

// Test Case: several nested structured control-flow loops with a deeply nested
// buffer allocation inside an if operation.
// Same behavior is an loop_nested_if_alloc: we have to insert deallocations
// before each yield in all loops recursively.

// CHECK-LABEL: func @loop_nested_alloc
func @loop_nested_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>,
  %res: memref<2xf32>) {
  %0 = alloc() : memref<2xf32>
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = scf.for %i2 = %lb to %ub step %step
      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
      %3 = scf.for %i3 = %lb to %ub step %step
        iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
        %4 = alloc() : memref<2xf32>
        %5 = cmpi "eq", %i, %ub : index
        %6 = scf.if %5 -> (memref<2xf32>) {
          %7 = alloc() : memref<2xf32>
          scf.yield %7 : memref<2xf32>
        } else {
          scf.yield %iterBuf3 : memref<2xf32>
        }
        scf.yield %6 : memref<2xf32>
      }
      scf.yield %3 : memref<2xf32>
    }
    scf.yield %2 : memref<2xf32>
  }
  "linalg.copy"(%1, %res) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

//      CHECK: %[[ALLOC0:.*]] = alloc()
// CHECK-NEXT: dealloc %[[ALLOC0]]
// CHECK-NEXT: %[[ALLOC1:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%arg3, %[[ALLOC1]])
// CHECK-NEXT: %[[VAL_7:.*]] = scf.for {{.*}} iter_args(%[[IALLOC0:.*]] = %[[ALLOC1]])
//      CHECK: %[[ALLOC2:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[IALLOC0]], %[[ALLOC2]])
// CHECK-NEXT: dealloc %[[IALLOC0]]
// CHECK-NEXT: %[[ALLOC3:.*]] = scf.for {{.*}} iter_args(%[[IALLOC1:.*]] = %[[ALLOC2]])
//      CHECK: %[[ALLOC5:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[IALLOC1]], %[[ALLOC5]])
// CHECK-NEXT: dealloc %[[IALLOC1]]

//      CHECK: %[[ALLOC6:.*]] = scf.for {{.*}} iter_args(%[[IALLOC2:.*]] = %[[ALLOC5]])
//      CHECK: %[[ALLOC8:.*]] = alloc()
// CHECK-NEXT: dealloc %[[ALLOC8]]
//      CHECK: %[[ALLOC9:.*]] = scf.if

//      CHECK: %[[ALLOC11:.*]] = alloc()
// CHECK-NEXT: %[[ALLOC12:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC11]], %[[ALLOC12]])
// CHECK-NEXT: dealloc %[[ALLOC11]]
// CHECK-NEXT: scf.yield %[[ALLOC12]]

//      CHECK: %[[ALLOC13:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[IALLOC2]], %[[ALLOC13]])
// CHECK-NEXT: scf.yield %[[ALLOC13]]

//      CHECK: dealloc %[[IALLOC2]]
// CHECK-NEXT: %[[ALLOC10:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC9]], %[[ALLOC10]])
// CHECK-NEXT: dealloc %[[ALLOC9]]
// CHECK-NEXT: scf.yield %[[ALLOC10]]

//      CHECK: %[[ALLOC7:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC6]], %[[ALLOC7]])
// CHECK-NEXT: dealloc %[[ALLOC6]]
// CHECK-NEXT: scf.yield %[[ALLOC7]]

//      CHECK: %[[ALLOC4:.*]] = alloc()
// CHECK-NEXT: linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
// CHECK-NEXT: dealloc %[[ALLOC3]]
// CHECK-NEXT: scf.yield %[[ALLOC4]]

//      CHECK: linalg.copy(%[[VAL_7]], %arg4)
// CHECK-NEXT: dealloc %[[VAL_7]]

// -----

// Test Case: explicit control-flow loop with a dynamically allocated buffer.
// The BufferPlacement transformation should fail on this explicit
// control-flow loop since they are not supported.

// CHECK-LABEL: func @loop_dynalloc
func @loop_dynalloc(
  %arg0 : i32,
  %arg1 : i32,
  %arg2: memref<?xf32>,
  %arg3: memref<?xf32>) {
  %const0 = constant 0 : i32
  br ^loopHeader(%const0, %arg2 : i32, memref<?xf32>)

^loopHeader(%i : i32, %buff : memref<?xf32>):
  %lessThan = cmpi "slt", %i, %arg1 : i32
  cond_br %lessThan,
    ^loopBody(%i, %buff : i32, memref<?xf32>),
    ^exit(%buff : memref<?xf32>)

^loopBody(%val : i32, %buff2: memref<?xf32>):
  %const1 = constant 1 : i32
  %inc = addi %val, %const1 : i32
  %size = std.index_cast %inc : i32 to index
  %alloc1 = alloc(%size) : memref<?xf32>
  br ^loopHeader(%inc, %alloc1 : i32, memref<?xf32>)

^exit(%buff3 : memref<?xf32>):
  "linalg.copy"(%buff3, %arg3) : (memref<?xf32>, memref<?xf32>) -> ()
  return
}

// expected-error@+1 {{Structured control-flow loops are supported only}}

// -----

// Test Case: explicit control-flow loop with a dynamically allocated buffer.
// The BufferPlacement transformation should fail on this explicit
// control-flow loop since they are not supported.

// CHECK-LABEL: func @do_loop_alloc
func @do_loop_alloc(
  %arg0 : i32,
  %arg1 : i32,
  %arg2: memref<2xf32>,
  %arg3: memref<2xf32>) {
  %const0 = constant 0 : i32
  br ^loopBody(%const0, %arg2 : i32, memref<2xf32>)

^loopBody(%val : i32, %buff2: memref<2xf32>):
  %const1 = constant 1 : i32
  %inc = addi %val, %const1 : i32
  %alloc1 = alloc() : memref<2xf32>
  br ^loopHeader(%inc, %alloc1 : i32, memref<2xf32>)

^loopHeader(%i : i32, %buff : memref<2xf32>):
  %lessThan = cmpi "slt", %i, %arg1 : i32
  cond_br %lessThan,
    ^loopBody(%i, %buff : i32, memref<2xf32>),
    ^exit(%buff : memref<2xf32>)

^exit(%buff3 : memref<2xf32>):
  "linalg.copy"(%buff3, %arg3) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// expected-error@+1 {{Structured control-flow loops are supported only}}

// -----

func @assumingOp(%arg0: !shape.witness, %arg2: memref<2xf32>, %arg3: memref<2xf32>) {
  // Confirm the alloc will be dealloc'ed in the block.
  %1 = shape.assuming %arg0 -> memref<2xf32> {
     %0 = alloc() : memref<2xf32>
    shape.assuming_yield %arg2 : memref<2xf32>
  }
  // Confirm the alloc will be returned and dealloc'ed after its use.
  %3 = shape.assuming %arg0 -> memref<2xf32> {
    %2 = alloc() : memref<2xf32>
    shape.assuming_yield %2 : memref<2xf32>
  }
  "linalg.copy"(%3, %arg3) : (memref<2xf32>, memref<2xf32>) -> ()
  return
}

// CHECK-LABEL: func @assumingOp(
// CHECK-SAME:     %[[ARG0:.*]]: !shape.witness,
// CHECK-SAME:     %[[ARG1:.*]]: memref<2xf32>,
// CHECK-SAME:     %[[ARG2:.*]]: memref<2xf32>) {
// CHECK:        %[[UNUSED_RESULT:.*]] = shape.assuming %[[ARG0]] -> (memref<2xf32>) {
// CHECK:         %[[ALLOC0:.*]] = alloc() : memref<2xf32>
// CHECK:         dealloc %[[ALLOC0]] : memref<2xf32>
// CHECK:         shape.assuming_yield %[[ARG1]] : memref<2xf32>
// CHECK:        }
// CHECK:        %[[ASSUMING_RESULT:.*]] = shape.assuming %[[ARG0]] -> (memref<2xf32>) {
// CHECK:         %[[TMP_ALLOC:.*]] = alloc() : memref<2xf32>
// CHECK:         %[[RETURNING_ALLOC:.*]] = alloc() : memref<2xf32>
// CHECK:         linalg.copy(%[[TMP_ALLOC]], %[[RETURNING_ALLOC]]) : memref<2xf32>, memref<2xf32>
// CHECK:         dealloc %[[TMP_ALLOC]] : memref<2xf32>
// CHECK:         shape.assuming_yield %[[RETURNING_ALLOC]] : memref<2xf32>
// CHECK:        }
// CHECK:        linalg.copy(%[[ASSUMING_RESULT:.*]], %[[ARG2]]) : memref<2xf32>, memref<2xf32>
// CHECK:        dealloc %[[ASSUMING_RESULT]] : memref<2xf32>
// CHECK:        return
// CHECK:       }