1// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \
2// RUN:   --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
3// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
4
5// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
6
7// Test Case: Dead operations in a single block.
8// BufferDeallocation expected behavior: It only inserts the two missing
9// DeallocOps after the last BufferBasedOp.
10
11// CHECK-LABEL: func @redundantOperations
12func.func @redundantOperations(%arg0: memref<2xf32>) {
13  %0 = memref.alloc() : memref<2xf32>
14  test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
15  %1 = memref.alloc() : memref<2xf32>
16  test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
17  return
18}
19
20//      CHECK: (%[[ARG0:.*]]: {{.*}})
21//      CHECK: %[[FIRST_ALLOC:.*]] = memref.alloc()
22//  CHECK-NOT: bufferization.dealloc
23//      CHECK: test.buffer_based in(%[[ARG0]]{{.*}}out(%[[FIRST_ALLOC]]
24//  CHECK-NOT: bufferization.dealloc
25//      CHECK: %[[SECOND_ALLOC:.*]] = memref.alloc()
26//  CHECK-NOT: bufferization.dealloc
27//      CHECK: test.buffer_based in(%[[FIRST_ALLOC]]{{.*}}out(%[[SECOND_ALLOC]]
28//      CHECK: bufferization.dealloc (%[[FIRST_ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
29//      CHECK: bufferization.dealloc (%[[SECOND_ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
30// CHECK-NEXT: return
31
32// TODO: The dealloc could be split in two to avoid runtime aliasing checks
33// since we can be sure at compile time that they will never alias.
34
35// -----
36
37// CHECK-LABEL: func @allocaIsNotDeallocated
38func.func @allocaIsNotDeallocated(%arg0: memref<2xf32>) {
39  %0 = memref.alloc() : memref<2xf32>
40  test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
41  %1 = memref.alloca() : memref<2xf32>
42  test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
43  return
44}
45
46//      CHECK: (%[[ARG0:.*]]: {{.*}})
47//      CHECK: %[[FIRST_ALLOC:.*]] = memref.alloc()
48// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}}out(%[[FIRST_ALLOC]]
49// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = memref.alloca()
50// CHECK-NEXT: test.buffer_based in(%[[FIRST_ALLOC]]{{.*}}out(%[[SECOND_ALLOC]]
51//      CHECK: bufferization.dealloc (%[[FIRST_ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
52// CHECK-NEXT: return
53
54// -----
55
56// Test Case: Inserting missing DeallocOp in a single block.
57
58// CHECK-LABEL: func @inserting_missing_dealloc_simple
59func.func @inserting_missing_dealloc_simple(
60  %arg0 : memref<2xf32>,
61  %arg1: memref<2xf32>) {
62  %0 = memref.alloc() : memref<2xf32>
63  test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
64  test.copy(%0, %arg1) : (memref<2xf32>, memref<2xf32>)
65  return
66}
67
68//      CHECK: %[[ALLOC0:.*]] = memref.alloc()
69//      CHECK: test.copy
70//      CHECK: bufferization.dealloc (%[[ALLOC0]] :{{.*}}) if (%true{{[0-9_]*}})
71
72// -----
73
74// Test Case: The ownership indicator is set to false for alloca
75
76// CHECK-LABEL: func @alloca_ownership_indicator_is_false
77func.func @alloca_ownership_indicator_is_false() {
78  %0 = memref.alloca() : memref<2xf32>
79  cf.br ^bb1(%0: memref<2xf32>)
80^bb1(%arg0 : memref<2xf32>):
81  return
82}
83
84//      CHECK:  %[[ALLOC0:.*]] = memref.alloca()
85// CHECK-NEXT:   cf.br ^bb1(%[[ALLOC0]], %false :
86// CHECK-NEXT: ^bb1([[A0:%.+]]: memref<2xf32>, [[COND0:%.+]]: i1):
87//      CHECK:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[A0]]
88//      CHECK:   bufferization.dealloc ([[BASE]] : {{.*}}) if ([[COND0]])
89// CHECK-NEXT:   return
90
91// -----
92
93func.func @dealloc_existing_clones(%arg0: memref<?x?xf64>, %arg1: memref<?x?xf64>) -> memref<?x?xf64> {
94  %0 = bufferization.clone %arg0 : memref<?x?xf64> to memref<?x?xf64>
95  %1 = bufferization.clone %arg1 : memref<?x?xf64> to memref<?x?xf64>
96  return %0 : memref<?x?xf64>
97}
98
99// CHECK-LABEL: func @dealloc_existing_clones
100//       CHECK: (%[[ARG0:.*]]: memref<?x?xf64>, %[[ARG1:.*]]: memref<?x?xf64>)
101//       CHECK: %[[RES0:.*]] = bufferization.clone %[[ARG0]]
102//       CHECK: %[[RES1:.*]] = bufferization.clone %[[ARG1]]
103//  CHECK-NEXT: bufferization.dealloc (%[[RES1]] :{{.*}}) if (%true{{[0-9_]*}})
104//   CHECK-NOT: retain
105//  CHECK-NEXT: return %[[RES0]]
106
107// TODO: The retain operand could be dropped to avoid runtime aliasing checks
108// since We can guarantee at compile-time that it will never alias with the
109// dealloc operand
110
111// -----
112
113memref.global "private" constant @__constant_4xf32 : memref<4xf32> = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]>
114
115func.func @op_without_aliasing_and_allocation() -> memref<4xf32> {
116  %0 = memref.get_global @__constant_4xf32 : memref<4xf32>
117  return %0 : memref<4xf32>
118}
119
120// CHECK-LABEL: func @op_without_aliasing_and_allocation
121//       CHECK:   [[GLOBAL:%.+]] = memref.get_global @__constant_4xf32
122//       CHECK:   [[RES:%.+]] = scf.if %false
123//       CHECK:     scf.yield [[GLOBAL]] :
124//       CHECK:     [[CLONE:%.+]] = bufferization.clone [[GLOBAL]]
125//       CHECK:     scf.yield [[CLONE]] :
126//       CHECK:   return [[RES]] :
127
128// -----
129
130// Allocations with "bufferization.manual_deallocation" are assigned an
131// ownership of "false".
132
133func.func @manual_deallocation(%c: i1, %f: f32, %idx: index) -> f32 {
134  %0 = memref.alloc() {bufferization.manual_deallocation} : memref<5xf32>
135  linalg.fill ins(%f : f32) outs(%0 : memref<5xf32>)
136  %1 = memref.alloc() : memref<5xf32>
137  linalg.fill ins(%f : f32) outs(%1 : memref<5xf32>)
138  %2 = arith.select %c, %0, %1 : memref<5xf32>
139  %3 = memref.load %2[%idx] : memref<5xf32>
140
141  // Only buffers that are under "manual deallocation" are allowed to be
142  // deallocated with memref.dealloc. For consistency reasons, the
143  // manual_deallocation attribute must also be specified. A runtime insertion
144  // is inserted to ensure that we do not have ownership. (This is not a
145  // bulletproof check, but covers some cases of invalid IR.)
146  memref.dealloc %0 {bufferization.manual_deallocation} : memref<5xf32>
147
148  return %3 : f32
149}
150
151// CHECK-LABEL: func @manual_deallocation(
152//       CHECK:   %[[true:.*]] = arith.constant true
153//       CHECK:   %[[manual_alloc:.*]] = memref.alloc() {bufferization.manual_deallocation} : memref<5xf32>
154//       CHECK:   %[[managed_alloc:.*]] = memref.alloc() : memref<5xf32>
155//       CHECK:   %[[selected:.*]] = arith.select
156//       CHECK:   cf.assert %[[true]], "expected that the block does not have ownership"
157//       CHECK:   memref.dealloc %[[manual_alloc]]
158//       CHECK:   bufferization.dealloc (%[[managed_alloc]] : memref<5xf32>) if (%[[true]])
159