1// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false \
2// RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
3// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \
4// RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
5
6// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
7// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
8
9// Test Case: Existing AllocOp with no users.
10// BufferDeallocation expected behavior: It should insert a DeallocOp right
11// before ReturnOp.
12
13func.func private @emptyUsesValue(%arg0: memref<4xf32>) {
14  %0 = memref.alloc() : memref<4xf32>
15  "test.read_buffer"(%0) : (memref<4xf32>) -> ()
16  return
17}
18
19// CHECK-LABEL: func private @emptyUsesValue(
20//       CHECK: [[ALLOC:%.*]] = memref.alloc()
21//       CHECK: bufferization.dealloc ([[ALLOC]] :
22//  CHECK-SAME:   if (%true{{[0-9_]*}})
23//   CHECK-NOT:   retain
24//  CHECK-NEXT: return
25
26// CHECK-DYNAMIC-LABEL: func private @emptyUsesValue(
27//  CHECK-DYNAMIC-SAME:     [[ARG0:%.+]]: memref<4xf32>)
28//       CHECK-DYNAMIC: [[ALLOC:%.*]] = memref.alloc()
29//  CHECK-DYNAMIC-NEXT: "test.read_buffer"
30//  CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
31//   CHECK-DYNAMIC-NOT:   retain
32//  CHECK-DYNAMIC-NEXT: return
33
34// -----
35
36func.func @emptyUsesValue(%arg0: memref<4xf32>) {
37  %0 = memref.alloc() : memref<4xf32>
38  "test.read_buffer"(%0) : (memref<4xf32>) -> ()
39  return
40}
41
42// CHECK-LABEL: func @emptyUsesValue(
43
44// CHECK-DYNAMIC-LABEL: func @emptyUsesValue(
45//       CHECK-DYNAMIC: [[ALLOC:%.*]] = memref.alloc()
46//       CHECK-DYNAMIC: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
47//   CHECK-DYNAMIC-NOT:   retain
48//  CHECK-DYNAMIC-NEXT: return
49
50// -----
51
52// Test Case: Dead operations in a single block.
53// BufferDeallocation expected behavior: It only inserts the two missing
54// DeallocOps after the last BufferBasedOp.
55
56func.func private @redundantOperations(%arg0: memref<2xf32>) {
57  %0 = memref.alloc() : memref<2xf32>
58  test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
59  %1 = memref.alloc() : memref<2xf32>
60  test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
61  return
62}
63
64// CHECK-LABEL: func private @redundantOperations
65//      CHECK: (%[[ARG0:.*]]: {{.*}})
66//      CHECK: %[[FIRST_ALLOC:.*]] = memref.alloc()
67// CHECK-NEXT: test.buffer_based
68//      CHECK: %[[SECOND_ALLOC:.*]] = memref.alloc()
69// CHECK-NEXT: test.buffer_based
70// CHECK-NEXT: bufferization.dealloc (%[[FIRST_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
71// CHECK-NEXT: bufferization.dealloc (%[[SECOND_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
72// CHECK-NEXT: return
73
74// CHECK-DYNAMIC-LABEL: func private @redundantOperations
75//      CHECK-DYNAMIC: (%[[ARG0:.*]]: memref{{.*}})
76//      CHECK-DYNAMIC: %[[FIRST_ALLOC:.*]] = memref.alloc()
77// CHECK-DYNAMIC-NEXT: test.buffer_based
78//      CHECK-DYNAMIC: %[[SECOND_ALLOC:.*]] = memref.alloc()
79// CHECK-DYNAMIC-NEXT: test.buffer_based
80// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[FIRST_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
81// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[SECOND_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
82// CHECK-DYNAMIC-NEXT: return
83
84// -----
85
86// Test Case: buffer deallocation escaping
87// BufferDeallocation expected behavior: It must not dealloc %arg1 and %x
88// since they are operands of return operation and should escape from
89// deallocating. It should dealloc %y after CopyOp.
90
91func.func private @memref_in_function_results(
92  %arg0: memref<5xf32>,
93  %arg1: memref<10xf32>,
94  %arg2: memref<5xf32>) -> (memref<10xf32>, memref<15xf32>) {
95  %x = memref.alloc() : memref<15xf32>
96  %y = memref.alloc() : memref<5xf32>
97  test.buffer_based in(%arg0: memref<5xf32>) out(%y: memref<5xf32>)
98  test.copy(%y, %arg2) : (memref<5xf32>, memref<5xf32>)
99  return %arg1, %x : memref<10xf32>, memref<15xf32>
100}
101
102// CHECK-LABEL: func private @memref_in_function_results
103//       CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>,
104//  CHECK-SAME: %[[RESULT:.*]]: memref<5xf32>)
105//       CHECK: %[[X:.*]] = memref.alloc()
106//       CHECK: %[[Y:.*]] = memref.alloc()
107//       CHECK: test.copy
108//  CHECK-NEXT: %[[V0:.+]] = scf.if %false
109//  CHECK-NEXT:   scf.yield %[[ARG1]]
110//  CHECK-NEXT: } else {
111//  CHECK-NEXT:   %[[CLONE:.+]] = bufferization.clone %[[ARG1]]
112//  CHECK-NEXT:   scf.yield %[[CLONE]]
113//  CHECK-NEXT: }
114//       CHECK: bufferization.dealloc (%[[Y]] : {{.*}}) if (%true{{[0-9_]*}})
115//   CHECK-NOT: retain
116//       CHECK: return %[[V0]], %[[X]]
117
118// CHECK-DYNAMIC-LABEL: func private @memref_in_function_results
119//       CHECK-DYNAMIC: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>,
120//  CHECK-DYNAMIC-SAME: %[[RESULT:.*]]: memref<5xf32>)
121//       CHECK-DYNAMIC: %[[X:.*]] = memref.alloc()
122//       CHECK-DYNAMIC: %[[Y:.*]] = memref.alloc()
123//       CHECK-DYNAMIC: test.copy
124//       CHECK-DYNAMIC: bufferization.dealloc (%[[Y]] : {{.*}}) if (%true{{[0-9_]*}})
125//   CHECK-DYNAMIC-NOT: retain
126//       CHECK-DYNAMIC: return %[[ARG1]], %[[X]], %false, %true
127
128// -----
129
130// CHECK-DYNAMIC-LABEL: func private @private_callee(
131//  CHECK-DYNAMIC-SAME:     %[[arg0:.*]]: memref<f32>) -> (memref<f32>, i1)
132//       CHECK-DYNAMIC:   %[[true:.*]] = arith.constant true
133//       CHECK-DYNAMIC:   %[[alloc:.*]] = memref.alloc() : memref<f32>
134//   CHECK-DYNAMIC-NOT:   bufferization.dealloc
135//       CHECK-DYNAMIC:   return %[[alloc]], %[[true]]
136func.func private @private_callee(%arg0: memref<f32>) -> memref<f32> {
137  %alloc = memref.alloc() : memref<f32>
138  return %alloc : memref<f32>
139}
140
141//       CHECK-DYNAMIC: func @caller() -> f32
142//       CHECK-DYNAMIC:   %[[true:.*]] = arith.constant true
143//       CHECK-DYNAMIC:   %[[alloc:.*]] = memref.alloc() : memref<f32>
144//       CHECK-DYNAMIC:   %[[call:.*]]:2 = call @private_callee(%[[alloc]])
145//       CHECK-DYNAMIC:   memref.load
146//       CHECK-DYNAMIC:   %[[base:.*]], %[[offset:.*]] = memref.extract_strided_metadata %[[call]]#0
147//       CHECK-DYNAMIC:   bufferization.dealloc (%[[alloc]], %[[base]] : {{.*}}) if (%[[true]], %[[call]]#1)
148//   CHECK-DYNAMIC-NOT:   retain
149func.func @caller() -> (f32) {
150  %alloc = memref.alloc() : memref<f32>
151  %ret = call @private_callee(%alloc) : (memref<f32>) -> memref<f32>
152
153  %val = memref.load %ret[] : memref<f32>
154  return %val : f32
155}
156