xref: /llvm-project/mlir/test/Dialect/Bufferization/Transforms/lower-deallocations.mlir (revision 662c6fc74c1ae7fdefd27524dfaeb7f3e9a1f553)
1// RUN: mlir-opt -verify-diagnostics -bufferization-lower-deallocations -split-input-file %s | FileCheck %s
2
3// CHECK-LABEL: func @conversion_dealloc_empty
4func.func @conversion_dealloc_empty() {
5  // CHECK-NOT: bufferization.dealloc
6  bufferization.dealloc
7  return
8}
9
10// -----
11
12func.func @conversion_dealloc_empty_but_retains(%arg0: memref<2xi32>, %arg1: memref<2xi32>) -> (i1, i1) {
13  %0:2 = bufferization.dealloc retain (%arg0, %arg1 : memref<2xi32>, memref<2xi32>)
14  return %0#0, %0#1 : i1, i1
15}
16
17// CHECK-LABEL: func @conversion_dealloc_empty
18//  CHECK-NEXT: [[FALSE:%.+]] = arith.constant false
19//  CHECK-NEXT: return [[FALSE]], [[FALSE]] :
20
21// -----
22
23// CHECK-NOT: func @deallocHelper
24// CHECK-LABEL: func @conversion_dealloc_simple
25// CHECK-SAME: [[ARG0:%.+]]: memref<2xf32>
26// CHECK-SAME: [[ARG1:%.+]]: i1
27func.func @conversion_dealloc_simple(%arg0: memref<2xf32>, %arg1: i1) {
28  bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
29  return
30}
31
32//      CHECK: scf.if [[ARG1]] {
33// CHECK-NEXT:   memref.dealloc [[ARG0]] : memref<2xf32>
34// CHECK-NEXT: }
35// CHECK-NEXT: return
36
37// -----
38
39func.func @conversion_dealloc_one_memref_and_multiple_retained(%arg0: memref<2xf32>, %arg1: memref<1xf32>, %arg2: i1, %arg3: memref<2xf32>) -> (i1, i1) {
40  %0:2 = bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg2) retain (%arg1, %arg3 : memref<1xf32>, memref<2xf32>)
41  return %0#0, %0#1 : i1, i1
42}
43
44// CHECK-LABEL: func @conversion_dealloc_one_memref_and_multiple_retained
45//  CHECK-SAME: ([[ARG0:%.+]]: memref<2xf32>, [[ARG1:%.+]]: memref<1xf32>, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xf32>)
46//   CHECK-DAG: [[M0:%.+]] = memref.extract_aligned_pointer_as_index [[ARG0]]
47//   CHECK-DAG: [[R0:%.+]] = memref.extract_aligned_pointer_as_index [[ARG1]]
48//   CHECK-DAG: [[R1:%.+]] = memref.extract_aligned_pointer_as_index [[ARG3]]
49//   CHECK-DAG: [[DOES_NOT_ALIAS_R0:%.+]] = arith.cmpi ne, [[M0]], [[R0]] : index
50//   CHECK-DAG: [[DOES_NOT_ALIAS_R1:%.+]] = arith.cmpi ne, [[M0]], [[R1]] : index
51//       CHECK: [[NOT_RETAINED:%.+]] = arith.andi [[DOES_NOT_ALIAS_R0]], [[DOES_NOT_ALIAS_R1]]
52//       CHECK: [[SHOULD_DEALLOC:%.+]] = arith.andi [[NOT_RETAINED]], [[ARG2]]
53//       CHECK: scf.if [[SHOULD_DEALLOC]]
54//       CHECK:   memref.dealloc [[ARG0]]
55//       CHECK: }
56//   CHECK-DAG: [[ALIASES_R0:%.+]] = arith.xori [[DOES_NOT_ALIAS_R0]], %true
57//   CHECK-DAG: [[ALIASES_R1:%.+]] = arith.xori [[DOES_NOT_ALIAS_R1]], %true
58//   CHECK-DAG: [[RES0:%.+]] = arith.andi [[ALIASES_R0]], [[ARG2]]
59//   CHECK-DAG: [[RES1:%.+]] = arith.andi [[ALIASES_R1]], [[ARG2]]
60//       CHECK: return [[RES0]], [[RES1]]
61
62// CHECK-NOT: func @dealloc_helper
63
64// -----
65
66func.func @conversion_dealloc_multiple_memrefs_and_retained(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1, %arg5: memref<2xf32>) -> (i1, i1) {
67  %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2, %arg5 : memref<1xf32>, memref<2xf32>)
68  return %0#0, %0#1 : i1, i1
69}
70
71// CHECK-LABEL: func @conversion_dealloc_multiple_memrefs_and_retained
72// CHECK-SAME: ([[ARG0:%.+]]: memref<2xf32>, [[ARG1:%.+]]: memref<5xf32>,
73// CHECK-SAME: [[ARG2:%.+]]: memref<1xf32>, [[ARG3:%.+]]: i1, [[ARG4:%.+]]: i1,
74// CHECK-SAME: [[ARG5:%.+]]: memref<2xf32>)
75//      CHECK: [[TO_DEALLOC_MR:%.+]] = memref.alloc() : memref<2xindex>
76//      CHECK: [[CONDS:%.+]] = memref.alloc() : memref<2xi1>
77//      CHECK: [[TO_RETAIN_MR:%.+]] = memref.alloc() : memref<2xindex>
78//  CHECK-DAG: [[V0:%.+]] = memref.extract_aligned_pointer_as_index [[ARG0]]
79//  CHECK-DAG: [[C0:%.+]] = arith.constant 0 : index
80//  CHECK-DAG: memref.store [[V0]], [[TO_DEALLOC_MR]][[[C0]]]
81//  CHECK-DAG: [[V1:%.+]] = memref.extract_aligned_pointer_as_index [[ARG1]]
82//  CHECK-DAG: [[C1:%.+]] = arith.constant 1 : index
83//  CHECK-DAG: memref.store [[V1]], [[TO_DEALLOC_MR]][[[C1]]]
84//  CHECK-DAG: [[C0:%.+]] = arith.constant 0 : index
85//  CHECK-DAG: memref.store [[ARG3]], [[CONDS]][[[C0]]]
86//  CHECK-DAG: [[C1:%.+]] = arith.constant 1 : index
87//  CHECK-DAG: memref.store [[ARG4]], [[CONDS]][[[C1]]]
88//  CHECK-DAG: [[V2:%.+]] = memref.extract_aligned_pointer_as_index [[ARG2]]
89//  CHECK-DAG: [[C0:%.+]] = arith.constant 0 : index
90//  CHECK-DAG: memref.store [[V2]], [[TO_RETAIN_MR]][[[C0]]]
91//  CHECK-DAG: [[V3:%.+]] = memref.extract_aligned_pointer_as_index [[ARG5]]
92//  CHECK-DAG: [[C1:%.+]] = arith.constant 1 : index
93//  CHECK-DAG: memref.store [[V3]], [[TO_RETAIN_MR]][[[C1]]]
94//  CHECK-DAG: [[CAST_DEALLOC:%.+]] = memref.cast [[TO_DEALLOC_MR]] : memref<2xindex> to memref<?xindex>
95//  CHECK-DAG: [[CAST_CONDS:%.+]] = memref.cast [[CONDS]] : memref<2xi1> to memref<?xi1>
96//  CHECK-DAG: [[CAST_RETAIN:%.+]] = memref.cast [[TO_RETAIN_MR]] : memref<2xindex> to memref<?xindex>
97//      CHECK: [[DEALLOC_CONDS:%.+]] = memref.alloc() : memref<2xi1>
98//      CHECK: [[RETAIN_CONDS:%.+]] = memref.alloc() : memref<2xi1>
99//      CHECK: [[CAST_DEALLOC_CONDS:%.+]] = memref.cast [[DEALLOC_CONDS]] : memref<2xi1> to memref<?xi1>
100//      CHECK: [[CAST_RETAIN_CONDS:%.+]] = memref.cast [[RETAIN_CONDS]] : memref<2xi1> to memref<?xi1>
101//      CHECK: call @dealloc_helper([[CAST_DEALLOC]], [[CAST_RETAIN]], [[CAST_CONDS]], [[CAST_DEALLOC_CONDS]], [[CAST_RETAIN_CONDS]])
102//      CHECK: [[C0:%.+]] = arith.constant 0 : index
103//      CHECK: [[SHOULD_DEALLOC_0:%.+]] = memref.load [[DEALLOC_CONDS]][[[C0]]]
104//      CHECK: scf.if [[SHOULD_DEALLOC_0]] {
105//      CHECK:   memref.dealloc %arg0
106//      CHECK: }
107//      CHECK: [[C1:%.+]] = arith.constant 1 : index
108//      CHECK: [[SHOULD_DEALLOC_1:%.+]] = memref.load [[DEALLOC_CONDS]][[[C1]]]
109//      CHECK: scf.if [[SHOULD_DEALLOC_1]]
110//      CHECK:   memref.dealloc [[ARG1]]
111//      CHECK: }
112//      CHECK: [[C0:%.+]] = arith.constant 0 : index
113//      CHECK: [[OWNERSHIP0:%.+]] = memref.load [[RETAIN_CONDS]][[[C0]]]
114//      CHECK: [[C1:%.+]] = arith.constant 1 : index
115//      CHECK: [[OWNERSHIP1:%.+]] = memref.load [[RETAIN_CONDS]][[[C1]]]
116//      CHECK: memref.dealloc [[TO_DEALLOC_MR]]
117//      CHECK: memref.dealloc [[TO_RETAIN_MR]]
118//      CHECK: memref.dealloc [[CONDS]]
119//      CHECK: memref.dealloc [[DEALLOC_CONDS]]
120//      CHECK: memref.dealloc [[RETAIN_CONDS]]
121//      CHECK: return [[OWNERSHIP0]], [[OWNERSHIP1]]
122
123//      CHECK: func private @dealloc_helper
124// CHECK-SAME: ([[TO_DEALLOC_MR:%.+]]: memref<?xindex>, [[TO_RETAIN_MR:%.+]]: memref<?xindex>,
125// CHECK-SAME: [[CONDS:%.+]]: memref<?xi1>, [[DEALLOC_CONDS_OUT:%.+]]: memref<?xi1>,
126// CHECK-SAME: [[RETAIN_CONDS_OUT:%.+]]: memref<?xi1>)
127//      CHECK:   [[TO_DEALLOC_SIZE:%.+]] = memref.dim [[TO_DEALLOC_MR]], %c0
128//      CHECK:   [[TO_RETAIN_SIZE:%.+]] = memref.dim [[TO_RETAIN_MR]], %c0
129//      CHECK:   scf.for [[ITER:%.+]] = %c0 to [[TO_RETAIN_SIZE]] step %c1 {
130// CHECK-NEXT:     memref.store %false, [[RETAIN_CONDS_OUT]][[[ITER]]]
131// CHECK-NEXT:   }
132//      CHECK:   scf.for [[OUTER_ITER:%.+]] = %c0 to [[TO_DEALLOC_SIZE]] step %c1 {
133//      CHECK:     [[TO_DEALLOC:%.+]] = memref.load [[TO_DEALLOC_MR]][[[OUTER_ITER]]]
134// CHECK-NEXT:     [[COND:%.+]] = memref.load [[CONDS]][[[OUTER_ITER]]]
135// CHECK-NEXT:     [[NO_RETAIN_ALIAS:%.+]] = scf.for [[ITER:%.+]] = %c0 to [[TO_RETAIN_SIZE]] step %c1 iter_args([[ITER_ARG:%.+]] = %true) -> (i1) {
136// CHECK-NEXT:       [[RETAIN_VAL:%.+]] = memref.load [[TO_RETAIN_MR]][[[ITER]]] : memref<?xindex>
137// CHECK-NEXT:       [[DOES_ALIAS:%.+]] = arith.cmpi eq, [[RETAIN_VAL]], [[TO_DEALLOC]] : index
138// CHECK-NEXT:       scf.if [[DOES_ALIAS]]
139// CHECK-NEXT:         [[RETAIN_COND:%.+]] = memref.load [[RETAIN_CONDS_OUT]][[[ITER]]]
140// CHECK-NEXT:         [[AGG_RETAIN_COND:%.+]] = arith.ori [[RETAIN_COND]], [[COND]] : i1
141// CHECK-NEXT:         memref.store [[AGG_RETAIN_COND]], [[RETAIN_CONDS_OUT]][[[ITER]]]
142// CHECK-NEXT:       }
143// CHECK-NEXT:       [[DOES_NOT_ALIAS:%.+]] = arith.cmpi ne, [[RETAIN_VAL]], [[TO_DEALLOC]] : index
144// CHECK-NEXT:       [[AGG_DOES_NOT_ALIAS:%.+]] = arith.andi [[ITER_ARG]], [[DOES_NOT_ALIAS]] : i1
145// CHECK-NEXT:       scf.yield [[AGG_DOES_NOT_ALIAS]] : i1
146// CHECK-NEXT:     }
147// CHECK-NEXT:     [[SHOULD_DEALLOC:%.+]] = scf.for [[ITER:%.+]] = %c0 to [[OUTER_ITER]] step %c1 iter_args([[ITER_ARG:%.+]] = [[NO_RETAIN_ALIAS]]) -> (i1) {
148// CHECK-NEXT:       [[OTHER_DEALLOC_VAL:%.+]] = memref.load [[ARG0]][[[ITER]]] : memref<?xindex>
149// CHECK-NEXT:       [[DOES_ALIAS:%.+]] = arith.cmpi ne, [[OTHER_DEALLOC_VAL]], [[TO_DEALLOC]] : index
150// CHECK-NEXT:       [[AGG_DOES_ALIAS:%.+]] = arith.andi [[ITER_ARG]], [[DOES_ALIAS]] : i1
151// CHECK-NEXT:       scf.yield [[AGG_DOES_ALIAS]] : i1
152// CHECK-NEXT:     }
153// CHECK-NEXT:     [[DEALLOC_COND:%.+]] = arith.andi [[SHOULD_DEALLOC]], [[COND]] : i1
154// CHECK-NEXT:     memref.store [[DEALLOC_COND]], [[DEALLOC_CONDS_OUT]][[[OUTER_ITER]]]
155// CHECK-NEXT:   }
156// CHECK-NEXT:   return
157
158// -----
159
160// This test check dealloc_helper function is generated on each nested symbol
161// table operation when needed and only generated once.
162module @conversion_nest_module_dealloc_helper {
163  func.func @top_level_func(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1, %arg5: memref<2xf32>) -> (i1, i1) {
164    %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2, %arg5 : memref<1xf32>, memref<2xf32>)
165    func.return %0#0, %0#1 : i1, i1
166  }
167  module @nested_module_not_need_dealloc_helper {
168    func.func @nested_module_not_need_dealloc_helper_func(%arg0: memref<2xf32>, %arg1: memref<1xf32>, %arg2: i1, %arg3: memref<2xf32>) -> (i1, i1) {
169      %0:2 = bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg2) retain (%arg1, %arg3 : memref<1xf32>, memref<2xf32>)
170      return %0#0, %0#1 : i1, i1
171    }
172  }
173  module @nested_module_need_dealloc_helper {
174    func.func @nested_module_need_dealloc_helper_func0(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1, %arg5: memref<2xf32>) -> (i1, i1) {
175      %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2, %arg5 : memref<1xf32>, memref<2xf32>)
176      func.return %0#0, %0#1 : i1, i1
177    }
178    func.func @nested_module_need_dealloc_helper_func1(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1, %arg5: memref<2xf32>) -> (i1, i1) {
179      %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2, %arg5 : memref<1xf32>, memref<2xf32>)
180      func.return %0#0, %0#1 : i1, i1
181    }
182  }
183}
184
185// CHECK:     module @conversion_nest_module_dealloc_helper {
186// CHECK:       func.func @top_level_func
187// CHECK:         call @dealloc_helper
188// CHECK:       module @nested_module_not_need_dealloc_helper {
189// CHECK:         func.func @nested_module_not_need_dealloc_helper_func
190// CHECK-NOT:       @dealloc_helper
191// CHECK:       module @nested_module_need_dealloc_helper {
192// CHECK:         func.func @nested_module_need_dealloc_helper_func0
193// CHECK:           call @dealloc_helper
194// CHECK:         func.func @nested_module_need_dealloc_helper_func1
195// CHECK:           call @dealloc_helper
196// CHECK:         func.func private @dealloc_helper
197// CHECK:       func.func private @dealloc_helper
198