1// RUN: mlir-opt -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 -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
9func.func private @f(%arg0: memref<f64>) -> memref<f64> {
10  return %arg0 : memref<f64>
11}
12
13func.func @function_call() {
14  %alloc = memref.alloc() : memref<f64>
15  %alloc2 = memref.alloc() : memref<f64>
16  %ret = call @f(%alloc) : (memref<f64>) -> memref<f64>
17  test.copy(%ret, %alloc2) : (memref<f64>, memref<f64>)
18  return
19}
20
21// CHECK-LABEL: func @function_call()
22//       CHECK: [[ALLOC0:%.+]] = memref.alloc(
23//  CHECK-NEXT: [[ALLOC1:%.+]] = memref.alloc(
24//  CHECK-NEXT: [[RET:%.+]] = call @f([[ALLOC0]]) : (memref<f64>) -> memref<f64>
25//  CHECK-NEXT: test.copy
26//  CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]
27// COM: the following dealloc operation should be split into three since we can
28// COM: be sure that the memrefs will never alias according to the buffer
29// COM: deallocation ABI, however, the local alias analysis is not powerful
30// COM: enough to detect this yet.
31//  CHECK-NEXT: bufferization.dealloc ([[ALLOC0]], [[ALLOC1]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, %true{{[0-9_]*}}, %true{{[0-9_]*}})
32
33// CHECK-DYNAMIC-LABEL: func @function_call()
34//       CHECK-DYNAMIC: [[ALLOC0:%.+]] = memref.alloc(
35//  CHECK-DYNAMIC-NEXT: [[ALLOC1:%.+]] = memref.alloc(
36//  CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[ALLOC0]]) : (memref<f64>) -> (memref<f64>, i1)
37//  CHECK-DYNAMIC-NEXT: test.copy
38//  CHECK-DYNAMIC-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
39//  CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC0]], [[ALLOC1]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, %true{{[0-9_]*}}, [[RET]]#1)
40
41// -----
42
43func.func @f(%arg0: memref<f64>) -> memref<f64> {
44  return %arg0 : memref<f64>
45}
46
47func.func @function_call_non_private() {
48  %alloc = memref.alloc() : memref<f64>
49  %alloc2 = memref.alloc() : memref<f64>
50  %ret = call @f(%alloc) : (memref<f64>) -> memref<f64>
51  test.copy(%ret, %alloc2) : (memref<f64>, memref<f64>)
52  return
53}
54
55// CHECK-LABEL: func @function_call_non_private
56//       CHECK: [[ALLOC0:%.+]] = memref.alloc(
57//       CHECK: [[ALLOC1:%.+]] = memref.alloc(
58//       CHECK: [[RET:%.+]] = call @f([[ALLOC0]]) : (memref<f64>) -> memref<f64>
59//  CHECK-NEXT: test.copy
60//  CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]
61//  CHECK-NEXT: bufferization.dealloc ([[ALLOC0]], [[ALLOC1]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, %true{{[0-9_]*}}, %true{{[0-9_]*}})
62//  CHECK-NEXT: return
63
64// CHECK-DYNAMIC-LABEL: func @function_call_non_private
65//       CHECK-DYNAMIC: [[ALLOC0:%.+]] = memref.alloc(
66//       CHECK-DYNAMIC: [[ALLOC1:%.+]] = memref.alloc(
67//       CHECK-DYNAMIC: [[RET:%.+]] = call @f([[ALLOC0]]) : (memref<f64>) -> memref<f64>
68//  CHECK-DYNAMIC-NEXT: test.copy
69//  CHECK-DYNAMIC-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]
70//  CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC0]], [[ALLOC1]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, %true{{[0-9_]*}}, %true{{[0-9_]*}})
71//  CHECK-DYNAMIC-NEXT: return
72
73// -----
74
75func.func private @f(%arg0: memref<f64>) -> memref<f64> {
76  return %arg0 : memref<f64>
77}
78
79func.func @function_call_requries_merged_ownership_mid_block(%arg0: i1) {
80  %alloc = memref.alloc() : memref<f64>
81  %alloc2 = memref.alloca() : memref<f64>
82  %0 = arith.select %arg0, %alloc, %alloc2 : memref<f64>
83  %ret = call @f(%0) : (memref<f64>) -> memref<f64>
84  test.copy(%ret, %alloc) : (memref<f64>, memref<f64>)
85  return
86}
87
88// CHECK-LABEL: func @function_call_requries_merged_ownership_mid_block
89//       CHECK:   [[ALLOC0:%.+]] = memref.alloc(
90//  CHECK-NEXT:   [[ALLOC1:%.+]] = memref.alloca(
91//  CHECK-NEXT:   [[SELECT:%.+]] = arith.select{{.*}}[[ALLOC0]], [[ALLOC1]]
92//  CHECK-NEXT:   [[RET:%.+]] = call @f([[SELECT]])
93//  CHECK-NEXT:   test.copy
94//  CHECK-NEXT:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]
95//  CHECK-NEXT:   bufferization.dealloc ([[ALLOC0]], [[BASE]] :
96//  CHECK-SAME:     if (%true{{[0-9_]*}}, %true{{[0-9_]*}})
97//   CHECK-NOT:     retain
98//  CHECK-NEXT:   return
99
100// CHECK-DYNAMIC-LABEL: func @function_call_requries_merged_ownership_mid_block
101//  CHECK-DYNAMIC-SAME: ([[ARG0:%.+]]: i1)
102//       CHECK-DYNAMIC:   [[ALLOC0:%.+]] = memref.alloc(
103//  CHECK-DYNAMIC-NEXT:   [[ALLOC1:%.+]] = memref.alloca(
104//  CHECK-DYNAMIC-NEXT:   [[SELECT:%.+]] = arith.select [[ARG0]], [[ALLOC0]], [[ALLOC1]]
105//  CHECK-DYNAMIC-NEXT:   [[RET:%.+]]:2 = call @f([[SELECT]])
106//  CHECK-DYNAMIC-NEXT:   test.copy
107//  CHECK-DYNAMIC-NEXT:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
108//  CHECK-DYNAMIC-NEXT:   bufferization.dealloc ([[ALLOC0]], [[BASE]] :
109//  CHECK-DYNAMIC-SAME:     if (%true{{[0-9_]*}}, [[RET]]#1)
110//   CHECK-DYNAMIC-NOT:     retain
111//  CHECK-DYNAMIC-NEXT:   return
112
113// TODO: the inserted clone is not necessary, we just have to know which of the
114// two allocations was selected, either by checking aliasing of the result at
115// runtime or by extracting the select condition using an OpInterface or by
116// hardcoding the select op
117
118// -----
119
120func.func private @f(memref<f32>) -> memref<f32>
121
122func.func @g(%arg0: memref<f32>) -> memref<f32> {
123  %0 = call @f(%arg0) : (memref<f32>) -> memref<f32>
124  return %0 : memref<f32>
125}
126
127// CHECK-LABEL: func private @f
128//  CHECK-SAME: (memref<f32>) -> memref<f32>
129//       CHECK: call @f({{.*}}) : (memref<f32>) -> memref<f32>
130
131// CHECK-DYNAMIC-LABEL: func private @f
132//  CHECK-DYNAMIC-SAME: (memref<f32>) -> memref<f32>
133//       CHECK-DYNAMIC: call @f({{.*}}) : (memref<f32>) -> memref<f32>
134
135// -----
136
137func.func @func_call_indirect(%m: memref<?xf32>, %f: (memref<?xf32>) -> (memref<?xf32>)) {
138  %0 = func.call_indirect %f(%m) : (memref<?xf32>) -> (memref<?xf32>)
139  return
140}
141
142// CHECK-LABEL: func @func_call_indirect(
143//       CHECK:   %[[true:.*]] = arith.constant true
144//       CHECK:   %[[call:.*]] = call_indirect {{.*}} : (memref<?xf32>) -> memref<?xf32>
145//       CHECK:   %[[base_call:.*]], %{{.*}}, %{{.*}}, %{{.*}} = memref.extract_strided_metadata %[[call]]
146//       CHECK:   bufferization.dealloc (%[[base_call]] : {{.*}}) if (%[[true]])
147
148// CHECK-DYNAMIC-LABEL: func @func_call_indirect(
149//       CHECK-DYNAMIC:   %[[true:.*]] = arith.constant true
150//       CHECK-DYNAMIC:   %[[call:.*]] = call_indirect {{.*}} : (memref<?xf32>) -> memref<?xf32>
151//       CHECK-DYNAMIC:   %[[base_call:.*]], %{{.*}}, %{{.*}}, %{{.*}} = memref.extract_strided_metadata %[[call]]
152//       CHECK-DYNAMIC:   bufferization.dealloc (%[[base_call]] : {{.*}}) if (%[[true]])
153